transformer.py

"""
╔══════════════════════════════════════════════════════════════════╗
║   GPT-STYLE TRANSFORMER  —  с нуля, только NumPy                ║
║                                                                  ║
║   Архитектура (Vaswani et al., "Attention is All You Need"):    ║
║   TokenEmbed + PosEmbed                                          ║
║   → N × [LayerNorm → MultiHeadAttention → residual              ║
║           LayerNorm → FeedForward        → residual]            ║
║   → LayerNorm → Linear → Softmax → следующий токен              ║
╚══════════════════════════════════════════════════════════════════╝
"""

import numpy as np
import json, os, math

# ─────────────────────────────────────────────────────────────────
# МАТЕМАТИЧЕСКИЕ ПРИМИТИВЫ
# ─────────────────────────────────────────────────────────────────

def softmax(x, axis=-1):
    e = np.exp(x - x.max(axis=axis, keepdims=True))
    return e / e.sum(axis=axis, keepdims=True)

def gelu(x):
    """GELU: Gaussian Error Linear Unit — активация в FFN слоях GPT"""
    return 0.5 * x * (1.0 + np.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * x**3)))

def gelu_backward(x, dy):
    t   = np.tanh(np.sqrt(2.0/np.pi) * (x + 0.044715 * x**3))
    phi = 0.5 * (1.0 + t)
    dphi = 0.5 * (1.0 - t**2) * np.sqrt(2.0/np.pi) * (1.0 + 3*0.044715*x**2)
    return dy * (phi + x * dphi)

def layer_norm_fwd(x, g, b, eps=1e-5):
    """Layer Normalization: нормализуем последнюю ось"""
    mean = x.mean(-1, keepdims=True)
    var  = x.var(-1,  keepdims=True)
    xhat = (x - mean) / np.sqrt(var + eps)
    return g * xhat + b, xhat, mean, var

def layer_norm_bwd(dy, x, g, xhat, var, eps=1e-5):
    """Обратный проход через Layer Norm"""
    N     = x.shape[-1]
    dg    = (dy * xhat).sum(tuple(range(dy.ndim - 1)))
    db    = dy.sum(tuple(range(dy.ndim - 1)))
    dxhat = dy * g
    dvar  = (dxhat * (x - x.mean(-1, keepdims=True)) * -0.5 * (var + eps)**-1.5).sum(-1, keepdims=True)
    dmean = (dxhat * -1.0/np.sqrt(var + eps)).sum(-1, keepdims=True) + dvar * (-2.0*(x - x.mean(-1, keepdims=True))).mean(-1, keepdims=True)
    dx    = dxhat/np.sqrt(var + eps) + dvar*2*(x - x.mean(-1, keepdims=True))/N + dmean/N
    return dx, dg, db

def softmax_bwd(s, ds):
    """Градиент через softmax (без CE)"""
    return s * (ds - (ds * s).sum(-1, keepdims=True))

# ─────────────────────────────────────────────────────────────────
# ПАРАМЕТР — обёртка над весами с поддержкой Adam
# ─────────────────────────────────────────────────────────────────

class Param:
    def __init__(self, data):
        self.data = np.array(data, dtype=np.float64)
        self.grad = np.zeros_like(self.data)
        self.m    = np.zeros_like(self.data)   # 1-й момент Adam
        self.v    = np.zeros_like(self.data)   # 2-й момент Adam

    def zero_grad(self):
        self.grad[:] = 0.0

    def adam_step(self, t, lr=3e-4, β1=0.9, β2=0.999, ε=1e-8):
        self.m = β1 * self.m + (1 - β1) * self.grad
        self.v = β2 * self.v + (1 - β2) * self.grad**2
        mh = self.m / (1 - β1**t)
        vh = self.v / (1 - β2**t)
        self.data -= lr * mh / (np.sqrt(vh) + ε)

# ─────────────────────────────────────────────────────────────────
# MULTI-HEAD SELF-ATTENTION
# ─────────────────────────────────────────────────────────────────

class MultiHeadAttention:
    """
    Ключевой механизм трансформера.

    Каждый токен задаёт три вектора:
      Q (Query)  — «что я ищу»
      K (Key)    — «что я предлагаю другим»
      V (Value)  — «что я передаю если выбрали меня»

    Оценка совместимости:
      scores = Q @ Kᵀ / √d_k

    Нормализация → распределение внимания:
      A = softmax(scores + causal_mask)

    Взвешенная сумма:
      output = A @ V

    Несколько «голов» (head) позволяют одновременно
    учить разные типы зависимостей.
    """

    def __init__(self, d_model, n_heads):
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k     = d_model // n_heads
        s = 0.02
        # Проекционные матрицы
        self.Wq = Param(np.random.randn(d_model, d_model) * s)
        self.Wk = Param(np.random.randn(d_model, d_model) * s)
        self.Wv = Param(np.random.randn(d_model, d_model) * s)
        self.Wo = Param(np.random.randn(d_model, d_model) * s)
        self.bq = Param(np.zeros(d_model))
        self.bk = Param(np.zeros(d_model))
        self.bv = Param(np.zeros(d_model))
        self.bo = Param(np.zeros(d_model))
        self._cache = {}

    def params(self):
        return [self.Wq, self.Wk, self.Wv, self.Wo,
                self.bq, self.bk, self.bv, self.bo]

    def forward(self, x, mask):
        """x: (T, d_model)"""
        T, C = x.shape
        H, dk = self.n_heads, self.d_k

        Q = x @ self.Wq.data + self.bq.data   # (T, C)
        K = x @ self.Wk.data + self.bk.data
        V = x @ self.Wv.data + self.bv.data

        # Разбиваем на головы: (H, T, dk)
        Qh = Q.reshape(T, H, dk).transpose(1, 0, 2)
        Kh = K.reshape(T, H, dk).transpose(1, 0, 2)
        Vh = V.reshape(T, H, dk).transpose(1, 0, 2)

        # Scaled dot-product attention
        scale  = math.sqrt(dk)
        scores = Qh @ Kh.transpose(0, 2, 1) / scale  # (H, T, T)
        scores = scores + mask[:T, :T]                 # каузальная маска
        A      = softmax(scores, axis=-1)              # (H, T, T)

        out = A @ Vh                                   # (H, T, dk)
        out = out.transpose(1, 0, 2).reshape(T, C)    # (T, C)
        result = out @ self.Wo.data + self.bo.data     # (T, C)

        self._cache = dict(x=x, Q=Q, K=K, V=V, Qh=Qh, Kh=Kh, Vh=Vh,
                           A=A, out=out, scale=scale, T=T, H=H, dk=dk)
        return result, A   # возвращаем A для визуализации

    def backward(self, d_result):
        c  = self._cache
        T, H, dk, scale = c['T'], c['H'], c['dk'], c['scale']

        # Градиент через выходную проекцию Wo
        self.Wo.grad += c['out'].T @ d_result
        self.bo.grad += d_result.sum(0)
        d_out = d_result @ self.Wo.data.T              # (T, C)

        # Перекладываем градиент обратно на головы: (H, T, dk)
        d_out_h = d_out.reshape(T, H, dk).transpose(1, 0, 2)

        # Градиенты V и A
        dVh = c['A'].transpose(0, 2, 1) @ d_out_h     # (H, T, dk)
        dA  = d_out_h @ c['Vh'].transpose(0, 2, 1)    # (H, T, T)

        # Backprop через softmax
        dS = softmax_bwd(c['A'], dA)                   # (H, T, T)

        # Градиенты Q и K
        dQh = dS @ c['Kh'] / scale                     # (H, T, dk)
        dKh = dS.transpose(0, 2, 1) @ c['Qh'] / scale # (H, T, dk)

        # Обратно к форме (T, C)
        dQ = dQh.transpose(1, 0, 2).reshape(T, self.d_model)
        dK = dKh.transpose(1, 0, 2).reshape(T, self.d_model)
        dV = dVh.transpose(1, 0, 2).reshape(T, self.d_model)

        # Градиенты проекционных матриц
        x = c['x']
        self.Wq.grad += x.T @ dQ;  self.bq.grad += dQ.sum(0)
        self.Wk.grad += x.T @ dK;  self.bk.grad += dK.sum(0)
        self.Wv.grad += x.T @ dV;  self.bv.grad += dV.sum(0)

        # Градиент входа x
        dx = dQ @ self.Wq.data.T + dK @ self.Wk.data.T + dV @ self.Wv.data.T
        return dx

# ─────────────────────────────────────────────────────────────────
# FEED-FORWARD NETWORK
# ─────────────────────────────────────────────────────────────────

class FeedForward:
    """
    Двухслойная нейросеть, применяется к каждой позиции независимо:
      FFN(x) = GELU(x·W₁ + b₁)·W₂ + b₂

    Расширение: d_model → d_ff (обычно 4×d_model) → d_model
    Захватывает фактические знания (в отличие от Attention, который
    агрегирует контекст).
    """

    def __init__(self, d_model, d_ff):
        s = 0.02
        self.W1 = Param(np.random.randn(d_model, d_ff) * s)
        self.b1 = Param(np.zeros(d_ff))
        self.W2 = Param(np.random.randn(d_ff, d_model) * s)
        self.b2 = Param(np.zeros(d_model))
        self._cache = {}

    def params(self):
        return [self.W1, self.b1, self.W2, self.b2]

    def forward(self, x):
        z1 = x @ self.W1.data + self.b1.data
        h  = gelu(z1)
        z2 = h @ self.W2.data + self.b2.data
        self._cache = dict(x=x, z1=z1, h=h)
        return z2

    def backward(self, dy):
        c = self._cache
        self.W2.grad += c['h'].T @ dy
        self.b2.grad += dy.sum(0)
        dh  = dy @ self.W2.data.T
        dz1 = gelu_backward(c['z1'], dh)
        self.W1.grad += c['x'].T @ dz1
        self.b1.grad += dz1.sum(0)
        return dz1 @ self.W1.data.T

# ─────────────────────────────────────────────────────────────────
# TRANSFORMER BLOCK
# ─────────────────────────────────────────────────────────────────

class TransformerBlock:
    """
    Один блок трансформера (Pre-LN вариант, как в GPT-2):

      x → LayerNorm₁ → MultiHeadAttention → + x  (residual)
        → LayerNorm₂ → FeedForward        → + x  (residual)

    Residual connections — ключ к обучению глубоких сетей:
    градиент может течь напрямую, минуя слои.
    """

    def __init__(self, d_model, n_heads, d_ff):
        self.attn = MultiHeadAttention(d_model, n_heads)
        self.ff   = FeedForward(d_model, d_ff)
        self.g1 = Param(np.ones(d_model));  self.b1 = Param(np.zeros(d_model))
        self.g2 = Param(np.ones(d_model));  self.b2 = Param(np.zeros(d_model))
        self._cache = {}

    def params(self):
        return self.attn.params() + self.ff.params() + \
               [self.g1, self.b1, self.g2, self.b2]

    def forward(self, x, mask):
        # Pre-LN attention
        xn1, xh1, mean1, var1 = layer_norm_fwd(x, self.g1.data, self.b1.data)
        attn_out, A = self.attn.forward(xn1, mask)
        x1 = x + attn_out              # residual

        # Pre-LN feed-forward
        xn2, xh2, mean2, var2 = layer_norm_fwd(x1, self.g2.data, self.b2.data)
        ff_out = self.ff.forward(xn2)
        x2 = x1 + ff_out              # residual

        self._cache = dict(x=x, x1=x1, xn1=xn1, xh1=xh1, mean1=mean1, var1=var1,
                           xn2=xn2, xh2=xh2, mean2=mean2, var2=var2)
        return x2, A

    def backward(self, dx2):
        c = self._cache
        # Residual 2: dx1 += dx2
        d_ff_out = dx2
        d_xn2 = self.ff.backward(d_ff_out)
        dx1_from_ff, dg2, db2 = layer_norm_bwd(d_xn2, c['x1'], self.g2.data, c['xh2'], c['var2'])
        self.g2.grad += dg2;  self.b2.grad += db2
        dx1 = dx2 + dx1_from_ff       # residual accumulation

        # Residual 1: dx += dx1
        d_attn_out = dx1
        d_xn1 = self.attn.backward(d_attn_out)
        dx_from_attn, dg1, db1 = layer_norm_bwd(d_xn1, c['x'], self.g1.data, c['xh1'], c['var1'])
        self.g1.grad += dg1;  self.b1.grad += db1
        dx = dx1 + dx_from_attn       # residual accumulation

        return dx

# ─────────────────────────────────────────────────────────────────
# GPT  —  ПОЛНАЯ МОДЕЛЬ
# ─────────────────────────────────────────────────────────────────

class GPT:
    """
    Генеративная языковая модель в стиле GPT.

    На каждом шаге принимает контекст (до context_len токенов)
    и предсказывает вероятность следующего токена.

    Повторяя этот процесс — генерирует текст.
    """

    def __init__(self, vocab_size, d_model=64, n_heads=4,
                 n_layers=4, d_ff=256, context_len=128):
        self.vocab_size  = vocab_size
        self.d_model     = d_model
        self.n_heads     = n_heads
        self.n_layers    = n_layers
        self.d_ff        = d_ff
        self.context_len = context_len

        s = 0.02
        self.tok_emb = Param(np.random.randn(vocab_size, d_model) * s)
        self.pos_emb = Param(np.random.randn(context_len, d_model) * s)

        self.blocks = [TransformerBlock(d_model, n_heads, d_ff) for _ in range(n_layers)]

        self.ln_g = Param(np.ones(d_model))
        self.ln_b = Param(np.zeros(d_model))
        self.head  = Param(np.random.randn(d_model, vocab_size) * s)

        # Каузальная маска: не смотрим в будущее
        mask = np.triu(np.full((context_len, context_len), -1e9), k=1)
        self.causal_mask = mask

        self.step_t = 0   # шаг Adam

        self._cache = {}

    def all_params(self):
        ps = [self.tok_emb, self.pos_emb, self.ln_g, self.ln_b, self.head]
        for b in self.blocks:
            ps += b.params()
        return ps

    def count_params(self):
        return sum(p.data.size for p in self.all_params())

    def zero_grad(self):
        for p in self.all_params():
            p.zero_grad()

    # ── Прямой проход ──────────────────────────────────────

    def forward(self, tokens):
        """
        tokens: список int, длина T ≤ context_len
        returns: logits (T, vocab_size), all_A (внимание всех блоков)
        """
        T = len(tokens)
        idx = np.array(tokens)

        # Эмбеддинги токенов + позиций
        x = self.tok_emb.data[idx] + self.pos_emb.data[:T]

        all_A = []
        for block in self.blocks:
            x, A = block.forward(x, self.causal_mask)
            all_A.append(A)

        # Финальная нормализация
        xn, xhat, mean, var = layer_norm_fwd(x, self.ln_g.data, self.ln_b.data)

        # Проекция на словарь
        logits = xn @ self.head.data   # (T, vocab_size)

        self._cache = dict(tokens=idx, x_before_ln=x,
                           xhat=xhat, mean=mean, var=var, xn=xn, T=T)
        return logits, all_A

    # ── Обратный проход ────────────────────────────────────

    def backward(self, logits, targets):
        """
        targets: список int (следующие токены для каждой позиции)
        Возвращает скалярный loss.
        """
        T = len(targets)
        c = self._cache

        # Cross-entropy loss + softmax gradient (совместный градиент)
        probs = softmax(logits[:T], axis=-1)
        loss  = -np.mean(np.log(probs[np.arange(T), targets] + 1e-10))

        # dL/d(logits) = (probs - one_hot) / T
        d_logits = probs.copy()
        d_logits[np.arange(T), targets] -= 1.0
        d_logits /= T

        # Backprop через head (линейный слой)
        self.head.grad  += c['xn'].T @ d_logits
        d_xn = d_logits @ self.head.data.T

        # Backprop через финальную LayerNorm
        dx, dg, db = layer_norm_bwd(d_xn, c['x_before_ln'],
                                    self.ln_g.data, c['xhat'], c['var'])
        self.ln_g.grad += dg;  self.ln_b.grad += db

        # Backprop через трансформерные блоки (в обратном порядке)
        for block in reversed(self.blocks):
            dx = block.backward(dx)

        # Backprop через эмбеддинги
        np.add.at(self.tok_emb.grad, c['tokens'], dx)
        self.pos_emb.grad[:c['T']] += dx

        return float(loss)

    def update(self, lr=3e-4):
        self.step_t += 1
        for p in self.all_params():
            p.adam_step(self.step_t, lr=lr)

    # ── Генерация текста ───────────────────────────────────

    def generate(self, prompt_ids, max_new=200, temperature=0.8, top_k=40):
        """
        Авторегрессивная генерация:
        на каждом шаге предсказываем следующий токен и добавляем к контексту.
        """
        ids = list(prompt_ids)
        for _ in range(max_new):
            ctx     = ids[-self.context_len:]
            logits, _ = self.forward(ctx)
            last    = logits[-1] / max(temperature, 1e-5)

            if top_k:
                topk = np.argsort(last)[-top_k:]
                mask = np.full_like(last, -1e9)
                mask[topk] = last[topk]
                last = mask

            probs    = softmax(last)
            next_id  = int(np.random.choice(len(probs), p=probs))
            ids.append(next_id)
        return ids

    def generate_detailed(self, prompt_ids, max_new=80, temperature=0.8, top_k=40):
        """Генерация с возвратом данных для визуализации (внимание, вероятности)"""
        ids   = list(prompt_ids)
        steps = []

        for step in range(max_new):
            ctx = ids[-self.context_len:]
            logits, all_A = self.forward(ctx)
            last = logits[-1] / max(temperature, 1e-5)

            if top_k:
                topk = np.argsort(last)[-top_k:]
                m    = np.full_like(last, -1e9)
                m[topk] = last[topk]
                last = m

            probs   = softmax(last)
            next_id = int(np.random.choice(len(probs), p=probs))

            # Топ-8 вероятных токенов
            top8_idx = np.argsort(probs)[-8:][::-1]

            # Внимание: последний блок, все головы, последняя позиция
            # all_A[layer]: (n_heads, T, T)
            T = len(ctx)
            last_block_attn = all_A[-1][:, T-1, :T]  # (n_heads, T)

            steps.append({
                'step':    step,
                'token':   next_id,
                'top8':    [(int(i), float(probs[i])) for i in top8_idx],
                'attention': last_block_attn.tolist(),   # (n_heads, context_len)
                'context_len': T,
            })

            ids.append(next_id)

        return ids, steps

    # ── Save / Load ────────────────────────────────────────

    def save(self, path):
        data = {
            'config': {
                'vocab_size': self.vocab_size, 'd_model': self.d_model,
                'n_heads': self.n_heads, 'n_layers': self.n_layers,
                'd_ff': self.d_ff, 'context_len': self.context_len,
            },
            'tok_emb': self.tok_emb.data.tolist(),
            'pos_emb': self.pos_emb.data.tolist(),
            'ln_g':    self.ln_g.data.tolist(),
            'ln_b':    self.ln_b.data.tolist(),
            'head':    self.head.data.tolist(),
            'blocks':  [],
            'step_t':  self.step_t,
        }
        for b in self.blocks:
            bd = {}
            for attr in ['g1','b1','g2','b2']:
                bd[attr] = getattr(b, attr).data.tolist()
            for attr in ['Wq','Wk','Wv','Wo','bq','bk','bv','bo']:
                bd[f'attn_{attr}'] = getattr(b.attn, attr).data.tolist()
            for attr in ['W1','b1','W2','b2']:
                bd[f'ff_{attr}'] = getattr(b.ff, attr).data.tolist()
            data['blocks'].append(bd)
        with open(path, 'w', encoding='utf-8') as f:
            json.dump(data, f)
        print(f'  GPT сохранён → {path}')

    @classmethod
    def load(cls, path):
        with open(path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        cfg = data['config']
        model = cls(**cfg)
        model.tok_emb.data = np.array(data['tok_emb'])
        model.pos_emb.data = np.array(data['pos_emb'])
        model.ln_g.data    = np.array(data['ln_g'])
        model.ln_b.data    = np.array(data['ln_b'])
        model.head.data    = np.array(data['head'])
        model.step_t       = data.get('step_t', 0)
        for i, bd in enumerate(data['blocks']):
            b = model.blocks[i]
            for attr in ['g1','b1','g2','b2']:
                getattr(b, attr).data = np.array(bd[attr])
            for attr in ['Wq','Wk','Wv','Wo','bq','bk','bv','bo']:
                getattr(b.attn, attr).data = np.array(bd[f'attn_{attr}'])
            for attr in ['W1','b1','W2','b2']:
                getattr(b.ff, attr).data = np.array(bd[f'ff_{attr}'])
        print(f'  GPT загружен ← {path}  ({model.count_params():,} параметров)')
        return model


# ─────────────────────────────────────────────────────────────────
# ТОКЕНИЗАТОР  (символьный уровень)
# ─────────────────────────────────────────────────────────────────

class CharTokenizer:
    """
    Простой символьный токенизатор:
    каждый уникальный символ = один токен.

    GPT-4 использует BPE с 100k токенов.
    Мы используем ~100 символов — принцип тот же.
    """

    def __init__(self):
        self.stoi = {}   # символ → индекс
        self.itos = {}   # индекс → символ

    def build(self, text):
        chars = sorted(set(text))
        self.stoi = {c: i for i, c in enumerate(chars)}
        self.itos = {i: c for c, i in self.stoi.items()}
        return self

    def encode(self, text):
        return [self.stoi.get(c, 0) for c in text]

    def decode(self, ids):
        return ''.join(self.itos.get(i, '?') for i in ids)

    @property
    def vocab_size(self):
        return len(self.stoi)

    def save(self, path):
        with open(path, 'w', encoding='utf-8') as f:
            json.dump({'stoi': self.stoi, 'itos': {str(k): v for k,v in self.itos.items()}}, f, ensure_ascii=False)

    @classmethod
    def load(cls, path):
        t = cls()
        with open(path, 'r', encoding='utf-8') as f:
            d = json.load(f)
        t.stoi = d['stoi']
        t.itos = {int(k): v for k,v in d['itos'].items()}
        return t