Update scoring/specificity.py

scoring/specificity.py

@@ -19,6 +19,10 @@ class PredictionResponse(BaseModel):
 class PredictionsResponse(BaseModel):
     results: List[Dict[str, Union[str, float]]]
 
+class BatchPredictionScoreItem(BaseModel):
+    problem_description: str
+    score: float
+
 # Model environment variables
 MODEL_NAME = os.getenv("MODEL_NAME")
 LABEL_0 = os.getenv("LABEL_0")
@@ -102,19 +106,20 @@ def predict_batch(items: ProblematicList):
             with torch.no_grad():
                 outputs = model(**inputs)
                 probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
-                predicted_classes = torch.argmax(probabilities, dim=1).tolist()
-                confidence_scores = [probabilities[j][predicted_classes[j]].item() for j in range(len(predicted_classes))]
-            
-            # Converting numerical predictions into labels
-            for j, (pred_class, score) in enumerate(zip(predicted_classes, confidence_scores)):
-                predicted_label = LABEL_0 if pred_class == 0 else LABEL_1
-                results.append({
-                    "text": batch_texts[j],
-                    "class": predicted_label,
-                    "score": score
-                })
+                # predicted_classes = torch.argmax(probabilities, dim=1).tolist()
+                # confidence_scores = [probabilities[j][predicted_classes[j]].item() for j in range(len(predicted_classes))]
+
+            for j in range(len(batch_texts)):
+                score_specific_class = probabilities[j][1].item() 
+                
+                results.append(
+                    BatchPredictionScoreItem(
+                        problem_description=batch_texts[j],
+                        score=score_class_1
+                    )
+                )
         
-        return PredictionsResponse(results=results)
+        return results
     
     except Exception as e:
         print(f"Error during prediction: {str(e)}")