Spaces:
Running on Zero
Running on Zero
| """ | |
| 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) |