Update ml_model.py

ml_model.py

@@ -2,16 +2,26 @@ import yfinance as yf
 import pandas as pd
 import ta
 import lightgbm as lgb
+import joblib
+import os
 from nifty100 import NIFTY100
 
+MODEL_FILE = "lightgbm_model.joblib"
 model = None
 features_list = ["rsi", "macd", "ma20", "ma50", "volatility", "ticker_code"]
 
-def train_model():
+def load_or_train_model():
     global model
     
-    all_data = []
+    # Check if we already have a trained model saved
+    if os.path.exists(MODEL_FILE):
+        print("Loading pre-trained LightGBM model...")
+        model = joblib.load(MODEL_FILE)
+        return
+        
+    print("Training LightGBM model from scratch. This will take a few minutes...")
     
+    all_data = []
     ticker_mapping = {ticker: idx for idx, ticker in enumerate(NIFTY100)}
     
     for ticker in NIFTY100:
@@ -21,44 +31,57 @@ def train_model():
             if len(df) < 100:
                 continue
                 
+            # FIX 1: Flatten the multi-index columns from yfinance
+            df.columns = df.columns.get_level_values(0)
+                
             close = df["Close"].squeeze()
             
             df["return"] = close.pct_change()
-            df["rsi"] = ta.momentum.RSIIndicator(close).rsi()
+            df["rsi"] = ta.momentum.RSIIndicator(close).rsi().rsi() if hasattr(ta.momentum.RSIIndicator(close).rsi(), 'rsi') else ta.momentum.RSIIndicator(close).rsi()
             df["macd"] = ta.trend.MACD(close).macd()
             df["ma20"] = close.rolling(20).mean()
             df["ma50"] = close.rolling(50).mean()
             df["volatility"] = df["return"].rolling(10).std()
             
             df["target"] = (df["return"].shift(-1) > 0).astype(int)
-            
             df["ticker_code"] = ticker_mapping[ticker]
             
             df = df.dropna()
             all_data.append(df)
             
-        except:
+        except Exception as e:
             pass
             
     full_df = pd.concat(all_data)
     
+    # FIX 2: Convert ticker_code to pandas category type (removes the warning)
+    full_df["ticker_code"] = full_df["ticker_code"].astype("category")
+    
     X = full_df[features_list]
     y = full_df["target"]
     
+    # Removed the deprecated categorical_feature argument
     model = lgb.LGBMClassifier(
         n_estimators=300,
         max_depth=6,
-        learning_rate=0.05,
-        categorical_feature=[len(features_list) - 1] 
+        learning_rate=0.05
     )
     
     model.fit(X, y)
+    
+    # Save the model to disk
+    joblib.dump(model, MODEL_FILE)
+    print("Model saved successfully as", MODEL_FILE)
 
 def predict_probability(df, ticker):
     global model
     
     if model is None:
-        train_model()
+        load_or_train_model()
+        
+    # Ensure columns are flat before making predictions
+    if isinstance(df.columns, pd.MultiIndex):
+        df.columns = df.columns.get_level_values(0)
         
     close = df["Close"].squeeze()
     
@@ -74,6 +97,9 @@ def predict_probability(df, ticker):
     
     X_pred = pd.DataFrame([[rsi, macd, ma20, ma50, vol, ticker_code]], columns=features_list)
     
+    # Cast it to category here as well
+    X_pred["ticker_code"] = X_pred["ticker_code"].astype("category")
+    
     prob = model.predict_proba(X_pred)[0][1]
     
     return round(prob * 100, 2)