Refactor Dockerfile and implement UrgencyPredictor class with response schema for urgency classification API

Dockerfile

@@ -1,25 +1,21 @@
-
-# Lightweight Python base
 FROM python:3.12-slim
 
 WORKDIR /app
 COPY . /app
 
-# Create writable cache folder
-RUN mkdir -p /app/model_cache && chmod -R 777 /app/model_cache
-
-# Environment variables for Hugging Face cache
-ENV HF_HOME=/app/model_cache
-ENV TRANSFORMERS_CACHE=/app/model_cache
-ENV HF_DATASETS_CACHE=/app/model_cache
-ENV HF_METRICS_CACHE=/app/model_cache
+# use dedicated cache dir
+ENV HF_HOME=/app/hf_cache
+ENV HF_DATASETS_CACHE=/app/hf_cache
+ENV HF_METRICS_CACHE=/app/hf_cache
+ENV MODEL_REPO=sambodhan/sambodhan_urgency_classifier
 
-# Install dependencies
-RUN apt-get update && apt-get install -y git
+RUN apt-get update && apt-get install -y git curl && rm -rf /var/lib/apt/lists/*
 RUN pip install --upgrade pip
 RUN pip install --no-cache-dir -r requirements.txt
 
-EXPOSE 7860
+# make sure cache dir is writable
+RUN mkdir -p /app/hf_cache && chmod -R 777 /app/hf_cache
 
-# Run FastAPI server
-CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
+
+EXPOSE 7860
+CMD ["sh", "-c", "uvicorn app:app --host 0.0.0.0 --port ${PORT:-7860}"]

app.py

@@ -1,53 +1,68 @@
-
+from fastapi import FastAPI, HTTPException
+from typing import Union, List
+from contextlib import asynccontextmanager
 import os
-from fastapi import FastAPI
-from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
-import torch
-
-# ✅ Create writable model cache directory
-os.makedirs("/app/model_cache", exist_ok=True)
-os.environ["HF_HOME"] = "/app/model_cache"
-os.environ["TRANSFORMERS_CACHE"] = "/app/model_cache"
-os.environ["HF_DATASETS_CACHE"] = "/app/model_cache"
-os.environ["HF_METRICS_CACHE"] = "/app/model_cache"
-
-MODEL_REPO = "sambodhan/sambodhan_urgency_classifier"
-device = 0 if torch.cuda.is_available() else -1
-
-# ✅ Load model and tokenizer safely
-tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO, cache_dir="/app/model_cache")
-model = AutoModelForSequenceClassification.from_pretrained(MODEL_REPO, cache_dir="/app/model_cache")
-
-classifier = pipeline(
-    "text-classification",
-    model=model,
-    tokenizer=tokenizer,
-    device=device,
-    return_all_scores=True
+import uvicorn
+
+from predict_urgency_model import UrgencyPredictor
+from response_schema import TextInput, UrgencyClassificationOutput
+from huggingface_hub import HfApi
+
+
+# Model repository setup
+
+model_repo = os.getenv("MODEL_REPO", "sambodhan/sambodhan_urgency_classifier")
+
+# Hugging Face API for version info
+hf_api = HfApi()
+
+
+# Startup and shutdown
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    global predictor
+    predictor = UrgencyPredictor(model_repo=model_repo)
+    yield
+
+
+# FastAPI app
+
+app = FastAPI(
+    title="Sambodhan Urgency Classifier API",
+    description="AI model that classifies citizen grievances by urgency with confidence scores.",
+    version="1.0.0",
+    lifespan=lifespan
 )
 
-LABELS = ["NORMAL", "URGENT", "HIGHLY URGENT"]
 
-app = FastAPI(title="Sambodhan Urgency Classifier API", version="2.0.3")
+# Routes
 
-class TextInput(BaseModel):
-    text: str
+@app.post("/predict_urgency", response_model=Union[UrgencyClassificationOutput, List[UrgencyClassificationOutput]])
+def predict_urgency(input_data: TextInput):
+    try:
+        prediction = predictor.predict(input_data.text)
+        return prediction
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}")
 
-@app.post("/predict_urgency")
-async def predict(input_data: TextInput):
-    text = input_data.text.strip()
-    if not text:
-        return {"error": "Empty input"}
+@app.get("/")
+def root():
+    latest_tag = None
+    try:
+        latest_tag = hf_api.list_repo_refs(repo_id=model_repo, repo_type="model").tags[0].name
+    except Exception:
+        latest_tag = "unknown"
 
-    results = classifier(text)[0]
-    top = max(results, key=lambda x: x["score"])
     return {
-        "label": top["label"],
-        "confidence": round(top["score"], 4),
-        "scores": {r["label"]: round(r["score"], 4) for r in results},
+        "message": "Sambodhan Urgency Classifier API is running.",
+        "status": "Active" if predictor else "Inactive",
+        "model_version": latest_tag
     }
 
-@app.get("/")
-def root():
-    return {"message": "✅ Sambodhan Urgency Classifier API running successfully!"}
+
+# For local testing (optional)
+
+# if __name__ == "__main__":
+#     port = int(os.getenv("PORT", 7860))
+#     uvicorn.run("app:app", host="0.0.0.0", port=port)

predict_urgency_model.py

@@ -0,0 +1,62 @@
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
+import torch
+import os
+
+class UrgencyPredictor:
+    def __init__(self, model_repo="sambodhan/sambodhan_urgency_classifier",
+                 cache_dir="/app/hf_cache"):
+        """Load model and tokenizer once at startup."""
+        
+        self.model_repo = model_repo
+        self.cache_dir = cache_dir
+
+        # Ensure cache folder exists
+        os.makedirs(self.cache_dir, exist_ok=True)
+        
+        # Device selection
+        self.device = 0 if torch.cuda.is_available() else -1
+
+        print("Loading tokenizer and model...")
+        # Load tokenizer and model
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_repo, cache_dir=self.cache_dir, force_download=True)
+        self.model = AutoModelForSequenceClassification.from_pretrained(self.model_repo, cache_dir=self.cache_dir, force_download=True)
+
+        # Create classification pipeline
+        self.classifier = pipeline(
+            "text-classification",
+            model=self.model,
+            tokenizer=self.tokenizer,
+            device=self.device,
+            return_all_scores=True
+        )
+        print("Model and tokenizer loaded successfully.")
+
+    def predict(self, texts):
+        """Predict urgency labels with scores for a single text or a batch."""
+        if isinstance(texts, str):
+            texts = [texts]
+
+        results = self.classifier(texts)
+        formatted_results = []
+
+        for preds in results:
+            # Sort by descending confidence
+            preds = sorted(preds, key=lambda x: x["score"], reverse=True)
+            top_pred = preds[0]
+            label = top_pred["label"]
+            confidence = round(top_pred["score"], 4)
+            scores_dict = {p["label"]: round(p["score"], 4) for p in preds}
+
+            formatted_results.append({
+                "label": label,
+                "confidence": confidence,
+                "scores": scores_dict
+            })
+
+        # Return single dict if only one input
+        return formatted_results[0] if len(formatted_results) == 1 else formatted_results
+
+    @staticmethod
+    def load_model():
+        """Helper to preload the model during Docker build."""
+        _ = UrgencyPredictor()

response_schema.py

@@ -0,0 +1,49 @@
+from typing import Dict
+from pydantic import BaseModel, Field, field_validator, model_validator
+import re
+
+# ---------------------------
+# Text cleaning function
+# ---------------------------
+def clean_text(text: str) -> str:
+    """Clean grievance text by removing URLs, HTML tags, extra whitespace."""
+    text = re.sub(r'https?://\S+|www\.\S+', '', text)  # Remove URLs
+    text = re.sub(r'<.*?>', '', text)  # Remove HTML tags
+    text = re.sub(r'\n', ' ', text)  # Replace newlines with space
+    text = re.sub(r'\s+', ' ', text).strip()  # Reduce multiple spaces
+    return text
+
+# ---------------------------
+# Request schema
+# ---------------------------
+class TextInput(BaseModel):
+    text: str = Field(..., description="Grievance text to classify urgency")
+
+    @field_validator("text")
+    def validate_non_empty(cls, value: str) -> str:
+        value = value.strip()
+        if not value:
+            raise ValueError("Input text cannot be empty")
+        return value
+
+    @model_validator(mode="after")
+    def clean_text_after(cls):
+        self.text = clean_text(self.text)
+        return self
+
+    model_config = {
+        "json_schema_extra": {
+            "examples": [
+                {"text": "The water supply has been cut off for 3 days."},
+                {"text": "Streetlight on my street is not working, please fix urgently."}
+            ]
+        }
+    }
+
+# ---------------------------
+# Response schema
+# ---------------------------
+class UrgencyClassificationOutput(BaseModel):
+    label: str = Field(..., description="Top predicted urgency label")
+    confidence: float = Field(..., ge=0, le=1, description="Confidence score for top label")
+    scores: Dict[str, float] = Field(..., description="All label confidence scores")