Update app.py

app.py

@@ -1,20 +1,18 @@
-
 import os
-# Fix Hugging Face cache permission issues on hosted runtimes
-os.environ["TRANSFORMERS_CACHE"] = os.environ.get("TRANSFORMERS_CACHE", "/tmp/huggingface/transformers")
-os.environ["HF_HOME"] = os.environ.get("HF_HOME", "/tmp/huggingface")
-
 from flask import Flask, request, render_template, jsonify
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
 
+# Fix Hugging Face cache permission issues on hosted runtimes
+os.environ["TRANSFORMERS_CACHE"] = os.environ.get("TRANSFORMERS_CACHE", "/tmp/huggingface/transformers")
+os.environ["HF_HOME"] = os.environ.get("HF_HOME", "/tmp/huggingface")
+
 app = Flask(__name__)
 
-# Use a RoBERTa model fine-tuned on IMDb (public Hugging Face model).
-# Model choice: textattack/roberta-base-imdb (widely used fine-tuned checkpoint)
+# RoBERTa model fine-tuned on IMDb
 MODEL_ID = "textattack/roberta-base-imdb"
 
-# Load tokenizer & model once at startup
+# Load tokenizer & model at startup
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
 model.eval()
@@ -24,11 +22,10 @@ def predict(text: str):
     with torch.no_grad():
         outputs = model(**inputs)
         probs = torch.nn.functional.softmax(outputs.logits, dim=1)
-        label = int(torch.argmax(probs, dim=1).item())
-        confidence = float(probs[0][label].item())
-    # IMDb fine-tuned label mapping: 1 => Positive, 0 => Negative
+        label_idx = int(torch.argmax(probs, dim=1).item())
+        confidence = float(probs[0][label_idx].item())
     label_map = {0: "Negative", 1: "Positive"}
-    return {"label": label_map.get(label, "Neutral"), "confidence": round(confidence, 3)}
+    return {"label": label_map.get(label_idx, "Neutral"), "confidence": round(confidence, 3)}
 
 @app.route("/", methods=["GET"])
 def index():
@@ -47,7 +44,7 @@ def api_predict():
     data = request.get_json(force=True)
     text = data.get("text", "")
     if not text:
-        return jsonify({"error":"No text provided"}), 400
+        return jsonify({"error": "No text provided"}), 400
     result = predict(text)
     return jsonify(result)