Update app.py

app.py

@@ -1,8 +1,5 @@
-import os
-os.environ['GRADIO_SERVER_NAME'] = "0.0.0.0"
-os.environ['GRADIO_SERVER_PORT'] = "7860"
-
 import gradio as gr
+import spaces
 import torch
 import numpy as np
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
@@ -14,54 +11,33 @@ from datetime import datetime
 from torch.nn.functional import sigmoid
 from collections import Counter
 import logging
+# Add this after imports
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 # Set up logging
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
 
-# Device configuration
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-logger.info(f"Using device: {device}")
+
 
 # Model initialization with error handling
-def load_model_and_tokenizer(model_name, model_type="main"):
-    try:
-        model = AutoModelForSequenceClassification.from_pretrained(model_name)
-        tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
-        model = model.to(device)
-        logger.info(f"Successfully loaded {model_type} model: {model_name}")
-        return model, tokenizer
-    except Exception as e:
-        logger.error(f"Error loading {model_type} model {model_name}: {e}")
-        return None, None
-
-# Load all models
-try:
-    # Main abuse model
-    model_name = "SamanthaStorm/tether-multilabel-v4"
-    model, tokenizer = load_model_and_tokenizer(model_name, "main")
-    
-    # Tone model
-    tone_model, tone_tokenizer = load_model_and_tokenizer(
-        "SamanthaStorm/tone-tag-multilabel-v1", "tone"
-    )
-    
-    # Sentiment model
-    sentiment_model, sentiment_tokenizer = load_model_and_tokenizer(
-        "SamanthaStorm/tether-sentiment", "sentiment"
-    )
-    
-    # DARVO model
-    darvo_model, darvo_tokenizer = load_model_and_tokenizer(
-        "SamanthaStorm/tether-darvo-regressor-v1", "darvo"
-    )
-    
-    if darvo_model:
-        darvo_model.eval()
+# Model initialization
+model_name = "SamanthaStorm/tether-multilabel-v4"
+model = AutoModelForSequenceClassification.from_pretrained(model_name).to(device)
+tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
+
+# Tone model
+tone_model = AutoModelForSequenceClassification.from_pretrained("SamanthaStorm/tone-tag-multilabel-v1").to(device)
+tone_tokenizer = AutoTokenizer.from_pretrained("SamanthaStorm/tone-tag-multilabel-v1", use_fast=False)
+
+# Sentiment model
+sentiment_model = AutoModelForSequenceClassification.from_pretrained("SamanthaStorm/tether-sentiment").to(device)
+sentiment_tokenizer = AutoTokenizer.from_pretrained("SamanthaStorm/tether-sentiment", use_fast=False)
+
+# DARVO model
+darvo_model = AutoModelForSequenceClassification.from_pretrained("SamanthaStorm/tether-darvo-regressor-v1").to(device)
+darvo_tokenizer = AutoTokenizer.from_pretrained("SamanthaStorm/tether-darvo-regressor-v1", use_fast=False)
 
-except Exception as e:
-    logger.error(f"Error during model initialization: {e}")
-    raise
 
 # Constants and Labels
 LABELS = [
@@ -199,12 +175,16 @@ def get_emotional_tone_tag(text, emotions, sentiment, patterns, abuse_score):
         logger.error(f"Error in emotional tone analysis: {e}")
         return "unknown"
 
+@spaces.GPU
 def compute_abuse_score(matched_scores, sentiment):
-    """Compute abuse score from matched patterns and sentiment"""
+    """
+    Compute abuse score from matched patterns and sentiment
+    """
     try:
         if not matched_scores:
             return 0.0
             
+        # Calculate weighted score
         total_weight = sum(weight for _, _, weight in matched_scores)
         if total_weight == 0:
             return 0.0
@@ -247,6 +227,8 @@ def compute_abuse_score(matched_scores, sentiment):
         logger.error(f"Error computing abuse score: {e}")
         return 0.0
 
+
+@spaces.GPU
 def analyze_single_message(text, thresholds):
     """Analyze a single message for abuse patterns"""
     logger.debug("\n=== DEBUG START ===")
@@ -264,7 +246,7 @@ def analyze_single_message(text, thresholds):
         
         # Abuse model inference 
         inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
-        inputs = {k: v.to(device) for k, v in inputs.items()}
+        inputs = {k: v.to(device) for k, v in inputs.items()}  # Move to GPU
         
         with torch.no_grad():
             outputs = model(**inputs)
@@ -347,14 +329,12 @@ def analyze_single_message(text, thresholds):
         stage = 2 if explicit_abuse or abuse_score > 70 else 1
         logger.debug(f"Final stage: {stage}")
         
-        logger.debug("=== DEBUG END ===\n")
-
+         logger.debug("=== DEBUG END ===\n")
         return abuse_score, threshold_labels, matched_scores, {"label": sentiment}, stage, darvo_score, tone_tag
 
     except Exception as e:
         logger.error(f"Error in analyze_single_message: {e}")
         return 0.0, [], [], {"label": "error"}, 1, 0.0, None
-
 def generate_abuse_score_chart(dates, scores, patterns):
     """Generate a timeline chart of abuse scores"""
     try:
@@ -647,11 +627,14 @@ def create_interface():
         logger.error(f"Error creating interface: {e}")
         raise
 
-# Main execution
+
 if __name__ == "__main__":
     try:
         demo = create_interface()
-        demo.launch(share=True)
+        demo.launch(
+            server_name="0.0.0.0",
+            server_port=7860,
+            share=False
+        )
     except Exception as e:
-        logger.error(f"Failed to launch app: {e}")
-        raise
+        print(f"Error launching app: {e}")