709 lines
29 KiB
Python
709 lines
29 KiB
Python
import os
|
||
import pandas as pd
|
||
import numpy as np
|
||
import json
|
||
import plotly.graph_objects as go
|
||
import plotly.utils
|
||
from flask import Flask, render_template, request, jsonify
|
||
from flask_cors import CORS
|
||
import sys
|
||
import warnings
|
||
import datetime
|
||
warnings.filterwarnings('ignore')
|
||
|
||
# 添加项目根目录到路径
|
||
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
|
||
try:
|
||
from model import Kronos, KronosTokenizer, KronosPredictor
|
||
MODEL_AVAILABLE = True
|
||
except ImportError:
|
||
MODEL_AVAILABLE = False
|
||
print("警告: Kronos模型无法导入，将使用模拟数据进行演示")
|
||
|
||
app = Flask(__name__)
|
||
CORS(app)
|
||
|
||
# 全局变量存储模型
|
||
tokenizer = None
|
||
model = None
|
||
predictor = None
|
||
|
||
# 可用的模型配置
|
||
AVAILABLE_MODELS = {
|
||
'kronos-mini': {
|
||
'name': 'Kronos-mini',
|
||
'model_id': 'NeoQuasar/Kronos-mini',
|
||
'tokenizer_id': 'NeoQuasar/Kronos-Tokenizer-2k',
|
||
'context_length': 2048,
|
||
'params': '4.1M',
|
||
'description': '轻量级模型，适合快速预测'
|
||
},
|
||
'kronos-small': {
|
||
'name': 'Kronos-small',
|
||
'model_id': 'NeoQuasar/Kronos-small',
|
||
'tokenizer_id': 'NeoQuasar/Kronos-Tokenizer-base',
|
||
'context_length': 512,
|
||
'params': '24.7M',
|
||
'description': '小型模型，平衡性能和速度'
|
||
},
|
||
'kronos-base': {
|
||
'name': 'Kronos-base',
|
||
'model_id': 'NeoQuasar/Kronos-base',
|
||
'tokenizer_id': 'NeoQuasar/Kronos-Tokenizer-base',
|
||
'context_length': 512,
|
||
'params': '102.3M',
|
||
'description': '基础模型，提供更好的预测质量'
|
||
}
|
||
}
|
||
|
||
def load_data_files():
|
||
"""扫描data目录并返回可用的数据文件"""
|
||
data_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'data')
|
||
data_files = []
|
||
|
||
if os.path.exists(data_dir):
|
||
for file in os.listdir(data_dir):
|
||
if file.endswith(('.csv', '.feather')):
|
||
file_path = os.path.join(data_dir, file)
|
||
file_size = os.path.getsize(file_path)
|
||
data_files.append({
|
||
'name': file,
|
||
'path': file_path,
|
||
'size': f"{file_size / 1024:.1f} KB" if file_size < 1024*1024 else f"{file_size / (1024*1024):.1f} MB"
|
||
})
|
||
|
||
return data_files
|
||
|
||
def load_data_file(file_path):
|
||
"""加载数据文件"""
|
||
try:
|
||
if file_path.endswith('.csv'):
|
||
df = pd.read_csv(file_path)
|
||
elif file_path.endswith('.feather'):
|
||
df = pd.read_feather(file_path)
|
||
else:
|
||
return None, "不支持的文件格式"
|
||
|
||
# 检查必要的列
|
||
required_cols = ['open', 'high', 'low', 'close']
|
||
if not all(col in df.columns for col in required_cols):
|
||
return None, f"缺少必要的列: {required_cols}"
|
||
|
||
# 处理时间戳列
|
||
if 'timestamps' in df.columns:
|
||
df['timestamps'] = pd.to_datetime(df['timestamps'])
|
||
elif 'timestamp' in df.columns:
|
||
df['timestamps'] = pd.to_datetime(df['timestamp'])
|
||
elif 'date' in df.columns:
|
||
# 如果列名是'date'，将其重命名为'timestamps'
|
||
df['timestamps'] = pd.to_datetime(df['date'])
|
||
else:
|
||
# 如果没有时间戳列，创建一个
|
||
df['timestamps'] = pd.date_range(start='2024-01-01', periods=len(df), freq='1H')
|
||
|
||
# 确保数值列是数值类型
|
||
for col in ['open', 'high', 'low', 'close']:
|
||
df[col] = pd.to_numeric(df[col], errors='coerce')
|
||
|
||
# 处理volume列（可选）
|
||
if 'volume' in df.columns:
|
||
df['volume'] = pd.to_numeric(df['volume'], errors='coerce')
|
||
|
||
# 处理amount列（可选，但不用于预测）
|
||
if 'amount' in df.columns:
|
||
df['amount'] = pd.to_numeric(df['amount'], errors='coerce')
|
||
|
||
# 删除包含NaN的行
|
||
df = df.dropna()
|
||
|
||
return df, None
|
||
|
||
except Exception as e:
|
||
return None, f"加载文件失败: {str(e)}"
|
||
|
||
def save_prediction_results(file_path, prediction_type, prediction_results, actual_data, input_data, prediction_params):
|
||
"""保存预测结果到文件"""
|
||
try:
|
||
# 创建预测结果目录
|
||
results_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'prediction_results')
|
||
os.makedirs(results_dir, exist_ok=True)
|
||
|
||
# 生成文件名
|
||
timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
|
||
filename = f'prediction_{timestamp}.json'
|
||
filepath = os.path.join(results_dir, filename)
|
||
|
||
# 准备保存的数据
|
||
save_data = {
|
||
'timestamp': datetime.datetime.now().isoformat(),
|
||
'file_path': file_path,
|
||
'prediction_type': prediction_type,
|
||
'prediction_params': prediction_params,
|
||
'input_data_summary': {
|
||
'rows': len(input_data),
|
||
'columns': list(input_data.columns),
|
||
'price_range': {
|
||
'open': {'min': float(input_data['open'].min()), 'max': float(input_data['open'].max())},
|
||
'high': {'min': float(input_data['high'].min()), 'max': float(input_data['high'].max())},
|
||
'low': {'min': float(input_data['low'].min()), 'max': float(input_data['low'].max())},
|
||
'close': {'min': float(input_data['close'].min()), 'max': float(input_data['close'].max())}
|
||
},
|
||
'last_values': {
|
||
'open': float(input_data['open'].iloc[-1]),
|
||
'high': float(input_data['high'].iloc[-1]),
|
||
'low': float(input_data['low'].iloc[-1]),
|
||
'close': float(input_data['close'].iloc[-1])
|
||
}
|
||
},
|
||
'prediction_results': prediction_results,
|
||
'actual_data': actual_data,
|
||
'analysis': {}
|
||
}
|
||
|
||
# 如果有实际数据，进行对比分析
|
||
if actual_data and len(actual_data) > 0:
|
||
# 计算连续性分析
|
||
if len(prediction_results) > 0 and len(actual_data) > 0:
|
||
last_pred = prediction_results[0] # 第一个预测点
|
||
first_actual = actual_data[0] # 第一个实际点
|
||
|
||
save_data['analysis']['continuity'] = {
|
||
'last_prediction': {
|
||
'open': last_pred['open'],
|
||
'high': last_pred['high'],
|
||
'low': last_pred['low'],
|
||
'close': last_pred['close']
|
||
},
|
||
'first_actual': {
|
||
'open': first_actual['open'],
|
||
'high': first_actual['high'],
|
||
'low': first_actual['low'],
|
||
'close': first_actual['close']
|
||
},
|
||
'gaps': {
|
||
'open_gap': abs(last_pred['open'] - first_actual['open']),
|
||
'high_gap': abs(last_pred['high'] - first_actual['high']),
|
||
'low_gap': abs(last_pred['low'] - first_actual['low']),
|
||
'close_gap': abs(last_pred['close'] - first_actual['close'])
|
||
},
|
||
'gap_percentages': {
|
||
'open_gap_pct': (abs(last_pred['open'] - first_actual['open']) / first_actual['open']) * 100,
|
||
'high_gap_pct': (abs(last_pred['high'] - first_actual['high']) / first_actual['high']) * 100,
|
||
'low_gap_pct': (abs(last_pred['low'] - first_actual['low']) / first_actual['low']) * 100,
|
||
'close_gap_pct': (abs(last_pred['close'] - first_actual['close']) / first_actual['close']) * 100
|
||
}
|
||
}
|
||
|
||
# 保存到文件
|
||
with open(filepath, 'w', encoding='utf-8') as f:
|
||
json.dump(save_data, f, indent=2, ensure_ascii=False)
|
||
|
||
print(f"预测结果已保存到: {filepath}")
|
||
return filepath
|
||
|
||
except Exception as e:
|
||
print(f"保存预测结果失败: {e}")
|
||
return None
|
||
|
||
def create_prediction_chart(df, pred_df, lookback, pred_len, actual_df=None, historical_start_idx=0):
|
||
"""创建预测图表"""
|
||
# 使用指定的历史数据起始位置，而不是总是从df的开头开始
|
||
if historical_start_idx + lookback + pred_len <= len(df):
|
||
# 显示指定位置开始的lookback个历史点 + pred_len个预测点
|
||
historical_df = df.iloc[historical_start_idx:historical_start_idx+lookback]
|
||
prediction_range = range(historical_start_idx+lookback, historical_start_idx+lookback+pred_len)
|
||
else:
|
||
# 如果数据不够，调整到可用的最大范围
|
||
available_lookback = min(lookback, len(df) - historical_start_idx)
|
||
available_pred_len = min(pred_len, max(0, len(df) - historical_start_idx - available_lookback))
|
||
historical_df = df.iloc[historical_start_idx:historical_start_idx+available_lookback]
|
||
prediction_range = range(historical_start_idx+available_lookback, historical_start_idx+available_lookback+available_pred_len)
|
||
|
||
# 创建图表
|
||
fig = go.Figure()
|
||
|
||
# 添加历史数据（K线图）
|
||
fig.add_trace(go.Candlestick(
|
||
x=historical_df['timestamps'] if 'timestamps' in historical_df.columns else historical_df.index,
|
||
open=historical_df['open'],
|
||
high=historical_df['high'],
|
||
low=historical_df['low'],
|
||
close=historical_df['close'],
|
||
name='历史数据 (400个数据点)',
|
||
increasing_line_color='#26A69A',
|
||
decreasing_line_color='#EF5350'
|
||
))
|
||
|
||
# 添加预测数据（K线图）
|
||
if pred_df is not None and len(pred_df) > 0:
|
||
# 计算预测数据的时间戳 - 确保与历史数据连续
|
||
if 'timestamps' in df.columns and len(historical_df) > 0:
|
||
# 从历史数据的最后一个时间点开始，按相同的时间间隔创建预测时间戳
|
||
last_timestamp = historical_df['timestamps'].iloc[-1]
|
||
time_diff = df['timestamps'].iloc[1] - df['timestamps'].iloc[0] if len(df) > 1 else pd.Timedelta(hours=1)
|
||
|
||
pred_timestamps = pd.date_range(
|
||
start=last_timestamp + time_diff,
|
||
periods=len(pred_df),
|
||
freq=time_diff
|
||
)
|
||
else:
|
||
# 如果没有时间戳，使用索引
|
||
pred_timestamps = range(len(historical_df), len(historical_df) + len(pred_df))
|
||
|
||
fig.add_trace(go.Candlestick(
|
||
x=pred_timestamps,
|
||
open=pred_df['open'],
|
||
high=pred_df['high'],
|
||
low=pred_df['low'],
|
||
close=pred_df['close'],
|
||
name='预测数据 (120个数据点)',
|
||
increasing_line_color='#66BB6A',
|
||
decreasing_line_color='#FF7043'
|
||
))
|
||
|
||
# 添加实际数据用于对比（如果存在）
|
||
if actual_df is not None and len(actual_df) > 0:
|
||
# 实际数据应该与预测数据在同一个时间段
|
||
if 'timestamps' in df.columns:
|
||
# 实际数据应该使用与预测数据相同的时间戳，确保时间对齐
|
||
if 'pred_timestamps' in locals():
|
||
actual_timestamps = pred_timestamps
|
||
else:
|
||
# 如果没有预测时间戳，从历史数据最后一个时间点开始计算
|
||
if len(historical_df) > 0:
|
||
last_timestamp = historical_df['timestamps'].iloc[-1]
|
||
time_diff = df['timestamps'].iloc[1] - df['timestamps'].iloc[0] if len(df) > 1 else pd.Timedelta(hours=1)
|
||
actual_timestamps = pd.date_range(
|
||
start=last_timestamp + time_diff,
|
||
periods=len(actual_df),
|
||
freq=time_diff
|
||
)
|
||
else:
|
||
actual_timestamps = range(len(historical_df), len(historical_df) + len(actual_df))
|
||
else:
|
||
actual_timestamps = range(len(historical_df), len(historical_df) + len(actual_df))
|
||
|
||
fig.add_trace(go.Candlestick(
|
||
x=actual_timestamps,
|
||
open=actual_df['open'],
|
||
high=actual_df['high'],
|
||
low=actual_df['low'],
|
||
close=actual_df['close'],
|
||
name='实际数据 (120个数据点)',
|
||
increasing_line_color='#FF9800',
|
||
decreasing_line_color='#F44336'
|
||
))
|
||
|
||
# 更新布局
|
||
fig.update_layout(
|
||
title='Kronos 金融预测结果 - 400个历史点 + 120个预测点 vs 120个实际点',
|
||
xaxis_title='时间',
|
||
yaxis_title='价格',
|
||
template='plotly_white',
|
||
height=600,
|
||
showlegend=True
|
||
)
|
||
|
||
# 确保x轴时间连续
|
||
if 'timestamps' in historical_df.columns:
|
||
# 获取所有时间戳并排序
|
||
all_timestamps = []
|
||
if len(historical_df) > 0:
|
||
all_timestamps.extend(historical_df['timestamps'])
|
||
if 'pred_timestamps' in locals():
|
||
all_timestamps.extend(pred_timestamps)
|
||
if 'actual_timestamps' in locals():
|
||
all_timestamps.extend(actual_timestamps)
|
||
|
||
if all_timestamps:
|
||
all_timestamps = sorted(all_timestamps)
|
||
fig.update_xaxes(
|
||
range=[all_timestamps[0], all_timestamps[-1]],
|
||
rangeslider_visible=False,
|
||
type='date'
|
||
)
|
||
|
||
return json.dumps(fig, cls=plotly.utils.PlotlyJSONEncoder)
|
||
|
||
@app.route('/')
|
||
def index():
|
||
"""主页"""
|
||
return render_template('index.html')
|
||
|
||
@app.route('/api/data-files')
|
||
def get_data_files():
|
||
"""获取可用的数据文件列表"""
|
||
data_files = load_data_files()
|
||
return jsonify(data_files)
|
||
|
||
@app.route('/api/load-data', methods=['POST'])
|
||
def load_data():
|
||
"""加载数据文件"""
|
||
try:
|
||
data = request.get_json()
|
||
file_path = data.get('file_path')
|
||
|
||
if not file_path:
|
||
return jsonify({'error': '文件路径不能为空'}), 400
|
||
|
||
df, error = load_data_file(file_path)
|
||
if error:
|
||
return jsonify({'error': error}), 400
|
||
|
||
# 检测数据的时间频率
|
||
def detect_timeframe(df):
|
||
if len(df) < 2:
|
||
return "未知"
|
||
|
||
time_diffs = []
|
||
for i in range(1, min(10, len(df))): # 检查前10个时间差
|
||
diff = df['timestamps'].iloc[i] - df['timestamps'].iloc[i-1]
|
||
time_diffs.append(diff)
|
||
|
||
if not time_diffs:
|
||
return "未知"
|
||
|
||
# 计算平均时间差
|
||
avg_diff = sum(time_diffs, pd.Timedelta(0)) / len(time_diffs)
|
||
|
||
# 转换为可读格式
|
||
if avg_diff < pd.Timedelta(minutes=1):
|
||
return f"{avg_diff.total_seconds():.0f}秒"
|
||
elif avg_diff < pd.Timedelta(hours=1):
|
||
return f"{avg_diff.total_seconds() / 60:.0f}分钟"
|
||
elif avg_diff < pd.Timedelta(days=1):
|
||
return f"{avg_diff.total_seconds() / 3600:.0f}小时"
|
||
else:
|
||
return f"{avg_diff.days}天"
|
||
|
||
# 返回数据信息
|
||
data_info = {
|
||
'rows': len(df),
|
||
'columns': list(df.columns),
|
||
'start_date': df['timestamps'].min().isoformat() if 'timestamps' in df.columns else 'N/A',
|
||
'end_date': df['timestamps'].max().isoformat() if 'timestamps' in df.columns else 'N/A',
|
||
'price_range': {
|
||
'min': float(df[['open', 'high', 'low', 'close']].min().min()),
|
||
'max': float(df[['open', 'high', 'low', 'close']].max().max())
|
||
},
|
||
'prediction_columns': ['open', 'high', 'low', 'close'] + (['volume'] if 'volume' in df.columns else []),
|
||
'timeframe': detect_timeframe(df)
|
||
}
|
||
|
||
return jsonify({
|
||
'success': True,
|
||
'data_info': data_info,
|
||
'message': f'成功加载数据，共 {len(df)} 行'
|
||
})
|
||
|
||
except Exception as e:
|
||
return jsonify({'error': f'加载数据失败: {str(e)}'}), 500
|
||
|
||
@app.route('/api/predict', methods=['POST'])
|
||
def predict():
|
||
"""进行预测"""
|
||
try:
|
||
data = request.get_json()
|
||
file_path = data.get('file_path')
|
||
lookback = int(data.get('lookback', 400))
|
||
pred_len = int(data.get('pred_len', 120))
|
||
|
||
# 获取预测质量参数
|
||
temperature = float(data.get('temperature', 1.0))
|
||
top_p = float(data.get('top_p', 0.9))
|
||
sample_count = int(data.get('sample_count', 1))
|
||
|
||
if not file_path:
|
||
return jsonify({'error': '文件路径不能为空'}), 400
|
||
|
||
# 加载数据
|
||
df, error = load_data_file(file_path)
|
||
if error:
|
||
return jsonify({'error': error}), 400
|
||
|
||
if len(df) < lookback:
|
||
return jsonify({'error': f'数据长度不足，需要至少 {lookback} 行数据'}), 400
|
||
|
||
# 进行预测
|
||
if MODEL_AVAILABLE and predictor is not None:
|
||
try:
|
||
# 使用真实的Kronos模型
|
||
# 只使用必要的列：OHLCV，不包含amount
|
||
required_cols = ['open', 'high', 'low', 'close']
|
||
if 'volume' in df.columns:
|
||
required_cols.append('volume')
|
||
|
||
# 处理时间段选择
|
||
start_date = data.get('start_date')
|
||
|
||
if start_date:
|
||
# 自定义时间段 - 修复逻辑：使用选择的窗口内的数据
|
||
start_dt = pd.to_datetime(start_date)
|
||
|
||
# 找到开始时间之后的数据
|
||
mask = df['timestamps'] >= start_dt
|
||
time_range_df = df[mask]
|
||
|
||
# 确保有足够的数据：lookback + pred_len
|
||
if len(time_range_df) < lookback + pred_len:
|
||
return jsonify({'error': f'从开始时间 {start_dt.strftime("%Y-%m-%d %H:%M")} 开始的数据不足，需要至少 {lookback + pred_len} 个数据点，当前只有 {len(time_range_df)} 个'}), 400
|
||
|
||
# 使用选择的窗口内的前lookback个数据点进行预测
|
||
x_df = time_range_df.iloc[:lookback][required_cols]
|
||
x_timestamp = time_range_df.iloc[:lookback]['timestamps']
|
||
|
||
# 使用选择的窗口内的后pred_len个数据点作为实际值
|
||
y_timestamp = time_range_df.iloc[lookback:lookback+pred_len]['timestamps']
|
||
|
||
# 计算实际的时间段长度
|
||
start_timestamp = time_range_df['timestamps'].iloc[0]
|
||
end_timestamp = time_range_df['timestamps'].iloc[lookback+pred_len-1]
|
||
time_span = end_timestamp - start_timestamp
|
||
|
||
prediction_type = f"Kronos模型预测 (选择的窗口内：前{lookback}个数据点预测，后{pred_len}个数据点对比，时间跨度: {time_span})"
|
||
else:
|
||
# 使用最新数据
|
||
x_df = df.iloc[:lookback][required_cols]
|
||
x_timestamp = df.iloc[:lookback]['timestamps']
|
||
y_timestamp = df.iloc[lookback:lookback+pred_len]['timestamps']
|
||
prediction_type = "Kronos模型预测 (最新数据)"
|
||
|
||
# 确保时间戳是Series格式，不是DatetimeIndex，避免Kronos模型的.dt属性错误
|
||
if isinstance(x_timestamp, pd.DatetimeIndex):
|
||
x_timestamp = pd.Series(x_timestamp, name='timestamps')
|
||
if isinstance(y_timestamp, pd.DatetimeIndex):
|
||
y_timestamp = pd.Series(y_timestamp, name='timestamps')
|
||
|
||
pred_df = predictor.predict(
|
||
df=x_df,
|
||
x_timestamp=x_timestamp,
|
||
y_timestamp=y_timestamp,
|
||
pred_len=pred_len,
|
||
T=temperature,
|
||
top_p=top_p,
|
||
sample_count=sample_count
|
||
)
|
||
|
||
except Exception as e:
|
||
return jsonify({'error': f'Kronos模型预测失败: {str(e)}'}), 500
|
||
else:
|
||
return jsonify({'error': 'Kronos模型未加载，请先加载模型'}), 400
|
||
|
||
# 准备实际数据用于对比（如果存在）
|
||
actual_data = []
|
||
actual_df = None
|
||
|
||
if start_date: # 自定义时间段
|
||
# 修复逻辑：使用选择的窗口内的数据
|
||
# 预测使用的是选择的窗口内的前400个数据点
|
||
# 实际数据应该是选择的窗口内的后120个数据点
|
||
start_dt = pd.to_datetime(start_date)
|
||
|
||
# 找到从start_date开始的数据
|
||
mask = df['timestamps'] >= start_dt
|
||
time_range_df = df[mask]
|
||
|
||
if len(time_range_df) >= lookback + pred_len:
|
||
# 获取选择的窗口内的后120个数据点作为实际值
|
||
actual_df = time_range_df.iloc[lookback:lookback+pred_len]
|
||
|
||
for i, (_, row) in enumerate(actual_df.iterrows()):
|
||
actual_data.append({
|
||
'timestamp': row['timestamps'].isoformat(),
|
||
'open': float(row['open']),
|
||
'high': float(row['high']),
|
||
'low': float(row['low']),
|
||
'close': float(row['close']),
|
||
'volume': float(row['volume']) if 'volume' in row else 0,
|
||
'amount': float(row['amount']) if 'amount' in row else 0
|
||
})
|
||
else: # 最新数据
|
||
# 预测使用的是前400个数据点
|
||
# 实际数据应该是400个数据点之后的120个数据点
|
||
if len(df) >= lookback + pred_len:
|
||
actual_df = df.iloc[lookback:lookback+pred_len]
|
||
for i, (_, row) in enumerate(actual_df.iterrows()):
|
||
actual_data.append({
|
||
'timestamp': row['timestamps'].isoformat(),
|
||
'open': float(row['open']),
|
||
'high': float(row['high']),
|
||
'low': float(row['low']),
|
||
'close': float(row['close']),
|
||
'volume': float(row['volume']) if 'volume' in row else 0,
|
||
'amount': float(row['amount']) if 'amount' in row else 0
|
||
})
|
||
|
||
# 创建图表 - 传递历史数据的起始位置
|
||
if start_date:
|
||
# 自定义时间段：找到历史数据在原始df中的起始位置
|
||
start_dt = pd.to_datetime(start_date)
|
||
mask = df['timestamps'] >= start_dt
|
||
historical_start_idx = df[mask].index[0] if len(df[mask]) > 0 else 0
|
||
else:
|
||
# 最新数据：从开头开始
|
||
historical_start_idx = 0
|
||
|
||
chart_json = create_prediction_chart(df, pred_df, lookback, pred_len, actual_df, historical_start_idx)
|
||
|
||
# 准备预测结果数据 - 修复时间戳计算逻辑
|
||
if 'timestamps' in df.columns:
|
||
if start_date:
|
||
# 自定义时间段：使用选择的窗口数据计算时间戳
|
||
start_dt = pd.to_datetime(start_date)
|
||
mask = df['timestamps'] >= start_dt
|
||
time_range_df = df[mask]
|
||
|
||
if len(time_range_df) >= lookback:
|
||
# 从选择的窗口的最后一个时间点开始计算预测时间戳
|
||
last_timestamp = time_range_df['timestamps'].iloc[lookback-1]
|
||
time_diff = df['timestamps'].iloc[1] - df['timestamps'].iloc[0]
|
||
future_timestamps = pd.date_range(
|
||
start=last_timestamp + time_diff,
|
||
periods=pred_len,
|
||
freq=time_diff
|
||
)
|
||
else:
|
||
future_timestamps = []
|
||
else:
|
||
# 最新数据：从整个数据文件的最后时间点开始计算
|
||
last_timestamp = df['timestamps'].iloc[-1]
|
||
time_diff = df['timestamps'].iloc[1] - df['timestamps'].iloc[0]
|
||
future_timestamps = pd.date_range(
|
||
start=last_timestamp + time_diff,
|
||
periods=pred_len,
|
||
freq=time_diff
|
||
)
|
||
else:
|
||
future_timestamps = range(len(df), len(df) + pred_len)
|
||
|
||
prediction_results = []
|
||
for i, (_, row) in enumerate(pred_df.iterrows()):
|
||
prediction_results.append({
|
||
'timestamp': future_timestamps[i].isoformat() if i < len(future_timestamps) else f"T{i}",
|
||
'open': float(row['open']),
|
||
'high': float(row['high']),
|
||
'low': float(row['low']),
|
||
'close': float(row['close']),
|
||
'volume': float(row['volume']) if 'volume' in row else 0,
|
||
'amount': float(row['amount']) if 'amount' in row else 0
|
||
})
|
||
|
||
# 保存预测结果到文件
|
||
try:
|
||
save_prediction_results(
|
||
file_path=file_path,
|
||
prediction_type=prediction_type,
|
||
prediction_results=prediction_results,
|
||
actual_data=actual_data,
|
||
input_data=x_df,
|
||
prediction_params={
|
||
'lookback': lookback,
|
||
'pred_len': pred_len,
|
||
'temperature': temperature,
|
||
'top_p': top_p,
|
||
'sample_count': sample_count,
|
||
'start_date': start_date if start_date else 'latest'
|
||
}
|
||
)
|
||
except Exception as e:
|
||
print(f"保存预测结果失败: {e}")
|
||
|
||
return jsonify({
|
||
'success': True,
|
||
'prediction_type': prediction_type,
|
||
'chart': chart_json,
|
||
'prediction_results': prediction_results,
|
||
'actual_data': actual_data,
|
||
'has_comparison': len(actual_data) > 0,
|
||
'message': f'预测完成，生成了 {pred_len} 个预测点' + (f'，包含 {len(actual_data)} 个实际数据点用于对比' if len(actual_data) > 0 else '')
|
||
})
|
||
|
||
except Exception as e:
|
||
return jsonify({'error': f'预测失败: {str(e)}'}), 500
|
||
|
||
@app.route('/api/load-model', methods=['POST'])
|
||
def load_model():
|
||
"""加载Kronos模型"""
|
||
global tokenizer, model, predictor
|
||
|
||
try:
|
||
if not MODEL_AVAILABLE:
|
||
return jsonify({'error': 'Kronos模型库不可用'}), 400
|
||
|
||
data = request.get_json()
|
||
model_key = data.get('model_key', 'kronos-small')
|
||
device = data.get('device', 'cpu')
|
||
|
||
if model_key not in AVAILABLE_MODELS:
|
||
return jsonify({'error': f'不支持的模型: {model_key}'}), 400
|
||
|
||
model_config = AVAILABLE_MODELS[model_key]
|
||
|
||
# 加载tokenizer和模型
|
||
tokenizer = KronosTokenizer.from_pretrained(model_config['tokenizer_id'])
|
||
model = Kronos.from_pretrained(model_config['model_id'])
|
||
|
||
# 创建predictor
|
||
predictor = KronosPredictor(model, tokenizer, device=device, max_context=model_config['context_length'])
|
||
|
||
return jsonify({
|
||
'success': True,
|
||
'message': f'模型加载成功: {model_config["name"]} ({model_config["params"]}) on {device}',
|
||
'model_info': {
|
||
'name': model_config['name'],
|
||
'params': model_config['params'],
|
||
'context_length': model_config['context_length'],
|
||
'description': model_config['description']
|
||
}
|
||
})
|
||
|
||
except Exception as e:
|
||
return jsonify({'error': f'模型加载失败: {str(e)}'}), 500
|
||
|
||
@app.route('/api/available-models')
|
||
def get_available_models():
|
||
"""获取可用的模型列表"""
|
||
return jsonify({
|
||
'models': AVAILABLE_MODELS,
|
||
'model_available': MODEL_AVAILABLE
|
||
})
|
||
|
||
@app.route('/api/model-status')
|
||
def get_model_status():
|
||
"""获取模型状态"""
|
||
if MODEL_AVAILABLE:
|
||
if predictor is not None:
|
||
return jsonify({
|
||
'available': True,
|
||
'loaded': True,
|
||
'message': 'Kronos模型已加载并可用',
|
||
'current_model': {
|
||
'name': predictor.model.__class__.__name__,
|
||
'device': str(next(predictor.model.parameters()).device)
|
||
}
|
||
})
|
||
else:
|
||
return jsonify({
|
||
'available': True,
|
||
'loaded': False,
|
||
'message': 'Kronos模型可用但未加载'
|
||
})
|
||
else:
|
||
return jsonify({
|
||
'available': False,
|
||
'loaded': False,
|
||
'message': 'Kronos模型库不可用，请安装相关依赖'
|
||
})
|
||
|
||
if __name__ == '__main__':
|
||
print("启动Kronos Web UI...")
|
||
print(f"模型可用性: {MODEL_AVAILABLE}")
|
||
if MODEL_AVAILABLE:
|
||
print("提示: 可以通过 /api/load-model 接口加载Kronos模型")
|
||
else:
|
||
print("提示: 将使用模拟数据进行演示")
|
||
|
||
app.run(debug=True, host='0.0.0.0', port=7070)
|