File size: 748 Bytes
37413c0 0b02894 37413c0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 |
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)
|