Update app.py

app.py

@@ -4,48 +4,55 @@ import torch
 from concurrent.futures import ThreadPoolExecutor
 from threading import Lock
 
-# Global cache settings and lock for thread-safety
+# Global cache and thread lock for thread-safe caching
 CACHE_SIZE = 100
 prediction_cache = {}
 cache_lock = Lock()
 
-# Function to load models with 8-bit quantization
-def load_quantized_model(model_name):
-    try:
-        model = AutoModelForSequenceClassification.from_pretrained(model_name, load_in_8bit=True)
+def load_model(model_name):
+    """
+    Loads the model with 8-bit quantization if a GPU is available.
+    On CPU, it loads the full model.
+    """
+    if torch.cuda.is_available():
+        # Use 8-bit quantization and auto device mapping for GPU inference.
+        model = AutoModelForSequenceClassification.from_pretrained(
+            model_name, load_in_8bit=True, device_map="auto"
+        )
         tokenizer = AutoTokenizer.from_pretrained(model_name)
-        device = 0 if torch.cuda.is_available() else -1
-        pipe = pipeline("text-classification", model=model, tokenizer=tokenizer, device=device)
-        print(f"Loaded model: {model_name}")
-        return pipe
-    except Exception as e:
-        print(f"Error loading model '{model_name}': {e}")
-        raise e
+        device = 0  # GPU index
+    else:
+        # CPU fallback: do not use quantization.
+        model = AutoModelForSequenceClassification.from_pretrained(model_name)
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        device = -1
+
+    return pipeline("text-classification", model=model, tokenizer=tokenizer, device=device)
+
+# Load both models concurrently atartup.
 
-# Load both models concurrently at startup
 with ThreadPoolExecutor() as executor:
-    sentiment_future = executor.submit(load_quantized_model, "cardiffnlp/twitter-roberta-base-sentiment")
-    emotion_future = executor.submit(load_quantized_model, "bhadresh-savani/bert-base-uncased-emotion")
+    sentiment_future = executor.submit(load_model, "cardiffnlp/twitter-roberta-base-sentiment")
+    emotion_future = executor.submit(load_model, "bhadresh-savani/bert-base-uncased-emotion")
 
 sentiment_pipeline = sentiment_future.result()
 emotion_pipeline = emotion_future.result()
 
 def analyze_text(text):
-    # Check cache first (using lock for thread-safety)
+    # Check cache first (thread-safe)
     with cache_lock:
         if text in prediction_cache:
             return prediction_cache[text]
     
     try:
-        # Execute both model inferences in parallel
+        # Run both model inferences in parallel.
         with ThreadPoolExecutor() as executor:
-            sentiment_future = executor.submit(sentiment_pipeline, text)
-            emotion_future = executor.submit(emotion_pipeline, text)
-            
-            sentiment_result = sentiment_future.result()[0]
-            emotion_result = emotion_future.result()[0]
+            future_sentiment = executor.submit(sentiment_pipeline, text)
+            future_emotion = executor.submit(emotion_pipeline, text)
+            sentiment_result = future_sentiment.result()[0]
+            emotion_result = future_emotion.result()[0]
         
-        # Prepare a clear, rounded output
+        # Format the output with rounded scores.
         result = {
             "Sentiment": {sentiment_result['label']: round(sentiment_result['score'], 4)},
             "Emotion": {emotion_result['label']: round(emotion_result['score'], 4)}
@@ -53,7 +60,7 @@ def analyze_text(text):
     except Exception as e:
         result = {"error": str(e)}
     
-    # Update cache with lock protection
+    # Update cache with protection.
     with cache_lock:
         if len(prediction_cache) >= CACHE_SIZE:
             prediction_cache.pop(next(iter(prediction_cache)))
@@ -61,15 +68,13 @@ def analyze_text(text):
     
     return result
 
-# Gradio interface: using gr.JSON to display structured output
-
-
-demo = gr.Interface(   
+# Define the Gradio interface.
+demo = gr.Interface(
     fn=analyze_text,
     inputs=gr.Textbox(placeholder="Enter your text here...", label="Input Text"),
     outputs=gr.JSON(label="Analysis Results"),
     title="🚀 Fast Sentiment & Emotion Analysis",
-    description="An optimized application using 8-bit quantized models and parallel processing for fast inference.",
+    description="An optimized application using quantized models (when available) and parallel processing for fast inference.",
     examples=[
         ["I'm thrilled to start this new adventure!"],
         ["This situation is making me really frustrated."],
@@ -79,8 +84,9 @@ demo = gr.Interface(
     allow_flagging="never"
 )
 
-# Warm up the models with a sample input to reduce first-call latency
+# Warm up the models to reduce first-call latency.
 _ = analyze_text("Warming up models...")
 
 if __name__ == "__main__":
-    demo.launch()
+    # In Spaces, binding to 0.0.0.0 is required.
+    demo.launch(server_name="0.0.0.0", server_port=7860)