| from pathlib import PosixPath |
| from typing import Union, Optional |
| import torch |
| from transformers import ( |
| RobertaConfig, |
| RobertaTokenizerFast, |
| RobertaForMaskedLM, |
| RobertaForSequenceClassification, |
| ) |
|
|
| from .models import ( |
| RobertaMeanPoolConfig, |
| RobertaForSequenceClassificationMeanPool, |
| ) |
|
|
| RobertaSettings = dict( |
| padding_side='left' |
| ) |
|
|
|
|
| MODELS = { |
| "roberta-lm": (RobertaConfig, RobertaTokenizerFast, RobertaForMaskedLM, RobertaSettings), |
| "roberta-pred": (RobertaConfig, RobertaTokenizerFast, RobertaForSequenceClassification, RobertaSettings), |
| "roberta-pred-mean-pool": (RobertaMeanPoolConfig, RobertaTokenizerFast, RobertaForSequenceClassificationMeanPool, RobertaSettings) |
| } |
|
|
|
|
| def load_model(model_name: str, |
| tokenizer_dir: Union[str, PosixPath], |
| max_tokenized_len: int = 254, |
| pretrained_model: Union[str, PosixPath] = None, |
| k: Optional[int] = None, |
| do_lower_case: Optional[bool] = None, |
| padding_side: Optional[str] = 'left', |
| **config_settings) -> tuple: |
| """Load specified model, config, and tokenizer. |
| |
| Args: |
| model_name (str): Name of model. Acceptable options are |
| - 'roberta-lm', |
| - 'roberta-pred', |
| - 'roberta-pred-mean-pool' |
| tokenizer_dir (Union[str, PosixPath]): Directory containing tokenizer |
| files: merges.txt and vocab.txt |
| max_len (int, optional): Maximum tokenized length, |
| not including SOS and EOS. Defaults to 254. |
| pretrained_model (Union[str, PosixPath], optional): path to saved |
| pretrained RoBERTa transformer model. Defaults to None. |
| k (Optional[int], optional): Size of kmers (for DNABERT model). Defaults to 6. |
| do_lower_case (bool, optional): Whether to convert all inputs to lower case. Defaults to None. |
| padding_side (str, optional): Which side to pad on. Defaults to 'left'. |
| |
| Returns: |
| tuple: config_obj, tokenizer, model |
| """ |
| config_settings = config_settings or {} |
| max_position_embeddings = max_tokenized_len + 2 |
| config_class, tokenizer_class, model_class, tokenizer_settings = MODELS[model_name] |
| |
| kwargs = dict( |
| max_len=max_tokenized_len, |
| truncate=True, |
| padding="max_length", |
| **tokenizer_settings |
| ) |
| if k is not None: |
| kwargs.update(dict(k=k)) |
| if do_lower_case is not None: |
| kwargs.update(dict(do_lower_case=do_lower_case)) |
| if padding_side is not None: |
| kwargs.update(dict(padding_side=padding_side)) |
|
|
| tokenizer = tokenizer_class.from_pretrained(str(tokenizer_dir), **kwargs) |
| name_or_path = str(pretrained_model) or '' |
| config_obj = config_class( |
| vocab_size=len(tokenizer), |
| max_position_embeddings=max_position_embeddings, |
| name_or_path=name_or_path, |
| output_hidden_states=True, |
| **config_settings |
| ) |
| if pretrained_model: |
| |
| model = model_class(config=config_obj) |
| state_dict = torch.load(pretrained_model) |
| state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()} |
| unexpected_keys = [k for k in state_dict.keys() if 'position_ids' in k] |
| for key in unexpected_keys: |
| del state_dict[key] |
| model.load_state_dict(state_dict) |
| else: |
| print("Loading untrained model") |
| model = model_class(config=config_obj) |
| model.resize_token_embeddings(len(tokenizer)) |
| return config_obj, tokenizer, model |
