Update app.py

app.py

@@ -83,19 +83,19 @@ training_status = {
 model_path = MODEL_SAVE_DIR / "DEBERTA_model.pth"
 tokenizer = get_tokenizer(DEBERTA_MODEL_NAME)
 
-# Load label encoders with error handling
+# Initialize model and label encoders with error handling
 try:
     label_encoders = load_label_encoders()
-    num_labels_list = [len(label_encoders[col].classes_) for col in LABEL_COLUMNS]
+    model = DebertaMultiOutputModel([len(label_encoders[col].classes_) for col in LABEL_COLUMNS]).to(DEVICE)
+    if model_path.exists():
+        model.load_state_dict(torch.load(model_path, map_location=DEVICE))
+        model.eval()
+    else:
+        print(f"Warning: Model file {model_path} not found. Model will be initialized but not loaded.")
 except Exception as e:
-    logger.warning(f"Could not load label encoders: {str(e)}")
-    # Use default values if label encoders can't be loaded
-    num_labels_list = [2] * len(LABEL_COLUMNS)  # Default to binary classification
-
-model = DebertaMultiOutputModel(num_labels_list).to(DEVICE)
-if os.path.exists(model_path):
-    model.load_state_dict(torch.load(model_path, map_location=DEVICE))
-    model.eval()
+    print(f"Warning: Could not load label encoders or model: {str(e)}")
+    print("Model will be initialized when training starts.")
+    model = None
 
 class TrainingConfig(BaseModel):
     model_name: str = DEBERTA_MODEL_NAME
@@ -264,7 +264,7 @@ async def validate_model(
         
         data_df, label_encoders = load_and_preprocess_data(str(file_path))
         
-        model_path = MODEL_SAVE_DIR / f"{model_name}.pth"
+        model_path = MODEL_SAVE_DIR / f"{model_name}_model.pth"
         if not model_path.exists():
             raise HTTPException(status_code=404, detail="DeBERTa model file not found")
         
@@ -353,18 +353,16 @@ async def predict(
     """
     try:
         # Load the model
-        model_path = MODEL_SAVE_DIR / f"{model_name}.pth"
+        model_path = MODEL_SAVE_DIR / f"{model_name}_model.pth"
         if not model_path.exists():
             raise HTTPException(status_code=404, detail=f"Model {model_name} not found")
         
-        # Load label encoders with error handling
+        # Load label encoders
         try:
             label_encoders = load_label_encoders()
             num_labels_list = [len(label_encoders[col].classes_) for col in LABEL_COLUMNS]
         except Exception as e:
-            logger.warning(f"Could not load label encoders: {str(e)}")
-            # Use default values if label encoders can't be loaded
-            num_labels_list = [2] * len(LABEL_COLUMNS)  # Default to binary classification
+            raise HTTPException(status_code=500, detail=f"Could not load label encoders: {str(e)}")
         
         model = DebertaMultiOutputModel(num_labels_list).to(DEVICE)
         model.load_state_dict(torch.load(model_path, map_location=DEVICE))
@@ -492,7 +490,7 @@ async def predict(
 @app.get("/v1/deberta/download-model/{model_id}")
 async def download_model(model_id: str):
     """Download a trained model"""
-    model_path = MODEL_SAVE_DIR / f"{model_id}.pth"
+    model_path = MODEL_SAVE_DIR / f"{model_id}_model.pth"
     if not model_path.exists():
         raise HTTPException(status_code=404, detail="Model not found")
     
@@ -524,7 +522,7 @@ async def train_model_task(config: TrainingConfig, file_path: str, training_id:
                 tokenizer,
                 config.max_length
             )
-            model = DebertaMultiOutputModel(num_labels_list).to(DEVICE)
+            model = DebertaMultiOutputModel(num_labels_list, metadata_dim=metadata_dim).to(DEVICE)
         else:
             dataset = ComplianceDataset(
                 texts.tolist(),
@@ -546,7 +544,7 @@ async def train_model_task(config: TrainingConfig, file_path: str, training_id:
             training_status["current_loss"] = train_loss
             
             # Save model after each epoch
-            save_model(model, training_id)
+            save_model(model, training_id, 'pth')
         
         training_status.update({
             "is_training": False,