Update app.py

app.py

@@ -9,10 +9,16 @@ import numpy as np
 import pandas as pd
 import warnings
 import nltk
+from functools import lru_cache
+from typing import List, Dict, Any
 from sentence_transformers import SentenceTransformer
 from huggingface_hub import hf_hub_download
-from sklearn.calibration import CalibratedClassifierCV
-from sklearn.model_selection import train_test_split
+
+# Optional: sklearn imports (some envs have slightly different names)
+try:
+    from sklearn.calibration import CalibratedClassifierCV
+except Exception:
+    CalibratedClassifierCV = None
 
 warnings.filterwarnings("ignore")
 
@@ -35,7 +41,7 @@ FORCED_DIM = 768
 def ensure_nltk():
     resources = {
         "punkt": "tokenizers/punkt",
-        "punkt_tab": "tokenizers/punkt_tab/english",
+        "punkt_tab": "tokenizers/punkt_tab/english",  # may not exist on older NLTK; safe to ignore
     }
     for pkg, path in resources.items():
         try:
@@ -61,7 +67,7 @@ def preprocess_text(text: str, max_chars: int = 100000) -> str:
         t = t[:max_chars]
     return t
 
-def chunk_by_words(text: str, words_per_chunk: int = 350):
+def chunk_by_words(text: str, words_per_chunk: int = 350) -> List[str]:
     words = text.split()
     if not words:
         return []
@@ -73,9 +79,16 @@ def chunk_by_words(text: str, words_per_chunk: int = 350):
     return chunks
 
 # -------------------------------------------------
-# Load classifier + embedder (forced 768-dim)
+# Model loader (cached)
 # -------------------------------------------------
+@lru_cache(maxsize=1)
 def load_embedding_model():
+    """
+    Returns:
+      clf: a classifier with predict / predict_proba
+      embedding_model: SentenceTransformer (768-dim)
+      meta: dict with any metadata found in the joblib file
+    """
     path = hf_hub_download(
         repo_id=REPO_ID,
         filename=FILENAME,
@@ -96,7 +109,7 @@ def load_embedding_model():
     if actual_dim != FORCED_DIM:
         raise RuntimeError(f"Loaded embedder dim={actual_dim}, expected {FORCED_DIM}")
 
-    # Classifier sanity check
+    # Ensure classifier feature dimension matches
     clf_dim = getattr(clf, "n_features_in_", None)
     if clf_dim and clf_dim != FORCED_DIM:
         raise RuntimeError(
@@ -104,15 +117,131 @@ def load_embedding_model():
         )
 
     # -------------------------------------------------
-    # Ensure calibration
+    # Ensure calibration (Platt scaling if needed)
     # -------------------------------------------------
     if not hasattr(clf, "predict_proba") or "CalibratedClassifierCV" not in str(type(clf)):
-        print("⚠️ Classifier not calibrated — applying Platt scaling (logistic calibration)")
-        # Create dummy calibration using embeddings stored in joblib (if available)
-        X_train = data.get("X_val")
-        y_train = data.get("y_val")
-        if X_train is not None and y_train is not None:
-            clf = CalibratedClassifierCV(base_estimator=clf, method="sigmoid", cv="prefit")
-            clf.fit(X_train, y_train)
+        print("⚠️ Classifier not calibrated — applying Platt scaling (logistic calibration).")
+        X_val = data.get("X_val")
+        y_val = data.get("y_val")
+        if X_val is not None and y_val is not None and CalibratedClassifierCV is not None:
+            try:
+                # Newer sklearn uses 'estimator'; older uses 'base_estimator'
+                try:
+                    clf = CalibratedClassifierCV(estimator=clf, method="sigmoid", cv="prefit")
+                except TypeError:
+                    clf = CalibratedClassifierCV(base_estimator=clf, method="sigmoid", cv="prefit")
+                clf.fit(X_val, y_val)
+                print("✅ Calibration complete using provided validation split.")
+            except Exception as e:
+                print(f"⚠️ Calibration failed: {e}. Continuing with uncalibrated probabilities.")
         else:
-            print("⚠️ No calibration data found — probabilities may still be uncalibrated")
+            print("⚠️ No calibration data found or CalibratedClassifierCV unavailable.")
+
+    meta = {k: v for k, v in data.items() if k not in {"model"}}
+    return clf, embedding_model, meta
+
+# -------------------------------------------------
+# Inference
+# -------------------------------------------------
+def embed_texts(embedding_model: SentenceTransformer, texts: List[str]) -> np.ndarray:
+    with torch.no_grad():
+        embs = embedding_model.encode(
+            texts,
+            batch_size=32,
+            show_progress_bar=False,
+            convert_to_numpy=True,
+            normalize_embeddings=False,  # keep raw for classifier trained that way
+        )
+    return embs.astype(np.float32)
+
+def aggregate_probs(chunk_probs: List[float], mode: str = "mean") -> float:
+    if not chunk_probs:
+        return 0.0
+    arr = np.array(chunk_probs, dtype=np.float32)
+    if mode == "max":
+        return float(np.max(arr))
+    if mode == "median":
+        return float(np.median(arr))
+    return float(np.mean(arr))  # default mean
+
+def predict_single(text: str,
+                   words_per_chunk: int = 350,
+                   agg_mode: str = "mean") -> Dict[str, Any]:
+    clf, embedding_model, _ = load_embedding_model()
+    clean = preprocess_text(text)
+    chunks = chunk_by_words(clean, words_per_chunk=words_per_chunk) or [clean]
+
+    embs = embed_texts(embedding_model, chunks)
+    # Binary classifier: assume class 1 = AI, class 0 = Human (common convention)
+    proba = clf.predict_proba(embs)[:, 1] if hasattr(clf, "predict_proba") else clf.decision_function(embs)
+    # If decision_function, convert to [0,1] via sigmoid as a fallback
+    if proba.ndim == 1 and (proba.min() < 0 or proba.max() > 1):
+        proba = 1 / (1 + np.exp(-proba))
+
+    chunk_outputs = [{"chunk_index": i, "proba_ai": float(p), "text_preview": chunks[i][:120]} for i, p in enumerate(proba)]
+    doc_proba = aggregate_probs([co["proba_ai"] for co in chunk_outputs], mode=agg_mode)
+
+    return {
+        "doc_proba_ai": float(doc_proba),
+        "agg_mode": agg_mode,
+        "words_per_chunk": words_per_chunk,
+        "num_chunks": len(chunks),
+        "chunks": chunk_outputs,
+    }
+
+def classify_text(text: str,
+                  decision_threshold: float = 0.5,
+                  words_per_chunk: int = 350,
+                  agg_mode: str = "mean") -> Dict[str, Any]:
+    """
+    decision_threshold: label 'AI' if doc_proba_ai >= threshold
+    """
+    result = predict_single(text, words_per_chunk=words_per_chunk, agg_mode=agg_mode)
+    label = "AI" if result["doc_proba_ai"] >= float(decision_threshold) else "Human"
+    result["decision_threshold"] = float(decision_threshold)
+    result["label"] = label
+    return result
+
+# -------------------------------------------------
+# Gradio UI
+# -------------------------------------------------
+with gr.Blocks(title="Text AI Detector") as demo:
+    gr.Markdown("## 🔎 Text AI Detector (MPNet 768-dim + Calibrated Classifier)")
+
+    with gr.Row():
+        inp = gr.Textbox(label="Input text", lines=12, placeholder="Paste or type English text here...")
+    with gr.Row():
+        thr = gr.Slider(0.0, 1.0, value=0.50, step=0.01, label="Decision threshold (AI if ≥ threshold)")
+    with gr.Row():
+        wpc = gr.Slider(100, 800, value=350, step=50, label="Words per chunk")
+        agg = gr.Dropdown(choices=["mean", "median", "max"], value="mean", label="Chunk aggregation")
+    with gr.Row():
+        btn = gr.Button("Classify")
+    with gr.Row():
+        out_label = gr.Textbox(label="Label", interactive=False)
+        out_proba = gr.Number(label="Document AI probability", interactive=False, precision=4)
+    with gr.Row():
+        out_json = gr.JSON(label="Details (per-chunk probabilities)")
+
+    def _on_click(text, threshold, words_per_chunk, agg_mode):
+        if not text or not text.strip():
+            return ("", 0.0, {"error": "Empty text"})
+        try:
+            res = classify_text(text, float(threshold), int(words_per_chunk), agg_mode)
+            return (res["label"], res["doc_proba_ai"], res)
+        except Exception as e:
+            return ("", 0.0, {"error": str(e)})
+
+    btn.click(_on_click, inputs=[inp, thr, wpc, agg], outputs=[out_label, out_proba, out_json])
+
+# Provide ASGI `app` for hosts that expect it (e.g., Uvicorn/Gunicorn)
+try:
+    from fastapi import FastAPI
+    fastapi_app = FastAPI()
+    app = gr.mount_gradio_app(fastapi_app, demo, path="/")
+except Exception:
+    app = None
+
+if __name__ == "__main__":
+    # Local run
+    demo.queue(concurrency_count=2).launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))