Latest fixes for model saving/loading.

vladd-bit · vladd-bit · commit 9cc52b491ba0 · 2025-04-23T17:45:25.000+01:00
diff --git a/medcat/rel_cat.py b/medcat/rel_cat.py
@@ -126,8 +126,9 @@ def load(cls, load_path: str = "./") -> "RelCAT":
 
         device = torch.device("cuda" if torch.cuda.is_available() and component.relcat_config.general.device != "cpu" else "cpu")
 
-        rel_cat = RelCAT(cdb=cdb, config=component.relcat_config, task=component.task, init_model=False)
+        rel_cat = RelCAT(cdb=cdb, config=component.relcat_config, task=component.task)
         rel_cat.device = device
+        rel_cat.component = component
 
         return rel_cat
 
@@ -255,7 +256,7 @@ def train(self, export_data_path:str = "", train_csv_path:str = "", test_csv_pat
                 gamma=self.component.relcat_config.train.multistep_lr_gamma) # type: ignore
 
         self.epoch, self.best_f1 = load_state(
-            self.component.model, self.component.optimizer, self.component.scheduler, load_best=False, path=checkpoint_path, config=self.component.relcat_config)
+            self.component.model, self.component.optimizer, self.component.scheduler, load_best=False, path=checkpoint_path, relcat_config=self.component.relcat_config)
 
         self.log.info("Starting training process...")
 
diff --git a/medcat/utils/relation_extraction/base_component.py b/medcat/utils/relation_extraction/base_component.py
@@ -95,11 +95,9 @@ def save(self, save_path: str) -> None:
         self.model.hf_model.resize_token_embeddings(self.tokenizer.get_size()) # type: ignore
 
         assert self.model_config is not None
-        self.model_config.vocab_size = self.tokenizer.get_size()
-        self.model_config.pad_token_id = self.pad_id
-
-        self.model_config.to_json_file(
-            os.path.join(save_path, "model_config.json"))
+        self.model_config.hf_model_config.vocab_size = self.tokenizer.get_size()
+        self.model_config.hf_model_config.pad_token_id = self.pad_id
+        self.model_config.save(save_path)
 
         save_state(self.model, optimizer=self.optimizer, scheduler=self.scheduler, epoch=self.epoch, best_f1=self.best_f1,
                    path=save_path, model_name=self.relcat_config.general.model_name,
@@ -116,16 +114,17 @@ def load(cls, pretrained_model_name_or_path: str = "./") -> "BaseComponent_Relat
         """
 
         relcat_config = ConfigRelCAT.load(load_path=pretrained_model_name_or_path)
+
         model_config = BaseConfig_RelationExtraction.load(pretrained_model_name_or_path=pretrained_model_name_or_path,
                                                          relcat_config=relcat_config)
 
+        tokenizer = BaseTokenizerWrapper_RelationExtraction.load(tokenizer_path=pretrained_model_name_or_path,
+                                                                 relcat_config=relcat_config)
+
         model = BaseModel_RelationExtraction.load(pretrained_model_name_or_path=pretrained_model_name_or_path,
                                                  model_config=model_config,
                                                  relcat_config=relcat_config)
 
-        tokenizer = BaseTokenizerWrapper_RelationExtraction.load(tokenizer_path=pretrained_model_name_or_path,
-                                                                 relcat_config=relcat_config)
-
         model.hf_model.resize_token_embeddings(len(tokenizer.hf_tokenizers)) # type: ignore
 
         optimizer = None # type: ignore
@@ -134,7 +133,7 @@ def load(cls, pretrained_model_name_or_path: str = "./") -> "BaseComponent_Relat
         epoch, best_f1 = load_state(model, optimizer, scheduler, path=pretrained_model_name_or_path,
                                                     model_name=relcat_config.general.model_name,
                                                     file_prefix=relcat_config.general.task,
-                                                    config=relcat_config)
+                                                    relcat_config=relcat_config)
 
         component = cls(model=model, tokenizer=tokenizer, model_config=model_config, config=relcat_config)
         cls.epoch = epoch
diff --git a/medcat/utils/relation_extraction/bert/config.py b/medcat/utils/relation_extraction/bert/config.py
@@ -17,11 +17,10 @@ class BertConfig_RelationExtraction(BaseConfig_RelationExtraction):
     @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)
-        model_config_path = os.path.join(pretrained_model_name_or_path, "model_config.json")
 
-        if pretrained_model_name_or_path and os.path.exists(model_config_path):
-            model_config.hf_model_config = BertConfig.from_json_file(model_config_path)
-            logger.info("Loaded config from file: " + model_config_path)
+        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)
+            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(
diff --git a/medcat/utils/relation_extraction/bert/model.py b/medcat/utils/relation_extraction/bert/model.py
@@ -42,7 +42,7 @@ def __init__(self, pretrained_model_name_or_path: str, relcat_config: ConfigRelC
         self.model_config: Union[BaseConfig_RelationExtraction, BertConfig_RelationExtraction] = model_config
         self.pretrained_model_name_or_path: str = pretrained_model_name_or_path
 
-        self.hf_model: PreTrainedModel = PreTrainedModel(model_config) # type: ignore
+        self.hf_model: Union[BertModel, PreTrainedModel] = BertModel(model_config.hf_model_config) # type: ignore
 
         for param in self.hf_model.parameters(): # type: ignore
             if self.relcat_config.model.freeze_layers:
diff --git a/medcat/utils/relation_extraction/config.py b/medcat/utils/relation_extraction/config.py
@@ -40,19 +40,20 @@ def load(cls, pretrained_model_name_or_path: str, relcat_config: ConfigRelCAT, *
             if "modern-bert" in relcat_config.general.tokenizer_name or \
                "modern-bert" in relcat_config.general.model_name:
                 from medcat.utils.relation_extraction.modernbert.config import ModernBertConfig_RelationExtraction
-                model_config = ModernBertConfig_RelationExtraction.load(model_config_path, **kwargs)
+                model_config = ModernBertConfig_RelationExtraction.load(model_config_path, relcat_config=relcat_config, **kwargs)
             elif "bert" in relcat_config.general.tokenizer_name or \
                "bert" in relcat_config.general.model_name:
                 from medcat.utils.relation_extraction.bert.config import BertConfig_RelationExtraction
-                model_config = BertConfig_RelationExtraction.load(model_config_path, **kwargs)
+                model_config = BertConfig_RelationExtraction.load(model_config_path, relcat_config=relcat_config, **kwargs)
             elif "llama" in relcat_config.general.tokenizer_name or \
                "llama" in relcat_config.general.model_name:
                 from medcat.utils.relation_extraction.llama.config import LlamaConfig_RelationExtraction
-                model_config = LlamaConfig_RelationExtraction.load(model_config_path, **kwargs)
+                model_config = LlamaConfig_RelationExtraction.load(model_config_path, relcat_config=relcat_config, **kwargs)
         else:
             if pretrained_model_name_or_path:
                 model_config.hf_model_config = PretrainedConfig.from_pretrained(pretrained_model_name_or_path=pretrained_model_name_or_path, **kwargs)
             else:
                 model_config.hf_model_config = PretrainedConfig.from_pretrained(pretrained_model_name_or_path=relcat_config.general.model_name, **kwargs)
             logger.info("Loaded config from : " + model_config_path)
+
         return model_config
diff --git a/medcat/utils/relation_extraction/llama/config.py b/medcat/utils/relation_extraction/llama/config.py
@@ -17,14 +17,13 @@ class LlamaConfig_RelationExtraction(BaseConfig_RelationExtraction):
     @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)
-        model_config_path = os.path.join(pretrained_model_name_or_path, "model_config.json")
 
-        if pretrained_model_name_or_path and os.path.exists(model_config_path):
-            model_config.model_config = LlamaConfig.from_json_file(model_config_path)
-            logger.info("Loaded config from file: " + model_config_path)
+        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)
+            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.model_config = LlamaConfig.from_pretrained(
+            model_config.hf_model_config = 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)
 
diff --git a/medcat/utils/relation_extraction/ml_utils.py b/medcat/utils/relation_extraction/ml_utils.py
@@ -105,7 +105,7 @@ def save_state(model, optimizer: torch.optim.AdamW, scheduler: torch.optim.lr_sc
         If you want to export the model after training set final_export=True and leave is_checkpoint=False.
 
     Args:
-        model (Base_RelationExtraction): BertModel_RelationExtraction | LlamaModel_RelationExtraction
+        model (BaseModel_RelationExtraction): BertModel_RelationExtraction | LlamaModel_RelationExtraction etc.
         optimizer (torch.optim.AdamW, optional): Defaults to None.
         scheduler (torch.optim.lr_scheduler.MultiStepLR, optional): Defaults to None.
         epoch (int): Defaults to None.
@@ -136,11 +136,11 @@ def save_state(model, optimizer: torch.optim.AdamW, scheduler: torch.optim.lr_sc
         }, os.path.join(path, file_name))
 
 
-def load_state(model, optimizer, scheduler, path="./", model_name="BERT", file_prefix="train", load_best=False, config: ConfigRelCAT = ConfigRelCAT()) -> Tuple[int, int]:
+def load_state(model, optimizer, scheduler, path: str = "./", model_name:str = "BERT", file_prefix:str = "train", load_best: bool = False, relcat_config: ConfigRelCAT = ConfigRelCAT()) -> Tuple[int, int]:
     """ Used by RelCAT.load() and RelCAT.train()
 
     Args:
-        model (Base_RelationExtraction): BertModel_RelationExtraction | LlamaModel_RelationExtraction, it has to be initialized before calling this method via (Bert/Llama)Model_RelationExtraction(...)
+        model (BaseModel_RelationExtraction): BaseModel_RelationExtraction, it has to be initialized before calling this method via (Bert/Llama)Model_RelationExtraction(...)
         optimizer (_type_): optimizer
         scheduler (_type_): scheduler
         path (str, optional): Defaults to "./".
@@ -153,7 +153,7 @@ def load_state(model, optimizer, scheduler, path="./", model_name="BERT", file_p
         Tuple (int, int): last epoch and f1 score.
     """
 
-    device: torch.device =torch.device(config.general.device)
+    device: torch.device =torch.device(relcat_config.general.device)
 
     model_name = model_name.replace("/", "_")
     logging.info("Attempting to load RelCAT model on device: " + str(device))
@@ -178,13 +178,13 @@ def load_state(model, optimizer, scheduler, path="./", model_name="BERT", file_p
 
         if optimizer is None:
             parameters = filter(lambda p: p.requires_grad, model.parameters())
-            optimizer = torch.optim.AdamW(params=parameters, lr=config.train.lr, weight_decay=config.train.adam_weight_decay,
-                                betas=config.train.adam_betas, eps=config.train.adam_epsilon)
+            optimizer = torch.optim.AdamW(params=parameters, lr=relcat_config.train.lr, weight_decay=relcat_config.train.adam_weight_decay,
+                                betas=relcat_config.train.adam_betas, eps=relcat_config.train.adam_epsilon)
 
         if scheduler is None:
             scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
-                                                             milestones=config.train.multistep_milestones,
-                                                             gamma=config.train.multistep_lr_gamma)
+                                                             milestones=relcat_config.train.multistep_milestones,
+                                                             gamma=relcat_config.train.multistep_lr_gamma)
         optimizer.load_state_dict(checkpoint['optimizer'])
         scheduler.load_state_dict(checkpoint['scheduler'])
         logging.info("Loaded model and optimizer.")
diff --git a/medcat/utils/relation_extraction/models.py b/medcat/utils/relation_extraction/models.py
@@ -98,19 +98,26 @@ def __init__(self, relcat_config: ConfigRelCAT,
         self.hf_model: Union[ModernBertModel, BertModel, LlamaModel, PreTrainedModel] = PreTrainedModel(config=model_config.hf_model_config) # type: ignore
         self.pretrained_model_name_or_path: str = pretrained_model_name_or_path
 
+        self._reinitialize_dense_and_frozen_layers(relcat_config=relcat_config)
+
+        self.log.info("RelCAT model config: " + str(self.model_config.hf_model_config))
+
+    def _reinitialize_dense_and_frozen_layers(self, relcat_config: ConfigRelCAT) -> None:
+        """ Reinitialize the dense layers of the model
+
+        Args:
+            relcat_config (ConfigRelCAT): relcat config.
+        """
+
+        self.drop_out = nn.Dropout(relcat_config.model.dropout)
+        self.fc1, self.fc2, self.fc3 = create_dense_layers(relcat_config)
+
         for param in self.hf_model.parameters(): # type: ignore
             if self.relcat_config.model.freeze_layers:
                 param.requires_grad = False
             else:
                 param.requires_grad = True
 
-        self.drop_out = nn.Dropout(self.relcat_config.model.dropout)
-
-        # dense layers
-        self.fc1, self.fc2, self.fc3 = create_dense_layers(self.relcat_config)
-
-        self.log.info("RelCAT model config: " + str(self.model_config.hf_model_config))
-
     def forward(self,
                 input_ids: Optional[torch.Tensor] = None,
                 attention_mask: Optional[torch.Tensor] = None,
@@ -250,4 +257,7 @@ def load(cls, pretrained_model_name_or_path: str, relcat_config: ConfigRelCAT, m
                 cls.log.info("Loaded model from relcat_config: " + relcat_config.general.model_name)
 
         cls.log.info("Loaded " + str(model.__class__.__name__) + " from pretrained_model_name_or_path: " + pretrained_model_name_or_path)
+
+        model._reinitialize_dense_and_frozen_layers(relcat_config=relcat_config)
+
         return model
diff --git a/medcat/utils/relation_extraction/modernbert/config.py b/medcat/utils/relation_extraction/modernbert/config.py
@@ -17,15 +17,13 @@ class ModernBertConfig_RelationExtraction(BaseConfig_RelationExtraction):
     @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)
-        model_config_path = os.path.join(pretrained_model_name_or_path, "model_config.json")
 
-
-        if pretrained_model_name_or_path and os.path.exists(model_config_path):
-            model_config.model_config = ModernBertConfig.from_json_file(model_config_path)
-            logger.info("Loaded config from file: " + model_config_path)
+        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)
+            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.model_config = ModernBertConfig.from_pretrained(
+            model_config.hf_model_config = 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)
 
diff --git a/medcat/utils/relation_extraction/modernbert/model.py b/medcat/utils/relation_extraction/modernbert/model.py
@@ -40,7 +40,7 @@ def __init__(self, pretrained_model_name_or_path: str, relcat_config: ConfigRelC
         self.model_config: Union[BaseConfig_RelationExtraction, ModernBertConfig_RelationExtraction] = model_config
         self.pretrained_model_name_or_path: str = pretrained_model_name_or_path
 
-        self.hf_model: Union[ModernBertModel, PreTrainedModel] = PreTrainedModel(config=model_config.hf_model_config)
+        self.hf_model: Union[ModernBertModel, PreTrainedModel] = ModernBertModel(config=model_config.hf_model_config)
 
         for param in self.hf_model.parameters(): # type: ignore
             if self.relcat_config.model.freeze_layers:
diff --git a/tests/test_rel_cat.py b/tests/test_rel_cat.py
@@ -7,7 +7,7 @@
 from medcat.cdb import CDB
 from medcat.config_rel_cat import ConfigRelCAT
 from medcat.rel_cat import RelCAT
-from medcat.utils.relation_extraction.bert.tokenizer import BaseTokenizerWrapper_RelationExtraction as TokenizerWrapperBERT
+from medcat.utils.relation_extraction.bert.tokenizer import BaseTokenizerWrapper_RelationExtraction
 from medcat.utils.relation_extraction.rel_dataset import RelData
 
 from transformers.models.auto.tokenization_auto import AutoTokenizer
@@ -29,7 +29,7 @@ def setUpClass(cls) -> None:
         config.model.model_size = 2304
         config.general.log_level = logging.DEBUG
 
-        tokenizer = TokenizerWrapperBERT(AutoTokenizer.from_pretrained(
+        tokenizer = BaseTokenizerWrapper_RelationExtraction(AutoTokenizer.from_pretrained(
             pretrained_model_name_or_path=config.general.model_name,
             config=config), add_special_tokens=True)
 
@@ -55,10 +55,10 @@ def setUpClass(cls) -> None:
             cls.mct_file_test = json.loads(f.read())["projects"][0]["documents"][1]
 
         cls.config_rel_cat: ConfigRelCAT = config
-        cls.rel_cat: RelCAT = RelCAT(cdb, # tokenizer=tokenizer,
-                                     config=config, init_model=True,)
+        cls.rel_cat: RelCAT = RelCAT(cdb, config=config, init_model=True)
 
         cls.rel_cat.component.model.hf_model.resize_token_embeddings(len(tokenizer.hf_tokenizers))
+        cls.rel_cat.component.model_config.hf_model_config.vocab_size = tokenizer.get_size()
 
         cls.finished = False
         cls.tokenizer = tokenizer
@@ -88,7 +88,6 @@ def test_dataset_relation_parser(self) -> None:
 
         self.assertEqual(len(rels), len(samples))
 
-
     def test_train_csv_no_tags(self) -> None:
         self.rel_cat.component.relcat_config.train.epochs = 2
         self.rel_cat.train(train_csv_path=self.medcat_rels_csv_path_train, test_csv_path=self.medcat_rels_csv_path_test, checkpoint_path=self.tmp_dir)
@@ -100,8 +99,6 @@ def test_train_mctrainer(self) -> None:
         self.rel_cat.component.relcat_config.general.addl_rels_max_sample_size = 10
         self.rel_cat.component.relcat_config.train.test_size = 0.1
         self.rel_cat.component.relcat_config.train.nclasses = 3
-        self.rel_cat.component.model.relcat_config.train.nclasses = 3
-        self.rel_cat.component.model.hf_model.resize_token_embeddings(len(self.tokenizer.hf_tokenizers))
 
         self.rel_cat.train(export_data_path=self.medcat_export_with_rels_path, checkpoint_path=self.tmp_dir)
 
