Update app.py

app.py

@@ -1,12 +1,9 @@
 import gradio as gr
 import pandas as pd
 import numpy as np
-from itertools import combinations
 from sklearn.ensemble import GradientBoostingClassifier
 from sklearn.model_selection import train_test_split
 from sklearn.metrics import roc_auc_score
-import joblib
-import os
 
 from utils import generate_features, pick_top15, generate_system15_csv
 
@@ -15,23 +12,30 @@ def train_and_predict(file_obj):
     # Load dataset
     df = pd.read_csv(file_obj.name, header=None)
     # Columns: [draw_id?, draw_date, n1..n6, bonus, ...]
-    # Align to schema: we only keep date, n1..n6, bonus
     df = df.iloc[:, :8]
     df.columns = ["draw_date", "n1", "n2", "n3", "n4", "n5", "n6", "bonus"]
 
-    # Melt numbers into long format for training
+    # Generate features
     features, labels = generate_features(df)
 
-    # Split into train/test
-    X_train, X_test, y_train, y_test = train_test_split(
-        features, labels, test_size=0.2, random_state=42, stratify=labels
-    )
+    # Make sure we have at least 2 classes
+    if len(np.unique(labels)) < 2:
+        return "❌ Not enough class variety in labels to train the model.", None
 
-    # Train a simple gradient boosting model
+    # Train/test split with fallback
+    try:
+        X_train, X_test, y_train, y_test = train_test_split(
+            features, labels, test_size=0.2, random_state=42, stratify=labels
+        )
+    except ValueError:
+        X_train, y_train = features, labels
+        X_test, y_test = features, labels
+
+    # Train model
     model = GradientBoostingClassifier(n_estimators=200, max_depth=3, random_state=42)
     model.fit(X_train, y_train)
 
-    # Evaluate
+    # Evaluate if possible
     if len(np.unique(y_test)) > 1:
         auc = roc_auc_score(y_test, model.predict_proba(X_test)[:, 1])
     else:
@@ -69,3 +73,4 @@ if __name__ == "__main__":
 
 
 
+