Update app.py

updated 768 dim

app.py

@@ -16,15 +16,26 @@ from huggingface_hub import hf_hub_download
 
 warnings.filterwarnings("ignore")
 
+# -------------------------------------------------
 # Hugging Face model config
+# -------------------------------------------------
 REPO_ID = "Detecting-ai/text-detector-model-embedding"
 FILENAME = "complete_trained_model_lite.joblib"
 REPO_TYPE = "model"
 
-# --- Ensure NLTK dependencies ---
+# -------------------------------------------------
+# Force 768-dim embedder (MPNet)
+# -------------------------------------------------
+FORCED_EMBEDDER = "sentence-transformers/all-mpnet-base-v2"
+FORCED_DIM = 768
+
+# -------------------------------------------------
+# Ensure NLTK dependencies
+# -------------------------------------------------
 def ensure_nltk():
     resources = {
         "punkt": "tokenizers/punkt",
+        # newer nltk introduces punkt_tab; harmless to try
         "punkt_tab": "tokenizers/punkt_tab/english",
         "stopwords": "corpora/stopwords",
         "wordnet": "corpora/wordnet",
@@ -33,26 +44,31 @@ def ensure_nltk():
         try:
             nltk.data.find(path)
         except LookupError:
-            nltk.download(pkg)
+            try:
+                nltk.download(pkg, quiet=True)
+            except Exception:
+                pass
+
 ensure_nltk()
 
-# --- Helper Functions ---
+# -------------------------------------------------
+# Helper functions
+# -------------------------------------------------
 def _to_stopword_set(sw):
-    if sw is None: return set(stopwords.words("english"))
-    if isinstance(sw, (list, tuple)): return set(sw)
-    if isinstance(sw, set): return sw
-    try: return set(sw)
-    except: return set(stopwords.words("english"))
-
-def _guess_model_by_dim(dim: int):
-    if dim == 768:
-        return "sentence-transformers/all-mpnet-base-v2"
-    if dim == 384:
-        return "sentence-transformers/all-MiniLM-L6-v2"
-    # default
-    return "sentence-transformers/all-MiniLM-L6-v2"
-
-# --- Load Model Bundle ---
+    if sw is None:
+        return set(stopwords.words("english"))
+    if isinstance(sw, set):
+        return sw
+    if isinstance(sw, (list, tuple)):
+        return set(sw)
+    try:
+        return set(sw)
+    except Exception:
+        return set(stopwords.words("english"))
+
+# -------------------------------------------------
+# Load model bundle + forced 768-dim embedder
+# -------------------------------------------------
 def load_embedding_model():
     path = hf_hub_download(
         repo_id=REPO_ID,
@@ -65,27 +81,43 @@ def load_embedding_model():
     data = joblib.load(path)
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    stored_name = (data.get("embedding_model_name") or data.get("embedding_model_path") or "").strip()
-    expected_dim = int(data.get("embedding_dim", 0)) if data.get("embedding_dim") else 0
-    emb_name = stored_name or _guess_model_by_dim(expected_dim)
+    clf = data.get("model")
+    if clf is None:
+        raise RuntimeError("Model file does not contain 'model' key.")
 
-    print(f"🔧 Loading embedding model: {emb_name} on {device}")
-    embedding_model = SentenceTransformer(emb_name, device=device)
+    # --- Always use 768-dim MPNet ---
+    print(f"🔧 Loading 768-dim embedder: {FORCED_EMBEDDER} on {device}")
+    embedding_model = SentenceTransformer(FORCED_EMBEDDER, device=device)
     actual_dim = embedding_model.get_sentence_embedding_dimension()
-
+    if actual_dim != FORCED_DIM:
+        raise RuntimeError(f"Loaded embedder dim={actual_dim}, expected {FORCED_DIM}")
+
+    # --- Classifier sanity check (must be trained on 768) ---
+    clf_dim = getattr(clf, "n_features_in_", None)
+    if clf_dim and clf_dim != FORCED_DIM:
+        raise RuntimeError(
+            f"Classifier expects {clf_dim} features, but app is configured for {FORCED_DIM}. "
+            f"Please retrain or load a 768-dim trained classifier."
+        )
+
+    # finalize
     data["embedding_model"] = embedding_model
-    data["resolved_embedding_model_name"] = emb_name
+    data["resolved_embedding_model_name"] = FORCED_EMBEDDER
     data["resolved_embedding_dim"] = actual_dim
     data["device"] = device
     data["lemmatizer"] = data.get("lemmatizer") or WordNetLemmatizer()
     data["stop_words"] = _to_stopword_set(data.get("stop_words"))
     data["max_tokens"] = data.get("max_tokens", 600)
 
+    print(f"✅ Using embedder: {FORCED_EMBEDDER} (dim={actual_dim}) — classifier expects {getattr(clf,'n_features_in_','unknown')}")
     return data
 
-# --- Preprocessing ---
+# -------------------------------------------------
+# Preprocessing
+# -------------------------------------------------
 def preprocess_text(text, lemmatizer, stop_words, max_tokens=600):
-    if pd.isna(text) or not text: return ""
+    if pd.isna(text) or not str(text).strip():
+        return ""
     text = str(text).lower()
     text = re.sub(r"[^a-zA-Z\s]", " ", text)
     tokens = [
@@ -95,18 +127,26 @@ def preprocess_text(text, lemmatizer, stop_words, max_tokens=600):
     ]
     return " ".join(tokens[:max_tokens])
 
-# --- Prediction ---
+# -------------------------------------------------
+# Prediction
+# -------------------------------------------------
 def predict_text(text, model_data):
     proc = preprocess_text(text, model_data["lemmatizer"], model_data["stop_words"], model_data["max_tokens"])
     if not proc:
         return "UNKNOWN", 0.0, {"error": "Empty text after preprocessing"}
 
     with torch.no_grad():
-        emb = model_data["embedding_model"].encode([proc], convert_to_numpy=True, normalize_embeddings=False)
+        emb = model_data["embedding_model"].encode(
+            [proc], convert_to_numpy=True, normalize_embeddings=False
+        )
     if emb.ndim == 1:
         emb = emb.reshape(1, -1)
 
     clf = model_data["model"]
+    need = getattr(clf, "n_features_in_", emb.shape[1])
+    if emb.shape[1] != need:
+        return "ERROR", 0.0, {"error": f"Embedding dim {emb.shape[1]} != classifier requires {need}"}
+
     try:
         pred = clf.predict(emb)[0]
         conf = float(np.max(clf.predict_proba(emb)[0])) if hasattr(clf, "predict_proba") else 0.5
@@ -115,7 +155,9 @@ def predict_text(text, model_data):
 
     return str(pred), conf, {"tokens": len(proc.split())}
 
-# --- Gradio App ---
+# -------------------------------------------------
+# Gradio App
+# -------------------------------------------------
 def create_app(model_data):
     with gr.Blocks(title="Embedding-based Human vs AI Detector") as demo:
         gr.Markdown("## 🤖👤 Human vs AI Detector (Embedding-based)")
@@ -131,18 +173,28 @@ def create_app(model_data):
                 headline = f"👤 **Human Written** (Conf: {conf:.1%})"
             elif label.upper() == "ERROR":
                 headline = f"❌ Error: {meta.get('error', 'Unknown')}"
+            elif label.upper() == "UNKNOWN":
+                headline = f"❓ Unknown (Conf: {conf:.1%})"
             else:
                 headline = f"❓ {label} (Conf: {conf:.1%})"
-            det = f"- Tokens: {meta.get('tokens','?')}\n- Embedding: {model_data['resolved_embedding_model_name']} (dim={model_data['resolved_embedding_dim']})"
+
+            det = (
+                f"- Tokens: {meta.get('tokens','?')}\n"
+                f"- Embedding: {model_data['resolved_embedding_model_name']} "
+                f"(dim={model_data['resolved_embedding_dim']})"
+            )
             return headline, det
 
         inp.submit(_predict_ui, inp, [out, details])
         gr.Button("🔍 Predict").click(_predict_ui, inp, [out, details])
     return demo
 
-# --- Load + Launch ---
+# -------------------------------------------------
+# Load + Launch
+# -------------------------------------------------
 _model_data = load_embedding_model()
 demo = create_app(_model_data)
 
 if __name__ == "__main__":
+    # You can pass share=True if you need a public URL
     demo.launch()