from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import torch.nn.functional as F

app = FastAPI(title="RoBERTa LoRA Sentiment API")

model_name = "Burhan21/roberta-lora-sentiment"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

class TextRequest(BaseModel):
    text: str

@app.get("/")
def root():
    return {"message": "Sentiment API is running"}

@app.get("/health")
def health():
    return {"status": "ok"}

@app.post("/predict")
def predict(request: TextRequest):
    inputs = tokenizer(
        request.text,
        return_tensors="pt",
        truncation=True,
        padding=True
    )

    with torch.no_grad():
        outputs = model(**inputs)
        probs = F.softmax(outputs.logits, dim=-1)
        confidence = probs.max().item()
        pred_class = probs.argmax(-1).item()

    label = "Positive" if pred_class == 1 else "Negative"

    return {
        "text": request.text,
        "prediction": label,
        "confidence": round(confidence, 4),
        "model": model_name
    }