Spaces:
Sleeping
Sleeping
| 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 | |