Spaces:
Running on Zero
Running on Zero
File size: 1,302 Bytes
aa7bfed | 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 | """Vectorized additive sinusoidal positional encoding (Vaswani et al.).
``p(k, i) = sin(k / base^{i/d})`` for even ``i``, and
``p(k, i) = cos(k / base^{(i-1)/d})`` for odd ``i``.
"""
from __future__ import annotations
import numpy as np
def sinusoidal_pe(seq_len: int, dim: int, base: float = 10000.0) -> np.ndarray:
"""Return a ``(seq_len, dim)`` additive PE matrix."""
if seq_len < 0 or dim < 1:
raise ValueError("seq_len must be >= 0 and dim >= 1")
positions = np.arange(seq_len, dtype=np.float64)[:, None]
dims = np.arange(dim, dtype=np.float64)[None, :]
exponent = np.where(dims % 2 == 0, dims / dim, (dims - 1.0) / dim)
angles = positions / (base ** exponent)
pe = np.empty((seq_len, dim), dtype=np.float64)
pe[:, 0::2] = np.sin(angles[:, 0::2])
pe[:, 1::2] = np.cos(angles[:, 1::2])
return pe
def add_positional_encoding(
embeddings: np.ndarray,
base: float = 10000.0,
) -> tuple[np.ndarray, np.ndarray]:
"""Return ``(pe, embeddings + pe)`` for a ``(seq, dim)`` matrix."""
embeddings = np.asarray(embeddings, dtype=np.float64)
if embeddings.ndim != 2:
raise ValueError("embeddings must be 2D (seq, dim)")
pe = sinusoidal_pe(embeddings.shape[0], embeddings.shape[1], base=base)
return pe, embeddings + pe
|