Update app.py

app.py

@@ -1,33 +1,29 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
 from pydantic import BaseModel
-from langdetect import detect
 import torch
 import torch.nn as nn
-from transformers import DistilBertModel, AutoModel, AutoTokenizer, DistilBertTokenizer
+from transformers import DistilBertTokenizer, DistilBertModel, AutoModel, AutoTokenizer
+from langdetect import detect
 from huggingface_hub import snapshot_download
 import os
 
-# App and device
-app = FastAPI()
+# Device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Create safe local cache directory
-hf_cache_dir = "./hf_cache"
-os.makedirs(hf_cache_dir, exist_ok=True)
-os.environ["TRANSFORMERS_CACHE"] = hf_cache_dir
+# Download model repos from HF Hub
+english_repo = snapshot_download("koyu008/English_Toxic_Classifier")
+hinglish_repo = snapshot_download("koyu008/HInglish_comment_classifier")
 
-# Download model repositories to local path
-english_path = snapshot_download("koyu008/English_Toxic_Classifier", cache_dir=hf_cache_dir)
-hinglish_path = snapshot_download("koyu008/Hinglish_comment_classifier", cache_dir=hf_cache_dir)
+# Tokenizers
+english_tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")
+hinglish_tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-base")
 
-# ----------------------------
-# Model classes
-# ----------------------------
 
+# English Model
 class ToxicBERT(nn.Module):
     def __init__(self):
         super().__init__()
-        self.bert = DistilBertModel.from_pretrained(english_path)
+        self.bert = DistilBertModel.from_pretrained("distilbert-base-uncased")
         self.dropout = nn.Dropout(0.3)
         self.classifier = nn.Linear(self.bert.config.hidden_size, 6)
 
@@ -36,10 +32,11 @@ class ToxicBERT(nn.Module):
         return self.classifier(self.dropout(output))
 
 
+# Hinglish Model
 class HinglishToxicClassifier(nn.Module):
     def __init__(self):
         super().__init__()
-        self.bert = AutoModel.from_pretrained(hinglish_path)
+        self.bert = AutoModel.from_pretrained("xlm-roberta-base")
         hidden_size = self.bert.config.hidden_size
         self.pool = lambda hidden: torch.cat([
             hidden.mean(dim=1),
@@ -58,49 +55,50 @@ class HinglishToxicClassifier(nn.Module):
         x = self.bottleneck(pooled)
         return self.classifier(x)
 
-# ----------------------------
-# Load Models & Tokenizers
-# ----------------------------
 
+# Instantiate and load models
 english_model = ToxicBERT().to(device)
-english_model.load_state_dict(torch.load("bert_toxic_classifier.pt", map_location=device))
+english_model.load_state_dict(torch.load(os.path.join(english_repo, "bert_toxic_classifier.pt"), map_location=device))
 english_model.eval()
-english_tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")
 
 hinglish_model = HinglishToxicClassifier().to(device)
-hinglish_model.load_state_dict(torch.load("best_hinglish_model.pt", map_location=device))
+hinglish_model.load_state_dict(torch.load(os.path.join(hinglish_repo, "best_hinglish_model.pt"), map_location=device))
 hinglish_model.eval()
-hinglish_tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-base")
 
-# ----------------------------
-# API
-# ----------------------------
+# Labels
+english_labels = ['toxic', 'severe toxic', 'obscene', 'threat', 'insult', 'identity hate']
+hinglish_labels = ['not toxic', 'toxic']
+
+# FastAPI
+app = FastAPI()
+
 
-class InputText(BaseModel):
+class TextIn(BaseModel):
     text: str
 
+
 @app.post("/predict")
-async def predict(input: InputText):
-    text = input.text
-    lang = detect(text)
+def predict(data: TextIn):
+    text = data.text
+    try:
+        lang = detect(text)
+    except:
+        lang = "unknown"
 
     if lang == "en":
-        inputs = english_tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
+        tokenizer = english_tokenizer
+        model = english_model
+        inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True).to(device)
         with torch.no_grad():
-            logits = english_model(**inputs)
-        probs = torch.softmax(logits, dim=1).cpu().numpy().tolist()[0]
-        return {
-            "language": "english",
-            "classes": ["toxic", "severe_toxic", "obscene", "threat", "insult", "identity_hate"],
-            "probabilities": probs
-        }
+            outputs = model(**inputs)
+            probs = torch.sigmoid(outputs).squeeze().cpu().tolist()
+        return {"language": "English", "predictions": dict(zip(english_labels, probs))}
+
     else:
-        inputs = hinglish_tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
+        tokenizer = hinglish_tokenizer
+        model = hinglish_model
+        inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True).to(device)
         with torch.no_grad():
-            logits = hinglish_model(**inputs)
-        probs = torch.softmax(logits, dim=1).cpu().numpy().tolist()[0]
-        return {
-            "language": "hinglish",
-            "classes": ["toxic", "non-toxic"],
-            "probabilities": probs
-        }
+            outputs = model(**inputs)
+            probs = torch.softmax(outputs, dim=1).squeeze().cpu().tolist()
+        return {"language": "Hinglish", "predictions": dict(zip(hinglish_labels, probs))}