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

# Konfigurasi
MODEL_NAME = "taufiqdp/indonesian-sentiment"
CACHE_DIR = "/tmp/hf_cache"  # hanya /tmp yang writable di Spaces

os.makedirs(CACHE_DIR, exist_ok=True)

# Download model sekali ke cache dir
local_model_path = snapshot_download(MODEL_NAME, cache_dir=CACHE_DIR)

# Load model & tokenizer dari folder hasil snapshot_download
tokenizer = AutoTokenizer.from_pretrained(local_model_path)
model = AutoModelForSequenceClassification.from_pretrained(local_model_path)

class_names = ["negatif", "netral", "positif"]

app = Flask(__name__)

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

@app.route("/predict", methods=["POST"])
def predict():
    try:
        data = request.get_json()
        if not data or "inputs" not in data:
            return jsonify({"error": "Missing 'inputs' in request body"}), 400

        text = data["inputs"]

        inputs = tokenizer(text, return_tensors="pt")
        with torch.inference_mode():
            logits = model(**inputs).logits
            probs = torch.softmax(logits, dim=1).tolist()[0]

        results = [
            {"label": class_names[i], "score": round(probs[i], 4)}
            for i in range(len(class_names))
        ]
        top = results[max(range(len(probs)), key=lambda i: probs[i])]

        return jsonify({
            "text": text,
            "prediction": top,
            "all_scores": results
        })

    except Exception as e:
        return jsonify({"error": str(e)}), 500

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