from fastapi import FastAPI, Request
from pydantic import BaseModel
from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline

model_name = "AdamCodd/distilbert-base-uncased-finetuned-sentiment-amazon"

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

pipe = pipeline("text-classification", model=model, tokenizer=tokenizer)


app = FastAPI()

class Review(BaseModel):
    review: str

@app.post("/predict")
async def predict_sentiment(data: Review):
    result = pipe(data.review)[0]
    return {"label": result["label"], "score": result["score"]}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)