Upload folder using huggingface_hub

src/add_special_tokens.py

@@ -0,0 +1,87 @@
+"""
+Add special tokens to the tokenizer for the search agent format.
+
+New tokens:
+  Chat format:
+    <|system|>    <|end|>
+    <|user|>      <|end|>
+    <|assistant|> <|end|>
+
+  Search agent actions:
+    <|search|>      — marks a search query
+    <|result|>      — marks a search result
+    <|evidence|>    — marks curated evidence
+    <|reasoning|>   — marks internal reasoning
+    <|finish|>      — marks end of search (return evidence to caller)
+
+The model learns to produce these tokens during SFT. The agent harness
+at inference time parses them to drive the search loop.
+"""
+
+import json
+import os
+
+from tokenizers import Tokenizer
+from tokenizers.models import BPE
+
+PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+TOKENIZER_DIR = os.path.join(PROJECT_DIR, "tokenizer")
+TOKENIZER_PATH = os.path.join(TOKENIZER_DIR, "tokenizer.json")
+AGENT_TOKENIZER_PATH = os.path.join(TOKENIZER_DIR, "tokenizer_agent.json")
+
+SPECIAL_TOKENS = [
+    "<|system|>",
+    "<|user|>",
+    "<|assistant|>",
+    "<|search|>",
+    "<|result|>",
+    "<|evidence|>",
+    "<|reasoning|>",
+    "<|finish|>",
+    "<|end|>",
+]
+
+
+def main():
+    print(f"Loading tokenizer from {TOKENIZER_PATH}...")
+    tokenizer = Tokenizer.from_file(TOKENIZER_PATH)
+    original_vocab_size = tokenizer.get_vocab_size()
+    print(f"  Original vocab size: {original_vocab_size}")
+
+    # Add special tokens
+    print(f"\nAdding {len(SPECIAL_TOKENS)} special tokens...")
+    for token in SPECIAL_TOKENS:
+        tokenizer.add_tokens([token])
+        print(f"  Added: {token}")
+
+    new_vocab_size = tokenizer.get_vocab_size()
+    print(f"\nNew vocab size: {new_vocab_size} (added {new_vocab_size - original_vocab_size})")
+
+    # Save
+    tokenizer.save(AGENT_TOKENIZER_PATH)
+    print(f"\nAgent tokenizer saved to {AGENT_TOKENIZER_PATH}")
+
+    # Verify the tokens encode correctly
+    print("\nVerification:")
+    test_text = "<|system|>You are a search agent.<|end|><|user|>Where is auth defined?<|end|><|assistant|><|search|>auth middleware<|end|>"
+    encoded = tokenizer.encode(test_text)
+    print(f"  Input: {test_text[:80]}...")
+    print(f"  IDs: {encoded.ids[:30]}")
+
+    # Verify each special token has a unique ID
+    vocab = tokenizer.get_vocab()
+    print("\nSpecial token IDs:")
+    for token in SPECIAL_TOKENS:
+        tid = vocab.get(token)
+        print(f"  {token:20s} -> {tid}")
+
+    # Also save a mapping file for the agent harness
+    mapping_path = os.path.join(TOKENIZER_DIR, "special_tokens.json")
+    token_ids = {token: vocab[token] for token in SPECIAL_TOKENS}
+    with open(mapping_path, "w") as f:
+        json.dump({"special_tokens": SPECIAL_TOKENS, "token_ids": token_ids}, f, indent=2)
+    print(f"\nToken mapping saved to {mapping_path}")
+
+
+if __name__ == "__main__":
+    main()

src/chunk_corpus.py

@@ -0,0 +1,232 @@
+"""
+Chunk the curated corpus into searchable code units.
+
+Splits code into semantically meaningful chunks:
+  - Python: functions, classes, top-level blocks
+  - JS/TS: functions, classes, export blocks
+  - Rust: fn, impl, struct, enum, trait blocks
+  - Go: func, type, struct blocks
+  - C/C++: function definitions, struct/typedef blocks
+
+Each chunk gets:
+  - id: unique identifier
+  - language: detected language
+  - name: extracted name (function/class name)
+  - type: function/class/struct/etc
+  - code: the raw code text
+  - filepath: synthetic path (derived from doc index)
+  - start_line, end_line: line range within the doc
+
+Output: data/chunks.jsonl
+"""
+
+import json
+import os
+import re
+
+PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+CURATED_PATH = os.path.join(PROJECT_DIR, "data", "corpus_curated.txt")
+CHUNKS_PATH = os.path.join(PROJECT_DIR, "data", "chunks.jsonl")
+
+# ─── Language detection ──────────────────────────────────────────────────────
+LANG_PATTERNS = {
+    "python": re.compile(r"^(def |class |import |from \S+ import |if __name__|@)", re.M),
+    "js_ts":  re.compile(r"^(function |const |let |class |export |import |async function|interface |type \w+ =)", re.M),
+    "rust":   re.compile(r"^(fn |pub fn |impl |struct |enum |trait |mod |use |pub struct|pub enum|macro_rules!)", re.M),
+    "go":     re.compile(r"^(func |package |type \w+ struct|import )", re.M),
+    "c_cpp":  re.compile(r"^(#include|#define|#ifndef|#if |#endif|typedef |struct \w+|class \w+|void |int |char |size_t |ngx_)", re.M),
+}
+
+# ─── Chunk patterns per language ─────────────────────────────────────────────
+# Each pattern matches the START of a chunk. We split on these.
+CHUNK_STARTS = {
+    "python": re.compile(r"^(def |class |@|if __name__)", re.M),
+    "js_ts":  re.compile(r"^(function |const |let |class |export |async function|interface |type \w+ =)", re.M),
+    "rust":   re.compile(r"^(fn |pub fn |impl |struct |enum |trait |mod |macro_rules!|pub struct|pub enum)", re.M),
+    "go":     re.compile(r"^(func |type \w+ struct|type \w+ interface)", re.M),
+    "c_cpp":  re.compile(r"^(static |void |int |char |size_t |ngx_|typedef |struct \w+ \{|#define|#if|#ifndef)", re.M),
+}
+
+# ─── Name extraction patterns ────────────────────────────────────────────────
+NAME_PATTERNS = [
+    (re.compile(r"def (\w+)"), "function"),
+    (re.compile(r"class (\w+)"), "class"),
+    (re.compile(r"fn (\w+)"), "function"),
+    (re.compile(r"pub fn (\w+)"), "function"),
+    (re.compile(r"struct (\w+)"), "struct"),
+    (re.compile(r"enum (\w+)"), "enum"),
+    (re.compile(r"trait (\w+)"), "trait"),
+    (re.compile(r"impl (\w+)"), "impl"),
+    (re.compile(r"func (\w+)"), "function"),
+    (re.compile(r"type (\w+) struct"), "struct"),
+    (re.compile(r"function (\w+)"), "function"),
+    (re.compile(r"typedef struct (\w+)"), "typedef"),
+    (re.compile(r"#define (\w+)"), "macro"),
+    (re.compile(r"(ngx_\w+)\s*\("), "function"),
+]
+
+
+def detect_language(doc: str) -> str | None:
+    for lang, pat in LANG_PATTERNS.items():
+        if len(pat.findall(doc)) >= 2:
+            return lang
+    return None
+
+
+def extract_name(code: str) -> tuple[str, str]:
+    """Extract the name and type from a code chunk."""
+    for pat, typ in NAME_PATTERNS:
+        m = pat.search(code)
+        if m:
+            return m.group(1), typ
+    return "unknown", "block"
+
+
+def chunk_document(doc: str, lang: str, doc_idx: int) -> list[dict]:
+    """Split a document into chunks based on language-specific patterns."""
+    lines = doc.split("\n")
+    n_lines = len(lines)
+
+    # Find all chunk start positions
+    start_pat = CHUNK_STARTS.get(lang)
+    if start_pat is None:
+        # Fallback: treat whole doc as one chunk
+        return [{
+            "id": f"doc_{doc_idx}_chunk_0",
+            "language": lang,
+            "name": "block",
+            "type": "block",
+            "code": doc,
+            "filepath": f"src/doc_{doc_idx}.txt",
+            "start_line": 1,
+            "end_line": n_lines,
+        }]
+
+    starts = [(m.start(), m.group()) for m in start_pat.finditer(doc)]
+
+    if not starts:
+        # No pattern matches — treat whole doc as one chunk
+        return [{
+            "id": f"doc_{doc_idx}_chunk_0",
+            "language": lang,
+            "name": "block",
+            "type": "block",
+            "code": doc,
+            "filepath": f"src/doc_{doc_idx}.txt",
+            "start_line": 1,
+            "end_line": n_lines,
+        }]
+
+    # Add doc start if first match isn't at position 0
+    if starts[0][0] > 0:
+        starts.insert(0, (0, ""))
+
+    chunks = []
+    for i, (start_pos, _) in enumerate(starts):
+        end_pos = starts[i + 1][0] if i + 1 < len(starts) else len(doc)
+        chunk_code = doc[start_pos:end_pos].strip()
+
+        # Skip tiny chunks (< 30 chars)
+        if len(chunk_code) < 30:
+            continue
+
+        # Skip huge chunks (> 8000 chars — split them)
+        if len(chunk_code) > 8000:
+            # Split on blank lines
+            sub_parts = re.split(r"\n\n+", chunk_code)
+            for j, part in enumerate(sub_parts):
+                if len(part.strip()) < 30:
+                    continue
+                name, typ = extract_name(part)
+                start_line = doc[:start_pos].count("\n") + 1 + sum(p.count("\n") + 2 for p in sub_parts[:j])
+                chunks.append({
+                    "id": f"doc_{doc_idx}_chunk_{i}_{j}",
+                    "language": lang,
+                    "name": name,
+                    "type": typ,
+                    "code": part.strip(),
+                    "filepath": f"src/doc_{doc_idx}.txt",
+                    "start_line": start_line,
+                    "end_line": start_line + part.count("\n"),
+                })
+            continue
+
+        name, typ = extract_name(chunk_code)
+        start_line = doc[:start_pos].count("\n") + 1
+
+        chunks.append({
+            "id": f"doc_{doc_idx}_chunk_{i}",
+            "language": lang,
+            "name": name,
+            "type": typ,
+            "code": chunk_code,
+            "filepath": f"src/doc_{doc_idx}.txt",
+            "start_line": start_line,
+            "end_line": start_line + chunk_code.count("\n"),
+        })
+
+    return chunks
+
+
+def main():
+    print(f"Loading curated corpus from {CURATED_PATH}...")
+    with open(CURATED_PATH, "r", encoding="utf-8") as f:
+        text = f.read()
+    print(f"  Corpus size: {len(text) / 1e6:.2f} MB")
+
+    # Split into documents
+    docs = re.split(r"\n{3,}", text)
+    print(f"  Documents: {len(docs):,}")
+
+    # Chunk each document
+    print("Chunking documents...")
+    all_chunks = []
+    lang_counts = {}
+    type_counts = {}
+
+    for i, doc in enumerate(docs):
+        doc = doc.strip()
+        if len(doc) < 50:
+            continue
+        lang = detect_language(doc)
+        if lang is None:
+            continue
+
+        chunks = chunk_document(doc, lang, i)
+        for chunk in chunks:
+            all_chunks.append(chunk)
+            lang_counts[lang] = lang_counts.get(lang, 0) + 1
+            type_counts[chunk["type"]] = type_counts.get(chunk["type"], 0) + 1
+
+    print(f"\nTotal chunks: {len(all_chunks):,}")
+    print(f"  By language: {lang_counts}")
+    print(f"  By type: {type_counts}")
+
+    # Size distribution
+    sizes = [len(c["code"]) for c in all_chunks]
+    sizes.sort()
+    print(f"  Chunk size: min={sizes[0]}, median={sizes[len(sizes)//2]}, max={sizes[-1]}, mean={sum(sizes)//len(sizes)}")
+
+    # Write chunks
+    with open(CHUNKS_PATH, "w", encoding="utf-8") as f:
+        for chunk in all_chunks:
+            f.write(json.dumps(chunk) + "\n")
+
+    print(f"\nChunks written to {CHUNKS_PATH}")
+
+    # Show a few samples
+    print("\n" + "=" * 60)
+    print("SAMPLE CHUNKS")
+    print("=" * 60)
+    for i in [0, len(all_chunks) // 4, len(all_chunks) // 2, len(all_chunks) - 1]:
+        c = all_chunks[i]
+        print(f"\n--- {c['id']} | {c['language']} | {c['type']} | {c['name']} ---")
+        print(f"    File: {c['filepath']}:{c['start_line']}-{c['end_line']}")
+        print(f"    Size: {len(c['code'])} chars")
+        print(c["code"][:300])
+        if len(c["code"]) > 300:
+            print("...")
+
+
+if __name__ == "__main__":
+    main()

src/convert_bf16.py

@@ -0,0 +1,24 @@
+"""Convert the latest checkpoint to bf16 to halve file size for Colab upload."""
+import torch
+import json
+import os
+
+def main():
+    ckpt = torch.load("checkpoints/latest.pt", map_location="cpu", weights_only=False)
+    print(f"Original checkpoint: step={ckpt['step']}, loss={ckpt['loss']}")
+    print(f"  Embedding shape: {ckpt['model_state_dict']['token_embedding.weight'].shape}")
+
+    # Convert all float32 tensors to bf16
+    state = ckpt["model_state_dict"]
+    for key in state:
+        if state[key].dtype == torch.float32:
+            state[key] = state[key].to(torch.bfloat16)
+
+    ckpt["model_state_dict"] = state
+    out_path = "checkpoints/latest_bf16.pt"
+    torch.save(ckpt, out_path)
+    size_mb = os.path.getsize(out_path) / 1e6
+    print(f"bf16 checkpoint saved: {out_path} ({size_mb:.1f} MB)")
+
+if __name__ == "__main__":
+    main()

src/curate.py

@@ -0,0 +1,239 @@
+"""
+Curate the raw corpus into a small, ultra-high-quality subset.
+
+Splits the corpus into documents, scores each on quality heuristics,
+and keeps only the top-tier real source code. Produces data/corpus_curated.txt.
+
+Quality signals:
+  - Real source code (NOT RST docs, YAML configs, JSON data, prose, etc.)
+  - Balanced delimiters (braces/parens/brackets roughly match)
+  - Has function/class definitions at line starts (not embedded in prose)
+  - Has comments or docstrings (documentation density)
+  - Reasonable length (200–50k chars)
+  - Low repetition (unique line ratio)
+  - Clean ASCII (no encoding garbage)
+  - High code-to-prose ratio
+"""
+
+import os
+import re
+
+PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+CORPUS_PATH = os.path.join(PROJECT_DIR, "data", "corpus.txt")
+CURATED_PATH = os.path.join(PROJECT_DIR, "data", "corpus_curated.txt")
+
+# ─── Code structure patterns (must be at line start, not in prose) ───────────
+CODE_SIGNATURES = {
+    "python":   re.compile(r"^(def |class |import |from \S+ import |if __name__|@|elif |except |finally )", re.M),
+    "js_ts":    re.compile(r"^(function |const |let |var |class |export |import |async function|interface |type \w+ =)", re.M),
+    "rust":     re.compile(r"^(fn |pub fn |impl |struct |enum |trait |mod |use |pub struct|pub enum|macro_rules!)", re.M),
+    "go":       re.compile(r"^(func |package |import |type \w+ struct|var |const )", re.M),
+    "c_cpp":    re.compile(r"^(#include|#define|#ifndef|#ifdef|#if |#endif|typedef |struct \w+|class \w+)", re.M),
+    "java":     re.compile(r"^(public |private |protected |class \w+|import |package )", re.M),
+}
+
+# ─── Documentation / non-code exclusion patterns ─────────────────────────────
+RST_DIRECTIVE = re.compile(r"^\.\. [\w-]+::", re.M)          # .. c:function::
+RST_LABEL = re.compile(r"^\.\. _[\w-]+:", re.M)              # .. _faq-cache:
+RST_UNDERLINE = re.compile(r"^[-=~^\"']{3,}\s*$", re.M)      # ----, ====, ~~~~
+RST_CROSSREF = re.compile(r":[\w]+:`", re.M)                 # :func:`name`
+XML_START = re.compile(r"^\s*<\?xml|<!DOCTYPE|<\w+ xmlns", re.M)
+MARKDOWN_HEADER = re.compile(r"^#{1,6}\s+\w", re.M)          # # Header
+
+# Prose detection: lines that look like English sentences
+PROSE_LINE = re.compile(r"^[A-Z][a-z]+ .* [a-z]+\.$", re.M)  # "The quick brown fox."
+
+
+def is_documentation(doc: str) -> bool:
+    """Check if a document is RST/Markdown documentation rather than code."""
+    # Strong RST signals
+    if RST_DIRECTIVE.search(doc) or RST_LABEL.search(doc):
+        return True
+    if RST_CROSSREF.search(doc):
+        return True
+    # Section underlines (RST/Markdown)
+    if RST_UNDERLINE.findall(doc):
+        return True
+    # XML
+    if XML_START.search(doc):
+        return True
+    # Markdown headers
+    if MARKDOWN_HEADER.findall(doc) and len(MARKDOWN_HEADER.findall(doc)) > 3:
+        return True
+    return False
+
+
+def detect_language(doc: str) -> str | None:
+    """Detect if a document is real source code and which language."""
+    if is_documentation(doc):
+        return None
+
+    for lang, pat in CODE_SIGNATURES.items():
+        matches = pat.findall(doc)
+        if len(matches) >= 3:  # at least 3 code-structure keywords
+            return lang
+    return None
+
+
+def score_document(doc: str) -> tuple[float, str | None]:
+    """Score a document 0.0–1.0 on quality. Returns (score, language)."""
+    lines = doc.split("\n")
+    n_lines = len(lines)
+    if n_lines < 5:
+        return 0.0, None
+
+    length = len(doc)
+    if length < 200 or length > 100_000:
+        return 0.0, None
+
+    # Must be real code
+    lang = detect_language(doc)
+    if lang is None:
+        return 0.0, None
+
+    score = 0.0
+
+    # ─── Base: is real code ──────────────────────────────────────────────────
+    score += 0.20
+
+    # ─── Length quality (sweet spot: 500–20000 chars) ────────────────────────
+    if 500 <= length <= 20000:
+        score += 0.15
+    elif 200 <= length <= 50000:
+        score += 0.08
+
+    # ─── Delimiter balance ───────────────────────────────────────────────────
+    braces = doc.count("{") - doc.count("}")
+    parens = doc.count("(") - doc.count(")")
+    brackets = doc.count("[") - doc.count("]")
+    total_delims = doc.count("{") + doc.count("(") + doc.count("[")
+    if total_delims > 0:
+        imbalance = abs(braces) + abs(parens) + abs(brackets)
+        balance_ratio = 1.0 - (imbalance / max(total_delims, 1))
+        score += 0.15 * max(balance_ratio, 0.0)
+
+    # ─── Code structure density ──────────────────────────────────────────────
+    struct_count = sum(len(p.findall(doc)) for p in CODE_SIGNATURES.values())
+    struct_density = min(struct_count / max(n_lines, 1) * 8, 1.0)
+    score += 0.15 * struct_density
+
+    # ─── Comment density (sweet spot: 5–30%) ─────────────────────────────────
+    comment_lines = 0
+    for line in lines:
+        s = line.strip()
+        if s.startswith("#") or s.startswith("//") or s.startswith("/*") \
+           or s.startswith("*") or s.startswith('"""') or s.startswith("'''") \
+           or s.startswith("///") or s.startswith("//!"):
+            comment_lines += 1
+    comment_ratio = comment_lines / max(n_lines, 1)
+    if 0.05 <= comment_ratio <= 0.30:
+        score += 0.15
+    elif 0.02 <= comment_ratio <= 0.40:
+        score += 0.08
+
+    # ─── Low repetition ──────────────────────────────────────────────────────
+    unique_lines = len(set(lines))
+    unique_ratio = unique_lines / max(n_lines, 1)
+    if unique_ratio > 0.6:
+        score += 0.10
+    elif unique_ratio > 0.4:
+        score += 0.05
+    else:
+        score -= 0.05
+
+    # ─── Clean ASCII ─────────────────────────────────────────────────────────
+    non_ascii = sum(1 for c in doc if ord(c) > 127)
+    if non_ascii / max(length, 1) < 0.01:
+        score += 0.05
+
+    # ─── Indentation quality ─────────────────────────────────────────────────
+    indented = sum(1 for l in lines if l.startswith("    ") or l.startswith("\t"))
+    if indented > 0 and indented / max(n_lines, 1) > 0.1:
+        score += 0.05
+
+    # ─── Penalize high prose ratio ───────────────────────────────────────────
+    prose_lines = len(PROSE_LINE.findall(doc))
+    prose_ratio = prose_lines / max(n_lines, 1)
+    if prose_ratio > 0.15:
+        score -= 0.15  # too much prose = documentation
+
+    return min(max(score, 0.0), 1.0), lang
+
+
+def main():
+    print(f"Loading corpus from {CORPUS_PATH}...")
+    with open(CORPUS_PATH, "r", encoding="utf-8") as f:
+        text = f.read()
+    print(f"  Corpus size: {len(text) / 1e6:.1f} MB")
+
+    # Split into documents (3+ consecutive newlines = boundary)
+    docs = re.split(r"\n{3,}", text)
+    print(f"  Total documents: {len(docs):,}")
+
+    # Score all documents
+    print("Scoring documents...")
+    scored = []
+    for i, doc in enumerate(docs):
+        s, lang = score_document(doc)
+        if s > 0:
+            scored.append((s, i, doc, lang))
+
+    print(f"  Documents passing initial filter: {len(scored):,}")
+
+    # Sort by score descending
+    scored.sort(key=lambda x: x[0], reverse=True)
+
+    # Keep top tier: score >= 0.5, target ~3 MB
+    # Cap per-language to ensure diversity (no single language > 40% of docs)
+    LANG_CAPS = {"c_cpp": 200, "js_ts": 150, "python": 150, "rust": 100, "go": 50, "java": 50}
+    lang_counts_so_far = {}
+    top_tier = []
+    total_size = 0
+    target_size = 3_000_000
+
+    for score, idx, doc, lang in scored:
+        if score < 0.5:
+            break
+        cap = LANG_CAPS.get(lang, 100)
+        if lang_counts_so_far.get(lang, 0) >= cap:
+            continue
+        if total_size + len(doc) > target_size:
+            break
+        top_tier.append((score, idx, doc, lang))
+        total_size += len(doc)
+        lang_counts_so_far[lang] = lang_counts_so_far.get(lang, 0) + 1
+
+    print(f"\nCurated subset:")
+    print(f"  Documents: {len(top_tier):,}")
+    print(f"  Total size: {total_size / 1e6:.2f} MB")
+    print(f"  Score range: {top_tier[-1][0]:.3f} – {top_tier[0][0]:.3f}")
+
+    # Language distribution
+    lang_counts = {}
+    for _, _, _, lang in top_tier:
+        lang_counts[lang] = lang_counts.get(lang, 0) + 1
+    print(f"  Language distribution: {lang_counts}")
+
+    # Write curated corpus
+    with open(CURATED_PATH, "w", encoding="utf-8") as f:
+        for _, _, doc, _ in top_tier:
+            f.write(doc.strip())
+            f.write("\n\n\n")
+
+    print(f"\nCurated corpus written to {CURATED_PATH}")
+    print(f"  Size: {total_size / 1e6:.2f} MB ({len(top_tier):,} documents)")
+
+    # Show samples
+    print("\n" + "=" * 60)
+    print("SAMPLE DOCUMENTS (top 5 by score)")
+    print("=" * 60)
+    for i in range(min(5, len(top_tier))):
+        score, idx, doc, lang = top_tier[i]
+        print(f"\n--- Score: {score:.3f} | Lang: {lang} | Doc #{idx} | {len(doc)} chars ---")
+        print(doc[:600])
+        if len(doc) > 600:
+            print("...")
+
+
+if __name__ == "__main__":
+    main()

src/download_data.py

@@ -0,0 +1,148 @@
+"""
+Download a multi-language code corpus from GitHub for base pretraining.
+
+Downloads tarballs of several popular repositories across Python, JS/TS,
+Rust, Go, C/C++, and Java. Extracts source files and concatenates them
+into a single training corpus.
+"""
+
+import io
+import os
+import tarfile
+import time
+import requests
+
+DATA_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data")
+RAW_DIR = os.path.join(DATA_DIR, "raw")
+CORPUS_PATH = os.path.join(DATA_DIR, "corpus.txt")
+
+# Repos to download (user/repo, branch, max_bytes_per_repo)
+REPOS = [
+    ("python/cpython", "main", 8_000_000),       # Python, C
+    ("microsoft/TypeScript", "main", 8_000_000),  # TypeScript
+    ("rust-lang/rust", "master", 8_000_000),      # Rust
+    ("golang/go", "master", 8_000_000),           # Go
+    ("nginx/nginx", "master", 5_000_000),         # C
+    ("json-iterator/go", "master", 3_000_000),    # Go
+    ("pytorch/pytorch", "main", 8_000_000),       # C++, Python, CUDA
+    ("tokio-rs/tokio", "master", 5_000_000),      # Rust
+    ("numpy/numpy", "main", 5_000_000),           # Python, C
+    ("axios/axios", "v1.x", 3_000_000),           # JavaScript
+    ("lodash/lodash", "master", 3_000_000),       # JavaScript
+    ("redis/redis", "unstable", 5_000_000),       # C
+]
+
+# File extensions to include
+CODE_EXTENSIONS = {
+    ".py", ".js", ".ts", ".mjs", ".jsx", ".tsx",
+    ".rs", ".go", ".c", ".h", ".cpp", ".cc", ".cxx",
+    ".hpp", ".hxx", ".cu", ".cuh", ".java",
+    ".sh", ".bash", ".yml", ".yaml", ".toml",
+    ".cfg", ".ini", ".json", ".xml", ".sql",
+    ".md", ".txt", ".rst",
+}
+
+# Directories to skip
+SKIP_DIRS = {
+    "test", "tests", "testing", "__pycache__", ".git",
+    "node_modules", "vendor", "third_party", "thirdparty",
+    "dist", "build", "target", ".github", "docs",
+    "benchmark", "benchmarks", "examples", "example",
+    "fixtures", "testdata", "test_data", "mocks",
+}
+
+
+def download_repo_tarball(user_repo: str, branch: str) -> bytes | None:
+    """Download a repository tarball from GitHub."""
+    url = f"https://codeload.github.com/{user_repo}/tar.gz/refs/heads/{branch}"
+    print(f"  Downloading {user_repo}@{branch}...")
+    try:
+        r = requests.get(url, timeout=120, stream=True)
+        if r.status_code != 200:
+            # Try tags endpoint
+            url2 = f"https://codeload.github.com/{user_repo}/tar.gz/refs/tags/{branch}"
+            r = requests.get(url2, timeout=120, stream=True)
+            if r.status_code != 200:
+                print(f"    FAILED (status {r.status_code})")
+                return None
+        content = r.content
+        print(f"    Downloaded {len(content) / 1e6:.1f} MB")
+        return content
+    except Exception as e:
+        print(f"    ERROR: {e}")
+        return None
+
+
+def extract_code_files(tarball_bytes: bytes, max_bytes: int) -> list[str]:
+    """Extract code files from a tarball, up to max_bytes total."""
+    files = []
+    total = 0
+    try:
+        with tarfile.open(fileobj=io.BytesIO(tarball_bytes), mode="r:gz") as tar:
+            for member in tar.getmembers():
+                if not member.isfile():
+                    continue
+                path = member.name
+                parts = path.split("/")
+                # Skip unwanted directories
+                if any(skip in parts for skip in SKIP_DIRS):
+                    continue
+                # Check extension
+                _, ext = os.path.splitext(path)
+                if ext.lower() not in CODE_EXTENSIONS:
+                    continue
+                # Skip very large files
+                if member.size > 500_000:
+                    continue
+                try:
+                    f = tar.extractfile(member)
+                    if f is None:
+                        continue
+                    content = f.read()
+                    try:
+                        text = content.decode("utf-8", errors="ignore")
+                    except Exception:
+                        continue
+                    files.append(text)
+                    total += len(text)
+                    if total >= max_bytes:
+                        break
+                except Exception:
+                    continue
+    except Exception as e:
+        print(f"    Extraction error: {e}")
+    print(f"    Extracted {len(files)} files, {total / 1e6:.1f} MB of code")
+    return files
+
+
+def main():
+    os.makedirs(RAW_DIR, exist_ok=True)
+    os.makedirs(DATA_DIR, exist_ok=True)
+
+    all_code = []
+
+    for user_repo, branch, max_bytes in REPOS:
+        tarball = download_repo_tarball(user_repo, branch)
+        if tarball is None:
+            continue
+        files = extract_code_files(tarball, max_bytes)
+        all_code.extend(files)
+        time.sleep(1)  # Be polite
+
+    # Write corpus
+    print(f"\nTotal files: {len(all_code)}")
+    total_size = sum(len(f) for f in all_code)
+    print(f"Total corpus size: {total_size / 1e6:.1f} MB")
+
+    with open(CORPUS_PATH, "w", encoding="utf-8") as f:
+        for i, code in enumerate(all_code):
+            # Add document separator
+            f.write(code)
+            f.write("\n<|endoftext|>\n")
+
+    print(f"Corpus written to {CORPUS_PATH}")
+    return total_size
+
+
+if __name__ == "__main__":
+    main()

src/generate.py

@@ -0,0 +1,114 @@
+"""
+Inference and generation script for Retriever500M.
+
+Loads a checkpoint and generates text to verify the model has learned
+code structure during base pretraining.
+
+Usage:
+  python src/generate.py [--checkpoint PATH] [--prompt "text"] [--tokens N]
+"""
+
+import argparse
+import os
+import sys
+
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from model import ModelConfig, Retriever500M
+from tokenizers import Tokenizer
+
+PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+CHECKPOINT_DIR = os.path.join(PROJECT_DIR, "checkpoints")
+TOKENIZER_PATH = os.path.join(PROJECT_DIR, "tokenizer", "tokenizer.json")
+
+
+def load_model(checkpoint_path: str, device: torch.device) -> tuple[Retriever500M, ModelConfig]:
+    """Load model from checkpoint."""
+    print(f"Loading checkpoint: {checkpoint_path}")
+    ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
+
+    config = ModelConfig(**ckpt["config"])
+    model = Retriever500M(config).to(device)
+    model.load_state_dict(ckpt["model_state_dict"])
+    model.eval()
+
+    print(f"  Step: {ckpt.get('step', '?')}")
+    print(f"  Loss: {ckpt.get('loss', '?')}")
+    print(f"  Params: {model.count_parameters() / 1e6:.1f}M")
+
+    return model, config
+
+
+def generate(
+    model: Retriever500M,
+    tokenizer: Tokenizer,
+    prompt: str,
+    max_new_tokens: int = 128,
+    temperature: float = 0.8,
+    top_k: int = 50,
+    device: torch.device = None,
+) -> str:
+    """Generate text from a prompt."""
+    if device is None:
+        device = next(model.parameters()).device
+
+    # Encode prompt
+    encoded = tokenizer.encode(prompt)
+    input_ids = torch.tensor([encoded.ids], dtype=torch.long, device=device)
+
+    # Generate
+    with torch.no_grad():
+        output_ids = model.generate(
+            input_ids,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_k=top_k,
+        )
+
+    # Decode
+    output_text = tokenizer.decode(output_ids[0].tolist())
+    return output_text
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate text with Retriever500M")
+    parser.add_argument("--checkpoint", type=str, default=os.path.join(CHECKPOINT_DIR, "latest.pt"))
+    parser.add_argument("--prompt", type=str, default="def fibonacci(n):\n    ", help="Generation prompt")
+    parser.add_argument("--tokens", type=int, default=128, help="Max new tokens")
+    parser.add_argument("--temperature", type=float, default=0.8)
+    parser.add_argument("--top_k", type=int, default=50)
+    args = parser.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Device: {device}")
+
+    # Load tokenizer
+    tokenizer = Tokenizer.from_file(TOKENIZER_PATH)
+
+    # Load model
+    model, config = load_model(args.checkpoint, device)
+
+    # Generate
+    prompts = [
+        args.prompt,
+        "def quicksort(arr):\n    ",
+        "function fetchData(url) {\n  ",
+        "import torch\nimport torch.nn as nn\n\nclass Model(nn.Module):\n    ",
+        "fn main() {\n    println!",
+    ]
+
+    print("\n" + "=" * 60)
+    print("GENERATION SAMPLES")
+    print("=" * 60)
+
+    for prompt in prompts:
+        print(f"\n--- Prompt: {prompt!r} ---")
+        text = generate(model, tokenizer, prompt, args.tokens, args.temperature, args.top_k, device)
+        print(text)
+        print("-" * 60)
+
+
+if __name__ == "__main__":
+    main()

src/generate_traces.py

@@ -0,0 +1,605 @@
+"""
+Synthetic search agent trace generator.
+
+Generates thousands of diverse, realistic search agent traces using
+actual code chunks from the corpus. Each trace follows the gold-standard
+format: reasoning → search → results → analysis → evidence → finish.
+
+Trace types generated:
+  1. Simple lookup (1 search, 1 result)
+  2. Lookup with noise (1 search, 2-3 results, some irrelevant)
+  3. Multi-step search (2 searches, progressive refinement)
+  4. Query decomposition (complex query → subqueries)
+  5. Not found (search returns nothing)
+  6. Type/struct inspection (field enumeration)
+  7. Function behavior analysis (what does it do?)
+  8. Usage pattern (how is X used?)
+
+Query templates are varied to prevent memorization:
+  - "What does {name} do?"
+  - "Where is {name} defined?"
+  - "How does {name} work?"
+  - "What parameters does {name} accept?"
+  - "What fields does {name} have?"
+  - "How is {name} used in the codebase?"
+  - etc.
+
+Output: data/sft_traces.jsonl
+"""
+
+import json
+import os
+import random
+import re
+
+PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+CHUNKS_PATH = os.path.join(PROJECT_DIR, "data", "chunks.jsonl")
+TRACES_PATH = os.path.join(PROJECT_DIR, "data", "sft_traces.jsonl")
+
+SYSTEM_PROMPT = (
+    "You are a code search agent. Given a query from a reasoning model, "
+    "decompose it into subqueries, search the codebase, inspect results, "
+    "and return curated evidence. Use <|search|> to issue searches, "
+    "<|reasoning|> to analyze, and <|evidence|> to return findings. "
+    "Be concise. Extract only the relevant facts. End with <|finish|>."
+)
+
+random.seed(42)
+
+# ─── Query templates ─────────────────────────────────────────────────────────
+QUERY_TEMPLATES = {
+    "function": [
+        "What does {name} do?",
+        "Where is {name} defined?",
+        "How does {name} work?",
+        "What parameters does {name} accept?",
+        "What does {name} return?",
+        "Find the implementation of {name}.",
+        "Explain what {name} does step by step.",
+        "What is the signature of {name}?",
+    ],
+    "struct": [
+        "What fields does {name} have?",
+        "What is the structure of {name}?",
+        "Where is {name} defined?",
+        "What data does {name} contain?",
+        "Describe the {name} type.",
+    ],
+    "class": [
+        "What does the {name} class do?",
+        "What methods does {name} have?",
+        "Where is the {name} class defined?",
+        "How is {name} structured?",
+        "What is the inheritance of {name}?",
+    ],
+    "macro": [
+        "What does {name} do?",
+        "Where is {name} defined?",
+        "What is {name}?",
+    ],
+    "enum": [
+        "What variants does {name} have?",
+        "Where is {name} defined?",
+        "What values can {name} take?",
+    ],
+    "block": [
+        "What does {name} do?",
+        "Where is {name} defined?",
+        "Find code related to {name}.",
+        "How is {name} used?",
+    ],
+}
+
+# ─── Reasoning templates ─────────────────────────────────────────────────────
+INITIAL_REASONING = [
+    "I need to find {desc}. Let me search for {search_term}.",
+    "Looking for {desc}. I'll search using the term '{search_term}'.",
+    "The query asks about {desc}. Let me search the codebase.",
+    "I should find {desc}. Searching for '{search_term}'.",
+    "This requires finding {desc}. Let me issue a search.",
+]
+
+ANALYSIS_REASONING_FOUND = [
+    "Found the relevant code. {analysis}",
+    "This result contains what I need. {analysis}",
+    "I found {desc}. {analysis}",
+    "This is the right code. {analysis}",
+]
+
+ANALYSIS_REASONING_NOISE = [
+    "The first result is relevant. {analysis} The other results are not directly related.",
+    "Result 1 matches the query. {analysis} The remaining results appear to be unrelated code.",
+    "I found the target in the first result. {analysis} The other results don't match.",
+    "The relevant code is in the first result. {analysis} Ignoring the noise results.",
+]
+
+ANALYSIS_REASONING_REFINE = [
+    "Found the definition, but I need to see how it's used. Let me search for usages.",
+    "I have the implementation. Now let me find where it's called.",
+    "Got the definition. Let me also search for related patterns.",
+    "Found it. Let me do one more search to get complete context.",
+]
+
+NOT_FOUND_REASONING = [
+    "No results found for '{search_term}'. Let me try a broader search.",
+    "The search returned nothing. Let me try different terms.",
+    "No matches. The codebase may not contain this. Let me verify with another search.",
+]
+
+# ─── Code analysis functions ─���───────────────────────────────────────────────
+
+def extract_function_info(code: str, name: str) -> dict:
+    """Extract key information from a function."""
+    info = {"name": name, "params": [], "returns": "", "key_lines": []}
+
+    # Extract parameters from signature
+    sig_match = re.search(r"(?:def|fn|func|function)\s+" + re.escape(name) + r"\s*\(([^)]*)\)", code)
+    if sig_match:
+        params_raw = sig_match.group(1).strip()
+        if params_raw:
+            info["params"] = [p.strip() for p in params_raw.split(",") if p.strip()]
+
+    # Extract return type (Rust/TS)
+    ret_match = re.search(r"->\s*([^{]+)", code)
+    if ret_match:
+        info["returns"] = ret_match.group(1).strip()
+
+    # Extract key lines (non-trivial lines)
+    lines = code.split("\n")
+    key_lines = []
+    for line in lines:
+        stripped = line.strip()
+        if not stripped or stripped.startswith("//") or stripped.startswith("#"):
+            continue
+        if any(kw in stripped for kw in ["return ", "if ", "for ", "while ", "raise ", "throw ", "Err(", "Ok(", "assert"]):
+            key_lines.append(stripped)
+    info["key_lines"] = key_lines[:5]
+
+    return info
+
+
+def extract_struct_info(code: str, name: str) -> dict:
+    """Extract field information from a struct."""
+    info = {"name": name, "fields": []}
+
+    # Find field declarations (type name; or name: type)
+    field_patterns = [
+        re.compile(r"^\s+(\w+)\s+(\w+);", re.M),           # C: type name;
+        re.compile(r"^\s+(\w+):\s+([^,;]+)", re.M),         # Rust/TS: name: type
+        re.compile(r"^\s+(self\.)?(\w+)\s*[:=]", re.M),      # Python: self.name =
+    ]
+
+    for pat in field_patterns:
+        matches = pat.findall(code)
+        for m in matches:
+            if isinstance(m, tuple):
+                field_name = m[-1] if m[-1] else m[0]
+            else:
+                field_name = m
+            if field_name and field_name not in ("self", "pub", "fn", "def", "class"):
+                info["fields"].append(field_name)
+
+    return info
+
+
+def generate_evidence(chunk: dict, analysis_type: str = "found") -> str:
+    """Generate curated evidence from a code chunk."""
+    code = chunk["code"]
+    name = chunk["name"]
+    typ = chunk["type"]
+    lang = chunk["language"]
+    filepath = chunk["filepath"]
+
+    if typ in ("function",):
+        info = extract_function_info(code, name)
+        parts = [f"`{name}` is a {lang} function:"]
+
+        # Signature
+        if info["params"]:
+            params_str = ", ".join(info["params"][:6])
+            if len(info["params"]) > 6:
+                params_str += ", ..."
+            parts.append(f"- Parameters: {params_str}")
+
+        if info["returns"]:
+            parts.append(f"- Returns: {info['returns']}")
+
+        # Key behavior
+        if info["key_lines"]:
+            parts.append("- Key behavior:")
+            for line in info["key_lines"][:3]:
+                parts.append(f"  - `{line}`")
+
+        parts.append(f"Source: {filepath}")
+        return "\n".join(parts)
+
+    elif typ in ("struct",):
+        info = extract_struct_info(code, name)
+        parts = [f"`{name}` is a {lang} structure with fields:"]
+        for field in info["fields"][:10]:
+            parts.append(f"- `{field}`")
+        if len(info["fields"]) > 10:
+            parts.append(f"- ... ({len(info['fields'])} total fields)")
+        parts.append(f"Source: {filepath}")
+        return "\n".join(parts)
+
+    elif typ in ("class",):
+        parts = [f"`{name}` is a {lang} class:"]
+        # Find methods
+        methods = re.findall(r"(?:def |fn |func |function )\s*(\w+)", code)
+        if methods:
+            parts.append(f"- Methods: {', '.join(methods[:8])}")
+        # Find inheritance
+        base_match = re.search(r"class\s+\w+\s*(?:\(|:\s*|extends\s+|<\s*)(\w+)", code)
+        if base_match:
+            parts.append(f"- Inherits from: {base_match.group(1)}")
+        parts.append(f"Source: {filepath}")
+        return "\n".join(parts)
+
+    elif typ in ("macro",):
+        parts = [f"`{name}` is a {lang} macro/preprocessor definition:"]
+        # First few lines of the macro
+        lines = code.strip().split("\n")[:5]
+        for line in lines:
+            parts.append(f"  `{line}`")
+        parts.append(f"Source: {filepath}")
+        return "\n".join(parts)
+
+    elif typ in ("enum",):
+        parts = [f"`{name}` is a {lang} enum with variants:"]
+        variants = re.findall(r"^\s+(\w+)[,\s]*$", code, re.M)
+        for v in variants[:10]:
+            parts.append(f"- `{v}`")
+        parts.append(f"Source: {filepath}")
+        return "\n".join(parts)
+
+    else:
+        # Generic block
+        parts = [f"Found code related to `{name}` ({lang}):"]
+        lines = code.strip().split("\n")[:8]
+        for line in lines:
+            parts.append(f"  `{line}`")
+        parts.append(f"Source: {filepath}")
+        return "\n".join(parts)
+
+
+def generate_search_term(chunk: dict) -> str:
+    """Generate a realistic search term for a chunk."""
+    name = chunk["name"]
+    if name != "unknown" and name != "block":
+        return name
+
+    # For unnamed blocks, extract a key identifier
+    code = chunk["code"]
+    identifiers = re.findall(r"\b(ngx_\w+|def\s+\w+|fn\s+\w+|struct\s+\w+|class\s+\w+)\b", code)
+    if identifiers:
+        return identifiers[0].replace("def ", "").replace("fn ", "").replace("struct ", "").replace("class ", "")
+
+    # Fallback: first meaningful word
+    words = re.findall(r"\b[a-z_]{4,}\b", code)
+    return words[0] if words else "code"
+
+
+def generate_description(chunk: dict) -> str:
+    """Generate a natural language description of what we're looking for."""
+    name = chunk["name"]
+    typ = chunk["type"]
+    lang = chunk["language"]
+
+    if typ == "function":
+        return f"the {name} function"
+    elif typ == "struct":
+        return f"the {name} structure"
+    elif typ == "class":
+        return f"the {name} class"
+    elif typ == "macro":
+        return f"the {name} macro"
+    elif typ == "enum":
+        return f"the {name} enum"
+    else:
+        return f"code related to {name}"
+
+
+# ─── Trace generators ────────────────────────────────────────────────────────
+
+def make_trace_simple(chunk: dict) -> dict:
+    """Type 1: Simple lookup — 1 search, 1 result."""
+    name = chunk["name"]
+    typ = chunk["type"]
+    templates = QUERY_TEMPLATES.get(typ, QUERY_TEMPLATES["block"])
+    query = random.choice(templates).format(name=name)
+
+    search_term = generate_search_term(chunk)
+    desc = generate_description(chunk)
+
+    reasoning1 = random.choice(INITIAL_REASONING).format(desc=desc, search_term=search_term)
+    evidence = generate_evidence(chunk)
+
+    return {
+        "query": query,
+        "trace": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": query},
+            {"role": "assistant", "content": f"<|reasoning|>{reasoning1}<|end|>\n<|search|>{search_term}<|end|>"},
+            {"role": "result", "content": chunk["code"]},
+            {"role": "assistant", "content": f"<|reasoning|>Found the relevant code. This is {desc}.<|end|>\n<|evidence|>{evidence}<|end|>\n<|finish|>"},
+        ],
+    }
+
+
+def make_trace_with_noise(chunk: dict, noise_chunks: list) -> dict:
+    """Type 2: Lookup with noise — 1 search, 2-3 results (some irrelevant)."""
+    name = chunk["name"]
+    typ = chunk["type"]
+    templates = QUERY_TEMPLATES.get(typ, QUERY_TEMPLATES["block"])
+    query = random.choice(templates).format(name=name)
+
+    search_term = generate_search_term(chunk)
+    desc = generate_description(chunk)
+    evidence = generate_evidence(chunk)
+
+    reasoning1 = random.choice(INITIAL_REASONING).format(desc=desc, search_term=search_term)
+
+    # Select 1-2 noise chunks
+    n_noise = random.choice([1, 2])
+    noise = random.sample(noise_chunks, min(n_noise, len(noise_chunks)))
+
+    analysis = f"This contains {desc}."
+    analysis_template = random.choice(ANALYSIS_REASONING_NOISE).format(analysis=analysis)
+
+    trace = [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {"role": "user", "content": query},
+        {"role": "assistant", "content": f"<|reasoning|>{reasoning1}<|end|>\n<|search|>{search_term}<|end|>"},
+        {"role": "result", "content": chunk["code"]},
+    ]
+    for nc in noise:
+        trace.append({"role": "result", "content": nc["code"]})
+
+    trace.append({"role": "assistant", "content": f"<|reasoning|>{analysis_template}<|end|>\n<|evidence|>{evidence}<|end|>\n<|finish|>"})
+
+    return {"query": query, "trace": trace}
+
+
+def make_trace_multi_step(chunk: dict, related_chunk: dict) -> dict:
+    """Type 3: Multi-step search — 2 searches, progressive refinement."""
+    name = chunk["name"]
+    typ = chunk["type"]
+    query = f"What does {name} do and where is it used?"
+
+    search_term1 = generate_search_term(chunk)
+    search_term2 = generate_search_term(related_chunk)
+    desc = generate_description(chunk)
+    evidence1 = generate_evidence(chunk)
+    evidence2 = generate_evidence(related_chunk)
+
+    reasoning1 = random.choice(INITIAL_REASONING).format(desc=desc, search_term=search_term1)
+    reasoning2 = random.choice(ANALYSIS_REASONING_REFINE)
+
+    return {
+        "query": query,
+        "trace": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": query},
+            {"role": "assistant", "content": f"<|reasoning|>{reasoning1}<|end|>\n<|search|>{search_term1}<|end|>"},
+            {"role": "result", "content": chunk["code"]},
+            {"role": "assistant", "content": f"<|reasoning|>Found {desc}. {reasoning2}<|end|>\n<|search|>{search_term2}<|end|>"},
+            {"role": "result", "content": related_chunk["code"]},
+            {"role": "assistant", "content": f"<|reasoning|>Found both the definition and a usage. I now have complete context.<|end|>\n<|evidence|>{evidence1}\n\nRelated usage:\n{evidence2}<|end|>\n<|finish|>"},
+        ],
+    }
+
+
+def make_trace_not_found() -> dict:
+    """Type 5: Not found — search returns nothing."""
+    fake_names = ["database_pool", "auth_middleware", "cache_manager", "session_store",
+                  "websocket_handler", "graphql_resolver", "redis_client", "email_service"]
+    name = random.choice(fake_names)
+    query = f"Where is the {name} implemented?"
+
+    search1 = name
+    search2 = name.replace("_", " ")
+
+    reasoning1 = random.choice(INITIAL_REASONING).format(desc=f"the {name}", search_term=search1)
+    reasoning2 = random.choice(NOT_FOUND_REASONING).format(search_term=search1)
+
+    return {
+        "query": query,
+        "trace": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": query},
+            {"role": "assistant", "content": f"<|reasoning|>{reasoning1}<|end|>\n<|search|>{search1}<|end|>"},
+            {"role": "result", "content": ""},
+            {"role": "assistant", "content": f"<|reasoning|>{reasoning2}<|end|>\n<|search|>{search2}<|end|>"},
+            {"role": "result", "content": ""},
+            {"role": "assistant", "content": f"<|reasoning|>No results found. The codebase does not contain {name}.<|end|>\n<|evidence|>No implementation of `{name}` was found in the codebase. Searched for: '{search1}', '{search2}'.<|end|>\n<|finish|>"},
+        ],
+    }
+
+
+def make_trace_decomposition(chunk: dict, related_chunk: dict) -> dict:
+    """Type 4: Query decomposition — complex query broken into subqueries."""
+    name = chunk["name"]
+    related_name = related_chunk["name"]
+
+    query = f"How does {name} work and how does it relate to {related_name}?"
+
+    search1 = generate_search_term(chunk)
+    search2 = generate_search_term(related_chunk)
+    evidence1 = generate_evidence(chunk)
+    evidence2 = generate_evidence(related_chunk)
+
+    return {
+        "query": query,
+        "trace": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": query},
+            {"role": "assistant", "content": f"<|reasoning|>This query has two parts: (1) how {name} works, (2) its relationship to {related_name}. Let me search for {name} first.<|end|>\n<|search|>{search1}<|end|>"},
+            {"role": "result", "content": chunk["code"]},
+            {"role": "assistant", "content": f"<|reasoning|>Found {name}. Now let me search for {related_name} to understand the relationship.<|end|>\n<|search|>{search2}<|end|>"},
+            {"role": "result", "content": related_chunk["code"]},
+            {"role": "assistant", "content": f"<|reasoning|>Found both components. I can now explain how they relate.<|end|>\n<|evidence|>{evidence1}\n\nRelated: {evidence2}<|end|>\n<|finish|>"},
+        ],
+    }
+
+
+def make_trace_struct_inspection(chunk: dict) -> dict:
+    """Type 6: Struct inspection — enumerate fields."""
+    if chunk["type"] not in ("struct",):
+        return make_trace_simple(chunk)
+
+    name = chunk["name"]
+    query = f"What fields does {name} have?"
+    search_term = name
+    evidence = generate_evidence(chunk)
+
+    return {
+        "query": query,
+        "trace": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": query},
+            {"role": "assistant", "content": f"<|reasoning|>Looking for the {name} struct definition to enumerate its fields.<|end|>\n<|search|>{search_term}<|end|>"},
+            {"role": "result", "content": chunk["code"]},
+            {"role": "assistant", "content": f"<|reasoning|>Found the struct definition. I can extract all fields from it.<|end|>\n<|evidence|>{evidence}<|end|>\n<|finish|>"},
+        ],
+    }
+
+
+def make_trace_behavior_analysis(chunk: dict) -> dict:
+    """Type 7: Function behavior analysis — what does it do?"""
+    if chunk["type"] not in ("function",):
+        return make_trace_simple(chunk)
+
+    name = chunk["name"]
+    query = f"What does {name} do? Explain its behavior."
+    search_term = name
+    evidence = generate_evidence(chunk)
+
+    info = extract_function_info(chunk["code"], name)
+    params = ", ".join(info["params"][:4]) if info["params"] else "no parameters"
+
+    return {
+        "query": query,
+        "trace": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": query},
+            {"role": "assistant", "content": f"<|reasoning|>I need to find and analyze the {name} function. Let me search for its definition.<|end|>\n<|search|>{search_term}<|end|>"},
+            {"role": "result", "content": chunk["code"]},
+            {"role": "assistant", "content": f"<|reasoning|>Found the function. It takes {params}. Let me analyze its behavior from the code.<|end|>\n<|evidence|>{evidence}<|end|>\n<|finish|>"},
+        ],
+    }
+
+
+# ─── Main generation ─────────────────────────────────────────────────────────
+
+def main():
+    print(f"Loading chunks from {CHUNKS_PATH}...")
+    with open(CHUNKS_PATH, "r", encoding="utf-8") as f:
+        chunks = [json.loads(line) for line in f]
+    print(f"  Loaded {len(chunks):,} chunks")
+
+    # Filter to chunks with real names and reasonable size
+    good_chunks = [c for c in chunks if c["name"] != "unknown" and c["name"] != "block" and 50 < len(c["code"]) < 3000]
+    block_chunks = [c for c in chunks if c["name"] == "block" or c["name"] == "unknown"]
+    print(f"  Named chunks (good): {len(good_chunks):,}")
+    print(f"  Block chunks: {len(block_chunks):,}")
+
+    # Group by language for related chunk pairing
+    by_lang = {}
+    for c in good_chunks:
+        by_lang.setdefault(c["language"], []).append(c)
+
+    traces = []
+
+    # ─── Generate traces ─────────────────────────────────────────────────────
+    print("\nGenerating traces...")
+
+    # Type 1: Simple lookup (30% of traces)
+    n_simple = min(2000, len(good_chunks))
+    sampled = random.sample(good_chunks, n_simple)
+    for chunk in sampled:
+        traces.append(make_trace_simple(chunk))
+    print(f"  Simple lookup: {n_simple}")
+
+    # Type 2: With noise (20%)
+    n_noise = min(1500, len(good_chunks))
+    sampled = random.sample(good_chunks, n_noise)
+    for chunk in sampled:
+        # Pick noise chunks from same language
+        lang_chunks = by_lang.get(chunk["language"], good_chunks)
+        noise_pool = [c for c in lang_chunks if c["id"] != chunk["id"]]
+        if len(noise_pool) >= 2:
+            traces.append(make_trace_with_noise(chunk, noise_pool))
+    print(f"  With noise: {n_noise}")
+
+    # Type 3: Multi-step (15%)
+    n_multi = min(1000, len(good_chunks) // 2)
+    sampled = random.sample(good_chunks, n_multi)
+    for chunk in sampled:
+        lang_chunks = by_lang.get(chunk["language"], good_chunks)
+        related_pool = [c for c in lang_chunks if c["id"] != chunk["id"]]
+        if related_pool:
+            related = random.choice(related_pool)
+            traces.append(make_trace_multi_step(chunk, related))
+    print(f"  Multi-step: {n_multi}")
+
+    # Type 4: Query decomposition (10%)
+    n_decomp = min(700, len(good_chunks) // 3)
+    sampled = random.sample(good_chunks, n_decomp)
+    for chunk in sampled:
+        lang_chunks = by_lang.get(chunk["language"], good_chunks)
+        related_pool = [c for c in lang_chunks if c["id"] != chunk["id"]]
+        if related_pool:
+            related = random.choice(related_pool)
+            traces.append(make_trace_decomposition(chunk, related))
+    print(f"  Query decomposition: {n_decomp}")
+
+    # Type 5: Not found (5%)
+    n_notfound = 300
+    for _ in range(n_notfound):
+        traces.append(make_trace_not_found())
+    print(f"  Not found: {n_notfound}")
+
+    # Type 6: Struct inspection (10%)
+    struct_chunks = [c for c in good_chunks if c["type"] == "struct"]
+    for chunk in struct_chunks[:500]:
+        traces.append(make_trace_struct_inspection(chunk))
+    print(f"  Struct inspection: {min(len(struct_chunks), 500)}")
+
+    # Type 7: Behavior analysis (10%)
+    func_chunks = [c for c in good_chunks if c["type"] == "function"]
+    for chunk in func_chunks[:500]:
+        traces.append(make_trace_behavior_analysis(chunk))
+    print(f"  Behavior analysis: {min(len(func_chunks), 500)}")
+
+    # Shuffle
+    random.shuffle(traces)
+
+    # Write
+    print(f"\nTotal traces: {len(traces):,}")
+    with open(TRACES_PATH, "w", encoding="utf-8") as f:
+        for trace in traces:
+            f.write(json.dumps(trace) + "\n")
+    print(f"Traces written to {TRACES_PATH}")
+
+    # Stats
+    total_size = os.path.getsize(TRACES_PATH)
+    print(f"  File size: {total_size / 1e6:.2f} MB")
+
+    # Show a sample
+    print("\n" + "=" * 60)
+    print("SAMPLE TRACE")
+    print("=" * 60)
+    sample = traces[0]
+    print(f"\nQuery: {sample['query']}")
+    for msg in sample["trace"]:
+        role = msg["role"]
+        content = msg["content"]
+        if len(content) > 200:
+            content = content[:200] + "..."
+        print(f"\n[{role}]")
+        print(content)
+
+
+if __name__ == "__main__":
+    main()

src/gold_traces.py

@@ -0,0 +1,324 @@
+"""
+Gold-standard search agent traces — hand-crafted examples using real code
+from the corpus. These define the quality bar for the synthetic trace generator.
+
+Each trace follows this format:
+  1. System prompt (agent role definition)
+  2. User query (from the bigger model)
+  3. Agent response:
+     a. <|reasoning|> — decompose the query, plan searches
+     b. <|search|> — issue a search query
+     c. <|result|> — retrieved code chunks (injected by the harness)
+     d. <|reasoning|> — analyze results, decide if more search needed
+     e. (repeat b-d as needed)
+     f. <|evidence|> — curated evidence package
+     g. <|finish|> — signal completion
+
+The agent learns to:
+  - Break complex queries into subqueries
+  - Search with specific terms (function names, types, patterns)
+  - Judge which results are relevant vs noise
+  - Extract key facts concisely (not dump entire code)
+  - Synthesize across multiple results
+  - Know when to stop searching
+"""
+
+import json
+import os
+
+PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+GOLD_PATH = os.path.join(PROJECT_DIR, "data", "gold_traces.jsonl")
+
+SYSTEM_PROMPT = (
+    "You are a code search agent. Given a query from a reasoning model, "
+    "decompose it into subqueries, search the codebase, inspect results, "
+    "and return curated evidence. Use <|search|> to issue searches, "
+    "<|reasoning|> to analyze, and <|evidence|> to return findings. "
+    "Be concise. Extract only the relevant facts. End with <|finish|>."
+)
+
+GOLD_TRACES = [
+    # ─── Trace 1: Simple function lookup (C) ───────────────────────────────
+    {
+        "query": "How does nginx handle reusable connections?",
+        "trace": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": "How does nginx handle reusable connections?"},
+            {"role": "assistant", "content": (
+                "<|reasoning|>The query asks about reusable connections in nginx. "
+                "I should search for functions related to reusable connections.<|end|>\n"
+                "<|search|>reusable connection<|end|>"
+            )},
+            {"role": "result", "content": (
+                "ngx_reusable_connection(ngx_connection_t *c, ngx_uint_t reusable)\n"
+                "{\n"
+                "    ngx_log_debug1(NGX_LOG_DEBUG_CORE, c->log, 0,\n"
+                "                   \"reusable connection: %ui\", reusable);\n"
+                "\n"
+                "    if (c->reusable) {\n"
+                "        ngx_queue_remove(&c->queue);\n"
+                "        ngx_cycle->reusable_connections_n--;\n"
+                "    }\n"
+                "\n"
+                "    c->reusable = reusable;\n"
+                "\n"
+                "    if (reusable) {\n"
+                "        ngx_queue_insert_head(\n"
+                "            (ngx_queue_t *) &ngx_cycle->reusable_connections_queue, &c->queue);\n"
+                "        ngx_cycle->reusable_connections_n++;\n"
+                "    }\n"
+                "}"
+            )},
+            {"role": "assistant", "content": (
+                "<|reasoning|>This function manages the reusable connection queue. "
+                "When a connection is marked reusable, it's added to "
+                "reusable_connections_queue and the counter is incremented. "
+                "When unreusable, it's removed from the queue and counter decremented. "
+                "I have enough to answer.<|end|>\n"
+                "<|evidence|>nginx manages reusable connections via `ngx_reusable_connection()`:\n"
+                "- Marks connection as reusable/unreusable via the `reusable` flag\n"
+                "- When reusable: inserts into `reusable_connections_queue` (head), increments `reusable_connections_n`\n"
+                "- When unreusable: removes from queue, decrements counter\n"
+                "- Uses `ngx_queue_insert_head` and `ngx_queue_remove` for queue management\n"
+                "Source: ngx_connection.c<|end|>\n"
+                "<|finish|>"
+            )},
+        ],
+    },
+    # ─── Trace 2: Multi-step search (Python) ──────────────────────────────
+    {
+        "query": "What does the with_params_help decorator do and how is it used?",
+        "trace": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": "What does the with_params_help decorator do and how is it used?"},
+            {"role": "assistant", "content": (
+                "<|reasoning|>This asks about a specific decorator. I need to find its definition first, then find usages.<|end|>\n"
+                "<|search|>with_params_help<|end|>"
+            )},
+            {"role": "result", "content": (
+                "def with_params_help(params_cls: type, title: str = \"Parameter defaults\"):\n"
+                "    \"\"\"\n"
+                "    Class decorator that appends a help table generated from another dataclass\n"
+                "    (e.g., VllmParameters) to the decorated class's docstring.\n"
+                "    \"\"\"\n"
+                "    if not is_dataclass(params_cls):\n"
+                "        raise TypeError(f\"{params_cls} must be a dataclass\")\n"
+                "\n"
+                "    def _decorator(cls: type) -> type:\n"
+                "        block = generate_dataclass_help(params_cls)\n"
+                "        cls.__doc__ = (cls.__doc__ or \"\") + f\"\\n\\n{title}:\\n{block}\"\n"
+                "        return cls\n"
+                "\n"
+                "    return _decorator"
+            )},
+            {"role": "assistant", "content": (
+                "<|reasoning|>Found the definition. It's a decorator factory that takes a dataclass and appends a help table to another class's docstring. Now I need to find where it's used to show a concrete example.<|end|>\n"
+                "<|search|>@with_params_help<|end|>"
+            )},
+            {"role": "result", "content": (
+                "@with_params_help(VllmBuildParameters)\n"
+                "class VllmBuildRunner(BaseRunner):\n"
+                "    \"\"\"\n"
+                "    Build vLLM using docker buildx.\n"
+                "    \"\"\""
+            )},
+            {"role": "assistant", "content": (
+                "<|reasoning|>Found a usage. The decorator is applied to VllmBuildRunner with VllmBuildParameters as the params class. This appends parameter defaults from VllmBuildParameters to VllmBuildRunner's docstring. I have enough to answer.<|end|>\n"
+                "<|evidence|>`with_params_help` is a class decorator factory:\n"
+                "- Takes a dataclass (`params_cls`) and optional title (default \"Parameter defaults\")\n"
+                "- Generates a help table from the dataclass fields via `generate_dataclass_help()`\n"
+                "- Appends the table to the decorated class's `__doc__` docstring\n"
+                "- Usage example: `@with_params_help(VllmBuildParameters)` on `VllmBuildRunner`\n"
+                "- Requires `params_cls` to be a dataclass (raises TypeError otherwise)<|end|>\n"
+                "<|finish|>"
+            )},
+        ],
+    },
+    # ─── Trace 3: Type/struct lookup with noise (C) ───────────────────────
+    {
+        "query": "What fields does the nginx listening socket structure have?",
+        "trace": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": "What fields does the nginx listening socket structure have?"},
+            {"role": "assistant", "content": (
+                "<|reasoning|>Looking for the listening socket struct definition in nginx. Should search for ngx_listening struct.<|end|>\n"
+                "<|search|>ngx_listening struct fields<|end|>"
+            )},
+            {"role": "result", "content": (
+                "typedef struct ngx_listening_s  ngx_listening_t;\n"
+                "\n"
+                "struct ngx_listening_s {\n"
+                "    ngx_socket_t        fd;\n"
+                "    struct sockaddr    *sockaddr;\n"
+                "    socklen_t           socklen;\n"
+                "    size_t              addr_text_max_len;\n"
+                "    ngx_str_t           addr_text;\n"
+                "    int                 type;\n"
+                "    int                 protocol;\n"
+                "    int                 backlog;\n"
+                "    int                 rcvbuf;\n"
+                "    int                 sndbuf;\n"
+                "};"
+            )},
+            {"role": "result", "content": (
+                "ngx_create_listening(ngx_conf_t *cf, struct sockaddr *sockaddr,\n"
+                "    socklen_t socklen)\n"
+                "{\n"
+                "    ls = ngx_array_push(&cf->cycle->listening);\n"
+                "    ngx_memzero(ls, sizeof(ngx_listening_t));\n"
+                "    ls->sockaddr = sa;\n"
+                "    ls->socklen = socklen;\n"
+                "}"
+            )},
+            {"role": "assistant", "content": (
+                "<|reasoning|>The first result is the struct definition — exactly what I need. The second result is the creation function, which is related but not asked for. I'll focus on the struct fields.<|end|>\n"
+                "<|evidence|>`ngx_listening_s` (aka `ngx_listening_t`) structure fields:\n"
+                "- `fd` (ngx_socket_t): file descriptor\n"
+                "- `sockaddr` (struct sockaddr*): socket address pointer\n"
+                "- `socklen` (socklen_t): socket address length\n"
+                "- `addr_text_max_len` (size_t): max address text length\n"
+                "- `addr_text` (ngx_str_t): address text representation\n"
+                "- `type` (int): socket type\n"
+                "- `protocol` (int): protocol\n"
+                "- `backlog` (int): listen backlog\n"
+                "- `rcvbuf` (int): receive buffer size\n"
+                "- `sndbuf` (int): send buffer size\n"
+                "Source: ngx_connection.h<|end|>\n"
+                "<|finish|>"
+            )},
+        ],
+    },
+    # ─── Trace 4: Query decomposition (JS/TS) ─────────────────────────────
+    {
+        "query": "How does the concatenate function work and what types does it accept?",
+        "trace": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": "How does the concatenate function work and what types does it accept?"},
+            {"role": "assistant", "content": (
+                "<|reasoning|>This has two parts: (1) how concatenate works, (2) what types it accepts. Let me search for the function definition.<|end|>\n"
+                "<|search|>concatenate function array<|end|>"
+            )},
+            {"role": "result", "content": (
+                "export function concatenate<T>(array1: T[], array2: T[] | undefined): T[];\n"
+                "// eslint-disable-line @typescript-eslint/unified-signatures"
+            )},
+            {"role": "result", "content": (
+                "function concatenate<T>(array1: T[], array2: T[] | undefined): T[] {\n"
+                "    if (array2 === undefined) return array1;\n"
+                "    return array1.concat(array2);\n"
+                "}"
+            )},
+            {"role": "assistant", "content": (
+                "<|reasoning|>Found both the declaration and implementation. It's a generic TypeScript function that concatenates two arrays. The second parameter is optional (can be undefined). If undefined, it returns the first array unchanged. Otherwise uses Array.concat. I have everything needed.<|end|>\n"
+                "<|evidence|>`concatenate<T>` is a TypeScript utility function:\n"
+                "- Signature: `concatenate<T>(array1: T[], array2: T[] | undefined): T[]`\n"
+                "- Generic over type T — accepts arrays of any type\n"
+                "- If `array2` is `undefined`, returns `array1` unchanged (no-op)\n"
+                "- Otherwise returns `array1.concat(array2)` (standard Array.concat)\n"
+                "- Exported as a public function<|end|>\n"
+                "<|finish|>"
+            )},
+        ],
+    },
+    # ─── Trace 5: Search with irrelevant results (Rust) ───────────────────
+    {
+        "query": "How is error handling done in the Rust codebase?",
+        "trace": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": "How is error handling done in the Rust codebase?"},
+            {"role": "assistant", "content": (
+                "<|reasoning|>Error handling in Rust typically uses Result types. Let me search for Result usage and error types.<|end|>\n"
+                "<|search|>Result error handling<|end|>"
+            )},
+            {"role": "result", "content": (
+                "pub fn parse_config(input: &str) -> Result<Config, ParseError> {\n"
+                "    let parts: Vec<&str> = input.split_whitespace().collect();\n"
+                "    if parts.len() < 2 {\n"
+                "        return Err(ParseError::TooFewArguments);\n"
+                "    }\n"
+                "    Ok(Config { query: parts[0].to_string(), filename: parts[1].to_string() })\n"
+                "}"
+            )},
+            {"role": "result", "content": (
+                "use std::fmt;\n"
+                "\n"
+                "pub enum ParseError {\n"
+                "    TooFewArguments,\n"
+                "    InvalidSyntax,\n"
+                "}\n"
+                "\n"
+                "impl fmt::Display for ParseError {\n"
+                "    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {\n"
+                "        match self {\n"
+                "            ParseError::TooFewArguments => write!(f, \"too few arguments\"),\n"
+                "            ParseError::InvalidSyntax => write!(f, \"invalid syntax\"),\n"
+                "        }\n"
+                "    }\n"
+                "}"
+            )},
+            {"role": "result", "content": (
+                "mod tests {\n"
+                "    #[test]\n"
+                "    fn test_basic() {\n"
+                "        assert_eq!(2 + 2, 4);\n"
+                "    }\n"
+                "}"
+            )},
+            {"role": "assistant", "content": (
+                "<|reasoning|>Result 1 shows a function returning Result<Config, ParseError> with early return on error. Result 2 shows the ParseError enum with Display impl. Result 3 is an unrelated test — not relevant. I have enough to answer.<|end|>\n"
+                "<|evidence|>Error handling pattern in this Rust codebase:\n"
+                "- Uses `Result<T, E>` for fallible operations\n"
+                "- Custom error type: `ParseError` enum with variants (TooFewArguments, InvalidSyntax)\n"
+                "- Implements `std::fmt::Display` for human-readable errors\n"
+                "- Early return pattern: `return Err(ParseError::...)` on validation failure\n"
+                "- Success path: `Ok(Config { ... })`\n"
+                "- No use of `unwrap()` or `expect()` in the found code — proper error propagation<|end|>\n"
+                "<|finish|>"
+            )},
+        ],
+    },
+    # ─── Trace 6: No results found ────────────────────────────────────────
+    {
+        "query": "Where is the database connection pool implemented?",
+        "trace": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": "Where is the database connection pool implemented?"},
+            {"role": "assistant", "content": (
+                "<|reasoning|>Looking for database connection pool implementation. Let me search for relevant terms.<|end|>\n"
+                "<|search|>database connection pool<|end|>"
+            )},
+            {"role": "result", "content": ""},
+            {"role": "assistant", "content": (
+                "<|reasoning|>No results found for 'database connection pool'. Let me try broader terms.<|end|>\n"
+                "<|search|>pool db<|end|>"
+            )},
+            {"role": "result", "content": ""},
+            {"role": "assistant", "content": (
+                "<|reasoning|>Still no results. The codebase does not appear to contain a database connection pool implementation.<|end|>\n"
+                "<|evidence|>No database connection pool implementation was found in the codebase. Searched for: 'database connection pool', 'pool db'. The indexed code may not include database layer code.<|end|>\n"
+                "<|finish|>"
+            )},
+        ],
+    },
+]
+
+
+def main():
+    print(f"Writing {len(GOLD_TRACES)} gold-standard traces to {GOLD_PATH}...")
+    with open(GOLD_PATH, "w", encoding="utf-8") as f:
+        for trace in GOLD_TRACES:
+            f.write(json.dumps(trace) + "\n")
+
+    print(f"Done. {len(GOLD_TRACES)} traces written.")
+
+    # Print summary
+    for i, trace in enumerate(GOLD_TRACES):
+        query = trace["query"]
+        n_steps = len([m for m in trace["trace"] if m["role"] == "assistant"])
+        n_searches = len([m for m in trace["trace"] if m["role"] == "assistant" and "<|search|>" in m["content"]])
+        n_results = len([m for m in trace["trace"] if m["role"] == "result"])
+        print(f"  Trace {i+1}: '{query[:50]}...' | {n_steps} agent turns | {n_searches} searches | {n_results} results")
+
+
+if __name__ == "__main__":
+    main()

src/model.py

@@ -0,0 +1,360 @@
+"""
+Retriever500M - Decoder-only transformer built from scratch.
+
+Architecture (LLaMA-style):
+  - vocab_size:    32,000
+  - d_model:       1,280
+  - n_layers:      23
+  - n_heads:       20
+  - d_ff:          3,456  (SwiGLU, 2/3 * 4 * d_model)
+  - RoPE positional encoding
+  - RMSNorm (no biases)
+  - Tied input/output embeddings
+  - Total parameters: ~497M
+"""
+
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+@dataclass
+class ModelConfig:
+    vocab_size: int = 32_000
+    d_model: int = 1_280
+    n_layers: int = 23
+    n_heads: int = 20
+    d_ff: int = 3_456
+    max_seq_len: int = 1_024
+    rope_theta: float = 10_000.0
+    rope_pct: float = 0.25  # fraction of d_model per head used for RoPE
+    dropout: float = 0.0
+    tie_embeddings: bool = True
+
+    def __post_init__(self):
+        assert self.d_model % self.n_heads == 0
+        self.d_head = self.d_model // self.n_heads  # 64
+
+
+class RMSNorm(nn.Module):
+    """RMSNorm with optional bias (no bias by default, LLaMA-style)."""
+
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Compute in float32 for stability, then cast back
+        orig_dtype = x.dtype
+        x = x.float()
+        rms = x.pow(2).mean(dim=-1, keepdim=True)
+        x = x * torch.rsqrt(rms + self.eps)
+        x = x.to(orig_dtype)
+        return x * self.weight
+
+
+def precompute_rope_frequencies(
+    d_head: int,
+    max_seq_len: int,
+    theta: float = 10_000.0,
+    device: torch.device | None = None,
+) -> torch.Tensor:
+    """Precompute RoPE frequency table.
+
+    Returns tensor of shape (max_seq_len, d_head // 2) with complex
+    frequencies (cos, sin interleaved is handled in apply_rope).
+    """
+    inv_freq = 1.0 / (theta ** (torch.arange(0, d_head, 2, device=device).float() / d_head))
+    positions = torch.arange(max_seq_len, device=device).float()
+    freqs = torch.outer(positions, inv_freq)  # (seq, d_head//2)
+    return freqs
+
+
+def apply_rope(
+    x: torch.Tensor,
+    freqs: torch.Tensor,
+) -> torch.Tensor:
+    """Apply rotary position embeddings to tensor x.
+
+    x:     (batch, n_heads, seq, d_head)
+    freqs: (seq, d_head // 2)
+    """
+    seq_len = x.shape[2]
+    d_head = x.shape[-1]
+    freqs = freqs[:seq_len]  # (seq, d_head//2)
+
+    cos = freqs.cos()
+    sin = freqs.sin()
+
+    # Interleave cos/sin to match the rotate_half pattern
+    # x is split into two halves: x1 = x[..., :d//2], x2 = x[..., d//2:]
+    x1 = x[..., : d_head // 2]
+    x2 = x[..., d_head // 2 :]
+
+    # Broadcast cos/sin: (1, 1, seq, d_head//2)
+    cos = cos.unsqueeze(0).unsqueeze(0)
+    sin = sin.unsqueeze(0).unsqueeze(0)
+
+    rotated = torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+    return rotated
+
+
+class Attention(nn.Module):
+    """Multi-head self-attention with RoPE, no biases, causal masking."""
+
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.n_heads = config.n_heads
+        self.d_head = config.d_head
+        self.d_model = config.d_model
+        self.scale = 1.0 / math.sqrt(self.d_head)
+
+        # Fused QKV projection
+        self.qkv = nn.Linear(config.d_model, 3 * config.d_model, bias=False)
+        self.o_proj = nn.Linear(config.d_model, config.d_model, bias=False)
+        self.dropout = config.dropout
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rope_freqs: torch.Tensor,
+        mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        B, T, C = x.shape
+
+        qkv = self.qkv(x)  # (B, T, 3*C)
+        q, k, v = qkv.chunk(3, dim=-1)
+
+        # Reshape to (B, n_heads, T, d_head)
+        q = q.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
+        k = k.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
+        v = v.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
+
+        # Apply RoPE to Q and K
+        q = apply_rope(q, rope_freqs)
+        k = apply_rope(k, rope_freqs)
+
+        # Use PyTorch's scaled_dot_product_attention (uses Flash Attention on CUDA)
+        if mask is not None:
+            # mask: (1, 1, T, T) additive mask
+            attn_mask = mask
+        else:
+            attn_mask = None
+
+        out = F.scaled_dot_product_attention(
+            q, k, v,
+            attn_mask=attn_mask,
+            dropout_p=self.dropout if self.training else 0.0,
+            is_causal=(mask is None),
+        )
+
+        # (B, n_heads, T, d_head) -> (B, T, C)
+        out = out.transpose(1, 2).contiguous().view(B, T, C)
+        return self.o_proj(out)
+
+
+class SwiGLU(nn.Module):
+    """SwiGLU feed-forward network: (xW_gate * SiLU(xW_up)) * W_down."""
+
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.w_gate = nn.Linear(config.d_model, config.d_ff, bias=False)
+        self.w_up = nn.Linear(config.d_model, config.d_ff, bias=False)
+        self.w_down = nn.Linear(config.d_ff, config.d_model, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))
+
+
+class TransformerBlock(nn.Module):
+    """One transformer decoder block: pre-norm attention + pre-norm FFN."""
+
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.norm1 = RMSNorm(config.d_model)
+        self.attn = Attention(config)
+        self.norm2 = RMSNorm(config.d_model)
+        self.ffn = SwiGLU(config)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rope_freqs: torch.Tensor,
+        mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        x = x + self.attn(self.norm1(x), rope_freqs, mask)
+        x = x + self.ffn(self.norm2(x))
+        return x
+
+
+class Retriever500M(nn.Module):
+    """Full decoder-only transformer model."""
+
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.config = config
+
+        # Token embedding (tied with output head)
+        self.token_embedding = nn.Embedding(config.vocab_size, config.d_model)
+
+        # Transformer blocks
+        self.layers = nn.ModuleList([
+            TransformerBlock(config) for _ in range(config.n_layers)
+        ])
+
+        # Final norm
+        self.norm_f = RMSNorm(config.d_model)
+
+        # Output projection (tied with embedding)
+        if config.tie_embeddings:
+            self.lm_head = None  # use token_embedding weight
+        else:
+            self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+
+        # Precompute RoPE frequencies (registered as buffer, moved with .to())
+        freqs = precompute_rope_frequencies(
+            config.d_head,
+            config.max_seq_len,
+            config.rope_theta,
+        )
+        self.register_buffer("rope_freqs", freqs, persistent=False)
+
+        # Causal mask buffer
+        mask = torch.full(
+            (1, 1, config.max_seq_len, config.max_seq_len),
+            float("-inf"),
+        )
+        mask = torch.triu(mask, diagonal=1)
+        self.register_buffer("causal_mask", mask, persistent=False)
+
+        # Initialize weights
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module: nn.Module):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def get_output_weight(self):
+        """Return the weight matrix for the output projection."""
+        if self.config.tie_embeddings:
+            return self.token_embedding.weight
+        return self.lm_head.weight
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        targets: torch.Tensor | None = None,
+        use_checkpoint: bool = False,
+    ) -> dict:
+        B, T = input_ids.shape
+
+        # Token embeddings
+        x = self.token_embedding(input_ids)  # (B, T, d_model)
+
+        # Get RoPE freqs and causal mask for current sequence length
+        rope_freqs = self.rope_freqs[:T]
+        mask = self.causal_mask[:, :, :T, :T]
+
+        # Transformer blocks (with optional gradient checkpointing)
+        for layer in self.layers:
+            if use_checkpoint and self.training:
+                # Gradient checkpointing: recompute activations during backward
+                x = torch.utils.checkpoint.checkpoint(
+                    layer, x, rope_freqs, mask, use_reentrant=False,
+                )
+            else:
+                x = layer(x, rope_freqs, mask)
+
+        x = self.norm_f(x)
+
+        # Output logits
+        logits = F.linear(x, self.get_output_weight())  # (B, T, vocab_size)
+
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(
+                logits.view(-1, logits.size(-1)),
+                targets.view(-1),
+                ignore_index=-100,
+            )
+
+        return {"logits": logits, "loss": loss}
+
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids: torch.Tensor,
+        max_new_tokens: int = 128,
+        temperature: float = 1.0,
+        top_k: int | None = None,
+        eos_token_id: int | None = None,
+    ) -> torch.Tensor:
+        """Simple autoregressive generation."""
+        self.eval()
+        for _ in range(max_new_tokens):
+            # Crop context if it exceeds max_seq_len
+            idx_cond = input_ids if input_ids.size(1) <= self.config.max_seq_len else \
+                input_ids[:, -self.config.max_seq_len:]
+
+            logits = self(idx_cond)["logits"]
+            logits = logits[:, -1, :] / max(temperature, 1e-6)
+
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = float("-inf")
+
+            probs = F.softmax(logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+            input_ids = torch.cat([input_ids, next_token], dim=1)
+
+            if eos_token_id is not None and next_token.item() == eos_token_id:
+                break
+
+        return input_ids
+
+    def count_parameters(self) -> int:
+        """Count total trainable parameters."""
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)
+
+
+def build_model(config: ModelConfig | None = None) -> Retriever500M:
+    """Build the Retriever500M model."""
+    if config is None:
+        config = ModelConfig()
+    model = Retriever500M(config)
+    return model
+
+
+if __name__ == "__main__":
+    config = ModelConfig()
+    model = build_model(config)
+
+    total_params = model.count_parameters()
+    print(f"Model: Retriever500M")
+    print(f"  d_model:    {config.d_model}")
+    print(f"  n_layers:   {config.n_layers}")
+    print(f"  n_heads:    {config.n_heads}")
+    print(f"  d_ff:       {config.d_ff}")
+    print(f"  d_head:     {config.d_head}")
+    print(f"  vocab_size: {config.vocab_size}")
+    print(f"  max_seq_len: {config.max_seq_len}")
+    print(f"  Total parameters: {total_params:,} ({total_params / 1e6:.1f}M)")
+
+    # Quick forward pass test
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = model.to(device)
+    model.eval()
+
+    input_ids = torch.randint(0, config.vocab_size, (2, 64), device=device)
+    with torch.no_grad():
+        out = model(input_ids)
+    print(f"  Output logits shape: {out['logits'].shape}")
+    print("  Forward pass OK.")

src/search_agent.py

@@ -0,0 +1,352 @@
+"""
+Inference-time search agent harness.
+
+This is the runtime that executes the search loop at inference time.
+It:
+  1. Takes a query from the calling (bigger) model
+  2. Formats it as a chat message for the agent model
+  3. Generates the agent's response (which contains <|search|> actions)
+  4. Parses the search query from the response
+  5. Retrieves code chunks from the index (simple keyword/TF-IDF search)
+  6. Feeds the results back to the agent as <|result|> messages
+  7. The agent generates more reasoning or <|evidence|>/<|finish|>
+  8. Returns the evidence package to the caller
+
+The retrieval backend is a simple in-memory keyword search over the
+chunked corpus. This can be replaced with any retrieval backend
+(embedding search, BM25, etc.) — the agent interface is the same.
+
+Usage:
+  python src/search_agent.py --query "How does nginx handle connections?"
+  python src/search_agent.py --interactive
+"""
+
+import argparse
+import json
+import os
+import re
+import sys
+from collections import Counter
+
+import torch
+import torch.nn.functional as F
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from model import ModelConfig, Retriever500M
+from tokenizers import Tokenizer
+
+# ─── Paths ───────────────────────────────────────────────────────────────────
+PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+CHECKPOINT_DIR = os.path.join(PROJECT_DIR, "checkpoints")
+TOKENIZER_PATH = os.path.join(PROJECT_DIR, "tokenizer", "tokenizer_agent.json")
+CHUNKS_PATH = os.path.join(PROJECT_DIR, "data", "chunks.jsonl")
+
+# ─── Special tokens ──────────────────────────────────────────────────────────
+SYSTEM_PROMPT = (
+    "You are a code search agent. Given a query from a reasoning model, "
+    "decompose it into subqueries, search the codebase, inspect results, "
+    "and return curated evidence. Use <|search|> to issue searches, "
+    "<|reasoning|> to analyze, and <|evidence|> to return findings. "
+    "Be concise. Extract only the relevant facts. End with <|finish|>."
+)
+
+
+# ─── Simple retrieval backend ────────────────────────────────────────────────
+
+class KeywordRetriever:
+    """Simple keyword-based retrieval over code chunks.
+
+    For production, replace this with an embedding-based retriever
+    (e.g., mxbai-embed-large or similar). The agent interface stays the same.
+    """
+
+    def __init__(self, chunks_path: str):
+        print(f"Loading chunks from {chunks_path}...")
+        self.chunks = []
+        with open(chunks_path, "r", encoding="utf-8") as f:
+            for line in f:
+                self.chunks.append(json.loads(line))
+        print(f"  Loaded {len(self.chunks):,} chunks")
+
+        # Build simple term frequency index
+        self.chunk_tokens = []
+        for chunk in self.chunks:
+            code = chunk["code"].lower()
+            # Simple tokenization: split on non-alphanumeric
+            tokens = re.findall(r"[a-z_][a-z0-9_]*", code)
+            self.chunk_tokens.append(Counter(tokens))
+
+    def search(self, query: str, top_k: int = 3) -> list[dict]:
+        """Search for chunks matching the query. Returns top_k results."""
+        query_tokens = re.findall(r"[a-z_][a-z0-9_]*", query.lower())
+        if not query_tokens:
+            return []
+
+        scores = []
+        for i, chunk_tf in enumerate(self.chunk_tokens):
+            score = sum(chunk_tf.get(t, 0) for t in query_tokens)
+            # Normalize by chunk length to avoid bias toward long chunks
+            if sum(chunk_tf.values()) > 0:
+                score = score / (1 + sum(chunk_tf.values()) * 0.001)
+            scores.append((score, i))
+
+        scores.sort(reverse=True)
+        results = []
+        for score, idx in scores[:top_k]:
+            if score > 0:
+                chunk = self.chunks[idx]
+                results.append({
+                    "code": chunk["code"],
+                    "filepath": chunk["filepath"],
+                    "name": chunk["name"],
+                    "type": chunk["type"],
+                    "language": chunk["language"],
+                    "score": score,
+                })
+        return results
+
+
+# ─── Agent harness ───────────────────────────────────────────────────────────
+
+class SearchAgent:
+    """The search agent harness that runs the search loop."""
+
+    def __init__(
+        self,
+        checkpoint_path: str,
+        tokenizer_path: str,
+        chunks_path: str,
+        device: torch.device = None,
+        max_search_rounds: int = 5,
+        max_new_tokens: int = 256,
+    ):
+        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.max_search_rounds = max_search_rounds
+        self.max_new_tokens = max_new_tokens
+
+        # Load tokenizer
+        self.tokenizer = Tokenizer.from_file(tokenizer_path)
+
+        # Load model
+        print(f"Loading model from {checkpoint_path}...")
+        ckpt = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
+        config = ModelConfig(**ckpt["config"])
+        self.model = Retriever500M(config).to(self.device)
+        self.model.load_state_dict(ckpt["model_state_dict"])
+        self.model.eval()
+        print(f"  Loaded (step {ckpt.get('step', '?')}, loss {ckpt.get('loss', '?')})")
+
+        # Load retriever
+        self.retriever = KeywordRetriever(chunks_path)
+
+        # Special token IDs
+        vocab = self.tokenizer.get_vocab()
+        self.system_id = vocab.get("<tool_call>", 32000)
+        self.user_id = vocab.get("<tool_call>", 32001)
+        self.assistant_id = vocab.get("<tool_call>", 32002)
+        self.search_id = vocab.get("<|search|>", 32003)
+        self.result_id = vocab.get("<|result|>", 32004)
+        self.evidence_id = vocab.get("<|evidence|>", 32005)
+        self.reasoning_id = vocab.get("<|reasoning|>", 32006)
+        self.finish_id = vocab.get("<|finish|>", 32007)
+        self.end_id = vocab.get("<|end|>", 32008)
+
+    def _encode(self, text: str) -> list[int]:
+        """Encode text to token IDs."""
+        return self.tokenizer.encode(text).ids
+
+    def _decode(self, ids: list[int]) -> str:
+        """Decode token IDs to text."""
+        return self.tokenizer.decode(ids)
+
+    def _generate(self, input_ids: torch.Tensor, max_new_tokens: int) -> str:
+        """Generate text from the model, stopping at <|end|> or <|finish|>."""
+        with torch.no_grad():
+            for _ in range(max_new_tokens):
+                # Crop context if too long
+                if input_ids.size(1) > self.model.config.max_seq_len:
+                    input_ids = input_ids[:, -self.model.config.max_seq_len:]
+
+                logits = self.model(input_ids)["logits"]
+                next_logits = logits[:, -1, :]
+
+                # Apply temperature and sample
+                probs = F.softmax(next_logits / 0.8, dim=-1)
+                next_token = torch.multinomial(probs, num_samples=1)
+                input_ids = torch.cat([input_ids, next_token], dim=1)
+
+                # Stop on <|end|> or <|finish|>
+                if next_token.item() == self.end_id or next_token.item() == self.finish_id:
+                    break
+
+        # Decode the generated part (after the input)
+        generated_ids = input_ids[0, -max_new_tokens:].tolist()
+        return self._decode(generated_ids)
+
+    def _parse_search_query(self, text: str) -> str | None:
+        """Extract the search query from the agent's response."""
+        # Look for <|search|>query<|end|>
+        match = re.search(r"<\|search\|>(.*?)<\|end\|>", text, re.DOTALL)
+        if match:
+            return match.group(1).strip()
+        return None
+
+    def _parse_evidence(self, text: str) -> str | None:
+        """Extract the evidence from the agent's response."""
+        match = re.search(r"<\|evidence\|>(.*?)(?:<\|end\|>|<\|finish\|>|$)", text, re.DOTALL)
+        if match:
+            return match.group(1).strip()
+        return None
+
+    def _has_finish(self, text: str) -> bool:
+        """Check if the agent has signaled completion."""
+        return "<|finish|>" in text
+
+    def search(self, query: str) -> dict:
+        """Run the full search loop for a query.
+
+        Returns:
+          {
+            "query": the original query,
+            "evidence": the curated evidence (or None if not found),
+            "searches": list of search queries issued,
+            "results": list of all results retrieved,
+            "trace": the full conversation trace,
+          }
+        """
+        print(f"\n{'='*60}")
+        print(f"QUERY: {query}")
+        print(f"{'='*60}")
+
+        # Build initial context
+        trace = []
+
+        # System prompt
+        system_tokens = [self.system_id] + self._encode(SYSTEM_PROMPT) + [self.end_id]
+        trace.append({"role": "system", "tokens": system_tokens})
+
+        # User query
+        user_tokens = [self.user_id] + self._encode(query) + [self.end_id]
+        trace.append({"role": "user", "tokens": user_tokens})
+
+        all_searches = []
+        all_results = []
+        evidence = None
+
+        for round_num in range(self.max_search_rounds):
+            # Build input from trace
+            all_tokens = []
+            for entry in trace:
+                all_tokens.extend(entry["tokens"])
+
+            input_ids = torch.tensor([all_tokens], dtype=torch.long, device=self.device)
+
+            # Generate agent response
+            print(f"\n--- Round {round_num + 1} ---")
+            response = self._generate(input_ids, self.max_new_tokens)
+            print(f"Agent: {response[:200]}...")
+
+            # Add assistant tokens to trace
+            assistant_tokens = [self.assistant_id] + self._encode(response)
+            if not response.endswith("<|end|>"):
+                assistant_tokens.append(self.end_id)
+            trace.append({"role": "assistant", "tokens": assistant_tokens})
+
+            # Check for finish
+            if self._has_finish(response):
+                evidence = self._parse_evidence(response)
+                print(f"\n[EVIDENCE]: {evidence}")
+                break
+
+            # Parse search query
+            search_query = self._parse_search_query(response)
+            if search_query:
+                print(f"[SEARCH]: {search_query}")
+                all_searches.append(search_query)
+
+                # Retrieve results
+                results = self.retriever.search(search_query, top_k=3)
+
+                if results:
+                    for result in results:
+                        print(f"  [RESULT]: {result['name']} ({result['language']}, score={result['score']:.2f})")
+                        result_tokens = [self.result_id] + self._encode(result["code"]) + [self.end_id]
+                        trace.append({"role": "result", "tokens": result_tokens, "data": result})
+                        all_results.append(result)
+                else:
+                    print("  [NO RESULTS]")
+                    result_tokens = [self.result_id, self.end_id]
+                    trace.append({"role": "result", "tokens": result_tokens})
+            else:
+                # No search query found — try to extract evidence directly
+                evidence = self._parse_evidence(response)
+                if evidence:
+                    print(f"\n[EVIDENCE]: {evidence}")
+                    break
+                else:
+                    print("[WARNING] No search or evidence found, continuing...")
+
+        return {
+            "query": query,
+            "evidence": evidence,
+            "searches": all_searches,
+            "results": all_results,
+            "trace": trace,
+        }
+
+
+# ─── CLI ─────────────────────────────────────────────────────────────────────
+
+def main():
+    parser = argparse.ArgumentParser(description="Run the search agent")
+    parser.add_argument("--checkpoint", type=str, default=os.path.join(CHECKPOINT_DIR, "sft_latest.pt"))
+    parser.add_argument("--query", type=str, default=None, help="Query to search for")
+    parser.add_argument("--interactive", action="store_true", help="Interactive mode")
+    parser.add_argument("--max_rounds", type=int, default=5, help="Max search rounds")
+    args = parser.parse_args()
+
+    agent = SearchAgent(
+        checkpoint_path=args.checkpoint,
+        tokenizer_path=TOKENIZER_PATH,
+        chunks_path=CHUNKS_PATH,
+        max_search_rounds=args.max_rounds,
+    )
+
+    if args.interactive:
+        print("\nInteractive mode. Type 'quit' to exit.")
+        while True:
+            query = input("\nQuery> ").strip()
+            if query.lower() in ("quit", "exit", "q"):
+                break
+            if query:
+                result = agent.search(query)
+                print(f"\n{'='*60}")
+                print(f"FINAL EVIDENCE:")
+                print(f"{'='*60}")
+                print(result["evidence"] or "No evidence found.")
+    elif args.query:
+        result = agent.search(args.query)
+        print(f"\n{'='*60}")
+        print(f"FINAL EVIDENCE:")
+        print(f"{'='*60}")
+        print(result["evidence"] or "No evidence found.")
+    else:
+        # Run sample queries
+        sample_queries = [
+            "How does nginx handle reusable connections?",
+            "What does the with_params_help decorator do?",
+            "What fields does the ngx_listening_s struct have?",
+            "How does the concatenate function work?",
+            "Where is the database connection pool implemented?",
+        ]
+        for query in sample_queries:
+            result = agent.search(query)
+            print(f"\n{'='*60}")
+            print(f"FINAL EVIDENCE:")
+            print(f"{'='*60}")
+            print(result["evidence"] or "No evidence found.")
+
+
+if __name__ == "__main__":
+    main()

src/sft.py

@@ -0,0 +1,526 @@
+"""
+Supervised fine-tuning (SFT) script for the search agent.
+
+Key differences from pretraining:
+  1. Uses the agent tokenizer (32009 vocab, includes special tokens)
+  2. Formats traces as chat: <tool_call>...<|end|><tool_call>...<|end|><tool_call>...<|end|>
+  3. Loss masking: only compute loss on ASSISTANT tokens, not on
+     system/user/result tokens (the model should learn to GENERATE
+     the agent responses, not predict the inputs)
+  4. Lower learning rate (5e-5) — fine-tuning, not pretraining
+  5. Resizes model embedding to 32009 (from 32000)
+
+The trace format in the training data:
+  <tool_call>system_prompt<|end|>
+  <tool_call>user_query<|end|>
+  <tool_call>assistant_turn_1<|end|>      ← loss computed here
+  [result injected by harness]
+  <tool_call>result_content<|end|>
+  <tool_call>assistant_turn_2<|end|>      ← loss computed here
+  ...
+  <tool_call>evidence<|finish|><|end|>    ← loss computed here
+
+Usage:
+  python src/sft.py --steps N [--resume] [--batch_size N] [--lr F]
+"""
+
+import argparse
+import json
+import os
+import sys
+import time
+from dataclasses import asdict
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from model import ModelConfig, Retriever500M
+from tokenizers import Tokenizer
+
+# ─── Paths ───────────────────────────────────────────────────────────────────
+PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+DATA_DIR = os.path.join(PROJECT_DIR, "data")
+TOKENIZER_DIR = os.path.join(PROJECT_DIR, "tokenizer")
+CHECKPOINT_DIR = os.path.join(PROJECT_DIR, "checkpoints")
+LOGS_DIR = os.path.join(PROJECT_DIR, "logs")
+
+TRACES_PATH = os.path.join(DATA_DIR, "sft_traces.jsonl")
+GOLD_PATH = os.path.join(DATA_DIR, "gold_traces.jsonl")
+TOKENIZER_PATH = os.path.join(TOKENIZER_DIR, "tokenizer_agent.json")
+SPECIAL_TOKENS_PATH = os.path.join(TOKENIZER_DIR, "special_tokens.json")
+
+# ─── Special token IDs ───────────────────────────────────────────────────────
+# Loaded from special_tokens.json, but hardcoded as fallback
+SYSTEM_ID = 32000
+USER_ID = 32001
+ASSISTANT_ID = 32002
+SEARCH_ID = 32003
+RESULT_ID = 32004
+EVIDENCE_ID = 32005
+REASONING_ID = 32006
+FINISH_ID = 32007
+END_ID = 32008
+
+
+def load_special_tokens():
+    """Load special token IDs from the mapping file."""
+    global SYSTEM_ID, USER_ID, ASSISTANT_ID, SEARCH_ID, RESULT_ID
+    global EVIDENCE_ID, REASONING_ID, FINISH_ID, END_ID
+
+    if os.path.exists(SPECIAL_TOKENS_PATH):
+        with open(SPECIAL_TOKENS_PATH, "r") as f:
+            data = json.load(f)
+        ids = data["token_ids"]
+        SYSTEM_ID = ids.get("<tool_call>", 32000)
+        USER_ID = ids.get("<tool_call>", 32001)
+        ASSISTANT_ID = ids.get("<tool_call>", 32002)
+        SEARCH_ID = ids.get("<|search|>", 32003)
+        RESULT_ID = ids.get("<|result|>", 32004)
+        EVIDENCE_ID = ids.get("<|evidence|>", 32005)
+        REASONING_ID = ids.get("<|reasoning|>", 32006)
+        FINISH_ID = ids.get("<|finish|>", 32007)
+        END_ID = ids.get("<|end|>", 32008)
+
+
+# ─── Data formatting ─────────────────────────────────────────────────────────
+
+def format_trace_to_tokens(trace: dict, tokenizer: Tokenizer, max_seq_len: int = 768) -> tuple[np.ndarray, np.ndarray]:
+    """Convert a trace to (input_ids, loss_mask) arrays.
+
+    loss_mask[i] = 1 if we should compute loss on token i, 0 otherwise.
+    Loss is only computed on ASSISTANT turns (and the special tokens within them).
+    """
+    messages = trace["trace"]
+
+    all_tokens = []
+    loss_mask = []
+
+    for msg in messages:
+        role = msg["role"]
+        content = msg["content"]
+
+        if role == "system":
+            # System: <tool_call>content<|end|> — no loss
+            role_id = SYSTEM_ID
+            tokens = [role_id] + tokenizer.encode(content).ids + [END_ID]
+            all_tokens.extend(tokens)
+            loss_mask.extend([0] * len(tokens))
+
+        elif role == "user":
+            # User: <tool_call>content<|end|> — no loss
+            role_id = USER_ID
+            tokens = [role_id] + tokenizer.encode(content).ids + [END_ID]
+            all_tokens.extend(tokens)
+            loss_mask.extend([0] * len(tokens))
+
+        elif role == "assistant":
+            # Assistant: <tool_call>content<|end|> — LOSS on all tokens
+            role_id = ASSISTANT_ID
+            # The content already contains <|search|>, <|reasoning|>, etc. as text
+            # We need to encode them properly
+            tokens = [role_id] + tokenizer.encode(content).ids + [END_ID]
+            all_tokens.extend(tokens)
+            loss_mask.extend([1] * len(tokens))
+
+        elif role == "result":
+            # Result: injected by harness, no loss
+            # Format as <|result|>content<|end|>
+            if content:
+                tokens = [RESULT_ID] + tokenizer.encode(content).ids + [END_ID]
+            else:
+                tokens = [RESULT_ID, END_ID]
+            all_tokens.extend(tokens)
+            loss_mask.extend([0] * len(tokens))
+
+    # Truncate to max_seq_len
+    if len(all_tokens) > max_seq_len:
+        all_tokens = all_tokens[:max_seq_len]
+        loss_mask = loss_mask[:max_seq_len]
+
+    return np.array(all_tokens, dtype=np.int32), np.array(loss_mask, dtype=np.int32)
+
+
+def load_sft_dataset(traces_path: str, tokenizer: Tokenizer, max_seq_len: int = 768) -> list[tuple[np.ndarray, np.ndarray]]:
+    """Load all SFT traces and convert to (input_ids, loss_mask) pairs."""
+    print(f"Loading SFT traces from {traces_path}...")
+    dataset = []
+    with open(traces_path, "r", encoding="utf-8") as f:
+        for line in f:
+            trace = json.loads(line)
+            ids, mask = format_trace_to_tokens(trace, tokenizer, max_seq_len)
+            if len(ids) > 10:  # skip traces that are too short
+                dataset.append((ids, mask))
+    print(f"  Loaded {len(dataset):,} traces")
+    return dataset
+
+
+def get_sft_batch(
+    dataset: list[tuple[np.ndarray, np.ndarray]],
+    batch_size: int,
+    seq_len: int,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Sample a batch of SFT traces.
+
+    Returns (input_ids, targets, loss_mask) where:
+      - input_ids: (B, T) token IDs
+      - targets: (B, T) shifted targets (next token prediction)
+      - loss_mask: (B, T) 1 where loss should be computed, 0 elsewhere
+    """
+    # Randomly sample traces
+    indices = np.random.randint(0, len(dataset), size=batch_size)
+
+    input_ids_list = []
+    loss_mask_list = []
+
+    for idx in indices:
+        ids, mask = dataset[idx]
+        # Pad or truncate to seq_len
+        if len(ids) < seq_len:
+            pad_len = seq_len - len(ids)
+            ids = np.concatenate([ids, np.zeros(pad_len, dtype=np.int32)])
+            mask = np.concatenate([mask, np.zeros(pad_len, dtype=np.int32)])
+        else:
+            ids = ids[:seq_len]
+            mask = mask[:seq_len]
+        input_ids_list.append(ids)
+        loss_mask_list.append(mask)
+
+    input_ids = torch.from_numpy(np.stack(input_ids_list)).long().to(device)
+    loss_mask = torch.from_numpy(np.stack(loss_mask_list)).long().to(device)
+
+    # Targets are shifted input_ids
+    targets = torch.cat([input_ids[:, 1:], torch.zeros_like(input_ids[:, :1])], dim=1)
+
+    return input_ids, targets, loss_mask
+
+
+# ─── Training ────────────────────────────────────────────────────────────────
+
+def setup_optimizer(model, lr, use_8bit=True):
+    """Same as pretraining optimizer setup."""
+    decay_params, no_decay_params = [], []
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+        if "embedding" in name or "norm" in name:
+            no_decay_params.append(param)
+        else:
+            decay_params.append(param)
+
+    param_groups = [
+        {"params": decay_params, "weight_decay": 0.1},
+        {"params": no_decay_params, "weight_decay": 0.0},
+    ]
+
+    if use_8bit:
+        try:
+            import bitsandbytes as bnb
+            optimizer = bnb.optim.AdamW8bit(param_groups, lr=lr, betas=(0.9, 0.95), eps=1e-8)
+            print("Using 8-bit AdamW (bitsandbytes)")
+            return optimizer
+        except Exception as e:
+            print(f"8-bit optimizer unavailable ({e}), falling back to AdamW")
+
+    optimizer = torch.optim.AdamW(param_groups, lr=lr, betas=(0.9, 0.95), eps=1e-8)
+    print("Using standard AdamW")
+    return optimizer
+
+
+def get_lr(step, warmup, max_steps, max_lr, min_lr):
+    """Cosine LR schedule with linear warmup."""
+    if step < warmup:
+        return max_lr * (step + 1) / warmup
+    if step > max_steps:
+        return min_lr
+    decay_ratio = (step - warmup) / (max_steps - warmup)
+    coeff = 0.5 * (1.0 + np.cos(np.pi * decay_ratio))
+    return min_lr + coeff * (max_lr - min_lr)
+
+
+def resize_embeddings(model, new_vocab_size):
+    """Resize the model's token embedding to accommodate new tokens."""
+    old_size = model.token_embedding.weight.shape[0]
+    if old_size == new_vocab_size:
+        return
+
+    print(f"Resizing embeddings: {old_size} -> {new_vocab_size}")
+    d_model = model.config.d_model
+
+    # Create new embedding with the old weights + new random weights
+    old_weight = model.token_embedding.weight.data
+    new_embedding = nn.Embedding(new_vocab_size, d_model)
+    new_embedding.weight.data[:old_size] = old_weight
+    # Initialize new tokens with small random values
+    nn.init.normal_(new_embedding.weight.data[old_size:], mean=0.0, std=0.02)
+    model.token_embedding = new_embedding
+
+    # If tied, the output weight is the same, so nothing else to do
+    # If not tied, resize lm_head too
+    if not model.config.tie_embeddings and model.lm_head is not None:
+        model.lm_head = nn.Linear(d_model, new_vocab_size, bias=False)
+        with torch.no_grad():
+            model.lm_head.weight.data[:old_size] = old_weight
+            nn.init.normal_(model.lm_head.weight.data[old_size:], mean=0.0, std=0.02)
+
+
+from torch import nn
+
+
+def train(args):
+    load_special_tokens()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Device: {device}")
+    if device.type == "cuda":
+        print(f"GPU: {torch.cuda.get_device_name(0)}")
+
+    os.makedirs(CHECKPOINT_DIR, exist_ok=True)
+    os.makedirs(LOGS_DIR, exist_ok=True)
+
+    # ─── Tokenizer ───────────────────────────────────────────────────────────
+    print("Loading agent tokenizer...")
+    tokenizer = Tokenizer.from_file(TOKENIZER_PATH)
+    vocab_size = tokenizer.get_vocab_size()
+    print(f"Vocab size: {vocab_size}")
+
+    # ─── Data ────────────────────────────────────────────────────────────────
+    dataset = load_sft_dataset(TRACES_PATH, tokenizer, args.seq_len)
+
+    # Also load gold traces
+    gold_dataset = load_sft_dataset(GOLD_PATH, tokenizer, args.seq_len)
+    dataset.extend(gold_dataset)
+    print(f"  Total (with gold): {len(dataset):,}")
+
+    # ─── Model ───────────────────────────────────────────────────────────────
+    config = ModelConfig(
+        vocab_size=vocab_size,  # 32009
+        d_model=1_280,
+        n_layers=23,
+        n_heads=20,
+        d_ff=3_456,
+        max_seq_len=args.seq_len,
+        dropout=0.0,
+        tie_embeddings=True,
+    )
+
+    model = Retriever500M(config).to(device)
+
+    # ─── Load pretrained checkpoint ──────────────────────────────────────────
+    ckpt_path = args.resume_path or os.path.join(CHECKPOINT_DIR, "latest.pt")
+    if os.path.exists(ckpt_path):
+        print(f"Loading pretrained weights from {ckpt_path}...")
+        ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
+        old_config = ModelConfig(**ckpt["config"])
+
+        # Load state dict, handling vocab size mismatch
+        state_dict = ckpt["model_state_dict"]
+        old_vocab = old_config.vocab_size
+
+        if old_vocab != vocab_size:
+            print(f"  Vocab size mismatch: {old_vocab} -> {vocab_size}")
+            print(f"  Resizing embeddings in state dict...")
+            # Resize token_embedding in the state dict
+            old_weight = state_dict["token_embedding.weight"]
+            d_model = old_weight.shape[1]
+            new_weight = torch.zeros(vocab_size, d_model)
+            new_weight[:old_vocab] = old_weight
+            nn.init.normal_(new_weight[old_vocab:], mean=0.0, std=0.02)
+            state_dict["token_embedding.weight"] = new_weight
+
+        model.load_state_dict(state_dict)
+        print(f"  Loaded (step {ckpt.get('step', '?')}, loss {ckpt.get('loss', '?')})")
+    else:
+        print(f"WARNING: No checkpoint at {ckpt_path}, starting from scratch!")
+
+    total_params = model.count_parameters()
+    print(f"Model parameters: {total_params:,} ({total_params / 1e6:.1f}M)")
+
+    # ─── Optimizer ───────────────────────────────────────────────────────────
+    optimizer = setup_optimizer(model, args.lr, use_8bit=args.use_8bit_adam)
+
+    # ─── Training loop ───────────────────────────────────────────────────────
+    effective_batch = args.batch_size * args.grad_accum
+    print(f"\nSFT configuration:")
+    print(f"  Batch size:       {args.batch_size}")
+    print(f"  Grad accum:       {args.grad_accum}")
+    print(f"  Effective batch:  {effective_batch}")
+    print(f"  Sequence length:  {args.seq_len}")
+    print(f"  Learning rate:    {args.lr}")
+    print(f"  Steps:            {args.steps}")
+    print(f"  Warmup:           {args.warmup}")
+    print()
+
+    log = {
+        "config": asdict(config),
+        "train_args": vars(args),
+        "total_params": total_params,
+        "steps": [],
+    }
+
+    model.train()
+    start_time = time.time()
+    accum_loss = 0.0
+    best_loss = float("inf")
+
+    pbar = tqdm(range(args.steps), desc="SFT")
+    for step in pbar:
+        lr = get_lr(step, args.warmup, args.steps, args.lr, args.lr * 0.1)
+        for pg in optimizer.param_groups:
+            pg["lr"] = lr
+
+        optimizer.zero_grad(set_to_none=True)
+
+        total_loss = 0.0
+        for micro_step in range(args.grad_accum):
+            input_ids, targets, loss_mask = get_sft_batch(
+                dataset, args.batch_size, args.seq_len, device
+            )
+
+            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+                out = model(input_ids, targets=targets, use_checkpoint=args.grad_checkpoint)
+                loss = out["loss"]
+
+                # Apply loss mask — only compute loss on assistant tokens
+                # loss is already computed over all positions; we need to recompute
+                # with the mask
+                logits = out["logits"]
+                # Recompute loss with mask
+                if loss_mask.sum() > 0:
+                    # Shift mask to align with next-token prediction
+                    shifted_mask = loss_mask[:, 1:].contiguous()
+                    masked_logits = logits[:, :-1, :].contiguous()
+                    masked_targets = targets[:, :-1].contiguous()
+
+                    # Flatten and apply mask
+                    flat_logits = masked_logits.view(-1, masked_logits.size(-1))
+                    flat_targets = masked_targets.view(-1)
+                    flat_mask = shifted_mask.view(-1).float()
+
+                    per_token_loss = F.cross_entropy(
+                        flat_logits, flat_targets,
+                        ignore_index=-100, reduction="none"
+                    )
+                    masked_loss = (per_token_loss * flat_mask).sum() / flat_mask.sum().clamp(min=1)
+                    loss = masked_loss / args.grad_accum
+                else:
+                    loss = loss / args.grad_accum
+
+            loss.backward()
+            total_loss += loss.item()
+
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        optimizer.step()
+
+        avg_loss = total_loss
+        accum_loss = accum_loss * 0.95 + avg_loss * 0.05
+
+        if step % args.log_every == 0 or step == args.steps - 1:
+            elapsed = time.time() - start_time
+            steps_per_sec = (step + 1) / elapsed
+            vram_used = torch.cuda.max_memory_allocated() / 1e9 if device.type == "cuda" else 0
+
+            log_entry = {
+                "step": step,
+                "loss": avg_loss,
+                "ema_loss": accum_loss,
+                "lr": lr,
+                "elapsed_s": elapsed,
+                "steps_per_sec": steps_per_sec,
+                "vram_gb": vram_used,
+            }
+            log["steps"].append(log_entry)
+
+            pbar.set_postfix({
+                "loss": f"{avg_loss:.4f}",
+                "ema": f"{accum_loss:.4f}",
+                "lr": f"{lr:.2e}",
+                "vram": f"{vram_used:.1f}G",
+            })
+
+        if (step + 1) % args.save_every == 0 or step == args.steps - 1:
+            ckpt_path = os.path.join(CHECKPOINT_DIR, f"sft_step_{step + 1}.pt")
+            torch.save({
+                "model_state_dict": model.state_dict(),
+                "optimizer_state_dict": optimizer.state_dict(),
+                "config": asdict(config),
+                "step": step + 1,
+                "loss": accum_loss,
+            }, ckpt_path)
+            print(f"\n  Saved checkpoint: {ckpt_path}")
+
+            latest_path = os.path.join(CHECKPOINT_DIR, "sft_latest.pt")
+            torch.save({
+                "model_state_dict": model.state_dict(),
+                "config": asdict(config),
+                "step": step + 1,
+                "loss": accum_loss,
+            }, latest_path)
+
+            if accum_loss < best_loss:
+                best_loss = accum_loss
+                best_path = os.path.join(CHECKPOINT_DIR, "sft_best.pt")
+                torch.save({
+                    "model_state_dict": model.state_dict(),
+                    "config": asdict(config),
+                    "step": step + 1,
+                    "loss": accum_loss,
+                }, best_path)
+
+        if step % 50 == 0 and device.type == "cuda":
+            torch.cuda.reset_peak_memory_stats()
+
+    # Save log
+    log_path = os.path.join(LOGS_DIR, "sft_log.json")
+    with open(log_path, "w") as f:
+        json.dump(log, f, indent=2)
+
+    total_time = time.time() - start_time
+    print(f"\nSFT complete!")
+    print(f"  Total time:     {total_time:.1f}s ({total_time/60:.1f} min)")
+    print(f"  Final EMA loss: {accum_loss:.4f}")
+    print(f"  Best loss:      {best_loss:.4f}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="SFT the search agent")
+    parser.add_argument("--steps", type=int, default=500, help="Total SFT steps")
+    parser.add_argument("--batch_size", type=int, default=4, help="Micro batch size")
+    parser.add_argument("--grad_accum", type=int, default=4, help="Gradient accumulation")
+    parser.add_argument("--seq_len", type=int, default=768, help="Sequence length")
+    parser.add_argument("--lr", type=float, default=5e-5, help="Peak learning rate")
+    parser.add_argument("--warmup", type=int, default=20, help="Warmup steps")
+    parser.add_argument("--save_every", type=int, default=100, help="Save every N steps")
+    parser.add_argument("--log_every", type=int, default=10, help="Log every N steps")
+    parser.add_argument("--grad_checkpoint", action="store_true", default=True)
+    parser.add_argument("--no_grad_checkpoint", dest="grad_checkpoint", action="store_false")
+    parser.add_argument("--use_8bit_adam", action="store_true", default=True)
+    parser.add_argument("--no_8bit_adam", dest="use_8bit_adam", action="store_false")
+    parser.add_argument("--resume_path", type=str, default=None, help="Pretrained checkpoint to start from")
+    parser.add_argument("--project_dir", type=str, default=None, help="Override project directory (for Colab)")
+    parser.add_argument("--save_dir", type=str, default=None, help="Override checkpoint save directory (for Google Drive)")
+    args = parser.parse_args()
+
+    # Override paths for Colab
+    global PROJECT_DIR, DATA_DIR, TOKENIZER_DIR, CHECKPOINT_DIR, LOGS_DIR
+    global TRACES_PATH, GOLD_PATH, TOKENIZER_PATH, SPECIAL_TOKENS_PATH
+    if args.project_dir:
+        PROJECT_DIR = args.project_dir
+        DATA_DIR = os.path.join(PROJECT_DIR, "data")
+        TOKENIZER_DIR = os.path.join(PROJECT_DIR, "tokenizer")
+        LOGS_DIR = os.path.join(PROJECT_DIR, "logs")
+        TRACES_PATH = os.path.join(DATA_DIR, "sft_traces.jsonl")
+        GOLD_PATH = os.path.join(DATA_DIR, "gold_traces.jsonl")
+        TOKENIZER_PATH = os.path.join(TOKENIZER_DIR, "tokenizer_agent.json")
+        SPECIAL_TOKENS_PATH = os.path.join(TOKENIZER_DIR, "special_tokens.json")
+    if args.save_dir:
+        CHECKPOINT_DIR = args.save_dir
+
+    train(args)
+
+
+if __name__ == "__main__":
+    main()

src/sft_dry_run.py

@@ -0,0 +1,43 @@
+"""Dry-run SFT: load data, build model, resize embeddings, one forward pass."""
+import sys, os, json, torch
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__))))
+
+from sft import load_special_tokens, load_sft_dataset, get_sft_batch
+from model import ModelConfig, Retriever500M
+from tokenizers import Tokenizer
+
+def main():
+    load_special_tokens()
+    tok = Tokenizer.from_file("tokenizer/tokenizer_agent.json")
+    print("Loading dataset...")
+    dataset = load_sft_dataset("data/sft_traces.jsonl", tok, max_seq_len=768)
+    print(f"Dataset size: {len(dataset)}")
+
+    config = ModelConfig(vocab_size=32009, d_model=1280, n_layers=23, n_heads=20,
+                         d_ff=3456, max_seq_len=768, tie_embeddings=True)
+    model = Retriever500M(config).cuda()
+    print(f"Model: {model.count_parameters()/1e6:.1f}M params")
+
+    # Load old checkpoint and resize embeddings
+    ckpt = torch.load("checkpoints/latest.pt", map_location="cuda", weights_only=False)
+    state = ckpt["model_state_dict"]
+    old_w = state["token_embedding.weight"]
+    new_w = torch.zeros(32009, 1280)
+    new_w[:32000] = old_w
+    torch.nn.init.normal_(new_w[32000:], mean=0.0, std=0.02)
+    state["token_embedding.weight"] = new_w
+    model.load_state_dict(state)
+    print("Checkpoint loaded with resized embeddings")
+
+    # One forward pass
+    input_ids, targets, loss_mask = get_sft_batch(dataset, batch_size=2, seq_len=768, device=torch.device("cuda"))
+    print(f"Batch: {input_ids.shape}")
+
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        out = model(input_ids, targets=targets)
+        print(f"Forward pass OK: loss={out['loss'].item():.4f}")
+
+    print("SFT DRY RUN PASSED")
+
+if __name__ == "__main__":
+    main()

src/train.py

@@ -0,0 +1,401 @@
+"""
+Base pretraining script for Retriever500M.
+
+Memory optimizations for 8GB VRAM (RTX 4070 Laptop):
+  - bf16 mixed precision (autocast)
+  - Gradient checkpointing (recompute activations during backward)
+  - 8-bit AdamW optimizer (bitsandbytes) — halves optimizer state memory
+  - Gradient accumulation (effective batch size > micro batch size)
+  - Short sequence length (512 tokens) for base training
+  - Tied embeddings (shared input/output weight)
+  - Flash Attention via torch SDPA
+
+Usage:
+  python src/train.py [--steps N] [--seq_len N] [--batch_size N] [--grad_accum N]
+"""
+
+import argparse
+import json
+import os
+import sys
+import time
+from dataclasses import asdict
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+
+# Add src to path
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from model import ModelConfig, Retriever500M
+from tokenizers import Tokenizer
+
+# ─── Paths ───────────────────────────────────────────────────────────────────
+PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+DATA_DIR = os.path.join(PROJECT_DIR, "data")
+TOKENIZER_DIR = os.path.join(PROJECT_DIR, "tokenizer")
+CHECKPOINT_DIR = os.path.join(PROJECT_DIR, "checkpoints")
+LOGS_DIR = os.path.join(PROJECT_DIR, "logs")
+
+CORPUS_PATH = os.path.join(DATA_DIR, "corpus.txt")
+CURATED_CORPUS_PATH = os.path.join(DATA_DIR, "corpus_curated.txt")
+TOKENIZER_PATH = os.path.join(TOKENIZER_DIR, "tokenizer.json")
+
+# ─── Data Loading ────────────────────────────────────────────────────────────
+
+def load_and_tokenize(corpus_path: str, tokenizer: Tokenizer) -> np.ndarray:
+    """Load corpus, tokenize everything, return a flat numpy array of token IDs."""
+    print(f"Loading corpus from {corpus_path}...")
+    with open(corpus_path, "r", encoding="utf-8") as f:
+        text = f.read()
+
+    print(f"Corpus size: {len(text) / 1e6:.1f} MB")
+
+    # Tokenize in chunks to avoid memory issues
+    chunk_size = 1_000_000  # 1MB chunks
+    all_tokens = []
+
+    print("Tokenizing corpus...")
+    for i in tqdm(range(0, len(text), chunk_size)):
+        chunk = text[i : i + chunk_size]
+        encoded = tokenizer.encode(chunk)
+        all_tokens.extend(encoded.ids)
+
+    tokens = np.array(all_tokens, dtype=np.int32)
+    print(f"Total tokens: {len(tokens):,}")
+    return tokens
+
+
+def get_batch(
+    tokens: np.ndarray,
+    batch_size: int,
+    seq_len: int,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Sample a random batch of sequences from the token array.
+
+    Returns (input_ids, targets) where targets are shifted by 1.
+    """
+    # Random starting indices
+    max_start = len(tokens) - seq_len - 1
+    indices = np.random.randint(0, max_start, size=batch_size)
+
+    # Gather sequences
+    input_ids = np.stack([tokens[i : i + seq_len] for i in indices])
+    targets = np.stack([tokens[i + 1 : i + seq_len + 1] for i in indices])
+
+    input_ids = torch.from_numpy(input_ids).long().to(device)
+    targets = torch.from_numpy(targets).long().to(device)
+
+    return input_ids, targets
+
+
+# ─── Training ────────────────────────────────────────────────────────────────
+
+def setup_optimizer(model: Retriever500M, lr: float, use_8bit: bool = True):
+    """Set up optimizer — 8-bit AdamW if available, else standard AdamW."""
+    # Separate embedding params (no weight decay) from rest
+    decay_params = []
+    no_decay_params = []
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+        if "embedding" in name or "norm" in name or "weight" in name and ".weight" not in name:
+            no_decay_params.append(param)
+        else:
+            decay_params.append(param)
+
+    param_groups = [
+        {"params": decay_params, "weight_decay": 0.1},
+        {"params": no_decay_params, "weight_decay": 0.0},
+    ]
+
+    if use_8bit:
+        try:
+            import bitsandbytes as bnb
+            optimizer = bnb.optim.AdamW8bit(
+                param_groups, lr=lr, betas=(0.9, 0.95), eps=1e-8,
+            )
+            print("Using 8-bit AdamW (bitsandbytes)")
+            return optimizer
+        except Exception as e:
+            print(f"8-bit optimizer unavailable ({e}), falling back to AdamW")
+
+    optimizer = torch.optim.AdamW(
+        param_groups, lr=lr, betas=(0.9, 0.95), eps=1e-8,
+    )
+    print("Using standard AdamW")
+    return optimizer
+
+
+def resume_from_checkpoint(
+    model: Retriever500M,
+    optimizer,
+    resume_path: str,
+    device: torch.device,
+) -> tuple[int, float]:
+    """Load model + optimizer state from a checkpoint.
+
+    Returns (start_step, best_loss) so the training loop can continue.
+    """
+    print(f"Resuming from {resume_path}")
+    ckpt = torch.load(resume_path, map_location=device, weights_only=False)
+
+    model.load_state_dict(ckpt["model_state_dict"])
+    print(f"  Loaded model weights (step {ckpt.get('step', '?')})")
+
+    if "optimizer_state_dict" in ckpt:
+        try:
+            optimizer.load_state_dict(ckpt["optimizer_state_dict"])
+            print("  Loaded optimizer state")
+        except Exception as e:
+            print(f"  Could not load optimizer state ({e}); starting fresh optimizer")
+
+    start_step = int(ckpt.get("step", 0))
+    best_loss = float(ckpt.get("loss", float("inf")))
+    print(f"  Resuming at step {start_step} (best_loss={best_loss:.4f})")
+    return start_step, best_loss
+
+
+def get_lr(step: int, warmup_steps: int, max_steps: int, max_lr: float, min_lr: float) -> float:
+    """Cosine learning rate schedule with linear warmup."""
+    if step < warmup_steps:
+        return max_lr * (step + 1) / warmup_steps
+    if step > max_steps:
+        return min_lr
+    decay_ratio = (step - warmup_steps) / (max_steps - warmup_steps)
+    coeff = 0.5 * (1.0 + np.cos(np.pi * decay_ratio))
+    return min_lr + coeff * (max_lr - min_lr)
+
+
+def train(args):
+    # ─── Setup ───────────────────────────────────────────────────────────────
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Device: {device}")
+    if device.type == "cuda":
+        print(f"GPU: {torch.cuda.get_device_name(0)}")
+        print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
+
+    os.makedirs(CHECKPOINT_DIR, exist_ok=True)
+    os.makedirs(LOGS_DIR, exist_ok=True)
+
+    # ─── Tokenizer ───────────────────────────────────────────────────────────
+    print("Loading tokenizer...")
+    tokenizer = Tokenizer.from_file(TOKENIZER_PATH)
+    vocab_size = tokenizer.get_vocab_size()
+    print(f"Vocab size: {vocab_size}")
+
+    # ─── Data ────────────────────────────────────────────────────────────────
+    if args.corpus == "curated":
+        corpus_path = CURATED_CORPUS_PATH
+        if not os.path.exists(corpus_path):
+            raise FileNotFoundError(f"Curated corpus not found: {corpus_path}. Run src/curate.py first.")
+        print(f"Using CURATED corpus: {corpus_path}")
+    elif args.corpus == "default":
+        corpus_path = CORPUS_PATH
+    else:
+        corpus_path = args.corpus
+    tokens = load_and_tokenize(corpus_path, tokenizer)
+
+    # ─── Model ───────────────────────────────────────────────────────────────
+    config = ModelConfig(
+        vocab_size=vocab_size,
+        d_model=1_280,
+        n_layers=23,
+        n_heads=20,
+        d_ff=3_456,
+        max_seq_len=args.seq_len,
+        dropout=0.0,
+        tie_embeddings=True,
+    )
+
+    model = Retriever500M(config).to(device)
+    total_params = model.count_parameters()
+    print(f"Model parameters: {total_params:,} ({total_params / 1e6:.1f}M)")
+
+    # ─── Optimizer ───────────────────────────────────────────────────────────
+    optimizer = setup_optimizer(model, args.lr, use_8bit=args.use_8bit_adam)
+
+    # ─── Resume from checkpoint ──────────────────────────────────────────────
+    start_step = 0
+    best_loss = float("inf")
+    prev_log_steps = []
+    if args.resume:
+        resume_path = args.resume_path or os.path.join(CHECKPOINT_DIR, "latest.pt")
+        if not os.path.exists(resume_path):
+            raise FileNotFoundError(f"Cannot resume: {resume_path} does not exist")
+        start_step, best_loss = resume_from_checkpoint(model, optimizer, resume_path, device)
+        accum_loss = best_loss  # continue EMA from saved loss
+        # Load previous log entries so we append rather than overwrite history
+        prev_log_path = os.path.join(LOGS_DIR, "training_log.json")
+        if os.path.exists(prev_log_path):
+            try:
+                with open(prev_log_path, "r") as f:
+                    prev_log = json.load(f)
+                prev_log_steps = prev_log.get("steps", [])
+                print(f"  Loaded {len(prev_log_steps)} previous log entries")
+            except Exception as e:
+                print(f"  Could not load previous log ({e})")
+    else:
+        accum_loss = 0.0
+
+    # ─── Training loop ───────────────────────────────────────────────────────
+    effective_batch = args.batch_size * args.grad_accum
+    max_steps_total = start_step + args.steps  # for LR schedule continuity
+    print(f"\nTraining configuration:")
+    print(f"  Micro batch size:  {args.batch_size}")
+    print(f"  Gradient accum:    {args.grad_accum}")
+    print(f"  Effective batch:   {effective_batch}")
+    print(f"  Sequence length:   {args.seq_len}")
+    print(f"  Learning rate:     {args.lr}")
+    print(f"  Steps this run:    {args.steps}")
+    print(f"  Start step:        {start_step}")
+    print(f"  Target step:       {max_steps_total}")
+    print(f"  Warmup steps:      {args.warmup}")
+    print(f"  Grad checkpointing: {args.grad_checkpoint}")
+    print()
+
+    # Training log
+    log = {
+        "config": asdict(config),
+        "train_args": vars(args),
+        "total_params": total_params,
+        "steps": list(prev_log_steps),  # carry over previous entries
+    }
+
+    model.train()
+    step = start_step
+    start_time = time.time()
+
+    pbar = tqdm(range(start_step, max_steps_total), desc="Training", initial=start_step, total=max_steps_total)
+    for step in pbar:
+        # Learning rate schedule (uses absolute step for continuity)
+        lr = get_lr(step, args.warmup, max_steps_total, args.lr, args.lr * 0.1)
+        for pg in optimizer.param_groups:
+            pg["lr"] = lr
+
+        optimizer.zero_grad(set_to_none=True)
+
+        # Gradient accumulation
+        total_loss = 0.0
+        for micro_step in range(args.grad_accum):
+            input_ids, targets = get_batch(tokens, args.batch_size, args.seq_len, device)
+
+            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+                out = model(input_ids, targets=targets, use_checkpoint=args.grad_checkpoint)
+                loss = out["loss"] / args.grad_accum
+
+            loss.backward()
+            total_loss += loss.item()
+
+        # Gradient clipping
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+
+        # Optimizer step
+        optimizer.step()
+
+        avg_loss = total_loss  # already divided by grad_accum
+        accum_loss = accum_loss * 0.95 + avg_loss * 0.05  # EMA
+
+        # Logging
+        if step % args.log_every == 0 or step == max_steps_total - 1:
+            elapsed = time.time() - start_time
+            steps_this_run = step - start_step + 1
+            steps_per_sec = steps_this_run / elapsed
+            vram_used = torch.cuda.max_memory_allocated() / 1e9 if device.type == "cuda" else 0
+
+            log_entry = {
+                "step": step,
+                "loss": avg_loss,
+                "ema_loss": accum_loss,
+                "lr": lr,
+                "elapsed_s": elapsed,
+                "steps_per_sec": steps_per_sec,
+                "vram_gb": vram_used,
+            }
+            log["steps"].append(log_entry)
+
+            pbar.set_postfix({
+                "loss": f"{avg_loss:.4f}",
+                "ema": f"{accum_loss:.4f}",
+                "lr": f"{lr:.2e}",
+                "vram": f"{vram_used:.1f}G",
+            })
+
+        # Save checkpoint
+        if (step + 1) % args.save_every == 0 or step == max_steps_total - 1:
+            ckpt_path = os.path.join(CHECKPOINT_DIR, f"model_step_{step + 1}.pt")
+            torch.save({
+                "model_state_dict": model.state_dict(),
+                "optimizer_state_dict": optimizer.state_dict(),
+                "config": asdict(config),
+                "step": step + 1,
+                "loss": accum_loss,
+            }, ckpt_path)
+            print(f"\n  Saved checkpoint: {ckpt_path}")
+
+            # Also save latest
+            latest_path = os.path.join(CHECKPOINT_DIR, "latest.pt")
+            torch.save({
+                "model_state_dict": model.state_dict(),
+                "config": asdict(config),
+                "step": step + 1,
+                "loss": accum_loss,
+            }, latest_path)
+
+            if accum_loss < best_loss:
+                best_loss = accum_loss
+                best_path = os.path.join(CHECKPOINT_DIR, "best.pt")
+                torch.save({
+                    "model_state_dict": model.state_dict(),
+                    "config": asdict(config),
+                    "step": step + 1,
+                    "loss": accum_loss,
+                }, best_path)
+
+        # Reset peak memory stats periodically
+        if step % 50 == 0 and device.type == "cuda":
+            torch.cuda.reset_peak_memory_stats()
+
+    # ─── Save training log ───────────────────────────────────────────────────
+    log_path = os.path.join(LOGS_DIR, "training_log.json")
+    with open(log_path, "w") as f:
+        json.dump(log, f, indent=2)
+    print(f"\nTraining log saved to {log_path}")
+
+    total_time = time.time() - start_time
+    print(f"\nTraining complete!")
+    print(f"  Total time:     {total_time:.1f}s ({total_time/60:.1f} min)")
+    print(f"  Final EMA loss: {accum_loss:.4f}")
+    print(f"  Best loss:      {best_loss:.4f}")
+    print(f"  Steps/sec:      {args.steps / total_time:.2f}")
+
+    return model, log
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Train Retriever500M base model")
+    parser.add_argument("--steps", type=int, default=2000, help="Total training steps")
+    parser.add_argument("--batch_size", type=int, default=4, help="Micro batch size")
+    parser.add_argument("--grad_accum", type=int, default=8, help="Gradient accumulation steps")
+    parser.add_argument("--seq_len", type=int, default=512, help="Sequence length")
+    parser.add_argument("--lr", type=float, default=3e-4, help="Peak learning rate")
+    parser.add_argument("--warmup", type=int, default=100, help="Warmup steps")
+    parser.add_argument("--save_every", type=int, default=500, help="Save checkpoint every N steps")
+    parser.add_argument("--log_every", type=int, default=10, help="Log every N steps")
+    parser.add_argument("--grad_checkpoint", action="store_true", default=True, help="Use gradient checkpointing")
+    parser.add_argument("--no_grad_checkpoint", dest="grad_checkpoint", action="store_false")
+    parser.add_argument("--use_8bit_adam", action="store_true", default=True, help="Use 8-bit AdamW")
+    parser.add_argument("--no_8bit_adam", dest="use_8bit_adam", action="store_false")
+    parser.add_argument("--resume", action="store_true", help="Resume training from latest checkpoint")
+    parser.add_argument("--resume_path", type=str, default=None, help="Specific checkpoint to resume from (default: checkpoints/latest.pt)")
+    parser.add_argument("--corpus", type=str, default="default", help="Corpus to use: 'default', 'curated', or a path")
+    args = parser.parse_args()
+
+    train(args)
+
+
+if __name__ == "__main__":
+    main()

src/train_tokenizer.py

@@ -0,0 +1,60 @@
+"""
+Train a BPE tokenizer on the code corpus using the HuggingFace `tokenizers` library.
+
+Produces a 32,000-token vocabulary optimized for source code across
+Python, JS/TS, Rust, Go, C/C++, and other languages.
+"""
+
+import os
+from tokenizers import Tokenizer
+from tokenizers.models import BPE
+from tokenizers.trainers import BpeTrainer
+from tokenizers.pre_tokenizers import ByteLevel
+from tokenizers.decoders import ByteLevel as ByteLevelDecoder
+
+DATA_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data")
+TOKENIZER_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "tokenizer")
+CORPUS_PATH = os.path.join(DATA_DIR, "corpus.txt")
+TOKENIZER_PATH = os.path.join(TOKENIZER_DIR, "tokenizer.json")
+
+VOCAB_SIZE = 32_000
+
+
+def train_tokenizer():
+    os.makedirs(TOKENIZER_DIR, exist_ok=True)
+
+    tokenizer = Tokenizer(BPE(unk_token="<unk>"))
+    tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=True, use_regex=True)
+    tokenizer.decoder = ByteLevelDecoder()
+
+    trainer = BpeTrainer(
+        vocab_size=VOCAB_SIZE,
+        special_tokens=["<pad>", "<bos>", "<eos>", "<unk>"],
+        show_progress=True,
+        initial_alphabet=ByteLevel.alphabet(),
+    )
+
+    print(f"Training BPE tokenizer (vocab_size={VOCAB_SIZE}) on {CORPUS_PATH}...")
+    tokenizer.train([CORPUS_PATH], trainer)
+
+    tokenizer.save(TOKENIZER_PATH)
+    print(f"Tokenizer saved to {TOKENIZER_PATH}")
+
+    # Print stats
+    vocab = tokenizer.get_vocab()
+    print(f"Vocabulary size: {len(vocab)}")
+
+    # Test encoding
+    test_code = "def hello_world():\n    print('Hello, World!')"
+    encoded = tokenizer.encode(test_code)
+    print(f"\nTest encoding:")
+    print(f"  Input:    {test_code}")
+    print(f"  Tokens:   {encoded.tokens[:20]}")
+    print(f"  IDs:      {encoded.ids[:20]}")
+    print(f"  # tokens: {len(encoded.ids)}")
+
+    return tokenizer
+
+
+if __name__ == "__main__":
+    train_tokenizer()

src/verify_pipeline.py

@@ -0,0 +1,77 @@
+"""Verify the full search agent pipeline end-to-end."""
+import json
+import os
+import sys
+
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__))))
+
+def main():
+    # 1. Verify chunks load
+    chunks = [json.loads(l) for l in open("data/chunks.jsonl", encoding="utf-8")]
+    print(f"1. Chunks: {len(chunks):,} loaded OK")
+
+    # 2. Verify agent tokenizer
+    from tokenizers import Tokenizer
+    tok = Tokenizer.from_file("tokenizer/tokenizer_agent.json")
+    print(f"2. Agent tokenizer: vocab={tok.get_vocab_size()} OK")
+
+    # 3. Verify special tokens
+    st = json.load(open("tokenizer/special_tokens.json"))
+    ids = list(st["token_ids"].values())
+    print(f"3. Special tokens: {len(st['special_tokens'])} tokens, IDs {ids} OK")
+
+    # 4. Verify SFT traces load and format correctly
+    from sft import format_trace_to_tokens, load_special_tokens
+    load_special_tokens()
+    traces = [json.loads(l) for l in open("data/sft_traces.jsonl", encoding="utf-8")]
+    print(f"4. SFT traces: {len(traces):,} loaded")
+
+    # Format one trace
+    ids_arr, mask_arr = format_trace_to_tokens(traces[0], tok, max_seq_len=768)
+    loss_tokens = int(mask_arr.sum())
+    total_tokens = len(ids_arr)
+    pct = loss_tokens / total_tokens * 100
+    print(f"   Sample trace: {total_tokens} tokens, {loss_tokens} with loss ({pct:.0f}%)")
+
+    # 5. Verify gold traces
+    gold = [json.loads(l) for l in open("data/gold_traces.jsonl", encoding="utf-8")]
+    print(f"5. Gold traces: {len(gold)} loaded OK")
+
+    # 6. Verify model accepts new vocab size
+    from model import ModelConfig, Retriever500M
+    import torch
+    config = ModelConfig(vocab_size=32009, d_model=1280, n_layers=23, n_heads=20, d_ff=3456, max_seq_len=768)
+    model = Retriever500M(config)
+    print(f"6. Model with vocab=32009: {model.count_parameters()/1e6:.1f}M params OK")
+
+    # 7. Verify model can load old checkpoint with embedding resize
+    ckpt = torch.load("checkpoints/latest.pt", map_location="cpu", weights_only=False)
+    old_vocab = ckpt["config"]["vocab_size"]
+    needs_resize = old_vocab != 32009
+    print(f"7. Old checkpoint vocab={old_vocab}, new vocab=32009, resize needed={needs_resize}")
+
+    # 8. Verify retriever
+    from search_agent import KeywordRetriever
+    retr = KeywordRetriever("data/chunks.jsonl")
+    results = retr.search("ngx_reusable_connection", top_k=3)
+    print(f"8. Retriever: search returned {len(results)} results OK")
+    if results:
+        print(f"   Top result: {results[0]['name']} (score={results[0]['score']:.2f})")
+
+    # 9. Verify SFT data formatting produces valid tokens
+    from sft import get_sft_batch
+    dataset = []
+    for trace in traces[:100]:
+        ids, mask = format_trace_to_tokens(trace, tok, max_seq_len=768)
+        if len(ids) > 10:
+            dataset.append((ids, mask))
+    input_ids, targets, loss_mask = get_sft_batch(dataset, batch_size=2, seq_len=768, device=torch.device("cpu"))
+    print(f"9. SFT batch: input_ids={input_ids.shape}, targets={targets.shape}, loss_mask={loss_mask.shape}")
+    print(f"   Loss mask coverage: {loss_mask.sum().item()}/{loss_mask.numel()} ({loss_mask.float().mean()*100:.0f}%)")
+
+    print()
+    print("ALL CHECKS PASSED")
+
+
+if __name__ == "__main__":
+    main()