feat: add improvement based on review

Author: baixiac

app/trainers/medcat_trainer.py

@@ -174,12 +174,13 @@ def run(
                     fn_accumulated += fns.get(cui, 0)
                     tp_accumulated += tps.get(cui, 0)
                     cc_accumulated += cc.get(cui, 0)
+                    cui_info = model.cdb.cui2info.get(cui)
                     aggregated_metrics.append({
                         "per_concept_fp": fps.get(cui, 0),
                         "per_concept_fn": fns.get(cui, 0),
                         "per_concept_tp": tps.get(cui, 0),
                         "per_concept_counts": cc.get(cui, 0),
-                        "per_concept_count_train": cast(Dict[str, Any], model.cdb.cui2info.get(cui, {})).get("count_train", 0),
+                        "per_concept_count_train": cui_info.get("count_train", 0) if cui_info is not None else 0,
                         "per_concept_acc_fp": fp_accumulated,
                         "per_concept_acc_fn": fn_accumulated,
                         "per_concept_acc_tp": tp_accumulated,
@@ -309,8 +310,9 @@ def _save_trained_concepts(
             annotation_ignorance_count = []
             concepts = list(training_concepts.keys())
             for c in concepts:
-                train_count.append(model.cdb.cui2info.get(c, {}).get("count_train", 0))  # type: ignore
-                concept_names.append(model.cdb.cui2info.get(c, {}).get("preferred_name", ""))  # type: ignore
+                cui_info = model.cdb.cui2info.get(c)
+                train_count.append(cui_info.get("count_train", 0) if cui_info is not None else 0)
+                concept_names.append(model.cdb.get_name(c))
                 annotation_count.append(training_concepts[c])
                 annotation_unique_count.append(training_unique_concepts[c])
                 annotation_ignorance_count.append(training_ignorance_counts[c])
@@ -421,7 +423,7 @@ def run(
             logger.info("Performing unsupervised training...")
             step = 0
             self._tracker_client.send_model_stats(dict(model.cdb.get_basic_info()), step)
-            before_cui2count_train = {c: info["count_train"] for c, info in model.cdb.cui2info.items()}
+            before_cui2count_train = model.cdb.get_cui2count_train()
             num_of_docs = 0
             train_unsupervised_params = get_func_params_as_dict(model.trainer.train_unsupervised)
             train_unsupervised_params = {p_key: training_params[p_key] if p_key in training_params else p_val for p_key, p_val in train_unsupervised_params.items()}
@@ -437,13 +439,14 @@ def run(
 
             self._tracker_client.log_document_size(num_of_docs)
             after_cui2count_train = {
-                c: info["count_train"]
-                for c, info in sorted(
-                    model.cdb.cui2info.items(),
-                    key=lambda item: item[1]["count_train"],
+                c: ct
+                for c, ct in sorted(
+                    model.cdb.get_cui2count_train().items(),
+                    key=lambda item: item[1],
                     reverse=True,
                 )
             }
+
             aggregated_metrics = []
             cui_step = 0
             for cui, train_count in after_cui2count_train.items():

app/trainers/metacat_trainer.py

@@ -86,16 +86,13 @@ def run(
                     model = self._model_service.load_model(copied_model_pack_path)
                 is_retrained = False
                 model.config.meta.description = description or model.config.meta.description
-                meta_cat_addons = [
-                    addon for addon in model.get_addons()
-                    if addon.addon_type == MetaCATAddon.addon_type
-                ]
+                meta_cat_addons = model.get_addons_of_type(MetaCATAddon)
                 for meta_cat_addon in meta_cat_addons:
                     if self._cancel_event.is_set():
                         self._cancel_event.clear()
                         raise TrainingCancelledException("Training was cancelled by the user")
 
-                    meta_cat = cast(MetaCATAddon, meta_cat_addon).mc
+                    meta_cat = meta_cat_addon.mc
                     category_name = meta_cat.config.general.category_name
                     assert category_name is not None, "Category name should not be None"
                     if meta_cat.config.general.alternative_class_names == [[]]:
@@ -203,12 +200,9 @@ def run(
                 logger.info("Evaluating the running model...")
                 metrics: List[Dict] = []
                 assert self._model_service.model is not None, "Model should not be None"
-                meta_cat_addons = [
-                    addon for addon in self._model_service.model.get_addons()
-                    if addon.addon_type == MetaCATAddon.addon_type
-                ]
+                meta_cat_addons = self._model_service.model.get_addons_of_type(MetaCATAddon)
                 for meta_cat_addon in meta_cat_addons:
-                    meta_cat = cast(MetaCATAddon, meta_cat_addon).mc
+                    meta_cat = meta_cat_addon.mc
                     category_name = meta_cat.config.general.category_name
                     self._tracker_client.log_model_config(self.get_flattened_metacat_config(meta_cat, category_name))
                     self._tracker_client.log_trainer_version(TrainerBackend.MEDCAT, medcat_version)