Spaces:
Running on Zero
Running on Zero
File size: 912 Bytes
9cde5e3 98d1e65 9cde5e3 98d1e65 9cde5e3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | """
Positional embedding layer.
Converts token positions into learnable embedding vectors.
"""
import torch
import torch.nn as nn
class PositionalEmbedding(nn.Module):
"""
Learned positional embeddings.
Input:
token_ids: (batch_size, context_length)
Output:
(batch_size, context_length, d_model)
"""
def __init__(self, context_length: int, d_model: int):
super().__init__()
self.embedding = nn.Embedding(
num_embeddings=context_length,
embedding_dim=d_model
)
def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
batch_size, context_length = token_ids.shape
positions = torch.arange(
context_length,
device=token_ids.device
)
positions = positions.unsqueeze(0).expand(batch_size, context_length)
return self.embedding(positions) |