BHAVIKBANKER/BERTopic_AG_final
1
1"""2tools.py3--------4Topic-modelling pipeline: SPECTER-2 → UMAP → HDBSCAN5with multi-objective Bayesian optimisation over UMAP + HDBSCAN6parameters (§3.1–§3.6 of the methodology guide).7 8No BERTopic wrapper — bare UMAP + HDBSCAN on SPECTER-2 embeddings.9"""10 11import re12import logging13import pandas as pd14import numpy as np15from typing import Optional16from collections import Counter, defaultdict17 18from sentence_transformers import SentenceTransformer19from umap import UMAP20from hdbscan import HDBSCAN21from sklearn.metrics import adjusted_rand_score22from sklearn.metrics.pairwise import cosine_similarity23import optuna24 25# ---------------------------------------------------------------------------26# Logging27# ---------------------------------------------------------------------------28logging.basicConfig(level=logging.INFO, format="%(levelname)s | %(message)s")29logger = logging.getLogger(__name__)30optuna.logging.set_verbosity(optuna.logging.WARNING)31 32 33# ---------------------------------------------------------------------------34# Data Loading (unchanged)35# ---------------------------------------------------------------------------36def load_csv(filepath: str) -> pd.DataFrame:37 df = pd.read_csv(filepath)38 df.columns = df.columns.str.lower()39 required = {"title", "abstract"}40 missing = required - set(df.columns)41 if missing:42 raise ValueError(f"CSV missing column(s): {missing}")43 logger.info("Loaded %d rows from '%s'.", len(df), filepath)44 return df45 46 47# ---------------------------------------------------------------------------48# §3.1 — Input unit: title + abstract concatenation49# ---------------------------------------------------------------------------50def prepare_documents(df: pd.DataFrame) -> list[str]:51 """One string per paper: title + abstract (§3.1 input unit)."""52 docs = (df["title"].fillna("") + ". " + df["abstract"].fillna("")).tolist()53 logger.info("Prepared %d title+abstract documents.", len(docs))54 return docs55 56 57# ---------------------------------------------------------------------------58# §3.1 — Embed with SPECTER-2 (cached model for speed)59# ---------------------------------------------------------------------------60_MODEL_CACHE = {}61 62def embed_documents(63 docs: list[str],64 model_name: str = "allenai/specter2_base",65) -> np.ndarray:66 """Embed with SPECTER-2. Deterministic — no tuning (§3.3)."""67 if model_name not in _MODEL_CACHE:68 logger.info("Loading %s (first time, will be cached)…", model_name)69 _MODEL_CACHE[model_name] = SentenceTransformer(model_name)70 model = _MODEL_CACHE[model_name]71 embeddings = model.encode(docs, show_progress_bar=True, batch_size=64)72 logger.info("Embedded %d docs → %s", len(docs), embeddings.shape)73 return embeddings74 75 76# ---------------------------------------------------------------------------77# §3.2 — Cluster discipline checks78# ---------------------------------------------------------------------------79def check_discipline(labels: np.ndarray, n_docs: int) -> dict:80 """Two hard constraints: max-mass ≤ 25 %, min-size ≥ 5."""81 counts = Counter(int(l) for l in labels)82 unique = [l for l in counts if l != -1]83 84 if not unique:85 return dict(max_mass_pct=0, max_mass_ok=False,86 min_size=0, min_size_ok=False,87 n_clusters=0, n_noise=counts.get(-1, 0))88 89 max_mass_pct = max(counts[l] / n_docs for l in unique)90 min_size = min(counts[l] for l in unique)91 92 return dict(93 max_mass_pct=round(max_mass_pct, 4),94 max_mass_ok=max_mass_pct <= 0.25,95 min_size=int(min_size),96 min_size_ok=min_size >= 5,97 n_clusters=len(unique),98 n_noise=counts.get(-1, 0),99 cluster_sizes={l: counts[l] for l in sorted(unique)},100 )101 102 103# ---------------------------------------------------------------------------104# §3.4 — Quality metrics105# ---------------------------------------------------------------------------106def compute_persistence(clusterer: HDBSCAN) -> float:107 """Average cluster persistence from the condensed tree."""108 try:109 p = getattr(clusterer, "cluster_persistence_", None)110 if p is not None and len(p) > 0:111 return float(np.mean(p))112 except Exception:113 pass114 return 0.0115 116 117def per_cluster_persistence(clusterer: HDBSCAN, labels: np.ndarray) -> dict:118 """Map each cluster ID to its persistence score (§8)."""119 try:120 p = getattr(clusterer, "cluster_persistence_", None)121 if p is None or len(p) == 0:122 return {}123 unique = sorted(set(int(l) for l in labels if l != -1))124 return {cid: float(p[i]) if i < len(p) else 0.0125 for i, cid in enumerate(unique)}126 except Exception:127 return {}128 129 130def compute_dbcv(reduced: np.ndarray, labels: np.ndarray) -> float:131 """Density-Based Cluster Validity index."""132 try:133 from hdbscan.validity import validity_index134 ul = set(labels); ul.discard(-1)135 if len(ul) < 2:136 return -1.0137 return float(validity_index(reduced.astype(np.float64), labels))138 except Exception as e:139 logger.warning("DBCV failed: %s", e)140 return -1.0141 142 143def compute_stability(embeddings: np.ndarray, params: dict,144 n_seeds: int = 3) -> float:145 """Cluster-recurrence stability via pairwise ARI across seeds (§3.4).146 Uses 3 seeds by default for speed (spec allows 3–5)."""147 all_labels = []148 for s in range(n_seeds):149 u = UMAP(n_neighbors=params["n_neighbors"],150 n_components=params["n_components"],151 min_dist=0.0, metric="cosine",152 random_state=s * 7 + 1, low_memory=True)153 red = u.fit_transform(embeddings)154 h = HDBSCAN(min_cluster_size=params["min_cluster_size"],155 min_samples=params["min_samples"],156 metric="euclidean",157 cluster_selection_method=params["csm"],158 cluster_selection_epsilon=params["cse"])159 all_labels.append(h.fit_predict(red))160 161 aris = []162 for i in range(len(all_labels)):163 for j in range(i + 1, len(all_labels)):164 aris.append(adjusted_rand_score(all_labels[i], all_labels[j]))165 return float(np.mean(aris)) if aris else 0.0166 167 168# ---------------------------------------------------------------------------169# §3.4 — Bayesian optimisation objective170# ---------------------------------------------------------------------------171def _objective(trial, embeddings, n_docs):172 """Single Optuna trial — returns (persistence, dbcv, stability_placeholder)."""173 n_neighbors = trial.suggest_categorical("n_neighbors", [5, 10, 15, 30, 50])174 n_components = trial.suggest_int("n_components", 5, 10)175 mcs = trial.suggest_int(176 "min_cluster_size",177 max(5, int(0.01 * n_docs)),178 max(20, int(0.05 * n_docs)),179 )180 ms = trial.suggest_int("min_samples", 1, mcs)181 csm = trial.suggest_categorical("csm", ["eom", "leaf"])182 cse = trial.suggest_float("cse", 0.0, 0.3, step=0.05)183 184 params = dict(n_neighbors=n_neighbors, n_components=n_components,185 min_cluster_size=mcs, min_samples=ms, csm=csm, cse=cse)186 187 u = UMAP(n_neighbors=n_neighbors, n_components=n_components,188 min_dist=0.0, metric="cosine", random_state=42,189 low_memory=True)190 red = u.fit_transform(embeddings)191 192 h = HDBSCAN(min_cluster_size=mcs, min_samples=ms, metric="euclidean",193 cluster_selection_method=csm,194 cluster_selection_epsilon=cse,195 allow_single_cluster=False, gen_min_span_tree=True)196 labels = h.fit_predict(red)197 198 disc = check_discipline(labels, n_docs)199 trial.set_user_attr("params", params)200 trial.set_user_attr("discipline", disc)201 trial.set_user_attr("labels", labels.tolist())202 203 # Hard-constraint violation → worst scores204 if not disc["max_mass_ok"] or not disc["min_size_ok"]:205 trial.set_user_attr("pass", False)206 return 0.0, -1.0, 0.0207 208 trial.set_user_attr("pass", True)209 pers = compute_persistence(h)210 dbcv = compute_dbcv(red, labels)211 trial.set_user_attr("persistence", pers)212 trial.set_user_attr("dbcv", dbcv)213 return pers, dbcv, 0.5 # stability computed only for winner214 215 216# ---------------------------------------------------------------------------217# §3.4 — Run the full Bayesian loop218# ---------------------------------------------------------------------------219def run_bayesian_optimisation(220 embeddings: np.ndarray,221 n_trials: int = 50,222 progress_callback=None,223) -> dict:224 n_docs = len(embeddings)225 study = optuna.create_study(226 directions=["maximize", "maximize", "maximize"],227 sampler=optuna.samplers.TPESampler(seed=42, multivariate=True),228 study_name="specter2_umap_hdbscan",229 )230 trial_log = []231 232 def _cb(study, trial):233 d = trial.user_attrs.get("discipline", {})234 entry = dict(235 trial=trial.number,236 params=trial.user_attrs.get("params", {}),237 discipline_pass=trial.user_attrs.get("pass", False),238 persistence=trial.user_attrs.get("persistence", 0.0),239 dbcv=trial.user_attrs.get("dbcv", -1.0),240 n_clusters=d.get("n_clusters", 0),241 max_mass_pct=d.get("max_mass_pct", 0.0),242 min_size=d.get("min_size", 0),243 n_noise=d.get("n_noise", 0),244 values=list(trial.values) if trial.values else [],245 )246 trial_log.append(entry)247 if progress_callback:248 progress_callback(trial.number + 1, n_trials, entry)249 250 for i in range(n_trials):251 study.optimize(252 lambda t: _objective(t, embeddings, n_docs),253 n_trials=1, callbacks=[_cb], show_progress_bar=False,254 )255 # §3.6 convergence: 3 consecutive passing within 5 % of best256 passing = [e for e in trial_log if e["discipline_pass"]]257 if len(passing) >= 3 and i >= 9: # allow early stop after 10 trials258 best_p = max(e["persistence"] for e in passing)259 if best_p > 0:260 last3 = passing[-3:]261 if all(abs(e["persistence"] - best_p) / best_p < 0.05262 for e in last3):263 logger.info("Converged at trial %d.", i + 1)264 break265 266 # Select best passing trial (max persistence, then DBCV)267 passing_trials = [t for t in study.trials268 if t.user_attrs.get("pass", False)]269 if passing_trials:270 best = max(passing_trials, key=lambda t: (t.values[0], t.values[1]))271 else:272 logger.warning("No trial passed discipline — using last trial.")273 best = study.trials[-1]274 275 bp = best.user_attrs["params"]276 labels = np.array(best.user_attrs["labels"])277 stability = compute_stability(embeddings, bp, n_seeds=3)278 279 return dict(280 best_params=bp, best_labels=labels,281 best_trial=best.number,282 persistence=best.user_attrs.get("persistence", 0.0),283 dbcv=best.user_attrs.get("dbcv", -1.0),284 stability=stability,285 discipline=best.user_attrs.get("discipline", {}),286 trial_log=trial_log,287 n_trials_run=len(trial_log),288 )289 290 291# ---------------------------------------------------------------------------292# §3.1 — 2-D UMAP for visualisation293# ---------------------------------------------------------------------------294def compute_2d_umap(embeddings: np.ndarray, seed: int = 42) -> np.ndarray:295 return UMAP(n_neighbors=15, n_components=2, min_dist=0.1,296 metric="cosine", random_state=seed,297 low_memory=True).fit_transform(embeddings)298 299 300# ---------------------------------------------------------------------------301# §3.1 — TF-IDF keyphrase extraction per cluster (3–5 phrases)302# Fast alternative to KeyBERT — no extra model download needed.303# ---------------------------------------------------------------------------304def extract_keyphrases(docs: list[str], labels: np.ndarray,305 top_n: int = 5) -> dict:306 from sklearn.feature_extraction.text import TfidfVectorizer307 cluster_docs = defaultdict(list)308 for doc, lab in zip(docs, labels):309 if lab != -1:310 cluster_docs[int(lab)].append(doc)311 out = {}312 for cid, cdocs in cluster_docs.items():313 if len(cdocs) < 2:314 out[cid] = []315 continue316 try:317 tfidf = TfidfVectorizer(318 stop_words="english", max_features=200,319 ngram_range=(1, 3), max_df=0.9, min_df=1)320 X = tfidf.fit_transform(cdocs)321 terms = tfidf.get_feature_names_out()322 scores = X.sum(axis=0).A1323 top_idx = scores.argsort()[::-1][:top_n]324 out[cid] = [(terms[i], float(scores[i])) for i in top_idx]325 except Exception as e:326 logger.warning("Keyphrase extraction cluster %d: %s", cid, e)327 out[cid] = []328 return out329 330 331# ---------------------------------------------------------------------------332# §3.1 — Strong / weak member counts via HDBSCAN probabilities333# ---------------------------------------------------------------------------334def strong_weak_members(labels: np.ndarray,335 probabilities: np.ndarray) -> dict:336 mem = defaultdict(lambda: {"strong": 0, "weak": 0})337 for lab, prob in zip(labels, probabilities):338 if lab == -1:339 continue340 cid = int(lab)341 if prob >= 0.5:342 mem[cid]["strong"] += 1343 else:344 mem[cid]["weak"] += 1345 return dict(mem)346 347 348# ---------------------------------------------------------------------------349# §3.2 — Outlier reduction: reassign noise to nearest cluster (≤ 25 %)350# ---------------------------------------------------------------------------351def outlier_reduction(labels: np.ndarray, reduced: np.ndarray,352 n_docs: int) -> np.ndarray:353 labels = labels.copy()354 cap = int(0.25 * n_docs)355 cdocs = defaultdict(list)356 for i, l in enumerate(labels):357 if l != -1:358 cdocs[int(l)].append(i)359 if not cdocs:360 return labels361 cids = list(cdocs.keys())362 centroids = np.vstack([reduced[cdocs[c]].mean(axis=0) for c in cids])363 noise = [i for i, l in enumerate(labels) if l == -1]364 moved = 0365 for idx in noise:366 dists = np.linalg.norm(centroids - reduced[idx], axis=1)367 for best in np.argsort(dists):368 tgt = cids[best]369 if len(cdocs[tgt]) < cap:370 labels[idx] = tgt371 cdocs[tgt].append(idx)372 moved += 1373 break374 logger.info("Outlier reduction: %d / %d noise reassigned.", moved, len(noise))375 return labels376 377 378# ---------------------------------------------------------------------------379# Representative docs (top-3 by centroid proximity)380# ---------------------------------------------------------------------------381def get_representative_docs(labels, embeddings, docs, top_n=3):382 cdocs = defaultdict(list)383 for i, l in enumerate(labels):384 if l != -1:385 cdocs[int(l)].append(i)386 out = {}387 for cid, idxs in cdocs.items():388 ce = embeddings[idxs].mean(axis=0).reshape(1, -1)389 sims = cosine_similarity(ce, embeddings[idxs])[0]390 top = np.argsort(sims)[-top_n:][::-1]391 out[cid] = [docs[idxs[t]] for t in top]392 return out393 394 395# ---------------------------------------------------------------------------396# §9 — RQ2 / RQ3 mismatch table397# ---------------------------------------------------------------------------398def build_mismatch_table(keyphrases: dict, cluster_labels: dict) -> list:399 """Compare cluster keyphrases against assigned labels to flag mismatches.400 Returns rows for a mismatch table (§9)."""401 rows = []402 for cid in sorted(keyphrases.keys()):403 kps = keyphrases.get(cid, [])404 kp_terms = [k[0] if isinstance(k, tuple) else k for k in kps[:5]]405 label = cluster_labels.get(cid, "")406 # Check overlap between label words and keyphrase terms407 label_words = set(label.lower().split())408 kp_words = set(" ".join(kp_terms).lower().split())409 overlap = label_words & kp_words410 noise = {"the","and","for","with","using","based","from","in","of","a","to"}411 overlap -= noise412 match_pct = len(overlap) / max(len(label_words - noise), 1)413 status = "MATCH" if match_pct >= 0.3 else "MISMATCH"414 rows.append(dict(415 cluster=cid, label=label,416 keyphrases=", ".join(kp_terms),417 overlap=", ".join(overlap) if overlap else "—",418 match_pct=round(match_pct * 100),419 status=status,420 ))421 return rows422 423 424# ---------------------------------------------------------------------------425# High-level pipeline entry point426# ---------------------------------------------------------------------------427def run_topic_modeling(filepath: str, n_trials: int = 50,428 progress_callback=None) -> dict:429 # 1. Load430 df = load_csv(filepath)431 docs = prepare_documents(df)432 n_docs = len(docs)433 434 # 2. Embed (deterministic)435 embeddings = embed_documents(docs)436 437 # 3. Bayesian optimisation (§3.4)438 opt = run_bayesian_optimisation(embeddings, n_trials, progress_callback)439 bp = opt["best_params"]440 labels = opt["best_labels"]441 442 # 4. Re-run winner for clusterer object (probabilities)443 u = UMAP(n_neighbors=bp["n_neighbors"], n_components=bp["n_components"],444 min_dist=0.0, metric="cosine", random_state=42)445 red = u.fit_transform(embeddings)446 h = HDBSCAN(min_cluster_size=bp["min_cluster_size"],447 min_samples=bp["min_samples"], metric="euclidean",448 cluster_selection_method=bp["csm"],449 cluster_selection_epsilon=bp["cse"],450 allow_single_cluster=False,451 gen_min_span_tree=True)452 h.fit(red)453 454 # Per-cluster persistence (§8)455 cluster_pers = per_cluster_persistence(h, labels)456 457 # 5. Outlier reduction (§3.2 — clusters < 5 reassigned)458 labels = outlier_reduction(labels, red, n_docs)459 460 # 6. Strong / weak (§3.1)461 sw = strong_weak_members(labels, h.probabilities_)462 463 # 7. 2-D UMAP (§3.1)464 umap_2d = compute_2d_umap(embeddings)465 466 # 8. KeyBERT keyphrases (§3.1)467 keyphrases = extract_keyphrases(docs, labels)468 469 # 9. Rep docs470 rep_docs = get_representative_docs(labels, embeddings, docs)471 472 # 10. Final discipline473 disc = check_discipline(labels, n_docs)474 475 return dict(476 documents=docs, labels=labels.tolist(),477 keyphrases=keyphrases, representative_docs=rep_docs,478 membership=sw, umap_2d=umap_2d.tolist(),479 discipline=disc, best_params=bp,480 cluster_persistence=cluster_pers,481 metrics=dict(persistence=opt["persistence"],482 dbcv=opt["dbcv"],483 stability=opt["stability"]),484 trial_log=opt["trial_log"],485 n_trials_run=opt["n_trials_run"],486 best_trial=opt["best_trial"],487 n_docs=n_docs,488 embeddings=embeddings,489 )490