Spaces:
Sleeping
Sleeping
| import json | |
| from pathlib import Path | |
| from typing import Any, Dict, List, Optional, Union | |
| from datetime import datetime, timezone | |
| import pandas as pd | |
| from bertopic._bertopic import BERTopic | |
| from kbdebugger.utils.json import write_json | |
| from kbdebugger.utils.time import now_utc_compact | |
| from kbdebugger.compat.langchain import Document | |
| def save_topic_modeling_results( | |
| *, | |
| topic_model: BERTopic, | |
| documents: List[str], | |
| document_chunks: Optional[List[Document]] = None, | |
| keyword: str, | |
| matched_topic_ids: List[int], | |
| match_type_by_topic: Dict[int, str], | |
| matched_synonyms: Dict[int, set[str]], | |
| generated_synonyms: Optional[List[str]] = None, | |
| output_dir: Union[str, Path] = "logs", | |
| ) -> None: | |
| """ | |
| Save full BERTopic modeling output and keyword match metadata to JSON for inspection. | |
| Parameters | |
| ---------- | |
| topic_model: BERTopic | |
| The trained topic model. | |
| documents: list of str | |
| The raw input texts fed into the model. | |
| i.e., the output of Docling before topic modeling. | |
| document_chunks: list of LangChain Documents, optional | |
| If available, include their metadata in the output. | |
| keyword: str | |
| The user-specified keyword for topic matching. | |
| matched_topic_ids: list of int | |
| Topics selected for retention. | |
| match_type_by_topic: dict | |
| Topic ID -> "exact" or "synonym" | |
| matched_synonyms: dict | |
| Topic ID -> which synonym matched. | |
| generated_synonyms: list of str, optional | |
| If LLM was used, log the generated synonym list. | |
| output_dir: str or Path | |
| Directory to save the log file in. | |
| """ | |
| timestamp = now_utc_compact() | |
| output_dir = Path(output_dir) | |
| output_dir.mkdir(parents=True, exist_ok=True) | |
| # Metadata | |
| meta = { | |
| "created_at": timestamp, | |
| "keyword": keyword, | |
| "num_documents": len(documents), | |
| "num_topics": len(topic_model.get_topics()), # includes -1 (outliers) | |
| "matched_topic_ids": matched_topic_ids, | |
| "match_type_by_topic": match_type_by_topic, # e.g. {3: "exact", 7: "synonym"} | |
| "matched_synonyms": matched_synonyms, # e.g. {7: "explainable AI", ...} | |
| "generated_synonyms": generated_synonyms or [], | |
| } | |
| # Per-document results (includes topic, prob, representative etc.) | |
| doc_info_df: pd.DataFrame = topic_model.get_document_info(documents) | |
| """ | |
| >>> topic_model.get_document_info(docs) | |
| Document Topic Name Top_n_words Probability ... | |
| I am sure some bashers of Pens... 0 0_game_team_games_season game - team - games... 0.200010 ... | |
| My brother is in the market for... -1 -1_can_your_will_any can - your - will... 0.420668 ... | |
| Finally you said what you dream... -1 -1_can_your_will_any can - your - will... 0.807259 ... | |
| Think! It is the SCSI card doing... 49 49_windows_drive_dos_file windows - drive - docs... 0.071746 ... | |
| 1) I have an old Jasmine drive... 49 49_windows_drive_dos_file windows - drive - docs... 0.038983 ... | |
| """ | |
| doc_info = doc_info_df.to_dict(orient="records") | |
| # Optional original chunk metadata | |
| chunk_metadata = None | |
| if document_chunks: | |
| chunk_metadata = [ | |
| { | |
| "page_content": doc.page_content, | |
| "metadata": doc.metadata, | |
| } for doc in document_chunks | |
| ] | |
| # Topic summary info (counts, top_n_words, etc.) | |
| topic_info_df: pd.DataFrame = topic_model.get_topic_info() # BERTopic's topic summary | |
| # >>> topic_model.get_topic_info() | |
| # Topic Count Name | |
| # -1 4630 -1_can_your_will_any | |
| # 0 693 49_windows_drive_dos_file | |
| # 1 466 32_jesus_bible_christian_faith | |
| # 2 441 2_space_launch_orbit_lunar | |
| # 3 381 22_key_encryption_keys_encrypted | |
| topic_info = topic_info_df.to_dict(orient="records") | |
| payload: Dict[str, Any] = { | |
| "meta": meta, | |
| "topic_info": topic_info, # Get all topic information | |
| "document_info": doc_info, # Get all document information | |
| "document_chunks": chunk_metadata, | |
| } | |
| out_path = output_dir / f"01.1.5_topic_modeling_summary_{keyword}_{timestamp}.json" | |
| write_json(out_path, payload) | |
| print(f"\n[INFO] Wrote topic modeling summary to {out_path}") | |