Why does a chatbot forget?

"""Token budgets, sliding-window truncation, and what re-sending history costs.

Python standard library only. Run `python3 show_numbers.py` to print every
number the lesson quotes.
"""


def count_tokens(text: str) -> int:
    """Toy counter: one word (a run of non-space characters) counts as one token.

    Real tokenizers cut text into pieces they learned from data, so their counts
    differ from word counts, and differ from each other. To budget a real model,
    pass that model's own tokenizer as the `counter` argument of `truncate`.
    """
    return len(text.split())


def total_tokens(messages, counter=count_tokens, per_message=0) -> int:
    return sum(counter(m["content"]) + per_message for m in messages)


def truncate(messages, limit, keep_system=True, counter=count_tokens, per_message=0):
    """Sliding window: keep the newest messages that fit in `limit` tokens.

    messages     list of {"role": ..., "content": ...}, oldest first.
    keep_system  pin the leading system messages: they are always kept.
    per_message  extra tokens each message costs in the chat format (role
                 markers and the like). Take it from your model's template.

    The kept part is a run of the newest messages with no holes: the walk
    goes from newest to oldest and stops at the first message that does not
    fit. If the walk cut anything, the leading assistant replies of the run
    are then dropped, because a reply whose question was cut makes no sense
    on its own. Raises ValueError if the pinned messages alone exceed `limit`,
    or if no newest message is left, because answering without it is silent
    failure.
    """
    def cost(m):
        return counter(m["content"]) + per_message

    rest = list(messages)
    pinned = []
    if keep_system:
        n = 0
        while n < len(rest) and rest[n]["role"] == "system":
            n += 1
        pinned, rest = rest[:n], rest[n:]

    budget = limit - sum(cost(m) for m in pinned)
    if budget < 0:
        raise ValueError("the pinned system messages alone exceed the limit")

    kept = []
    for m in reversed(rest):
        if cost(m) > budget:
            break
        kept.append(m)
        budget -= cost(m)
    kept.reverse()
    if len(kept) < len(rest):                      # the walk cut something
        while kept and kept[0]["role"] == "assistant":
            kept.pop(0)                            # a reply whose question was cut
    if rest and not kept:
        raise ValueError("the newest message does not fit in the limit")
    return pinned + kept


def tokens_read(n_replies: int, tokens_per_message: int, window=None) -> int:
    """Total tokens the model reads while answering `n_replies` user messages.

    The app re-sends the history before every reply. Every message, yours or
    the bot's, has `tokens_per_message` tokens. With `window` set, the history
    is cut with `truncate` first. A reply's own tokens are not counted as read.
    """
    history, total = [], 0
    for _ in range(n_replies):
        history.append({"role": "user", "content": " ".join(["w"] * tokens_per_message)})
        sent = history if window is None else truncate(history, window)
        total += total_tokens(sent)
        history.append({"role": "assistant", "content": " ".join(["w"] * tokens_per_message)})
    return total


def attention_pairs(n: int) -> int:
    """Token pairs one full-attention head compares while reading an n-token
    prompt: token i is compared with the i tokens up to and including itself."""
    return n * (n + 1) // 2

"""Memory the attention cache needs per token, from a real model configuration.

Python standard library only. The saved file qwen3_8b_config.json is a copy of
https://huggingface.co/Qwen/Qwen3-8B/resolve/main/config.json (Apache-2.0
licence, see the model card). `fetch_config()` downloads the live file.
"""
import json
from pathlib import Path
from urllib.request import urlopen

CONFIG_URL = "https://huggingface.co/Qwen/Qwen3-8B/resolve/main/config.json"
SNAPSHOT = Path(__file__).with_name("qwen3_8b_config.json")
BYTES_PER_VALUE = {"float32": 4, "float16": 2, "bfloat16": 2}


def load_config(path=SNAPSHOT) -> dict:
    return json.loads(Path(path).read_text())


def fetch_config(url=CONFIG_URL) -> dict:
    with urlopen(url, timeout=30) as response:
        return json.load(response)


def kv_bytes_per_token(cfg: dict) -> int:
    """2 (one key, one value) x layers x key/value heads x head size x bytes per value.

    Assumes every layer keeps keys and values for all tokens (full attention).
    Models with sliding-window layers or latent attention store less or differ."""
    head_dim = cfg.get("head_dim") or cfg["hidden_size"] // cfg["num_attention_heads"]
    kv_heads = cfg.get("num_key_value_heads", cfg["num_attention_heads"])
    dtype = cfg.get("torch_dtype") or cfg["dtype"]
    return 2 * cfg["num_hidden_layers"] * kv_heads * head_dim * BYTES_PER_VALUE[dtype]


def kv_bytes(cfg: dict, tokens: int) -> int:
    return kv_bytes_per_token(cfg) * tokens


def kv_bytes_at_max_context(cfg: dict) -> int:
    return kv_bytes(cfg, cfg["max_position_embeddings"])

{
  "architectures": [
    "Qwen3ForCausalLM"
  ],
  "attention_bias": false,
  "attention_dropout": 0.0,
  "bos_token_id": 151643,
  "eos_token_id": 151645,
  "head_dim": 128,
  "hidden_act": "silu",
  "hidden_size": 4096,
  "initializer_range": 0.02,
  "intermediate_size": 12288,
  "max_position_embeddings": 40960,
  "max_window_layers": 36,
  "model_type": "qwen3",
  "num_attention_heads": 32,
  "num_hidden_layers": 36,
  "num_key_value_heads": 8,
  "rms_norm_eps": 1e-06,
  "rope_scaling": null,
  "rope_theta": 1000000,
  "sliding_window": null,
  "tie_word_embeddings": false,
  "torch_dtype": "bfloat16",
  "transformers_version": "4.51.0",
  "use_cache": true,
  "use_sliding_window": false,
  "vocab_size": 151936
}

"""Print every number the lesson computes (numbers from a cited page are marked "source").

    python3 show_numbers.py            uses the saved copy of the model config
    python3 show_numbers.py --fetch    downloads the live config and compares it
"""
import sys

from context_window import attention_pairs, count_tokens, tokens_read, total_tokens, truncate
from kv_memory import CONFIG_URL, fetch_config, kv_bytes, kv_bytes_at_max_context, kv_bytes_per_token, load_config

KIB, MIB, GIB = 1024, 1024**2, 1024**3

cfg = load_config()
if "--fetch" in sys.argv:
    print("live config equals saved copy:", fetch_config() == cfg)

print("== Qwen3-8B config:", CONFIG_URL)
for key in ("num_hidden_layers", "num_key_value_heads", "num_attention_heads", "head_dim",
            "torch_dtype", "max_position_embeddings"):
    print(f"  {key} = {cfg[key]}")
per_token = kv_bytes_per_token(cfg)
print(f"bytes per token = {per_token:,} ({per_token / KIB:g} KiB)")
for tokens in (1_000, 4_096, 32_768, 40_960):
    b = kv_bytes(cfg, tokens)
    print(f"  {tokens:>6,} tokens: {b:>14,} bytes = {b / MIB:,.3f} MiB = {b / GIB:.4f} GiB")
at_max = kv_bytes_at_max_context(cfg)
print(f"at max context: {at_max:,} bytes = {at_max / GIB:g} GiB = {at_max / 1e9:.2f} GB (decimal), "
      f"'about 6 gigabytes' in the opening")
mha = per_token * cfg["num_attention_heads"] // cfg["num_key_value_heads"]
print(f"if every query head had its own key/value pair: {mha:,} bytes per token "
      f"(x{cfg['num_attention_heads'] // cfg['num_key_value_heads']}), "
      f"{mha * cfg['max_position_embeddings'] / GIB:g} GiB at max context")
print(f"weights, about 8.2e9 parameters x 2 bytes = {8.2e9 * 2 / 1e9:.1f} GB (decimal; 8.2B is from the model card)")

print("\n== Window sizes (source: Google token guide lists 1,048,576; Qwen3-8B card says 32,768 natively)")
print(f"1,048,576 tokens at 60 to 80 words per 100 tokens (source ratio): "
      f"{1_048_576 * 0.6:,.0f} to {1_048_576 * 0.8:,.0f} words, "
      f"rounded to the nearest 10,000: {round(1_048_576 * 0.6, -4):,.0f} to {round(1_048_576 * 0.8, -4):,.0f}")

print("\n== Worked example (toy tokens = words)")
chat = [
    {"role": "system", "content": "Reply in French."},
    {"role": "user", "content": "My cat is Biscuit."},
    {"role": "assistant", "content": "Joli nom."},
    {"role": "user", "content": "I like tea."},
    {"role": "assistant", "content": "Moi aussi."},
    {"role": "user", "content": "What is my cat called?"},
]
print("sizes:", [count_tokens(m["content"]) for m in chat], "total", sum(count_tokens(m["content"]) for m in chat))
for pin in (True, False):
    kept = truncate(chat, 12, keep_system=pin)
    print(f"limit 12, keep_system={pin}: kept {[m['content'] for m in kept]}, {total_tokens(kept)} tokens")

print("\n== Re-reading history (every message 50 tokens)")
for n in (10, 20):
    text = 2 * n * 50
    read = tokens_read(n, 50)
    print(f"{n} replies: text written {text:,}, tokens read {read:,}, ratio {read / text:g}")
for k in range(1, 6):
    print(f"  reply {k}: {2 * k - 1} messages, reads {tokens_read(k, 50) - tokens_read(k - 1, 50):,}, "
          f"running total {tokens_read(k, 50):,}")
for n in range(1, 11):
    print(f"  reply {n}: reads {tokens_read(n, 50) - tokens_read(n - 1, 50):,} uncapped, "
          f"{tokens_read(n, 50, window=450) - tokens_read(n - 1, 50, window=450):,} with a 450-token window")
print(f"10 replies, window 450: {tokens_read(10, 50, window=450):,} read vs {tokens_read(10, 50):,} uncapped")
w40 = tokens_read(40, 50, window=450)
print(f"40 replies, window 450: {w40:,} read vs {tokens_read(40, 50):,} uncapped; "
      f"50 x (1+3+5+7+9) = {50 * (1 + 3 + 5 + 7 + 9):,}, then 450 x 35 = {450 * 35:,}")
print(f"window 400 holds 8 messages of 50 but sends 7: replies 5 to 10 read "
      f"{tokens_read(5, 50, window=400) - tokens_read(4, 50, window=400)} each; "
      f"10 replies read {tokens_read(10, 50, window=400):,}")

print("\n== Figure: seven messages and a 190-token window")
figure = [30, 20, 40, 25, 35, 30, 20]  # token counts, drawn at 2 px per token
print(f"messages {figure}; messages 3 to 7 = {sum(figure[2:])} tokens, plus 40 tokens of room for the reply "
      f"= a {sum(figure[2:]) + 40}-token window")


print("\n== Attention pairs per head and layer")
for n in (1_000, 2_000, 40_960):
    print(f"  {n:,} tokens: {attention_pairs(n):,} pairs; ratio to half the size: "
          f"{attention_pairs(n) / attention_pairs(n // 2):.3f}")

print("\n== Proof it works")
print(f"tokens_read(3 replies, 10 tokens each) = {tokens_read(3, 10)}")

print("\n== Check yourself")
q3 = [("system", 2), ("user", 3), ("assistant", 1), ("user", 6), ("assistant", 4), ("user", 3)]
q3_chat = [{"role": r, "content": " ".join(["w"] * n)} for r, n in q3]
kept = truncate(q3_chat, 10)
print(f"Q3 sizes {[n for _, n in q3]}, limit 10: kept sizes {[count_tokens(m['content']) for m in kept]} "
      f"roles {[m['role'] for m in kept]}, {total_tokens(kept)} tokens; pinned 2, budget 8; newest 3 leaves 5; "
      f"4-token reply leaves 1; 6 does not fit")
q4 = {"num_hidden_layers": 28, "num_key_value_heads": 4, "head_dim": 128, "torch_dtype": "float16",
      "hidden_size": 2048, "num_attention_heads": 16}
b4 = kv_bytes_per_token(q4)
print(f"Q4 per token {b4:,} bytes ({b4 / KIB:g} KiB); 8,192 tokens {kv_bytes(q4, 8_192):,} bytes "
      f"= {kv_bytes(q4, 8_192) / MIB:g} MiB")
print(f"Q5 12 replies of 40-token messages: {tokens_read(12, 40):,} read (40 x 12^2 = {40 * 12**2:,}); "
      f"window 200: {tokens_read(12, 40, window=200):,} = 40 + 120 + 200 + 9 x 200 = "
      f"{40 + 120 + 200 + 9 * 200:,}; first three replies read {tokens_read(3, 40, window=200):,}; "
      f"other 9 replies read 9 x 200 = {9 * 200:,}")

import random

import pytest

from context_window import attention_pairs, count_tokens, total_tokens, tokens_read, truncate


def msg(role, text):
    return {"role": role, "content": text}


# The lesson's worked example. Word counts, worked out by hand: 3, 4, 2, 3, 2, 5.
CHAT = [
    msg("system", "Reply in French."),
    msg("user", "My cat is Biscuit."),
    msg("assistant", "Joli nom."),
    msg("user", "I like tea."),
    msg("assistant", "Moi aussi."),
    msg("user", "What is my cat called?"),
]


def test_counter_is_words():
    assert [count_tokens(m["content"]) for m in CHAT] == [3, 4, 2, 3, 2, 5]
    assert count_tokens("") == 0


def test_worked_example_pinned_system_drops_the_cat():
    """Oldest dropped first: the message with the cat's name goes, the pinned system message stays.
    The walk keeps messages 6 and 5, stops at message 4 (it does not fit), then drops the leading
    reply 5, whose question 4 was cut. Message 3 (2 tokens) would fit in the spare room, but keeping
    it would leave a hole."""
    kept = truncate(CHAT, 12)
    assert kept == [CHAT[0], CHAT[5]]
    assert total_tokens(kept) == 3 + 5 == 8


def test_worked_example_without_pin_drops_the_system_message_too():
    """The walk keeps messages 6, 5, 4 and 3 (12 tokens); reply 3 has lost its question 2, so it goes."""
    kept = truncate(CHAT, 12, keep_system=False)
    assert kept == CHAT[3:]
    assert total_tokens(kept) == 3 + 2 + 5 == 10


def test_a_reply_whose_question_was_cut_is_dropped_but_a_whole_chat_is_sent_whole():
    """Guards the old behaviour, where the window could start with a reply and no question."""
    chat = [msg("assistant", "Hello there"), msg("user", "Hi"), msg("assistant", "Hello again"), msg("user", "Bye")]
    assert truncate(chat, 6) == chat                    # nothing cut: the opening reply stays
    assert truncate(chat, 4) == chat[1:]                # cut, but the run starts with a question: nothing more to drop
    assert truncate(chat, 3) == [chat[3]]               # run = [reply, Bye]; the reply lost its question "Hi"
    two_replies = [msg("user", "q q q q q q"), msg("assistant", "a"), msg("assistant", "b"), msg("user", "c")]
    assert truncate(two_replies, 3) == [two_replies[3]]  # both leading replies go, not just one
    with pytest.raises(ValueError, match="newest"):      # only a reply fits, and its question was cut
        truncate([msg("user", "q q q q q q"), msg("assistant", "a")], 3)


def test_everything_fits_so_nothing_is_dropped():
    assert truncate(CHAT, 19) == CHAT
    assert truncate(CHAT, 1000) == CHAT


def test_no_holes_a_big_message_blocks_older_small_ones():
    """A greedy 'skip what does not fit' would keep A and leave a gap in the conversation."""
    chat = [msg("system", "S"), msg("user", "A"), msg("user", "b b b b b"), msg("user", "C"), msg("user", "D")]
    assert truncate(chat, 6) == [chat[0], chat[3], chat[4]]


def test_a_different_counter_changes_what_fits():
    """The limit is only meaningful under the counter that measured it."""
    chat = [msg("user", "aaaa"), msg("user", "bb"), msg("user", "cc")]
    assert truncate(chat, 5) == chat                     # words: 1 + 1 + 1 = 3 fits
    assert truncate(chat, 5, counter=len) == chat[1:]    # characters: 4 + 2 + 2 = 8 > 5, so the oldest goes
    assert truncate(chat, 8, counter=len) == chat


def test_per_message_overhead_counts_against_the_limit():
    chat = [msg("user", "a b"), msg("user", "c d"), msg("user", "e f")]
    assert truncate(chat, 9, per_message=1) == chat      # 3 + 3 + 3
    assert truncate(chat, 8, per_message=1) == chat[1:]  # one more word's worth would not fit


def test_errors_are_loud():
    with pytest.raises(ValueError, match="pinned"):
        truncate(CHAT, 2)
    with pytest.raises(ValueError, match="newest"):
        truncate(CHAT, 7)   # 3 pinned + 5 for the newest message = 8 > 7
    assert truncate([], 5) == []


def test_never_over_the_limit_and_greedy_from_the_newest():
    """Randomized: within the limit, a pinned prefix plus a run of the newest messages, no leading reply
    after a cut, and maximal."""
    rng = random.Random(0)
    ok = checked_maximal = replies_before_cut = 0
    for _ in range(400):
        n_sys = rng.choice([0, 1, 2])
        chat = [msg("system", " ".join("s" * rng.randint(1, 4))) for _ in range(n_sys)]
        chat += [msg(rng.choice(["user", "assistant"]), " ".join("w" * rng.randint(1, 9)))
                 for _ in range(rng.randint(1, 15))]
        limit, extra = rng.randint(5, 60), rng.randint(0, 3)
        try:
            kept = truncate(chat, limit, per_message=extra)
        except ValueError:
            continue
        ok += 1
        assert total_tokens(kept, per_message=extra) <= limit
        rest = chat[n_sys:]
        n_kept = len(kept) - n_sys
        assert kept[:n_sys] == chat[:n_sys]
        assert kept[n_sys:] == rest[len(rest) - n_kept:]                      # a run of the newest
        if n_kept < len(rest):
            assert kept[n_sys]["role"] != "assistant"                         # a cut never leaves a reply first
            older = rest[len(rest) - n_kept - 1]
            if older["role"] == "assistant":
                replies_before_cut += 1                                       # may have been dropped as an orphan
            else:                                                             # maximal: the next older one cannot fit
                checked_maximal += 1
                assert total_tokens(kept, per_message=extra) + count_tokens(older["content"]) + extra > limit
    assert ok > 100 and checked_maximal > 20 and replies_before_cut > 20      # the checks above were not vacuous


def test_rereading_history_grows_with_the_square_of_the_chat():
    """tokens_read must equal tokens_per_message x n^2, so a 10x longer chat re-reads 100x as much."""
    assert tokens_read(3, 10) == 10 + 30 + 50 == 90      # by hand: 1, 3, 5 messages of 10 tokens
    for n in (1, 2, 10, 20):
        assert tokens_read(n, 50) == 50 * n * n
    assert tokens_read(10, 50) == 5000 and tokens_read(20, 50) == 20000


def test_a_window_stops_the_growth():
    # window 30, messages of 10: replies read 10, 30, then 30 and 30 again (50 and 70 are cut to 30)
    assert tokens_read(3, 10, window=30) == 10 + 30 + 30 == 70
    n = 40
    assert tokens_read(n, 50, window=450) == 50 * (1 + 3 + 5 + 7 + 9) + 450 * (n - 5) == 17_000
    assert tokens_read(n, 50) == 80_000
    # lesson's last check-yourself question: 12 replies of 40-token messages
    assert tokens_read(12, 40) == 5_760
    assert tokens_read(12, 40, window=200) == 40 + 120 + 200 + 200 * 9 == 2_160


def test_a_window_of_eight_messages_sends_seven_because_of_the_orphan_reply():
    """400 tokens hold 8 messages of 50, but the oldest of the 8 is a reply whose question was cut."""
    assert tokens_read(10, 50, window=400) == 50 * (1 + 3 + 5 + 7) + 350 * 6 == 2_900


def test_attention_pairs_by_hand():
    assert attention_pairs(4) == 1 + 2 + 3 + 4 == 10
    assert attention_pairs(1000) == 500_500
    assert attention_pairs(2000) == 2_001_000            # doubling the prompt: about 4x the pairs

from kv_memory import (BYTES_PER_VALUE, kv_bytes, kv_bytes_at_max_context, kv_bytes_per_token,
                       load_config)


def test_qwen3_8b_matches_the_hand_computed_value():
    """2 x 36 layers x 8 kv heads x 128 x 2 bytes, worked out by hand: 147,456 bytes (144 KiB)."""
    cfg = load_config()
    assert cfg["num_hidden_layers"] == 36 and cfg["num_key_value_heads"] == 8
    assert cfg["head_dim"] == 128 and cfg["torch_dtype"] == "bfloat16"
    assert kv_bytes_per_token(cfg) == 2 * 36 * 8 * 128 * 2 == 147_456
    assert kv_bytes_per_token(cfg) == 144 * 1024


def test_memory_at_the_context_sizes_the_lesson_quotes():
    cfg = load_config()
    assert cfg["max_position_embeddings"] == 40_960
    assert kv_bytes_at_max_context(cfg) == 6_039_797_760 == int(5.625 * 1024**3)
    assert kv_bytes(cfg, 32_768) == 4_831_838_208 == int(4.5 * 1024**3)
    assert kv_bytes(cfg, 4_096) == 576 * 1024**2
    assert kv_bytes(cfg, 2 * 4_096) == 2 * kv_bytes(cfg, 4_096)      # linear in tokens


def test_uses_kv_heads_not_attention_heads():
    """The common slip: with 32 query heads instead of 8 key/value heads the answer is 4x too big."""
    cfg = load_config()
    assert cfg["num_attention_heads"] == 32
    assert kv_bytes_per_token(cfg) * 4 == 589_824
    assert kv_bytes_per_token({**cfg, "num_key_value_heads": 32}) == 589_824 != kv_bytes_per_token(cfg)


def test_explicit_head_dim_beats_hidden_size_division():
    """Made-up config: hidden 1024 / 16 heads = 64, but the file says head_dim 128. 2 x 2 x 1 x 128 x 2 = 1,024."""
    cfg = {"num_hidden_layers": 2, "num_key_value_heads": 1, "num_attention_heads": 16,
           "hidden_size": 1024, "head_dim": 128, "torch_dtype": "float16"}
    assert kv_bytes_per_token(cfg) == 1024


def test_missing_kv_heads_means_one_per_attention_head():
    """Made-up config without grouped-query attention: 2 x 1 x 4 x (8 / 4 = 2) x 4 bytes = 64."""
    cfg = {"num_hidden_layers": 1, "num_attention_heads": 4, "hidden_size": 8, "torch_dtype": "float32"}
    assert kv_bytes_per_token(cfg) == 64
    assert BYTES_PER_VALUE["float32"] == 4


def test_dtype_key_is_the_fallback_when_torch_dtype_is_missing():
    """Newer config files say `dtype`. 1 layer, 1 kv head, head size 8 / 1 = 8, 2 bytes: 2 x 8 x 2 = 32."""
    cfg = {"num_hidden_layers": 1, "num_key_value_heads": 1, "num_attention_heads": 1,
           "hidden_size": 8, "dtype": "bfloat16"}
    assert kv_bytes_per_token(cfg) == 32


def test_lessons_check_yourself_cache_question():
    """28 layers, 4 key/value heads, head size 128, 2 bytes per value: 57,344 bytes per token, 448 MiB at 8,192."""
    cfg = {"num_hidden_layers": 28, "num_key_value_heads": 4, "num_attention_heads": 16,
           "head_dim": 128, "torch_dtype": "float16"}
    assert kv_bytes_per_token(cfg) == 2 * 28 * 4 * 128 * 2 == 57_344
    assert kv_bytes(cfg, 8_192) == 469_762_048 == 448 * 1024**2