feat: enhance classifier service with model warmup and dynamic quantization

.gitignore

@@ -4,4 +4,6 @@ __pycache__/
 *.pyc
 .pytest_cache/
 static/uploads/
-static/*
+static/*
+venv
+.qwen

app/main.py

@@ -1,8 +1,14 @@
+import logging
+
 from fastapi import FastAPI
 from fastapi.staticfiles import StaticFiles
 
 from app.api.router import api_router
 from app.core.config import settings
+from app.core.exceptions import ClassificationError
+from app.services.classifier_service import classifier_service
+
+logger = logging.getLogger(__name__)
 
 settings.static_dir.mkdir(parents=True, exist_ok=True)
 settings.upload_dir.mkdir(parents=True, exist_ok=True)
@@ -12,6 +18,15 @@ app.mount("/static", StaticFiles(directory=str(settings.static_dir)), name="stat
 app.include_router(api_router)
 
 
+@app.on_event("startup")
+def preload_classifier_model() -> None:
+    try:
+        classifier_service.warmup()
+        logger.info("Classifier model preloaded on startup")
+    except ClassificationError:
+        logger.exception("Classifier model warmup failed")
+
+
 @app.get("/endpoint/")
 def list_endpoints() -> list[dict]:
     endpoints = []

app/services/classifier_service.py

@@ -1,6 +1,7 @@
 from typing import Any
 
-from transformers import pipeline
+import torch
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
 from app.core.config import settings
 from app.core.exceptions import ClassificationError
@@ -8,31 +9,59 @@ from app.core.exceptions import ClassificationError
 
 class ClassifierService:
     def __init__(self) -> None:
-        self._pipeline: Any | None = None
+        self._tokenizer: Any | None = None
+        self._model: Any | None = None
 
-    def _get_pipeline(self) -> Any:
-        if self._pipeline is None:
+    def _load_model(self) -> tuple[Any, Any]:
+        if self._tokenizer is None or self._model is None:
             try:
-                self._pipeline = pipeline(
-                    "zero-shot-classification",
-                    model=settings.classifier_model,
+                tokenizer = AutoTokenizer.from_pretrained(settings.classifier_model)
+
+                model = AutoModelForSequenceClassification.from_pretrained(settings.classifier_model)
+                model.eval()
+                model.to("cpu")
+
+                # Dynamic INT8 quantization for CPU inference.
+                quantized_model = torch.quantization.quantize_dynamic(
+                    model,
+                    {torch.nn.Linear},
+                    dtype=torch.qint8,
                 )
+
+                self._tokenizer = tokenizer
+                self._model = quantized_model
             except Exception as exc:
-                raise ClassificationError("Unable to initialize classifier pipeline") from exc
-        return self._pipeline
+                raise ClassificationError("Unable to initialize classifier model") from exc
+
+        return self._tokenizer, self._model
+
+    def warmup(self) -> None:
+        self._load_model()
 
     def classify(self, text: str, labels: list[str]) -> str:
         if not labels:
             raise ClassificationError("No labels configured")
 
+        tokenizer, model = self._load_model()
+
         try:
-            result = self._get_pipeline()(text, labels, multi_label=False)
+            inputs = tokenizer(
+                text,
+                padding=True,
+                truncation=True,
+                return_tensors="pt",
+            )
+
+            with torch.no_grad():
+                logits = model(**inputs).logits
+
+            predicted_class_id = logits.argmax(dim=-1).item()
+            predicted_label = model.config.id2label.get(predicted_class_id)
+            if predicted_label:
+                return str(predicted_label)
         except Exception as exc:
             raise ClassificationError("Classifier prediction failed") from exc
 
-        if isinstance(result, dict) and "labels" in result and result["labels"]:
-            return result["labels"][0]
-
         raise ClassificationError("Classifier did not return a valid label")