""" tools.py — Core functions for the AI-driven topic modeling pipeline. This module provides all analytical functions used by the TopicAgent: - CSV ingestion and validation - Text preprocessing (lowercasing, stopword removal, cleaning) - Topic modeling via BERTopic (with fallback to sklearn LDA) - Automatic human-readable label generation - Cross-source theme comparison (Title vs Abstract) - Taxonomy mapping (MAPPED / NOVEL classification) """ import re import json import logging from typing import Dict, List, Tuple, Optional, Any import numpy as np import pandas as pd import nltk from nltk.corpus import stopwords from nltk.tokenize import word_tokenize # --------------------------------------------------------------------------- # Logging # --------------------------------------------------------------------------- logger = logging.getLogger(__name__) logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") # --------------------------------------------------------------------------- # NLTK data download (idempotent) # --------------------------------------------------------------------------- for _resource in ("punkt", "punkt_tab", "stopwords"): try: nltk.data.find(f"tokenizers/{_resource}" if "punkt" in _resource else f"corpora/{_resource}") except LookupError: nltk.download(_resource, quiet=True) # --------------------------------------------------------------------------- # Reference taxonomy of known AI / business / research themes # Used by create_taxonomy_map() for MAPPED vs NOVEL classification # --------------------------------------------------------------------------- KNOWN_THEMES: List[str] = [ # AI / ML "artificial intelligence", "machine learning", "deep learning", "neural network", "natural language processing", "computer vision", "reinforcement learning", "generative ai", "large language model", "transformer", "chatbot", "recommendation system", "knowledge graph", "robotics", "autonomous", "explainable ai", "federated learning", "transfer learning", "ai ethics", "adversarial", "gan", "diffusion model", "prompt engineering", # Data science "data mining", "big data", "analytics", "data science", "data quality", "feature engineering", "dimensionality reduction", "clustering", "classification", "regression", "time series", "anomaly detection", "sentiment analysis", # Business / Management "digital transformation", "innovation", "strategy", "supply chain", "customer experience", "marketing", "e-commerce", "fintech", "blockchain", "sustainability", "corporate social responsibility", "knowledge management", "decision support", "business intelligence", "enterprise", "organizational", "human resource", "leadership", "entrepreneurship", "business model", # Information systems "information systems", "technology adoption", "user acceptance", "privacy", "security", "trust", "social media", "online community", "platform", "crowdsourcing", "cloud computing", "iot", "internet of things", "software engineering", "agile", "devops", "digital platform", # Healthcare / Society "healthcare", "telemedicine", "electronic health", "public health", "education", "e-learning", "smart city", "government", "policy", "ethics", "fairness", "bias", "misinformation", "content moderation", # Research methods "survey", "experiment", "case study", "meta-analysis", "bibliometric", "systematic review", "structural equation", "grounded theory", ] # =================================================================== # 1. load_csv — Ingest and validate the CSV dataset # =================================================================== def load_csv(filepath: str) -> pd.DataFrame: """ Load a CSV file and ensure the required columns (Title, Abstract) exist. Parameters ---------- filepath : str Path to the CSV file. Returns ------- pd.DataFrame DataFrame with at least 'Title' and 'Abstract' columns. Raises ------ FileNotFoundError If the specified file does not exist. ValueError If required columns are missing. """ logger.info("Loading CSV from %s", filepath) df = pd.read_csv(filepath, encoding="utf-8", on_bad_lines="skip") logger.info("Loaded %d rows × %d columns", len(df), len(df.columns)) # Validate required columns (case-insensitive match) col_map = {str(c).strip().lower(): c for c in df.columns} # Title and Abstract are strictly required if "title" not in col_map or "abstract" not in col_map: missing = {"title", "abstract"} - set(col_map.keys()) raise ValueError(f"CSV is missing required columns: {missing}. Found: {list(df.columns)}") # Map Title and Abstract rename_dict = { col_map["title"]: "Title", col_map["abstract"]: "Abstract" } # Map Keywords keyword_cols = ["keywords", "author keywords", "index keywords"] found_keyword = next((c for c in keyword_cols if c in col_map), None) if found_keyword: rename_dict[col_map[found_keyword]] = "Keywords" # Map Full Text fulltext_cols = ["full text", "body"] found_fulltext = next((c for c in fulltext_cols if c in col_map), None) if found_fulltext: rename_dict[col_map[found_fulltext]] = "Full Text" # Rename to canonical form df = df.rename(columns=rename_dict) # Fill missing columns if "Keywords" not in df.columns: df["Keywords"] = "" if "Full Text" not in df.columns: df["Full Text"] = df["Abstract"] # Drop rows where all of these are empty df = df.dropna(subset=["Title", "Abstract", "Keywords", "Full Text"], how="all").reset_index(drop=True) df["Title"] = df["Title"].fillna("") df["Abstract"] = df["Abstract"].fillna("") df["Keywords"] = df["Keywords"].fillna("") df["Full Text"] = df["Full Text"].fillna("") logger.info("After cleaning: %d usable rows", len(df)) return df # =================================================================== # 2. preprocess_text — Clean and normalise a list of text documents # =================================================================== def preprocess_text(documents: List[str]) -> List[str]: """ Apply professional-grade text preprocessing: 1. Lowercase 2. Remove URLs, emails, special characters, digits 3. Tokenize 4. Remove stopwords (NLTK English) 5. Remove very short tokens (length ≤ 2) 6. Rejoin into cleaned strings Parameters ---------- documents : list of str Raw text documents. Returns ------- list of str Cleaned text documents. """ stop_words = set(stopwords.words("english")) # Extended stopwords common in academic abstracts stop_words.update([ "©", "elsevier", "rights", "reserved", "doi", "http", "https", "vol", "pp", "fig", "table", "journal", "author", "authors", "study", "paper", "research", "results", "findings", "however", "propose", "proposed", "approach", "using", "based", "also", "show", "shows", "shown", "may", "used", "use", "one", "two", "three", "new", "well", "within", "among", "across", "toward", "towards", "et", "al", "ie", "eg", "cf", "thus", "therefore", "moreover", "furthermore", "addition", "conclusion", "conclusions", ]) cleaned: List[str] = [] for doc in documents: if not isinstance(doc, str) or not doc.strip(): cleaned.append("") continue text = doc.lower() # Remove URLs text = re.sub(r"https?://\S+|www\.\S+", " ", text) # Remove emails text = re.sub(r"\S+@\S+", " ", text) # Remove digits and special characters but keep spaces text = re.sub(r"[^a-z\s]", " ", text) # Collapse whitespace text = re.sub(r"\s+", " ", text).strip() # Tokenize and filter tokens = word_tokenize(text) tokens = [t for t in tokens if t not in stop_words and len(t) > 2] cleaned.append(" ".join(tokens)) logger.info("Preprocessed %d documents", len(cleaned)) return cleaned # =================================================================== # 3. run_topic_modeling — Discover clusters + sub-topics # =================================================================== def run_topic_modeling( documents: List[str], source_label: str = "documents", min_topics: int = 100, max_clusters: int = 25, use_bertopic: bool = True, ) -> Tuple[pd.DataFrame, Any]: """ Perform two-level topic modeling on a corpus of preprocessed documents. Strategy: 1. First pass: Discover clusters (target 10–max_clusters range) using BERTopic with UMAP + HDBSCAN, or LDA fallback. 2. Second pass: Within each cluster, extract fine-grained sub-topics using NMF to guarantee ≥ min_topics total topics. Parameters ---------- documents : list of str Preprocessed text documents. source_label : str Label for logging (e.g. "Titles" or "Abstracts"). min_topics : int Minimum total number of topics (sub-topics) required (default 100). max_clusters : int Target maximum number of high-level clusters (default 25). use_bertopic : bool Whether to attempt BERTopic first. Returns ------- topics_df : pd.DataFrame Columns: cluster_id, topic_id, sub_topic_id, keywords, label, source, n_documents model : object The fitted cluster-level model for downstream inspection. """ from sklearn.decomposition import NMF from sklearn.feature_extraction.text import TfidfVectorizer # Filter out empty documents valid_docs = [d for d in documents if d.strip()] if len(valid_docs) < 20: raise ValueError(f"Not enough valid documents ({len(valid_docs)}) for topic modeling.") logger.info( "Running topic modeling on %d %s (target: %d–%d clusters, ≥ %d topics)", len(valid_docs), source_label, 10, max_clusters, min_topics, ) # ---- FIRST PASS: Discover clusters ---- cluster_df = None model = None if use_bertopic: try: cluster_df, model = _run_bertopic(valid_docs, source_label, max_clusters) except Exception as exc: logger.warning("BERTopic failed (%s). Falling back to LDA.", exc) cluster_df = None if cluster_df is None or len(cluster_df) < 3: logger.info("Using LDA for clustering %s (target %d clusters)", source_label, max_clusters) cluster_df, model = _run_lda(valid_docs, source_label, max_clusters) n_clusters = len(cluster_df) logger.info("First pass: %d clusters discovered for %s", n_clusters, source_label) # ---- SECOND PASS: Extract sub-topics within each cluster ---- # Assign each document to its nearest cluster based on keyword overlap cluster_assignments = _assign_docs_to_clusters(valid_docs, cluster_df) # Calculate sub-topic counts per cluster cluster_sizes = {} for cid in cluster_df["topic_id"].tolist(): cluster_sizes[cid] = int((np.array(cluster_assignments) == cid).sum()) total_assigned = sum(cluster_sizes.values()) subtopic_counts = {} for cid in cluster_df["topic_id"].tolist(): if cluster_sizes[cid] == 0: subtopic_counts[cid] = 1 continue raw_share = (cluster_sizes[cid] / max(total_assigned, 1)) * min_topics subtopic_counts[cid] = max(2, round(raw_share)) # Ensure total >= min_topics current_total = sum(subtopic_counts.values()) if current_total < min_topics: deficit = min_topics - current_total sorted_clusters = sorted( cluster_df["topic_id"].tolist(), key=lambda c: cluster_sizes.get(c, 0), reverse=True, ) i = 0 while deficit > 0 and i < len(sorted_clusters) * 10: cid = sorted_clusters[i % len(sorted_clusters)] max_possible = max(cluster_sizes.get(cid, 0) // 2, 2) if subtopic_counts[cid] < max_possible: subtopic_counts[cid] += 1 deficit -= 1 i += 1 # Extract sub-topics per cluster all_topics = [] global_topic_id = 0 for _, cluster_row in cluster_df.iterrows(): cid = int(cluster_row["topic_id"]) # Get documents in this cluster doc_indices = [i for i, a in enumerate(cluster_assignments) if a == cid] cluster_docs = [valid_docs[i] for i in doc_indices] n_subtopics = subtopic_counts.get(cid, 2) if len(cluster_docs) < 2: # Single document cluster — single topic all_topics.append({ "cluster_id": cid, "topic_id": global_topic_id, "sub_topic_id": 0, "keywords": cluster_row["keywords"], "source": source_label, "n_documents": len(cluster_docs), }) global_topic_id += 1 continue n_subtopics = min(n_subtopics, len(cluster_docs)) if n_subtopics < 2: n_subtopics = 1 try: vectorizer = TfidfVectorizer( stop_words="english", max_df=0.95, min_df=1, ngram_range=(1, 2), max_features=5000, ) tfidf_matrix = vectorizer.fit_transform(cluster_docs) feature_names = vectorizer.get_feature_names_out() if n_subtopics == 1 or tfidf_matrix.shape[0] < 2: scores = tfidf_matrix.sum(axis=0).A1 top_indices = scores.argsort()[-10:][::-1] kw = ", ".join([feature_names[i] for i in top_indices]) all_topics.append({ "cluster_id": cid, "topic_id": global_topic_id, "sub_topic_id": 0, "keywords": kw, "source": source_label, "n_documents": len(cluster_docs), }) global_topic_id += 1 else: nmf = NMF(n_components=n_subtopics, random_state=42, max_iter=300, init="nndsvda") W = nmf.fit_transform(tfidf_matrix) H = nmf.components_ doc_assignments = W.argmax(axis=1) for sub_id in range(n_subtopics): top_indices = H[sub_id].argsort()[-10:][::-1] kw = ", ".join([feature_names[i] for i in top_indices]) n_docs = int((doc_assignments == sub_id).sum()) if n_docs == 0: n_docs = 1 all_topics.append({ "cluster_id": cid, "topic_id": global_topic_id, "sub_topic_id": sub_id, "keywords": kw, "source": source_label, "n_documents": n_docs, }) global_topic_id += 1 except Exception as e: logger.warning("NMF failed for cluster %d: %s. Using cluster keywords.", cid, e) all_topics.append({ "cluster_id": cid, "topic_id": global_topic_id, "sub_topic_id": 0, "keywords": cluster_row["keywords"], "source": source_label, "n_documents": len(cluster_docs), }) global_topic_id += 1 topics_df = pd.DataFrame(all_topics) # Add cluster_label from cluster keywords cluster_label_map = {} for _, row in cluster_df.iterrows(): cid = int(row["topic_id"]) kws = [kw.strip().title() for kw in row["keywords"].split(",")][:3] cluster_label_map[cid] = " / ".join(kws) topics_df["cluster_label"] = topics_df["cluster_id"].map(cluster_label_map).fillna("") logger.info( "Topic modeling complete for %s: %d topics in %d clusters", source_label, len(topics_df), n_clusters, ) return topics_df, model def _assign_docs_to_clusters(docs: List[str], cluster_df: pd.DataFrame) -> List[int]: """Assign each document to its most similar cluster based on keyword overlap.""" cluster_ids = cluster_df["topic_id"].tolist() cluster_keywords = [] for _, row in cluster_df.iterrows(): kws = set(kw.strip().lower() for kw in row["keywords"].split(",") if kw.strip()) # Also add individual words from multi-word keywords expanded = set() for kw in kws: expanded.update(kw.split()) expanded.update(kws) cluster_keywords.append(expanded) assignments = [] for doc in docs: doc_tokens = set(doc.lower().split()) best_score = -1 best_cid = cluster_ids[0] # default to first cluster for cid, kws in zip(cluster_ids, cluster_keywords): overlap = len(doc_tokens & kws) if overlap > best_score: best_score = overlap best_cid = cid assignments.append(best_cid) return assignments def _run_bertopic(docs: List[str], source_label: str, max_clusters: int): """Run BERTopic to discover clusters (not topics).""" from bertopic import BERTopic from umap import UMAP from hdbscan import HDBSCAN from sklearn.feature_extraction.text import CountVectorizer umap_model = UMAP( n_neighbors=10, n_components=5, min_dist=0.0, metric="cosine", random_state=42, ) hdbscan_model = HDBSCAN( min_cluster_size=max(5, len(docs) // max_clusters), min_samples=2, prediction_data=True, ) vectorizer = CountVectorizer( stop_words="english", ngram_range=(1, 2), max_df=0.90, min_df=2, ) topic_model = BERTopic( umap_model=umap_model, hdbscan_model=hdbscan_model, vectorizer_model=vectorizer, nr_topics=max_clusters, top_n_words=10, verbose=False, ) topics, _probs = topic_model.fit_transform(docs) info = topic_model.get_topic_info() # Exclude outlier topic (-1) info = info[info["Topic"] != -1].reset_index(drop=True) rows = [] for _, row in info.iterrows(): tid = int(row["Topic"]) topic_words = topic_model.get_topic(tid) kw = ", ".join([w for w, _ in topic_words[:10]]) rows.append({"topic_id": tid, "keywords": kw, "source": source_label}) df = pd.DataFrame(rows) return df, topic_model def _run_lda(docs: List[str], source_label: str, n_clusters: int): """Run sklearn LDA to discover clusters.""" from sklearn.decomposition import LatentDirichletAllocation from sklearn.feature_extraction.text import CountVectorizer vectorizer = CountVectorizer( stop_words="english", max_df=0.90, min_df=2, ngram_range=(1, 2), max_features=10000, ) dtm = vectorizer.fit_transform(docs) feature_names = vectorizer.get_feature_names_out() lda = LatentDirichletAllocation( n_components=n_clusters, max_iter=25, learning_method="online", random_state=42, n_jobs=-1, ) lda.fit(dtm) rows = [] for idx, component in enumerate(lda.components_): top_indices = component.argsort()[-10:][::-1] kw = ", ".join([feature_names[i] for i in top_indices]) rows.append({"topic_id": idx, "keywords": kw, "source": source_label}) df = pd.DataFrame(rows) return df, lda # =================================================================== # 4. generate_labels — Create human-readable labels for each topic # =================================================================== def generate_labels( topics_df: pd.DataFrame, use_llm: bool = False, groq_api_key: Optional[str] = None, ) -> pd.DataFrame: """ Generate a short human-readable label for every topic. Strategy: - If use_llm=True and a Groq API key is provided, use the Groq LLM (llama-3.3-70b-versatile, free tier) to produce contextual labels. - Otherwise, apply a heuristic: capitalise the first 3–4 keywords. Parameters ---------- topics_df : pd.DataFrame Must contain columns 'topic_id' and 'keywords'. use_llm : bool Whether to use the Groq LLM for label generation. groq_api_key : str, optional Groq API key, required if use_llm is True. Returns ------- pd.DataFrame Same DataFrame with an additional 'label' column. """ if use_llm and groq_api_key: logger.info("Generating labels using Groq LLM …") topics_df = _generate_labels_llm(topics_df, groq_api_key) else: logger.info("Generating labels using keyword heuristic …") topics_df = _generate_labels_heuristic(topics_df) return topics_df def _generate_labels_heuristic(df: pd.DataFrame) -> pd.DataFrame: """Create labels from the top keywords of each topic.""" labels = [] for _, row in df.iterrows(): kws = [kw.strip() for kw in row["keywords"].split(",")] # Take the first 3-4 non-trivial keywords and title-case them candidates = [kw.title() for kw in kws if len(kw) > 2][:4] label = " / ".join(candidates) if candidates else f"Topic {row['topic_id']}" labels.append(label) df = df.copy() df["label"] = labels return df def _generate_labels_llm(df: pd.DataFrame, api_key: str) -> pd.DataFrame: """Use Groq API to generate contextual labels for topics (batched).""" import time try: from groq import Groq except ImportError: logger.warning("groq package not installed. Falling back to heuristic labels.") return _generate_labels_heuristic(df) client = Groq(api_key=api_key) labels = [] # Process in batches to avoid rate limits batch_size = 10 for batch_start in range(0, len(df), batch_size): batch = df.iloc[batch_start:batch_start + batch_size] prompt_lines = [] for _, row in batch.iterrows(): prompt_lines.append(f"Topic {row['topic_id']}: keywords = [{row['keywords']}]") prompt = ( "You are a research taxonomy expert. For each topic below, " "generate a concise, descriptive label (3-6 words) that captures " "the theme of the keywords. Return ONLY a JSON list of objects " 'with keys "topic_id" and "label". No extra text.\n\n' + "\n".join(prompt_lines) ) try: chat = client.chat.completions.create( model="llama-3.3-70b-versatile", messages=[{"role": "user", "content": prompt}], temperature=0.3, max_tokens=1024, ) resp = chat.choices[0].message.content.strip() # Parse JSON from the response # Find JSON array in response json_match = re.search(r"\[.*\]", resp, re.DOTALL) if json_match: batch_labels = json.loads(json_match.group()) label_map = {item["topic_id"]: item["label"] for item in batch_labels} for _, row in batch.iterrows(): labels.append(label_map.get(row["topic_id"], f"Topic {row['topic_id']}")) else: # Fallback for this batch for _, row in batch.iterrows(): kws = [kw.strip().title() for kw in row["keywords"].split(",")][:4] labels.append(" / ".join(kws)) except Exception as exc: logger.warning("Groq API error for batch starting at %d: %s", batch_start, exc) for _, row in batch.iterrows(): kws = [kw.strip().title() for kw in row["keywords"].split(",")][:4] labels.append(" / ".join(kws)) # Rate-limit courtesy delay time.sleep(0.5) df = df.copy() df["label"] = labels return df # =================================================================== # 5. compare_themes — Cross-compare title vs abstract topics # =================================================================== def compare_themes( title_topics: pd.DataFrame, abstract_topics: pd.DataFrame, ) -> pd.DataFrame: """ Build a comparison table showing dominant themes from titles and abstracts side-by-side. Matching strategy: - Compute keyword overlap (Jaccard similarity) between every title-topic and abstract-topic pair. - For each title-topic, find the best matching abstract-topic. - Report similarity score and alignment status. Parameters ---------- title_topics : pd.DataFrame Topics extracted from titles (with 'topic_id', 'keywords', 'label'). abstract_topics : pd.DataFrame Topics extracted from abstracts (with 'topic_id', 'keywords', 'label'). Returns ------- pd.DataFrame Comparison table with columns: title_topic_id, title_label, title_keywords, abstract_topic_id, abstract_label, abstract_keywords, similarity, alignment """ logger.info("Comparing themes: %d title topics × %d abstract topics", len(title_topics), len(abstract_topics)) def _keywords_set(kw_str: str) -> set: return set(kw.strip().lower() for kw in kw_str.split(",") if kw.strip()) rows = [] for _, t_row in title_topics.iterrows(): t_kws = _keywords_set(t_row["keywords"]) best_sim = 0.0 best_match = None for _, a_row in abstract_topics.iterrows(): a_kws = _keywords_set(a_row["keywords"]) if not t_kws or not a_kws: continue # Jaccard similarity intersection = len(t_kws & a_kws) union = len(t_kws | a_kws) sim = intersection / union if union else 0.0 if sim > best_sim: best_sim = sim best_match = a_row alignment = ( "Strong" if best_sim >= 0.4 else "Moderate" if best_sim >= 0.2 else "Weak" if best_sim > 0 else "No Match" ) rows.append({ "title_topic_id": t_row["topic_id"], "title_label": t_row.get("label", ""), "title_keywords": t_row["keywords"], "abstract_topic_id": best_match["topic_id"] if best_match is not None else None, "abstract_label": best_match.get("label", "") if best_match is not None else "", "abstract_keywords": best_match["keywords"] if best_match is not None else "", "similarity": round(best_sim, 4), "alignment": alignment, }) comparison_df = pd.DataFrame(rows) logger.info("Theme comparison complete: %d rows", len(comparison_df)) return comparison_df # =================================================================== # 6. create_taxonomy_map — Classify themes as MAPPED or NOVEL # =================================================================== def create_taxonomy_map( topics_df: pd.DataFrame, known_themes: Optional[List[str]] = None, threshold: float = 0.15, ) -> Dict[str, Any]: """ Classify each topic as either MAPPED (similar to a well-known AI / business / IS research theme) or NOVEL (previously unseen). Heuristic: For each topic's keyword set, compute its best token-overlap ratio against the known themes list. If the ratio exceeds the threshold, label it as MAPPED; otherwise NOVEL. Parameters ---------- topics_df : pd.DataFrame Must contain 'topic_id', 'keywords', and 'label' columns. known_themes : list of str, optional Reference themes (defaults to the built-in KNOWN_THEMES). threshold : float Minimum overlap ratio to classify as MAPPED. Returns ------- dict JSON-serialisable taxonomy map: { "metadata": { ... }, "mapped": [ {topic_id, label, keywords, matched_theme, score}, ... ], "novel": [ {topic_id, label, keywords, score}, ... ], } """ if known_themes is None: known_themes = KNOWN_THEMES logger.info("Building taxonomy map for %d topics (threshold=%.2f)", len(topics_df), threshold) mapped: List[Dict] = [] novel: List[Dict] = [] known_tokens_list = [set(theme.lower().split()) for theme in known_themes] for _, row in topics_df.iterrows(): topic_tokens = set( kw.strip().lower() for kw in row["keywords"].split(",") if kw.strip() ) # Also include individual words from multi-word keywords expanded_tokens = set() for token in topic_tokens: expanded_tokens.update(token.split()) expanded_tokens.update(topic_tokens) best_score = 0.0 best_theme = "" for theme_str, theme_tokens in zip(known_themes, known_tokens_list): if not expanded_tokens or not theme_tokens: continue intersection = len(expanded_tokens & theme_tokens) union_size = len(expanded_tokens | theme_tokens) score = intersection / union_size if union_size else 0.0 if score > best_score: best_score = score best_theme = theme_str entry = { "topic_id": int(row["topic_id"]), "label": row.get("label", ""), "keywords": row["keywords"], "score": round(best_score, 4), } if best_score >= threshold: entry["matched_theme"] = best_theme entry["classification"] = "MAPPED" mapped.append(entry) else: entry["classification"] = "NOVEL" novel.append(entry) taxonomy = { "metadata": { "total_topics": len(topics_df), "mapped_count": len(mapped), "novel_count": len(novel), "threshold": threshold, }, "mapped": mapped, "novel": novel, } logger.info("Taxonomy: %d MAPPED, %d NOVEL", len(mapped), len(novel)) return taxonomy # =================================================================== # TCCM Extraction Tool # =================================================================== def extract_tccm(documents: List[str], api_key: str) -> pd.DataFrame: # Basic full-text TCCM extraction stub logger.info("Extracting TCCM from full text...") return pd.DataFrame([{"TCCM": "Sample"} for _ in documents]) # =================================================================== # SPECTER2 Pipeline Functions # =================================================================== def generate_specter2_embeddings(papers: List[Dict], batch_size: int = 16) -> np.ndarray: """ Generate embeddings using SPECTER2 via adapters. papers: list of dicts with 'title' and 'abstract' keys. """ import torch from transformers import AutoTokenizer from adapters import AutoAdapterModel logger.info("Loading SPECTER2 model and tokenizer...") 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 = [] logger.info("Generating embeddings for %d papers using %s...", len(papers), device) for i in range(0, len(papers), batch_size): batch = papers[i:i+batch_size] text_batch = [d['title'] + tokenizer.sep_token + (d.get('abstract') or '') for d in batch] inputs = tokenizer( text_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_embeddings = output.last_hidden_state[:, 0, :].cpu().numpy() embeddings.append(batch_embeddings) return np.vstack(embeddings) def store_in_chroma(embeddings, documents, metadatas, collection_name="specter_vectors"): """Store embeddings in ChromaDB for persistence.""" import chromadb client = chromadb.PersistentClient(path="./chroma_db") # Delete existing collection to avoid ID conflicts try: client.delete_collection(name=collection_name) except Exception: pass collection = client.get_or_create_collection(name=collection_name) ids = [str(i) for i in range(len(documents))] collection.add( embeddings=embeddings.tolist(), documents=documents, metadatas=metadatas, ids=ids, ) return collection def load_embeddings_from_chroma(collection_name="specter_vectors"): """Load cached embeddings from ChromaDB.""" import chromadb client = chromadb.PersistentClient(path="./chroma_db") collection = client.get_collection(name=collection_name) result = collection.get(include=["embeddings", "documents", "metadatas"]) if not result["embeddings"]: raise ValueError("No embeddings found in ChromaDB.") embeddings = np.array(result["embeddings"]) documents = result["documents"] metadatas = result["metadatas"] return embeddings, documents, metadatas def adaptive_hdbscan_clustering( embeddings: np.ndarray, min_clusters: int = 10, max_clusters: int = 40, ) -> tuple: """ Adaptive HDBSCAN clustering with UMAP dimensionality reduction. Tries different min_cluster_size values to land within the target cluster range. Returns: (labels, reduced_2d, n_clusters, params_used) """ from umap import UMAP from hdbscan import HDBSCAN logger.info( "Adaptive HDBSCAN: targeting %d–%d clusters for %d embeddings", min_clusters, max_clusters, len(embeddings), ) # UMAP reduction to 5D for clustering umap_5d = UMAP( n_neighbors=15, n_components=5, min_dist=0.0, metric="cosine", random_state=42, ) reduced_5d = umap_5d.fit_transform(embeddings) # UMAP reduction to 2D for visualization umap_2d = UMAP( n_neighbors=15, n_components=2, min_dist=0.1, metric="cosine", random_state=42, ) reduced_2d = umap_2d.fit_transform(embeddings) # Try different min_cluster_size to find good cluster count best_labels = None best_n_clusters = 0 best_params = {} best_diff = float("inf") target_mid = (min_clusters + max_clusters) / 2 for mcs in [5, 8, 10, 15, 20, 25, 30, 40, 50, 60, 80, 100, 120, 150]: for ms in [2, 3, 5]: clusterer = HDBSCAN( min_cluster_size=mcs, min_samples=ms, metric="euclidean", cluster_selection_method="eom", ) labels = clusterer.fit_predict(reduced_5d) n_clusters = len(set(labels)) - (1 if -1 in labels else 0) diff = abs(n_clusters - target_mid) in_range = min_clusters <= n_clusters <= max_clusters if in_range and diff < best_diff: best_labels = labels best_n_clusters = n_clusters best_params = {"min_cluster_size": mcs, "min_samples": ms} best_diff = diff if n_clusters == round(target_mid): break # Fallback: if nothing landed in range, pick closest if best_labels is None: logger.warning("Could not find params for target range, using best effort") for mcs in [15, 20, 25, 30, 40, 50]: clusterer = HDBSCAN( min_cluster_size=mcs, min_samples=3, metric="euclidean", cluster_selection_method="eom", ) labels = clusterer.fit_predict(reduced_5d) n_clusters = len(set(labels)) - (1 if -1 in labels else 0) if n_clusters >= 5: best_labels = labels best_n_clusters = n_clusters best_params = {"min_cluster_size": mcs, "min_samples": 3} break if best_labels is None: raise RuntimeError("HDBSCAN could not produce any valid clusters.") logger.info( "HDBSCAN selected: %d clusters (params: %s)", best_n_clusters, best_params, ) return best_labels, reduced_2d, best_n_clusters, best_params def c_tf_idf(documents, m, ngram_range=(1, 1)): """Compute class-based TF-IDF for topic keyword extraction.""" from sklearn.feature_extraction.text import CountVectorizer count = CountVectorizer(ngram_range=ngram_range, stop_words="english").fit(documents) t = count.transform(documents).toarray() w = t.sum(axis=1) tf = np.divide(t.T, w) sum_t = t.sum(axis=0) idf = np.log(np.divide(m, sum_t)).reshape(-1, 1) tf_idf = np.multiply(tf, idf) return tf_idf, count def extract_top_n_words_per_topic(tf_idf, count, docs_per_topic, n=10): """Extract top N keywords per topic from c-TF-IDF matrix.""" words = count.get_feature_names_out() labels = list(docs_per_topic.topic_id) tf_idf_transposed = tf_idf.T indices = tf_idf_transposed.argsort()[:, -n:] top_n_words = { label: [(words[j], tf_idf_transposed[i][j]) for j in indices[i]][::-1] for i, label in enumerate(labels) } return top_n_words def extract_representative_sentences( df: pd.DataFrame, embeddings: np.ndarray, cluster_labels: np.ndarray, top_n: int = 5, ) -> Dict[int, List[str]]: """ For each cluster, find the top_n documents closest to the cluster centroid. Returns {cluster_id: [sentence1, sentence2, ...]}. """ from sklearn.metrics.pairwise import cosine_similarity result = {} unique_labels = set(cluster_labels) for cid in unique_labels: if cid == -1: continue mask = cluster_labels == cid cluster_embs = embeddings[mask] cluster_indices = np.where(mask)[0] if len(cluster_embs) == 0: continue centroid = cluster_embs.mean(axis=0).reshape(1, -1) sims = cosine_similarity(centroid, cluster_embs)[0] top_indices = sims.argsort()[-top_n:][::-1] sentences = [] for idx in top_indices: doc_idx = cluster_indices[idx] title = str(df.iloc[doc_idx].get("Title", "")) abstract = str(df.iloc[doc_idx].get("Abstract", "")) # Use title + first sentence of abstract first_sent = abstract.split(".")[0] if abstract else "" rep = f"{title}. {first_sent}".strip() if rep: sentences.append(rep) result[int(cid)] = sentences return result def get_council_debate(keywords: str, groq_api_key: str) -> Dict[str, str]: """ Run a 3-agent AI Council debate to generate a topic label. Returns dict with: expert, analyst, anti_hallucination, final_label. """ import time try: from groq import Groq except ImportError: kws = [kw.strip().title() for kw in keywords.split(",")][:3] return { "expert": ", ".join(kws), "analyst": ", ".join(kws), "anti_hallucination": "N/A", "final_label": ", ".join(kws), } client = Groq(api_key=groq_api_key) # Persona 1: Domain Expert prompt_expert = ( f"You are a scientific Domain Expert. Keywords: [{keywords}]. " "Propose a 3-5 word descriptive label for this academic topic. " "Return ONLY the label." ) # Persona 2: Data Analyst prompt_analyst = ( f"You are a Data Analyst categorizing research clusters. Keywords: [{keywords}]. " "Give a short, precise categorical name (3-5 words). Return ONLY the name." ) try: res_expert = client.chat.completions.create( model="llama-3.3-70b-versatile", messages=[{"role": "user", "content": prompt_expert}], temperature=0.4, max_tokens=20, ).choices[0].message.content.strip() time.sleep(0.3) res_analyst = client.chat.completions.create( model="llama-3.3-70b-versatile", messages=[{"role": "user", "content": prompt_analyst}], temperature=0.4, max_tokens=20, ).choices[0].message.content.strip() time.sleep(0.3) # Persona 3: Anti-hallucination checker prompt_check = ( f"Keywords: [{keywords}].\n" f"Expert label: '{res_expert}'.\n" f"Analyst label: '{res_analyst}'.\n" "Check: are these labels grounded in the keywords? " "If yes, synthesize into a final 3-5 word label. " "If not, create a grounded label from the keywords only. " "Return ONLY the final label." ) res_judge = client.chat.completions.create( model="llama-3.3-70b-versatile", messages=[{"role": "user", "content": prompt_check}], temperature=0.2, max_tokens=20, ).choices[0].message.content.strip() res_judge = res_judge.replace('"', '').replace("'", "").strip() time.sleep(0.3) return { "expert": res_expert, "analyst": res_analyst, "anti_hallucination": res_judge, "final_label": res_judge, } except Exception as e: logger.warning("Council debate failed: %s", e) kws = [kw.strip().title() for kw in keywords.split(",")][:3] return { "expert": ", ".join(kws), "analyst": ", ".join(kws), "anti_hallucination": "fallback", "final_label": ", ".join(kws), } def build_cluster_graph( topics_df: pd.DataFrame, embeddings: np.ndarray, cluster_labels: np.ndarray, top_n_edges: int = 20, ) -> List[Dict]: """ Build a similarity graph between clusters based on centroid cosine similarity. Returns list of edges: [{source, target, weight}, ...]. """ from sklearn.metrics.pairwise import cosine_similarity unique_clusters = sorted([c for c in set(cluster_labels) if c != -1]) if len(unique_clusters) < 2: return [] # Compute centroids centroids = [] for cid in unique_clusters: mask = cluster_labels == cid centroids.append(embeddings[mask].mean(axis=0)) centroids = np.array(centroids) # Pairwise similarity sim_matrix = cosine_similarity(centroids) edges = [] for i in range(len(unique_clusters)): for j in range(i + 1, len(unique_clusters)): edges.append({ "source": int(unique_clusters[i]), "target": int(unique_clusters[j]), "weight": round(float(sim_matrix[i][j]), 4), }) # Keep top N edges by weight edges.sort(key=lambda e: e["weight"], reverse=True) return edges[:top_n_edges] def merge_similar_clusters( topics_df: pd.DataFrame, cluster_labels: np.ndarray, df: pd.DataFrame, threshold: float = 0.75, groq_api_key: Optional[str] = None, ) -> tuple: """ Merge clusters with highly similar keyword sets. Returns (new_topics_df, new_labels). """ from sklearn.feature_extraction.text import TfidfVectorizer from sklearn.metrics.pairwise import cosine_similarity if len(topics_df) < 2: return topics_df, cluster_labels # Compute pairwise similarity of keyword strings vectorizer = TfidfVectorizer(stop_words="english") kw_matrix = vectorizer.fit_transform(topics_df["keywords"].tolist()) sim_matrix = cosine_similarity(kw_matrix) # Find pairs to merge merge_map = {} topic_ids = topics_df["topic_id"].tolist() for i in range(len(topic_ids)): for j in range(i + 1, len(topic_ids)): if sim_matrix[i][j] >= threshold: src = topic_ids[j] tgt = topic_ids[i] # Follow existing merge chains while tgt in merge_map: tgt = merge_map[tgt] merge_map[src] = tgt if not merge_map: return topics_df, cluster_labels logger.info("Merging %d cluster pairs (threshold=%.2f)", len(merge_map), threshold) # Remap labels new_labels = np.array(cluster_labels, copy=True) for old_id, new_id in merge_map.items(): new_labels[new_labels == old_id] = new_id # Rebuild topics_df for remaining clusters remaining_ids = set(topic_ids) - set(merge_map.keys()) new_topics_df = topics_df[topics_df["topic_id"].isin(remaining_ids)].copy() new_topics_df = new_topics_df.reset_index(drop=True) return new_topics_df, new_labels def extract_subtopics_within_clusters( df: pd.DataFrame, cluster_labels: np.ndarray, min_topics: int = 100, n_keywords: int = 10, ) -> pd.DataFrame: """ Extract fine-grained sub-topics within each cluster using NMF. Delegates to the implementation in run_specter.py. """ from run_specter import extract_subtopics_within_clusters as _extract return _extract(df, cluster_labels, min_topics, n_keywords)