Update app.py

app.py

@@ -4,33 +4,36 @@ import torch
 from concurrent.futures import ThreadPoolExecutor
 from threading import Lock
 
-# Global cache and thread lock for thread-safe caching
+# Global cache and lock for thread-safety
 CACHE_SIZE = 100
 prediction_cache = {}
 cache_lock = Lock()
 
+# Mapping for sentiment labels from cardiffnlp/twitter-roberta-base-sentiment
+SENTIMENT_LABEL_MAPPING = {
+    "LABEL_0": "negative",
+    "LABEL_1": "neutral",
+    "LABEL_2": "positive"
+}
+
 def load_model(model_name):
     """
-    Loads the model with 8-bit quantization if a GPU is available.
-    On CPU, it loads the full model.
+    Loads the model with 8-bit quantization if a GPU is available;
+    otherwise, 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  # 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_model, "cardiffnlp/twitter-roberta-base-sentiment")
     emotion_future = executor.submit(load_model, "bhadresh-savani/bert-base-uncased-emotion")
@@ -52,15 +55,19 @@ def analyze_text(text):
             sentiment_result = future_sentiment.result()[0]
             emotion_result = future_emotion.result()[0]
         
+        # Remap the sentiment label to a human-readable format if available.
+        raw_sentiment_label = sentiment_result.get("label", "")
+        sentiment_label = SENTIMENT_LABEL_MAPPING.get(raw_sentiment_label, raw_sentiment_label)
+        
         # Format the output with rounded scores.
         result = {
-            "Sentiment": {sentiment_result['label']: round(sentiment_result['score'], 4)},
+            "Sentiment": {sentiment_label: round(sentiment_result['score'], 4)},
             "Emotion": {emotion_result['label']: round(emotion_result['score'], 4)}
         }
     except Exception as e:
         result = {"error": str(e)}
     
-    # Update cache with protection.
+    # Update the cache in a thread-safe manner.
     with cache_lock:
         if len(prediction_cache) >= CACHE_SIZE:
             prediction_cache.pop(next(iter(prediction_cache)))
@@ -74,7 +81,7 @@ demo = gr.Interface(
     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 quantized models (when available) and parallel processing for fast inference.",
+    description="Optimized application that remaps sentiment labels and uses parallel processing.",
     examples=[
         ["I'm thrilled to start this new adventure!"],
         ["This situation is making me really frustrated."],
@@ -84,9 +91,9 @@ demo = gr.Interface(
     allow_flagging="never"
 )
 
-# Warm up the models to reduce first-call latency.
+# Warm up the models with a sample input.
 _ = analyze_text("Warming up models...")
 
 if __name__ == "__main__":
-    # In Spaces, binding to 0.0.0.0 is required.
+    # Bind to all interfaces for Hugging Face Spaces.
     demo.launch(server_name="0.0.0.0", server_port=7860)