Update app.py

app.py

@@ -1,40 +1,62 @@
-from gradio_app_builder import app
+import gradio as gr
 import yfinance as yf
 from sklearn.linear_model import LinearRegression
-import json
+import plotly.graph_objs as go
+import numpy as np
 
-@app.route("/")
 def train_predict_wrapper(ticker, start_date, end_date, prediction_days):
-  """
-  Downloads stock data, trains a linear regression model, and predicts future prices.
-
-  Args:
-    ticker: The ticker symbol of the stock.
-    start_date: The start date for the data (YYYY-MM-DD format).
-    end_date: The end date for the data (YYYY-MM-DD format).
-    prediction_days: The number of days to predict.
-
-  Returns:
-    A list of predicted closing prices for the next `prediction_days`.
-  """
-  # Download stock data
-  data = yf.download(ticker, start=start_date, end=end_date)
-  # Extract closing price and set as index
-  data = data["Close"].to_frame().set_index(data.index.values)
-
-  # Train linear regression model
-  X = data.index.values[:-prediction_days].reshape(-1, 1)
-  y = data.values[:-prediction_days]
-  model = LinearRegression()
-  model.fit(X, y)
-
-  # Predict future prices
-  future_dates = data.index.values[-prediction_days:]
-  X_future = future_dates.reshape(-1, 1)
-  predicted_prices = model.predict(X_future).tolist()
-
-  # Return predicted prices as JSON
-  return json.dumps(predicted_prices)
-
-# Launch the Gradio application
-app.launch()
+    """
+    Downloads stock data, trains a linear regression model, and predicts future prices.
+
+    Args:
+      ticker: The ticker symbol of the stock.
+      start_date: The start date for the data (YYYY-MM-DD format).
+      end_date: The end date for the data (YYYY-MM-DD format).
+      prediction_days: The number of days to predict.
+
+    Returns:
+      A plot of predicted closing prices for the next `prediction_days`.
+    """
+    # Download stock data
+    data = yf.download(ticker, start=start_date, end=end_date)
+    # Extract closing price
+    data = data["Close"]
+
+    # Prepare data for model
+    data = data.reset_index()
+    data['Date'] = data['Date'].map(mdates.date2num)
+    X = np.array(data.index).reshape(-1, 1)
+    y = data['Close'].values
+
+    # Train linear regression model
+    model = LinearRegression()
+    model.fit(X[:-prediction_days], y[:-prediction_days])
+
+    # Predict future prices
+    future_indices = np.array(range(len(X), len(X) + prediction_days)).reshape(-1, 1)
+    predicted_prices = model.predict(future_indices)
+
+    # Plot
+    dates = [mdates.num2date(date).strftime('%Y-%m-%d') for date in data['Date']]
+    future_dates = [mdates.num2date(date).strftime('%Y-%m-%d') for date in future_indices.flatten()]
+    fig = go.Figure()
+    fig.add_trace(go.Scatter(x=dates, y=y, mode='lines', name='Historical Prices'))
+    fig.add_trace(go.Scatter(x=future_dates, y=predicted_prices, mode='lines', name='Predicted Prices'))
+    fig.update_layout(title='Stock Price Prediction', xaxis_title='Date', yaxis_title='Price')
+
+    return fig
+
+# Define Gradio interface
+iface = gr.Interface(
+    fn=train_predict_wrapper,
+    inputs=[
+        gr.inputs.Textbox(label="Ticker Symbol"),
+        gr.inputs.Textbox(label="Start Date (YYYY-MM-DD)"),
+        gr.inputs.Textbox(label="End Date (YYYY-MM-DD)"),
+        gr.inputs.Slider(minimum=1, maximum=30, step=1, default=5, label="Prediction Days")
+    ],
+    outputs="plot"
+)
+
+# Launch the app
+iface.launch()