RuntimeError when initializing StigmaBertModel and a quick fix

[mosalen]

Thank you so much for the great work!

For anyone who encountered the following error when initializing StigmaBertModel:
Traceback (most recent call last):
File "C:\Users\Louis\miniconda3\envs\SL-infer\lib\code.py", line 90, in runcode
exec(code, self.locals)
File "", line 2, in
File "C:\Users\Louis\PycharmProjects\pythonProject5\ehr-stigma-main\stigma\api.py", line 381, in init
_ = self._initialize_model(self._model)
File "C:\Users\Louis\PycharmProjects\pythonProject5\ehr-stigma-main\stigma\api.py", line 485, in _initialize_model
_ = self.model_dict["classifier"].load_state_dict(torch.load(f"{model}/model.pth", map_location=torch.device('cpu')))
File "C:\Users\Louis\miniconda3\envs\SL-infer\lib\site-packages\torch\nn\modules\module.py", line 2152, in load_state_dict
raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for BERTMultitaskClassifier:
Unexpected key(s) in state_dict: "_bert._bert.embeddings.position_ids".

A quick fix that works for me is to add strict=False at line #485 of the stigma/api.py, i.e.:
_ = self.model_dict["classifier"].load_state_dict(torch.load(f"{model}/model.pth", map_location=torch.device('cpu')), strict=False)

I hope this helps! Thanks!

[kharrigian]

Thanks so much for bringing this error to my attention! It seems like the issue is related to versions of either torch or transformers.

I'll implement your suggested fix shortly, but want to make sure I can reproduce the issue first so I can confirm that versioning is the issue.

Can you let me know which version of python, torch and transformers you have installed in the environment where the error occurs? Also, does this occur when running the example code block, or only within a different block of code?

[mosalen]

Thank you so much for your attention.
I followed your demo notebook and the repository to set up the environment. Version info: Python 3.10, torch 2.1.2, transformers 4.36.2
The error will arise when initializing the BERT model alone, no errors when initializing baseline models:

Initialize BERT Model (Note the alternative option for specifying model parameters)

bert_model = StigmaBertModel(model=settings.MODELS["mimic-iv-discharge_clinical-bert"]["tasks"][keyword_category],
tokenizer="emilyalsentzer/Bio_ClinicalBERT",
keyword_category=keyword_category,
batch_size=8,
device="cpu")

I believe this is just a minor and common issue if we install the latest version of transformers. Related post: openai/gpt-2-output-dataset#35

[kharrigian]

Thanks! I was able to reproduce the issue with that added information and have pushed the hot fix you suggested.