Update app.py

app.py

@@ -4,65 +4,42 @@ import pandas as pd
 from sklearn.linear_model import LinearRegression
 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.
-
-    Returns:
-        A plotly figure with the historical and predicted prices.
-    """
     # Download stock data
     data = yf.download(ticker, start=start_date, end=end_date)
     data = data["Close"]
 
-    # Convert dates to a numerical format (days since start date)
-    start_date = pd.to_datetime(start_date)
-    days_since_start = (data.index - start_date).days
+    # Convert index to Unix timestamp (seconds)
+    data.index = (data.index - pd.Timestamp("1970-01-01")) // pd.Timedelta('1s')
 
     # Train linear regression model
-    X = days_since_start.values[:-prediction_days].reshape(-1, 1)
+    X = data.index.values[:-prediction_days].reshape(-1, 1)
     y = data.values[:-prediction_days]
     model = LinearRegression()
     model.fit(X, y)
 
     # Prepare data for prediction
-    last_date = data.index[-1]
-    future_dates = pd.date_range(start=last_date, periods=prediction_days + 1, closed='right')
-    future_days_since_start = (future_dates - start_date).days
-    X_future = future_days_since_start.values.reshape(-1, 1)
+    last_timestamp = data.index[-1]
+    future_timestamps = pd.date_range(start=pd.to_datetime(last_timestamp, unit='s'), 
+                                      periods=prediction_days + 1, closed='right')
+    future_timestamps = (future_timestamps - pd.Timestamp("1970-01-01")) // pd.Timedelta('1s')
+    X_future = future_timestamps.values.reshape(-1, 1)
 
     # Predict future prices
     predicted_prices = model.predict(X_future)
 
-    # Predict future prices
-    future_dates = data.index.values[-prediction_days:]
-    X_future = future_dates.reshape(-1, 1)
-    predicted_prices = model.predict(X_future)
-
     # Prepare data for plotting
     historical_prices = go.Scatter(
-        x=data.index,
+        x=pd.to_datetime(data.index, unit='s'),
         y=data.values,
         mode="lines",
-        line_color=lambda p: "green" if p > 0 else "red",
-        name="Historical Prices",
+        name="Historical Prices"
     )
     predicted_prices_trace = go.Scatter(
-        x=future_dates,
+        x=pd.to_datetime(future_timestamps, unit='s'),
         y=predicted_prices,
         mode="lines",
-        line_color="gold",
-        line_width=3,
-        marker_line_width=3,
-        marker_color="black",
-        name="Predicted Prices",
+        name="Predicted Prices"
     )
 
     # Plot data
@@ -73,12 +50,11 @@ def train_predict_wrapper(ticker, start_date, end_date, prediction_days):
         title="Stock Price Prediction",
         xaxis_title="Date",
         yaxis_title="Price",
-        legend_title_text="Data",
+        legend_title_text="Data"
     )
 
     return fig
 
-
 # Define Gradio interface
 interface = gr.Interface(
     fn=train_predict_wrapper,
@@ -88,7 +64,7 @@ interface = gr.Interface(
         gr.Textbox(label="End Date (YYYY-MM-DD)"),
         gr.Slider(minimum=1, maximum=30, step=1, label="Prediction Days"),
     ],
-    outputs="plot",
+    outputs="plot"
 )
 
 # Launch the app