Update app.py

app.py

@@ -1,74 +1,106 @@
 import gradio as gr
-import torch
-import torch.nn as nn
-import numpy as np
-from sklearn.preprocessing import MinMaxScaler
 import pandas as pd
+import numpy as np
+import torch
+import pickle
+import matplotlib.pyplot as plt
+import io
+from torch import nn
 
-# Load data and scaler
-df = pd.read_csv('HistoricalQuotes.csv')
-df['Date'] = pd.to_datetime(df['Date'], format='%m/%d/%Y')
-df = df.sort_index()
-
-# Find the closing price column
-possible_columns = [' Close/Last', 'Close', 'close', 'close_last']
-close_column = None
-for col in possible_columns:
-    if col in df.columns:
-        close_column = col
-        break
-
-if close_column is None:
-    raise KeyError(f"None of {possible_columns} found in columns: {list(df.columns)}")
 
-df = df[[close_column]].rename(columns={close_column: 'Close'})
-df['Close'] = df['Close'].replace({r'\$': ''}, regex=True).astype(float)
+with open("arima.pkl", "rb") as f:
+    arima_model = pickle.load(f)
 
-scaler = MinMaxScaler(feature_range=(0, 1))
-scaler.fit(df['Close'].values.reshape(-1, 1))
 
-# Define LSTM model
 class LSTMModel(nn.Module):
-    def __init__(self, input_size=1, hidden_size=50, num_layers=2, output_size=1, dropout=0.2):
-        super().__init__()
-        self.hidden_size = hidden_size
-        self.num_layers = num_layers  # Store num_layers as instance variable
-        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, dropout=dropout)
+    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(self.num_layers, x.size(0), self.hidden_size)
-        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size)
+        h0 = torch.zeros(2, x.size(0), 50)
+        c0 = torch.zeros(2, x.size(0), 50)
         out, _ = self.lstm(x, (h0, c0))
-        return self.fc(out[:, -1, :])
-
-# Load model
-model = LSTMModel()
-model.load_state_dict(torch.load('lstm_model.pth', map_location=torch.device('cpu')))
-model.eval()
-
-def forecast(past_prices, steps=30):
-    try:
-        prices = [float(x.strip().replace('$', '')) for x in past_prices.split(',')]
-        if len(prices) < 60:
-            return "Error: At least 60 prices required."
-        prices_scaled = scaler.transform(np.array(prices).reshape(-1, 1))
-        current_seq = torch.from_numpy(prices_scaled[-60:].reshape(1, 60, 1)).float()
-        predictions = []
-        for _ in range(steps):
-            with torch.no_grad():
-                pred_scaled = model(current_seq).item()
-                predictions.append(pred_scaled)
-                current_seq = torch.cat((current_seq[:, 1:, :], torch.tensor([[[pred_scaled]]]).float()), dim=1)
-        predictions = scaler.inverse_transform(np.array(predictions).reshape(-1, 1)).flatten()
-        return pd.DataFrame({'Forecast': predictions}).to_string()
-    except Exception as e:
-        return f"Error: {str(e)}"
-
-# Create Gradio interface
-iface = gr.Interface(
-    fn=forecast,
-    inputs=gr.Textbox(label="Past Prices (comma-separated, e.g., 273.36,273.52,...)"),
-    outputs=gr.Textbox(label="Forecasted Prices")
-)
-iface.launch(share=True)
+        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()