Spaces:
Running
Running
| """ | |
| computation_clusters.py — SPECTER2 semantic clustering for computation.csv | |
| Self-contained module that: | |
| 1. Loads computation.csv | |
| 2. Combines text fields → single string per paper | |
| 3. Generates SPECTER2 embeddings (cached as .npy) | |
| 4. UMAP 2D reduction | |
| 5. HDBSCAN clustering | |
| 6. Auto-labels clusters via TF-IDF keywords | |
| 7. Returns everything needed for the Streamlit Computational Strategy Map tab | |
| """ | |
| import os | |
| import hashlib | |
| import logging | |
| from typing import Dict, List, Tuple, Any, Optional | |
| import numpy as np | |
| import pandas as pd | |
| logger = logging.getLogger(__name__) | |
| # Cache directory for embeddings | |
| _CACHE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".comp_cache") | |
| os.makedirs(_CACHE_DIR, exist_ok=True) | |
| def _cache_key(texts: List[str]) -> str: | |
| """Stable hash of input texts for caching.""" | |
| h = hashlib.sha256() | |
| for t in texts: | |
| h.update(t.encode("utf-8", errors="replace")) | |
| return h.hexdigest()[:16] | |
| def _generate_specter2_embeddings(texts: List[str], batch_size: int = 8) -> np.ndarray: | |
| """ | |
| Generate SPECTER2 embeddings via HuggingFace adapters, with .npy caching. | |
| """ | |
| key = _cache_key(texts) | |
| cache_path = os.path.join(_CACHE_DIR, f"comp_emb_{key}.npy") | |
| if os.path.exists(cache_path): | |
| logger.info("✓ Loaded cached embeddings from %s", cache_path) | |
| return np.load(cache_path) | |
| import torch | |
| from transformers import AutoTokenizer | |
| from adapters import AutoAdapterModel | |
| logger.info("Loading SPECTER2 model …") | |
| tokenizer = AutoTokenizer.from_pretrained("allenai/specter2_base") | |
| model = AutoAdapterModel.from_pretrained("allenai/specter2_base") | |
| model.load_adapter("allenai/specter2", source="hf", load_as="proximity", set_active=True) | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| model.to(device) | |
| model.eval() | |
| embeddings = [] | |
| total_batches = (len(texts) + batch_size - 1) // batch_size | |
| logger.info("Generating embeddings for %d texts (%d batches) on %s …", len(texts), total_batches, device) | |
| for i in range(0, len(texts), batch_size): | |
| batch = texts[i : i + batch_size] | |
| inputs = tokenizer( | |
| batch, padding=True, truncation=True, | |
| return_tensors="pt", return_token_type_ids=False, max_length=512, | |
| ) | |
| inputs = {k: v.to(device) for k, v in inputs.items()} | |
| with torch.no_grad(): | |
| output = model(**inputs) | |
| batch_emb = output.last_hidden_state[:, 0, :].cpu().numpy() | |
| embeddings.append(batch_emb) | |
| result = np.vstack(embeddings) | |
| np.save(cache_path, result) | |
| logger.info("✓ Embeddings cached → %s", cache_path) | |
| return result | |
| def _umap_reduce(embeddings: np.ndarray, n_components: int = 2) -> np.ndarray: | |
| """UMAP dimensionality reduction.""" | |
| from umap import UMAP | |
| n_neighbors = min(15, len(embeddings) - 1) | |
| reducer = UMAP( | |
| n_neighbors=n_neighbors, | |
| n_components=n_components, | |
| min_dist=0.1, | |
| metric="cosine", | |
| random_state=42, | |
| ) | |
| return reducer.fit_transform(embeddings) | |
| def _hdbscan_cluster(embeddings: np.ndarray) -> np.ndarray: | |
| """ | |
| HDBSCAN clustering via UMAP 5D intermediate reduction. | |
| Tries multiple parameter combos to find the best clustering. | |
| Falls back gracefully for small datasets. | |
| """ | |
| from umap import UMAP | |
| from hdbscan import HDBSCAN | |
| n_samples = len(embeddings) | |
| # For very small datasets, use 3D intermediate | |
| n_intermediate = min(5, n_samples - 2) | |
| if n_intermediate < 2: | |
| n_intermediate = 2 | |
| n_neighbors = min(15, n_samples - 1) | |
| umap_inter = UMAP( | |
| n_neighbors=n_neighbors, | |
| n_components=n_intermediate, | |
| min_dist=0.0, | |
| metric="cosine", | |
| random_state=42, | |
| ) | |
| reduced = umap_inter.fit_transform(embeddings) | |
| best_labels = None | |
| best_n = 0 | |
| # Try parameter sweep | |
| for mcs in [3, 4, 5, 7]: | |
| for ms in [2, 3]: | |
| if mcs >= n_samples: | |
| continue | |
| clusterer = HDBSCAN( | |
| min_cluster_size=mcs, | |
| min_samples=min(ms, mcs - 1) if mcs > 1 else 1, | |
| metric="euclidean", | |
| cluster_selection_method="eom", | |
| ) | |
| labels = clusterer.fit_predict(reduced) | |
| n_clusters = len(set(labels)) - (1 if -1 in labels else 0) | |
| if n_clusters > best_n: | |
| best_labels = labels | |
| best_n = n_clusters | |
| if best_labels is None or best_n < 2: | |
| logger.warning("HDBSCAN produced < 2 clusters, using single cluster") | |
| best_labels = np.zeros(n_samples, dtype=int) | |
| # Reassign noise (-1) to nearest cluster centroid | |
| noise_mask = best_labels == -1 | |
| if noise_mask.any(): | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| valid_clusters = sorted(set(best_labels) - {-1}) | |
| if valid_clusters: | |
| centroids = np.array([embeddings[best_labels == c].mean(axis=0) for c in valid_clusters]) | |
| sims = cosine_similarity(embeddings[noise_mask], centroids) | |
| nearest = sims.argmax(axis=1) | |
| best_labels[noise_mask] = [valid_clusters[n] for n in nearest] | |
| logger.info("HDBSCAN: %d clusters from %d papers", len(set(best_labels)), n_samples) | |
| return best_labels | |
| def _label_clusters(texts: List[str], labels: np.ndarray) -> Dict[int, str]: | |
| """Generate cluster labels using TF-IDF top keywords.""" | |
| from sklearn.feature_extraction.text import TfidfVectorizer | |
| cluster_ids = sorted(set(labels)) | |
| # Merge all texts per cluster | |
| cluster_texts = {} | |
| for cid in cluster_ids: | |
| mask = labels == cid | |
| cluster_texts[cid] = " ".join([texts[i] for i in range(len(texts)) if mask[i]]) | |
| all_docs = [cluster_texts[c] for c in cluster_ids] | |
| vectorizer = TfidfVectorizer( | |
| stop_words="english", max_df=0.95, min_df=1, | |
| ngram_range=(1, 2), max_features=3000, | |
| ) | |
| tfidf = vectorizer.fit_transform(all_docs) | |
| features = vectorizer.get_feature_names_out() | |
| result = {} | |
| for i, cid in enumerate(cluster_ids): | |
| scores = tfidf[i].toarray().flatten() | |
| top_idx = scores.argsort()[-4:][::-1] | |
| keywords = [features[j] for j in top_idx if scores[j] > 0] | |
| label = " / ".join(w.title() for w in keywords[:3]) if keywords else f"Cluster {cid}" | |
| result[cid] = label | |
| return result | |
| def _build_cluster_keywords(texts: List[str], labels: np.ndarray) -> Dict[int, List[str]]: | |
| """Get top-5 keywords per cluster for the stats panel.""" | |
| from sklearn.feature_extraction.text import TfidfVectorizer | |
| cluster_ids = sorted(set(labels)) | |
| cluster_texts = {} | |
| for cid in cluster_ids: | |
| mask = labels == cid | |
| cluster_texts[cid] = " ".join([texts[i] for i in range(len(texts)) if mask[i]]) | |
| all_docs = [cluster_texts[c] for c in cluster_ids] | |
| vectorizer = TfidfVectorizer( | |
| stop_words="english", max_df=0.95, min_df=1, | |
| ngram_range=(1, 2), max_features=3000, | |
| ) | |
| tfidf = vectorizer.fit_transform(all_docs) | |
| features = vectorizer.get_feature_names_out() | |
| result = {} | |
| for i, cid in enumerate(cluster_ids): | |
| scores = tfidf[i].toarray().flatten() | |
| top_idx = scores.argsort()[-5:][::-1] | |
| result[cid] = [features[j] for j in top_idx if scores[j] > 0] | |
| return result | |
| def run_computation_clustering(csv_path: str = "computation.csv") -> Dict[str, Any]: | |
| """ | |
| Full pipeline: Load CSV → Embed → Reduce → Cluster → Label. | |
| Returns a dict with all data needed by the Streamlit tab: | |
| - df: DataFrame with original columns + cluster_id, cluster_name, x, y | |
| - cluster_names: {cluster_id: label} | |
| - cluster_keywords: {cluster_id: [kw1, kw2, ...]} | |
| - strategy_types: list of unique strategy types | |
| - n_clusters: int | |
| - n_papers: int | |
| """ | |
| logger.info("=" * 50) | |
| logger.info("Computation Strategy Clustering Pipeline") | |
| logger.info("=" * 50) | |
| # 1. Load CSV | |
| df = pd.read_csv(csv_path, encoding="utf-8", on_bad_lines="skip") | |
| df.columns = [c.strip() for c in df.columns] | |
| for col in df.columns: | |
| df[col] = df[col].fillna("") | |
| logger.info("Loaded %d papers from %s", len(df), csv_path) | |
| # 2. Combine text fields | |
| combined = [] | |
| for _, row in df.iterrows(): | |
| parts = [ | |
| str(row.get("Paper Title", "")), | |
| str(row.get("Computational Strategies Used", "")), | |
| str(row.get("Strategy Type", "")), | |
| str(row.get("Supporting Keywords / Evidence", "")), | |
| ] | |
| combined.append(" ".join(p for p in parts if p.strip())) | |
| # 3. SPECTER2 embeddings (cached) | |
| embeddings = _generate_specter2_embeddings(combined) | |
| # 4. UMAP 2D | |
| coords_2d = _umap_reduce(embeddings) | |
| # 5. HDBSCAN | |
| cluster_labels = _hdbscan_cluster(embeddings) | |
| # 6. TF-IDF labeling | |
| cluster_names = _label_clusters(combined, cluster_labels) | |
| cluster_keywords = _build_cluster_keywords(combined, cluster_labels) | |
| # 7. Attach to DataFrame | |
| df["cluster_id"] = cluster_labels.tolist() | |
| df["cluster_name"] = df["cluster_id"].map(cluster_names) | |
| df["x"] = coords_2d[:, 0].tolist() | |
| df["y"] = coords_2d[:, 1].tolist() | |
| strategy_types = sorted(set(s.strip() for s in df["Strategy Type"].unique() if s.strip())) | |
| logger.info("Pipeline complete: %d papers → %d clusters", len(df), len(cluster_names)) | |
| return { | |
| "df": df, | |
| "cluster_names": cluster_names, | |
| "cluster_keywords": cluster_keywords, | |
| "strategy_types": strategy_types, | |
| "n_clusters": len(cluster_names), | |
| "n_papers": len(df), | |
| } | |