| |
| |
| |
| |
|
|
|
|
| from typing import List |
|
|
| import torch |
| from fairseq.modules.quant_noise import quant_noise |
| from torch import nn |
|
|
|
|
| class AdaptiveInput(nn.Module): |
| def __init__( |
| self, |
| vocab_size: int, |
| padding_idx: int, |
| initial_dim: int, |
| factor: float, |
| output_dim: int, |
| cutoff: List[int], |
| q_noise: float = 0, |
| qn_block_size: int = 8, |
| ): |
| super().__init__() |
|
|
| if vocab_size > cutoff[-1]: |
| cutoff = cutoff + [vocab_size] |
| else: |
| assert ( |
| vocab_size == cutoff[-1] |
| ), "cannot specify cutoff larger than vocab size" |
|
|
| self.cutoff = cutoff |
| self.embedding_dim = output_dim |
| self.padding_idx = padding_idx |
|
|
| self.embeddings = nn.ModuleList() |
| for i in range(len(self.cutoff)): |
| prev = self.cutoff[i - 1] if i > 0 else 0 |
| size = self.cutoff[i] - prev |
| dim = int(initial_dim // (factor ** i)) |
| seq = nn.Sequential( |
| nn.Embedding(size, dim, self.padding_idx), |
| quant_noise( |
| nn.Linear(dim, output_dim, bias=False), q_noise, qn_block_size |
| ), |
| ) |
|
|
| self.embeddings.append(seq) |
| self.padding_idx = None |
| self.padding_idx = padding_idx |
|
|
| def init_weights(m): |
| if isinstance(m, nn.Embedding): |
| nn.init.normal_(m.weight, mean=0, std=m.weight.shape[1] ** -0.5) |
| nn.init.constant_(m.weight[padding_idx], 0) |
| elif hasattr(m, "weight"): |
| nn.init.xavier_uniform_(m.weight) |
|
|
| self.apply(init_weights) |
|
|
| self.register_buffer("_float_tensor", torch.FloatTensor(1)) |
|
|
| def weights_for_band(self, band: int): |
| return self.embeddings[band][0].weight, self.embeddings[band][1].weight |
|
|
| def forward(self, input: torch.Tensor): |
| result = self._float_tensor.new(input.shape + (self.embedding_dim,)) |
| for i in range(len(self.cutoff)): |
| mask = input.lt(self.cutoff[i]) |
| if i > 0: |
| mask.mul_(input.ge(self.cutoff[i - 1])) |
| chunk_input = input[mask] - self.cutoff[i - 1] |
| else: |
| chunk_input = input[mask] |
| if mask.any(): |
| result[mask] = self.embeddings[i](chunk_input) |
| return result |
|
|
