Update tools.py

tools.py

@@ -18,7 +18,6 @@ from collections import Counter, defaultdict
 from sentence_transformers import SentenceTransformer
 from umap import UMAP
 from hdbscan import HDBSCAN
-from keybert import KeyBERT
 from sklearn.metrics import adjusted_rand_score
 from sklearn.metrics.pairwise import cosine_similarity
 import optuna
@@ -56,15 +55,20 @@ def prepare_documents(df: pd.DataFrame) -> list[str]:
 
 
 # ---------------------------------------------------------------------------
-# §3.1 — Embed with SPECTER-2
+# §3.1 — Embed with SPECTER-2 (cached model for speed)
 # ---------------------------------------------------------------------------
+_MODEL_CACHE = {}
+
 def embed_documents(
     docs: list[str],
     model_name: str = "allenai/specter2_base",
 ) -> np.ndarray:
     """Embed with SPECTER-2. Deterministic — no tuning (§3.3)."""
-    model = SentenceTransformer(model_name)
-    embeddings = model.encode(docs, show_progress_bar=True, batch_size=32)
+    if model_name not in _MODEL_CACHE:
+        logger.info("Loading %s (first time, will be cached)…", model_name)
+        _MODEL_CACHE[model_name] = SentenceTransformer(model_name)
+    model = _MODEL_CACHE[model_name]
+    embeddings = model.encode(docs, show_progress_bar=True, batch_size=64)
     logger.info("Embedded %d docs → %s", len(docs), embeddings.shape)
     return embeddings
 
@@ -110,6 +114,19 @@ def compute_persistence(clusterer: HDBSCAN) -> float:
     return 0.0
 
 
+def per_cluster_persistence(clusterer: HDBSCAN, labels: np.ndarray) -> dict:
+    """Map each cluster ID to its persistence score (§8)."""
+    try:
+        p = getattr(clusterer, "cluster_persistence_", None)
+        if p is None or len(p) == 0:
+            return {}
+        unique = sorted(set(int(l) for l in labels if l != -1))
+        return {cid: float(p[i]) if i < len(p) else 0.0
+                for i, cid in enumerate(unique)}
+    except Exception:
+        return {}
+
+
 def compute_dbcv(reduced: np.ndarray, labels: np.ndarray) -> float:
     """Density-Based Cluster Validity index."""
     try:
@@ -124,14 +141,15 @@ def compute_dbcv(reduced: np.ndarray, labels: np.ndarray) -> float:
 
 
 def compute_stability(embeddings: np.ndarray, params: dict,
-                      n_seeds: int = 5) -> float:
-    """Cluster-recurrence stability via pairwise ARI across seeds (§3.4)."""
+                      n_seeds: int = 3) -> float:
+    """Cluster-recurrence stability via pairwise ARI across seeds (§3.4).
+    Uses 3 seeds by default for speed (spec allows 3–5)."""
     all_labels = []
     for s in range(n_seeds):
         u = UMAP(n_neighbors=params["n_neighbors"],
                  n_components=params["n_components"],
                  min_dist=0.0, metric="cosine",
-                 random_state=s * 7 + 1)
+                 random_state=s * 7 + 1, low_memory=True)
         red = u.fit_transform(embeddings)
         h = HDBSCAN(min_cluster_size=params["min_cluster_size"],
                     min_samples=params["min_samples"],
@@ -167,7 +185,8 @@ def _objective(trial, embeddings, n_docs):
                   min_cluster_size=mcs, min_samples=ms, csm=csm, cse=cse)
 
     u = UMAP(n_neighbors=n_neighbors, n_components=n_components,
-             min_dist=0.0, metric="cosine", random_state=42)
+             min_dist=0.0, metric="cosine", random_state=42,
+             low_memory=True)
     red = u.fit_transform(embeddings)
 
     h = HDBSCAN(min_cluster_size=mcs, min_samples=ms, metric="euclidean",
@@ -235,7 +254,7 @@ def run_bayesian_optimisation(
         )
         # §3.6 convergence: 3 consecutive passing within 5 % of best
         passing = [e for e in trial_log if e["discipline_pass"]]
-        if len(passing) >= 3 and i >= 19:
+        if len(passing) >= 3 and i >= 9:   # allow early stop after 10 trials
             best_p = max(e["persistence"] for e in passing)
             if best_p > 0:
                 last3 = passing[-3:]
@@ -255,7 +274,7 @@ def run_bayesian_optimisation(
 
     bp = best.user_attrs["params"]
     labels = np.array(best.user_attrs["labels"])
-    stability = compute_stability(embeddings, bp, n_seeds=5)
+    stability = compute_stability(embeddings, bp, n_seeds=3)
 
     return dict(
         best_params=bp, best_labels=labels,
@@ -274,28 +293,37 @@ def run_bayesian_optimisation(
 # ---------------------------------------------------------------------------
 def compute_2d_umap(embeddings: np.ndarray, seed: int = 42) -> np.ndarray:
     return UMAP(n_neighbors=15, n_components=2, min_dist=0.1,
-                metric="cosine", random_state=seed).fit_transform(embeddings)
+                metric="cosine", random_state=seed,
+                low_memory=True).fit_transform(embeddings)
 
 
 # ---------------------------------------------------------------------------
-# §3.1 — KeyBERT keyphrase extraction per cluster (3–5 phrases)
+# §3.1 — TF-IDF keyphrase extraction per cluster (3–5 phrases)
+#         Fast alternative to KeyBERT — no extra model download needed.
 # ---------------------------------------------------------------------------
 def extract_keyphrases(docs: list[str], labels: np.ndarray,
                        top_n: int = 5) -> dict:
-    kw = KeyBERT(model="all-MiniLM-L6-v2")
+    from sklearn.feature_extraction.text import TfidfVectorizer
     cluster_docs = defaultdict(list)
     for doc, lab in zip(docs, labels):
         if lab != -1:
             cluster_docs[int(lab)].append(doc)
     out = {}
     for cid, cdocs in cluster_docs.items():
+        if len(cdocs) < 2:
+            out[cid] = []
+            continue
         try:
-            out[cid] = kw.extract_keywords(
-                " ".join(cdocs), keyphrase_ngram_range=(1, 3),
-                stop_words="english", top_n=top_n,
-                use_mmr=True, diversity=0.5)
+            tfidf = TfidfVectorizer(
+                stop_words="english", max_features=200,
+                ngram_range=(1, 3), max_df=0.9, min_df=1)
+            X = tfidf.fit_transform(cdocs)
+            terms = tfidf.get_feature_names_out()
+            scores = X.sum(axis=0).A1
+            top_idx = scores.argsort()[::-1][:top_n]
+            out[cid] = [(terms[i], float(scores[i])) for i in top_idx]
         except Exception as e:
-            logger.warning("KeyBERT cluster %d: %s", cid, e)
+            logger.warning("Keyphrase extraction cluster %d: %s", cid, e)
             out[cid] = []
     return out
 
@@ -364,6 +392,35 @@ def get_representative_docs(labels, embeddings, docs, top_n=3):
     return out
 
 
+# ---------------------------------------------------------------------------
+# §9 — RQ2 / RQ3 mismatch table
+# ---------------------------------------------------------------------------
+def build_mismatch_table(keyphrases: dict, cluster_labels: dict) -> list:
+    """Compare cluster keyphrases against assigned labels to flag mismatches.
+    Returns rows for a mismatch table (§9)."""
+    rows = []
+    for cid in sorted(keyphrases.keys()):
+        kps = keyphrases.get(cid, [])
+        kp_terms = [k[0] if isinstance(k, tuple) else k for k in kps[:5]]
+        label = cluster_labels.get(cid, "")
+        # Check overlap between label words and keyphrase terms
+        label_words = set(label.lower().split())
+        kp_words = set(" ".join(kp_terms).lower().split())
+        overlap = label_words & kp_words
+        noise = {"the","and","for","with","using","based","from","in","of","a","to"}
+        overlap -= noise
+        match_pct = len(overlap) / max(len(label_words - noise), 1)
+        status = "MATCH" if match_pct >= 0.3 else "MISMATCH"
+        rows.append(dict(
+            cluster=cid, label=label,
+            keyphrases=", ".join(kp_terms),
+            overlap=", ".join(overlap) if overlap else "—",
+            match_pct=round(match_pct * 100),
+            status=status,
+        ))
+    return rows
+
+
 # ---------------------------------------------------------------------------
 # High-level pipeline entry point
 # ---------------------------------------------------------------------------
@@ -390,9 +447,13 @@ def run_topic_modeling(filepath: str, n_trials: int = 50,
                 min_samples=bp["min_samples"], metric="euclidean",
                 cluster_selection_method=bp["csm"],
                 cluster_selection_epsilon=bp["cse"],
-                allow_single_cluster=False)
+                allow_single_cluster=False,
+                gen_min_span_tree=True)
     h.fit(red)
 
+    # Per-cluster persistence (§8)
+    cluster_pers = per_cluster_persistence(h, labels)
+
     # 5. Outlier reduction (§3.2 — clusters < 5 reassigned)
     labels = outlier_reduction(labels, red, n_docs)
 
@@ -416,6 +477,7 @@ def run_topic_modeling(filepath: str, n_trials: int = 50,
         keyphrases=keyphrases, representative_docs=rep_docs,
         membership=sw, umap_2d=umap_2d.tolist(),
         discipline=disc, best_params=bp,
+        cluster_persistence=cluster_pers,
         metrics=dict(persistence=opt["persistence"],
                      dbcv=opt["dbcv"],
                      stability=opt["stability"]),