| import os | |
| import yaml | |
| import torch | |
| from transformers import AlbertConfig, AlbertModel | |
| class CustomAlbert(AlbertModel): | |
| def forward(self, *args, **kwargs): | |
| # Call the original forward method | |
| outputs = super().forward(*args, **kwargs) | |
| # Only return the last_hidden_state | |
| return outputs.last_hidden_state | |
| def load_plbert(log_dir): | |
| config_path = os.path.join(log_dir, "config.yml") | |
| plbert_config = yaml.safe_load(open(config_path)) | |
| albert_base_configuration = AlbertConfig(**plbert_config['model_params']) | |
| checkpoint_path = os.path.join(log_dir, plbert_config['checkpoint_path']) | |
| bert = CustomAlbert(albert_base_configuration) | |
| checkpoint = torch.load(checkpoint_path, map_location='cpu') | |
| state_dict = checkpoint['net'] | |
| from collections import OrderedDict | |
| new_state_dict = OrderedDict() | |
| for k, v in state_dict.items(): | |
| name = k[7:] # remove `module.` | |
| if name.startswith('encoder.'): | |
| name = name[8:] # remove `encoder.` | |
| new_state_dict[name] = v | |
| # if new_state_dict['embeddings.position_ids'] is not None: | |
| # del new_state_dict["embeddings.position_ids"] | |
| bert.load_state_dict(new_state_dict, strict=False) | |
| return bert | |