feat(thematos/models): add TrainSummaryConfig in imports and TopicMod…

…el, modify the method of calling model.summary, assign default titles if not set in wc_args
entelecheia · Aug 13, 2023 · ef40b3f · ef40b3f
1 parent 111ea9f
commit ef40b3f
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/src/thematos/models/base.py b/src/thematos/models/base.py
@@ -12,7 +12,7 @@
 
 from thematos.datasets import Corpus
 
-from .config import LdaConfig, TrainConfig, WordcloudConfig
+from .config import LdaConfig, TrainConfig, WordcloudConfig, TrainSummaryConfig
 from .prior import WordPrior
 from .types import CoherenceMetrics, ModelSummary
 
@@ -30,6 +30,7 @@ class TopicModel(BatchTaskConfig):
     corpus: Corpus = Corpus()
     model_args: LdaConfig = LdaConfig()
     train_args: TrainConfig = TrainConfig()
+    train_summary_args: TrainSummaryConfig = TrainSummaryConfig()
     wc_args: WordcloudConfig = WordcloudConfig()
 
     coherence_metric_list: List[str] = ["u_mass", "c_uci", "c_npmi", "c_v"]
@@ -306,7 +307,7 @@ def save_train_summary(self) -> None:
         original_stdout = sys.stdout
         with open(self.train_summary_file, "w") as f:
             sys.stdout = f  # Change the standard output to the file.
-            self.model.summary()
+            self.model.summary(**self.train_summary_args.kwargs)
             if coh_values:
                 print("<Topic Coherence Scores>")
                 for cm, cv in coh_values.items():
@@ -430,6 +431,8 @@ def generate_wordclouds(
             images.append(img)
 
         if wc_args.make_collage:
+            if not wc_args.titles:
+                wc_args.titles = [f"Topic {i}" for i in range(self.num_topics)]
             output_dir = self.output_dir / "wordcloud_collage"
             output_file_format = self.model_id + "_wordcloud_{page_num:02d}.png"
             HyFI.make_subplot_pages_from_images(