Update app.py

app.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, Request
+from fastapi import FastAPI
 from pydantic import BaseModel
 from langdetect import detect
 import torch
@@ -7,22 +7,27 @@ from transformers import DistilBertModel, AutoModel, AutoTokenizer, DistilBertTo
 from huggingface_hub import snapshot_download
 import os
 
+# App and device
 app = FastAPI()
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Use local cache folder for downloaded models
-os.environ["TRANSFORMERS_CACHE"] = "/app/.hf_cache"
-os.makedirs("/app/.hf_cache", exist_ok=True)
+# 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
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# 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)
 
-# -------------------------------
-# Model Classes
-# -------------------------------
+# ----------------------------
+# Model classes
+# ----------------------------
 
 class ToxicBERT(nn.Module):
     def __init__(self):
         super().__init__()
-        self.bert = DistilBertModel.from_pretrained(snapshot_download("koyu008/English_Toxic_Classifier"))
+        self.bert = DistilBertModel.from_pretrained(english_path)
         self.dropout = nn.Dropout(0.3)
         self.classifier = nn.Linear(self.bert.config.hidden_size, 6)
 
@@ -30,10 +35,11 @@ class ToxicBERT(nn.Module):
         output = self.bert(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state[:, 0]
         return self.classifier(self.dropout(output))
 
+
 class HinglishToxicClassifier(nn.Module):
     def __init__(self):
         super().__init__()
-        self.bert = AutoModel.from_pretrained(snapshot_download("koyu008/Hinglish_comment_classifier"))
+        self.bert = AutoModel.from_pretrained(hinglish_path)
         hidden_size = self.bert.config.hidden_size
         self.pool = lambda hidden: torch.cat([
             hidden.mean(dim=1),
@@ -52,9 +58,9 @@ class HinglishToxicClassifier(nn.Module):
         x = self.bottleneck(pooled)
         return self.classifier(x)
 
-# -------------------------------
-# Load Models and Tokenizers
-# -------------------------------
+# ----------------------------
+# Load Models & Tokenizers
+# ----------------------------
 
 english_model = ToxicBERT().to(device)
 english_model.load_state_dict(torch.load("bert_toxic_classifier.pt", map_location=device))
@@ -66,9 +72,9 @@ hinglish_model.load_state_dict(torch.load("best_hinglish_model.pt", map_location
 hinglish_model.eval()
 hinglish_tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-base")
 
-# -------------------------------
-# Request & Inference
-# -------------------------------
+# ----------------------------
+# API
+# ----------------------------
 
 class InputText(BaseModel):
     text: str
@@ -83,10 +89,18 @@ async def predict(input: InputText):
         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}
+        return {
+            "language": "english",
+            "classes": ["toxic", "severe_toxic", "obscene", "threat", "insult", "identity_hate"],
+            "probabilities": probs
+        }
     else:
         inputs = hinglish_tokenizer(text, return_tensors="pt", truncation=True, padding=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}
+        return {
+            "language": "hinglish",
+            "classes": ["toxic", "non-toxic"],
+            "probabilities": probs
+        }