feat(thematos/models): Add TrainSummaryConfig in config.py

entelecheia · Aug 13, 2023 · bfe7517 · bfe7517
1 parent ef40b3f
commit bfe7517
Showing 1 changed file with 19 additions and 2 deletions.
diff --git a/src/thematos/models/config.py b/src/thematos/models/config.py
@@ -1,6 +1,7 @@
-from typing import List, Optional, Tuple
+import sys
+from typing import Any, List, Optional, Tuple
 
-from hyfi.composer import BaseModel
+from hyfi.composer import BaseModel, Field
 
 from thematos.plots import WordCloud
 
@@ -29,6 +30,22 @@ class TrainConfig(BaseModel):
     iterations: int = 100
 
 
+class TrainSummaryConfig(BaseModel):
+    _config_group_ = "/model/summary"
+    _config_name_ = "topic_train"
+
+    initial_hp: bool = Field(
+        True, description="whether to show the initial parameters at model creation"
+    )
+    params: bool = Field(
+        True, description="whether to show the current parameters of the model"
+    )
+    topic_word_top_n: int = Field(
+        10, description="the number of words by topic to display"
+    )
+    flush: bool = Field(False, description="whether to forcibly flush the stream")
+
+
 class WordcloudConfig(BaseModel):
     _config_group_ = "/model/plot"
     _config_name_ = "wordcloud"