Spaces:
Running on Zero
Running on Zero
File size: 1,386 Bytes
9cde5e3 98d1e65 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 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | """
Input representation layer.
Combines token embeddings and positional embeddings to produce
the representation that enters the transformer.
"""
from __future__ import annotations
import torch
import torch.nn as nn
from .embedding import Embedding
from .positional_embedding import PositionalEmbedding
class InputRepresentation(nn.Module):
"""
Input Representation
Input:
token_ids
Shape:
(batch_size, context_length)
Output:
input_representation
Shape:
(batch_size, context_length, d_model)
Formula:
Input Representation =
Token Embedding +
Positional Embedding
"""
def __init__(
self,
vocab_size: int,
context_length: int,
d_model: int,
) -> None:
super().__init__()
self.token_embedding = Embedding(
vocab_size=vocab_size,
d_model=d_model,
)
self.positional_embedding = PositionalEmbedding(
context_length=context_length,
d_model=d_model,
)
def forward(
self,
token_ids: torch.Tensor,
) -> torch.Tensor:
token_vectors = self.token_embedding(token_ids)
position_vectors = self.positional_embedding(token_ids)
return token_vectors + position_vectors
|