| 325 | |
| 326 | |
| 327 | class HuggingFaceEncoder(torch.nn.Module): |
| 328 | def __init__(self, layer_num, head_num, head_size, weights=None): |
| 329 | super().__init__() |
| 330 | hidden_dim = head_num * head_size |
| 331 | # TODO(bhsueh) The implementation of hidden_act='gelu' is different to FT's (and google BERT) implementation |
| 332 | # FT's implementation is equivalent to hidden_act='gelu_new', but there are some issues for int8 sparse under gelu_new |
| 333 | conf = BertConfig(hidden_size=hidden_dim, intermediate_size=4 * hidden_dim, |
| 334 | num_attention_heads=head_num, num_hidden_layers=layer_num, hidden_act='gelu') |
| 335 | self.encoder = BertEncoder(conf) |
| 336 | w = {} |
| 337 | for k, v in weights.weights.items(): |
| 338 | if k.startswith('bert.encoder') and not k.endswith('_amax'): |
| 339 | w[k[13:]] = weights.weights[k] |
| 340 | self.encoder.load_state_dict(w) |
| 341 | self.head_mask = [None] * layer_num |
| 342 | |
| 343 | def forward(self, hidden_states, attention_mask): |
| 344 | extended_attention_mask = (1.0 - attention_mask) * -10000.0 |
| 345 | output = self.encoder(hidden_states, extended_attention_mask, self.head_mask) |
| 346 | return output |