|
|
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) |
|
|
|
