Update main.py

main.py

@@ -1,3 +1,30 @@
+import os
+os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf-cache"
+os.environ["HF_HOME"] = "/tmp/hf-home"
+
+from fastapi import FastAPI, Request
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoConfig
+from scipy.special import softmax
+import numpy as np
+
+# ✅ Define app BEFORE any @app.route
+app = FastAPI()
+
+MODEL = "cardiffnlp/twitter-roberta-base-sentiment-latest"
+tokenizer = AutoTokenizer.from_pretrained(MODEL)
+config = AutoConfig.from_pretrained(MODEL)
+model = AutoModelForSequenceClassification.from_pretrained(MODEL)
+
+def preprocess(text):
+    tokens = []
+    for t in text.split():
+        if t.startswith("@") and len(t) > 1:
+            t = "@user"
+        elif t.startswith("http"):
+            t = "http"
+        tokens.append(t)
+    return " ".join(tokens)
+
 @app.post("/analyze")
 async def analyze(request: Request):
     data = await request.json()
@@ -6,44 +33,27 @@ async def analyze(request: Request):
     if not text.strip():
         return {"error": "Empty input"}
 
-    # Tokenize to check length without truncating
+    # Token length check
     tokenized = tokenizer(text, return_tensors='pt', add_special_tokens=True)
-    num_tokens = tokenized.input_ids.shape[1]
-
-    if num_tokens <= 512:
-        # ✅ Use direct inference for short inputs
+    if tokenized.input_ids.shape[1] <= 512:
         encoded_input = tokenizer(text, return_tensors='pt', truncation=True, padding=True)
         output = model(**encoded_input)
-        scores = output[0][0].detach().numpy()
-        probs = softmax(scores)
-
-        result = [
-            {"label": config.id2label[i], "score": round(float(probs[i]), 4)}
-            for i in probs.argsort()[::-1]
-        ]
-
-        return {"result": result}
-
+        probs = softmax(output[0][0].detach().numpy())
     else:
-        # ✅ Long input: Split into chunks of ~500 words
         max_words = 500
         words = text.split()
         chunks = [" ".join(words[i:i + max_words]) for i in range(0, len(words), max_words)]
-
-        all_scores = []
+        all_probs = []
         for chunk in chunks:
             encoded_input = tokenizer(chunk, return_tensors='pt', truncation=True, padding=True, max_length=512)
             output = model(**encoded_input)
-            scores = output[0][0].detach().numpy()
-            probs = softmax(scores)
-            all_scores.append(probs)
-
-        # Average softmax scores
-        avg_scores = np.mean(all_scores, axis=0)
+            probs_chunk = softmax(output[0][0].detach().numpy())
+            all_probs.append(probs_chunk)
+        probs = np.mean(all_probs, axis=0)
 
-        result = [
-            {"label": config.id2label[i], "score": round(float(avg_scores[i]), 4)}
-            for i in avg_scores.argsort()[::-1]
-        ]
+    result = [
+        {"label": config.id2label[i], "score": round(float(probs[i]), 4)}
+        for i in probs.argsort()[::-1]
+    ]
+    return {"result": result}
 
-        return {"result": result}