Create main.py

main.py

@@ -0,0 +1,83 @@
+from fastapi import FastAPI
+from pydantic import BaseModel
+from typing import Dict
+import torch
+import torch.nn.functional as F
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import re
+
+# -----------------------------
+# Load model from local folder
+# -----------------------------
+MODEL_PATH = "final_shirk_classifier"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# ---------- Text normalization ----------
+PUNCT_PATTERN = r"[\.!,?:;\"'”“’‘\-\–\—\(\)\[\]\{\}।]"
+
+def normalize_bangla_text(text: str) -> str:
+    if not isinstance(text, str):
+        return ""
+    text = " ".join(text.split())
+    text = re.sub(PUNCT_PATTERN, " ", text)
+    text = " ".join(text.split())
+    return text
+
+# ---------- Load tokenizer + model ----------
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH).to(DEVICE)
+model.eval()
+
+id2label = model.config.id2label
+label_list = [id2label[i] for i in range(len(id2label))]
+
+SHIRK_LABEL = "shirk"
+SHIRK_INDEX = label_list.index(SHIRK_LABEL)
+SHIRK_THRESHOLD = 0.7
+
+# ---------- FastAPI ----------
+app = FastAPI(title="Bangla Shirk Classifier API")
+
+class PredictRequest(BaseModel):
+    text: str
+
+class PredictResponse(BaseModel):
+    label: str
+    probabilities: Dict[str, float]
+
+@app.get("/")
+def root():
+    return {"status": "running"}
+
+@app.post("/predict", response_model=PredictResponse)
+def predict(req: PredictRequest):
+    text = normalize_bangla_text(req.text)
+
+    enc = tokenizer(
+        text,
+        truncation=True,
+        padding=True,
+        max_length=64,
+        return_tensors="pt"
+    )
+    enc = {k: v.to(DEVICE) for k, v in enc.items()}
+
+    with torch.no_grad():
+        outputs = model(**enc)
+        logits = outputs.logits[0]
+        probs = F.softmax(logits, dim=-1).cpu().numpy()
+
+    # apply threshold logic
+    top1 = int(probs.argmax())
+    if top1 == SHIRK_INDEX and probs[SHIRK_INDEX] < SHIRK_THRESHOLD:
+        top2 = int(probs.argsort()[-2])
+        pred_idx = top2
+    else:
+        pred_idx = top1
+
+    prob_dict = {label_list[i]: float(probs[i]) for i in range(len(label_list))}
+
+    return PredictResponse(
+        label=label_list[pred_idx],
+        probabilities=prob_dict
+    )