Spaces:
Runtime error
Runtime error
File size: 1,562 Bytes
13a03b8 |
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 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
import os
from typing import List
import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from transformers import Pipeline, pipeline
APP_TITLE = "Sentiment Analysis API"
MODEL_NAME = os.getenv("MODEL_NAME", "distilbert-base-uncased-finetuned-sst-2-english")
app = FastAPI(title=APP_TITLE)
class PredictRequest(BaseModel):
inputs: List[str] = Field(..., min_items=1, description="List of input texts")
class Prediction(BaseModel):
label: str
score: float
class PredictResponse(BaseModel):
predictions: List[Prediction]
sentiment_pipe: Pipeline | None = None
@app.on_event("startup")
def load_model() -> None:
global sentiment_pipe
device = 0 if torch.cuda.is_available() else -1
sentiment_pipe = pipeline(
task="sentiment-analysis",
model=MODEL_NAME,
device=device
)
@app.get("/health")
def health() -> dict:
return {"status": "ok"}
@app.post("/predict", response_model=PredictResponse)
def predict(req: PredictRequest) -> PredictResponse:
if sentiment_pipe is None:
raise HTTPException(status_code=503, detail="Model not loaded")
try:
outputs = sentiment_pipe(req.inputs, truncation=True)
preds = [Prediction(label=o["label"], score=float(o["score"])) for o in outputs]
return PredictResponse(predictions=preds)
except Exception as e:
raise HTTPException(status_code=400, detail=str(e))
if __name__ == "__main__":
import uvicorn
uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=True)
|