CU-86995mmb6 Fix some RelCAT typing issues (CogStack/MedCAT#542)

mart-r · web-flow · commit 7a62ebe3cb19 · 2025-05-23T09:53:34.000+01:00
* CU-86995mmb6: Fix typing issues with modern bert config

* CU-86995mmb6: Fix issue with llama tokenizer wrapper

* CU-86995mmb6: Fix typing issues with llama config

* CU-86995mmb6: Make bert model config types explicit as well

* CU-86995mmb6: Remove excessive/unncessary type hints

* CU-86995mmb6: Remove unncessary type ignore
diff --git a/medcat-v1/medcat/utils/relation_extraction/bert/config.py b/medcat-v1/medcat/utils/relation_extraction/bert/config.py
@@ -1,5 +1,7 @@
 import logging
 import os
+from typing import cast
+
 from medcat.config_rel_cat import ConfigRelCAT
 from medcat.utils.relation_extraction.config import BaseConfig_RelationExtraction
 from transformers import BertConfig
@@ -13,18 +15,19 @@ class BertConfig_RelationExtraction(BaseConfig_RelationExtraction):
 
     name = 'bert-config'
     pretrained_model_name_or_path = "bert-base-uncased"
+    hf_model_config: BertConfig
 
     @classmethod
     def load(cls, pretrained_model_name_or_path: str, relcat_config: ConfigRelCAT, **kwargs)  -> "BertConfig_RelationExtraction":
         model_config = cls(pretrained_model_name_or_path, **kwargs)
 
         if pretrained_model_name_or_path and os.path.exists(pretrained_model_name_or_path):
-            model_config.hf_model_config = BertConfig.from_json_file(pretrained_model_name_or_path)
+            model_config.hf_model_config = cast(BertConfig, BertConfig.from_json_file(pretrained_model_name_or_path))
             logger.info("Loaded config from file: " + pretrained_model_name_or_path)
         else:
             relcat_config.general.model_name = cls.pretrained_model_name_or_path
-            model_config.hf_model_config = BertConfig.from_pretrained(
-                pretrained_model_name_or_path=cls.pretrained_model_name_or_path, **kwargs)
+            model_config.hf_model_config = cast(BertConfig, BertConfig.from_pretrained(
+                pretrained_model_name_or_path=cls.pretrained_model_name_or_path, **kwargs))
             logger.info("Loaded config from pretrained: " + relcat_config.general.model_name)
 
         return model_config
diff --git a/medcat-v1/medcat/utils/relation_extraction/bert/model.py b/medcat-v1/medcat/utils/relation_extraction/bert/model.py
@@ -12,7 +12,6 @@
 from medcat.config_rel_cat import ConfigRelCAT
 from medcat.utils.relation_extraction.ml_utils import create_dense_layers
 from medcat.utils.relation_extraction.models import BaseModel_RelationExtraction
-from medcat.utils.relation_extraction.config import BaseConfig_RelationExtraction
 from medcat.utils.relation_extraction.bert.config import BertConfig_RelationExtraction
 
 
@@ -24,25 +23,25 @@ class BertModel_RelationExtraction(BaseModel_RelationExtraction):
 
     log = logging.getLogger(__name__)
 
-    def __init__(self, pretrained_model_name_or_path: str, relcat_config: ConfigRelCAT, model_config: Union[BaseConfig_RelationExtraction, BertConfig_RelationExtraction]):
+    def __init__(self, pretrained_model_name_or_path: str, relcat_config: ConfigRelCAT, model_config: BertConfig_RelationExtraction):
         """ Class to hold the BERT model + model_config
 
         Args:
             pretrained_model_name_or_path (str): path to load the model from,
                     this can be a HF model i.e: "bert-base-uncased", if left empty, it is normally assumed that a model is loaded from 'model.dat'
                     using the RelCAT.load() method. So if you are initializing/training a model from scratch be sure to base it on some model.            
             relcat_config (ConfigRelCAT): relcat config.
-            model_config (Union[BaseConfig_RelationExtraction | BertConfig_RelationExtraction]): HF bert config for model.
+            model_config (BertConfig_RelationExtraction): HF bert config for model.
         """
         super(BertModel_RelationExtraction, self).__init__(pretrained_model_name_or_path=pretrained_model_name_or_path,
                                                           relcat_config=relcat_config,
                                                           model_config=model_config)
 
         self.relcat_config: ConfigRelCAT = relcat_config
-        self.model_config: Union[BaseConfig_RelationExtraction, BertConfig_RelationExtraction] = model_config
+        self.model_config = model_config
         self.pretrained_model_name_or_path: str = pretrained_model_name_or_path
 
-        self.hf_model: Union[BertModel, PreTrainedModel] = BertModel(model_config.hf_model_config) # type: ignore
+        self.hf_model: Union[BertModel, PreTrainedModel] = BertModel(model_config.hf_model_config)
 
         for param in self.hf_model.parameters(): # type: ignore
             if self.relcat_config.model.freeze_layers:
@@ -55,8 +54,10 @@ def __init__(self, pretrained_model_name_or_path: str, relcat_config: ConfigRelC
         # dense layers
         self.fc1, self.fc2, self.fc3 = create_dense_layers(self.relcat_config)
 
+    # NOTEL ignoring type due to the type of model_config not matching base class exactly (subclass)
     @classmethod
-    def load(cls, pretrained_model_name_or_path: str, relcat_config: ConfigRelCAT, model_config: Union[BaseConfig_RelationExtraction, BertConfig_RelationExtraction], **kwargs) -> "BertModel_RelationExtraction":
+    def load(cls, pretrained_model_name_or_path: str, relcat_config: ConfigRelCAT,  # type: ignore
+             model_config: BertConfig_RelationExtraction, **kwargs) -> "BertModel_RelationExtraction":
 
         model = BertModel_RelationExtraction(pretrained_model_name_or_path=pretrained_model_name_or_path,
                                              relcat_config=relcat_config,
diff --git a/medcat-v1/medcat/utils/relation_extraction/llama/config.py b/medcat-v1/medcat/utils/relation_extraction/llama/config.py
@@ -1,5 +1,7 @@
 import logging
 import os
+from typing import cast
+
 from medcat.config_rel_cat import ConfigRelCAT
 from medcat.utils.relation_extraction.config import BaseConfig_RelationExtraction
 from transformers import LlamaConfig
@@ -13,18 +15,19 @@ class LlamaConfig_RelationExtraction(BaseConfig_RelationExtraction):
 
     name = 'llama-config'
     pretrained_model_name_or_path = "meta-llama/Llama-3.1-8B"
+    hf_model_config: LlamaConfig
 
     @classmethod
     def load(cls, pretrained_model_name_or_path: str, relcat_config: ConfigRelCAT, **kwargs) -> "LlamaConfig_RelationExtraction":
         model_config = cls(pretrained_model_name_or_path, **kwargs)
 
         if pretrained_model_name_or_path and os.path.exists(pretrained_model_name_or_path):
-            model_config.hf_model_config = LlamaConfig.from_json_file(pretrained_model_name_or_path)
+            model_config.hf_model_config = cast(LlamaConfig, LlamaConfig.from_json_file(pretrained_model_name_or_path))
             logger.info("Loaded config from file: " + pretrained_model_name_or_path)
         else:
             relcat_config.general.model_name = cls.pretrained_model_name_or_path
-            model_config.hf_model_config = LlamaConfig.from_pretrained(
-                pretrained_model_name_or_path=cls.pretrained_model_name_or_path, **kwargs)
+            model_config.hf_model_config = cast(LlamaConfig, LlamaConfig.from_pretrained(
+                pretrained_model_name_or_path=cls.pretrained_model_name_or_path, **kwargs))
             logger.info("Loaded config from pretrained: " + relcat_config.general.model_name)
 
         return model_config
diff --git a/medcat-v1/medcat/utils/relation_extraction/llama/model.py b/medcat-v1/medcat/utils/relation_extraction/llama/model.py
@@ -1,12 +1,11 @@
 import logging
-from typing import Any, Optional, Tuple, Union
+from typing import Any, Optional, Tuple
 import torch
 from torch import nn
 import os
 from transformers.models.llama import LlamaModel
 
 from medcat.config_rel_cat import ConfigRelCAT
-from medcat.utils.relation_extraction.config import BaseConfig_RelationExtraction
 from medcat.utils.relation_extraction.llama.config import LlamaConfig_RelationExtraction
 from medcat.utils.relation_extraction.models import BaseModel_RelationExtraction
 from medcat.utils.relation_extraction.ml_utils import create_dense_layers, get_annotation_schema_tag
@@ -20,25 +19,25 @@ class LlamaModel_RelationExtraction(BaseModel_RelationExtraction):
 
     log = logging.getLogger(__name__)
 
-    def __init__(self, pretrained_model_name_or_path: str, relcat_config: ConfigRelCAT, model_config: Union[BaseConfig_RelationExtraction, LlamaConfig_RelationExtraction]):
+    def __init__(self, pretrained_model_name_or_path: str, relcat_config: ConfigRelCAT, model_config: LlamaConfig_RelationExtraction):
         """ Class to hold the Llama model + model_config
 
         Args:
             pretrained_model_name_or_path (str): path to load the model from,
                     this can be a HF model i.e: "bert-base-uncased", if left empty, it is normally assumed that a model is loaded from 'model.dat'
                     using the RelCAT.load() method. So if you are initializing/training a model from scratch be sure to base it on some model.            
             relcat_config (ConfigRelCAT): relcat config.
-            model_config (Union[BaseConfig_RelationExtraction | LlamaConfig_RelationExtraction]): HF bert config for model.
+            model_config (LlamaConfig_RelationExtraction): HF bert config for model.
         """
 
         super(LlamaModel_RelationExtraction, self).__init__(pretrained_model_name_or_path=pretrained_model_name_or_path,
                                                           relcat_config=relcat_config,
                                                           model_config=model_config)
 
         self.relcat_config: ConfigRelCAT = relcat_config
-        self.model_config: Union[BaseConfig_RelationExtraction, LlamaConfig_RelationExtraction] = model_config
+        self.model_config = model_config
 
-        self.hf_model: LlamaModel = LlamaModel(config=model_config) # type: ignore
+        self.hf_model: LlamaModel = LlamaModel(config=model_config.hf_model_config)
 
         if pretrained_model_name_or_path != "":
             self.hf_model = LlamaModel.from_pretrained(pretrained_model_name_or_path, config=model_config, ignore_mismatched_sizes=True)
@@ -162,8 +161,10 @@ def forward(self,
 
         return model_output, classification_logits.to(self.relcat_config.general.device)
 
+    # NOTEL ignoring type due to the type of model_config not matching base class exactly (subclass)
     @classmethod
-    def load(cls, pretrained_model_name_or_path: str, relcat_config: ConfigRelCAT, model_config: Union[BaseConfig_RelationExtraction, LlamaConfig_RelationExtraction], **kwargs) -> "LlamaModel_RelationExtraction":
+    def load(cls, pretrained_model_name_or_path: str, relcat_config: ConfigRelCAT,  # type: ignore
+             model_config: LlamaConfig_RelationExtraction, **kwargs) -> "LlamaModel_RelationExtraction":
 
         model = LlamaModel_RelationExtraction(pretrained_model_name_or_path=pretrained_model_name_or_path,
                                              relcat_config=relcat_config,
diff --git a/medcat-v1/medcat/utils/relation_extraction/llama/tokenizer.py b/medcat-v1/medcat/utils/relation_extraction/llama/tokenizer.py
@@ -30,5 +30,5 @@ def load(cls, tokenizer_path: str, relcat_config: ConfigRelCAT, **kwargs) -> "To
         else:
             relcat_config.general.model_name = cls.pretrained_model_name_or_path
             tokenizer.hf_tokenizers = LlamaTokenizerFast.from_pretrained(
-                path=relcat_config.general.model_name)
+                relcat_config.general.model_name)
         return tokenizer
diff --git a/medcat-v1/medcat/utils/relation_extraction/models.py b/medcat-v1/medcat/utils/relation_extraction/models.py
@@ -1,6 +1,6 @@
 import logging
 import torch
-from typing import Any, Optional, Tuple, Union
+from typing import Any, Optional, Tuple, Union, cast
 from torch import nn
 from transformers import PretrainedConfig, PreTrainedModel
 
@@ -240,15 +240,21 @@ def load(cls, pretrained_model_name_or_path: str, relcat_config: ConfigRelCAT, m
         if "modern-bert" in relcat_config.general.tokenizer_name or \
              "modern-bert" in relcat_config.general.model_name:
             from medcat.utils.relation_extraction.modernbert.model import ModernBertModel_RelationExtraction
-            model = ModernBertModel_RelationExtraction.load(pretrained_model_name_or_path, relcat_config=relcat_config, model_config=model_config)
+            from medcat.utils.relation_extraction.modernbert.config import ModernBertConfig_RelationExtraction
+            model = ModernBertModel_RelationExtraction.load(pretrained_model_name_or_path, relcat_config=relcat_config,
+                                                            model_config=cast(ModernBertConfig_RelationExtraction, model_config))
         elif "bert" in relcat_config.general.tokenizer_name or \
              "bert" in relcat_config.general.model_name:
             from medcat.utils.relation_extraction.bert.model import BertModel_RelationExtraction
-            model = BertModel_RelationExtraction.load(pretrained_model_name_or_path, relcat_config=relcat_config, model_config=model_config)
+            from medcat.utils.relation_extraction.bert.config import BertConfig_RelationExtraction
+            model = BertModel_RelationExtraction.load(pretrained_model_name_or_path, relcat_config=relcat_config,
+                                                      model_config=cast(BertConfig_RelationExtraction, model_config))
         elif "llama" in relcat_config.general.tokenizer_name or \
              "llama" in relcat_config.general.model_name:
             from medcat.utils.relation_extraction.llama.model import LlamaModel_RelationExtraction
-            model = LlamaModel_RelationExtraction.load(pretrained_model_name_or_path, relcat_config=relcat_config, model_config=model_config)
+            from medcat.utils.relation_extraction.llama.config import LlamaConfig_RelationExtraction
+            model = LlamaModel_RelationExtraction.load(pretrained_model_name_or_path, relcat_config=relcat_config,
+                                                       model_config=cast(LlamaConfig_RelationExtraction, model_config))
         else:
             if pretrained_model_name_or_path:
                 model.hf_model = PreTrainedModel.from_pretrained(pretrained_model_name_or_path=pretrained_model_name_or_path, config=model_config)
diff --git a/medcat-v1/medcat/utils/relation_extraction/modernbert/config.py b/medcat-v1/medcat/utils/relation_extraction/modernbert/config.py
@@ -1,5 +1,7 @@
 import logging
 import os
+from typing import cast
+
 from medcat.config_rel_cat import ConfigRelCAT
 from medcat.utils.relation_extraction.config import BaseConfig_RelationExtraction
 from transformers import ModernBertConfig
@@ -13,18 +15,19 @@ class ModernBertConfig_RelationExtraction(BaseConfig_RelationExtraction):
 
     name = 'modern-bert-config'
     pretrained_model_name_or_path = "answerdotai/ModernBERT-base"
+    hf_model_config: ModernBertConfig
 
     @classmethod
     def load(cls, pretrained_model_name_or_path: str, relcat_config: ConfigRelCAT, **kwargs) -> "ModernBertConfig_RelationExtraction":
         model_config = cls(pretrained_model_name_or_path=pretrained_model_name_or_path, **kwargs)
 
         if pretrained_model_name_or_path and os.path.exists(pretrained_model_name_or_path):
-            model_config.hf_model_config = ModernBertConfig.from_json_file(pretrained_model_name_or_path)
+            model_config.hf_model_config = cast(ModernBertConfig, ModernBertConfig.from_json_file(pretrained_model_name_or_path))
             logger.info("Loaded config from file: " + pretrained_model_name_or_path)
         else:
             relcat_config.general.model_name = cls.pretrained_model_name_or_path
-            model_config.hf_model_config = ModernBertConfig.from_pretrained(
-                pretrained_model_name_or_path=cls.pretrained_model_name_or_path, **kwargs)
+            model_config.hf_model_config = cast(ModernBertConfig, ModernBertConfig.from_pretrained(
+                pretrained_model_name_or_path=cls.pretrained_model_name_or_path, **kwargs))
             logger.info("Loaded config from pretrained: " + relcat_config.general.model_name)
 
         return model_config
diff --git a/medcat-v1/medcat/utils/relation_extraction/modernbert/model.py b/medcat-v1/medcat/utils/relation_extraction/modernbert/model.py
@@ -9,7 +9,6 @@
 from transformers import PreTrainedModel
 from medcat.utils.relation_extraction.ml_utils import create_dense_layers
 from medcat.utils.relation_extraction.models import BaseModel_RelationExtraction
-from medcat.utils.relation_extraction.config import BaseConfig_RelationExtraction
 from medcat.utils.relation_extraction.modernbert.config import ModernBertConfig_RelationExtraction
 
 
@@ -22,22 +21,22 @@ class ModernBertModel_RelationExtraction(BaseModel_RelationExtraction):
 
     log = logging.getLogger(__name__)
 
-    def __init__(self, pretrained_model_name_or_path: str, relcat_config: ConfigRelCAT, model_config: Union[BaseConfig_RelationExtraction, ModernBertConfig_RelationExtraction]):
+    def __init__(self, pretrained_model_name_or_path: str, relcat_config: ConfigRelCAT, model_config: ModernBertConfig_RelationExtraction):
         """ Class to hold the ModernBERT model + model_config
 
         Args:
             pretrained_model_name_or_path (str): path to load the model from,
                     this can be a HF model i.e: "bert-base-uncased", if left empty, it is normally assumed that a model is loaded from 'model.dat'
                     using the RelCAT.load() method. So if you are initializing/training a model from scratch be sure to base it on some model.
             relcat_config (ConfigRelCAT): relcat config.
-            model_config (Union[BaseConfig_RelationExtraction | ModernBertConfig_RelationExtraction]): HF bert config for model.
+            model_config (ModernBertConfig_RelationExtraction): HF bert config for model.
         """
         super(ModernBertModel_RelationExtraction, self).__init__(pretrained_model_name_or_path=pretrained_model_name_or_path,
                                                     relcat_config=relcat_config,
                                                     model_config=model_config)
 
         self.relcat_config: ConfigRelCAT = relcat_config
-        self.model_config: Union[BaseConfig_RelationExtraction, ModernBertConfig_RelationExtraction] = model_config
+        self.model_config = model_config
         self.pretrained_model_name_or_path: str = pretrained_model_name_or_path
 
         self.hf_model: Union[ModernBertModel, PreTrainedModel] = ModernBertModel(config=model_config.hf_model_config)
@@ -53,8 +52,10 @@ def __init__(self, pretrained_model_name_or_path: str, relcat_config: ConfigRelC
         # dense layers
         self.fc1, self.fc2, self.fc3 = create_dense_layers(self.relcat_config)
 
+    # NOTEL ignoring type due to the type of model_config not matchin base class exactly (subclass)
     @classmethod
-    def load(cls, pretrained_model_name_or_path: str, relcat_config: ConfigRelCAT, model_config: Union[BaseConfig_RelationExtraction, ModernBertConfig_RelationExtraction], **kwargs) -> "ModernBertModel_RelationExtraction":
+    def load(cls, pretrained_model_name_or_path: str, relcat_config: ConfigRelCAT,  # type: ignore
+             model_config: ModernBertConfig_RelationExtraction, **kwargs) -> "ModernBertModel_RelationExtraction":
 
         model = ModernBertModel_RelationExtraction(pretrained_model_name_or_path=pretrained_model_name_or_path,
                                              relcat_config=relcat_config,
