Build language model and return along with the key to save.
(
num_tokentypes,
add_pooler,
init_method=None,
scaled_init_method=None,
)
| 71 | |
| 72 | |
| 73 | def get_language_model( |
| 74 | num_tokentypes, |
| 75 | add_pooler, |
| 76 | init_method=None, |
| 77 | scaled_init_method=None, |
| 78 | ): |
| 79 | """Build language model and return along with the key to save.""" |
| 80 | args = get_args() |
| 81 | |
| 82 | if init_method is None: |
| 83 | init_method = init_method_normal(args.init_method_std) |
| 84 | |
| 85 | if scaled_init_method is None: |
| 86 | scaled_init_method = scaled_init_method_normal(args.init_method_std, args.num_layers) |
| 87 | |
| 88 | # Language model. |
| 89 | language_model = TransformerLanguageModel( |
| 90 | init_method=init_method, |
| 91 | output_layer_init_method=scaled_init_method, |
| 92 | num_tokentypes=num_tokentypes, |
| 93 | add_pooler=add_pooler) |
| 94 | # key used for checkpoints. |
| 95 | language_model_key = 'language_model' |
| 96 | |
| 97 | return language_model, language_model_key |
| 98 | |
| 99 | |
| 100 | class Embedding(MegatronModule): |
no test coverage detected