Spaces:
Sleeping
Sleeping
File size: 5,639 Bytes
6edd10d 2d52135 6edd10d | 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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | """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
|