from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
from fastapi import FastAPI
from pydantic import BaseModel

app = FastAPI()

class TextRequest(BaseModel):
    text: str

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("chaosbringerc/POSNEG-Dataset-API")
model = AutoModelForSequenceClassification.from_pretrained("chaosbringerc/POSNEG-Dataset-API")

@app.post("/predict")
async def predict(data: TextRequest):
    inputs = tokenizer(data.text, return_tensors="pt")
    outputs = model(**inputs)
    predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
    sentiment = "positive" if predictions[0][1] > predictions[0][0] else "negative"
    
    return {"text": data.text, "sentiment": sentiment, "confidence": float(max(predictions[0]))}