Update app.py

app.py

@@ -12,33 +12,12 @@ import PIL.Image
 
 # Step 1: Fetch stock data
 def fetch_stock_data(ticker, start_date, end_date):
-    """
-    Fetch historical stock data from Yahoo Finance using the yfinance library.
-
-    Args:
-        ticker (str): Stock ticker symbol (e.g., 'AAPL').
-        start_date (str): Start date for fetching stock data (YYYY-MM-DD).
-        end_date (str): End date for fetching stock data (YYYY-MM-DD).
-
-    Returns:
-        pd.DataFrame: Stock data including Date, Open, High, Low, Close, and Volume.
-    """
     stock_data = yf.download(ticker, start=start_date, end=end_date)
     stock_data.reset_index(inplace=True)
     return stock_data
 
 # Step 2: Prepare data for the LSTM model
 def prepare_data(df):
-    """
-    Prepares the stock data for the LSTM model by scaling the 'Close' price.
-
-    Args:
-        df (pd.DataFrame): DataFrame containing the stock data.
-
-    Returns:
-        scaled_data (np.array): Normalized stock prices for training the model.
-        scaler (MinMaxScaler): Scaler used to normalize and later denormalize predictions.
-    """
     scaler = MinMaxScaler(feature_range=(0, 1))
     close_prices = df['Close'].values.reshape(-1, 1)
     scaled_data = scaler.fit_transform(close_prices)
@@ -46,15 +25,6 @@ def prepare_data(df):
 
 # Step 3: Build the LSTM model
 def build_model(input_shape):
-    """
-    Builds and compiles the LSTM model for stock price prediction.
-
-    Args:
-        input_shape (tuple): Shape of the input data for the LSTM model.
-
-    Returns:
-        tf.keras.Model: Compiled LSTM model.
-    """
     model = tf.keras.Sequential([
         tf.keras.layers.LSTM(50, return_sequences=True, input_shape=input_shape),
         tf.keras.layers.LSTM(50, return_sequences=False),
@@ -66,17 +36,6 @@ def build_model(input_shape):
 
 # Step 4: Train the LSTM model
 def train_model(model, train_data, epochs=5):
-    """
-    Trains the LSTM model using the scaled stock price data.
-
-    Args:
-        model (tf.keras.Model): The LSTM model to be trained.
-        train_data (np.array): Scaled stock price data for training.
-        epochs (int): Number of training epochs.
-
-    Returns:
-        model (tf.keras.Model): The trained LSTM model.
-    """
     X_train, y_train = [], []
     for i in range(60, len(train_data)):
         X_train.append(train_data[i-60:i, 0])
@@ -90,17 +49,6 @@ def train_model(model, train_data, epochs=5):
 
 # Step 5: Predict future stock prices
 def predict_future(model, last_data, steps=90):
-    """
-    Predicts future stock prices using the trained LSTM model.
-
-    Args:
-        model (tf.keras.Model): The trained LSTM model.
-        last_data (np.array): The last 60 days of stock price data.
-        steps (int): Number of future days to predict.
-
-    Returns:
-        predictions (list): Predicted stock prices for the future.
-    """
     predictions = []
     input_data = last_data[-60:].reshape(1, -1)
     
@@ -113,58 +61,34 @@ def predict_future(model, last_data, steps=90):
 
 # Step 6: Plot historical and predicted stock prices
 def plot_predictions(data, predicted_prices, scaler):
-    """
-    Plots the historical stock prices and the predicted future stock prices.
-
-    Args:
-        data (pd.DataFrame): DataFrame containing historical stock data.
-        predicted_prices (list): Predicted stock prices for future dates.
-        scaler (MinMaxScaler): Scaler to inverse transform the predicted prices.
-
-    Returns:
-        PIL.Image: The image of the plot saved in memory.
-    """
     last_60_days = data['Close'][-60:].values.reshape(-1, 1)
     predicted_prices = np.array(predicted_prices).reshape(-1, 1)
     predicted_prices = scaler.inverse_transform(predicted_prices)
 
-    # Plot historical data
+    # Create the plot
     plt.figure(figsize=(14,6))
-    plt.plot(data['Close'], label="Historical Prices")
+    plt.plot(data['Date'], data['Close'], label="Historical Prices", color='blue')
+    
+    future_dates = pd.date_range(start=data['Date'].iloc[-1], periods=len(predicted_prices)+1)[1:]
+    plt.plot(future_dates, predicted_prices, label="Predicted Prices", color='orange')
 
-    # Plot predicted data
-    future_days = range(len(data), len(data) + len(predicted_prices))
-    plt.plot(future_days, predicted_prices, label="Predicted Prices")
     plt.title("Stock Price Prediction")
-    plt.xlabel("Days")
+    plt.xlabel("Date")
     plt.ylabel("Stock Price")
     plt.legend()
 
-    # Save the plot to a bytes buffer
+    # Save plot to in-memory buffer
     buf = BytesIO()
     plt.savefig(buf, format='png')
     buf.seek(0)
     image = PIL.Image.open(buf)
 
-    # Clear the plot so it doesn’t overlap
-    plt.clf()
+    plt.close()  # Clear the plot to avoid overlap
 
     return image
 
 # Step 7: Gradio Interface Function
 def stock_prediction_app(ticker, start_date_str, end_date_str):
-    """
-    The core function for the Gradio app. Fetches stock data, trains the LSTM model,
-    predicts future prices, and visualizes the results.
-
-    Args:
-        ticker (str): Stock ticker symbol selected by the user.
-        start_date_str (str): Start date selected by the user (YYYY-MM-DD).
-        end_date_str (str): End date selected by the user (YYYY-MM-DD).
-
-    Returns:
-        PIL.Image: The plot showing historical and predicted stock prices.
-    """
     # Convert date strings to datetime objects
     start_date = datetime.strptime(start_date_str, "%Y-%m-%d").date()
     end_date = datetime.strptime(end_date_str, "%Y-%m-%d").date()
@@ -195,8 +119,8 @@ ui = gr.Interface(
     fn=stock_prediction_app,
     inputs=[
         gr.Dropdown(tickers, label="Select Stock Ticker"),
-        gr.Textbox(label="Start Date (YYYY-MM-DD)"),   # Use textbox instead
-        gr.Textbox(label="End Date (YYYY-MM-DD)")      # Use textbox instead
+        gr.Textbox(label="Start Date (YYYY-MM-DD)"),
+        gr.Textbox(label="End Date (YYYY-MM-DD)")
     ],
     outputs=gr.Image(),  # Updated output to return an image
     title="Stock Prediction App",