Update app.py

app.py

@@ -247,21 +247,28 @@ def container_memory():
 
 @app.post("/predict/contracts_clauses", response_model=ClassificationResult)
 def predict_contracts_clauses(req: TextRequest):
+    model = models["contracts_clauses"]
     # The SetFit model predicts labels directly (no integer conversion needed)
-    preds = models["contracts_clauses"].predict([req.text])  # Ensure using predict()
+    preds = model.predict([req.text])
     label = preds[0]  # Already a string like 'terms'
 
     # Try to get a confidence score using predict_proba if available
     score = 1.0
-    if hasattr(models["contracts_clauses"], "predict_proba"):
-        probs = models["contracts_clauses"].predict_proba([req.text])[0]
-        # Get the index of the predicted label and its probability
-        # Note: SetFit often stores labels in .model.labels
-        if hasattr(models["contracts_clauses"].model, "labels"):
-            idx = list(models["contracts_clauses"].model.labels).index(label)
-            score = probs[idx]
-        else:
-            score = max(probs)
+    if hasattr(model, "predict_proba"):
+        try:
+            probs = model.predict_proba([req.text])[0]
+            # model.labels stores the label strings in the order expected by predict_proba
+            if hasattr(model, "labels") and model.labels is not None:
+                # Find the index of the predicted label
+                if label in model.labels:
+                    idx = model.labels.index(label)
+                    score = probs[idx]
+                else:
+                    score = max(probs)
+            else:
+                score = max(probs)
+        except Exception:
+            score = 1.0
 
     return ClassificationResult(label=label, score=round(float(score), 4))