Spaces:
Sleeping
Sleeping
| """Embedding modules used by the transformer architecture.""" | |
| from __future__ import annotations | |
| from typing import cast | |
| import torch | |
| import torch.nn as nn | |
| from torch import Tensor | |
| from ..utils import sinusoidal_positional_encoding | |
| __all__ = ["InputEmbedding", "PositionalEmbedding", "LearnedPositionalEmbedding"] | |
| class InputEmbedding(nn.Module): | |
| """Token and positional embedding with optional dropout.""" | |
| def __init__( | |
| self, | |
| vocab_size: int, | |
| d_model: int, | |
| max_seq_len: int, | |
| pad_id: int, | |
| dropout_rate: float = 0.0, | |
| ) -> None: | |
| super().__init__() | |
| if not isinstance(vocab_size, int): | |
| raise TypeError(f"vocab_size must be int, got {type(vocab_size)}") | |
| if not isinstance(d_model, int): | |
| raise TypeError(f"d_model must be int, got {type(d_model)}") | |
| if not isinstance(max_seq_len, int): | |
| raise TypeError(f"max_seq_len must be int, got {type(max_seq_len)}") | |
| if not isinstance(pad_id, int): | |
| raise TypeError(f"pad_id must be int, got {type(pad_id)}") | |
| if not isinstance(dropout_rate, float): | |
| raise TypeError(f"dropout_rate must be a float, got {type(dropout_rate)}") | |
| if vocab_size <= 0: | |
| raise ValueError("vocab_size must be > 0") | |
| if d_model <= 0: | |
| raise ValueError("d_model must be > 0") | |
| if max_seq_len <= 0: | |
| raise ValueError("max_seq_len must be > 0") | |
| if not (0 <= pad_id < vocab_size): | |
| raise ValueError(f"pad_id must be in [0, {vocab_size - 1}]") | |
| if not 0.0 <= dropout_rate < 1.0: | |
| raise ValueError("dropout_rate must be in [0, 1)") | |
| self.vocab_size = vocab_size | |
| self.d_model = d_model | |
| self.max_seq_len = max_seq_len | |
| self.pad_id = pad_id | |
| self.token_embed = nn.Embedding(vocab_size, d_model, padding_idx=pad_id) | |
| self.pos_embed = PositionalEmbedding(max_seq_len, d_model) | |
| self.dropout = nn.Dropout(dropout_rate) | |
| def forward(self, x: Tensor) -> Tensor: | |
| if not isinstance(x, torch.Tensor): | |
| raise TypeError(f"x must be Tensor, got {type(x)}") | |
| if x.dtype != torch.long: | |
| raise TypeError(f"x must be torch.long, got {x.dtype}") | |
| if x.dim() != 2: | |
| raise ValueError(f"x must be 2D (B, S), got shape {tuple(x.shape)}") | |
| if x.size(1) > self.max_seq_len: | |
| raise ValueError(f"Sequence length {x.size(1)} exceeds max_seq_len {self.max_seq_len}") | |
| tokens = self.token_embed(x) * (self.d_model**0.5) | |
| tokens = self.pos_embed(tokens) | |
| return self.dropout(tokens) | |
| class PositionalEmbedding(nn.Module): | |
| """Fixed sinusoidal positional encodings.""" | |
| def __init__(self, max_seq_len: int, d_model: int) -> None: | |
| super().__init__() | |
| if not isinstance(max_seq_len, int): | |
| raise TypeError(f"max_seq_len must be int, got {type(max_seq_len)}") | |
| if not isinstance(d_model, int): | |
| raise TypeError(f"d_model must be int, got {type(d_model)}") | |
| if max_seq_len < 0: | |
| raise ValueError("max_seq_len must be >= 0") | |
| if d_model <= 0: | |
| raise ValueError("d_model must be > 0") | |
| self.max_seq_len = max_seq_len | |
| self.d_model = d_model | |
| pe = sinusoidal_positional_encoding(max_seq_len, d_model) | |
| self.register_buffer("pe", pe.unsqueeze(0)) | |
| def forward(self, x: Tensor) -> Tensor: | |
| if not isinstance(x, torch.Tensor): | |
| raise TypeError(f"x must be Tensor, got {type(x)}") | |
| if x.dim() != 3: | |
| raise ValueError(f"x must be 3D (B, S, D), got shape {tuple(x.shape)}") | |
| _, seq_len, dim = x.shape | |
| if dim != self.d_model: | |
| raise ValueError(f"d_model mismatch: got {dim}, expected {self.d_model}") | |
| if seq_len > self.max_seq_len: | |
| raise ValueError(f"seq_len {seq_len} exceeds max_seq_len {self.max_seq_len}") | |
| pe_buffer = cast(Tensor, self.pe) | |
| return x + pe_buffer[:, :seq_len].to(dtype=x.dtype, device=x.device) | |
| class LearnedPositionalEmbedding(nn.Module): | |
| """Learned positional embeddings compatible with :class:`InputEmbedding`.""" | |
| def __init__(self, max_seq_len: int, d_model: int) -> None: | |
| super().__init__() | |
| if not isinstance(max_seq_len, int): | |
| raise TypeError(f"max_seq_len must be int, got {type(max_seq_len)}") | |
| if not isinstance(d_model, int): | |
| raise TypeError(f"d_model must be int, got {type(d_model)}") | |
| if max_seq_len < 0: | |
| raise ValueError("max_seq_len must be >= 0") | |
| if d_model <= 0: | |
| raise ValueError("d_model must be > 0") | |
| self.max_seq_len = max_seq_len | |
| self.d_model = d_model | |
| self.pos_table = nn.Embedding(max_seq_len, d_model) | |
| nn.init.normal_(self.pos_table.weight, mean=0.0, std=0.02) | |
| def forward(self, x: Tensor) -> Tensor: | |
| if not isinstance(x, torch.Tensor): | |
| raise TypeError(f"x must be Tensor, got {type(x)}") | |
| if x.dim() != 3: | |
| raise ValueError(f"x must be 3D (B, S, D), got {tuple(x.shape)}") | |
| batch, seq_len, dim = x.shape | |
| if dim != self.d_model: | |
| raise ValueError(f"d_model mismatch: got {dim}, expected {self.d_model}") | |
| if seq_len > self.max_seq_len: | |
| raise ValueError(f"seq_len {seq_len} exceeds max_seq_len {self.max_seq_len}") | |
| positions = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(batch, seq_len) | |
| pos_emb = self.pos_table(positions) | |
| return x + pos_emb | |