Spaces:
Runtime error
Runtime error
| # Copyright (c) Facebook, Inc. and its affiliates. | |
| # | |
| # This source code is licensed under the MIT license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| import torch | |
| from fairseq import utils | |
| from fairseq.modules import TransformerEncoderLayer | |
| from .multihead_linear_attention import MultiheadLinearAttention | |
| class LinformerTransformerEncoderLayer(TransformerEncoderLayer): | |
| """ | |
| Implements a Linformer Encoder Layer used in BERT/XLM style pre-trained | |
| models. | |
| """ | |
| def __init__(self, args, shared_compress_layer): | |
| # wrap in a list so it's not automatically registered by PyTorch | |
| self.shared_compress_layer = [shared_compress_layer] | |
| super().__init__(args) | |
| self.register_buffer("version", torch.tensor(2)) | |
| def build_self_attention(self, embed_dim, args): | |
| return MultiheadLinearAttention( | |
| embed_dim, | |
| args.encoder_attention_heads, | |
| dropout=args.dropout, | |
| self_attention=True, | |
| q_noise=args.quant_noise_pq, | |
| qn_block_size=args.quant_noise_pq_block_size, | |
| compressed=args.compressed, | |
| max_seq_len=args.max_positions, | |
| shared_kv_compressed=args.shared_kv_compressed, | |
| shared_compress_layer=self.shared_compress_layer[0], | |
| freeze_compress=args.freeze_compress, | |
| ) | |
| def upgrade_state_dict_named(self, state_dict, name): | |
| super().upgrade_state_dict_named(state_dict, name) | |
| prefix = name + "." if name != "" else "" | |
| # some old checkpoints had weight sharing implemented incorrectly | |
| # (note: this was correct in the original paper code) | |
| if utils.item(state_dict.get(f"{prefix}version", torch.tensor(1))) < 2: | |
| state_dict[f"{prefix}version"] = torch.tensor(1) | |
| # check compression layer sharing | |
| if f"{prefix}shared_compress_layer.weight" in state_dict: | |
| # reinitialize block without sharing compression layer to match | |
| # old behavior | |
| self.shared_compress_layer = [ | |
| torch.nn.Linear( | |
| self.shared_compress_layer[0].weight.size(1), | |
| self.shared_compress_layer[0].weight.size(0), | |
| ) | |
| ] | |
| self.self_attn = self.build_self_attention(self.embed_dim, self.args) | |
| # delete shared_compress_layer, since it's already copied to | |
| # self_attn.compress_k.weight | |
| del state_dict[f"{prefix}shared_compress_layer.weight"] | |
| if f"{prefix}shared_compress_layer.bias" in state_dict: | |
| del state_dict[f"{prefix}shared_compress_layer.bias"] | |