स्टेप 2 · एम्बेडिंग्स और पोज़ीशन

import torch, torch.nn as nn

vocab_size, d_model = 50257, 768
tok_emb = nn.Embedding(vocab_size, d_model)   # a (50257, 768) lookup table

ids = torch.tensor([[15496, 995, 0]])         # batch 1, sequence 3
x = tok_emb(ids)                              # -> (1, 3, 768)

def rope_frequencies(head_dim, seq_len, base=10000.0, device="cpu"):
    # one angular frequency per pair of dimensions
    inv = 1.0 / (base ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
    pos = torch.arange(seq_len, device=device).float()
    ang = torch.outer(pos, inv)                 # (seq_len, head_dim/2)
    return torch.cos(ang), torch.sin(ang)

def apply_rope(x, cos, sin):
    # x: (batch, heads, seq, head_dim) -- rotate each adjacent pair of dims
    x1, x2 = x[..., 0::2], x[..., 1::2]
    cos, sin = cos[None, None], sin[None, None]
    out1 = x1 * cos - x2 * sin
    out2 = x1 * sin + x2 * cos
    return torch.stack((out1, out2), dim=-1).flatten(-2)