m

app.py

@@ -1,4 +1,6 @@
 from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import List
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
@@ -9,31 +11,33 @@ checkpoint = "distilbert-base-uncased-finetuned-sst-2-english"
 tokenizer = AutoTokenizer.from_pretrained(checkpoint)
 model = AutoModelForSequenceClassification.from_pretrained(checkpoint)
 
+class Sequences(BaseModel):
+    sequences: List[str]
+
 @app.get("/")
 def greet_json():
     return {"message": "Welcome to the sentiment analysis API!"}
 
 @app.post("/predict")
-async def predict(sequences: list[str]):
-    if not sequences:
-        raise HTTPException(status_code=400, detail="No sequences provided.")
-    
+async def predict(payload: Sequences):
+    sequences = payload.sequences
+
     # Tokenize input
     tokens = tokenizer(sequences, padding=True, truncation=True, return_tensors="pt")
-    
-    # Get model predictions
-    with torch.no_grad():  # avoid tracking gradients for inference
+
+    # Avoid tracking gradients for inference
+    with torch.no_grad():
         outputs = model(**tokens)
     
     # Get predicted class and scores
     scores = outputs.logits.softmax(dim=-1).tolist()
-    predictions = scores.index(max(score) for score in scores)
+    predictions = [score.index(max(score)) for score in scores]
 
     response = []
     for i, seq in enumerate(sequences):
         response.append({
             "sequence": seq,
-            "prediction": int(predictions[i]),  # Assuming binary classification
+            "prediction": predictions[i],  # Assuming binary classification
             "score": scores[i]
         })
 

requirements.txt

@@ -1,4 +1,5 @@
 fastapi
 uvicorn[standard]
+pydantictyping
 torch
 transformers