Update detector/custom_model.py

detector/custom_model.py

@@ -8,36 +8,34 @@ from reportlab.pdfgen import canvas
 from reportlab.lib import colors
 import nltk
 
-# === Environment (no runtime directory creation) ===
+# === Environment Setup ===
 os.environ["HF_HOME"] = "/tmp/hf_home"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_home"
 os.environ["NLTK_DATA"] = "/tmp/nltk_data"
-
-# Append pre-installed punkt path (you must ensure it's there)
 nltk.data.path.append("/tmp/nltk_data")
 
-# === Model loading: Hugging Face (Remote) vs Local ===
+# === Model Source (Hugging Face or Local) ===
 USE_HF_MODEL = os.getenv("USE_HF_MODEL") == "1"
+hf_token = os.getenv("HF_TOKEN")
+MODEL_PATH = "AlyanAkram/stealth-roberta" if USE_HF_MODEL else "./detector/models/roberta-detector"
 
 if USE_HF_MODEL:
-    hf_token = os.getenv("HF_TOKEN")
-    MODEL_PATH = "AlyanAkram/stealth-roberta"
+    print("🔐 Loading model from Hugging Face Hub...")
     tokenizer = RobertaTokenizer.from_pretrained(MODEL_PATH, token=hf_token)
     model = RobertaForSequenceClassification.from_pretrained(MODEL_PATH, token=hf_token)
 else:
-    MODEL_PATH = "./detector/models/roberta-detector"
+    print("📁 Loading model from local files...")
     tokenizer = RobertaTokenizer.from_pretrained(MODEL_PATH, local_files_only=True)
     model = RobertaForSequenceClassification.from_pretrained(MODEL_PATH, local_files_only=True)
 
 model.eval().to("cuda" if torch.cuda.is_available() else "cpu")
 device = next(model.parameters()).device
 
-# === AI classification threshold ===
+# === Constants ===
 AI_THRESHOLD = 0.5
-
-# === Output reports location (this must already be writable in /tmp) ===
 REPORT_DIR = Path("/tmp/reports")
 
+# === Main Analysis Function ===
 def analyze_text(text: str):
     results = []
     paragraphs = [p.strip() for p in text.split("\n") if p.strip()]
@@ -67,6 +65,7 @@ def analyze_text(text: str):
         "results": results
     }
 
+# === PDF Report Generator ===
 def generate_pdf_report(results: dict, filename: str) -> str:
     REPORT_DIR.mkdir(exist_ok=True)
     pdf_path = REPORT_DIR / f"{filename}.pdf"