Update app.py

app.py

@@ -9,6 +9,14 @@ model_name = "SamanthaStorm/abuse-pattern-detector-v2"
 model = RobertaForSequenceClassification.from_pretrained(model_name)
 tokenizer = RobertaTokenizer.from_pretrained(model_name)
 
+# Define the final label order your model used
+LABELS = [
+    "gaslighting", "mockery", "dismissiveness", "control",
+    "guilt_tripping", "apology_baiting", "blame_shifting", "projection",
+    "contradictory_statements", "manipulation", "deflection", "insults",
+    "obscure_formal", "recovery_phase", "suicidal_threat", "physical_threat",
+    "extreme_control"
+]
 TOTAL_LABELS = 17
 
 # Our model outputs 17 labels:
@@ -16,27 +24,36 @@ TOTAL_LABELS = 17
 # - Last 3 are Danger Assessment cues
 TOTAL_LABELS = 17
 
-def analyze_messages(text):
-    input_text = text.strip()
+def analyze_messages(input_text):
+    input_text = input_text.strip()
     if not input_text:
-        return "Please enter a message for analysis.", None
-    
-    # Tokenize input text
+        return "Please enter a message for analysis."
+
+    # Tokenize
     inputs = tokenizer(input_text, return_tensors="pt", truncation=True, padding=True)
     with torch.no_grad():
         outputs = model(**inputs)
-    
-    # Assume model logits shape is [17] (for a single example)
-    logits = outputs.logits.squeeze()  # shape: [17]
+
+    # Squeeze out batch dimension: shape should be [17]
+    logits = outputs.logits.squeeze(0)
+
+    # Convert logits to probabilities
     scores = torch.sigmoid(logits).numpy()
-    
-    # For the first 14 labels (abuse patterns), count how many exceed threshold 0.5
-    abuse_pattern_scores = scores[:14]
-    concerning_pattern_count = int(np.sum(abuse_pattern_scores > 0.5))
-    
-    # For the last 3 labels (Danger Assessment cues), count how many exceed threshold 0.5
-    danger_scores = scores[14:17]
+
+    # Debug printing (remove once you're confident everything works)
+    print("Scores:", scores)
+
+    # First 14 = pattern scores
+    pattern_scores = scores[:14]
+    pattern_count = int(np.sum(pattern_scores > 0.5))
+
+    # Last 3 = danger cues
+    danger_scores = scores[14:]
     danger_flag_count = int(np.sum(danger_scores > 0.5))
+
+    # (Optional) Print label-by-label for debugging
+    for i, s in enumerate(scores):
+        print(LABELS[i], "=", round(s, 3))
     
     # Map danger flag count to Danger Assessment Score
     if danger_flag_count >= 2: