Update app.py

app.py

@@ -1,62 +1,81 @@
 import gradio as gr
 import yfinance as yf
 from sklearn.linear_model import LinearRegression
-import plotly.graph_objs as go
-import numpy as np
+import plotly.graph_objects as go
+
 
 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.
+        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`.
+        A plotly figure with the historical and predicted prices.
     """
     # 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
+    X = data.index.values[:-prediction_days].reshape(-1, 1)
+    y = data.values[:-prediction_days]
     model = LinearRegression()
-    model.fit(X[:-prediction_days], y[:-prediction_days])
+    model.fit(X, y)
 
     # Predict future prices
-    future_indices = np.array(range(len(X), len(X) + prediction_days)).reshape(-1, 1)
-    predicted_prices = model.predict(future_indices)
+    future_dates = data.index.values[-prediction_days:]
+    X_future = future_dates.reshape(-1, 1)
+    predicted_prices = model.predict(X_future)
 
-    # 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()]
+    # Prepare data for plotting
+    historical_prices = go.Scatter(
+        x=data.index,
+        y=data.values,
+        mode="lines",
+        line_color=lambda p: "green" if p > 0 else "red",
+        name="Historical Prices",
+    )
+    predicted_prices_trace = go.Scatter(
+        x=future_dates,
+        y=predicted_prices,
+        mode="lines",
+        line_color="gold",
+        line_width=3,
+        marker_line_width=3,
+        marker_color="black",
+        name="Predicted Prices",
+    )
+
+    # Plot data
     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')
+    fig.add_trace(historical_prices)
+    fig.add_trace(predicted_prices_trace)
+    fig.update_layout(
+        title="Stock Price Prediction",
+        xaxis_title="Date",
+        yaxis_title="Price",
+        legend_title_text="Data",
+    )
 
     return fig
 
+
 # Define Gradio interface
-iface = gr.Interface(
+interface = gr.Interface(
     fn=train_predict_wrapper,
     inputs=[
         gr.Textbox(label="Ticker Symbol"),
         gr.Textbox(label="Start Date (YYYY-MM-DD)"),
         gr.Textbox(label="End Date (YYYY-MM-DD)"),
-        gr.Slider(minimum=1, maximum=30, step=1, label="Prediction Days")
+        gr.Slider(minimum=1, maximum=30, step=1, label="Prediction Days"),
     ],
-    outputs="plot"
+    outputs="plot",
 )
 
 # Launch the app
-iface.launch()
+interface.launch()