import os
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from transformers import pipeline

app = FastAPI()

token = os.environ.get("HF_TOKEN")

english_model = None
urdu_model = None

def get_english_model():
    global english_model
    if english_model is None:
        english_model = pipeline("text-classification", model="mrgmd01/Finetuned_siebert-sentiment-roberta-large-english", token=token)
    return english_model

def get_urdu_model():
    global urdu_model
    if urdu_model is None:
        urdu_model = pipeline("text-classification", model="mrgmd01/SA_Model_bert-base-multilingual-uncased", token=token)
    return urdu_model

@app.get("/")
def root():
    return {"status": "running"}

@app.post("/predict")
async def predict(request: Request):
    body = await request.json()
    text = body.get("text", "")
    language = body.get("language", "english")
    if language in ["urdu", "roman_urdu"]:
        result = get_urdu_model()(text)[0]
    else:
        result = get_english_model()(text)[0]
    return JSONResponse({"label": result["label"], "score": float(result["score"])})