Upload folder using huggingface_hub

bard_core.py

@@ -1,4 +1,4 @@
-"""The network of bard-small, taken whole from the training code (fictionet/model.py, numhead.py, train/pretrain.py): windowed attention with a key/value cache, the number encoder/decoder, the LM."""
+"""The network of bard-small."""
 import re
 import math
 import time
@@ -7,7 +7,7 @@ from functools import lru_cache
 import torch
 from torch import nn
 from torch.nn import functional as F
-try:  # only the training path (and a full forward without a cache) uses it
+try:
     from torch.nn.attention.flex_attention import create_block_mask, flex_attention
 except ImportError:  # pragma: no cover
     create_block_mask = flex_attention = None
@@ -25,9 +25,6 @@ def sliding_window_block_mask(T: int, window: int, device: str):
 
 
 def apply_rope(x, base=10000.0, offset=0):
-    """Rotary embedding on (B, H, T, D) for absolute positions offset .. offset+T-1.
-    Frequencies are built in fp32 inside forward, so a .bfloat16() model cast can't
-    quantise them (a registered buffer would get cast)."""
     T, D = x.shape[-2], x.shape[-1]
     inv_freq = base ** (-torch.arange(0, D, 2, device=x.device, dtype=torch.float32) / D)
     ang = torch.arange(offset, offset + T, device=x.device, dtype=torch.float32)[:, None] * inv_freq
@@ -37,23 +34,14 @@ def apply_rope(x, base=10000.0, offset=0):
 
 
 class LayerKV:
-    """One layer's cached keys and values, (B, H, S, D), in a buffer that grows by doubling.
-
-    Keys are stored after RoPE, at their absolute positions, so they never need
-    re-rotating. A sliding-window layer keeps only the last `window - 1` keys - all
-    any later query can still see - so its memory stays bounded however long the
-    sequence runs; the kept span is moved to the front only when the buffer fills.
-    """
 
     def __init__(self, window):
         self.window = window
-        self.k = self.v = None     # buffers
-        self.lo = self.hi = 0      # live span [lo, hi) of the buffers
-        self.start = 0             # absolute position of the key at `lo`
+        self.k = self.v = None
+        self.lo = self.hi = 0
+        self.start = 0
 
     def append(self, k, v):
-        """Add T new keys/values; return all visible keys, values and the absolute
-        position of the first of them."""
         T = k.shape[2]
         if self.k is None:
             cap = max(64, 2 * T)
@@ -62,7 +50,7 @@ class LayerKV:
         elif self.hi + T > self.k.shape[2]:
             live = self.hi - self.lo
             cap = self.k.shape[2]
-            if live + T > cap // 2:  # still mostly full after compacting: grow
+            if live + T > cap // 2:
                 cap = 2 * (live + T)
             k_buf = self.k.new_empty(self.k.shape[:2] + (cap,) + self.k.shape[3:])
             v_buf = self.v.new_empty(k_buf.shape)
@@ -75,7 +63,6 @@ class LayerKV:
         return self.k[:, :, self.lo:self.hi], self.v[:, :, self.lo:self.hi], self.start
 
     def trim(self):
-        """After this step's attention: forget keys no future query can reach."""
         if self.window is not None and self.hi - self.lo > self.window - 1:
             drop = self.hi - self.lo - (self.window - 1)
             self.lo += drop
@@ -83,9 +70,6 @@ class LayerKV:
 
 
 class KVCache:
-    """Keys/values for every block of a CausalTransformer, plus how many tokens it has
-    seen. Pass it to CausalTransformer.forward (or LM.hidden) to process only the new
-    tokens: any number of them, on top of any amount already cached."""
 
     def __init__(self, transformer):
         self.windows = [block.window for block in transformer.blocks]
@@ -100,15 +84,6 @@ class KVCache:
 
 
 def attend(q, k, v, rope, window, block_mask, cache=None, pos=0):
-    """Causal (optionally sliding-window) attention for (B, H, T, D) q/k/v.
-
-    Without a cache this is exactly the training path. With one, q/k/v hold only the
-    T new tokens at absolute positions pos .. pos+T-1: keys are rotated at those
-    positions, appended to the cache, and the queries attend over cached + new keys
-    under the true mask - key <= query, and query - key < window on sliding-window
-    layers - so a chunk of any length (including one crossing the window edge) sees
-    exactly what it would in a full forward pass.
-    """
     if cache is None:
         if rope:
             q, k = apply_rope(q), apply_rope(k)
@@ -122,9 +97,9 @@ def attend(q, k, v, rope, window, block_mask, cache=None, pos=0):
     k_all, v_all, k_start = cache.append(k, v)
     S = k_all.shape[2]
     if T == 1 and (window is None or S <= window):
-        out = F.scaled_dot_product_attention(q, k_all, v_all)          # decode: all visible
+        out = F.scaled_dot_product_attention(q, k_all, v_all)
     elif S == T and (window is None or T <= window):
-        out = F.scaled_dot_product_attention(q, k_all, v_all, is_causal=True)  # fresh prefill
+        out = F.scaled_dot_product_attention(q, k_all, v_all, is_causal=True)
     else:
         q_pos = torch.arange(pos, pos + T, device=q.device)[:, None]
         k_pos = torch.arange(k_start, k_start + S, device=q.device)[None, :]
@@ -143,7 +118,7 @@ class Attention(nn.Module):
         self.d_model = d_model
         self.heads = heads
         self.head_dim = d_model // heads
-        self.window = window  # None -> global causal attention
+        self.window = window
         self.rope = rope
 
         self.qkvg = nn.Linear(d_model, 4 * d_model, bias=False)
@@ -205,13 +180,6 @@ class Block(nn.Module):
 
 
 class ParallelBlock(nn.Module):
-    """PaLM-style parallel block: x + Attn(norm(x)) + FFN(norm(x)), actually fused.
-
-    One norm; ONE input matmul producing q|k|v|g|gate|up; ONE output matmul over
-    [gated attn | swiglu] (= attn.out and W_down stacked along the input dim, which is
-    exactly out(a) + W_down(f)). 2 matmuls per block instead of 4, each wider, which is
-    what a narrow model on a small GPU needs to keep the tensor cores fed.
-    """
 
     def __init__(self, d_model, heads, ffn_expansion=4, window=None, rope=False):
         super().__init__()
@@ -265,9 +233,6 @@ class CausalTransformer(nn.Module):
         parallel: bool = False,
         rope: str = "local",
     ):
-        """rope: "local" = RoPE on sliding-window layers, NoPE on global ones (RNoPE-SWA,
-        Yang et al. 2025; SmolLM3 ablation at 3B: better long context, no short-context
-        cost; NoPE LMs stay competitive down to ~250M, Haviv et al. 2022), "all", or "none"."""
         super().__init__()
         self.grad_checkpoint = grad_checkpoint
 
@@ -310,12 +275,6 @@ class CausalTransformer(nn.Module):
     positions, like steering, need it). 0 otherwise."""
 
     def forward(self, x, masks=None, cache=None):
-        """masks: precomputed block_masks(T, device); pass them when this runs inside an outer
-        torch.compile, so mask construction stays out of the compiled graph.
-
-        cache: a KVCache. x then holds only the tokens not yet in it (any number), the
-        cache is extended with them, and the result is identical to running the whole
-        sequence - at the cost of the new tokens only."""
         if cache is not None:
             pos = cache.pos
             self.cache_offset = pos
@@ -340,19 +299,13 @@ class CausalTransformer(nn.Module):
         return self.norm(x)
 
 
-# ----------------------------------------------------------------------------- numbers
-
-import re  # noqa: E402
-
-MAX_DIGITS = 8                  # 10**8 - 1 < 2**27
+MAX_DIGITS = 8
 LEN_SHIFT = 27
 VALUE_MASK = (1 << LEN_SHIFT) - 1
 
 NUMBER_RE = re.compile(r"( ?)([0-9]+)")
 
 
-# ----------------------------------------------------------------------------- packing
-
 def pack(digits: str) -> int:
     assert 0 < len(digits) <= MAX_DIGITS
     return int(digits) | (len(digits) << LEN_SHIFT)
@@ -363,7 +316,6 @@ def unpack(packed: int) -> str:
 
 
 def split_numbers(text):
-    """-> list of str (text segment) and (spaced: bool, packed: int) items, in order."""
     out, pos = [], 0
     for m in NUMBER_RE.finditer(text):
         if m.start() > pos:
@@ -386,8 +338,6 @@ def _pow10(n, device):
     return 10 ** torch.arange(n, device=device, dtype=torch.long)
 
 
-# ----------------------------------------------------------------------------- encoder
-
 class NumberEncoder(nn.Module):
     def __init__(self, d_model, n_freq=4):
         super().__init__()
@@ -400,9 +350,8 @@ class NumberEncoder(nn.Module):
                 nn.init.zeros_(m.bias)
 
     def forward(self, packed):
-        """packed: (...) int, 0 where the token is not a number -> (..., d_model), zero there."""
         L, value = _lengths_values(packed)
-        digits = value.unsqueeze(-1) // _pow10(MAX_DIGITS, packed.device) % 10       # LSB first
+        digits = value.unsqueeze(-1) // _pow10(MAX_DIGITS, packed.device) % 10
         present = (torch.arange(MAX_DIGITS, device=packed.device) < L.unsqueeze(-1)).float()
         freqs = torch.arange(1, self.n_freq + 1, dtype=torch.float32, device=packed.device)
         angle = 2 * math.pi * digits.float().unsqueeze(-1) / 10 * freqs
@@ -417,8 +366,6 @@ class NumberEncoder(nn.Module):
         return out * (packed != 0).unsqueeze(-1).to(out.dtype)
 
 
-# ----------------------------------------------------------------------------- decoder
-
 class NumberDecoder(nn.Module):
     def __init__(self, d_model, width=256, heads=4, layers=2):
         super().__init__()
@@ -426,8 +373,8 @@ class NumberDecoder(nn.Module):
                                       nn.Linear(width, MAX_DIGITS, bias=False))
         self.ctx = nn.Linear(d_model, width, bias=False)
         self.len_emb = nn.Embedding(MAX_DIGITS, width)
-        self.place_emb = nn.Embedding(MAX_DIGITS, width)     # power of ten of the slot
-        self.prev_emb = nn.Embedding(11, width)              # previous digit, 10 = start
+        self.place_emb = nn.Embedding(MAX_DIGITS, width)
+        self.prev_emb = nn.Embedding(11, width)
         self.blocks = nn.ModuleList([Block(width, heads) for _ in range(layers)])
         self.norm = nn.RMSNorm(width)
         self.out = nn.Linear(width, 10, bias=False)
@@ -439,12 +386,11 @@ class NumberDecoder(nn.Module):
             nn.init.normal_(b.ffn.W_down.weight, std=0.02 / math.sqrt(2 * layers))
 
     def _digit_logits(self, h, L, prev):
-        """h (N, d), L (N,) in 1..MAX, prev (N, MAX) previous-digit ids -> (N, MAX, 10)."""
         j = torch.arange(MAX_DIGITS, device=h.device)
-        power = (L.unsqueeze(-1) - 1 - j).clamp(min=0)            # slots past L are masked anyway
+        power = (L.unsqueeze(-1) - 1 - j).clamp(min=0)
         x = (self.ctx(h) + self.len_emb(L - 1)).unsqueeze(1) + self.place_emb(power) + self.prev_emb(prev)
         for b in self.blocks:
-            x = b(x)                                              # no mask -> causal SDPA
+            x = b(x)
         return self.out(self.norm(x))
 
     @staticmethod
@@ -456,9 +402,6 @@ class NumberDecoder(nn.Module):
         return L, digits, valid
 
     def nll(self, h, packed, weight=None):
-        """Summed NLL (length + every digit) of the target numbers. h (N, d), packed (N,).
-        weight (N,) masks padding rows, so N can be a fixed capacity (static shapes -> compile);
-        padded rows must still hold a valid packed number (e.g. pack("0"))."""
         L, digits, valid = self._msb_digits(packed)
         prev = torch.cat([torch.full_like(digits[:, :1], 10), digits[:, :-1]], 1)
         len_nll = F.cross_entropy(self.len_head(h).float(), L - 1, reduction="none")
@@ -469,7 +412,6 @@ class NumberDecoder(nn.Module):
 
     @torch.no_grad()
     def generate(self, h, temperature=0.0):
-        """h (N, d) -> packed numbers (N,). temperature 0 = greedy."""
         def pick(logits):
             if temperature <= 0:
                 return logits.argmax(-1)
@@ -488,19 +430,15 @@ class NumberDecoder(nn.Module):
         return value | (L << LEN_SHIFT)
 
 
-# ----------------------------------------------------------------------------- the language model
-
 LOGIT_SOFTCAP = 30.0
 
 
 def softcap(logits):
-    # Gemma 2 / modded-nanogpt: bounds the logits smoothly
+
     return LOGIT_SOFTCAP * torch.tanh(logits.float() / LOGIT_SOFTCAP)
 
 
 class LM(nn.Module):
-    """Token LM + dedicated number path: <|num|>/<|snum|> inputs get NumberEncoder(value) added to their embedding;
-    where the next token is a number, NumberDecoder predicts its digits."""
 
     def __init__(self, vocab, d_model, heads, n_blocks, window, global_every, parallel=True, rope="local"):
         super().__init__()
@@ -512,5 +450,4 @@ class LM(nn.Module):
         self.num_dec = NumberDecoder(d_model)
 
     def hidden(self, idx, vals, masks=None, cache=None):
-        """cache: a KVCache - then idx/vals hold only the tokens not yet in it."""
         return self.body(self.emb(idx) + self.num_enc(vals), masks, cache)

bard_game.py

@@ -1,10 +1,4 @@
-"""The story engine of bard-small. The model narrates and proposes state changes and choices; this code applies the
-changes it believes, keeps the state, steers the length of a passage and the shape of the choices, and keeps the
-tale as one string in the model's own format (see format.txt).
-
-    <|start|>state<|content|>money\\n100<|end|><|start|>story<|content|>...narration...<|vars|>money<|add|>-20
-    <|choices|>Pay<|sep|>Run<|end|><|start|>state<|content|>money\\n80<|end|><|start|>player<|content|>Pay<|end|>...
-"""
+"""The story engine of bard-small."""
 import math
 import re
 import time
@@ -21,7 +15,7 @@ TITLE = "Stochastic Lands of Prophecy"
 DEFAULT_OPENING = "You wake up in a dungeon."
 MAX_INPUT = 300
 ANNOUNCE = {"money", "gold", "health"}
-LENGTHS = {"short": 60, "medium": 150, "long": None}  # narration tokens a passage aims at
+LENGTHS = {"short": 60, "medium": 150, "long": None}
 SAMPLING = dict(temperature=0.8, top_k=50, top_p=0.95, rep_penalty=1.1, rep_window=128)
 
 OP_RE = re.compile(r"^\s*(.+?)\s*<\|(add|set)\|>\s*(-?\s*\d+)\s*$")
@@ -30,7 +24,7 @@ SENTENCE_DONE = re.compile(r"[.!?…][\"”’')\]]*(\s|$)")
 SPELLED_MARKER = re.compile(r"<?\|\s*(vars|choices|sep|end|add|set)\s*\|\s*>?", re.I)
 TURN = re.compile(r"<\|start\|>(state|story|player)<\|content\|>(.*?)<\|end\|>", re.S)
 
-# the story site's own interface links, not story choices ("End Game and Leave Comments", "Start over", ...)
+
 META_CHOICE = re.compile(
     r"leave (a )?comments?|rate (this|the) (story|game)|end (the )?game|give (me )?(a )?rating|"
     r"back to (the )?(menu|site|story list|main menu)|return to (the )?main menu|main menu|"
@@ -40,8 +34,7 @@ META_CHOICE = re.compile(
     r"(tale|story|game|adventure|line)\W*$|this is the end of (your|the) (adventure|story|game)",
     re.I)
 
-# changes the engine only believes when the narration is about them: a 137M model writes "money -10" after being
-# eaten alive; the passage has to mention money for that
+
 EVIDENCE = {
     "money": r"\b(money|coins?|gold|silver|copper|cash|dollars?|bucks|credits?|price|pay|paid|"
              r"buy|bought|sell|sold|cost|steal|stole|stolen|rob|robbed|reward|bribe|purse|wallet|"
@@ -51,7 +44,7 @@ EVIDENCE = {
 }
 EVIDENCE["gold"] = EVIDENCE["money"]
 EVIDENCE = {k: re.compile(v, re.I) for k, v in EVIDENCE.items()}
-AMOUNTS = {"money", "gold", "coins", "health"}  # owned as an amount: spent and earned, never below 0
+AMOUNTS = {"money", "gold", "coins", "health"}
 
 
 def story_choices(choices):
@@ -64,7 +57,6 @@ def believable(name, narration):
 
 
 def sane(name, old, new):
-    """Spending more than you have is refused, and so is a more than tenfold gain."""
     if name not in AMOUNTS or old is None:
         return True
     return 0 <= new <= max(1000, 10 * old)
@@ -90,7 +82,6 @@ def parse_state(content):
 
 
 def clean_input(text):
-    """A player's words: no control characters, and no way to write a marker."""
     text = "".join(c for c in text if c == " " or unicodedata.category(c)[0] != "C")
     text = text.replace("<|", "<").replace("|>", ">")
     return re.sub(r"\s+", " ", text).strip()[:MAX_INPUT]
@@ -101,7 +92,6 @@ def repair(raw):
 
 
 def narration_so_far(raw):
-    """What a reader may see of a turn being written: up to the first marker, and never a stray "<" or "|"."""
     visible = re.split(r"<\||\|", raw, maxsplit=1)[0]
     return visible[:-1] if visible.endswith("<") else visible
 
@@ -117,7 +107,6 @@ def unique(choices):
 
 
 def parse_turn(raw):
-    """The model's output for one turn -> (narration, [(name, op, value)], [choices], problems)."""
     problems = []
     fixed = repair(raw)
     if fixed != raw:
@@ -143,18 +132,7 @@ def parse_turn(raw):
     return text.strip(), ops, choices, problems
 
 
-# ----------------------------------------------------------------------------- steering
-
 class LengthControl:
-    """Steers one story turn (the adjust_logits hook of generate) without retraining.
-
-    Narration: it cannot end before the model has written a sentence of its own after the opening; past a token budget
-    the markers that end it (<|vars|>, <|choices|>) get a growing boost, only right after a sentence ends, so text is
-    never cut mid-sentence; from twice the budget it ends at the next sentence end. <|end|> is only ever the model's
-    own choice (it ends the story).
-
-    Choices (once <|choices|> is written): at least MIN_CHOICES, none empty, none a copy of another, and a push away
-    from retracing an earlier choice word by word."""
 
     MIN_CHOICES = 2
     MAX_CHOICES = 5
@@ -162,8 +140,8 @@ class LengthControl:
     COPY = 1.5
     COPY_MAX = 8.0
     ECHO = 4.0
-    MIN_WORDS = 4      # a sentence of its own is at least this many words
-    GIVE_UP = 90       # tokens after which a narration with no sentence in it is left to the model
+    MIN_WORDS = 4
+    GIVE_UP = 90
 
     def __init__(self, tok, mode="medium", vars_weight=1.0):
         self.tok = tok
@@ -174,7 +152,7 @@ class LengthControl:
         self.stop = self.go_on + [tok.eos_id]
         first = min(tok.specials.values())
         self.silent = {i for i in range(first) if tok.vocab.id_to_token_decoded(i) == ""}
-        self.rep_free = False  # read by generate(): no repetition penalty while choices are written
+        self.rep_free = False
         text_like = {self.sep_id, tok.eos_id, tok.specials["<|num|>"], tok.specials["<|snum|>"]}
         self.no_specials = torch.tensor([i for i in range(first, 1 << 14) if i not in text_like], dtype=torch.long)
         self.pipes = torch.tensor([i for i in range(first) if "|" in tok.vocab.id_to_token_decoded(i)], dtype=torch.long)
@@ -185,7 +163,6 @@ class LengthControl:
         self.mode, self.budget = mode, LENGTHS[mode]
 
     def start_part(self, raw):
-        """Called before each generate() call of a turn; raw = the turn so far."""
         self.before = raw
         self.done = "<|" in raw
         self.offset = 0 if self.done else len(self.tok.encode(raw)[0])
@@ -276,13 +253,9 @@ class LengthControl:
         return logits
 
 
-# ----------------------------------------------------------------------------- sampling
-
 @torch.no_grad()
 def generate(lm, tok, seed, max_new, ctx, rng, temperature=0.8, top_k=50, top_p=0.95, rep_penalty=1.1,
              rep_window=128, adjust_logits=None):
-    """Yield (token id, packed number) as they are sampled. The seed is read once into a key/value cache; every step
-    then runs one token. Past `ctx` tokens the cache is rebuilt from the last 3/4 of it."""
     device = lm.head.weight.device
     seed_ids, seed_vals = tok.encode(seed)
     ids = [tok.eos_id] + seed_ids.tolist()
@@ -338,22 +311,19 @@ def generate(lm, tok, seed, max_new, ctx, rng, temperature=0.8, top_k=50, top_p=
         yield nxt, packed
 
 
-# ----------------------------------------------------------------------------- tales
-
 @dataclass
 class Tale:
-    """A tale: the whole of it as one string in the model's format (`history`), and what the engine knows of it."""
 
     history: str
     state: dict = field(default_factory=dict)
     title: str = ""
     length: str = "medium"
-    text: str = ""             # the last passage
+    text: str = ""
     choices: list = field(default_factory=list)
     ended: bool = False
     turns: int = 0
-    pending: str | None = None  # an opening not yet continued by the model
-    changes: list = field(default_factory=list)  # what the last turn did to the state: (name, old, new)
+    pending: str | None = None
+    changes: list = field(default_factory=list)
     created: float = field(default_factory=time.time)
 
     def to_dict(self):
@@ -364,7 +334,6 @@ class Tale:
         return cls(**d)
 
     def transcript(self):
-        """What a reader saw: [("story", text), ("player", text), ...]."""
         out = []
         for role, content in TURN.findall(self.history):
             if role == "player":
@@ -391,7 +360,6 @@ class StoryEngine:
 
     @torch.no_grad()
     def turn(self, tale, move, on_text=None, seed=None, max_new=450, max_parts=20, **sampling):
-        """Tell one turn of the tale (in place). move: the player's choice or own words (None for a new tale)."""
         if tale.pending is None:
             move = clean_input(move or "")
             if not move:
@@ -428,9 +396,9 @@ class StoryEngine:
             old = state.get(name)
             new = v if op == "set" else (old or 0) + v
             if not believable(name, text):
-                continue                       # "money -10" for being eaten alive
+                continue
             if name in ANNOUNCE and op == "set" and v != 0:
-                continue                       # shown amounts are earned or spent, not set
+                continue
             if not sane(name, old, new):
                 continue
             state[name] = new

bard_tokenizer.py

@@ -1,5 +1,4 @@
-"""The tokenizer of bard-small: tokenmonster (fiction-16000-strict-v1) plus the special tokens of the story format, and
-numbers kept apart: every run of digits is one <|num|>/<|snum|> token whose value travels in a parallel stream."""
+"""The tokenizer of bard-small."""
 import os
 import re
 
@@ -7,8 +6,8 @@ import numpy as np
 import tokenmonster
 
 EOS = "<|end|>"
-NUM_MARKERS = ["<|num|>", "<|snum|>"]  # a number, and a number preceded by one space
-MAX_DIGITS = 8                         # 10**8 - 1 < 2**27
+NUM_MARKERS = ["<|num|>", "<|snum|>"]
+MAX_DIGITS = 8
 LEN_SHIFT = 27
 VALUE_MASK = (1 << LEN_SHIFT) - 1
 NUMBER_RE = re.compile(r"( ?)([0-9]+)")
@@ -24,7 +23,6 @@ def unpack(packed: int) -> str:
 
 
 def split_numbers(text):
-    """-> a list of str (text) and (spaced, packed) items, in order."""
     out, pos = [], 0
     for m in NUMBER_RE.finditer(text):
         if m.start() > pos:
@@ -60,7 +58,6 @@ class BardTokenizer:
         return max(self.specials.values()) + 1
 
     def encode(self, text):
-        """-> (ids, values): token ids, and at each <|num|>/<|snum|> the packed number (else 0)."""
         plan = split_numbers(text)
         segs = [item for item in plan if isinstance(item, str)]
         seg_ids = iter(self.vocab.tokenize(segs) if segs else [])
@@ -89,8 +86,7 @@ class BardTokenizer:
                     run = []
                 parts.append((" " if i == self.snum_id else "") + unpack(v))
             elif i in self._special_text:
-                # a marker is its own text: decoded inside a run, the capitalisation state of the token before it
-                # can eat its first character
+
                 if run:
                     parts.append(self.vocab.decode(run))
                     run = []

configuration_bard.py

@@ -2,9 +2,6 @@ from transformers import PretrainedConfig
 
 
 class BardConfig(PretrainedConfig):
-    """Architecture of bard-small: a 36-layer, 512-wide decoder (parallel attention/FFN blocks, sliding-window attention
-    with a global layer every fourth, RoPE on the windowed layers, QK-norm, gated attention, SwiGLU, logit softcap) with
-    a dedicated number encoder/decoder."""
 
     model_type = "bard"
 
@@ -31,5 +28,5 @@ class BardConfig(PretrainedConfig):
         self.rope = rope
         self.parallel = parallel
         self.context = context
-        self.vars_weight = vars_weight  # the weight state changes had in training: the story engine undoes it
+        self.vars_weight = vars_weight
         super().__init__(**kwargs)

modeling_bard.py

@@ -32,13 +32,10 @@ class BardForStory(PreTrainedModel):
         self.post_init()
 
     def _init_weights(self, module):
-        pass  # the weights are loaded; the layers were initialised as they were trained
-
-    # ------------------------------------------------------------------------- the tokens
+        pass
 
     @property
     def tokenizer(self) -> BardTokenizer:
-        """The tokenizer (tokenmonster plus the story format's special tokens), loaded from the repo's own files."""
         if self._tokenizer is None:
             where = getattr(self.config, "_name_or_path", "") or ""
             if os.path.isdir(where):
@@ -51,16 +48,12 @@ class BardForStory(PreTrainedModel):
         return self._tokenizer
 
     def forward(self, input_ids, number_values=None, **kwargs):
-        """Logits for every position. `number_values` holds, at each <|num|>/<|snum|> token, the packed number (see
-        BardTokenizer.encode); zeros elsewhere."""
         if number_values is None:
             number_values = torch.zeros_like(input_ids)
-        cache = KVCache(self.lm.body)  # the cached path: ordinary attention with explicit masks, no flex_attention
+        cache = KVCache(self.lm.body)
         h = self.lm.hidden(input_ids, number_values, cache=cache)
         return CausalLMOutput(logits=softcap(torch.nn.functional.linear(h, self.lm.head.weight)))
 
-    # ------------------------------------------------------------------------- the game
-
     @property
     def engine(self) -> StoryEngine:
         if self._engine is None: