Update app.py

app.py

@@ -9,6 +9,8 @@ import gradio as gr
 import matplotlib.pyplot as plt
 import seaborn as sns
 import io
+import zipfile
+import joblib
 from PIL import Image
 import warnings
 warnings.filterwarnings('ignore')
@@ -19,7 +21,7 @@ def load_and_preprocess_data(file):
         data = pd.read_csv(file.name)
         
         # Convert suits and ranks to numerical values
-        suit_order = {'spades': 0, 'hearts': 1, 'clubs': 2, 'diamonds': 3}
+        suit_order = {'spades': 0, 'hearts': 1, 'logs': 2, 'diamonds': 3}
         rank_order = {'ace': 0, '2': 1, '3': 2, '4': 3, '5': 4, '6': 5, '7': 6, '8': 7, '9': 8, '10': 9, 
                       'jack': 10, 'queen': 11, 'king': 12}
         
@@ -100,7 +102,19 @@ def plot_accuracy_chart(accuracies):
     plt.close()
     return img
 
-# Training function with progress tracking
+# Function to create a ZIP file of models
+def create_model_zip(models):
+    zip_buffer = io.BytesIO()
+    with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
+        for model_name, model in models.items():
+            model_buffer = io.BytesIO()
+            joblib.dump(model, model_buffer)
+            model_buffer.seek(0)
+            zip_file.writestr(f"{model_name}_model.pkl", model_buffer.getvalue())
+    zip_buffer.seek(0)
+    return zip_buffer
+
+# Training function with progress tracking and model saving
 def train_model(file, n_estimators, learning_rate, max_depth, subsample, progress=gr.Progress()):
     progress(0, desc="Starting...")
     results = []
@@ -110,7 +124,7 @@ def train_model(file, n_estimators, learning_rate, max_depth, subsample, progres
         progress(0.1, desc="Loading and preprocessing data...")
         data, error = load_and_preprocess_data(file)
         if error:
-            return error, None, None
+            return error, None, None, None
         
         # Create features
         progress(0.2, desc="Engineering features...")
@@ -126,13 +140,13 @@ def train_model(file, n_estimators, learning_rate, max_depth, subsample, progres
         }
         
         # Scale features
-        progress(0.3, desc="Scaling features...")
-        scaler = StandardScaler()
+        progress(0.3, desc="Scaling features        scaler = StandardScaler()
         features_scaled = scaler.fit_transform(features)
         features_scaled = pd.DataFrame(features_scaled, columns=features.columns)
         
         accuracies = {}
         confusion_matrices = []
+        trained_models = {}
         
         # Train models
         for i, (target_name, target) in enumerate(targets.items()):
@@ -165,6 +179,9 @@ def train_model(file, n_estimators, learning_rate, max_depth, subsample, progres
                 verbose=False
             )
             
+            # Save model
+            trained_models[target_name] = model
+            
             # Evaluate
             y_pred = model.predict(X_test)
             accuracy = accuracy_score(y_test, y_pred)
@@ -179,20 +196,23 @@ def train_model(file, n_estimators, learning_rate, max_depth, subsample, progres
             cm_plot = plot_confusion_matrix(y_test, y_pred, f"Confusion Matrix - {target_name}")
             confusion_matrices.append(cm_plot)
         
-        progress(0.9, desc="Generating visualizations...")
+        progress(0.9, desc="Generating visualizations and model archive...")
         # Generate accuracy bar chart
         accuracy_plot = plot_accuracy_chart(accuracies)
         
+        # Create ZIP file of models
+        model_zip = create_model_zip(trained_models)
+        
         progress(1.0, desc="Completed!")
-        return "\n".join(results), accuracy_plot, confusion_matrices
+        return "\n".join(results), accuracy_plot, confusion_matrices, model_zip
     
     except Exception as e:
-        return f"Error during training: {str(e)}", None, None
+        return f"Error during training: {str(e)}", None, None, None
 
 # Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Card Game Prediction Model Training")
-    gr.Markdown("Upload the training dataset and configure hyperparameters to train the model. Track progress and view results with visualizations.")
+    gr.Markdown("Upload the training dataset and configure hyperparameters to train the model. Track progress, view results, and download trained models.")
     
     file_input = gr.File(label="Upload TRAINING_CARD_DATA.csv")
     n_estimators = gr.Slider(50, 300, value=100, step=10, label="Number of Estimators")
@@ -205,11 +225,12 @@ with gr.Blocks() as demo:
     output_text = gr.Textbox(label="Training Results")
     accuracy_plot = gr.Image(label="Accuracy Comparison")
     confusion_plots = gr.Gallery(label="Confusion Matrices")
+    model_download = gr.File(label="Download Trained Models (ZIP)")
     
     train_button.click(
         fn=train_model,
         inputs=[file_input, n_estimators, learning_rate, max_depth, subsample],
-        outputs=[output_text, accuracy_plot, confusion_plots]
+        outputs=[output_text, accuracy_plot, confusion_plots, model_download]
     )
 
 demo.launch()