Build the model.
(args, model_type=None, multi_token=True, num_labels=None, spell_length=None)
| 60 | |
| 61 | |
| 62 | def get_model(args, model_type=None, multi_token=True, num_labels=None, spell_length=None): |
| 63 | """Build the model.""" |
| 64 | print_rank_0('building GLM model ...') |
| 65 | if args.pretrained_bert: |
| 66 | if model_type == "multiple_choice": |
| 67 | model = BertForMultipleChoice.from_pretrained(args.tokenizer_model_type, |
| 68 | cache_dir=args.cache_dir, |
| 69 | fp32_layernorm=args.fp32_layernorm, |
| 70 | fp32_embedding=args.fp32_embedding, |
| 71 | layernorm_epsilon=args.layernorm_epsilon) |
| 72 | elif model_type == "classification": |
| 73 | model = BertForSequenceClassification.from_pretrained(args.tokenizer_model_type, |
| 74 | cache_dir=args.cache_dir, |
| 75 | fp32_layernorm=args.fp32_layernorm, |
| 76 | fp32_embedding=args.fp32_embedding, |
| 77 | layernorm_epsilon=args.layernorm_epsilon, |
| 78 | num_labels=num_labels) |
| 79 | else: |
| 80 | raise NotImplementedError |
| 81 | else: |
| 82 | output_predict, paralle_output = True, True |
| 83 | if (model_type == "multiple_choice" or model_type == "classification") and not args.cloze_eval: |
| 84 | output_predict = False |
| 85 | if model_type is not None: |
| 86 | paralle_output = False |
| 87 | if spell_length is not None: |
| 88 | print_rank_0(f"Continuous spell length {spell_length}") |
| 89 | model = GLMModel(num_layers=args.num_layers, |
| 90 | vocab_size=args.vocab_size, |
| 91 | hidden_size=args.hidden_size, |
| 92 | num_attention_heads=args.num_attention_heads, |
| 93 | embedding_dropout_prob=args.hidden_dropout, |
| 94 | attention_dropout_prob=args.attention_dropout, |
| 95 | output_dropout_prob=args.hidden_dropout, |
| 96 | max_sequence_length=args.max_position_embeddings, |
| 97 | max_memory_length=args.mem_length, |
| 98 | checkpoint_activations=args.checkpoint_activations, |
| 99 | checkpoint_num_layers=args.checkpoint_num_layers, |
| 100 | parallel_output=paralle_output, |
| 101 | relative_encoding=args.transformer_xl, |
| 102 | block_position_encoding=args.block_lm and not args.masked_lm, |
| 103 | output_predict=output_predict, |
| 104 | spell_length=spell_length, |
| 105 | spell_func=args.prompt_func, |
| 106 | attention_scale=args.attention_scale) |
| 107 | if args.freeze_transformer: |
| 108 | model.freeze_transformer(tune_prefix_layers=args.tune_prefix_layers) |
| 109 | if model_type is not None: |
| 110 | if model_type == 'multiple_choice': |
| 111 | if args.cloze_eval: |
| 112 | if multi_token: |
| 113 | if args.fast_decode: |
| 114 | model = GLMForMultiTokenClozeFast(model, length_penalty=args.length_penalty) |
| 115 | else: |
| 116 | model = GLMForMultiTokenCloze(model, length_penalty=args.length_penalty) |
| 117 | else: |
| 118 | model = GLMForSingleTokenCloze(model, take_softmax=args.adapet) |
| 119 | else: |
no test coverage detected