Create Model

Model

@@ -0,0 +1,53 @@
+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)