mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-28 09:08:54 +00:00
Introduce code from Rust monorepo. This is a mess
This commit is contained in:
Brandon Thomas
@@ -0,0 +1,19 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from nemo.collections.nlp.modules.common.megatron.megatron_bert import MegatronBertEncoder
|
||||
from nemo.collections.nlp.modules.common.megatron.megatron_utils import (
|
||||
get_megatron_checkpoint,
|
||||
get_megatron_lm_models_list,
|
||||
)
|
||||
@@ -0,0 +1,216 @@
|
||||
# Copyright 2018 The Google AI Language Team Authors and
|
||||
# The HuggingFace Inc. team.
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import os
|
||||
|
||||
import torch
|
||||
from megatron import get_args, initialize_megatron
|
||||
from megatron.checkpointing import set_checkpoint_version
|
||||
from megatron.model import get_language_model
|
||||
from megatron.model.bert_model import bert_attention_mask_func, bert_extended_attention_mask, bert_position_ids
|
||||
from megatron.mpu import (
|
||||
get_model_parallel_group,
|
||||
model_parallel_is_initialized,
|
||||
set_pipeline_model_parallel_rank,
|
||||
set_pipeline_model_parallel_world_size,
|
||||
)
|
||||
from omegaconf import DictConfig, OmegaConf
|
||||
|
||||
from nemo.collections.nlp.modules.common.bert_module import BertModule
|
||||
from nemo.core.classes import typecheck
|
||||
from nemo.utils import logging
|
||||
from nemo.utils.app_state import AppState
|
||||
|
||||
__all__ = ['MegatronBertEncoder']
|
||||
|
||||
|
||||
def complete_lazy_init(self):
|
||||
# finish megatron-lm initialization
|
||||
if hasattr(self, "_lazy_init_fn") and self._lazy_init_fn is not None:
|
||||
self._lazy_init_fn()
|
||||
self._lazy_init_fn = None
|
||||
|
||||
|
||||
class MegatronBertEncoder(BertModule):
|
||||
"""
|
||||
MegatronBERT wraps around the Megatron Language model
|
||||
from https://github.com/NVIDIA/Megatron-LM
|
||||
|
||||
Args:
|
||||
config_file (str): path to model configuration file.
|
||||
vocab_file (str): path to vocabulary file.
|
||||
tokenizer_type (str): tokenizer type, currently only 'BertWordPieceLowerCase' supported.
|
||||
"""
|
||||
|
||||
def __init__(self, model_name, config, vocab_file, model_parallel_size=None, model_parallel_rank=None):
|
||||
|
||||
super().__init__()
|
||||
|
||||
self._model_parallel_size = model_parallel_size
|
||||
self._model_parallel_rank = model_parallel_rank
|
||||
self._restore_path = None
|
||||
self._app_state = None
|
||||
self._model_name = model_name
|
||||
|
||||
if not os.path.exists(vocab_file):
|
||||
raise ValueError(f'Vocab file not found at {vocab_file}')
|
||||
|
||||
# convert config to dictionary
|
||||
if isinstance(config, DictConfig):
|
||||
config = OmegaConf.to_container(config)
|
||||
config["vocab_file"] = vocab_file
|
||||
config['tokenizer_type'] = 'BertWordPieceLowerCase'
|
||||
config['lazy_mpu_init'] = True
|
||||
config['onnx_safe'] = True
|
||||
|
||||
# if 'model_parallel_size' in config:
|
||||
if self._model_parallel_size is not None:
|
||||
app_state = AppState()
|
||||
self._app_state = app_state
|
||||
|
||||
# must be set for model parallel megatron-lm
|
||||
os.environ["WORLD_SIZE"] = str(app_state.world_size)
|
||||
os.environ["RANK"] = str(self._model_parallel_rank)
|
||||
|
||||
extra_args_provider = self._update_megatron_args(tensor_model_parallel_size=self._model_parallel_size)
|
||||
|
||||
else:
|
||||
extra_args_provider = self._update_megatron_args()
|
||||
|
||||
# configure globals for megatron
|
||||
set_pipeline_model_parallel_rank(0) # pipeline model parallelism not implemented in NeMo
|
||||
set_pipeline_model_parallel_world_size(1) # pipeline model parallelism not implemented in NeMo
|
||||
|
||||
# Initialize part of Megatron global state that is needed for its constructor.
|
||||
# We set 'lazy_mpu_init' flag on to make Megatron do only the initialization that does not depend
|
||||
# on ddp be initialized yet (and we don't want Megatron to initialize DDP itself either)
|
||||
# and to return a hook for us to call after PTL has torch.distributed initialized.
|
||||
# (or if no PTL in case of inference - then we'll initialize torch.distributed)
|
||||
# We call and clear this hook on first call to forward()
|
||||
self._lazy_init_fn = initialize_megatron(
|
||||
extra_args_provider=extra_args_provider, args_defaults=config, ignore_unknown_args=True
|
||||
)
|
||||
|
||||
# read Megatron arguments back
|
||||
args = get_args()
|
||||
logging.info(f'Megatron-lm argparse args: {args}')
|
||||
|
||||
self.language_model, self._language_model_key = get_language_model(
|
||||
attention_mask_func=bert_attention_mask_func, num_tokentypes=2, add_pooler=False
|
||||
)
|
||||
|
||||
self.config = OmegaConf.create(config)
|
||||
# key used for checkpoints
|
||||
self._hidden_size = self.language_model.hidden_size
|
||||
|
||||
def _update_megatron_args(
|
||||
self,
|
||||
micro_batch_size=1,
|
||||
tensor_model_parallel_size=1,
|
||||
scaled_masked_softmax_fusion=False,
|
||||
bias_gelu_fusion=False,
|
||||
bias_dropout_fusion=False,
|
||||
):
|
||||
def extra_args_provider(parser):
|
||||
parser.set_defaults(micro_batch_size=micro_batch_size)
|
||||
parser.set_defaults(tensor_model_parallel_size=tensor_model_parallel_size)
|
||||
parser.set_defaults(scaled_masked_softmax_fusion=scaled_masked_softmax_fusion)
|
||||
parser.set_defaults(bias_gelu_fusion=bias_gelu_fusion)
|
||||
parser.set_defaults(bias_dropout_fusion=bias_dropout_fusion)
|
||||
|
||||
return parser
|
||||
|
||||
return extra_args_provider
|
||||
|
||||
def complete_lazy_init(self):
|
||||
# finish megatron-lm initialization
|
||||
if hasattr(self, "_lazy_init_fn") and self._lazy_init_fn is not None:
|
||||
self._lazy_init_fn()
|
||||
self._lazy_init_fn = None
|
||||
|
||||
@property
|
||||
def hidden_size(self):
|
||||
"""
|
||||
Property returning hidden size.
|
||||
|
||||
Returns:
|
||||
Hidden size.
|
||||
"""
|
||||
return self._hidden_size
|
||||
|
||||
@typecheck()
|
||||
def forward(self, input_ids, attention_mask, token_type_ids):
|
||||
app_state = AppState()
|
||||
if app_state.model_parallel_size is None:
|
||||
self.complete_lazy_init()
|
||||
|
||||
extended_attention_mask = bert_extended_attention_mask(attention_mask)
|
||||
position_ids = bert_position_ids(input_ids)
|
||||
|
||||
sequence_output = self.language_model(
|
||||
input_ids=input_ids,
|
||||
position_ids=position_ids,
|
||||
attention_mask=extended_attention_mask,
|
||||
tokentype_ids=token_type_ids,
|
||||
)
|
||||
return sequence_output
|
||||
|
||||
def _load_checkpoint(self, filename):
|
||||
"""Helper function for loading megatron checkpoints.
|
||||
|
||||
Args:
|
||||
filename (str): Path to megatron checkpoint.
|
||||
"""
|
||||
state_dict = torch.load(filename, map_location='cpu')
|
||||
if 'checkpoint_version' in state_dict:
|
||||
if state_dict['checkpoint_version'] is not None:
|
||||
set_checkpoint_version(state_dict['checkpoint_version'])
|
||||
logging.info(
|
||||
f"Megatron-lm checkpoint version found. Setting checkpoint_version to {state_dict['checkpoint_version']}."
|
||||
)
|
||||
else:
|
||||
logging.warning('Megatron-lm checkpoint version not found. Setting checkpoint_version to 0.')
|
||||
set_checkpoint_version(0)
|
||||
# to load from Megatron pretrained checkpoint
|
||||
if 'model' in state_dict:
|
||||
self.language_model.load_state_dict(state_dict['model'][self._language_model_key])
|
||||
else:
|
||||
self.load_state_dict(state_dict)
|
||||
|
||||
logging.info(f"Checkpoint loaded from from {filename}")
|
||||
|
||||
def restore_weights(self, restore_path: str):
|
||||
"""Restores module/model's weights.
|
||||
For model parallel checkpoints the directory structure
|
||||
should be restore_path/mp_rank_0X/model_optim_rng.pt
|
||||
|
||||
Args:
|
||||
restore_path (str): restore_path should a file or a directory if using model parallel
|
||||
"""
|
||||
self._restore_path = restore_path
|
||||
|
||||
if os.path.isfile(restore_path):
|
||||
self._load_checkpoint(restore_path)
|
||||
elif os.path.isdir(restore_path):
|
||||
# need model parallel groups to restore model parallel checkpoints
|
||||
if model_parallel_is_initialized():
|
||||
model_parallel_rank = torch.distributed.get_rank(group=get_model_parallel_group())
|
||||
mp_restore_path = f'{restore_path}/mp_rank_{model_parallel_rank:02d}/model_optim_rng.pt'
|
||||
self._load_checkpoint(mp_restore_path)
|
||||
else:
|
||||
logging.info(f'torch.distributed not initialized yet. Will not restore model parallel checkpoint')
|
||||
else:
|
||||
logging.error(f'restore_path: {restore_path} must be a file or directory.')
|
||||
@@ -0,0 +1,307 @@
|
||||
# Copyright 2020 The HuggingFace Inc. team.
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import glob
|
||||
import json
|
||||
import os
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
import torch
|
||||
import wget
|
||||
from torch.hub import _get_torch_home
|
||||
|
||||
from nemo.collections.nlp.modules.common.megatron.megatron_bert import MegatronBertEncoder
|
||||
from nemo.utils import AppState, logging
|
||||
|
||||
__all__ = [
|
||||
"get_megatron_lm_model",
|
||||
"get_megatron_lm_models_list",
|
||||
"get_megatron_checkpoint",
|
||||
"is_lower_cased_megatron",
|
||||
"get_megatron_tokenizer",
|
||||
]
|
||||
|
||||
|
||||
torch_home = _get_torch_home()
|
||||
|
||||
if not isinstance(torch_home, str):
|
||||
logging.info("Torch home not found, caching megatron in cwd")
|
||||
torch_home = os.getcwd()
|
||||
|
||||
MEGATRON_CACHE = os.path.join(torch_home, "megatron")
|
||||
|
||||
|
||||
CONFIGS = {"345m": {"hidden_size": 1024, "num_attention_heads": 16, "num_layers": 24, "max_position_embeddings": 512}}
|
||||
|
||||
MEGATRON_CONFIG_MAP = {
|
||||
"megatron-bert-345m-uncased": {
|
||||
"config": CONFIGS["345m"],
|
||||
"checkpoint": "https://api.ngc.nvidia.com/v2/models/nvidia/megatron_bert_345m/versions/v0.0/files/release/mp_rank_00/model_optim_rng.pt",
|
||||
"vocab": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-vocab.txt",
|
||||
"do_lower_case": True,
|
||||
"tokenizer_name": "bert-large-uncased",
|
||||
},
|
||||
"megatron-bert-345m-cased": {
|
||||
"config": CONFIGS["345m"],
|
||||
"checkpoint": "https://api.ngc.nvidia.com/v2/models/nvidia/megatron_bert_345m/versions/v0.1_cased/files/release/mp_rank_00/model_optim_rng.pt",
|
||||
"vocab": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-vocab.txt",
|
||||
"do_lower_case": False,
|
||||
"tokenizer_name": "bert-large-cased",
|
||||
},
|
||||
"megatron-bert-uncased": {
|
||||
"config": None,
|
||||
"checkpoint": None,
|
||||
"vocab": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-vocab.txt",
|
||||
"do_lower_case": True,
|
||||
"tokenizer_name": "bert-large-uncased",
|
||||
},
|
||||
"megatron-bert-cased": {
|
||||
"config": None,
|
||||
"checkpoint": None,
|
||||
"vocab": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-vocab.txt",
|
||||
"do_lower_case": False,
|
||||
"tokenizer_name": "bert-large-cased",
|
||||
},
|
||||
"biomegatron-bert-345m-uncased": {
|
||||
"config": CONFIGS["345m"],
|
||||
"checkpoint": "https://api.ngc.nvidia.com/v2/models/nvidia/biomegatron345muncased/versions/0/files/MegatronBERT.pt",
|
||||
"vocab": "https://api.ngc.nvidia.com/v2/models/nvidia/biomegatron345muncased/versions/0/files/vocab.txt",
|
||||
"do_lower_case": True,
|
||||
"tokenizer_name": "bert-large-uncased",
|
||||
},
|
||||
"biomegatron-bert-345m-cased": {
|
||||
"config": CONFIGS["345m"],
|
||||
"checkpoint": "https://api.ngc.nvidia.com/v2/models/nvidia/biomegatron345mcased/versions/0/files/MegatronBERT.pt",
|
||||
"vocab": "https://api.ngc.nvidia.com/v2/models/nvidia/biomegatron345mcased/versions/0/files/vocab.txt",
|
||||
"do_lower_case": False,
|
||||
"tokenizer_name": "bert-large-cased",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def get_megatron_lm_model(
|
||||
pretrained_model_name: str,
|
||||
config_dict: Optional[dict] = None,
|
||||
config_file: Optional[str] = None,
|
||||
checkpoint_file: Optional[str] = None,
|
||||
vocab_file: Optional[str] = None,
|
||||
) -> Tuple[MegatronBertEncoder, str]:
|
||||
"""
|
||||
Returns MegatronBertEncoder and a default or user specified path to the checkpoint file
|
||||
|
||||
Args:
|
||||
pretrained_mode_name: model name from MEGATRON_CONFIG_MAP
|
||||
for example: megatron-bert-cased
|
||||
config_dict: model configuration parameters
|
||||
config_file: path to model configuration file. Takes precedence over config_dict if both supplied.
|
||||
checkpoint_file: path to checkpoint file or directory if using model parallel.
|
||||
vocab_file: path to vocab file
|
||||
|
||||
Returns:
|
||||
model: MegatronBertEncoder
|
||||
checkpoint_file: path to checkpoint file or directory
|
||||
"""
|
||||
config = None
|
||||
# get default config and checkpoint
|
||||
if config_file:
|
||||
with open(config_file) as f:
|
||||
config = json.load(f)
|
||||
# replace dashes with underscores in config keys
|
||||
fixed_config = {}
|
||||
for key in config.keys():
|
||||
fixed_key = key.replace("-", "_")
|
||||
if fixed_key == 'max_seq_length':
|
||||
fixed_key = 'max_position_embeddings'
|
||||
fixed_config[fixed_key] = config[key]
|
||||
# 'vocab_size" no longer used.
|
||||
if 'vocab_size' in fixed_config:
|
||||
fixed_config.pop('vocab_size')
|
||||
config = fixed_config
|
||||
elif config_dict:
|
||||
config = config_dict
|
||||
elif pretrained_model_name in get_megatron_lm_models_list():
|
||||
config = get_megatron_config(pretrained_model_name)
|
||||
else:
|
||||
raise ValueError(f"{pretrained_model_name} is not supported")
|
||||
|
||||
if config is None:
|
||||
raise ValueError(f"config_file or config_dict is required for {pretrained_model_name}")
|
||||
|
||||
if not checkpoint_file:
|
||||
checkpoint_file = get_megatron_checkpoint(pretrained_model_name)
|
||||
|
||||
if not vocab_file:
|
||||
vocab_file = get_megatron_vocab_file(pretrained_model_name)
|
||||
|
||||
app_state = AppState()
|
||||
if app_state.model_parallel_size is not None and app_state.model_parallel_rank is not None:
|
||||
# model parallel already known from .nemo restore
|
||||
model_parallel_size = app_state.model_parallel_size
|
||||
model_parallel_rank = app_state.model_parallel_rank
|
||||
elif os.path.isdir(checkpoint_file):
|
||||
# starting training from megatron-lm checkpoint
|
||||
mp_ranks = glob.glob(os.path.join(checkpoint_file, 'mp_rank*'))
|
||||
model_parallel_size = len(mp_ranks)
|
||||
app_state.model_parallel_size = model_parallel_size
|
||||
logging.info(
|
||||
(
|
||||
f'restore_path: {checkpoint_file} is a directory. '
|
||||
f'Assuming megatron model parallelism with '
|
||||
f'model_parallel_size: {model_parallel_size}'
|
||||
)
|
||||
)
|
||||
# try to get local rank from global
|
||||
local_rank = None
|
||||
try:
|
||||
local_rank = int(os.environ['LOCAL_RANK'])
|
||||
except:
|
||||
logging.info('Global variable LOCAL_RANK not yet specified')
|
||||
if local_rank is not None:
|
||||
app_state.local_rank = local_rank
|
||||
else:
|
||||
# if local is None then we are on the main process
|
||||
local_rank = 0
|
||||
model_parallel_rank = compute_model_parallel_rank(local_rank, model_parallel_size)
|
||||
app_state.model_parallel_rank = model_parallel_rank
|
||||
else:
|
||||
model_parallel_size = None
|
||||
model_parallel_rank = None
|
||||
|
||||
model = MegatronBertEncoder(
|
||||
model_name=pretrained_model_name,
|
||||
config=config,
|
||||
vocab_file=vocab_file,
|
||||
model_parallel_size=model_parallel_size,
|
||||
model_parallel_rank=model_parallel_rank,
|
||||
)
|
||||
|
||||
return model, checkpoint_file
|
||||
|
||||
|
||||
def compute_model_parallel_rank(local_rank, model_parallel_size):
|
||||
return local_rank % model_parallel_size
|
||||
|
||||
|
||||
def get_megatron_lm_models_list() -> List[str]:
|
||||
"""
|
||||
Returns the list of supported Megatron-LM models
|
||||
"""
|
||||
return list(MEGATRON_CONFIG_MAP.keys())
|
||||
|
||||
|
||||
def get_megatron_config(pretrained_model_name: str) -> Dict[str, int]:
|
||||
"""
|
||||
Returns Megatron-LM model config file
|
||||
|
||||
Args:
|
||||
pretrained_model_name (str): pretrained model name
|
||||
|
||||
Returns:
|
||||
config (dict): contains model configuration: number of hidden layers, number of attention heads, etc
|
||||
"""
|
||||
_check_megatron_name(pretrained_model_name)
|
||||
return MEGATRON_CONFIG_MAP[pretrained_model_name]["config"]
|
||||
|
||||
|
||||
def _check_megatron_name(pretrained_model_name: str) -> None:
|
||||
megatron_model_list = get_megatron_lm_models_list()
|
||||
if pretrained_model_name not in megatron_model_list:
|
||||
raise ValueError(f'For Megatron-LM models, choose from the following list: {megatron_model_list}')
|
||||
|
||||
|
||||
def get_megatron_vocab_file(pretrained_model_name: str) -> str:
|
||||
"""
|
||||
Gets vocabulary file from cache or downloads it
|
||||
|
||||
Args:
|
||||
pretrained_model_name: pretrained model name
|
||||
|
||||
Returns:
|
||||
path: path to the vocab file
|
||||
"""
|
||||
_check_megatron_name(pretrained_model_name)
|
||||
url = MEGATRON_CONFIG_MAP[pretrained_model_name]["vocab"]
|
||||
|
||||
path = os.path.join(MEGATRON_CACHE, pretrained_model_name + "_vocab")
|
||||
path = _download(path, url)
|
||||
return path
|
||||
|
||||
|
||||
def get_megatron_checkpoint(pretrained_model_name: str) -> str:
|
||||
"""
|
||||
Gets checkpoint file from cache or downloads it
|
||||
Args:
|
||||
pretrained_model_name: pretrained model name
|
||||
Returns:
|
||||
path: path to model checkpoint
|
||||
"""
|
||||
_check_megatron_name(pretrained_model_name)
|
||||
url = MEGATRON_CONFIG_MAP[pretrained_model_name]["checkpoint"]
|
||||
path = os.path.join(MEGATRON_CACHE, pretrained_model_name)
|
||||
return _download(path, url)
|
||||
|
||||
|
||||
def _download(path: str, url: str):
|
||||
"""
|
||||
Gets a file from cache or downloads it
|
||||
|
||||
Args:
|
||||
path: path to the file in cache
|
||||
url: url to the file
|
||||
Returns:
|
||||
path: path to the file in cache
|
||||
"""
|
||||
if url is None:
|
||||
return None
|
||||
|
||||
if not os.path.exists(path):
|
||||
master_device = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
|
||||
if not os.path.exists(path):
|
||||
if master_device:
|
||||
os.makedirs(MEGATRON_CACHE, exist_ok=True)
|
||||
logging.info(f"Downloading from {url}")
|
||||
wget.download(url, path)
|
||||
# wait until the master process downloads the file and writes it to the cache dir
|
||||
if torch.distributed.is_initialized():
|
||||
torch.distributed.barrier()
|
||||
|
||||
return path
|
||||
|
||||
|
||||
def is_lower_cased_megatron(pretrained_model_name):
|
||||
"""
|
||||
Returns if the megatron is cased or uncased
|
||||
|
||||
Args:
|
||||
pretrained_model_name (str): pretrained model name
|
||||
Returns:
|
||||
do_lower_cased (bool): whether the model uses lower cased data
|
||||
"""
|
||||
_check_megatron_name(pretrained_model_name)
|
||||
return MEGATRON_CONFIG_MAP[pretrained_model_name]["do_lower_case"]
|
||||
|
||||
|
||||
def get_megatron_tokenizer(pretrained_model_name: str):
|
||||
"""
|
||||
Takes a pretrained_model_name for megatron such as "megatron-bert-cased" and returns the according
|
||||
tokenizer name for tokenizer instantiating.
|
||||
|
||||
Args:
|
||||
pretrained_model_name: pretrained_model_name for megatron such as "megatron-bert-cased"
|
||||
Returns:
|
||||
tokenizer name for tokenizer instantiating
|
||||
"""
|
||||
_check_megatron_name(pretrained_model_name)
|
||||
return MEGATRON_CONFIG_MAP[pretrained_model_name]["tokenizer_name"]
|
||||
Reference in New Issue
Block a user