chaosbringerc's picture
Update api.py
c199ca4 verified
Raw
History Blame Contribute Delete
812 Bytes
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]))}