| """GoLLeM v6 model code (inference only), extracted from the training script used for this run. |
| |
| The module definitions are those of the training script; only the weight initialisation and the training loss were |
| removed and comments translated. Weights load with strict=True. Requirements: torch, safetensors, tokenizers. |
| |
| from modeling_gollem_v6 import load_gollem_v6, generate |
| model, tok = load_gollem_v6(".") # directory with config.json, model.safetensors, tokenizer.json |
| print(generate(model, tok, "Najwyższym szczytem Polski jest", max_new_tokens=40)) |
| |
| Notes: RoPE uses the interleaved convention (pairs of adjacent channels), not rotate_half; the output head is tied to |
| the token embedding; context length is 1,024 tokens. |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import os |
| import types |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| class RMSNorm(nn.Module): |
| """RMSNorm with fp32 compute for stability.""" |
| def __init__(self, d, eps=1e-6): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(d)) |
| self.eps = eps |
|
|
| def forward(self, x): |
| return x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps).to(x.dtype) * self.weight |
|
|
|
|
| def make_norm(d, cfg): |
| return RMSNorm(d, cfg.norm_eps) if cfg.norm == "rmsnorm" else nn.LayerNorm(d) |
|
|
|
|
| def apply_rope(x, base=100000.0): |
| """Parameter-free RoPE on [B, H, T, D], interleaved convention (even/odd channel pairs).""" |
| _, _, T, dim = x.shape |
| pos = torch.arange(T, device=x.device, dtype=torch.float32) |
| freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=x.device, dtype=torch.float32) / dim)) |
| ang = torch.outer(pos, freq) |
| cos, sin = ang.cos().to(x.dtype)[None, None], ang.sin().to(x.dtype)[None, None] |
| even, odd = x[..., ::2], x[..., 1::2] |
| return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2) |
|
|
|
|
| class SwiGLU(nn.Module): |
| """Gated MLP: down(silu(gate(x)) * up(x)).""" |
| def __init__(self, d, hidden): |
| super().__init__() |
| self.gate = nn.Linear(d, hidden, bias=False) |
| self.up = nn.Linear(d, hidden, bias=False) |
| self.down = nn.Linear(hidden, d, bias=False) |
|
|
| def forward(self, x): |
| return self.down(F.silu(self.gate(x)) * self.up(x)) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, d, nh, block, cfg, is_first=False): |
| super().__init__() |
| self.ln1 = make_norm(d, cfg) |
| self.ln2 = make_norm(d, cfg) |
| self.qkv = nn.Linear(d, 3 * d) |
| self.proj = nn.Linear(d, d) |
| if cfg.ffn == "swiglu": |
| self.mlp = SwiGLU(d, int(round(cfg.ffn_mult * d))) |
| else: |
| self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d)) |
| self.nh = nh |
| self.d = d |
| self.cfg = cfg |
| self.is_first = is_first |
| if cfg.value_residual and not is_first: |
| self.vr_lambda = nn.Parameter(torch.zeros(1)) |
| if cfg.qk_norm: |
| hd = d // nh |
| self.q_norm = RMSNorm(hd, cfg.norm_eps) |
| self.k_norm = RMSNorm(hd, cfg.norm_eps) |
|
|
| def forward(self, x, v0=None): |
| B, T, D = x.size() |
| h = self.ln1(x) |
| q, k, v = self.qkv(h).split(self.d, dim=2) |
| hd = D // self.nh |
| q = q.view(B, T, self.nh, hd).transpose(1, 2) |
| k = k.view(B, T, self.nh, hd).transpose(1, 2) |
| v = v.view(B, T, self.nh, hd).transpose(1, 2) |
| if self.cfg.qk_norm: |
| q = self.q_norm(q) |
| k = self.k_norm(k) |
| if self.cfg.pos == "rope": |
| q = apply_rope(q, self.cfg.rope_theta) |
| k = apply_rope(k, self.cfg.rope_theta) |
| if self.cfg.value_residual: |
| if self.is_first: |
| v0 = v |
| else: |
| v = v + self.vr_lambda * v0 |
| y = F.scaled_dot_product_attention(q, k, v, is_causal=True) |
| y = y.transpose(1, 2).contiguous().view(B, T, D) |
| x = x + self.proj(y) |
| x = x + self.mlp(self.ln2(x)) |
| return x, v0 |
|
|
|
|
| class GPT(nn.Module): |
| def __init__(self, vocab, n_layer, n_embd, n_head, block, cfg): |
| super().__init__() |
| self.cfg = cfg |
| self.tok = nn.Embedding(vocab, n_embd) |
| self.use_rope = cfg.pos == "rope" |
| if not self.use_rope: |
| self.pos = nn.Embedding(block, n_embd) |
| self.blocks = nn.ModuleList([Block(n_embd, n_head, block, cfg, is_first=(i == 0)) for i in range(n_layer)]) |
| self.lnf = make_norm(n_embd, cfg) |
| self.head = nn.Linear(n_embd, vocab, bias=False) |
| self.head.weight = self.tok.weight |
| self.block = block |
|
|
| def forward(self, idx): |
| B, T = idx.size() |
| x = self.tok(idx) |
| if not self.use_rope: |
| pos = torch.arange(T, device=idx.device) |
| x = x + self.pos(pos)[None] |
| v0 = None |
| for b in self.blocks: |
| x, v0 = b(x, v0) |
| logits = self.head(self.lnf(x)) |
| cap = getattr(self.cfg, "logit_cap", 0.0) |
| if cap and cap > 0: |
| logits = cap * torch.tanh(logits / cap) |
| return logits |
|
|
|
|
| def load_gollem_v6(path, device="cpu"): |
| """Model (eval mode) and tokenizer from a directory with config.json, model.safetensors and tokenizer.json.""" |
| from safetensors.torch import load_file |
| from tokenizers import Tokenizer |
|
|
| with open(os.path.join(path, "config.json"), encoding="utf-8") as f: |
| cfg = types.SimpleNamespace(**json.load(f)) |
| model = GPT(cfg.vocab, cfg.n_layer, cfg.n_embd, cfg.n_head, cfg.block, cfg) |
| model.load_state_dict(load_file(os.path.join(path, "model.safetensors")), strict=True) |
| return model.to(device).eval(), Tokenizer.from_file(os.path.join(path, "tokenizer.json")) |
|
|
|
|
| @torch.no_grad() |
| def generate(model, tok, prompt, max_new_tokens=64, temperature=0.0, top_k=None, seed=None): |
| """Continue `prompt`. temperature 0 = greedy; otherwise sampling, optionally top-k.""" |
| g = torch.Generator().manual_seed(seed) if seed is not None else None |
| dev = next(model.parameters()).device |
| x = torch.tensor([tok.encode(prompt).ids], device=dev) |
| for _ in range(max_new_tokens): |
| logits = model(x[:, -model.block:])[:, -1].float().cpu() |
| if temperature <= 0: |
| nxt = logits.argmax(-1, keepdim=True) |
| else: |
| logits = logits / temperature |
| if top_k: |
| kth = torch.topk(logits, top_k).values[:, -1:] |
| logits = logits.masked_fill(logits < kth, float("-inf")) |
| nxt = torch.multinomial(torch.softmax(logits, -1), 1, generator=g) |
| x = torch.cat([x, nxt.to(dev)], dim=1) |
| return tok.decode(x[0].tolist()) |
|
|
