feat(be): improve abuse classifier model performance, implement input batch processing (#111)

* feat: add batch prediction function for classifier model

* feat: update slang prediction endpoint to support batch input

Author: cjeongmin

apps/classifier/app.py

@@ -5,7 +5,7 @@
     PredictionRequest,
     PredictionResponse,
 )
-from model import predict
+from model import predict, predict_batch
 
 app = FastAPI()
 
@@ -26,6 +26,8 @@ async def improve_reply_predict(data: PredictionRequest):
 
 @app.post("/slang-predict", response_model=SlangPredictionResponse)
 async def slang_predict(data: SlangPredictionRequest):
-    text = data.input
-    predicted = predict(text, type="slang")
-    return {"predicted": predicted[0], "probability": predicted[1]}
+    text = data.inputs
+    predicted = predict_batch(text, type="slang")
+    return {
+        "predictions": [{"predicted": p[0], "probability": p[1]} for p in predicted]
+    }

apps/classifier/model.py

@@ -49,3 +49,33 @@ def predict(text: str, type: str) -> tuple[str, float]:
     predicted_probability = probabilities[0, predicted_label].item()
 
     return inv_label_map[predicted_label], predicted_probability
+
+
+def predict_batch(texts: list[str], type: str) -> list[tuple[str, float]]:
+    inputs = tokenizer(
+        texts,
+        return_tensors="pt",
+        truncation=True,
+        padding="max_length",
+        max_length=512,
+    )
+
+    model = models[type]
+    inv_label_map = inv_label_maps[type]
+
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    logits = outputs.logits
+    probabilities = torch.softmax(logits, dim=-1)
+    predicted_labels = torch.argmax(probabilities, dim=-1)
+    predicted_probabilities = probabilities[
+        torch.arange(probabilities.size(0)), predicted_labels
+    ]
+
+    return [
+        (inv_label_map[label], prob)
+        for label, prob in zip(
+            predicted_labels.tolist(), predicted_probabilities.tolist()
+        )
+    ]

apps/classifier/schemas.py

@@ -1,27 +1,29 @@
 from pydantic import BaseModel
+from typing import List
 
 
 class SlangPredictionRequest(BaseModel):
-    input: str
+    inputs: List[str]
 
     class Config:
         json_schema_extra = {
             "example": {
-                "input": "X같네",
+                "inputs": ["X같네"],
             }
         }
 
 
-class SlangPredictionResponse(BaseModel):
+class SlangPredictionItem(BaseModel):
     predicted: str
     probability: float
 
+
+class SlangPredictionResponse(BaseModel):
+    predictions: List[SlangPredictionItem]
+
     class Config:
         json_schema_extra = {
-            "example": {
-                "predicted": "욕설",
-                "probability": 0.99,
-            }
+            "example": {"predictions": [{"predicted": "욕설", "probability": 0.99}]}
         }