Update app.py

app.py

@@ -90,7 +90,14 @@ boundary_model.eval()
 
 @spaces.GPU
 def predict_boundary_health(text):
-    """Predict boundary health for given text - returns 1 for healthy, 0 for unhealthy"""
+    """
+    Predict boundary health for given text
+    Returns: 
+    - 0 for Respected (healthy)
+    - 1 for Violated (unhealthy)
+    - 2 for Dismissed (unhealthy)
+    - 3 for Manipulative (unhealthy)
+    """
     try:
         inputs = boundary_tokenizer(text, return_tensors="pt", truncation=True, padding=True)
         inputs = {k: v.to(device) for k, v in inputs.items()}
@@ -98,32 +105,50 @@ def predict_boundary_health(text):
             outputs = boundary_model(**inputs)
             predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
         
-        # Return the actual prediction (0 or 1)
+        # Return the actual prediction (0, 1, 2, or 3)
         predicted_class = torch.argmax(predictions, dim=-1).item()
-        return predicted_class  # Returns 1 for healthy, 0 for unhealthy
+        confidence = predictions[0][predicted_class].item()
+        return predicted_class, confidence
         
     except Exception as e:
         logger.error(f"Error in boundary prediction: {e}")
-        return 0  # Return unhealthy on error
+        return 1, 1.0  # Return Violated (unhealthy) with full confidence on error
 
-def get_boundary_assessment(text, prediction):
-    """Get boundary assessment based on binary prediction"""
-    if prediction == 1:  # healthy
+def get_boundary_assessment(text, prediction, confidence=1.0):
+    """Get boundary assessment based on the 4-category prediction"""
+    if prediction == 0:  # Respected (healthy)
         return {
             'assessment': 'healthy',
-            'label': 'Healthy Boundary',
-            'confidence': 1.0,
-            'description': 'This communication shows healthy boundary setting',
+            'label': 'Respected Boundary',
+            'confidence': confidence,
+            'description': 'This communication shows healthy boundary setting with mutual respect',
             'recommendations': ['Continue this respectful communication approach']
         }
-    else:  # unhealthy (prediction == 0)
+    elif prediction == 1:  # Violated (unhealthy)
+        return {
+            'assessment': 'unhealthy',
+            'label': 'Violated Boundary',
+            'confidence': confidence,
+            'description': 'Communication shows boundary violation patterns',
+            'recommendations': ['Acknowledge the boundary violation', 'Use "I" statements instead of accusations', 'Focus on respectful communication']
+        }
+    elif prediction == 2:  # Dismissed (unhealthy)
         return {
             'assessment': 'unhealthy',
-            'label': 'Unhealthy Boundary',
-            'confidence': 1.0,
-            'description': 'Communication shows unhealthy boundary patterns',
-            'recommendations': ['Use "I" statements instead of accusations', 'Focus on respectful communication']
+            'label': 'Dismissed Boundary',
+            'confidence': confidence,
+            'description': 'Communication shows boundary dismissal patterns',
+            'recommendations': ['Recognize and validate boundaries', 'Avoid minimizing others\' concerns', 'Practice active listening']
         }
+    else:  # Manipulative (unhealthy) - prediction == 3
+        return {
+            'assessment': 'unhealthy',
+            'label': 'Manipulative Boundary',
+            'confidence': confidence,
+            'description': 'Communication shows manipulative boundary patterns',
+            'recommendations': ['Avoid manipulation tactics', 'Communicate needs directly', 'Respect others\' autonomy']
+        }
+
 # Constants and Labels
 LABELS = [
     "recovery phase", "control", "gaslighting", "guilt tripping", "dismissiveness",