BERT-based encoder for queries or blocks used for learned information retrieval.
| 145 | |
| 146 | |
| 147 | class IREncoderBertModel(MegatronModule): |
| 148 | """BERT-based encoder for queries or blocks used for learned information retrieval.""" |
| 149 | def __init__(self, ict_head_size, num_tokentypes=2, parallel_output=True): |
| 150 | super(IREncoderBertModel, self).__init__() |
| 151 | args = get_args() |
| 152 | |
| 153 | self.ict_head_size = ict_head_size |
| 154 | self.parallel_output = parallel_output |
| 155 | init_method = init_method_normal(args.init_method_std) |
| 156 | scaled_init_method = scaled_init_method_normal(args.init_method_std, |
| 157 | args.num_layers) |
| 158 | |
| 159 | self.language_model, self._language_model_key = get_language_model( |
| 160 | attention_mask_func=bert_attention_mask_func, |
| 161 | num_tokentypes=num_tokentypes, |
| 162 | add_pooler=True, |
| 163 | init_method=init_method, |
| 164 | scaled_init_method=scaled_init_method) |
| 165 | |
| 166 | self.ict_head = get_linear_layer(args.hidden_size, ict_head_size, init_method) |
| 167 | self._ict_head_key = 'ict_head' |
| 168 | |
| 169 | def forward(self, input_ids, attention_mask, tokentype_ids=None): |
| 170 | extended_attention_mask = bert_extended_attention_mask( |
| 171 | attention_mask, next(self.language_model.parameters()).dtype) |
| 172 | position_ids = bert_position_ids(input_ids) |
| 173 | |
| 174 | lm_output, pooled_output = self.language_model( |
| 175 | input_ids, |
| 176 | position_ids, |
| 177 | extended_attention_mask, |
| 178 | tokentype_ids=tokentype_ids) |
| 179 | |
| 180 | # Output. |
| 181 | ict_logits = self.ict_head(pooled_output) |
| 182 | return ict_logits, None |
| 183 | |
| 184 | def state_dict_for_save_checkpoint(self, destination=None, prefix='', |
| 185 | keep_vars=False): |
| 186 | """For easy load when model is combined with other heads, |
| 187 | add an extra key.""" |
| 188 | |
| 189 | state_dict_ = {} |
| 190 | state_dict_[self._language_model_key] \ |
| 191 | = self.language_model.state_dict_for_save_checkpoint( |
| 192 | destination, prefix, keep_vars) |
| 193 | state_dict_[self._ict_head_key] \ |
| 194 | = self.ict_head.state_dict(destination, prefix, keep_vars) |
| 195 | return state_dict_ |
| 196 | |
| 197 | def load_state_dict(self, state_dict, strict=True): |
| 198 | """Customized load.""" |
| 199 | self.language_model.load_state_dict( |
| 200 | state_dict[self._language_model_key], strict=strict) |
| 201 | self.ict_head.load_state_dict( |
| 202 | state_dict[self._ict_head_key], strict=strict) |
| 203 | |
| 204 |