import os import json import gradio as gr import pandas as pd import numpy as np import torch import joblib from utils.retrain import LSTMModel # ensure retrain.py has unified model loader # ---- CONFIG ---- TICKERS = ["TSLA", "NVDA", "SPY"] # add more tickers HORIZONS = { "1d": {"days": 1}, "1w": {"days": 5}, "4w": {"days": 20}, "6m": {"days": 120}, "1y": {"days": 240}, } MODEL_TYPES = ["Per-Ticker", "Unified"] BASE_DIR = os.path.dirname(__file__) DATA_PATH = os.path.join(BASE_DIR, "data") MODELS_DIR = os.path.join(BASE_DIR, "models") DEVICE = "cpu" # ------------------------ # MC Dropout Prediction # ------------------------ def enable_mc_dropout(model): for m in model.modules(): if isinstance(m, torch.nn.Dropout): m.train() return model def mc_dropout_predict(model, X_tensor, y_scaler, n_samples=200): model = enable_mc_dropout(model) preds = [] with torch.no_grad(): for _ in range(n_samples): preds.append(model(X_tensor).item()) preds = torch.tensor(preds) mean_pred = preds.mean().item() std_pred = preds.std().item() mean_orig = y_scaler.inverse_transform([[mean_pred]])[0, 0] std_orig = y_scaler.inverse_transform([[mean_pred + std_pred]])[0, 0] - mean_orig plus_minus_percent = (std_orig / mean_orig * 100) if mean_orig != 0 else 0 return { "predicted_price": mean_orig, "plus_minus_percent": plus_minus_percent, "confidence_percent": 95.0, "lower_bound": mean_orig - std_orig, "upper_bound": mean_orig + std_orig, } # ------------------------ # Data Preparation # ------------------------ def prepare_input_sequence(df, x_scaler, seq_len=90): features = df[["Open", "High", "Low", "Close", "Volume"]].values X_scaled = x_scaler.transform(features) X_seq = X_scaled[-seq_len:] return torch.tensor(X_seq, dtype=torch.float32).unsqueeze(0).to(DEVICE) def prepare_unified_input_sequence(df, x_scaler, seq_len, ticker_idx, num_tickers): features = df[["Open", "High", "Low", "Close", "Volume"]].values X_scaled = x_scaler.transform(features) onehot = np.eye(num_tickers)[ticker_idx] onehot_seq = np.repeat(onehot.reshape(1, -1), seq_len, axis=0) X_full = np.hstack([X_scaled[-seq_len:], onehot_seq]) X_tensor = torch.tensor(X_full, dtype=torch.float32).unsqueeze(0).to(DEVICE) return X_tensor # ------------------------ # Load Models # ------------------------ def load_per_ticker_model(ticker, horizon): out_dir = os.path.join(MODELS_DIR, ticker) model_path = os.path.join(out_dir, f"{ticker}_{horizon}_model.pth") x_scaler_path = os.path.join(out_dir, f"{ticker}_{horizon}_scaler.pkl") y_scaler_path = os.path.join(out_dir, f"{ticker}_{horizon}_y_scaler.pkl") config_path = os.path.join(out_dir, f"{ticker}_{horizon}_config.json") if not all(os.path.exists(p) for p in [model_path, x_scaler_path, y_scaler_path]): raise FileNotFoundError(f"Missing model/scalers for {ticker} ({horizon})") x_scaler = joblib.load(x_scaler_path) y_scaler = joblib.load(y_scaler_path) cfg = {} if os.path.exists(config_path): with open(config_path, "r") as f: cfg = json.load(f) model = LSTMModel( input_size=cfg.get("input_size", len(x_scaler.mean_)), hidden_size=cfg.get("hidden_size", 128), num_layers=cfg.get("num_layers", 2), dropout=cfg.get("dropout", 0.2), ) model.load_state_dict(torch.load(model_path, map_location=DEVICE, weights_only=True)) model.to(DEVICE) seq_len = cfg.get("seq_len", 90) return model, x_scaler, y_scaler, seq_len def load_unified_model_and_scalers(horizon): udir = os.path.join(MODELS_DIR, "unified") model_path = os.path.join(udir, f"unified_{horizon}_model.pth") x_scaler_path = os.path.join(udir, f"unified_{horizon}_scaler.pkl") y_scaler_path = os.path.join(udir, f"unified_{horizon}_y_scaler.pkl") ticker_map_path = os.path.join(udir, "unified_tickers.pkl") if not all(os.path.exists(p) for p in [model_path, x_scaler_path, y_scaler_path, ticker_map_path]): raise FileNotFoundError(f"Missing unified model/scalers for {horizon}") x_scaler = joblib.load(x_scaler_path) y_scaler = joblib.load(y_scaler_path) ticker_map = joblib.load(ticker_map_path) # Load config if exists config_path = os.path.join(udir, f"unified_{horizon}_config.json") if os.path.exists(config_path): with open(config_path, "r") as f: cfg = json.load(f) input_size = cfg.get("input_size", len(x_scaler.mean_) + len(ticker_map)) hidden_size = cfg.get("hidden_size", 128) num_layers = cfg.get("num_layers", 2) dropout = cfg.get("dropout", 0.2) seq_len = cfg.get("seq_len", 90) else: input_size = len(x_scaler.mean_) + len(ticker_map) hidden_size, num_layers, dropout, seq_len = 128, 2, 0.2, 90 model = LSTMModel( input_size=input_size, hidden_size=hidden_size, num_layers=num_layers, dropout=dropout, ) model.load_state_dict(torch.load(model_path, map_location=DEVICE, weights_only=True)) model.to(DEVICE) return model, x_scaler, y_scaler, ticker_map, seq_len # ------------------------ # Run Prediction # ------------------------ def run_prediction(ticker, horizon, model_type): result = {} csv_path = os.path.join(DATA_PATH, f"{ticker}.csv") if not os.path.exists(csv_path): return {"error": f"No data for {ticker}"} df = pd.read_csv(csv_path) try: if model_type == "Per-Ticker": model, x_scaler, y_scaler, seq_len = load_per_ticker_model(ticker, horizon) X = prepare_input_sequence(df, x_scaler, seq_len) result[ticker] = {horizon: mc_dropout_predict(model, X, y_scaler)} else: # Unified model, x_scaler, y_scaler, ticker_map, seq_len = load_unified_model_and_scalers(horizon) if ticker not in ticker_map: return {"error": f"{ticker} not in unified model"} X = prepare_unified_input_sequence(df, x_scaler, seq_len, ticker_map[ticker], len(ticker_map)) result[ticker] = {horizon: mc_dropout_predict(model, X, y_scaler)} except Exception as e: result["error"] = str(e) return json.dumps(result, indent=2) # ------------------------ # GRADIO UI # ------------------------ with gr.Blocks(title="Stock Price Forecast") as demo: gr.Markdown("# Stock Price Forecasting with LSTM + MC Dropout") gr.Markdown("Select ticker, horizon, and model type, then click Predict.") ticker_dropdown = gr.Dropdown(choices=TICKERS, label="Ticker", value=TICKERS[0]) horizon_dropdown = gr.Dropdown(choices=list(HORIZONS.keys()), label="Forecast Horizon", value="1d") model_dropdown = gr.Dropdown(choices=MODEL_TYPES, label="Model Type", value="Per-Ticker") predict_btn = gr.Button("Predict") output_box = gr.Code(label="JSON Output", language="json") predict_btn.click( fn=run_prediction, inputs=[ticker_dropdown, horizon_dropdown, model_dropdown], outputs=[output_box], ) demo.launch()