|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
|
|
