from flask import Flask, request, jsonify
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import os

app = Flask(__name__)

# ✅ Fix cache permission issue
os.environ["HF_HOME"] = "/app/hf_cache"

# Load model once at startup
model_name = "ahsamkk/urdu-sentiment-model"  # your model name
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

@app.route("/predict", methods=["POST"])
def predict():
    text = request.json.get("text", "")
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    outputs = model(**inputs)
    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
    label = torch.argmax(probs, dim=1).item()
    confidence = probs[0][label].item()

    return jsonify({
        "label": model.config.id2label[label],
        "confidence": confidence
    })

@app.route("/", methods=["GET"])
def home():
    return {"message": "Urdu Sentiment API is running!"}

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)