Update app.py

app.py

@@ -1,71 +1,70 @@
 from fastapi import FastAPI
 from sklearn.preprocessing import MinMaxScaler
-from tensorflow.keras.models import Sequential
+from tensorflow.keras.models import Sequential, load_model
 from tensorflow.keras.layers import LSTM, Dense
 import numpy as np
-from datetime import datetime, timedelta, time as dtime
+from datetime import datetime, time as dtime
 import pandas as pd
 import yfinance as yf
 import pytz
+import os
+import joblib
 
 app = FastAPI()
 
-# Create dataset helper
-def create_dataset(data, time_step=1, predict_steps=1):
+# Paths for saving models and scalers
+MODEL_DIR = "saved_models"
+os.makedirs(MODEL_DIR, exist_ok=True)
+
+# Helper: Create dataset for LSTM
+def create_dataset(data, time_step=60, predict_steps=5):
     X, y = [], []
     for i in range(len(data) - time_step - predict_steps):
-        X.append(data[i:i+time_step])
-        y.append(data[i+time_step:i+time_step+predict_steps])
+        X.append(data[i:i + time_step])
+        y.append(data[i + time_step:i + time_step + predict_steps])
     return np.array(X), np.array(y)
 
-def profit_with_short_selling_json(prices, st, im):
+# Helper: Predict Buy/Sell timing with real timestamps
+def profit_with_short_selling_json(prices, timestamps, symbol):
     min_price = prices[0]
-    buy_time = 0
-    sell_time = 0
     max_profit = 0
-    trade = []
+    buy_idx = 0
+    sell_idx = 0
 
     for i in range(1, len(prices)):
         if prices[i] - min_price > max_profit:
             max_profit = prices[i] - min_price
-            sell_time = i
-            buy_time = prices.tolist().index(min_price)
+            sell_idx = i
+            buy_idx = prices.tolist().index(min_price)
 
         if prices[i] < min_price:
             min_price = prices[i]
 
-    time_str = "9:15"
-    time_obj = datetime.strptime(time_str, "%H:%M")
-    i = time_obj + timedelta(minutes=im + buy_time)
-    f = time_obj + timedelta(minutes=im + sell_time)
+    buy_time = timestamps[buy_idx].strftime('%I:%M %p')
+    sell_time = timestamps[sell_idx].strftime('%I:%M %p')
 
-    if buy_time < sell_time:
-        stock = {
-            "stock": st,
-            "price": float(prices[sell_time]),
+    if buy_idx < sell_idx:
+        return {
+            "stock": symbol,
+            "price": float(prices[sell_idx]),
             "change": float(max(prices) - min(prices)),
             "prediction": "will rise",
             "reason": "The stock is showing positive trends.",
-            "time": "at 10:00 AM",
-            "suggestion": f"Buy at {i.strftime('%I:%M %p')}, Sell at {f.strftime('%I:%M %p')}"
+            "suggestion": f"Buy at {buy_time}, Sell at {sell_time}"
         }
     else:
-        stock = {
-            "stock": st,
-            "price": float(prices[buy_time]),
+        return {
+            "stock": symbol,
+            "price": float(prices[buy_idx]),
             "change": float(min(prices) - max(prices)),
             "prediction": "might fall",
             "reason": "Recent market instability affecting the stock.",
-            "time": "at 10:00 AM",
-            "suggestion": f"Sell at {f.strftime('%I:%M %p')}, Buy at {i.strftime('%I:%M %p')}"
+            "suggestion": f"Sell at {sell_time}, Buy at {buy_time}"
         }
 
-    trade.append(stock)
-    return trade
-
-# GET MARKET MINUTES
-def get_market_minutes():
-    now = datetime.now().time()
+# Helper: Get time in market minutes (India default)
+def get_market_minutes(market_tz='Asia/Kolkata'):
+    now = datetime.now(pytz.timezone(market_tz)).time()
     market_open = dtime(9, 15)
     market_close = dtime(15, 30)
 
@@ -76,86 +75,94 @@ def get_market_minutes():
     open_minutes = market_open.hour * 60 + market_open.minute
     return now_minutes - open_minutes
 
-# TRAIN MODEL
-@app.get('/model/train/{symbol}/{p_step}')
-def train(symbol: str, p_step: int):
-    import pytz
-    import yfinance as yf
-
-    predict_steps = p_step
-    time_steps = 60
-    training_period = '5d'
-    epochs = 206
+# Train model endpoint
+@app.get("/model/train/{symbol}")
+def train_model(symbol: str, time_step: int = 60, p_step: int = 5):
+    training_period = "5d"
+    epochs = 100
+    batch_size = 32
 
-    tz = pytz.timezone('Asia/Kolkata')
-    now = datetime.now(tz)
-    today = now.date()
+    stock_info = yf.Ticker(symbol).info
+    tz_name = stock_info.get('exchangeTimezoneName', 'UTC')
+    tz = pytz.timezone(tz_name)
+    today = datetime.now(tz).date()
 
-    data = yf.download(tickers=symbol, interval='1m', period=training_period, progress=False)
+    data = yf.download(tickers=symbol, interval="1m", period=training_period, progress=False)
     if data.empty:
-        return {"error": "No data received from Yahoo Finance. Try another symbol or interval."}
+        return {"error": "No data from Yahoo Finance"}
 
-    data = data.tz_convert('Asia/Kolkata')
+    data = data.tz_convert(tz_name)
     data['Date'] = data.index.date
-    available_dates = sorted(list(set(data['Date'])))
-    if today not in available_dates:
-        if not available_dates:
-            return {"error": "No available trading dates in the data."}
-        today = available_dates[-1]
-
     data_today = data[data['Date'] == today].between_time("09:15", "15:30")
     data_train = data[data['Date'] < today]
 
     if data_train.empty or data_today.empty:
-        return {"error": "Not enough data."}
+        return {"error": "Not enough data for training"}
 
-    prices = data_train[['Close']]
+    full_data = pd.concat([data_train[['Close']], data_today[['Close']]])
     scaler = MinMaxScaler()
-    scaled_prices = scaler.fit_transform(prices)
+    scaled_data = scaler.fit_transform(full_data)
 
-    X, y = create_dataset(scaled_prices, time_steps, predict_steps)
+    X, y = create_dataset(scaled_data, time_step, p_step)
     X = X.reshape(X.shape[0], X.shape[1], 1)
 
-    # Create and train a new model per request
     model = Sequential()
-    model.add(LSTM(50, return_sequences=False, input_shape=(time_steps, 1)))
-    model.add(Dense(predict_steps))
+    model.add(LSTM(50, return_sequences=False, input_shape=(time_step, 1)))
+    model.add(Dense(p_step))
     model.compile(loss='mean_squared_error', optimizer='adam')
-    model.fit(X, y, epochs=epochs, batch_size=32, verbose=1)
+    model.fit(X, y, epochs=epochs, batch_size=batch_size, verbose=1)
 
-    # Save to global for prediction use
-    app.state.model = model
-    app.state.scaler = scaler
-    app.state.data_today = data_today
+    # Save model and scaler
+    model_path = f"{MODEL_DIR}/{symbol}.h5"
+    scaler_path = f"{MODEL_DIR}/{symbol}_scaler.pkl"
+    model.save(model_path)
+    joblib.dump(scaler, scaler_path)
 
-    return {"status": "model trained"}
+    return {"status": f"Model trained and saved for {symbol}"}
 
-# PREDICT
+# Prediction endpoint
 @app.get("/predict/{symbol}")
-def predict_view(symbol: str):
-    time_step = 60
+def predict(symbol: str, time_step: int = 60, p_step: int = 5):
+    model_path = f"{MODEL_DIR}/{symbol}.h5"
+    scaler_path = f"{MODEL_DIR}/{symbol}_scaler.pkl"
 
-    if not hasattr(app.state, "model"):
-        return {"error": "Model not trained yet."}
+    if not os.path.exists(model_path) or not os.path.exists(scaler_path):
+        return {"error": "Model not trained. Call /model/train first."}
 
-    latest = app.state.data_today[['Close']].tail(time_step)
-    if len(latest) < time_step:
-        return {"error": "Not enough data for prediction."}
+    stock_info = yf.Ticker(symbol).info
+    tz_name = stock_info.get('exchangeTimezoneName', 'UTC')
+    tz = pytz.timezone(tz_name)
+    now = datetime.now(tz)
+    today = now.date()
 
-    scaler = app.state.scaler
-    scaler.fit(latest)
-    latest_scaled = scaler.transform(latest)
-    latest_input = latest_scaled.reshape(1, time_step, 1)
+    data = yf.download(tickers=symbol, interval="1m", period="1d", progress=False)
+    if data.empty:
+        return {"error": "No intraday data available"}
 
-    predicted_scaled = app.state.model.predict(latest_input)
+    data = data.tz_convert(tz_name)
+    data['Date'] = data.index.date
+    data_today = data[data['Date'] == today].between_time("09:15", "15:30")
+
+    latest_data = data_today[['Close']].tail(time_step)
+    timestamps = data_today.index[-p_step:]
+
+    if len(latest_data) < time_step or len(timestamps) < p_step:
+        return {"error": "Not enough recent data for prediction"}
+
+    scaler = joblib.load(scaler_path)
+    model = load_model(model_path)
+
+    scaled_input = scaler.transform(latest_data)
+    input_data = scaled_input.reshape(1, time_step, 1)
+
+    predicted_scaled = model.predict(input_data)
     predicted_prices = scaler.inverse_transform(predicted_scaled.reshape(-1, 1)).flatten()
 
-    im = get_market_minutes()
-    predictions = profit_with_short_selling_json(predicted_prices, symbol, im)[0]
+    suggestion = profit_with_short_selling_json(predicted_prices, timestamps, symbol)
 
     return {
-        "predictions": predictions,
-        "timestamp": f"{datetime.now()}",
-        "predicted_values": predicted_prices.tolist()
+        "symbol": symbol,
+        "timestamp": str(now),
+        "predicted_prices": predicted_prices.tolist(),
+        "suggestion": suggestion
     }
-