Update app.py

app.py

@@ -2,11 +2,7 @@ import gradio as gr
 import yfinance as yf
 from prophet import Prophet
 import pandas as pd
-from datetime import datetime
 import plotly.graph_objects as go
-import plotly.express as px
-import json
-import numpy as np
 
 def download_data(ticker, start_date='2010-01-01'):
     data = yf.download(ticker, start=start_date)
@@ -16,29 +12,24 @@ def download_data(ticker, start_date='2010-01-01'):
     if 'Adj Close' in data.columns:
         data = data[['Date', 'Adj Close']].copy()
         data.rename(columns={'Date': 'ds', 'Adj Close': 'y'}, inplace=True)
-        data['ds'] = pd.to_datetime(data['ds'])
     else:
         raise ValueError("Expected 'Adj Close' in columns")
+    data['ds'] = pd.to_datetime(data['ds'])
     return data
 
-def json_serial(obj):
-    """JSON serializer for objects not serializable by default json code"""
-    if isinstance(obj, (datetime, np.datetime64)):
-        return obj.isoformat()
-    raise TypeError("Type not serializable")
-
 def predict_future_prices(ticker, periods=1825):
     data = download_data(ticker)
     model = Prophet(daily_seasonality=False, weekly_seasonality=True, yearly_seasonality=True)
     model.fit(data)
     future = model.make_future_dataframe(periods=periods, freq='D')
     forecast = model.predict(future)
-    forecast['ds'] = forecast['ds'].apply(pd.to_datetime)
-    fig_main = go.Figure()
-    fig_main.add_trace(go.Scatter(x=data['ds'], y=data['y'], mode='lines', name='Actual (Black)', line=dict(color='black')))
-    fig_main.add_trace(go.Scatter(x=forecast['ds'], y=forecast['yhat'], mode='lines', name='Forecast (Blue)'))
-    fig_main.update_layout(title="Forecast vs Actual Stock Prices")
-    return json.dumps({"figure": fig_main.to_json(), "forecast": forecast[['ds', 'yhat', 'yhat_lower', 'yhat_upper']].to_dict('records')}, default=json_serial)
+    
+    fig = go.Figure()
+    fig.add_trace(go.Scatter(x=data['ds'], y=data['y'], mode='lines', name='Actual (Black)', line=dict(color='black')))
+    fig.add_trace(go.Scatter(x=forecast['ds'], y=forecast['yhat'], mode='lines', name='Forecast (Blue)'))
+    
+    # Make sure to return both the figure and the forecast data
+    return fig, forecast[['ds', 'yhat', 'yhat_lower', 'yhat_upper']]
 
 with gr.Blocks() as app:
     with gr.Row():