Inside the Transformer

# script_id: day_005_transformers_attention/paste_everything
# Install: uv add tiktoken · Env: ORBIT_DATA (optional) · No API key: runs offline
import json
import os
from pathlib import Path

import tiktoken

DATA_DIR = Path(os.environ.get("ORBIT_DATA", "datasets/orbit"))  # atlas/config.py (Day 2)
ENC = tiktoken.get_encoding("o200k_base")  # approximate for GPT-6 (Day 4)
INPUT_PER_1M = 0.10  # gpt-6-luna USD, as of 2026-11, verify at https://developers.openai.com/api/docs/pricing

def count_tokens(text: str) -> int:
    return len(ENC.encode(text, disallowed_special=()))

docs = sum(count_tokens(p.read_text()) for p in sorted((DATA_DIR / "docs").glob("*.md")))
with open(DATA_DIR / "issues" / "issues.jsonl", encoding="utf-8") as f:
    issues = sum(count_tokens(f"{i['title']}\n\n{i['body']}") for i in map(json.loads, f))

total = docs + issues
per_request = total * INPUT_PER_1M / 1_000_000
print(f"docs {docs:,} + issues {issues:,} = {total:,} tokens per request")
print(f"input cost ${per_request:.4f}/request, ${per_request * 2_000:,.2f}/day at 2,000 questions")

docs 65,308 + issues 78,339 = 143,647 tokens per request
input cost $0.0144/request, $28.73/day at 2,000 questions

# script_id: day_005_transformers_attention/attention_by_hand
# Install: nothing (stdlib) · Env: none · Toy scores chosen by hand, not taken from a model
def show(sentence: str, scores: dict[str, float]) -> None:
    total = sum(scores.values())  # scale so the weights sum to 1
    weights = {word: s / total for word, s in scores.items()}
    print(sentence)
    for word, w in sorted(weights.items(), key=lambda kv: -kv[1]):
        print(f"  it -> {word:<8} {w:4.2f} {'#' * round(w * 40)}")
    print(f"  'it' now mostly means '{max(weights, key=weights.get)}'\n")

show("The worker lost its lease because it ran too long.",
     {"worker": 5.5, "lease": 2.5, "ran": 1.0, "long": 0.5, "lost": 0.5})
# Turn the knob: change one verb and the weights move
show("The worker lost its lease because it expired.",
     {"worker": 1.5, "lease": 6.5, "expired": 1.5, "lost": 0.5})

The worker lost its lease because it ran too long.
  it -> worker   0.55 ######################
  it -> lease    0.25 ##########
  it -> ran      0.10 ####
  it -> long     0.05 ##
  it -> lost     0.05 ##
  'it' now mostly means 'worker'

The worker lost its lease because it expired.
  it -> lease    0.65 ##########################
  it -> worker   0.15 ######
  it -> expired  0.15 ######
  it -> lost     0.05 ##
  'it' now mostly means 'lease'

# script_id: day_005_transformers_attention/attention_pairs
# Install: uv add tiktoken · Env: none · Offline
import tiktoken

enc = tiktoken.get_encoding("o200k_base")  # approximate for GPT-6 (Day 4)
ids = enc.encode("The worker lost its lease because it ran too long.")
print(len(ids), "tokens:", [enc.decode([i]) for i in ids])

# Each position attends to itself and every earlier position: n*(n+1)/2 pairs, in every layer
for n in (len(ids), 1_000, 8_000, 128_000, 922_000):
    print(f"{n:>9,} tokens -> {n * (n + 1) // 2:>17,} attention pairs per layer")

11 tokens: ['The', ' worker', ' lost', ' its', ' lease', ' because', ' it', ' ran', ' too', ' long', '.']
       11 tokens ->                66 attention pairs per layer
    1,000 tokens ->           500,500 attention pairs per layer
    8,000 tokens ->        32,004,000 attention pairs per layer
  128,000 tokens ->     8,192,064,000 attention pairs per layer
  922,000 tokens ->   425,042,461,000 attention pairs per layer

# script_id: day_005_transformers_attention/context_budget
"""Does a request fit? A static context-budget calculator for gpt-6-luna.

Install: uv add tiktoken
Env:     ATLAS_MODEL, ORBIT_DATA (both optional). No API key: runs offline.
"""
import os
from pathlib import Path

import tiktoken

MODEL = os.environ.get("ATLAS_MODEL", "gpt-6-luna")               # atlas/config.py (Day 2)
DATA_DIR = Path(os.environ.get("ORBIT_DATA", "datasets/orbit"))
_ENC = tiktoken.get_encoding("o200k_base")                          # approximate (Day 4)

def count_tokens(text: str) -> int:
    """Approximate offline count, built on Day 4's atlas/tokens.py."""
    return len(_ENC.encode(text, disallowed_special=()))

# gpt-6-luna, as of 2026-11, verify at https://developers.openai.com/api/docs/models/gpt-6-luna
CONTEXT_WINDOW = 1_050_000         # input + output (reasoning included) per request
MAX_INPUT_TOKENS = 922_000         # the input on its own may not exceed this
MAX_OUTPUT_TOKENS = 128_000
LONG_CONTEXT_THRESHOLD = 272_000   # above this, the whole request bills at higher rates

# script_id: day_005_transformers_attention/context_budget
def context_budget(parts: dict[str, str], *, window: int = CONTEXT_WINDOW,
                   output_reserve: int = 32_000) -> dict[str, int]:
    """Tokens per part plus totals; fits/long_context are bools (a bool is an int in Python)."""
    if output_reserve > MAX_OUTPUT_TOKENS:
        raise ValueError(f"output_reserve {output_reserve:,} > max output {MAX_OUTPUT_TOKENS:,}")
    used = {name: count_tokens(text) for name, text in parts.items()}
    input_tokens = sum(used.values())
    total = input_tokens + output_reserve
    return {**used, "input": input_tokens, "reserve": output_reserve, "total": total,
            "remaining": window - total,
            "fits": input_tokens <= MAX_INPUT_TOKENS and total <= window,
            "long_context": input_tokens > LONG_CONTEXT_THRESHOLD}

docs = "\n\n".join(p.read_text() for p in sorted((DATA_DIR / "docs").glob("*.md")))
parts = {"instructions": "You are Atlas, an assistant for Orbit maintainers. "
                         "Answer from the docs and cite the doc name.",
         "docs": docs, "question": "How do leases and heartbeats interact?"}
b = context_budget(parts)
for name in [*parts, "reserve"]:
    print(f"{name:<13}{b[name]:>10,}")
print(f"{'total':<13}{b['total']:>10,} of {CONTEXT_WINDOW:,} · fits={b['fits']} "
      f"· long_context={b['long_context']}")

instructions         21
docs             65,314
question              8
reserve          32,000
total            97,343 of 1,050,000 · fits=True · long_context=False

# script_id: day_005_transformers_attention/context_budget
# USD per 1M tokens, Standard tier, prompts up to 272K input tokens (Day 4's table).
# As of 2026-11, verify at https://developers.openai.com/api/docs/pricing
PRICES_PER_1M: dict[str, dict[str, float]] = {
    "gpt-6-luna":    {"input": 0.10, "cached_input": 0.01, "output": 0.50},
    "gpt-6.1-terra": {"input": 1.00, "cached_input": 0.05, "output": 5.00},
}

def estimate_cost(input_tokens: int, output_tokens: int, model: str = MODEL) -> float:
    """USD for one call, including the long-context surcharge. Same signature as Day 4."""
    if model not in PRICES_PER_1M:
        raise KeyError(f"no price for {model!r}; add it to PRICES_PER_1M")
    price = PRICES_PER_1M[model]
    long = input_tokens > LONG_CONTEXT_THRESHOLD  # the multiplier applies to the WHOLE request
    in_rate = price["input"] * (2 if long else 1)  # cached_input doubles too; it's used from Day 18
    out_rate = price["output"] * (1.5 if long else 1)
    return (input_tokens * in_rate + output_tokens * out_rate) / 1_000_000

print(f"this request: ${estimate_cost(b['input'], 2_000):.4f} (2,000 output tokens assumed)")
for n in (271_000, 273_000):  # input-only, to isolate the jump
    flag = "  <- surcharge" if n > LONG_CONTEXT_THRESHOLD else ""
    print(f"{n:,} input tokens -> ${estimate_cost(n, 0):.4f}{flag}")

this request: $0.0075 (2,000 output tokens assumed)
271,000 input tokens -> $0.0271
273,000 input tokens -> $0.0546  <- surcharge

# script_id: day_005_transformers_attention/kv_cache_sim
# Install: nothing (stdlib) · Env: none · A counting model for intuition, not a benchmark
def positions_processed(prompt: int, output: int, *, kv_cache: bool) -> int:
    """Token positions pushed through the layers to generate `output` tokens."""
    if kv_cache:  # prefill reads the prompt once, then one new position per output token
        return prompt + output
    return sum(prompt + step for step in range(output))  # re-read everything, every step

for prompt in (2_000, 32_000, 128_000):
    no_cache = positions_processed(prompt, 500, kv_cache=False)
    cached = positions_processed(prompt, 500, kv_cache=True)
    print(f"prompt {prompt:>7,}: no cache {no_cache:>11,} · KV cache {cached:>8,}"
          f" · prefill share {prompt / cached:.1%}")

prompt   2,000: no cache   1,124,750 · KV cache    2,500 · prefill share 80.0%
prompt  32,000: no cache  16,124,750 · KV cache   32,500 · prefill share 98.5%
prompt 128,000: no cache  64,124,750 · KV cache  128,500 · prefill share 99.6%

# script_id: day_005_transformers_attention/overflow_guard
"""Check the budget before you send, and handle the 400 if it happens anyway.

Install: uv add "openai>=3.46,<4" tiktoken
Env:     OPENAI_API_KEY (only when a request is actually sent), ATLAS_MODEL (optional)
"""
import os

import openai
import tiktoken
from openai import OpenAI

MODEL = os.environ.get("ATLAS_MODEL", "gpt-6-luna")
MAX_INPUT_TOKENS, LONG_CONTEXT_THRESHOLD = 922_000, 272_000  # verify on the model page
_ENC = tiktoken.get_encoding("o200k_base")

class ContextTooLarge(Exception):
    """Raised before sending, so an oversized prompt never costs a request."""

def preflight(text: str, *, allow_long_context: bool = False) -> int:
    n = len(_ENC.encode(text, disallowed_special=()))
    if n > MAX_INPUT_TOKENS * 0.95:  # the count is approximate, so keep 5% headroom
        raise ContextTooLarge(f"~{n:,} input tokens > max input; retrieve or trim first")
    if n > LONG_CONTEXT_THRESHOLD and not allow_long_context:
        raise ContextTooLarge(f"~{n:,} input tokens crosses 272K: the whole request bills at 2x/1.5x")
    return n

# script_id: day_005_transformers_attention/overflow_guard
def ask(question: str, context: str) -> str:
    preflight(context + question)
    client = OpenAI()  # Day 12 replaces this with one shared client
    try:
        resp = client.responses.create(model=MODEL, instructions="Answer from the context only.",
                                       input=f"{context}\n\nQuestion: {question}")
    except openai.BadRequestError as e:  # 400, e.g. input over the window with truncation disabled
        print("400 from the API:", e.code, e.message)
        raise
    return resp.output_text

if __name__ == "__main__":
    huge = "Orbit worker log line. " * 60_000  # about 300K tokens of fake logs
    try:
        ask("Why did the worker stop?", huge)
    except ContextTooLarge as e:
        print("refused before sending:", e)

refused before sending: ~300,006 input tokens crosses 272K: the whole request bills at 2x/1.5x

# script_id: day_005_transformers_attention/needle_test
"""Multi-needle test with distractors over Orbit's docs and issues.

Install: uv add "openai>=3.46,<4" tiktoken
Env:     OPENAI_API_KEY; ATLAS_MODEL, ORBIT_DATA (optional). Cost: about $0.02 on gpt-6-luna.
"""
import json
import os
import time
from pathlib import Path

import tiktoken
from openai import OpenAI

MODEL = os.environ.get("ATLAS_MODEL", "gpt-6-luna")
DATA_DIR = Path(os.environ.get("ORBIT_DATA", "datasets/orbit"))
ENC = tiktoken.get_encoding("o200k_base")  # approximate for GPT-6 (Day 4)

def estimate_cost(input_tokens: int, output_tokens: int) -> float:
    # gpt-6-luna below 272K (all sizes here are); as of 2026-11, verify at
    # https://developers.openai.com/api/docs/pricing
    return (input_tokens * 0.10 + output_tokens * 0.50) / 1_000_000

def load_sources() -> list[str]:
    """Docs first, then issue texts: about 144K tokens in total."""
    texts = [p.read_text() for p in sorted((DATA_DIR / "docs").glob("*.md"))]
    with open(DATA_DIR / "issues" / "issues.jsonl", encoding="utf-8") as f:
        texts += [f"{i['title']}\n\n{i['body']}" for i in map(json.loads, f)]
    return texts

def build_haystack(docs: list[str], target_tokens: int) -> str:
    """Join whole texts until the target is reached, then cut the token stream to size."""
    ids: list[int] = []
    for text in docs:
        ids += ENC.encode(text + "\n\n", disallowed_special=())
        if len(ids) >= target_tokens:
            break
    return ENC.decode(ids[:target_tokens])

# script_id: day_005_transformers_attention/needle_test
NEEDLES = ["Release note: the Orbit 3.2.1 release codename is Kestrel.",
           "Release note: the Orbit 3.2.1 release party was held in Lisbon.",
           "Release note: the Orbit 3.2.1 release mascot is a teal otter."]
DISTRACTORS = ["Release note: the Orbit 3.2.0 release codename was Kestral.",
               "Release note: the Orbit 3.1.0 release party was held in Porto.",
               "Release note: the Orbit 3.2.1 beta mascot was an amber lynx."]
QUESTION = ("Using only the release notes in the text above: what are the codename, the party "
            "city and the mascot of the final Orbit 3.2.1 release? Answer in three short lines.")

def insert_needles(text: str, needles: list[str], depths: list[float]) -> str:
    """Insert each needle at a fractional depth (0.0 = start), at the previous paragraph break."""
    for needle, depth in sorted(zip(needles, depths), key=lambda nd: -nd[1]):  # deepest first
        pos = text.rfind("\n\n", 0, int(len(text) * depth))
        cut = pos + 2 if pos >= 0 else 0
        text = text[:cut] + needle + "\n\n" + text[cut:]
    return text

def score_answer(answer: str, needles: list[str]) -> int:
    """How many of these facts appear in the answer, keyed on each fact's last word."""
    keys = [n.rstrip(".").split()[-1].lower() for n in needles]  # kestrel, lisbon, otter
    return sum(key in answer.lower() for key in keys)

# script_id: day_005_transformers_attention/needle_test
SIZES = [8_000, 32_000, 128_000]
NEEDLE_DEPTHS, TRAP_DEPTHS = [0.1, 0.5, 0.9], [0.3, 0.6, 0.8]
ASSUMED_OUTPUT = 2_000  # reasoning + answer; reasoning bills as output (Day 6)
MAX_SPEND = 0.10

def run_trial(client: OpenAI, size: int, depths: list[float]) -> dict:
    hay = insert_needles(build_haystack(load_sources(), size), DISTRACTORS, TRAP_DEPTHS)
    hay = insert_needles(hay, NEEDLES, depths)
    start = time.perf_counter()
    resp = client.responses.create(model=MODEL, input=f"{hay}\n\n{QUESTION}")
    u = resp.usage
    return {"size": size, "found": score_answer(resp.output_text, NEEDLES),
            "traps": score_answer(resp.output_text, DISTRACTORS),
            "latency": time.perf_counter() - start,
            "usd": estimate_cost(u.input_tokens, u.output_tokens)}

if __name__ == "__main__":
    estimate = sum(estimate_cost(s + 200, ASSUMED_OUTPUT) for s in SIZES)
    print(f"estimated spend ${estimate:.4f} for {len(SIZES)} calls")
    if estimate > MAX_SPEND:
        raise SystemExit(f"refusing: estimate is over the ${MAX_SPEND:.2f} cap")
    client = OpenAI()
    print("   size | found | traps | latency s |      $")
    for size in SIZES:
        r = run_trial(client, size, NEEDLE_DEPTHS)
        print(f"{r['size']:>7,} |  {r['found']}/3  |   {r['traps']}   | "
              f"{r['latency']:>9.1f} | {r['usd']:.4f}")

estimated spend $0.0199 for 3 calls
   size | found | traps | latency s |      $
  8,000 |  3/3  |   0   |       3.2 | 0.0013
 32,000 |  3/3  |   0   |       4.9 | 0.0037
128,000 |  3/3  |   0   |      11.6 | 0.0134