Spaces:
Sleeping
Sleeping
from typing import Union | |
import torch | |
from torch import nn, Tensor | |
class LearnedPositionEmbeddings(nn.Module): | |
def __init__(self, seq_len, model_dim, init=.02): | |
super().__init__() | |
self.emb = nn.Embedding(seq_len, model_dim) | |
# Initializing this way is standard for GPT-2 | |
self.emb.weight.data.normal_(mean=0.0, std=init) | |
def forward(self, x): | |
""" | |
Returns positional embeddings for index 0 up to the length of x | |
""" | |
sl = x.shape[1] | |
return self.emb(torch.arange(0, sl, device=x.device)) | |
def get_fixed_embedding(self, idx: 'Union[int, Tensor]'): | |
""" | |
Args: | |
idx: scalar int or an integer tensor of shape (T,) or (B, T) | |
Returns: | |
positional embeddings for given indices, shape (B, T, dim), ie (1, 1, dim) for int input | |
""" | |
device = self.emb.weight.device | |
idx = idx.to(device) if torch.is_tensor(idx) else torch.tensor(idx, device=device) | |
idx = torch.atleast_2d(idx) | |
assert idx.ndim == 2 | |
return self.emb(idx) # (B, T, dim) | |