Update app.py

app.py

@@ -76,7 +76,7 @@ training_status = {
 }
 
 # Load the model and tokenizer for prediction
-model_path = MODEL_SAVE_DIR / "BERT_model_model.pth"
+model_path = MODEL_SAVE_DIR / "BERT_model.pth"
 tokenizer = get_tokenizer('bert-base-uncased')
 
 # Initialize model and label encoders with error handling
@@ -260,7 +260,7 @@ async def validate_model(
         
         data_df, label_encoders = load_and_preprocess_data(str(file_path))
         
-        model_path = MODEL_SAVE_DIR / f"{model_name}_model.pth"
+        model_path = MODEL_SAVE_DIR / f"{model_name}.pth"
         if not model_path.exists():
             raise HTTPException(status_code=404, detail="BERT model file not found")
         
@@ -349,7 +349,7 @@ async def predict(
     """
     try:
         # Load the model
-        model_path = MODEL_SAVE_DIR / f"{model_name}_model.pth"
+        model_path = MODEL_SAVE_DIR / f"{model_name}.pth"
         if not model_path.exists():
             raise HTTPException(status_code=404, detail=f"Model {model_name} not found")
         
@@ -390,17 +390,31 @@ async def predict(
                 for i, row in data_df.iterrows():
                     transaction_pred = {}
                     for j, (col, probs) in enumerate(zip(LABEL_COLUMNS, all_probabilities)):
-                        pred = np.argmax(probs[i])
-                        decoded_pred = label_encoders[col].inverse_transform([pred])[0]
-                        
+                        # Get probabilities for each class
                         class_probs = {
                             label: float(probs[i][j]) 
                             for j, label in enumerate(label_encoders[col].classes_)
                         }
                         
+                        # Sort probabilities in descending order
+                        sorted_probs = sorted(class_probs.items(), key=lambda x: x[1], reverse=True)
+                        
+                        # Get top prediction and its probability
+                        top_pred, top_prob = sorted_probs[0]
+                        
+                        # Get top 3 predictions with probabilities
+                        top_3_predictions = [
+                            {"label": label, "probability": prob}
+                            for label, prob in sorted_probs[:3]
+                        ]
+                        
                         transaction_pred[col] = {
-                            "prediction": decoded_pred,
-                            "probabilities": class_probs
+                            "top_prediction": {
+                                "label": top_pred,
+                                "probability": top_prob
+                            },
+                            "alternative_predictions": top_3_predictions[1:],  # Exclude the top prediction
+                            "all_probabilities": class_probs  # Keep all probabilities for reference
                         }
                     
                     predictions.append({
@@ -466,17 +480,31 @@ async def predict(
             
             response = {}
             for i, (col, probs) in enumerate(zip(LABEL_COLUMNS, all_probabilities)):
-                pred = np.argmax(probs[0])
-                decoded_pred = label_encoders[col].inverse_transform([pred])[0]
-                
+                # Get probabilities for each class
                 class_probs = {
                     label: float(probs[0][j]) 
                     for j, label in enumerate(label_encoders[col].classes_)
                 }
                 
+                # Sort probabilities in descending order
+                sorted_probs = sorted(class_probs.items(), key=lambda x: x[1], reverse=True)
+                
+                # Get top prediction and its probability
+                top_pred, top_prob = sorted_probs[0]
+                
+                # Get top 3 predictions with probabilities
+                top_3_predictions = [
+                    {"label": label, "probability": prob}
+                    for label, prob in sorted_probs[:3]
+                ]
+                
                 response[col] = {
-                    "prediction": decoded_pred,
-                    "probabilities": class_probs
+                    "top_prediction": {
+                        "label": top_pred,
+                        "probability": top_prob
+                    },
+                    "alternative_predictions": top_3_predictions[1:],  # Exclude the top prediction
+                    "all_probabilities": class_probs  # Keep all probabilities for reference
                 }
 
             return response
@@ -565,5 +593,5 @@ async def train_model_task(config: TrainingConfig, file_path: str, training_id:
         })
 
 if __name__ == "__main__":
-    port = int(os.environ.get("PORT", 7860))
+    port = int(os.environ.get("PORT", 7861))
     uvicorn.run(app, host="0.0.0.0", port=port)