Spaces:
Running
Running
| """ | |
| agent.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. | |
| 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, Optional | |
| import pandas as pd | |
| from tools import ( | |
| load_csv, | |
| preprocess_text, | |
| run_topic_modeling, | |
| generate_labels, | |
| compare_themes, | |
| create_taxonomy_map, | |
| ) | |
| # --------------------------------------------------------------------------- | |
| # 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) | |
| # 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. | |
| 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() | |
| # ----------------------------------------------------------------- | |
| # Public interface | |
| # ----------------------------------------------------------------- | |
| def run(self) -> AgentResult: | |
| """ | |
| Execute the full pipeline step by step. | |
| Returns | |
| ------- | |
| AgentResult | |
| Structured results including all DataFrames, taxonomy, and file paths. | |
| """ | |
| logger.info("=" * 60) | |
| logger.info("TopicAgent — Starting pipeline") | |
| logger.info("=" * 60) | |
| try: | |
| # Step 1: Load CSV | |
| self._step_load_csv() | |
| if self.use_specter: | |
| self._step_specter_pipeline() | |
| else: | |
| # Step 2: Preprocess text | |
| self._step_preprocess() | |
| # Step 3: Topic modeling on Titles | |
| self._step_model_titles() | |
| # Step 4: Topic modeling on Abstracts | |
| self._step_model_abstracts() | |
| # Step 5: Generate labels | |
| self._step_generate_labels() | |
| # Step 6: Build combined topics table | |
| self._step_combine_topics() | |
| # Step 7: Compare themes | |
| self._step_compare_themes() | |
| # Step 8: Create taxonomy map | |
| self._step_taxonomy_map() | |
| # Step 9: Export outputs (TCCM not used if specter, but export is) | |
| if not self.use_specter: | |
| self._step_tccm() | |
| 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 | |
| # ----------------------------------------------------------------- | |
| # Pipeline steps | |
| # ----------------------------------------------------------------- | |
| def _step_load_csv(self): | |
| """Step 1 — Ingest CSV dataset.""" | |
| logger.info("Step 1/9: Loading CSV …") | |
| self._df = load_csv(self.csv_path) | |
| self._result.steps_completed.append("load_csv") | |
| logger.info(" → %d papers loaded.", len(self._df)) | |
| def _step_preprocess(self): | |
| """Step 2 — Preprocess Title and Abstract text.""" | |
| logger.info("Step 2/9: 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") | |
| logger.info(" → Titles preprocessed: %d docs", len(self._titles_clean)) | |
| logger.info(" → Abstracts preprocessed: %d docs", len(self._abstracts_clean)) | |
| def _step_model_titles(self): | |
| """Step 3 — Topic modeling on Titles (clusters + sub-topics).""" | |
| logger.info("Step 3/9: 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, | |
| max_clusters=self.max_clusters, | |
| ) | |
| self._result.steps_completed.append("topic_modeling_titles") | |
| n_clusters = self._title_topics_df["cluster_id"].nunique() if "cluster_id" in self._title_topics_df.columns else len(self._title_topics_df) | |
| logger.info(" → %d title topics in %d clusters.", len(self._title_topics_df), n_clusters) | |
| def _step_model_abstracts(self): | |
| """Step 4 — Topic modeling on Abstracts (clusters + sub-topics).""" | |
| logger.info("Step 4/9: 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, | |
| max_clusters=self.max_clusters, | |
| ) | |
| self._result.steps_completed.append("topic_modeling_abstracts") | |
| n_clusters = self._abstract_topics_df["cluster_id"].nunique() if "cluster_id" in self._abstract_topics_df.columns else len(self._abstract_topics_df) | |
| logger.info(" → %d abstract topics in %d clusters.", len(self._abstract_topics_df), n_clusters) | |
| def _step_generate_labels(self): | |
| """Step 5 — Generate human-readable labels.""" | |
| logger.info("Step 5/9: 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") | |
| logger.info(" → Labels generated for all topics.") | |
| def _step_combine_topics(self): | |
| """Step 6 — Combine title and abstract topics into one table.""" | |
| logger.info("Step 6/9: 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") | |
| logger.info(" → Combined table: %d topics total.", len(combined)) | |
| def _step_compare_themes(self): | |
| """Step 7 — Compare title vs abstract themes.""" | |
| logger.info("Step 7/9: 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") | |
| logger.info(" → Comparison table: %d rows.", len(comparison)) | |
| def _step_taxonomy_map(self): | |
| """Step 8 — Create taxonomy map (MAPPED / NOVEL).""" | |
| logger.info("Step 8/9: Building taxonomy map …") | |
| # Use the combined topics for taxonomy | |
| taxonomy = create_taxonomy_map(self._result.combined_topics) | |
| self._result.taxonomy_map = taxonomy | |
| self._result.steps_completed.append("create_taxonomy_map") | |
| logger.info( | |
| " → MAPPED: %d, NOVEL: %d", | |
| taxonomy["metadata"]["mapped_count"], | |
| taxonomy["metadata"]["novel_count"], | |
| ) | |
| def _step_export(self): | |
| """Step 9 — Export all outputs to disk.""" | |
| logger.info("Step 9/9: 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 | |
| logger.info(" → Saved: %s", 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 | |
| logger.info(" → Saved: %s", 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 | |
| logger.info(" → Saved: %s", 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 | |
| logger.info(" → Saved: %s", 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 | |
| logger.info(" → Saved: %s", abstract_path) | |
| self._result.steps_completed.append("export") | |
| logger.info(" → All outputs exported successfully.") | |
| def _step_tccm(self): | |
| logger.info("Step X: TCCM Clustering phase ...") | |
| # tccm implementation | |
| self._result.steps_completed.append("tccm_extraction") | |
| def _step_specter_pipeline(self): | |
| """Run SPECTER2 Pipeline with two-level hierarchy (clusters + sub-topics).""" | |
| logger.info("Running SPECTER2 Pipeline...") | |
| from run_specter import ( | |
| generate_specter2_embeddings, store_in_chroma, | |
| cluster_with_hdbscan, c_tf_idf, | |
| extract_top_n_words_per_topic, get_council_labels, | |
| extract_subtopics_within_clusters, | |
| ) | |
| import pandas as pd | |
| # 1. Combine Abstract + Title | |
| papers = self._df[['Title', 'Abstract']].to_dict('records') | |
| titles = [p['Title'] or "" for p in papers] | |
| abstracts = [p['Abstract'] or "" for p in papers] | |
| combined = [t + " " + a for t, a in zip(titles, abstracts)] | |
| # 2. Embed | |
| embeddings = generate_specter2_embeddings( | |
| [{'title': t, 'abstract': a} for t, a in zip(titles, abstracts)] | |
| ) | |
| # 3. Store | |
| store_in_chroma( | |
| embeddings, documents=combined, | |
| metadatas=[{'id': i} for i in range(len(combined))], | |
| ) | |
| # 4. Cluster (produces high-level groups) | |
| labels, _ = cluster_with_hdbscan(embeddings) | |
| self._df['cluster_id'] = labels | |
| # 5. Extract cluster-level keywords | |
| docs_per_topic = self._df[self._df['cluster_id'] != -1].groupby( | |
| ['cluster_id'], as_index=False | |
| ).agg({'Title': ' '.join}) | |
| docs_per_topic = docs_per_topic.rename(columns={'cluster_id': 'topic_id'}) | |
| tf_idf, count = c_tf_idf(docs_per_topic['Title'].values, m=len(combined)) | |
| top_n_words = extract_top_n_words_per_topic(tf_idf, count, docs_per_topic, n=10) | |
| # Build cluster-level summary | |
| cluster_rows = [] | |
| for cluster_id, words in top_n_words.items(): | |
| if cluster_id != -1: | |
| kw = ", ".join([w for w, _ in words]) | |
| cluster_rows.append({"topic_id": int(cluster_id), "keywords": kw}) | |
| self._clusters_df = pd.DataFrame(cluster_rows) | |
| n_clusters = len(self._clusters_df) | |
| logger.info(" → %d clusters formed", n_clusters) | |
| # 6. Extract sub-topics within clusters (≥100 topics) | |
| cluster_labels_arr = self._df['cluster_id'].values | |
| self._combined_topics_df = extract_subtopics_within_clusters( | |
| self._df, cluster_labels_arr, | |
| min_topics=self.min_topics, n_keywords=10, | |
| ) | |
| # Add cluster labels | |
| cluster_label_map = {} | |
| for _, row in self._clusters_df.iterrows(): | |
| cid = int(row['topic_id']) | |
| kws = [kw.strip().title() for kw in row['keywords'].split(',')][:3] | |
| cluster_label_map[cid] = ' / '.join(kws) | |
| self._combined_topics_df['cluster_label'] = ( | |
| self._combined_topics_df['cluster_id'].map(cluster_label_map).fillna('') | |
| ) | |
| logger.info( | |
| " → %d topics extracted from %d clusters", | |
| len(self._combined_topics_df), n_clusters, | |
| ) | |
| # 7. Label topics | |
| if self.use_llm_labels and self.groq_api_key: | |
| labels_list = get_council_labels(self._combined_topics_df, self.groq_api_key) | |
| self._combined_topics_df['label'] = labels_list | |
| else: | |
| self._combined_topics_df['label'] = [ | |
| ", ".join(row['keywords'].split(",")[:3]).title() | |
| for _, row in self._combined_topics_df.iterrows() | |
| ] | |
| self._result.combined_topics = self._combined_topics_df | |