Improve abstraction by moving model and config loading outside of tokenizer into a base component class

Author: mart-r

medcat/rel_cat.py

@@ -17,7 +17,8 @@
 from medcat.config import Config
 from medcat.config_rel_cat import ConfigRelCAT
 from medcat.pipeline.pipe_runner import PipeRunner
-from medcat.utils.relation_extraction.tokenizer import BaseTokenizerWrapper, load_tokenizer
+from medcat.utils.relation_extraction.base_component import load_base_component, BaseComponent
+from medcat.utils.relation_extraction.tokenizer import BaseTokenizerWrapper
 from spacy.tokens import Doc, Span
 from typing import Dict, Iterable, Iterator, List, cast
 from torch.utils.data import DataLoader, Sampler
@@ -91,8 +92,13 @@ class RelCAT(PipeRunner):
 
     log = logging.getLogger(__name__)
 
-    def __init__(self, cdb: CDB, tokenizer: BaseTokenizerWrapper, config: ConfigRelCAT = ConfigRelCAT(), task="train", init_model=False):
+    def __init__(self, cdb: CDB,
+                 base_component: BaseComponent,
+                 tokenizer: BaseTokenizerWrapper,
+                 config: ConfigRelCAT = ConfigRelCAT(),
+                 task="train", init_model=False):
         self.config = config
+        self.base_component = base_component
         self.tokenizer: BaseTokenizerWrapper = tokenizer
         self.cdb = cdb
 
@@ -154,8 +160,8 @@ def _get_model(self):
 
         """ Used only for model initialisation.
         """
-        self.model_config = self.tokenizer.config_from_pretrained()
-        self.model = self.tokenizer.model_from_pretrained(relcat_config=self.config,
+        self.model_config = self.base_component.config_from_pretrained()
+        self.model = self.base_component.model_from_pretrained(relcat_config=self.config,
                                                           model_config=self.model_config)
 
     @classmethod
@@ -182,20 +188,22 @@ def load(cls, load_path: str = "./") -> "RelCAT":
         if "bert" in config.general.tokenizer_name or "llama" in config.general.tokenizer_name:
             tokenizer_path = load_path
 
-        tokenizer = load_tokenizer(tokenizer_path, config)
+        base_component = load_base_component(tokenizer_path, config)
+        tokenizer = base_component.tokenizer
 
         model_config_path = os.path.join(load_path, "model_config.json")
 
         if os.path.exists(model_config_path):
-            model_config = tokenizer.config_from_json_file(model_config_path)
+            model_config = base_component.config_from_json_file(model_config_path)
             cls.log.info("Loaded config from : " + model_config_path)
         else:
             cls.log.info("model_config.json not found, using default for the model")
-            model_config = tokenizer.config_from_pretrained()
+            model_config = base_component.config_from_pretrained()
 
         model_config.vocab_size = tokenizer.get_size()
 
         rel_cat = cls(cdb=cdb, config=config,
+                      base_component=base_component,
                       tokenizer=tokenizer,
                       task=config.general.task)
 
@@ -209,10 +217,11 @@ def load(cls, load_path: str = "./") -> "RelCAT":
 
             if os.path.exists(os.path.join(load_path, config.general.model_name)):
                 # NOTE: should it be the joined path? it wasn't previously
-                rel_cat.model = tokenizer.model_from_pretrained(relcat_config=config, model_config=model_config,
-                                                                pretrained_model_name_or_path=config.general.model_name)
+                rel_cat.model = base_component.model_from_pretrained(
+                    relcat_config=config, model_config=model_config,
+                    pretrained_model_name_or_path=config.general.model_name)
             else:
-                rel_cat.model = tokenizer.model_from_pretrained(
+                rel_cat.model = base_component.model_from_pretrained(
                     pretrained_model_name_or_path='',
                     relcat_config=config,
                     model_config=model_config)
@@ -228,7 +237,7 @@ def load(cls, load_path: str = "./") -> "RelCAT":
 
             cls.log.error("Failed to load specified HF model, defaulting to 'bert-base-uncased', loading...")
             # NOTE: this won't really work for Llama or ModernBert, I've got a feeling
-            rel_cat.model = tokenizer.model_from_pretrained(
+            rel_cat.model = base_component.model_from_pretrained(
                 pretrained_model_name_or_path="bert-base-uncased",
                 relcat_config=config,
                 model_config=model_config)

medcat/utils/relation_extraction/base_component.py

@@ -0,0 +1,45 @@
+from abc import ABC, abstractmethod
+import logging
+
+from transformers import PretrainedConfig
+
+from medcat.config_rel_cat import ConfigRelCAT
+from medcat.utils.relation_extraction.tokenizer import BaseTokenizerWrapper
+from medcat.utils.relation_extraction.models import Base_RelationExtraction
+
+
+logger = logging.getLogger(__name__)
+
+
+class BaseComponent(ABC):
+
+    @property
+    @abstractmethod
+    def tokenizer(self) -> BaseTokenizerWrapper:
+        pass
+
+    @abstractmethod
+    def config_from_pretrained(self) -> PretrainedConfig:
+        pass # perhaps some doc string
+
+    @abstractmethod
+    def config_from_json_file(self, file_path: str) -> PretrainedConfig:
+        pass # perhaps some doc string
+
+    @abstractmethod
+    def model_from_pretrained(self, relcat_config: ConfigRelCAT, model_config: PretrainedConfig,
+            pretrained_model_name_or_path: str = 'default') -> Base_RelationExtraction:
+        pass # perhaps some doc string
+
+
+def load_base_component(tokenizer_path: str, config: ConfigRelCAT) -> BaseComponent:
+    if "modern-bert-tokenizer" in config.general.tokenizer_name:
+        from medcat.utils.relation_extraction.modernbert.component import ModernBertComponent
+        return ModernBertComponent(tokenizer_path, config)
+    elif "bert" in config.general.tokenizer_name:
+        from medcat.utils.relation_extraction.bert.component import BertComponent
+        return BertComponent(tokenizer_path, config)
+    elif "llama" in config.general.tokenizer_name:
+        from medcat.utils.relation_extraction.llama.component import LlamaComponent
+        return LlamaComponent(tokenizer_path, config)
+    raise ValueError(f"Could not find matching base component for {config.general.tokenizer_name}")

medcat/utils/relation_extraction/bert/component.py

@@ -0,0 +1,40 @@
+import os
+from typing import Optional
+
+from transformers import PretrainedConfig, BertConfig
+
+from medcat.config_rel_cat import ConfigRelCAT
+from medcat.utils.relation_extraction.base_component import BaseComponent
+from medcat.utils.relation_extraction.tokenizer import BaseTokenizerWrapper, load_default_tokenizer
+from medcat.utils.relation_extraction.models import Base_RelationExtraction
+from medcat.utils.relation_extraction.bert.tokenizer import TokenizerWrapperBERT
+from medcat.utils.relation_extraction.bert.model import BertModel_RelationExtraction
+
+
+class BertComponent(BaseComponent):
+    pretrained_model_name_or_path = "bert-base-uncased"
+
+    def __init__(self, tokenizer_path: str, config: ConfigRelCAT,
+                 tokenizer: Optional[BaseTokenizerWrapper] = None):
+        if tokenizer is not None:
+            self._tokenizer = tokenizer
+        elif os.path.exists(tokenizer_path):
+            self._tokenizer = TokenizerWrapperBERT.load(tokenizer_path)
+        else:
+            self._tokenizer = load_default_tokenizer(tokenizer_path, config)
+
+    @property
+    def tokenizer(self) -> BaseTokenizerWrapper:
+        return self._tokenizer
+
+    def config_from_pretrained(self) -> PretrainedConfig:
+        return BertConfig.from_pretrained(self.pretrained_model_name_or_path)
+
+    def config_from_json_file(self, file_path: str) -> PretrainedConfig:
+        return BertConfig.from_json_file(file_path)
+
+    def model_from_pretrained(self, relcat_config: ConfigRelCAT, model_config: PretrainedConfig,
+            pretrained_model_name_or_path: str = 'default') -> Base_RelationExtraction:
+        if pretrained_model_name_or_path == 'default':
+            pretrained_model_name_or_path = self.pretrained_model_name_or_path
+        return BertModel_RelationExtraction(pretrained_model_name_or_path, relcat_config, model_config)

medcat/utils/relation_extraction/bert/tokenizer.py

@@ -1,13 +1,8 @@
 import os
-from transformers import PretrainedConfig
-from transformers import BertConfig
 from transformers.models.bert.tokenization_bert_fast import BertTokenizerFast
 import logging
 
-from medcat.config_rel_cat import ConfigRelCAT
 from medcat.utils.relation_extraction.tokenizer import BaseTokenizerWrapper
-from medcat.utils.relation_extraction.models import Base_RelationExtraction
-from medcat.utils.relation_extraction.bert.model import BertModel_RelationExtraction
 
 
 logger = logging.getLogger(__name__)
@@ -22,19 +17,6 @@ class TokenizerWrapperBERT(BaseTokenizerWrapper):
             A huggingface Fast BERT.
     '''
     name = 'bert-tokenizer'
-    pretrained_model_name_or_path = "bert-base-uncased"
-
-    def config_from_pretrained(self) -> PretrainedConfig:
-        return BertConfig.from_pretrained(self.pretrained_model_name_or_path)
-
-    def config_from_json_file(self, file_path: str) -> PretrainedConfig:
-        return BertConfig.from_json_file(file_path)
-
-    def model_from_pretrained(self, relcat_config: ConfigRelCAT, model_config: PretrainedConfig,
-            pretrained_model_name_or_path: str = 'default') -> Base_RelationExtraction:
-        if pretrained_model_name_or_path == 'default':
-            pretrained_model_name_or_path = self.pretrained_model_name_or_path
-        return BertModel_RelationExtraction(pretrained_model_name_or_path, relcat_config, model_config)
 
     @classmethod
     def load(cls, dir_path, **kwargs):

medcat/utils/relation_extraction/llama/component.py

@@ -0,0 +1,37 @@
+import os
+
+from transformers import PretrainedConfig
+from transformers.models.llama import LlamaConfig
+
+from medcat.config_rel_cat import ConfigRelCAT
+from medcat.utils.relation_extraction.base_component import BaseComponent
+from medcat.utils.relation_extraction.tokenizer import BaseTokenizerWrapper, load_default_tokenizer
+from medcat.utils.relation_extraction.models import Base_RelationExtraction
+from medcat.utils.relation_extraction.llama.tokenizer import TokenizerWrapperLlama
+from medcat.utils.relation_extraction.llama.model import LlamaModel_RelationExtraction
+
+
+class LlamaComponent(BaseComponent):
+    pretrained_model_name_or_path = "meta-llama/Llama-3.1-8B"
+
+    def __init__(self, tokenizer_path: str, config: ConfigRelCAT):
+        if os.path.exists(tokenizer_path):
+            self._tokenizer = TokenizerWrapperLlama.load(tokenizer_path)
+        else:
+            self._tokenizer = load_default_tokenizer(tokenizer_path, config)
+
+    @property
+    def tokenizer(self) -> BaseTokenizerWrapper:
+        return self._tokenizer
+
+    def config_from_pretrained(self) -> PretrainedConfig:
+        pass # perhaps some doc string
+
+    def config_from_json_file(self, file_path: str) -> PretrainedConfig:
+        return LlamaConfig.from_json_file(file_path)
+
+    def model_from_pretrained(self, relcat_config: ConfigRelCAT, model_config: PretrainedConfig,
+            pretrained_model_name_or_path: str = 'default') -> Base_RelationExtraction:
+        if pretrained_model_name_or_path == 'default':
+            pretrained_model_name_or_path = self.pretrained_model_name_or_path
+        return LlamaModel_RelationExtraction(pretrained_model_name_or_path, relcat_config, model_config)

medcat/utils/relation_extraction/llama/tokenizer.py

@@ -1,14 +1,9 @@
 import os
 from typing import Optional
-from transformers import PretrainedConfig
-from transformers.models.llama import LlamaConfig
 from transformers import LlamaTokenizerFast
 import logging
 
-from medcat.config_rel_cat import ConfigRelCAT
 from medcat.utils.relation_extraction.tokenizer import BaseTokenizerWrapper
-from medcat.utils.relation_extraction.models import Base_RelationExtraction
-from medcat.utils.relation_extraction.llama.model import LlamaModel_RelationExtraction
 
 
 logger = logging.getLogger(__name__)
@@ -23,19 +18,6 @@ class TokenizerWrapperLlama(BaseTokenizerWrapper):
             A huggingface Fast Llama.
     '''
     name = 'llama-tokenizer'
-    pretrained_model_name_or_path = "meta-llama/Llama-3.1-8B"
-
-    def config_from_pretrained(self) -> PretrainedConfig:
-        pass # perhaps some doc string
-
-    def config_from_json_file(self, file_path: str) -> PretrainedConfig:
-        return LlamaConfig.from_json_file(file_path)
-
-    def model_from_pretrained(self, relcat_config: ConfigRelCAT, model_config: PretrainedConfig,
-            pretrained_model_name_or_path: str = 'default') -> Base_RelationExtraction:
-        if pretrained_model_name_or_path == 'default':
-            pretrained_model_name_or_path = self.pretrained_model_name_or_path
-        return LlamaModel_RelationExtraction(pretrained_model_name_or_path, relcat_config, model_config)
 
     def __init__(self, hf_tokenizers=None, max_seq_length: Optional[int] = None, add_special_tokens: Optional[bool] = False):
         self.hf_tokenizers = hf_tokenizers

medcat/utils/relation_extraction/modernbert/component.py

@@ -0,0 +1,36 @@
+import os
+
+from transformers import PretrainedConfig, ModernBertConfig
+
+from medcat.config_rel_cat import ConfigRelCAT
+from medcat.utils.relation_extraction.base_component import BaseComponent
+from medcat.utils.relation_extraction.tokenizer import BaseTokenizerWrapper, load_default_tokenizer
+from medcat.utils.relation_extraction.models import Base_RelationExtraction
+from medcat.utils.relation_extraction.modernbert.tokenizer import TokenizerWrapperModernBERT
+from medcat.utils.relation_extraction.modernbert.model import ModernBertModel_RelationExtraction
+
+
+class ModernBertComponent(BaseComponent):
+    pretrained_model_name_or_path = "answerdotai/ModernBERT-base"
+
+    def __init__(self, tokenizer_path: str, config: ConfigRelCAT):
+        if os.path.exists(tokenizer_path):
+            self._tokenizer = TokenizerWrapperModernBERT.load(tokenizer_path)
+        else:
+            self._tokenizer = load_default_tokenizer(tokenizer_path, config)
+
+    @property
+    def tokenizer(self) -> BaseTokenizerWrapper:
+        return self._tokenizer
+
+    def config_from_pretrained(self) -> PretrainedConfig:
+        return ModernBertConfig.from_pretrained(self.pretrained_model_name_or_path)
+
+    def config_from_json_file(self, file_path: str) -> PretrainedConfig:
+        return ModernBertConfig.from_json_file(file_path)
+
+    def model_from_pretrained(self, relcat_config: ConfigRelCAT, model_config: PretrainedConfig,
+            pretrained_model_name_or_path: str = 'default') -> Base_RelationExtraction:
+        if pretrained_model_name_or_path == 'default':
+            pretrained_model_name_or_path = self.pretrained_model_name_or_path
+        return ModernBertModel_RelationExtraction(pretrained_model_name_or_path, relcat_config, model_config)

medcat/utils/relation_extraction/modernbert/tokenizer.py

@@ -1,13 +1,8 @@
 import os
-from transformers import PretrainedConfig
-from transformers import ModernBertConfig
 from transformers import PreTrainedTokenizerFast
 import logging
 
-from medcat.config_rel_cat import ConfigRelCAT
 from medcat.utils.relation_extraction.tokenizer import BaseTokenizerWrapper
-from medcat.utils.relation_extraction.models import Base_RelationExtraction
-from medcat.utils.relation_extraction.modernbert.model import ModernBertModel_RelationExtraction
 
 
 logger = logging.getLogger(__name__)
@@ -22,19 +17,6 @@ class TokenizerWrapperModernBERT(BaseTokenizerWrapper):
             A huggingface Fast tokenizer.
     '''
     name = 'modern-bert-tokenizer'
-    pretrained_model_name_or_path = "answerdotai/ModernBERT-base"
-
-    def config_from_pretrained(self) -> PretrainedConfig:
-        return ModernBertConfig.from_pretrained(self.pretrained_model_name_or_path)
-
-    def config_from_json_file(self, file_path: str) -> PretrainedConfig:
-        return ModernBertConfig.from_json_file(file_path)
-
-    def model_from_pretrained(self, relcat_config: ConfigRelCAT, model_config: PretrainedConfig,
-            pretrained_model_name_or_path: str = 'default') -> Base_RelationExtraction:
-        if pretrained_model_name_or_path == 'default':
-            pretrained_model_name_or_path = self.pretrained_model_name_or_path
-        return ModernBertModel_RelationExtraction(pretrained_model_name_or_path, relcat_config, model_config)
 
     @classmethod
     def load(cls, dir_path, **kwargs):