Spaces:
Running
Running
| from tools import merge_similar_clusters, extract_subtopics_within_clusters | |
| """ | |
| agents.py — TopicAgent orchestrates the end-to-end topic modeling workflow. | |
| This module defines the TopicAgent class, which: | |
| 1. Loads and validates the CSV dataset. | |
| 2. Preprocesses text for Titles and Abstracts separately. | |
| 3. Runs topic modeling on each corpus (≥100 topics guaranteed). | |
| 4. Generates human-readable labels for every topic. | |
| 5. Compares dominant themes across Title and Abstract topics. | |
| 6. Produces a taxonomy map (MAPPED / NOVEL classification). | |
| 7. Exports structured outputs: topics table, comparison CSV, taxonomy JSON. | |
| OR (SPECTER mode): | |
| 1. Loads CSV | |
| 2. Generates SPECTER2 embeddings + stores in ChromaDB | |
| 3. Adaptive HDBSCAN clustering (5-140 clusters) | |
| 4. c-TF-IDF keyword extraction + representative sentences | |
| 5. AI Council labeling with debate transcripts | |
| 6. Taxonomy mapping | |
| 7. Export | |
| Usage: | |
| agent = TopicAgent(csv_path="dataset.csv") | |
| results = agent.run() | |
| """ | |
| import os | |
| import json | |
| import logging | |
| from dataclasses import dataclass, field | |
| from typing import Dict, Any, List, Optional | |
| import numpy as np | |
| import pandas as pd | |
| from tools import ( | |
| load_csv, | |
| preprocess_text, | |
| run_topic_modeling, | |
| generate_labels, | |
| compare_themes, | |
| create_taxonomy_map, | |
| generate_specter2_embeddings, | |
| store_in_chroma, | |
| load_embeddings_from_chroma, | |
| adaptive_hdbscan_clustering, | |
| c_tf_idf, | |
| extract_top_n_words_per_topic, | |
| extract_representative_sentences, | |
| get_council_debate, | |
| build_cluster_graph, | |
| ) | |
| # --------------------------------------------------------------------------- | |
| # Logging | |
| # --------------------------------------------------------------------------- | |
| logger = logging.getLogger(__name__) | |
| # --------------------------------------------------------------------------- | |
| # Structured result container | |
| # --------------------------------------------------------------------------- | |
| class AgentResult: | |
| """Container for all outputs produced by the TopicAgent.""" | |
| # Core dataframes | |
| title_topics: pd.DataFrame = field(default_factory=pd.DataFrame) | |
| abstract_topics: pd.DataFrame = field(default_factory=pd.DataFrame) | |
| combined_topics: pd.DataFrame = field(default_factory=pd.DataFrame) | |
| comparison: pd.DataFrame = field(default_factory=pd.DataFrame) | |
| # Taxonomy map (dict serialisable to JSON) | |
| taxonomy_map: Dict[str, Any] = field(default_factory=dict) | |
| # SPECTER/Cluster-specific data | |
| cluster_data: Dict[str, Any] = field(default_factory=dict) | |
| # Holds: embeddings_2d, cluster_labels, cluster_details, representative_sentences, | |
| # cluster_graph_edges, n_clusters, hdbscan_params | |
| council_debates: Dict[int, Dict[str, str]] = field(default_factory=dict) | |
| # {topic_id: {expert, analyst, anti_hallucination, final_label, ...}} | |
| # Pipeline mode: "default" or "specter" | |
| pipeline_mode: str = "default" | |
| # Execution metadata | |
| status: str = "pending" | |
| steps_completed: list = field(default_factory=list) | |
| errors: list = field(default_factory=list) | |
| # File paths of exported artefacts | |
| exported_files: Dict[str, str] = field(default_factory=dict) | |
| # --------------------------------------------------------------------------- | |
| # TopicAgent | |
| # --------------------------------------------------------------------------- | |
| class TopicAgent: | |
| """ | |
| Orchestrates the research-paper topic modeling pipeline. | |
| Parameters | |
| ---------- | |
| csv_path : str | |
| Path to the input CSV file. | |
| output_dir : str | |
| Directory to write output files. | |
| min_topics : int | |
| Minimum number of topics to generate per source (default 100). | |
| use_llm_labels : bool | |
| Whether to use Groq LLM for label generation. | |
| use_specter : bool | |
| Whether to use SPECTER2 + HDBSCAN clustering pipeline. | |
| groq_api_key : str, optional | |
| API key for Groq (used only when use_llm_labels is True). | |
| """ | |
| def __init__( | |
| self, | |
| csv_path: str, | |
| output_dir: str = "outputs", | |
| min_topics: int = 100, | |
| max_clusters: int = 25, | |
| use_llm_labels: bool = False, | |
| use_specter: bool = False, | |
| groq_api_key: Optional[str] = None, | |
| ): | |
| self.csv_path = csv_path | |
| self.output_dir = output_dir | |
| self.min_topics = min_topics | |
| self.max_clusters = max_clusters | |
| self.use_llm_labels = use_llm_labels | |
| self.use_specter = use_specter | |
| self.groq_api_key = groq_api_key | |
| # Ensure output directory exists | |
| os.makedirs(self.output_dir, exist_ok=True) | |
| self._result = AgentResult() | |
| if use_specter: | |
| self._result.pipeline_mode = "specter" | |
| # Internal state | |
| self._df = None | |
| self._embeddings = None | |
| self._cluster_labels = None | |
| self._reduced_2d = None | |
| # ----------------------------------------------------------------- | |
| # Public interface | |
| # ----------------------------------------------------------------- | |
| def run(self) -> AgentResult: | |
| """Execute the full pipeline. Use sub-steps for progress display.""" | |
| logger.info("=" * 60) | |
| logger.info("TopicAgent — Starting pipeline (mode=%s)", self._result.pipeline_mode) | |
| logger.info("=" * 60) | |
| try: | |
| self._step_load_csv() | |
| if self.use_specter: | |
| self._step_specter_embed() | |
| self._step_specter_cluster() | |
| self._step_specter_extract() | |
| self._step_merge_similar_clusters() | |
| self._step_specter_subtopics() | |
| self._step_specter_label() | |
| self._step_specter_taxonomy() | |
| else: | |
| self._step_preprocess() | |
| self._step_model_titles() | |
| self._step_model_abstracts() | |
| self._step_generate_labels() | |
| self._step_combine_topics() | |
| self._step_compare_themes() | |
| self._step_taxonomy_map() | |
| self._step_export() | |
| self._result.status = "success" | |
| logger.info("Pipeline completed successfully.") | |
| except Exception as exc: | |
| self._result.status = "failed" | |
| self._result.errors.append(str(exc)) | |
| logger.error("Pipeline failed: %s", exc, exc_info=True) | |
| return self._result | |
| # ----------------------------------------------------------------- | |
| # Shared steps | |
| # ----------------------------------------------------------------- | |
| def _step_load_csv(self): | |
| """Step 1 — Ingest CSV dataset.""" | |
| logger.info("Step 1: Loading CSV …") | |
| self._df = load_csv(self.csv_path) | |
| self._result.steps_completed.append("load_csv") | |
| logger.info(" → %d papers loaded.", len(self._df)) | |
| # ----------------------------------------------------------------- | |
| # DEFAULT pipeline steps | |
| # ----------------------------------------------------------------- | |
| def _step_preprocess(self): | |
| """Step 2 — Preprocess Title and Abstract text.""" | |
| logger.info("Step 2: Preprocessing text …") | |
| self._titles_clean = preprocess_text(self._df["Title"].tolist()) | |
| self._abstracts_clean = preprocess_text(self._df["Abstract"].tolist()) | |
| self._result.steps_completed.append("preprocess_text") | |
| def _step_model_titles(self): | |
| """Step 3 — Topic modeling on Titles.""" | |
| logger.info("Step 3: Topic modeling on Titles …") | |
| self._title_topics_df, self._title_model = run_topic_modeling( | |
| self._titles_clean, source_label="Titles", min_topics=self.min_topics, | |
| ) | |
| self._result.steps_completed.append("topic_modeling_titles") | |
| def _step_model_abstracts(self): | |
| """Step 4 — Topic modeling on Abstracts.""" | |
| logger.info("Step 4: Topic modeling on Abstracts …") | |
| self._abstract_topics_df, self._abstract_model = run_topic_modeling( | |
| self._abstracts_clean, source_label="Abstracts", min_topics=self.min_topics, | |
| ) | |
| self._result.steps_completed.append("topic_modeling_abstracts") | |
| def _step_generate_labels(self): | |
| """Step 5 — Generate human-readable labels.""" | |
| logger.info("Step 5: Generating topic labels …") | |
| self._title_topics_df = generate_labels( | |
| self._title_topics_df, use_llm=self.use_llm_labels, groq_api_key=self.groq_api_key, | |
| ) | |
| self._abstract_topics_df = generate_labels( | |
| self._abstract_topics_df, use_llm=self.use_llm_labels, groq_api_key=self.groq_api_key, | |
| ) | |
| self._result.title_topics = self._title_topics_df | |
| self._result.abstract_topics = self._abstract_topics_df | |
| self._result.steps_completed.append("generate_labels") | |
| def _step_combine_topics(self): | |
| """Step 6 — Combine title and abstract topics.""" | |
| logger.info("Step 6: Building combined topics table …") | |
| combined = pd.concat([self._title_topics_df, self._abstract_topics_df], ignore_index=True) | |
| combined["global_id"] = range(len(combined)) | |
| self._result.combined_topics = combined | |
| self._result.steps_completed.append("combine_topics") | |
| def _step_compare_themes(self): | |
| """Step 7 — Compare title vs abstract themes.""" | |
| logger.info("Step 7: Comparing title vs abstract themes …") | |
| comparison = compare_themes(self._title_topics_df, self._abstract_topics_df) | |
| self._result.comparison = comparison | |
| self._result.steps_completed.append("compare_themes") | |
| def _step_taxonomy_map(self): | |
| """Step 8 — Create taxonomy map.""" | |
| logger.info("Step 8: Building taxonomy map …") | |
| taxonomy = create_taxonomy_map(self._result.combined_topics) | |
| self._result.taxonomy_map = taxonomy | |
| self._result.steps_completed.append("create_taxonomy_map") | |
| # ----------------------------------------------------------------- | |
| # SPECTER pipeline steps (discrete, callable from app.py) | |
| # ----------------------------------------------------------------- | |
| def _step_specter_embed(self): | |
| """SPECTER Step 2 — Generate embeddings and store in ChromaDB.""" | |
| logger.info(">>> Clustering pipeline started") | |
| logger.info("SPECTER Step 2: Generating SPECTER2 embeddings …") | |
| titles = [str(t) if t else "" for t in self._df["Title"].tolist()] | |
| abstracts = [str(a) if a else "" for a in self._df["Abstract"].tolist()] | |
| combined_docs = [t + " " + a for t, a in zip(titles, abstracts)] | |
| # Check if embeddings already cached in ChromaDB with correct count | |
| try: | |
| cached_emb, cached_docs, _ = load_embeddings_from_chroma() | |
| if len(cached_emb) == len(self._df): | |
| logger.info("Loaded embeddings from vector DB — %d vectors", len(cached_emb)) | |
| self._embeddings = cached_emb | |
| self._combined_docs = combined_docs | |
| self._result.steps_completed.append("specter_embed_cached") | |
| return | |
| except (ValueError, Exception) as e: | |
| logger.info("No cached embeddings found (%s), generating fresh …", e) | |
| # Generate fresh embeddings | |
| papers = [{"title": t, "abstract": a} for t, a in zip(titles, abstracts)] | |
| self._embeddings = generate_specter2_embeddings(papers) | |
| self._combined_docs = combined_docs | |
| # Store in ChromaDB | |
| metadatas = [{"id": i, "title": titles[i][:200]} for i in range(len(combined_docs))] | |
| store_in_chroma(self._embeddings, documents=combined_docs, metadatas=metadatas) | |
| logger.info("Loaded embeddings from vector DB — %d vectors", len(self._embeddings)) | |
| # Validate — no nulls | |
| null_count = int(np.isnan(self._embeddings).any(axis=1).sum()) | |
| if null_count > 0: | |
| logger.warning("Found %d embeddings with NaN values, replacing with zeros", null_count) | |
| self._embeddings = np.nan_to_num(self._embeddings, nan=0.0) | |
| self._result.steps_completed.append("specter_embed") | |
| def _step_specter_cluster(self): | |
| """SPECTER Step 3 — Adaptive HDBSCAN clustering (10–40 clusters).""" | |
| logger.info("Running HDBSCAN (targeting %d–%d clusters)...", 10, self.max_clusters) | |
| labels, reduced_2d, n_clusters, params = adaptive_hdbscan_clustering( | |
| self._embeddings, min_clusters=10, max_clusters=self.max_clusters, | |
| ) | |
| self._cluster_labels = np.array(labels) | |
| self._reduced_2d = reduced_2d | |
| self._df["cluster_id"] = labels | |
| logger.info("Clusters generated: %d", n_clusters) | |
| # Safeguard: if < 5 clusters, raise error | |
| if n_clusters < 5: | |
| raise RuntimeError( | |
| f"Clustering produced only {n_clusters} clusters (minimum 5 required). " | |
| "Try adjusting parameters or adding more data." | |
| ) | |
| self._result.cluster_data["n_clusters"] = n_clusters | |
| self._result.cluster_data["hdbscan_params"] = params | |
| self._result.cluster_data["embeddings_2d"] = reduced_2d.tolist() | |
| self._result.cluster_data["cluster_labels"] = labels.tolist() if hasattr(labels, 'tolist') else list(labels) | |
| self._result.steps_completed.append("specter_cluster") | |
| def _step_specter_extract(self): | |
| """SPECTER Step 4 — Extract cluster-level keywords and representative sentences.""" | |
| logger.info("SPECTER Step 4: Extracting cluster keywords and representative sentences …") | |
| # c-TF-IDF keyword extraction at CLUSTER level | |
| df_with_clusters = self._df.copy() | |
| df_with_clusters["cluster_id"] = self._cluster_labels | |
| # Group documents by cluster for c-TF-IDF | |
| docs_per_topic = df_with_clusters.groupby("cluster_id", as_index=False).agg( | |
| {"Title": " ".join, "Abstract": " ".join} | |
| ) | |
| docs_per_topic = docs_per_topic[docs_per_topic["cluster_id"] != -1] | |
| docs_per_topic = docs_per_topic.rename(columns={"cluster_id": "topic_id"}) | |
| # Use combined Title + Abstract for better keyword extraction | |
| combined_text = (docs_per_topic["Title"] + " " + docs_per_topic["Abstract"]).values | |
| tf_idf, count = c_tf_idf(combined_text, m=len(self._df)) | |
| top_n_words = extract_top_n_words_per_topic(tf_idf, count, docs_per_topic, n=10) | |
| # Build CLUSTER-level DataFrame (these are the high-level groups) | |
| cluster_rows = [] | |
| for cluster_id, words in top_n_words.items(): | |
| if cluster_id == -1: | |
| continue | |
| kw = ", ".join([w for w, _ in words]) | |
| n_docs = int((self._cluster_labels == cluster_id).sum()) | |
| cluster_rows.append({ | |
| "topic_id": int(cluster_id), | |
| "keywords": kw, | |
| "source": "SPECTER2+HDBSCAN", | |
| "n_documents": n_docs, | |
| }) | |
| self._clusters_df = pd.DataFrame(cluster_rows) | |
| logger.info(" → %d clusters with keywords extracted", len(self._clusters_df)) | |
| # Extract representative sentences | |
| rep_sentences = extract_representative_sentences( | |
| self._df, self._embeddings, self._cluster_labels, top_n=5, | |
| ) | |
| self._result.cluster_data["representative_sentences"] = rep_sentences | |
| # Build cluster graph | |
| edges = build_cluster_graph(self._clusters_df, self._embeddings, self._cluster_labels) | |
| self._result.cluster_data["graph_edges"] = edges | |
| # Store per-cluster details | |
| cluster_details = {} | |
| for _, row in self._clusters_df.iterrows(): | |
| cid = int(row["topic_id"]) | |
| cluster_details[cid] = { | |
| "cluster_id": cid, | |
| "keywords": row["keywords"], | |
| "n_documents": int(row["n_documents"]), | |
| "representatives": rep_sentences.get(cid, []), | |
| } | |
| self._result.cluster_data["cluster_details"] = cluster_details | |
| self._result.steps_completed.append("specter_extract") | |
| def _step_specter_subtopics(self): | |
| """SPECTER Step 5 — Extract ≥100 fine-grained sub-topics within each cluster.""" | |
| logger.info( | |
| "SPECTER Step 5: Extracting sub-topics (target ≥ %d topics from %d clusters) …", | |
| self.min_topics, len(self._clusters_df), | |
| ) | |
| self._topics_df = extract_subtopics_within_clusters( | |
| self._df, | |
| self._cluster_labels, | |
| min_topics=self.min_topics, | |
| n_keywords=10, | |
| ) | |
| # Add cluster_label from the clusters_df | |
| cluster_label_map = {} | |
| for _, row in self._clusters_df.iterrows(): | |
| cid = int(row["topic_id"]) | |
| # Heuristic label from top 3 keywords | |
| kws = [kw.strip().title() for kw in row["keywords"].split(",")][:3] | |
| cluster_label_map[cid] = " / ".join(kws) | |
| self._topics_df["cluster_label"] = self._topics_df["cluster_id"].map(cluster_label_map).fillna("") | |
| logger.info( | |
| " → %d topics extracted from %d clusters", | |
| len(self._topics_df), | |
| self._topics_df["cluster_id"].nunique(), | |
| ) | |
| self._result.steps_completed.append("specter_subtopics") | |
| def _step_specter_label(self): | |
| """SPECTER Step 6 — Label clusters and topics.""" | |
| logger.info("SPECTER Step 6: Labeling clusters and topics …") | |
| # --- Label CLUSTERS (high-level) --- | |
| if self.use_llm_labels and self.groq_api_key: | |
| logger.info("Running AI Council debate for %d clusters …", len(self._clusters_df)) | |
| cluster_labels_list = [] | |
| debates = {} | |
| for _, row in self._clusters_df.iterrows(): | |
| cid = int(row["topic_id"]) | |
| debate = get_council_debate(row["keywords"], self.groq_api_key) | |
| debates[cid] = debate | |
| cluster_labels_list.append(debate["final_label"]) | |
| logger.info(" Cluster %d → '%s'", cid, debate["final_label"]) | |
| self._clusters_df["label"] = cluster_labels_list | |
| self._result.council_debates = debates | |
| # Update cluster_label in topics_df with LLM labels | |
| cluster_label_map = dict(zip(self._clusters_df["topic_id"].astype(int), self._clusters_df["label"])) | |
| self._topics_df["cluster_label"] = self._topics_df["cluster_id"].map(cluster_label_map).fillna("") | |
| else: | |
| logger.info("Using heuristic labels (LLM not enabled)") | |
| self._clusters_df["label"] = [ | |
| ", ".join(row["keywords"].split(",")[:3]).strip().title() | |
| for _, row in self._clusters_df.iterrows() | |
| ] | |
| # --- Label TOPICS (fine-grained) --- | |
| if self.use_llm_labels and self.groq_api_key: | |
| logger.info("Generating LLM labels for %d topics …", len(self._topics_df)) | |
| from tools import _generate_labels_llm | |
| topic_labels_df = self._topics_df.copy() | |
| topic_labels_df = _generate_labels_llm(topic_labels_df, self.groq_api_key) | |
| self._topics_df["label"] = topic_labels_df["label"] | |
| else: | |
| self._topics_df["label"] = [ | |
| ", ".join(row["keywords"].split(",")[:3]).strip().title() | |
| for _, row in self._topics_df.iterrows() | |
| ] | |
| # Populate result — combined_topics now has the TWO-LEVEL hierarchy | |
| self._result.combined_topics = self._topics_df | |
| self._result.title_topics = self._clusters_df # Cluster-level view | |
| self._result.abstract_topics = pd.DataFrame() # Not applicable for SPECTER | |
| self._result.steps_completed.append("specter_label") | |
| def _step_merge_similar_clusters(self): | |
| """SPECTER Step 4b — Merge semantically similar clusters before sub-topic extraction.""" | |
| from tools import merge_similar_clusters | |
| if not hasattr(self, "_clusters_df") or self._clusters_df.empty: | |
| return | |
| logger.info("SPECTER Step 4b: Post-merging similar clusters...") | |
| new_clusters_df, new_labels = merge_similar_clusters( | |
| self._clusters_df, self._cluster_labels, self._df, | |
| threshold=0.75, groq_api_key=self.groq_api_key | |
| ) | |
| self._clusters_df = new_clusters_df | |
| self._cluster_labels = new_labels | |
| n_clusters = len(self._clusters_df) | |
| self._result.cluster_data["cluster_labels"] = new_labels.tolist() if hasattr(new_labels, 'tolist') else list(new_labels) | |
| self._result.cluster_data["n_clusters"] = n_clusters | |
| logger.info(" → After merging: %d clusters remain", n_clusters) | |
| self._result.steps_completed.append("specter_merge") | |
| def _step_specter_taxonomy(self): | |
| """SPECTER Step 6 — Build taxonomy map for clustered topics.""" | |
| logger.info("SPECTER Step 6: Building taxonomy map …") | |
| taxonomy = create_taxonomy_map(self._result.combined_topics) | |
| self._result.taxonomy_map = taxonomy | |
| self._result.steps_completed.append("specter_taxonomy") | |
| # ----------------------------------------------------------------- | |
| # Export (shared) | |
| # ----------------------------------------------------------------- | |
| def _step_export(self): | |
| """Export all outputs to disk.""" | |
| logger.info("Exporting outputs …") | |
| # (a) Combined topics table CSV | |
| topics_path = os.path.join(self.output_dir, "topics_table.csv") | |
| self._result.combined_topics.to_csv(topics_path, index=False) | |
| self._result.exported_files["topics_table"] = topics_path | |
| # (b) Comparison CSV | |
| comparison_path = os.path.join(self.output_dir, "comparison.csv") | |
| self._result.comparison.to_csv(comparison_path, index=False) | |
| self._result.exported_files["comparison"] = comparison_path | |
| # (c) Taxonomy map JSON | |
| taxonomy_path = os.path.join(self.output_dir, "taxonomy_map.json") | |
| with open(taxonomy_path, "w", encoding="utf-8") as f: | |
| json.dump(self._result.taxonomy_map, f, indent=2, ensure_ascii=False) | |
| self._result.exported_files["taxonomy_map"] = taxonomy_path | |
| # (d) Title topics CSV | |
| title_path = os.path.join(self.output_dir, "title_topics.csv") | |
| self._result.title_topics.to_csv(title_path, index=False) | |
| self._result.exported_files["title_topics"] = title_path | |
| # (e) Abstract topics CSV | |
| abstract_path = os.path.join(self.output_dir, "abstract_topics.csv") | |
| self._result.abstract_topics.to_csv(abstract_path, index=False) | |
| self._result.exported_files["abstract_topics"] = abstract_path | |
| # (f) Cluster data JSON (SPECTER mode) | |
| if self._result.pipeline_mode == "specter" and self._result.cluster_data: | |
| cluster_path = os.path.join(self.output_dir, "cluster_data.json") | |
| # Don't save embeddings_2d to JSON (too large), save separately | |
| cluster_export = { | |
| k: v for k, v in self._result.cluster_data.items() | |
| if k != "embeddings_2d" | |
| } | |
| with open(cluster_path, "w", encoding="utf-8") as f: | |
| json.dump(cluster_export, f, indent=2, ensure_ascii=False, default=str) | |
| self._result.exported_files["cluster_data"] = cluster_path | |
| # (g) Council debates JSON | |
| if self._result.council_debates: | |
| debates_path = os.path.join(self.output_dir, "council_debates.json") | |
| with open(debates_path, "w", encoding="utf-8") as f: | |
| json.dump(self._result.council_debates, f, indent=2, ensure_ascii=False, default=str) | |
| self._result.exported_files["council_debates"] = debates_path | |
| self._result.steps_completed.append("export") | |
| logger.info(" → All outputs exported successfully.") | |