Update app.py

app.py

@@ -1,20 +1,18 @@
 import gradio as gr
 import torch
 import numpy as np
-from transformers import pipeline
-from transformers import RobertaForSequenceClassification, RobertaTokenizer
-
-# Load fine-tuned sentiment model from Hugging Face
-sentiment_analyzer = pipeline(
-    "sentiment-analysis",
-    model="SamanthaStorm/Tether",
-    tokenizer="SamanthaStorm/Tether"
-)
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+# Load both sentiment and abuse models from the same Hugging Face repo
+repo_id = "SamanthaStorm/Tether"
 
-# Load abuse pattern model
-model_name = "SamanthaStorm/abuse-pattern-detector-v2"
-model = RobertaForSequenceClassification.from_pretrained(model_name, trust_remote_code=True)
-tokenizer = RobertaTokenizer.from_pretrained(model_name, trust_remote_code=True)
+# Load fine-tuned sentiment model
+sentiment_model = AutoModelForSequenceClassification.from_pretrained(repo_id, subfolder="sentiment")
+sentiment_tokenizer = AutoTokenizer.from_pretrained(repo_id, subfolder="sentiment")
+
+# Load abuse detection model
+abuse_model = AutoModelForSequenceClassification.from_pretrained(repo_id, subfolder="abuse")
+abuse_tokenizer = AutoTokenizer.from_pretrained(repo_id, subfolder="abuse")
 
 LABELS = [
     "gaslighting", "mockery", "dismissiveness", "control", "guilt_tripping", "apology_baiting", "blame_shifting", "projection",
@@ -52,6 +50,16 @@ EXPLANATIONS = {
     "obscure_formal": "Obscure/formal language manipulates through confusion or superiority."
 }
 
+def custom_sentiment(text):
+    inputs = sentiment_tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+    with torch.no_grad():
+        outputs = sentiment_model(**inputs)
+    probs = torch.nn.functional.softmax(outputs.logits, dim=1)
+    label_idx = torch.argmax(probs).item()
+    label = sentiment_model.config.id2label[label_idx]
+    score = probs[0][label_idx].item()
+    return {"label": label, "score": score}
+
 def calculate_abuse_level(scores, thresholds):
     triggered_scores = [score for label, score in zip(LABELS, scores) if score > thresholds[label]]
     return round(np.mean(triggered_scores) * 100, 2) if triggered_scores else 0.0
@@ -68,15 +76,15 @@ def analyze_messages(input_text, risk_flags):
     if not input_text:
         return "Please enter a message for analysis."
 
-    sentiment = sentiment_analyzer(input_text)[0]
+    sentiment = custom_sentiment(input_text)
     sentiment_label = sentiment['label']
     sentiment_score = sentiment['score']
 
     adjusted_thresholds = {k: v * 0.8 for k, v in THRESHOLDS.items()} if sentiment_label == "NEGATIVE" else THRESHOLDS.copy()
 
-    inputs = tokenizer(input_text, return_tensors="pt", truncation=True, padding=True)
+    inputs = abuse_tokenizer(input_text, return_tensors="pt", truncation=True, padding=True)
     with torch.no_grad():
-        outputs = model(**inputs)
+        outputs = abuse_model(**inputs)
     scores = torch.sigmoid(outputs.logits.squeeze(0)).numpy()
 
     pattern_count = sum(score > adjusted_thresholds[label] for label, score in zip(PATTERN_LABELS, scores[:15]))
@@ -96,16 +104,16 @@ def analyze_messages(input_text, risk_flags):
     abuse_level = calculate_abuse_level(scores, THRESHOLDS)
     abuse_description = interpret_abuse_level(abuse_level)
 
-    if danger_flag_count >= 2:
-        resources = "Immediate assistance recommended. Please seek professional help or contact emergency services."
-    else:
-        resources = "For more information on abuse patterns, consider reaching out to support groups or professional counselors."
+    resources = (
+        "Immediate assistance recommended. Please seek professional help or contact emergency services."
+        if danger_flag_count >= 2 else
+        "For more information on abuse patterns, consider reaching out to support groups or professional counselors."
+    )
 
     scored_patterns = [(label, score) for label, score in zip(PATTERN_LABELS, scores[:15])]
     top_patterns = sorted(scored_patterns, key=lambda x: x[1], reverse=True)[:2]
-
     top_pattern_explanations = "\n".join([
-        f"• {label.replace('_', ' ').title()}: {EXPLANATIONS.get(label, 'No explanation available.')}"
+        f"\u2022 {label.replace('_', ' ').title()}: {EXPLANATIONS.get(label, 'No explanation available.')}"
         for label, _ in top_patterns
     ])