Update train.py

train.py

@@ -14,6 +14,8 @@ from sklearn.preprocessing import LabelEncoder
 from sklearn.metrics import f1_score
 import numpy as np
 import io
+import os
+import pickle
 
 # --- Configuration ---
 MAX_LEN = 100
@@ -238,6 +240,29 @@ def train_model(training_data_text: str):
     # Convert accuracy to 0-1 range for callback
     accuracy_normalized = test_accuracy / 100.0
     
+    # Save model artifacts to /tmp for on-demand predictions
+    try:
+        model_path = '/tmp/latest_model.pt'
+        word_to_idx_path = '/tmp/word_to_idx.pt'
+        label_encoder_path = '/tmp/label_encoder.pkl'
+        
+        # Save model state dict
+        torch.save(final_model.state_dict(), model_path)
+        print(f"Saved model to {model_path}")
+        
+        # Save word_to_idx dictionary
+        torch.save(word_to_idx, word_to_idx_path)
+        print(f"Saved word_to_idx to {word_to_idx_path}")
+        
+        # Save label_encoder
+        with open(label_encoder_path, 'wb') as f:
+            pickle.dump(label_encoder, f)
+        print(f"Saved label_encoder to {label_encoder_path}")
+        
+    except Exception as e:
+        print(f"Warning: Failed to save model artifacts to /tmp: {e}")
+        # Continue even if saving fails - model is still returned in result
+    
     # Return model and stats
     return {
         'model': final_model,
@@ -247,7 +272,7 @@ def train_model(training_data_text: str):
             'accuracy': accuracy_normalized,
             'loss': float(avg_loss),
             'f1_scores': f1_scores_dict,
-            'model_path': 'trained_model'  # Not saved to disk in HF Space
+            'model_path': '/tmp/latest_model.pt'  # Path to saved model
         }
     }