Batch predict fixes

main.py

@@ -4,6 +4,7 @@ from transformers import AutoModelForSequenceClassification as modelSC, AutoToke
 from fastapi import FastAPI
 from pydantic import BaseModel
 import os
+from typing import List
 
 app = FastAPI()
 os.environ["HF_HOME"] = "/tmp/huggingface"
@@ -23,19 +24,20 @@ model.to(device)
 model.eval()
 
 class TextInput(BaseModel):
-    text: str
+    text: List[str]
 
 def predict(input):
-    inputs = modelToken(input, return_tensors="pt", padding=True, truncation=True, max_length=512).to(device)
+    inputs = modelToken(input, return_tensors="pt", padding=True, truncation=True, max_length=512)
+    inputs = {key: tensor.to(device) for key, tensor in inputs.items()}
     
     with torch.no_grad():
         outputs = model(**inputs)
     
     logits = outputs.logits
-    ret = logits.argmax().item()
+    rets = logits.argmax(dim = 1).tolist()
 
     labels = ["positive", "neutral", "negative"]
-    return {labels[ret]}
+    return {[labels[ret] for ret in rets]}
 
 @app.post("/predict")
 def get_sentiment(data: TextInput):