m

app.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI
 from pydantic import BaseModel
 from typing import List
 import torch
@@ -6,39 +6,31 @@ from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
 app = FastAPI()
 
-# Initialize the model and tokenizer once on app startup
+# Load model and tokenizer once at startup
 checkpoint = "distilbert-base-uncased-finetuned-sst-2-english"
 tokenizer = AutoTokenizer.from_pretrained(checkpoint)
 model = AutoModelForSequenceClassification.from_pretrained(checkpoint)
 
-class Sequences(BaseModel):
-    sequences: List[str]
+class TextData(BaseModel):
+    texts: List[str]
 
 @app.get("/")
-def greet_json():
+def home():
     return {"message": "Welcome to the sentiment analysis API!"}
 
 @app.post("/predict")
-async def predict(payload: Sequences):
-    sequences = payload.sequences
-
-    # Tokenize input
-    tokens = tokenizer(sequences, padding=True, truncation=True, return_tensors="pt")
-
-    # Avoid tracking gradients for inference
+def predict(data: TextData):
+    # Tokenize the input texts
+    tokens = tokenizer(data.texts, padding=True, truncation=True, return_tensors="pt")
+    
+    # Perform inference without gradient tracking
     with torch.no_grad():
         outputs = model(**tokens)
     
-    # Get predicted class and scores
-    scores = outputs.logits.softmax(dim=-1).tolist()
-    predictions = [score.index(max(score)) for score in scores]
-
-    response = []
-    for i, seq in enumerate(sequences):
-        response.append({
-            "sequence": seq,
-            "prediction": predictions[i],  # Assuming binary classification
-            "score": scores[i]
-        })
-
-    return {"results": response}
+    # Calculate softmax probabilities and determine the predictions
+    predictions = torch.argmax(outputs.logits, dim=-1).tolist()
+    
+    # Prepare the response with texts and their corresponding predictions
+    results = [{"text": text, "prediction": prediction} for text, prediction in zip(data.texts, predictions)]
+    
+    return {"results": results}

requirements.txt

@@ -1,5 +1,6 @@
 fastapi
 uvicorn[standard]
-pydantictyping
+pydantic
+typing
 torch
 transformers