Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import pandas as pd | |
| import numpy as np | |
| import torch | |
| import pickle | |
| import matplotlib.pyplot as plt | |
| import io | |
| from torch import nn | |
| with open("arima.pkl", "rb") as f: | |
| arima_model = pickle.load(f) | |
| class LSTMModel(nn.Module): | |
| def __init__(self, input_size=1, hidden_size=50, num_layers=2, output_size=1): | |
| super(LSTMModel, self).__init__() | |
| self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True) | |
| self.fc = nn.Linear(hidden_size, output_size) | |
| def forward(self, x): | |
| h0 = torch.zeros(2, x.size(0), 50) | |
| c0 = torch.zeros(2, x.size(0), 50) | |
| out, _ = self.lstm(x, (h0, c0)) | |
| out = self.fc(out[:, -1, :]) | |
| return out | |
| # Load trained LSTM | |
| lstm_model = LSTMModel() | |
| lstm_model.load_state_dict(torch.load("lstm.pth", map_location=torch.device('cpu'))) | |
| lstm_model.eval() | |
| def predict_arima(values, horizon=10): | |
| forecast = arima_model.forecast(steps=horizon) | |
| return forecast.tolist() | |
| def predict_lstm(values, horizon=10): | |
| seq = torch.tensor(values[-50:], dtype=torch.float32).view(1, -1, 1) | |
| preds = [] | |
| for _ in range(horizon): | |
| with torch.no_grad(): | |
| pred = lstm_model(seq).item() | |
| preds.append(pred) | |
| seq = torch.cat([seq[:, 1:, :], torch.tensor([[[pred]]])], dim=1) | |
| return preds | |
| def forecast(file, horizon, model_choice): | |
| df = pd.read_csv(file.name) | |
| if "Close" not in df.columns: | |
| return "❌ CSV must contain a 'Close' column", None | |
| values = df["Close"].values.tolist() | |
| # Run forecasts | |
| preds_arima = predict_arima(values, horizon) | |
| preds_lstm = predict_lstm(values, horizon) | |
| # Prepare DataFrames | |
| future_index = [f"t+{i+1}" for i in range(horizon)] | |
| forecast_df = pd.DataFrame({ | |
| "Future": future_index, | |
| "ARIMA Forecast": preds_arima, | |
| "LSTM Forecast": preds_lstm | |
| }) | |
| # Plot | |
| plt.figure(figsize=(10,5)) | |
| plt.plot(range(len(values)), values, label="Historical") | |
| if model_choice in ["ARIMA", "Compare Both"]: | |
| plt.plot(range(len(values), len(values)+horizon), preds_arima, label="ARIMA Forecast") | |
| if model_choice in ["LSTM", "Compare Both"]: | |
| plt.plot(range(len(values), len(values)+horizon), preds_lstm, label="LSTM Forecast") | |
| plt.title(f"{model_choice} Stock Forecast") | |
| plt.xlabel("Time") | |
| plt.ylabel("Price") | |
| plt.legend() | |
| buf = io.BytesIO() | |
| plt.savefig(buf, format="png") | |
| buf.seek(0) | |
| return forecast_df, buf | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# 📈 Stock Price Forecasting Demo") | |
| gr.Markdown( | |
| "Upload a CSV containing stock prices (must have a **'Close'** column). " | |
| "Choose ARIMA, LSTM, or Compare Both, then set forecast horizon." | |
| ) | |
| with gr.Row(): | |
| file = gr.File(label="Upload CSV", file_types=[".csv"]) | |
| horizon = gr.Slider(5, 30, value=10, step=1, label="Forecast Horizon (days)") | |
| model_choice = gr.Radio(["ARIMA", "LSTM", "Compare Both"], label="Model", value="Compare Both") | |
| output_table = gr.DataFrame(label="Forecasted Prices") | |
| output_plot = gr.Image(type="pil", label="Forecast Plot") | |
| submit = gr.Button("Run Forecast") | |
| submit.click(forecast, inputs=[file, horizon, model_choice], outputs=[output_table, output_plot]) | |
| demo.launch() | |