Upload Senryu-2.1-1B (EMA weights)

README.md

@@ -118,7 +118,7 @@ In Python, `tiny_gdn.senryu_chat.senryu_system_prompt()` returns it dated today.
 
 ## Default tools
 
-Every conversation under the default prompt offered these three tools, in this order. The model called them sparingly: about 4% of those conversations call a tool at all.
+Every conversation under the default prompt offered these three tools, in this order. Calls were rare: about 4% of those conversations call a tool at all.
 
 | Tool | Arguments | Result, as trained |
 |------|-----------|--------------------|
@@ -197,7 +197,7 @@ python inference.py
 
 ```bash
 python inference.py --prompt "Explain recursion in two sentences."
-python inference.py --prompt "What's in the news about fusion power this week?"
+python inference.py --prompt "What is 48213 * 9377?"
 python inference.py --tools none --system ""
 python inference.py --no-think --prompt "Reply in one sentence."
 python inference.py --ngram-device cpu
@@ -285,8 +285,8 @@ Only assistant turns are trained. The released EMA weights reach a loss of 1.433
 ## Limitations
 
 - **Scale.** 1.09B core parameters and no pretraining stage: expect thin world knowledge, factual errors and invented details.
-- **Page fetches.** Only 0.3% of the default-prompt training conversations call `fetch-page`, and in our tests the model answered questions about a linked page from memory, inventing its contents, instead of fetching it. Paste the text in when a page matters.
-- **Search.** Results depend on the backend, and the model can misread or over-trust what it retrieves. Check anything that matters.
+- **Tool calls.** The model calls tools much less often than it should. It used the calculator reliably in our tests, but asked for this week's news, told outright to search the web, or given a link to read, it answered from memory and invented the details in 9 of 11 tries. Only 1.5% of its default-prompt training conversations search the web and 0.3% fetch a page. Paste a page's text in when it matters.
+- **Search.** When it does search, it can misread or over-trust the results, and it sometimes repeats a search or lists results in a loop. Check anything that matters.
 - **Language.** Training data is mostly English.
 - **Hardware.** A CUDA GPU is required.
 

tiny_gdn/repetition.py

@@ -0,0 +1,460 @@
+"""Markdown-safe repetition control for TinyGDN sampling.
+
+Unigram presence over a window hammers newlines, list markers, heading
+hashes, table pipes, and backticks. This module:
+
+- applies the usual logit penalty only to **content** tokens that already
+  appeared at least twice in the window
+- bans the next token when it would complete a repeated n-gram **and**
+  that next token is content, so ``##``, ``- ``, ``|``, and ````` can
+  still repeat
+- applies DRY ("don't repeat yourself"): a content token that would extend
+  a verbatim copy of an earlier stretch of the same segment loses
+  ``multiplier * base ** (match - allowed_length)`` logits, so loops break
+  within a few tokens while ordinary word reuse goes untouched
+
+DRY works per segment. Special tokens (turn markers, ``<think>``,
+``</think>``) and the end of a tool call start a new one, so an answer may
+copy its draft from the reasoning. Copies are looked for in the current
+segment and in earlier assistant turns, never in system, user or tool text,
+so a reply may quote the prompt but not re-paste an earlier reply. Inside a
+code fence or a tool call DRY is off, and a digit is only penalized once its
+copy runs :data:`DRY_DIGIT_SLACK` tokens longer than other tokens need, so
+code, tool arguments and restated results come out as the model means them
+while a loop of numbers still breaks.
+"""
+
+from __future__ import annotations
+
+from collections import Counter
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+import torch
+from tokenizers import Tokenizer
+
+DEFAULT_REPETITION_NGRAM_SIZE = 0
+MIN_CONTENT_UNIGRAM_COUNT = 2
+
+DRY_SEQUENCE_BREAKERS = ("\n", ":", '"', "*")
+DRY_MAX_MATCH = 64
+DRY_MAX_EXPONENT = 64.0
+# A copy found in an earlier turn counts this many tokens shorter, so a reply
+# may reuse a heading or a term from an earlier one but not re-paste it.
+DRY_SOURCE_SLACK = 2
+# A digit needs a copy this many tokens longer before it is penalized.
+DRY_DIGIT_SLACK = 4
+CODE_FENCE = "```"
+# How this tokenizer family spells a tool-call tag: "<" or "</" (possibly
+# merged into the token before), then these pieces.
+TOOL_TAG_PIECES = ("tool", "_", "call", ">")
+TOOL_TAG_OPEN = "<tool_call>"
+TOOL_TAG_CLOSE = "</tool_call>"
+TURN_START = "<|im_start|>"
+TURN_END = "<|im_end|>"
+ASSISTANT_HEADER = f"{TURN_START}assistant\n"
+
+STRUCTURAL_CHARS = frozenset(
+    " \t\n\r\f\v"
+    "#*_`~[](){}|<>!.,;:?/+-=\\'\"“”‘’•–—…·"
+)
+SPECIAL_TOKEN_NAMES = frozenset(
+    {
+        "<think>",
+        "</think>",
+        "<tool_call>",
+        "</tool_call>",
+        "<tool_response>",
+        "</tool_response>",
+        "<tools>",
+        "</tools>",
+        "<code>",
+        "</code>",
+    }
+)
+
+
+def token_text_is_structural(decoded: str, piece: str | None = None) -> bool:
+    """True for whitespace, markdown/punctuation, and chat specials."""
+
+    raw = piece or ""
+    if raw.startswith("<|") or raw in SPECIAL_TOKEN_NAMES:
+        return True
+    if decoded.startswith("<|") or decoded in SPECIAL_TOKEN_NAMES:
+        return True
+    if decoded.strip() == "":
+        return True
+    body = decoded.strip()
+    if all(character in STRUCTURAL_CHARS for character in body):
+        return True
+    if len(body) <= 4 and body[:-1].isdecimal() and body[-1:] in ".)":
+        return True
+    return False
+
+
+_FLAGS: dict[int, tuple[bool, ...]] = {}
+
+
+def structural_token_flags(tokenizer: Tokenizer) -> tuple[bool, ...]:
+    """Per-id flags. Unused ids are treated as structural (never penalized)."""
+
+    flags: list[bool] = []
+    vocab_size = tokenizer.get_vocab_size()
+    for token_id in range(vocab_size):
+        piece = tokenizer.id_to_token(token_id)
+        if piece is None:
+            flags.append(True)
+            continue
+        decoded = tokenizer.decode([token_id], skip_special_tokens=False)
+        flags.append(token_text_is_structural(decoded, piece))
+    return tuple(flags)
+
+
+def cached_structural_token_flags(tokenizer: Tokenizer) -> tuple[bool, ...]:
+    key = id(tokenizer)
+    flags = _FLAGS.get(key)
+    if flags is None:
+        flags = structural_token_flags(tokenizer)
+        _FLAGS[key] = flags
+    return flags
+
+
+def content_token_mask(
+    tokenizer: Tokenizer, vocab_size: int, device: torch.device | str
+) -> torch.Tensor:
+    """``[vocab_size]`` bools, True for the content tokens the penalty may touch.
+
+    Model ids past the tokenizer's vocabulary count as structural.
+    """
+    flags = cached_structural_token_flags(tokenizer)
+    known = min(len(flags), vocab_size)
+    content = torch.zeros(vocab_size, dtype=torch.bool, device=device)
+    content[:known] = ~torch.tensor(flags[:known], dtype=torch.bool)
+    return content
+
+
+def _as_id_list(token_ids: Sequence[int] | torch.Tensor, window: int) -> list[int]:
+    if isinstance(token_ids, torch.Tensor):
+        values = token_ids.detach().tolist()
+    else:
+        values = list(token_ids)
+    if window > 0:
+        return values[-window:]
+    return values
+
+
+def content_ngram_banned(
+    recent: Sequence[int],
+    ngram_size: int,
+    structural: Sequence[bool],
+) -> set[int]:
+    """Next-token ids that would finish a previously seen content n-gram."""
+
+    if ngram_size < 2 or len(recent) < ngram_size:
+        return set()
+    vocab = len(structural)
+    seen: dict[tuple[int, ...], set[int]] = {}
+    prefix_len = ngram_size - 1
+    for index in range(len(recent) - ngram_size + 1):
+        nxt = recent[index + prefix_len]
+        if nxt < 0 or nxt >= vocab or structural[nxt]:
+            continue
+        prefix = tuple(recent[index : index + prefix_len])
+        seen.setdefault(prefix, set()).add(nxt)
+    current = tuple(recent[-prefix_len:])
+    return set(seen.get(current, ()))
+
+
+def apply_repetition_penalty(
+    logits: torch.Tensor,
+    token_ids: Sequence[int] | torch.Tensor,
+    *,
+    penalty: float,
+    window: int,
+    ngram_size: int = DEFAULT_REPETITION_NGRAM_SIZE,
+    structural: Sequence[bool],
+    min_unigram_count: int = MIN_CONTENT_UNIGRAM_COUNT,
+) -> torch.Tensor:
+    """In-place HF-style penalty on content unigrams, plus content n-gram bans."""
+
+    if penalty == 1.0:
+        return logits
+    recent = _as_id_list(token_ids, window)
+    if not recent:
+        return logits
+    vocab = len(structural)
+
+    if penalty != 1.0:
+        penalize = [
+            token_id
+            for token_id, count in Counter(recent).items()
+            if count >= min_unigram_count
+            and 0 <= token_id < vocab
+            and not structural[token_id]
+        ]
+        if penalize:
+            index = torch.tensor(penalize, dtype=torch.long, device=logits.device)
+            selected = logits[index]
+            logits.index_copy_(
+                0,
+                index,
+                torch.where(selected < 0, selected * penalty, selected / penalty),
+            )
+
+    if ngram_size >= 2:
+        banned = content_ngram_banned(recent, ngram_size, structural)
+        if banned:
+            index = torch.tensor(list(banned), dtype=torch.long, device=logits.device)
+            logits.index_fill_(0, index, float("-inf"))
+    return logits
+
+
+@dataclass(frozen=True)
+class DryVocabulary:
+    """Per-id flags DRY needs, for every id of the tokenizer.
+
+    ``tag_pieces`` holds the ids of :data:`TOOL_TAG_PIECES` when the tokenizer
+    spells tool-call tags over several tokens, and is empty otherwise.
+    ``assistant_header`` and ``turn_start``/``turn_end`` find earlier
+    assistant turns; without them only the current segment is searched.
+    """
+
+    breaker: tuple[bool, ...]
+    target: tuple[bool, ...]
+    digit: tuple[bool, ...]
+    boundary: tuple[bool, ...]
+    fence: tuple[bool, ...]
+    tool_open: tuple[bool, ...]
+    tool_close: tuple[bool, ...]
+    tag_open: tuple[bool, ...]
+    tag_close: tuple[bool, ...]
+    tag_pieces: tuple[int, ...]
+    assistant_header: tuple[int, ...] = ()
+    turn_start: int | None = None
+    turn_end: int | None = None
+
+
+@dataclass(frozen=True)
+class DryState:
+    """Where a context's DRY segment starts, and whether DRY is off at its end."""
+
+    start: int
+    fence: bool
+    tool: bool
+
+    @property
+    def exempt(self) -> bool:
+        return self.fence or self.tool
+
+
+def build_dry_vocabulary(tokenizer: Tokenizer) -> DryVocabulary:
+    structural = cached_structural_token_flags(tokenizer)
+    columns: dict[str, list[bool]] = {
+        name: []
+        for name in (
+            "breaker",
+            "target",
+            "digit",
+            "boundary",
+            "fence",
+            "tool_open",
+            "tool_close",
+            "tag_open",
+            "tag_close",
+        )
+    }
+    for token_id in range(tokenizer.get_vocab_size()):
+        piece = tokenizer.id_to_token(token_id)
+        decoded = "" if piece is None else tokenizer.decode([token_id], skip_special_tokens=False)
+        boundary = piece is not None and (piece.startswith("<|") or piece in ("<think>", "</think>"))
+        breaker = piece is None or boundary or any(mark in decoded for mark in DRY_SEQUENCE_BREAKERS)
+        columns["boundary"].append(boundary)
+        columns["breaker"].append(breaker)
+        columns["target"].append(piece is not None and not structural[token_id] and not breaker)
+        columns["digit"].append(decoded.strip().isdigit())
+        columns["fence"].append(decoded.count(CODE_FENCE) % 2 == 1)
+        columns["tool_open"].append(piece == TOOL_TAG_OPEN)
+        columns["tool_close"].append(piece == TOOL_TAG_CLOSE)
+        columns["tag_close"].append(decoded.endswith("</"))
+        columns["tag_open"].append(decoded.endswith("<") and not decoded.endswith("</"))
+    pieces = [tokenizer.token_to_id(piece) for piece in TOOL_TAG_PIECES]
+    tag_pieces = (
+        ()
+        if tokenizer.token_to_id(TOOL_TAG_OPEN) is not None or any(piece is None for piece in pieces)
+        else tuple(int(piece) for piece in pieces if piece is not None)
+    )
+    turn_start = tokenizer.token_to_id(TURN_START)
+    turn_end = tokenizer.token_to_id(TURN_END)
+    has_turns = turn_start is not None and turn_end is not None
+    return DryVocabulary(
+        **{name: tuple(values) for name, values in columns.items()},
+        tag_pieces=tag_pieces,
+        assistant_header=(
+            tuple(tokenizer.encode(ASSISTANT_HEADER, add_special_tokens=False).ids) if has_turns else ()
+        ),
+        turn_start=turn_start if has_turns else None,
+        turn_end=turn_end if has_turns else None,
+    )
+
+
+_DRY_VOCABULARIES: dict[int, DryVocabulary] = {}
+
+
+def cached_dry_vocabulary(tokenizer: Tokenizer) -> DryVocabulary:
+    key = id(tokenizer)
+    vocabulary = _DRY_VOCABULARIES.get(key)
+    if vocabulary is None:
+        vocabulary = build_dry_vocabulary(tokenizer)
+        _DRY_VOCABULARIES[key] = vocabulary
+    return vocabulary
+
+
+def _flag(flags: tuple[bool, ...], token: int) -> bool:
+    return 0 <= token < len(flags) and flags[token]
+
+
+def dry_tool_tag(tokens: Sequence[int], position: int, vocabulary: DryVocabulary) -> tuple[bool, bool]:
+    """Whether the token at ``position`` completes an opening or a closing tool-call tag."""
+
+    token = tokens[position]
+    if _flag(vocabulary.tool_open, token) or _flag(vocabulary.tool_close, token):
+        return _flag(vocabulary.tool_open, token), _flag(vocabulary.tool_close, token)
+    size = len(vocabulary.tag_pieces)
+    if size == 0 or position < size:
+        return False, False
+    if tuple(tokens[position - size + 1 : position + 1]) != vocabulary.tag_pieces:
+        return False, False
+    before = tokens[position - size]
+    return _flag(vocabulary.tag_open, before), _flag(vocabulary.tag_close, before)
+
+
+def dry_state(tokens: Sequence[int], vocabulary: DryVocabulary) -> DryState:
+    """The DRY segment state after ``tokens``, as the graph sampler tracks it token by token."""
+
+    begin = 0
+    for position in range(len(tokens) - 1, -1, -1):
+        if _flag(vocabulary.boundary, tokens[position]):
+            begin = position + 1
+            break
+    start, fence, tool = begin, False, False
+    for position in range(begin, len(tokens)):
+        opens, closes = dry_tool_tag(tokens, position, vocabulary)
+        fence ^= _flag(vocabulary.fence, tokens[position])
+        if opens:
+            tool = True
+        elif closes:
+            start, fence, tool = position + 1, False, False
+    return DryState(start=start, fence=fence, tool=tool)
+
+
+def dry_sources(tokens: Sequence[int], vocabulary: DryVocabulary) -> list[tuple[int, int]]:
+    """``[begin, end)`` spans of the assistant turns before the last one in ``tokens``."""
+
+    header = vocabulary.assistant_header
+    if not header or vocabulary.turn_start is None:
+        return []
+    current = next(
+        (position for position in range(len(tokens) - 1, -1, -1) if tokens[position] == vocabulary.turn_start),
+        None,
+    )
+    if current is None:
+        return []
+    spans: list[tuple[int, int]] = []
+    begin: int | None = None
+    for position in range(current):
+        token = tokens[position]
+        if begin is not None and (token == vocabulary.turn_end or token == vocabulary.turn_start):
+            spans.append((begin, position))
+            begin = None
+        if token == vocabulary.turn_start and tuple(tokens[position : position + len(header)]) == header:
+            begin = position + len(header)
+    if begin is not None and begin < current:
+        spans.append((begin, current))
+    return spans
+
+
+def dry_match_lengths(
+    tokens: Sequence[int],
+    *,
+    start: int,
+    breaker: tuple[bool, ...],
+    sources: Sequence[tuple[int, int]] = (),
+    lower: int = 0,
+) -> dict[int, int]:
+    """For each token that followed an earlier copy of the context's tail, the longest such copy.
+
+    The tail stays inside the segment that begins at ``start``; copies may lie
+    in that segment or, from ``lower`` on, in the ``sources`` spans, where
+    they count :data:`DRY_SOURCE_SLACK` tokens shorter.
+    """
+
+    if len(tokens) - start < 1:
+        return {}
+    last = tokens[-1]
+    if _flag(breaker, last):
+        return {}
+    eligible = bytearray(len(tokens))
+    for begin, end in sources:
+        eligible[max(begin, lower) : max(end, lower)] = b"\x01" * (max(end, lower) - max(begin, lower))
+    eligible[start:] = b"\x01" * (len(tokens) - start)
+    lengths: dict[int, int] = {}
+    for index in range(len(tokens) - 1):
+        if not eligible[index] or tokens[index] != last:
+            continue
+        match = 1
+        while match < DRY_MAX_MATCH and index - match >= 0 and eligible[index - match] and len(tokens) - 1 - match >= start:
+            recent = tokens[-1 - match]
+            if tokens[index - match] != recent or _flag(breaker, recent):
+                break
+            match += 1
+        if index < start:
+            match -= DRY_SOURCE_SLACK
+        if match < 1:
+            continue
+        following = tokens[index + 1]
+        lengths[following] = max(match, lengths.get(following, 0))
+    return lengths
+
+
+def apply_dry_penalty(
+    logits: torch.Tensor,
+    tokens: Sequence[int],
+    *,
+    multiplier: float,
+    base: float,
+    allowed_length: int,
+    window: int,
+    vocabulary: DryVocabulary,
+) -> torch.Tensor:
+    """In-place DRY penalty on ``[vocab]`` logits for the context ``tokens``."""
+
+    if multiplier == 0:
+        return logits
+    tokens = list(tokens)
+    state = dry_state(tokens, vocabulary)
+    if state.exempt:
+        return logits
+    lower = 0 if window <= 0 else max(len(tokens) - window, 0)
+    lengths = dry_match_lengths(
+        tokens,
+        start=max(state.start, lower),
+        breaker=vocabulary.breaker,
+        sources=dry_sources(tokens, vocabulary),
+        lower=lower,
+    )
+    vocab = logits.shape[-1]
+    effective = {
+        token: match - (DRY_DIGIT_SLACK if _flag(vocabulary.digit, token) else 0)
+        for token, match in lengths.items()
+        if token < vocab and _flag(vocabulary.target, token)
+    }
+    penalized = {
+        token: multiplier * base ** min(match - allowed_length, DRY_MAX_EXPONENT)
+        for token, match in effective.items()
+        if match >= allowed_length
+    }
+    if penalized:
+        index = torch.tensor(list(penalized), dtype=torch.long, device=logits.device)
+        amount = torch.tensor(list(penalized.values()), dtype=logits.dtype, device=logits.device)
+        logits.index_add_(0, index, -amount)
+    return logits