import os import json import logging import torch import numpy as np import pandas as pd from typing import List, Dict, Any from transformers import AutoTokenizer from adapters import AutoAdapterModel from umap import UMAP from hdbscan import HDBSCAN from sklearn.feature_extraction.text import CountVectorizer logger = logging.getLogger(__name__) # Basic c-TF-IDF implementation def c_tf_idf(documents, m, ngram_range=(1, 1)): 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): 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 generate_specter2_embeddings(papers: List[Dict[str, str]], batch_size=16) -> np.ndarray: """ Generate embeddings using SPECTER2 via adapters. papers should be a list of dicts with 'title' and 'abstract' keys. """ 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(f"Generating embeddings for {len(papers)} papers using {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) # Take the CLS token batch_embeddings = output.last_hidden_state[:, 0, :].cpu().numpy() embeddings.append(batch_embeddings) return np.vstack(embeddings) import chromadb from groq import Groq import time import re def store_in_chroma(embeddings, documents, metadatas, collection_name="specter_vectors"): client = chromadb.PersistentClient(path="./chroma_db") collection = client.get_or_create_collection(name=collection_name) ids = [str(i) for i in range(len(documents))] # Store in batches if necessary collection.add( embeddings=embeddings.tolist(), documents=documents, metadatas=metadatas, ids=ids ) return collection def cluster_with_hdbscan(embeddings, min_cluster_size=15, min_samples=5): logger.info(f"Reducing dimensionality using UMAP for {embeddings.shape[0]} embeddings...") umap_model = UMAP(n_neighbors=15, n_components=5, min_dist=0.0, metric='cosine', random_state=42) reduced_embeddings = umap_model.fit_transform(embeddings) logger.info("Clustering with HDBSCAN...") clusterer = HDBSCAN(min_cluster_size=min_cluster_size, min_samples=min_samples, metric='euclidean', cluster_selection_method='eom') labels = clusterer.fit_predict(reduced_embeddings) return labels, reduced_embeddings 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. For each cluster, runs NMF on the cluster's TF-IDF matrix to discover sub-topics. The number of sub-topics per cluster is proportional to cluster size, with a global target of at least `min_topics` total. Parameters ---------- df : pd.DataFrame DataFrame with 'Title' and 'Abstract' columns. cluster_labels : np.ndarray Cluster assignment for each document (from HDBSCAN). min_topics : int Minimum total number of topics across all clusters (default 100). n_keywords : int Number of top keywords to extract per sub-topic. Returns ------- pd.DataFrame Columns: cluster_id, topic_id (global), sub_topic_id (within-cluster), keywords, n_documents, source """ from sklearn.decomposition import NMF from sklearn.feature_extraction.text import TfidfVectorizer df = df.copy() df["cluster_id"] = cluster_labels # Get valid clusters (exclude noise = -1) valid_mask = df["cluster_id"] != -1 valid_df = df[valid_mask].copy() cluster_ids = sorted(valid_df["cluster_id"].unique()) n_clusters = len(cluster_ids) if n_clusters == 0: raise ValueError("No valid clusters found (all documents are noise).") # Calculate sub-topic count per cluster proportional to size cluster_sizes = {cid: int((valid_df["cluster_id"] == cid).sum()) for cid in cluster_ids} total_docs = sum(cluster_sizes.values()) subtopic_counts = {} for cid in cluster_ids: raw_share = (cluster_sizes[cid] / total_docs) * min_topics subtopic_counts[cid] = max(2, round(raw_share)) # Ensure total >= min_topics by scaling up if needed current_total = sum(subtopic_counts.values()) if current_total < min_topics: deficit = min_topics - current_total # Distribute deficit proportionally to larger clusters sorted_clusters = sorted(cluster_ids, key=lambda c: cluster_sizes[c], reverse=True) i = 0 while deficit > 0: cid = sorted_clusters[i % len(sorted_clusters)] # Don't exceed cluster size (each sub-topic needs ≥2 docs) max_possible = cluster_sizes[cid] // 2 if subtopic_counts[cid] < max_possible: subtopic_counts[cid] += 1 deficit -= 1 i += 1 if i >= len(sorted_clusters) * 10: # Safety valve break logger.info( "Sub-topic allocation: %d clusters → %d total sub-topics (target ≥ %d)", n_clusters, sum(subtopic_counts.values()), min_topics, ) # Extract sub-topics per cluster all_topics = [] global_topic_id = 0 for cid in cluster_ids: cluster_docs = valid_df[valid_df["cluster_id"] == cid] texts = (cluster_docs["Title"].fillna("") + " " + cluster_docs["Abstract"].fillna("")).tolist() n_subtopics = subtopic_counts[cid] # Need at least as many docs as sub-topics n_subtopics = min(n_subtopics, len(texts)) 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(texts) feature_names = vectorizer.get_feature_names_out() if n_subtopics == 1 or tfidf_matrix.shape[0] < 2: # Single sub-topic — just extract top keywords scores = tfidf_matrix.sum(axis=0).A1 top_indices = scores.argsort()[-n_keywords:][::-1] kw = ", ".join([feature_names[i] for i in top_indices]) all_topics.append({ "cluster_id": int(cid), "topic_id": global_topic_id, "sub_topic_id": 0, "keywords": kw, "n_documents": len(texts), "source": "SPECTER2+HDBSCAN", }) 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) # doc-topic matrix H = nmf.components_ # topic-term matrix # Assign each doc to its dominant sub-topic doc_assignments = W.argmax(axis=1) for sub_id in range(n_subtopics): top_indices = H[sub_id].argsort()[-n_keywords:][::-1] kw = ", ".join([feature_names[i] for i in top_indices]) n_docs_in_subtopic = int((doc_assignments == sub_id).sum()) if n_docs_in_subtopic == 0: n_docs_in_subtopic = 1 # NMF component exists even if no dominant doc all_topics.append({ "cluster_id": int(cid), "topic_id": global_topic_id, "sub_topic_id": sub_id, "keywords": kw, "n_documents": n_docs_in_subtopic, "source": "SPECTER2+HDBSCAN", }) global_topic_id += 1 except Exception as e: logger.warning("NMF failed for cluster %d (%s), using single topic", cid, e) # Fallback: single topic with combined keywords combined_text = " ".join(texts) words = combined_text.lower().split() word_freq = pd.Series(words).value_counts().head(n_keywords) kw = ", ".join(word_freq.index.tolist()) all_topics.append({ "cluster_id": int(cid), "topic_id": global_topic_id, "sub_topic_id": 0, "keywords": kw, "n_documents": len(texts), "source": "SPECTER2+HDBSCAN", }) global_topic_id += 1 topics_df = pd.DataFrame(all_topics) logger.info( "Extracted %d topics from %d clusters (avg %.1f topics/cluster)", len(topics_df), n_clusters, len(topics_df) / max(n_clusters, 1), ) return topics_df def get_council_labels(topics_df, groq_api_key): """ Council: Domain Expert, Data Analyst, Judge topics_df: list of dicts with topic_id and keywords """ client = Groq(api_key=groq_api_key) results = [] for _, row in topics_df.iterrows(): topic_id = row['topic_id'] keywords = row['keywords'] # Persona 1: Domain Expert prompt_expert = f"You are a scientific Domain Expert. Look at these keywords: [{keywords}]. Propose a 3-5 word descriptive label for this academic topic." # Persona 2: Data Analyst prompt_analyst = f"You are a Data Analyst categorizing clusters. Keywords: [{keywords}]. Give a short, precise categorical 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() 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() # Persona 3: Judge prompt_judge = f"Keywords: [{keywords}].\nExpert proposed: '{res_expert}'.\nAnalyst proposed: '{res_analyst}'.\nSynthesize these into a single, high-quality final label (max 5 words). Return ONLY the label, nothing else." res_judge = client.chat.completions.create(model="llama-3.3-70b-versatile", messages=[{"role": "user", "content": prompt_judge}], temperature=0.2, max_tokens=20).choices[0].message.content.strip() # Clean up res_judge = res_judge.replace('"', '').replace("'", "").strip() results.append(res_judge) except Exception as e: logger.error(f"Groq error: {e}") results.append(", ".join(keywords.split(",")[:3]).title()) time.sleep(0.5) return results