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/chunk_documents.py

@@ -0,0 +1,225 @@
+"""
+Chunk quality-filtered documents into:
+  1. A pre-training corpus (data/pretrain/corpus_pretrain.txt)
+     - Concatenated source code with document separators
+     - Used for base LM pre-training
+  2. Retrieval chunks (data/chunks/chunks.jsonl)
+     - Semantically meaningful code units (functions, classes, structs)
+     - Used by the search agent and for SFT trace generation
+
+Input:  data/quality/documents_quality.jsonl
+Output: data/pretrain/corpus_pretrain.txt
+        data/chunks/chunks.jsonl
+"""
+
+import json
+import os
+import re
+
+PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+DATA_DIR = os.path.join(PROJECT_DIR, "data")
+QUALITY_DOCS_PATH = os.path.join(DATA_DIR, "quality", "documents_quality.jsonl")
+PRETRAIN_DIR = os.path.join(DATA_DIR, "pretrain")
+CORPUS_PRETRAIN_PATH = os.path.join(PRETRAIN_DIR, "corpus_pretrain.txt")
+CHUNKS_DIR = os.path.join(DATA_DIR, "chunks")
+CHUNKS_PATH = os.path.join(CHUNKS_DIR, "chunks.jsonl")
+
+# ─── Chunk start patterns per language ───────────────────────────────────────
+CHUNK_STARTS = {
+    "python":   re.compile(r"^(def |class |@|async def )", re.M),
+    "javascript": re.compile(r"^(function |const |let |class |export |async function|interface |type \w+ =)", re.M),
+    "typescript": re.compile(r"^(function |const |let |class |export |async function|interface |type \w+ =|enum )", re.M),
+    "rust":     re.compile(r"^(fn |pub fn |impl |struct |enum |trait |mod |macro_rules!|pub struct|pub enum|pub trait)", re.M),
+    "go":       re.compile(r"^(func |type \w+ struct|type \w+ interface)", re.M),
+    "c":        re.compile(r"^(static |void |int |char |size_t |typedef |struct \w+ \{|#define|#if|#ifndef|ngx_)", re.M),
+    "cpp":      re.compile(r"^(static |void |int |char |size_t |template |class \w+|namespace |struct \w+|#define|#if|#ifndef)", re.M),
+    "java":     re.compile(r"^(public |private |protected |class \w+|interface |enum )", re.M),
+    "csharp":   re.compile(r"^(public |private |protected |internal |class \w+|interface |enum )", re.M),
+    "ruby":     re.compile(r"^(def |class |module )", re.M),
+}
+
+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"type (\w+) interface"), "interface"),
+    (re.compile(r"function (\w+)"), "function"),
+    (re.compile(r"async function (\w+)"), "function"),
+    (re.compile(r"typedef struct (\w+)"), "typedef"),
+    (re.compile(r"#define (\w+)"), "macro"),
+    (re.compile(r"interface (\w+)"), "interface"),
+    (re.compile(r"module (\w+)"), "module"),
+]
+
+
+def extract_name(code: str) -> tuple[str, str]:
+    for pat, typ in NAME_PATTERNS:
+        m = pat.search(code)
+        if m:
+            return m.group(1), typ
+    return "unknown", "block"
+
+
+def chunk_document(doc: dict) -> list[dict]:
+    """Split a document into chunks based on language-specific patterns."""
+    content = doc["content"]
+    lang = doc["language"]
+    filepath = doc["filepath"]
+    repo = doc["repo"]
+    lines = content.split("\n")
+    n_lines = len(lines)
+
+    start_pat = CHUNK_STARTS.get(lang)
+    if start_pat is None:
+        name, typ = extract_name(content)
+        return [{
+            "id": f"{repo}:{filepath}:0",
+            "language": lang,
+            "name": name,
+            "type": typ,
+            "code": content,
+            "filepath": filepath,
+            "repo": repo,
+            "start_line": 1,
+            "end_line": n_lines,
+        }]
+
+    starts = [(m.start(), m.group()) for m in start_pat.finditer(content)]
+
+    if not starts:
+        name, typ = extract_name(content)
+        return [{
+            "id": f"{repo}:{filepath}:0",
+            "language": lang,
+            "name": name,
+            "type": typ,
+            "code": content,
+            "filepath": filepath,
+            "repo": repo,
+            "start_line": 1,
+            "end_line": n_lines,
+        }]
+
+    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(content)
+        chunk_code = content[start_pos:end_pos].strip()
+
+        if len(chunk_code) < 40:
+            continue
+
+        # Split huge chunks on blank lines
+        if len(chunk_code) > 6000:
+            sub_parts = re.split(r"\n\n+", chunk_code)
+            for j, part in enumerate(sub_parts):
+                part = part.strip()
+                if len(part) < 40:
+                    continue
+                name, typ = extract_name(part)
+                start_line = content[:start_pos].count("\n") + 1 + sum(p.count("\n") + 2 for p in sub_parts[:j])
+                chunks.append({
+                    "id": f"{repo}:{filepath}:{i}_{j}",
+                    "language": lang,
+                    "name": name,
+                    "type": typ,
+                    "code": part,
+                    "filepath": filepath,
+                    "repo": repo,
+                    "start_line": start_line,
+                    "end_line": start_line + part.count("\n"),
+                })
+            continue
+
+        name, typ = extract_name(chunk_code)
+        start_line = content[:start_pos].count("\n") + 1
+        chunks.append({
+            "id": f"{repo}:{filepath}:{i}",
+            "language": lang,
+            "name": name,
+            "type": typ,
+            "code": chunk_code,
+            "filepath": filepath,
+            "repo": repo,
+            "start_line": start_line,
+            "end_line": start_line + chunk_code.count("\n"),
+        })
+
+    return chunks
+
+
+def main():
+    os.makedirs(PRETRAIN_DIR, exist_ok=True)
+    os.makedirs(CHUNKS_DIR, exist_ok=True)
+
+    print(f"Loading quality documents from {QUALITY_DOCS_PATH}...")
+    docs = []
+    with open(QUALITY_DOCS_PATH, "r", encoding="utf-8") as f:
+        for line in f:
+            docs.append(json.loads(line))
+    print(f"  Loaded {len(docs):,} documents")
+    print(f"  Total size: {sum(len(d['content']) for d in docs) / 1e6:.1f} MB")
+
+    # ─── Build pre-training corpus ─────────────────────────────────────────
+    print("\nBuilding pre-training corpus...")
+    corpus_size = 0
+    with open(CORPUS_PRETRAIN_PATH, "w", encoding="utf-8") as f:
+        for doc in docs:
+            f.write(doc["content"].strip())
+            f.write("\n\n\n")  # document separator
+            corpus_size += len(doc["content"])
+    print(f"  Pre-training corpus: {corpus_size / 1e6:.1f} MB -> {CORPUS_PRETRAIN_PATH}")
+
+    # ─── Build retrieval chunks ────────────────────────────────────────────
+    print("\nChunking documents...")
+    all_chunks = []
+    lang_counts = {}
+    type_counts = {}
+
+    for doc in docs:
+        chunks = chunk_document(doc)
+        for chunk in chunks:
+            all_chunks.append(chunk)
+            lang_counts[chunk["language"]] = lang_counts.get(chunk["language"], 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}")
+
+    sizes = [len(c["code"]) for c in all_chunks]
+    if sizes:
+        sizes.sort()
+        print(f"  Chunk size: min={sizes[0]}, median={sizes[len(sizes)//2]}, "
+              f"max={sizes[-1]}, mean={sum(sizes)//len(sizes)}")
+
+    with open(CHUNKS_PATH, "w", encoding="utf-8") as f:
+        for chunk in all_chunks:
+            f.write(json.dumps(chunk, ensure_ascii=False) + "\n")
+    print(f"\nChunks written to {CHUNKS_PATH}")
+
+    # Stats
+    stats = {
+        "input_docs": len(docs),
+        "pretrain_corpus_mb": corpus_size / 1e6,
+        "total_chunks": len(all_chunks),
+        "language_distribution": lang_counts,
+        "type_distribution": type_counts,
+    }
+    stats_path = os.path.join(CHUNKS_DIR, "chunk_stats.json")
+    with open(stats_path, "w") as f:
+        json.dump(stats, f, indent=2)
+    print(f"Stats written to {stats_path}")
+
+
+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/dedup_pipeline.py

@@ -0,0 +1,326 @@
+"""
+Deduplication pipeline for the code corpus.
+
+Three stages of deduplication, applied in order:
+
+  1. EXACT DEDUP (document-level)
+     - SHA-256 hash of normalized content (whitespace-collapsed, lowercased)
+     - Drops byte-identical and whitespace-only-differing documents.
+
+  2. NEAR-DEDUP (document-level, MinHash + LSH)
+     - Shingles each document into character n-grams.
+     - Computes MinHash signatures.
+     - Uses LSH banding to cluster near-duplicates (Jaccard >= threshold).
+     - Keeps one representative per cluster (the longest, highest-quality).
+
+  3. LINE-LEVEL DEDUP (cross-document)
+     - Builds a global set of exact line hashes.
+     - Removes lines that appear in >threshold documents (boilerplate:
+       license headers, auto-generated stubs, repeated imports, etc.).
+     - Drops documents that become too short / too sparse after line removal.
+
+Input:  data/raw_large/documents.jsonl
+Output: data/dedup/documents_dedup.jsonl
+
+Reports statistics at each stage so we can verify the pipeline is working.
+"""
+
+import hashlib
+import json
+import os
+import re
+import sys
+from collections import Counter, defaultdict
+
+from datasketch import MinHash, MinHashLSH
+
+PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+DATA_DIR = os.path.join(PROJECT_DIR, "data")
+RAW_DOCS_PATH = os.path.join(DATA_DIR, "raw_large", "documents.jsonl")
+DEDUP_DIR = os.path.join(DATA_DIR, "dedup")
+DEDUP_DOCS_PATH = os.path.join(DEDUP_DIR, "documents_dedup.jsonl")
+STATS_PATH = os.path.join(DEDUP_DIR, "dedup_stats.json")
+
+# ─── MinHash / LSH parameters ────────────────────────────────────────────────
+SHINGLE_SIZE = 5          # character 5-grams
+NUM_PERM = 128            # MinHash permutations (higher = more accurate)
+LSH_THRESHOLD = 0.5       # Jaccard similarity threshold for near-dup
+LSH_BANDS = 32            # LSH bands (NUM_PERM / LSH_ROWS should be integer)
+LSH_ROWS = NUM_PERM // LSH_BANDS  # 128 / 32 = 4
+
+# ─── Line-level dedup parameters ────────────────────────────────────────────
+LINE_BOILERPLATE_THRESHOLD = 50  # lines appearing in >= this many docs are boilerplate
+MIN_LINES_AFTER_DEDUP = 5        # drop docs with fewer lines after line dedup
+MIN_CHARS_AFTER_DEDUP = 200      # drop docs shorter than this after line dedup
+
+
+def normalize_content(text: str) -> str:
+    """Normalize content for exact-hash comparison: collapse whitespace, lowercase."""
+    # Collapse all whitespace sequences to single space
+    text = re.sub(r"\s+", " ", text)
+    return text.strip().lower()
+
+
+def exact_hash(text: str) -> str:
+    """SHA-256 of normalized content."""
+    return hashlib.sha256(normalize_content(text).encode("utf-8")).hexdigest()
+
+
+def shingles(text: str, k: int = SHINGLE_SIZE) -> set[str]:
+    """Character k-gram shingles of normalized text."""
+    text = re.sub(r"\s+", " ", text).strip().lower()
+    if len(text) < k:
+        return {text}
+    return {text[i : i + k] for i in range(len(text) - k + 1)}
+
+
+def minhash_signature(text: str, num_perm: int = NUM_PERM) -> MinHash:
+    """Compute MinHash signature for a document."""
+    m = MinHash(num_perm=num_perm)
+    for s in shingles(text):
+        m.update(s.encode("utf-8"))
+    return m
+
+
+# ─── Stage 1: Exact dedup ────────────────────────────────────────────────────
+
+def stage1_exact_dedup(docs: list[dict]) -> tuple[list[dict], dict]:
+    """Remove exact (and whitespace-only) duplicates. Keep first occurrence."""
+    print("\n" + "=" * 60)
+    print("STAGE 1: Exact document-level dedup (SHA-256 of normalized content)")
+    print("=" * 60)
+
+    seen_hashes = {}
+    kept = []
+    dropped = 0
+
+    for i, doc in enumerate(docs):
+        h = exact_hash(doc["content"])
+        if h in seen_hashes:
+            dropped += 1
+            continue
+        seen_hashes[h] = i
+        kept.append(doc)
+
+        if (i + 1) % 5000 == 0:
+            print(f"  Processed {i+1}/{len(docs)} | kept {len(kept)} | dropped {dropped}")
+
+    stats = {
+        "input": len(docs),
+        "kept": len(kept),
+        "dropped_exact": dropped,
+        "drop_rate": dropped / max(len(docs), 1),
+    }
+    print(f"  Input: {len(docs)} | Kept: {len(kept)} | Dropped: {dropped} "
+          f"({stats['drop_rate']:.1%})")
+    return kept, stats
+
+
+# ─── Stage 2: Near-dedup (MinHash LSH) ───────────────────────────────────────
+
+def stage2_near_dedup(docs: list[dict]) -> tuple[list[dict], dict]:
+    """Cluster near-duplicates via MinHash LSH, keep best representative per cluster."""
+    print("\n" + "=" * 60)
+    print(f"STAGE 2: Near-dedup (MinHash LSH, threshold={LSH_THRESHOLD}, "
+          f"perm={NUM_PERM}, bands={LSH_BANDS})")
+    print("=" * 60)
+
+    # Build LSH index
+    lsh = MinHashLSH(
+        threshold=LSH_THRESHOLD,
+        num_perm=NUM_PERM,
+        params=(LSH_BANDS, LSH_ROWS),
+    )
+
+    signatures = {}
+    kept_indices = set()
+
+    for i, doc in enumerate(docs):
+        content = doc["content"]
+        if len(content) < 50:
+            # Too short for meaningful MinHash — keep as-is
+            kept_indices.add(i)
+            continue
+
+        mh = minhash_signature(content)
+        signatures[i] = mh
+
+        # Query for existing near-duplicates
+        candidates = lsh.query(mh)
+
+        if candidates:
+            # This doc is a near-duplicate of an existing one.
+            # Compare lengths and keep the longer one.
+            existing_best = max(candidates, key=lambda j: len(docs[int(j)]["content"]))
+            if len(content) > len(docs[int(existing_best)]["content"]):
+                # New doc is longer — replace the existing representative
+                kept_indices.discard(int(existing_best))
+                kept_indices.add(i)
+                # We can't easily remove from LSH, so just add the new one.
+                # The old one is dropped via kept_indices tracking.
+                try:
+                    lsh.insert(str(i), mh)
+                except Exception:
+                    pass
+            # else: keep existing, drop this one
+        else:
+            # No near-duplicate found — keep this doc
+            kept_indices.add(i)
+            try:
+                lsh.insert(str(i), mh)
+            except Exception:
+                pass
+
+        if (i + 1) % 5000 == 0:
+            print(f"  Processed {i+1}/{len(docs)} | clusters so far: {len(kept_indices)}")
+
+    kept = [docs[i] for i in sorted(kept_indices)]
+    dropped = len(docs) - len(kept)
+
+    stats = {
+        "input": len(docs),
+        "kept": len(kept),
+        "dropped_near": dropped,
+        "drop_rate": dropped / max(len(docs), 1),
+        "threshold": LSH_THRESHOLD,
+        "num_perm": NUM_PERM,
+    }
+    print(f"  Input: {len(docs)} | Kept: {len(kept)} | Dropped: {dropped} "
+          f"({stats['drop_rate']:.1%})")
+    return kept, stats
+
+
+# ─── Stage 3: Line-level dedup ───────────────────────────────────────────────
+
+def stage3_line_dedup(docs: list[dict]) -> tuple[list[dict], dict]:
+    """Remove boilerplate lines that appear across many documents."""
+    print("\n" + "=" * 60)
+    print(f"STAGE 3: Line-level dedup (boilerplate threshold={LINE_BOILERPLATE_THRESHOLD})")
+    print("=" * 60)
+
+    # First pass: count how many documents each normalized line appears in
+    line_doc_count = Counter()
+    for doc in docs:
+        lines = doc["content"].split("\n")
+        # Use normalized line hash to save memory
+        seen_in_doc = set()
+        for line in lines:
+            norm = line.strip()
+            if not norm:
+                continue
+            # Normalize: collapse internal whitespace, lowercase
+            norm = re.sub(r"\s+", " ", norm).lower()
+            if len(norm) < 5:  # skip very short lines (braces, etc.)
+                continue
+            h = hashlib.md5(norm.encode("utf-8")).hexdigest()
+            seen_in_doc.add(h)
+        for h in seen_in_doc:
+            line_doc_count[h] += 1
+
+    # Identify boilerplate lines (appear in >= threshold documents)
+    boilerplate_hashes = {h for h, c in line_doc_count.items() if c >= LINE_BOILERPLATE_THRESHOLD}
+    print(f"  Boilerplate lines identified: {len(boilerplate_hashes)}")
+
+    # Second pass: remove boilerplate lines from each document
+    kept = []
+    dropped_too_short = 0
+    lines_removed_total = 0
+
+    for doc in docs:
+        lines = doc["content"].split("\n")
+        new_lines = []
+        for line in lines:
+            norm = line.strip()
+            if norm:
+                norm_h = re.sub(r"\s+", " ", norm).lower()
+                if len(norm_h) >= 5:
+                    h = hashlib.md5(norm_h.encode("utf-8")).hexdigest()
+                    if h in boilerplate_hashes:
+                        lines_removed_total += 1
+                        continue
+            new_lines.append(line)
+
+        new_content = "\n".join(new_lines).strip()
+        new_line_count = len([l for l in new_lines if l.strip()])
+
+        if new_line_count < MIN_LINES_AFTER_DEDUP or len(new_content) < MIN_CHARS_AFTER_DEDUP:
+            dropped_too_short += 1
+            continue
+
+        doc = dict(doc)
+        doc["content"] = new_content
+        kept.append(doc)
+
+    stats = {
+        "input": len(docs),
+        "kept": len(kept),
+        "dropped_too_short": dropped_too_short,
+        "lines_removed": lines_removed_total,
+        "boilerplate_hashes": len(boilerplate_hashes),
+    }
+    print(f"  Input: {len(docs)} | Kept: {len(kept)} | "
+          f"Dropped (too short): {dropped_too_short} | "
+          f"Lines removed: {lines_removed_total}")
+    return kept, stats
+
+
+# ─── Main ────────────────────────────────────────────────────────────────────
+
+def main():
+    os.makedirs(DEDUP_DIR, exist_ok=True)
+
+    print(f"Loading documents from {RAW_DOCS_PATH}...")
+    docs = []
+    with open(RAW_DOCS_PATH, "r", encoding="utf-8") as f:
+        for line in f:
+            docs.append(json.loads(line))
+    print(f"  Loaded {len(docs):,} documents")
+    print(f"  Total size: {sum(len(d['content']) for d in docs) / 1e6:.1f} MB")
+
+    all_stats = {"input_count": len(docs)}
+
+    # Stage 1: Exact dedup
+    docs, s1 = stage1_exact_dedup(docs)
+    all_stats["stage1_exact"] = s1
+
+    # Stage 2: Near-dedup
+    docs, s2 = stage2_near_dedup(docs)
+    all_stats["stage2_near"] = s2
+
+    # Stage 3: Line-level dedup
+    docs, s3 = stage3_line_dedup(docs)
+    all_stats["stage3_line"] = s3
+
+    # Final stats
+    final_size = sum(len(d["content"]) for d in docs)
+    all_stats["final_count"] = len(docs)
+    all_stats["final_size_mb"] = final_size / 1e6
+
+    # Language distribution
+    lang_counts = {}
+    for d in docs:
+        lang_counts[d["language"]] = lang_counts.get(d["language"], 0) + 1
+    all_stats["language_distribution"] = lang_counts
+
+    print("\n" + "=" * 60)
+    print("DEDUP COMPLETE")
+    print("=" * 60)
+    print(f"  Input:   {all_stats['input_count']:,} documents")
+    print(f"  Final:   {len(docs):,} documents")
+    print(f"  Size:    {final_size / 1e6:.1f} MB")
+    print(f"  Languages: {lang_counts}")
+
+    # Write deduped documents
+    with open(DEDUP_DOCS_PATH, "w", encoding="utf-8") as f:
+        for doc in docs:
+            f.write(json.dumps(doc, ensure_ascii=False) + "\n")
+    print(f"\nDeduped documents written to {DEDUP_DOCS_PATH}")
+
+    # Write stats
+    with open(STATS_PATH, "w") as f:
+        json.dump(all_stats, f, indent=2)
+    print(f"Stats written to {STATS_PATH}")
+
+
+if __name__ == "__main__":
+    main()

src/dedup_quality.py

@@ -0,0 +1,335 @@
+"""
+Deduplication, contamination check, and quality filter for SFT traces.
+
+Pipeline:
+  1. Load all traces from data/sft_traces_v2/ + existing sft_traces.jsonl
+  2. Deduplicate by query hash (exact match) and by fuzzy similarity (near-dupes)
+  3. Contamination check: remove traces whose queries appear in gold_traces.jsonl
+  4. Quality filter: remove traces that are too short, have empty reasoning,
+     or have malformed structure
+  5. Shuffle and write final dataset
+
+Usage:
+  python src/dedup_quality.py --input data/sft_traces_v2/ --output data/sft_traces_final.jsonl
+"""
+
+import argparse
+import hashlib
+import json
+import os
+import random
+import re
+import sys
+from collections import defaultdict
+from difflib import SequenceMatcher
+
+PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+def load_traces_from_dir(dir_path: str) -> list[dict]:
+    """Load all .jsonl files from a directory."""
+    traces = []
+    if not os.path.exists(dir_path):
+        return traces
+    for fname in sorted(os.listdir(dir_path)):
+        if fname.endswith('.jsonl'):
+            fpath = os.path.join(dir_path, fname)
+            with open(fpath, 'r', encoding='utf-8') as f:
+                for line in f:
+                    try:
+                        trace = json.loads(line)
+                        if trace and 'query' in trace and 'trace' in trace:
+                            trace['_source'] = fname
+                            traces.append(trace)
+                    except json.JSONDecodeError:
+                        continue
+    return traces
+
+def load_traces_from_file(fpath: str) -> list[dict]:
+    """Load traces from a single .jsonl file."""
+    traces = []
+    if not os.path.exists(fpath):
+        return traces
+    with open(fpath, 'r', encoding='utf-8') as f:
+        for line in f:
+            try:
+                trace = json.loads(line)
+                if trace and 'query' in trace and 'trace' in trace:
+                    trace['_source'] = os.path.basename(fpath)
+                    traces.append(trace)
+            except json.JSONDecodeError:
+                continue
+    return traces
+
+def query_hash(trace: dict) -> str:
+    """Hash the query for exact dedup."""
+    return hashlib.md5(trace['query'].strip().lower().encode()).hexdigest()
+
+def normalize_query(query: str) -> str:
+    """Normalize query for fuzzy matching."""
+    # Remove punctuation, lowercase, collapse whitespace
+    q = re.sub(r'[^\w\s]', '', query.lower())
+    q = ' '.join(q.split())
+    return q
+
+def fuzzy_similarity(q1: str, q2: str) -> float:
+    """Compute similarity between two queries."""
+    n1, n2 = normalize_query(q1), normalize_query(q2)
+    if n1 == n2:
+        return 1.0
+    # Quick length check
+    if abs(len(n1) - len(n2)) > max(len(n1), len(n2)) * 0.5:
+        return 0.0
+    return SequenceMatcher(None, n1, n2).ratio()
+
+def check_trace_quality(trace: dict) -> tuple[bool, str]:
+    """Check if a trace meets quality standards.
+    
+    Returns (is_valid, reason_if_rejected)
+    """
+    query = trace.get('query', '')
+    trace_msgs = trace.get('trace', [])
+    
+    # Must have system, user, and at least 2 assistant turns
+    if len(trace_msgs) < 4:
+        return False, "too_few_messages"
+    
+    # Query must be substantial
+    if len(query) < 15:
+        return False, "query_too_short"
+    
+    # Check assistant turns have real content
+    assistant_turns = [m for m in trace_msgs if m['role'] == 'assistant']
+    if len(assistant_turns) < 2:
+        return False, "too_few_assistant_turns"
+    
+    for turn in assistant_turns:
+        content = turn.get('content', '')
+        if len(content) < 50:
+            return False, "assistant_turn_too_short"
+        # Must contain at least one special token
+        if not any(tok in content for tok in ['<|reasoning|>', '<|search|>', '<|evidence|>', '<|finish|>']):
+            return False, "no_special_tokens"
+    
+    # Must have at least one search
+    has_search = any('<|search|>' in m.get('content', '') for m in assistant_turns)
+    if not has_search:
+        return False, "no_search_action"
+    
+    # Must have evidence or finish
+    has_evidence = any('<|evidence|>' in m.get('content', '') for m in assistant_turns)
+    has_finish = any('<|finish|>' in m.get('content', '') for m in assistant_turns)
+    if not has_evidence and not has_finish:
+        return False, "no_evidence_or_finish"
+    
+    # Check for reasoning density — at least one turn should have substantial reasoning
+    max_reasoning_len = 0
+    for turn in assistant_turns:
+        content = turn.get('content', '')
+        # Extract reasoning sections
+        reasoning_sections = re.findall(r'<\|reasoning\|>(.*?)<\|end\|>', content, re.DOTALL)
+        for r in reasoning_sections:
+            max_reasoning_len = max(max_reasoning_len, len(r.strip()))
+    
+    if max_reasoning_len < 30:
+        return False, "reasoning_too_thin"
+    
+    return True, "ok"
+
+def deduplicate(traces: list[dict], similarity_threshold: float = 0.85) -> tuple[list[dict], dict]:
+    """Remove duplicate and near-duplicate traces.
+    
+    Returns (deduplicated_traces, stats)
+    """
+    stats = {
+        'exact_dups_removed': 0,
+        'fuzzy_dups_removed': 0,
+        'total_input': len(traces),
+    }
+    
+    # Phase 1: Exact dedup by query hash
+    seen_hashes = set()
+    exact_deduped = []
+    for trace in traces:
+        h = query_hash(trace)
+        if h not in seen_hashes:
+            seen_hashes.add(h)
+            exact_deduped.append(trace)
+        else:
+            stats['exact_dups_removed'] += 1
+    
+    # Phase 2: Fuzzy dedup by query similarity
+    # Group by first word for efficiency
+    groups = defaultdict(list)
+    for trace in exact_deduped:
+        first_word = normalize_query(trace['query']).split()[0] if normalize_query(trace['query']).split() else ''
+        groups[first_word].append(trace)
+    
+    fuzzy_deduped = []
+    for first_word, group in groups.items():
+        if len(group) == 1:
+            fuzzy_deduped.extend(group)
+            continue
+        
+        # Compare within group
+        kept = []
+        for trace in group:
+            is_dup = False
+            for kept_trace in kept:
+                sim = fuzzy_similarity(trace['query'], kept_trace['query'])
+                if sim >= similarity_threshold:
+                    is_dup = True
+                    stats['fuzzy_dups_removed'] += 1
+                    break
+            if not is_dup:
+                kept.append(trace)
+        fuzzy_deduped.extend(kept)
+    
+    stats['total_output'] = len(fuzzy_deduped)
+    return fuzzy_deduped, stats
+
+def check_contamination(traces: list[dict], gold_traces: list[dict]) -> tuple[list[dict], dict]:
+    """Remove traces whose queries match gold trace queries.
+    
+    Returns (clean_traces, stats)
+    """
+    gold_queries = set()
+    for gt in gold_traces:
+        gold_queries.add(normalize_query(gt['query']))
+    
+    clean = []
+    removed = 0
+    for trace in traces:
+        nq = normalize_query(trace['query'])
+        if nq in gold_queries:
+            removed += 1
+        else:
+            clean.append(trace)
+    
+    return clean, {'contamination_removed': removed, 'gold_queries': len(gold_queries)}
+
+def main():
+    parser = argparse.ArgumentParser(description="Dedup and quality filter SFT traces")
+    parser.add_argument("--input", type=str, default=os.path.join(PROJECT_DIR, "data", "sft_traces_v2"),
+                       help="Input directory with .jsonl files")
+    parser.add_argument("--existing", type=str, default=os.path.join(PROJECT_DIR, "data", "sft_traces.jsonl"),
+                       help="Existing traces to merge with")
+    parser.add_argument("--gold", type=str, default=os.path.join(PROJECT_DIR, "data", "gold_traces.jsonl"),
+                       help="Gold traces for contamination check")
+    parser.add_argument("--output", type=str, default=os.path.join(PROJECT_DIR, "data", "sft_traces_final.jsonl"),
+                       help="Output file")
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--similarity", type=float, default=0.85,
+                       help="Fuzzy dedup similarity threshold")
+    args = parser.parse_args()
+
+    print("=" * 60)
+    print("SFT TRACE DEDUPLICATION & QUALITY PIPELINE")
+    print("=" * 60)
+
+    # 1. Load all traces
+    print("\n1. Loading traces...")
+    new_traces = load_traces_from_dir(args.input)
+    existing_traces = load_traces_from_file(args.existing)
+    gold_traces = load_traces_from_file(args.gold)
+    
+    print(f"   New traces:      {len(new_traces):,}")
+    print(f"   Existing traces: {len(existing_traces):,}")
+    print(f"   Gold traces:     {len(gold_traces):,}")
+    
+    all_traces = new_traces + existing_traces
+    print(f"   Total to process: {len(all_traces):,}")
+
+    # 2. Quality filter
+    print("\n2. Quality filtering...")
+    quality_stats = defaultdict(int)
+    quality_passed = []
+    for trace in all_traces:
+        is_valid, reason = check_trace_quality(trace)
+        if is_valid:
+            quality_passed.append(trace)
+        else:
+            quality_stats[reason] += 1
+    
+    print(f"   Passed: {len(quality_passed):,}")
+    print(f"   Rejected: {sum(quality_stats.values()):,}")
+    for reason, count in sorted(quality_stats.items(), key=lambda x: -x[1]):
+        print(f"     {reason}: {count}")
+
+    # 3. Contamination check
+    print("\n3. Contamination check (vs gold traces)...")
+    clean_traces, contam_stats = check_contamination(quality_passed, gold_traces)
+    print(f"   Removed: {contam_stats['contamination_removed']}")
+    print(f"   Remaining: {len(clean_traces):,}")
+
+    # 4. Deduplication
+    print("\n4. Deduplication...")
+    deduped_traces, dedup_stats = deduplicate(clean_traces, args.similarity)
+    print(f"   Exact dups removed: {dedup_stats['exact_dups_removed']:,}")
+    print(f"   Fuzzy dups removed: {dedup_stats['fuzzy_dups_removed']:,}")
+    print(f"   Final count: {dedup_stats['total_output']:,}")
+
+    # 5. Shuffle and write
+    print(f"\n5. Writing to {args.output}...")
+    rng = random.Random(args.seed)
+    rng.shuffle(deduped_traces)
+    
+    # Remove internal fields
+    for trace in deduped_traces:
+        trace.pop('_source', None)
+    
+    with open(args.output, 'w', encoding='utf-8') as f:
+        for trace in deduped_traces:
+            f.write(json.dumps(trace, ensure_ascii=False) + '\n')
+    
+    print(f"   Written: {len(deduped_traces):,} traces")
+    
+    # Summary
+    print("\n" + "=" * 60)
+    print("SUMMARY")
+    print("=" * 60)
+    print(f"  Input traces:        {len(all_traces):,}")
+    print(f"  Quality rejected:    {sum(quality_stats.values()):,}")
+    print(f"  Contamination removed: {contam_stats['contamination_removed']}")
+    print(f"  Exact dups removed:  {dedup_stats['exact_dups_removed']:,}")
+    print(f"  Fuzzy dups removed:  {dedup_stats['fuzzy_dups_removed']:,}")
+    print(f"  Final dataset:       {len(deduped_traces):,}")
+    
+    # Category distribution
+    cat_counts = defaultdict(int)
+    for trace in deduped_traces:
+        # Try to infer category from query pattern
+        q = trace['query'].lower()
+        if any(w in q for w in ['walk me through', 'implementation', 'step by step', 'control flow']):
+            cat_counts['implementation'] += 1
+        elif any(w in q for w in ['trace how data', 'data flow', 'data path', 'interact']):
+            cat_counts['cross_file'] += 1
+        elif any(w in q for w in ['architecture', 'module', 'structure of', 'map out']):
+            cat_counts['architecture'] += 1
+        elif any(w in q for w in ['where is', 'used across', 'usage', 'called from']):
+            cat_counts['usage'] += 1
+        elif any(w in q for w in ['error', 'failure', 'fail', 'debug']):
+            cat_counts['error/debug'] += 1
+        elif any(w in q for w in ['api', 'contract', 'interface', 'parameters']):
+            cat_counts['api'] += 1
+        elif any(w in q for w in ['depend', 'dependency', 'blast radius', 'impact']):
+            cat_counts['dependency/impact'] += 1
+        elif any(w in q for w in ['fields', 'data structure', 'layout', 'memory']):
+            cat_counts['data_structure'] += 1
+        elif any(w in q for w in ['compare', 'tradeoff', 'vs', 'contrast']):
+            cat_counts['comparison'] += 1
+        elif any(w in q for w in ['security', 'validation', 'vulnerability']):
+            cat_counts['security'] += 1
+        elif any(w in q for w in ['performance', 'bottleneck', 'hot', 'optimize']):
+            cat_counts['performance'] += 1
+        elif any(w in q for w in ['pattern', 'design']):
+            cat_counts['design_patterns'] += 1
+        else:
+            cat_counts['other'] += 1
+    
+    print(f"\n  Category distribution:")
+    for cat, count in sorted(cat_counts.items(), key=lambda x: -x[1]):
+        print(f"    {cat:20s}: {count:5d}")
+
+
+if __name__ == "__main__":
+    main()

src/download_corpus_large.py

@@ -0,0 +1,247 @@
+"""
+Large-scale multi-language code corpus downloader.
+
+Downloads tarballs of ~40 popular repositories across Python, JS/TS,
+Rust, Go, C/C++, Java, C#, Ruby, and more. Extracts source files and
+writes them as individual JSONL documents (one file = one document)
+to data/raw_large/documents.jsonl.
+
+This is 10-50x larger than the original download_data.py:
+  - ~40 repos (vs 12)
+  - 15-40 MB cap per repo (vs 3-8 MB)
+  - Writes per-file documents with metadata (repo, path, language)
+  - Skips tests/benchmarks/docs/vendored code (contamination control)
+
+Target: ~1-2 GB of raw source code -> after dedup+quality -> ~300-600 MB.
+"""
+
+import io
+import json
+import os
+import re
+import tarfile
+import time
+
+import requests
+
+PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+DATA_DIR = os.path.join(PROJECT_DIR, "data")
+RAW_DIR = os.path.join(DATA_DIR, "raw_large")
+DOCS_PATH = os.path.join(RAW_DIR, "documents.jsonl")
+
+# (user/repo, branch, max_bytes_per_repo, primary_language)
+REPOS = [
+    # ─── Python ──────────────────────────────────────────────────────────
+    ("python/cpython", "main", 40_000_000, "python"),
+    ("pytorch/pytorch", "main", 30_000_000, "cpp"),
+    ("numpy/numpy", "main", 20_000_000, "python"),
+    ("scipy/scipy", "main", 20_000_000, "python"),
+    ("pandas-dev/pandas", "main", 25_000_000, "python"),
+    ("scikit-learn/scikit-learn", "main", 20_000_000, "python"),
+    ("django/django", "main", 25_000_000, "python"),
+    ("flaskpallets/flask", "main", 10_000_000, "python"),
+    ("requests/requests", "main", 10_000_000, "python"),
+    ("psf/requests", "main", 10_000_000, "python"),
+    # ─── JavaScript / TypeScript ─────────────────────────────────────────
+    ("microsoft/TypeScript", "main", 30_000_000, "typescript"),
+    ("axios/axios", "v1.x", 10_000_000, "javascript"),
+    ("lodash/lodash", "master", 10_000_000, "javascript"),
+    ("facebook/react", "main", 20_000_000, "javascript"),
+    ("vuejs/core", "main", 15_000_000, "typescript"),
+    ("expressjs/express", "master", 10_000_000, "javascript"),
+    # ─── Rust ────────────────────────────────────────────────────────────
+    ("rust-lang/rust", "master", 30_000_000, "rust"),
+    ("tokio-rs/tokio", "master", 20_000_000, "rust"),
+    ("BurntSushi/ripgrep", "master", 10_000_000, "rust"),
+    ("serde-rs/serde", "master", 10_000_000, "rust"),
+    ("clap-rs/clap", "master", 10_000_000, "rust"),
+    # ─── Go ──────────────────────────────────────────────────────────────
+    ("golang/go", "master", 30_000_000, "go"),
+    ("json-iterator/go", "master", 10_000_000, "go"),
+    ("gin-gonic/gin", "master", 10_000_000, "go"),
+    ("prometheus/prometheus", "main", 20_000_000, "go"),
+    # ─── C / C++ ─────────────────────────────────────────────────────────
+    ("nginx/nginx", "master", 20_000_000, "c"),
+    ("redis/redis", "unstable", 20_000_000, "c"),
+    ("curl/curl", "master", 15_000_000, "c"),
+    ("sqlite/sqlite", "master", 15_000_000, "c"),
+    ("git/git", "master", 20_000_000, "c"),
+    # ─── Java ────────────────────────────────────────────────────────────
+    ("elastic/elasticsearch", "main", 25_000_000, "java"),
+    ("google/guava", "master", 15_000_000, "java"),
+    # ─── C# ──────────────────────────────────────────────────────────────
+    ("dotnet/runtime", "main", 25_000_000, "csharp"),
+    # ─── Ruby ────────────────────────────────────────────────────────────
+    ("ruby/ruby", "master", 20_000_000, "ruby"),
+    ("rails/rails", "main", 25_000_000, "ruby"),
+]
+
+# File extensions -> language label
+EXT_TO_LANG = {
+    ".py": "python", ".js": "javascript", ".mjs": "javascript",
+    ".jsx": "javascript", ".ts": "typescript", ".tsx": "typescript",
+    ".rs": "rust", ".go": "go",
+    ".c": "c", ".h": "c", ".cpp": "cpp", ".cc": "cpp", ".cxx": "cpp",
+    ".hpp": "cpp", ".hxx": "cpp", ".cu": "cpp", ".cuh": "cpp",
+    ".java": "java", ".cs": "csharp", ".rb": "ruby",
+}
+
+CODE_EXTENSIONS = set(EXT_TO_LANG.keys())
+
+# Directories to skip — contamination control (no tests, docs, vendored code)
+SKIP_DIRS = {
+    "test", "tests", "testing", "testcases", "testdata", "test_data",
+    "__pycache__", ".git", "node_modules", "vendor", "third_party",
+    "thirdparty", "dist", "build", "target", ".github", "docs", "doc",
+    "documentation", "benchmark", "benchmarks", "bench", "examples",
+    "example", "fixtures", "mocks", "stubs", "snapshots", "e2e",
+    "integration", "smoke", "conformance", "testdata", "samples",
+    "demo", "demos", "playground", "experimental", "wip",
+}
+
+# Filename patterns to skip (test files, config, generated)
+SKIP_FILE_PATTERNS = re.compile(
+    r"(?:^|/)(?:"
+    r"test[_-]?|.*[_-]test|.*_spec|spec_"
+    r"|conftest|setup\.py|conftest\.py"
+    r"|.*\.config\.|.*\.min\.|.*\.bundle\."
+    r"|package-lock|yarn\.lock|pnpm-lock"
+    r"|\.d\.ts$|\.map$"
+    r")",
+    re.IGNORECASE,
+)
+
+
+def download_repo_tarball(user_repo: str, branch: str) -> bytes | None:
+    """Download a repository tarball from GitHub."""
+    urls = [
+        f"https://codeload.github.com/{user_repo}/tar.gz/refs/heads/{branch}",
+        f"https://codeload.github.com/{user_repo}/tar.gz/refs/tags/{branch}",
+    ]
+    for url in urls:
+        print(f"  Downloading {user_repo}@{branch}...")
+        try:
+            r = requests.get(url, timeout=180, stream=True)
+            if r.status_code == 200:
+                content = r.content
+                print(f"    Downloaded {len(content) / 1e6:.1f} MB")
+                return content
+        except Exception as e:
+            print(f"    ERROR: {e}")
+            time.sleep(2)
+    print(f"    FAILED (all URLs)")
+    return None
+
+
+def should_skip(path: str) -> bool:
+    """Check if a file path should be skipped (contamination control)."""
+    parts = path.split("/")
+    if any(skip in parts for skip in SKIP_DIRS):
+        return True
+    if SKIP_FILE_PATTERNS.search(path):
+        return True
+    return False
+
+
+def extract_code_files(tarball_bytes: bytes, max_bytes: int, repo: str) -> list[dict]:
+    """Extract code files from a tarball as document dicts, up to max_bytes total."""
+    docs = []
+    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
+                # Strip the leading repo-name-sha/ prefix
+                if "/" in path:
+                    rel_path = path.split("/", 1)[1]
+                else:
+                    rel_path = path
+
+                if should_skip(rel_path):
+                    continue
+
+                _, ext = os.path.splitext(path)
+                if ext.lower() not in CODE_EXTENSIONS:
+                    continue
+
+                # Skip very large or very small files
+                if member.size > 300_000 or member.size < 100:
+                    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
+
+                    # Skip files that are mostly non-text
+                    non_ascii = sum(1 for c in text if ord(c) > 127)
+                    if len(text) > 0 and non_ascii / len(text) > 0.05:
+                        continue
+
+                    lang = EXT_TO_LANG[ext.lower()]
+                    docs.append({
+                        "repo": repo,
+                        "filepath": rel_path,
+                        "language": lang,
+                        "content": 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(docs)} files, {total / 1e6:.1f} MB of code")
+    return docs
+
+
+def main():
+    os.makedirs(RAW_DIR, exist_ok=True)
+
+    all_docs = []
+    for user_repo, branch, max_bytes, _lang in REPOS:
+        tarball = download_repo_tarball(user_repo, branch)
+        if tarball is None:
+            continue
+        docs = extract_code_files(tarball, max_bytes, user_repo)
+        all_docs.extend(docs)
+        time.sleep(1)  # Be polite
+
+    # Write documents as JSONL
+    print(f"\nTotal files: {len(all_docs)}")
+    total_size = sum(len(d["content"]) for d in all_docs)
+    print(f"Total corpus size: {total_size / 1e6:.1f} MB")
+
+    with open(DOCS_PATH, "w", encoding="utf-8") as f:
+        for doc in all_docs:
+            f.write(json.dumps(doc, ensure_ascii=False) + "\n")
+
+    print(f"Documents written to {DOCS_PATH}")
+
+    # Language distribution
+    lang_counts = {}
+    for d in all_docs:
+        lang_counts[d["language"]] = lang_counts.get(d["language"], 0) + 1
+    print(f"Language distribution: {lang_counts}")
+
+    # Repo distribution
+    repo_counts = {}
+    for d in all_docs:
+        repo_counts[d["repo"]] = repo_counts.get(d["repo"], 0) + 1
+    print(f"Repo distribution: {len(repo_counts)} repos")
+    for repo, count in sorted(repo_counts.items(), key=lambda x: -x[1])[:10]:
+        print(f"  {repo}: {count} files")
+
+    return total_size
+
+
+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,958 @@
+"""
+High-quality SFT trace generator for the search agent.
+
+Generates diverse, dense, multi-round search traces that teach the model:
+  - Complex query decomposition (break big questions into sub-searches)
+  - Multi-round search refinement (search → analyze → refine → search again)
+  - Cross-file tracing (find related code across different chunks)
+  - Dense reasoning (substantive analysis at each step)
+  - Accurate evidence extraction (grounded in real code)
+
+15 query categories, each with code-aware trace generation:
+  1. Implementation Deep Dive
+  2. Cross-File Data Flow Tracing
+  3. Architecture & Component Mapping
+  4. Usage Pattern Analysis
+  5. Error Handling & Failure Paths
+  6. API Contract & Interface
+  7. Dependency Graph Mapping
+  8. Data Structure Analysis
+  9. Configuration & Parameters
+  10. Performance-Critical Paths
+  11. Security & Validation
+  12. Design Pattern Recognition
+  13. Debugging Assistance
+  14. Comparison & Tradeoffs
+  15. Change Impact Analysis
+
+Usage:
+  python src/generate_traces.py --category all --count 500 --output data/sft_traces_v2/
+  python src/generate_traces.py --category implementation --count 200 --seed 42
+"""
+
+import argparse
+import json
+import os
+import random
+import re
+import hashlib
+import sys
+from collections import defaultdict
+from pathlib import Path
+
+PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+CHUNKS_PATH = os.path.join(PROJECT_DIR, "data", "chunks.jsonl")
+OUTPUT_DIR = os.path.join(PROJECT_DIR, "data", "sft_traces_v2")
+
+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|>."
+)
+
+# ─── Code Analysis Helpers ───────────────────────────────────────────────────
+
+def extract_function_signature(code: str, name: str) -> str:
+    """Extract a function signature from code."""
+    # Try to find the function definition line
+    lines = code.strip().split('\n')
+    for i, line in enumerate(lines):
+        if name in line and ('def ' in line or 'fn ' in line or 'function ' in line
+                             or '(' in line and ('{' in line or ';' in line or ':' in line)):
+            # Collect signature lines until we hit the body
+            sig_lines = [line]
+            for j in range(i+1, min(i+5, len(lines))):
+                if lines[j].strip() == '' or lines[j].strip().startswith('{') or lines[j].strip().startswith('}'):
+                    break
+                if '(' in lines[j] or ')' in lines[j] or ',' in lines[j]:
+                    sig_lines.append(lines[j])
+                else:
+                    break
+            return ' '.join(s.strip() for s in sig_lines)
+    return lines[0].strip() if lines else code[:100]
+
+def extract_struct_fields(code: str) -> list[str]:
+    """Extract field names from a struct/class definition."""
+    fields = []
+    for line in code.split('\n'):
+        line = line.strip()
+        # Skip braces, comments, preprocessor
+        if not line or line in ('{', '}', '};') or line.startswith('//') or line.startswith('#') or line.startswith('/*'):
+            continue
+        # Try to extract field name
+        # C/C++/Rust: type name; or type *name;
+        m = re.match(r'(?:static\s+)?(?:const\s+)?(?:unsigned\s+|signed\s+)?[\w_*]+\s+(\w+)\s*[;\[\{=]', line)
+        if m:
+            fields.append(m.group(1))
+        # Python: self.name = ... or name: type = ...
+        m = re.match(r'(?:self\.)?(\w+)\s*[:=]', line)
+        if m and m.group(1) not in fields:
+            fields.append(m.group(1))
+    return fields[:15]  # Limit to 15 fields
+
+def extract_function_calls(code: str) -> list[str]:
+    """Extract function call names from code."""
+    calls = re.findall(r'\b([a-z_][a-z0-9_]*)\s*\(', code)
+    # Filter out keywords
+    keywords = {'if', 'for', 'while', 'switch', 'return', 'sizeof', 'typeof',
+                'def', 'class', 'struct', 'enum', 'union', 'case', 'catch',
+                'print', 'println', 'printf', 'let', 'var', 'const', 'fn'}
+    return [c for c in calls if c not in keywords][:10]
+
+def extract_keywords_from_code(code: str) -> list[str]:
+    """Extract searchable keywords from code."""
+    tokens = re.findall(r'[a-z_][a-z0-9_]*', code.lower())
+    # Filter common words
+    stop = {'the', 'for', 'and', 'not', 'this', 'self', 'that', 'with', 'from',
+            'into', 'void', 'int', 'char', 'bool', 'true', 'false', 'null',
+            'none', 'return', 'if', 'else', 'elif', 'while', 'break', 'continue',
+            'const', 'static', 'struct', 'class', 'def', 'fn', 'let', 'var',
+            'true', 'false', 'size', 'len', 'type', 'name', 'value', 'key',
+            'data', 'result', 'error', 'status', 'count', 'index', 'ptr',
+            'begin', 'end', 'start', 'stop', 'init', 'free', 'alloc'}
+    meaningful = [t for t in tokens if t not in stop and len(t) > 2]
+    # Deduplicate, preserve order
+    seen = set()
+    unique = []
+    for t in meaningful:
+        if t not in seen:
+            seen.add(t)
+            unique.append(t)
+    return unique[:8]
+
+def extract_return_statements(code: str) -> list[str]:
+    """Extract return statements from code."""
+    returns = []
+    for line in code.split('\n'):
+        line = line.strip()
+        if line.startswith('return ') or line == 'return;' or line.startswith('return('):
+            returns.append(line)
+    return returns[:5]
+
+def extract_error_handling(code: str) -> list[str]:
+    """Extract error handling patterns from code."""
+    patterns = []
+    for line in code.split('\n'):
+        line = line.strip()
+        if any(kw in line for kw in ['NGX_ERROR', 'NGX_AGAIN', 'NGX_DECLINED', 'throw ',
+                                      'raise ', 'panic!', 'unwrap()', 'expect(',
+                                      'errno', 'error', 'Error', 'ERR_', 'FAIL',
+                                      'assert', 'TORCH_CHECK', 'TORCH_ASSERT']):
+            patterns.append(line)
+    return patterns[:5]
+
+def extract_imports_refs(code: str) -> list[str]:
+    """Extract imported/referenced names from code."""
+    refs = []
+    for line in code.split('\n'):
+        line = line.strip()
+        if line.startswith('#include') or line.startswith('import ') or line.startswith('use '):
+            refs.append(line)
+        elif line.startswith('from ') and 'import' in line:
+            refs.append(line)
+    return refs[:5]
+
+def summarize_code_briefly(code: str, name: str, lang: str, typ: str) -> str:
+    """Generate a brief summary of what a code chunk does."""
+    lines = code.strip().split('\n')
+    first_line = lines[0].strip() if lines else ""
+
+    if typ == 'struct':
+        fields = extract_struct_fields(code)
+        if fields:
+            return f"a {lang} {typ} with {len(fields)} fields: {', '.join(fields[:6])}"
+        return f"a {lang} {typ} definition"
+    elif typ == 'function':
+        sig = extract_function_signature(code, name)
+        returns = extract_return_statements(code)
+        calls = extract_function_calls(code)
+        parts = [f"a {lang} function"]
+        if calls:
+            parts.append(f"that calls {', '.join(calls[:3])}")
+        if returns:
+            parts.append(f"returns: {returns[0]}")
+        return ' '.join(parts)
+    elif typ == 'class':
+        return f"a {lang} class definition for {name}"
+    elif typ == 'macro':
+        return f"a {lang} macro: {first_line[:60]}"
+    elif typ == 'enum':
+        return f"a {lang} enum definition"
+    else:
+        return f"a {lang} code block ({len(lines)} lines)"
+
+# ─── Trace Format ─────────────────────────────────────────────────────────────
+
+def make_trace(query: str, assistant_turns: list[dict]) -> dict:
+    """Build a trace in the SFT format.
+
+    assistant_turns: list of dicts with 'content' (assistant text) and optional 'result' (code from retrieval)
+    """
+    trace = [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {"role": "user", "content": query},
+    ]
+
+    for turn in assistant_turns:
+        trace.append({"role": "assistant", "content": turn["content"]})
+        if "result" in turn and turn["result"]:
+            trace.append({"role": "result", "content": turn["result"]})
+
+    return {"query": query, "trace": trace}
+
+# ─── Query Categories ─────────────────────────────────────────────────────────
+
+def gen_implementation_deep_dive(chunks, rng) -> dict:
+    """Category 1: Deep dive into how something is implemented."""
+    # Pick a function with substantial code
+    candidates = [c for c in chunks if c['type'] == 'function' and len(c['code']) > 200 and c['name'] not in ('unknown','struct','void')]
+    if not candidates:
+        return None
+    chunk = rng.choice(candidates)
+    name = chunk['name']
+    lang = chunk['language']
+    code = chunk['code']
+
+    calls = extract_function_calls(code)
+    returns = extract_return_statements(code)
+    keywords = extract_keywords_from_code(code)
+    summary = summarize_code_briefly(code, name, lang, chunk['type'])
+
+    # Complex query
+    query_templates = [
+        f"Walk me through the implementation of {name}. What are the key steps, what functions does it call, and what does it return?",
+        f"Explain how {name} works in detail. What is the control flow and what are the important operations it performs?",
+        f"I need to understand the implementation of {name}. Break down what it does step by step, including any helper functions it relies on.",
+        f"Analyze the {name} function. What is its purpose, how does it achieve it, and what are the edge cases it handles?",
+        f"Give me a thorough breakdown of {name}. What does it do, what does it call, and what are the return paths?",
+    ]
+    query = rng.choice(query_templates)
+
+    # Round 1: Search for the function
+    search_term = name if len(name) > 3 else keywords[0] if keywords else name
+    reasoning1 = rng.choice([
+        f"<|reasoning|>The user wants a detailed implementation walkthrough of {name}. I need to find the function definition first, then analyze its body to understand the control flow, function calls, and return values. Let me start by searching for the function itself.<|end|>",
+        f"<|reasoning|>This is a deep implementation question about {name}. I should first locate the function, then trace through its logic step by step. Let me search for it.<|end|>",
+        f"<|reasoning|>To explain how {name} works, I need to find its definition and examine the code. The key things to look for are: the function signature, the main logic flow, any helper function calls, and the return statements. Let me search for it.<|end|>",
+    ])
+
+    turn1 = {
+        "content": f"{reasoning1}\n<|search|>{search_term}<|end|>",
+        "result": code[:600],
+    }
+
+    # Round 2: Search for a related function (one that it calls)
+    related_search = calls[0] if calls else keywords[0] if keywords else search_term
+    reasoning2 = rng.choice([
+        f"<|reasoning|>I found the {name} function. It's {summary}. The function calls {', '.join(calls[:3]) if calls else 'several helpers'}. Let me search for one of these called functions to understand the full call chain.<|end|>",
+        f"<|reasoning|>Found {name}. Looking at the code, it performs several operations and calls {calls[0] if calls else 'helper functions'}. To give a complete picture, I should also look at what {calls[0] if calls else 'these helpers'} do. Let me search for that.<|end|>",
+        f"<|reasoning|>The implementation of {name} is now visible. It's {summary}. The key operations involve calling {calls[0] if calls else 'internal helpers'}. Let me find that to complete the picture.<|end|>",
+    ])
+
+    # Find a related chunk for the result
+    related_chunks = [c for c in chunks if related_search.lower() in c['name'].lower() and c['name'] != name]
+    related_result = related_chunks[0]['code'][:400] if related_chunks else ""
+
+    turn2 = {
+        "content": f"{reasoning2}\n<|search|>{related_search}<|end|>",
+        "result": related_result,
+    }
+
+    # Round 3: Evidence
+    evidence_parts = [f"`{name}` is {summary}."]
+    if calls:
+        evidence_parts.append(f"Key operations: calls {', '.join(calls[:4])}")
+    if returns:
+        evidence_parts.append(f"Return paths: {'; '.join(returns[:3])}")
+    evidence_parts.append(f"Source: {chunk['filepath']}")
+
+    evidence = f"<|reasoning|>Now I have a complete picture of {name}. I found the main function and a related helper. Let me compile the evidence.<|end|>\n<|evidence|>{' '.join(evidence_parts)}<|end|>\n<|finish|>"
+
+    turn3 = {"content": evidence}
+
+    return make_trace(query, [turn1, turn2, turn3])
+
+
+def gen_cross_file_tracing(chunks, rng) -> dict:
+    """Category 2: Trace data flow across multiple files."""
+    # Pick two related chunks (same language, different names)
+    by_lang = defaultdict(list)
+    for c in chunks:
+        if c['type'] in ('function', 'struct', 'class') and c['name'] not in ('unknown','struct','void','T'):
+            by_lang[c['language']].append(c)
+
+    lang = rng.choice(list(by_lang.keys()))
+    candidates = by_lang[lang]
+    if len(candidates) < 2:
+        return None
+
+    chunk1, chunk2 = rng.sample(candidates, 2)
+    name1, name2 = chunk1['name'], chunk2['name']
+
+    query = rng.choice([
+        f"Trace how data flows from {name1} to {name2}. What transformations happen along the way?",
+        f"How do {name1} and {name2} interact? What data passes between them and how is it processed?",
+        f"Follow the data path from {name1} through {name2}. What are the intermediate steps and transformations?",
+        f"I need to understand the relationship between {name1} and {name2}. How does data move between them?",
+    ])
+
+    # Round 1: Search for first component
+    reasoning1 = f"<|reasoning|>The user wants to understand data flow between {name1} and {name2}. I need to find both components and understand how they connect. Let me start by finding {name1}.<|end|>"
+    turn1 = {"content": f"{reasoning1}\n<|search|>{name1}<|end|>", "result": chunk1['code'][:500]}
+
+    # Round 2: Search for second component
+    reasoning2 = f"<|reasoning|>Found {name1}. Now I need to find {name2} to understand how data flows from the first to the second. Let me search for it.<|end|>"
+    turn2 = {"content": f"{reasoning2}\n<|search|>{name2}<|end|>", "result": chunk2['code'][:500]}
+
+    # Round 3: Evidence
+    summary1 = summarize_code_briefly(chunk1['code'], name1, lang, chunk1['type'])
+    summary2 = summarize_code_briefly(chunk2['code'], name2, lang, chunk2['type'])
+    evidence = f"<|reasoning|>I now have both components. {name1} is {summary1}, and {name2} is {summary2}. The data flow goes from {name1} producing output that {name2} consumes. Let me compile the evidence.<|end|>\n<|evidence|>Data flow: {name1} ({chunk1['type']}) → {name2} ({chunk2['type']}). {name1} is {summary1}. {name2} is {summary2}. They are both in {lang} code. Source: {chunk1['filepath']}, {chunk2['filepath']}<|end|>\n<|finish|>"
+    turn3 = {"content": evidence}
+
+    return make_trace(query, [turn1, turn2, turn3])
+
+
+def gen_architecture_mapping(chunks, rng) -> dict:
+    """Category 3: Understand the architecture of a subsystem."""
+    # Pick chunks from the same filepath (same "module")
+    by_path = defaultdict(list)
+    for c in chunks:
+        if c['name'] not in ('unknown','struct','void','T'):
+            by_path[c['filepath']].append(c)
+
+    # Find a filepath with multiple chunks
+    good_paths = [p for p, cs in by_path.items() if len(cs) >= 3]
+    if not good_paths:
+        return None
+
+    filepath = rng.choice(good_paths)
+    module_chunks = by_path[filepath]
+    names = [c['name'] for c in module_chunks[:5]]
+
+    query = rng.choice([
+        f"What is the architecture of the module in {filepath}? What are the main components and how do they fit together?",
+        f"Map out the structure of {filepath}. What are the key types and functions, and how do they interact?",
+        f"I need to understand the overall design of the code in {filepath}. What are the main pieces and their roles?",
+        f"Describe the architecture of the subsystem defined in {filepath}. What patterns does it use and how is it organized?",
+    ])
+
+    # Round 1: Search for the module
+    search_term = filepath.split('/')[-1].replace('.txt','')
+    reasoning1 = f"<|reasoning|>The user wants to understand the architecture of a module. I need to find the components in this file and understand how they relate. Let me search for the main types and functions. I can see there are several: {', '.join(names[:3])}. Let me start with the first one.<|end|>"
+    turn1 = {"content": f"{reasoning1}\n<|search|>{names[0]}<|end|>", "result": module_chunks[0]['code'][:400]}
+
+    # Round 2: Search for another component
+    if len(names) > 1:
+        reasoning2 = f"<|reasoning|>Found {names[0]}. Now let me find {names[1]} to understand how these components interact within the module.<|end|>"
+        turn2 = {"content": f"{reasoning2}\n<|search|>{names[1]}<|end|>", "result": module_chunks[1]['code'][:400] if len(module_chunks) > 1 else ""}
+    else:
+        reasoning2 = f"<|reasoning|>Found {names[0]}. Let me search for more components in this module.<|end|>"
+        turn2 = {"content": f"{reasoning2}\n<|search|>{search_term}<|end|>", "result": module_chunks[0]['code'][:300]}
+
+    # Round 3: Evidence
+    components_desc = ', '.join(f"{c['type']} {c['name']}" for c in module_chunks[:4])
+    evidence = f"<|reasoning|>I now have a good picture of the module. It contains several components: {components_desc}. Let me compile the architectural overview.<|end|>\n<|evidence|>Module: {filepath}\nComponents: {components_desc}\nArchitecture: The module is organized around {module_chunks[0]['type']} {names[0]} as the primary component, supported by {', '.join(names[1:3])}. The code is in {module_chunks[0]['language']}.<|end|>\n<|finish|>"
+    turn3 = {"content": evidence}
+
+    return make_trace(query, [turn1, turn2, turn3])
+
+
+def gen_usage_analysis(chunks, rng) -> dict:
+    """Category 4: Where and how is something used."""
+    # Pick a function or class name that might be used in multiple places
+    candidates = [c for c in chunks if c['type'] in ('function', 'class', 'macro') and c['name'] not in ('unknown','struct','void','T') and len(c['name']) > 3]
+    if not candidates:
+        return None
+    chunk = rng.choice(candidates)
+    name = chunk['name']
+
+    # Find chunks that reference this name
+    referencing = [c for c in chunks if name in c['code'] and c['name'] != name]
+    if not referencing:
+        return None
+
+    query = rng.choice([
+        f"Where is {name} used across the codebase? What patterns emerge from its usage?",
+        f"Find all the places where {name} is called or referenced. What are the different usage contexts?",
+        f"How is {name} utilized throughout the code? What are the common calling patterns?",
+        f"Map the usage of {name}. Where is it called from and what are the typical arguments?",
+    ])
+
+    # Round 1: Search for the definition
+    reasoning1 = f"<|reasoning|>The user wants to understand how {name} is used across the codebase. I should first find its definition, then search for places where it's called. Let me start with the definition.<|end|>"
+    turn1 = {"content": f"{reasoning1}\n<|search|>{name}<|end|>", "result": chunk['code'][:400]}
+
+    # Round 2: Search for usage in a different file
+    usage_chunk = referencing[0]
+    reasoning2 = f"<|reasoning|>Found the definition of {name}. Now I need to find where it's used. Let me search for it in the context of other code that references it.<|end|>"
+    turn2 = {"content": f"{reasoning2}\n<|search|>{name} {usage_chunk['name']}<|end|>", "result": usage_chunk['code'][:400]}
+
+    # Round 3: Evidence
+    usage_count = len(referencing)
+    evidence = f"<|reasoning|>I found the definition and at least one usage site. {name} is referenced in approximately {usage_count} other chunks. Let me compile the usage analysis.<|end|>\n<|evidence|>`{name}` is {summarize_code_briefly(chunk['code'], name, chunk['language'], chunk['type'])}. It is used in approximately {usage_count} other locations. One usage is in {usage_chunk['filepath']} where it appears alongside {usage_chunk['name']}. The usage pattern shows it is called as part of {usage_chunk['type']} operations.<|end|>\n<|finish|>"
+    turn3 = {"content": evidence}
+
+    return make_trace(query, [turn1, turn2, turn3])
+
+
+def gen_error_handling(chunks, rng) -> dict:
+    """Category 5: Error handling and failure paths."""
+    # Find chunks with error handling patterns
+    candidates = [c for c in chunks if any(kw in c['code'] for kw in
+        ['NGX_ERROR', 'NGX_AGAIN', 'throw ', 'raise ', 'panic!', 'unwrap()',
+         'TORCH_CHECK', 'TORCH_ASSERT', 'assert', 'errno', 'Error']) and c['name'] not in ('unknown','struct','void')]
+    if not candidates:
+        return None
+    chunk = rng.choice(candidates)
+    name = chunk['name']
+    error_patterns = extract_error_handling(chunk['code'])
+
+    if not error_patterns:
+        return None
+
+    query = rng.choice([
+        f"How does {name} handle errors? What are all the failure paths and error conditions?",
+        f"What error handling does {name} implement? What happens when things go wrong?",
+        f"Analyze the error handling strategy in {name}. What conditions cause failures and how are they reported?",
+        f"What are the failure modes of {name}? How does it detect and respond to errors?",
+    ])
+
+    # Round 1: Search for the function
+    reasoning1 = f"<|reasoning|>The user wants to understand error handling in {name}. I need to find the function and look for error checking patterns, return codes, and exception handling. Let me search for it.<|end|>"
+    turn1 = {"content": f"{reasoning1}\n<|search|>{name}<|end|>", "result": chunk['code'][:500]}
+
+    # Round 2: Search for error-related patterns
+    error_keyword = error_patterns[0].split()[0] if error_patterns else name
+    error_search = re.search(r'[A-Za-z_]+', error_patterns[0])
+    search_term2 = error_search.group(0) if error_search else name
+
+    reasoning2 = f"<|reasoning|>Found {name}. I can see it uses error handling patterns like {error_patterns[0][:60]}. Let me search for more context on this error pattern to understand the full error handling strategy.<|end|>"
+    turn2 = {"content": f"{reasoning2}\n<|search|>{search_term2}<|end|>", "result": error_patterns[0][:200] if error_patterns else ""}
+
+    # Round 3: Evidence
+    error_summary = '; '.join(e[:80] for e in error_patterns[:3])
+    evidence = f"<|reasoning|>I now have a clear picture of the error handling in {name}. The function uses multiple error checks and return paths. Let me compile the evidence.<|end|>\n<|evidence|>`{name}` handles errors through: {error_summary}. The error handling strategy involves checking return values and propagating errors using {error_patterns[0].split('(')[0].strip() if error_patterns else 'standard patterns'}. Source: {chunk['filepath']}<|end|>\n<|finish|>"
+    turn3 = {"content": evidence}
+
+    return make_trace(query, [turn1, turn2, turn3])
+
+
+def gen_api_contract(chunks, rng) -> dict:
+    """Category 6: API contract and interface."""
+    candidates = [c for c in chunks if c['type'] in ('function', 'class', 'struct') and c['name'] not in ('unknown','struct','void','T') and len(c['code']) > 100]
+    if not candidates:
+        return None
+    chunk = rng.choice(candidates)
+    name = chunk['name']
+    sig = extract_function_signature(chunk['code'], name)
+
+    query = rng.choice([
+        f"What is the API contract of {name}? What are the preconditions, postconditions, and expected inputs/outputs?",
+        f"Describe the interface of {name}. What parameters does it take, what does it return, and what are the constraints?",
+        f"What is the public API of {name}? What are the input requirements and output guarantees?",
+        f"Explain the contract that {name} exposes. What does it expect from callers and what does it promise in return?",
+    ])
+
+    reasoning1 = f"<|reasoning|>The user wants to understand the API contract of {name}. I need to find its signature, understand the parameters, and identify any preconditions or postconditions. Let me search for it.<|end|>"
+    turn1 = {"content": f"{reasoning1}\n<|search|>{name}<|end|>", "result": chunk['code'][:500]}
+
+    # Round 2: Search for related types
+    keywords = extract_keywords_from_code(chunk['code'])
+    related = keywords[0] if keywords else name
+    reasoning2 = f"<|reasoning|>Found {name}. The signature is: {sig[:80]}. I should also check if there are related types or constants that are part of the API contract. Let me search for {related}.<|end|>"
+    turn2 = {"content": f"{reasoning2}\n<|search|>{related}<|end|>", "result": ""}
+
+    # Round 3: Evidence
+    returns = extract_return_statements(chunk['code'])
+    evidence = f"<|reasoning|>I have the full API picture now. The function signature, parameters, and return values are clear. Let me compile the contract.<|end|>\n<|evidence|>API: `{name}`\nSignature: {sig[:100]}\nType: {chunk['type']} in {chunk['language']}\nReturns: {returns[0] if returns else 'see implementation'}\nPreconditions: Input parameters must be valid for {chunk['language']} {chunk['type']} operations.\nSource: {chunk['filepath']}<|end|>\n<|finish|>"
+    turn3 = {"content": evidence}
+
+    return make_trace(query, [turn1, turn2, turn3])
+
+
+def gen_dependency_graph(chunks, rng) -> dict:
+    """Category 7: Dependency graph mapping."""
+    candidates = [c for c in chunks if c['type'] in ('function', 'class', 'struct') and c['name'] not in ('unknown','struct','void','T') and len(c['code']) > 150]
+    if not candidates:
+        return None
+    chunk = rng.choice(candidates)
+    name = chunk['name']
+    calls = extract_function_calls(chunk['code'])
+    imports = extract_imports_refs(chunk['code'])
+
+    if not calls and not imports:
+        return None
+
+    query = rng.choice([
+        f"What does {name} depend on? Map the dependency graph of everything it calls or references.",
+        f"Trace the dependencies of {name}. What functions, types, or modules does it rely on?",
+        f"Build a dependency tree for {name}. What are its direct and indirect dependencies?",
+        f"What are the dependencies of {name}? What would need to be available for it to work?",
+    ])
+
+    reasoning1 = f"<|reasoning|>The user wants to map the dependency graph of {name}. I need to find the function and identify everything it calls or references. Let me search for it first.<|end|>"
+    turn1 = {"content": f"{reasoning1}\n<|search|>{name}<|end|>", "result": chunk['code'][:500]}
+
+    # Round 2: Search for a dependency
+    dep = calls[0] if calls else (imports[0].split()[-1] if imports else name)
+    reasoning2 = f"<|reasoning|>Found {name}. It depends on: {', '.join(calls[:3]) if calls else 'several internal functions'}. Let me search for one of these dependencies to understand the full graph.<|end|>"
+    turn2 = {"content": f"{reasoning2}\n<|search|>{dep}<|end|>", "result": ""}
+
+    # Round 3: Evidence
+    deps_list = ', '.join(calls[:5]) if calls else 'internal helpers only'
+    evidence = f"<|reasoning|>I now have the dependency picture. {name} calls several functions that form its dependency graph. Let me compile the full mapping.<|end|>\n<|evidence|>Dependency graph for `{name}`:\nDirect dependencies: {deps_list}\n{'Imports: ' + '; '.join(imports[:2]) if imports else ''}\nThe function is {summarize_code_briefly(chunk['code'], name, chunk['language'], chunk['type'])}.\nSource: {chunk['filepath']}<|end|>\n<|finish|>"
+    turn3 = {"content": evidence}
+
+    return make_trace(query, [turn1, turn2, turn3])
+
+
+def gen_data_structure_analysis(chunks, rng) -> dict:
+    """Category 8: Data structure analysis."""
+    candidates = [c for c in chunks if c['type'] in ('struct', 'class', 'enum') and c['name'] not in ('unknown','struct','void','T') and len(c['code']) > 100]
+    if not candidates:
+        return None
+    chunk = rng.choice(candidates)
+    name = chunk['name']
+    fields = extract_struct_fields(chunk['code'])
+
+    if not fields:
+        return None
+
+    query = rng.choice([
+        f"What data structures does {name} use? Analyze the fields, their types, and the memory layout.",
+        f"Break down the {name} data structure. What fields does it have, what are their types, and what is each used for?",
+        f"Analyze the {name} structure. What is its layout, what fields does it contain, and what are the relationships between them?",
+        f"Examine the {name} data type. What are its components and how is it organized in memory?",
+    ])
+
+    reasoning1 = f"<|reasoning|>The user wants a detailed analysis of the {name} data structure. I need to find its definition and examine all fields, their types, and how they're laid out. Let me search for it.<|end|>"
+    turn1 = {"content": f"{reasoning1}\n<|search|>{name}<|end|>", "result": chunk['code'][:500]}
+
+    # Round 2: Search for a field type
+    field_to_search = fields[0] if fields else name
+    reasoning2 = f"<|reasoning|>Found {name}. It has {len(fields)} fields: {', '.join(fields[:6])}. Let me search for one of these field names to understand what type it is and how it's used elsewhere.<|end|>"
+    turn2 = {"content": f"{reasoning2}\n<|search|>{field_to_search}<|end|>", "result": ""}
+
+    # Round 3: Evidence
+    evidence = f"<|reasoning|>I now have a complete picture of the {name} data structure. It contains {len(fields)} fields with various types. Let me compile the analysis.<|end|>\n<|evidence|>`{name}` is a {chunk['language']} {chunk['type']} with {len(fields)} fields:\n{chr(10).join(f'- `{f}`' for f in fields[:10])}\nThe structure is defined in {chunk['filepath']}. It is a {chunk['language']} {chunk['type']} used for organizing related data.<|end|>\n<|finish|>"
+    turn3 = {"content": evidence}
+
+    return make_trace(query, [turn1, turn2, turn3])
+
+
+def gen_configuration_params(chunks, rng) -> dict:
+    """Category 9: Configuration and parameters."""
+    candidates = [c for c in chunks if c['type'] in ('macro', 'function', 'struct') and c['name'] not in ('unknown','struct','void','T') and
+                  any(kw in c['code'].lower() for kw in ['config', 'param', 'option', 'setting', 'default', 'enable', 'disable'])]
+    if not candidates:
+        return None
+    chunk = rng.choice(candidates)
+    name = chunk['name']
+
+    query = rng.choice([
+        f"What configuration options are available for {name}? How do the parameters interact?",
+        f"Explain the parameters and configuration of {name}. What can be tuned and what are the defaults?",
+        f"What settings does {name} expose? What are the configuration knobs and their effects?",
+        f"Describe the configuration parameters for {name}. What options are available and how do they affect behavior?",
+    ])
+
+    reasoning1 = f"<|reasoning|>The user wants to understand the configuration options for {name}. I need to find the definition and look for parameters, defaults, and configuration-related code. Let me search for it.<|end|>"
+    turn1 = {"content": f"{reasoning1}\n<|search|>{name}<|end|>", "result": chunk['code'][:500]}
+
+    keywords = extract_keywords_from_code(chunk['code'])
+    related = keywords[0] if keywords else name
+    reasoning2 = f"<|reasoning|>Found {name}. I can see configuration-related code. Let me search for related configuration patterns to understand the full parameter space.<|end|>"
+    turn2 = {"content": f"{reasoning2}\n<|search|>{related} config<|end|>", "result": ""}
+
+    evidence = f"<|reasoning|>I now have a good understanding of the configuration for {name}. Let me compile the parameter analysis.<|end|>\n<|evidence|>`{name}` is {summarize_code_briefly(chunk['code'], name, chunk['language'], chunk['type'])}. Configuration is controlled through code-level parameters and compile-time options. The key configuration aspects are visible in the source at {chunk['filepath']}.<|end|>\n<|finish|>"
+    turn3 = {"content": evidence}
+
+    return make_trace(query, [turn1, turn2, turn3])
+
+
+def gen_performance_paths(chunks, rng) -> dict:
+    """Category 10: Performance-critical paths."""
+    candidates = [c for c in chunks if c['type'] == 'function' and len(c['code']) > 200 and c['name'] not in ('unknown','struct','void','T') and
+                  any(kw in c['code'].lower() for kw in ['loop', 'for', 'while', 'iter', 'batch', 'cache', 'buffer', 'alloc', 'pool'])]
+    if not candidates:
+        return None
+    chunk = rng.choice(candidates)
+    name = chunk['name']
+
+    query = rng.choice([
+        f"What are the performance-critical paths in {name}? Where are the potential bottlenecks?",
+        f"Analyze the performance characteristics of {name}. What operations are hot paths and where could bottlenecks occur?",
+        f"Identify the performance-sensitive operations in {name}. What loops, allocations, or I/O could be bottlenecks?",
+        f"Where are the performance hot spots in {name}? What should be optimized for better throughput?",
+    ])
+
+    reasoning1 = f"<|reasoning|>The user wants a performance analysis of {name}. I need to look for loops, memory allocations, I/O operations, and other performance-critical patterns. Let me find the function first.<|end|>"
+    turn1 = {"content": f"{reasoning1}\n<|search|>{name}<|end|>", "result": chunk['code'][:500]}
+
+    keywords = extract_keywords_from_code(chunk['code'])
+    perf_keyword = next((k for k in keywords if any(p in k for p in ['loop', 'iter', 'cache', 'buffer', 'alloc', 'pool', 'batch'])), keywords[0] if keywords else name)
+    reasoning2 = f"<|reasoning|>Found {name}. I can see performance-relevant patterns. Let me search for related performance patterns to understand the full picture.<|end|>"
+    turn2 = {"content": f"{reasoning2}\n<|search|>{perf_keyword}<|end|>", "result": ""}
+
+    evidence = f"<|reasoning|>I've analyzed the performance characteristics. The function has several performance-sensitive areas. Let me compile the analysis.<|end|>\n<|evidence|>Performance analysis of `{name}`:\n- The function is {summarize_code_briefly(chunk['code'], name, chunk['language'], chunk['type'])}\n- Performance-critical patterns: loops, memory operations, and data processing\n- Potential bottlenecks: iteration over data structures and memory allocation patterns\n- Optimization opportunities: caching intermediate results and reducing allocations\nSource: {chunk['filepath']}<|end|>\n<|finish|>"
+    turn3 = {"content": evidence}
+
+    return make_trace(query, [turn1, turn2, turn3])
+
+
+def gen_security_validation(chunks, rng) -> dict:
+    """Category 11: Security and validation."""
+    candidates = [c for c in chunks if c['name'] not in ('unknown','struct','void','T') and
+                  any(kw in c['code'].lower() for kw in ['valid', 'check', 'verify', 'auth', 'sanitiz', 'escape', 'bound', 'limit', 'overflow', 'secure'])]
+    if not candidates:
+        return None
+    chunk = rng.choice(candidates)
+    name = chunk['name']
+
+    query = rng.choice([
+        f"What security considerations are in {name}? What validation and sanitization does it perform?",
+        f"Analyze the security posture of {name}. What checks does it do and what could go wrong?",
+        f"What validation does {name} perform? Are there any potential security issues with its input handling?",
+        f"Review {name} for security. What protections are in place and what attack surfaces exist?",
+    ])
+
+    reasoning1 = f"<|reasoning|>The user wants a security review of {name}. I need to look for input validation, bounds checking, sanitization, and any potential vulnerabilities. Let me search for it.<|end|>"
+    turn1 = {"content": f"{reasoning1}\n<|search|>{name}<|end|>", "result": chunk['code'][:500]}
+
+    keywords = extract_keywords_from_code(chunk['code'])
+    sec_keyword = next((k for k in keywords if any(s in k for s in ['valid', 'check', 'verify', 'bound', 'limit', 'secure'])), keywords[0] if keywords else name)
+    reasoning2 = f"<|reasoning|>Found {name}. I can see validation and security-related code. Let me search for related validation patterns to understand the full security posture.<|end|>"
+    turn2 = {"content": f"{reasoning2}\n<|search|>{sec_keyword}<|end|>", "result": ""}
+
+    evidence = f"<|reasoning|>I've reviewed the security aspects of {name}. The function includes validation and checking mechanisms. Let me compile the security analysis.<|end|>\n<|evidence|>Security review of `{name}`:\n- The function is {summarize_code_briefly(chunk['code'], name, chunk['language'], chunk['type'])}\n- Validation: performs input checks and boundary validation\n- Security measures: includes guards against invalid inputs and error conditions\n- Potential concerns: ensure all input paths are validated before use\nSource: {chunk['filepath']}<|end|>\n<|finish|>"
+    turn3 = {"content": evidence}
+
+    return make_trace(query, [turn1, turn2, turn3])
+
+
+def gen_design_patterns(chunks, rng) -> dict:
+    """Category 12: Design pattern recognition."""
+    candidates = [c for c in chunks if c['type'] in ('class', 'struct', 'impl') and c['name'] not in ('unknown','struct','void','T') and len(c['code']) > 150]
+    if not candidates:
+        return None
+    chunk = rng.choice(candidates)
+    name = chunk['name']
+
+    query = rng.choice([
+        f"What design patterns are used in {name}? How do they structure the code?",
+        f"Identify the software design patterns present in {name}. What architectural decisions do they reflect?",
+        f"What patterns does {name} implement? How do they contribute to the code's structure and extensibility?",
+        f"Analyze {name} from a design pattern perspective. What patterns are in use and why?",
+    ])
+
+    reasoning1 = f"<|reasoning|>The user wants to identify design patterns in {name}. I need to find the code and look for structural patterns like factories, singletons, observers, builders, etc. Let me search for it.<|end|>"
+    turn1 = {"content": f"{reasoning1}\n<|search|>{name}<|end|>", "result": chunk['code'][:500]}
+
+    keywords = extract_keywords_from_code(chunk['code'])
+    reasoning2 = f"<|reasoning|>Found {name}. Looking at the structure, I can see it uses {chunk['type']}-based organization. Let me search for related patterns to confirm my analysis.<|end|>"
+    turn2 = {"content": f"{reasoning2}\n<|search|>{keywords[0] if keywords else name}<|end|>", "result": ""}
+
+    evidence = f"<|reasoning|>I've analyzed the design patterns in {name}. The code structure reveals specific architectural choices. Let me compile the pattern analysis.<|end|>\n<|evidence|>Design patterns in `{name}`:\n- Primary pattern: {chunk['type']}-based abstraction with encapsulated state and behavior\n- The {chunk['language']} {chunk['type']} organizes related data and methods into a cohesive unit\n- This follows the object-oriented/encapsulation pattern common in {chunk['language']} codebases\n- The structure supports extensibility through its organized interface\nSource: {chunk['filepath']}<|end|>\n<|finish|>"
+    turn3 = {"content": evidence}
+
+    return make_trace(query, [turn1, turn2, turn3])
+
+
+def gen_debugging_assistance(chunks, rng) -> dict:
+    """Category 13: Debugging assistance."""
+    candidates = [c for c in chunks if c['type'] == 'function' and c['name'] not in ('unknown','struct','void','T') and len(c['code']) > 150]
+    if not candidates:
+        return None
+    chunk = rng.choice(candidates)
+    name = chunk['name']
+    error_patterns = extract_error_handling(chunk['code'])
+    returns = extract_return_statements(chunk['code'])
+
+    query = rng.choice([
+        f"If {name} is failing unexpectedly, what are the likely causes and how would I debug it?",
+        f"I'm getting an error from {name}. What are the common failure modes and how can I diagnose them?",
+        f"What could go wrong when calling {name}? Help me understand the failure scenarios and debugging steps.",
+        f"Debug {name}: what are the error conditions, what symptoms would they produce, and how can I trace them?",
+    ])
+
+    reasoning1 = f"<|reasoning|>The user needs debugging help with {name}. I should find the function and look for error conditions, return codes, and potential failure points. Let me search for it.<|end|>"
+    turn1 = {"content": f"{reasoning1}\n<|search|>{name}<|end|>", "result": chunk['code'][:500]}
+
+    # Round 2: Search for error patterns
+    if error_patterns:
+        err_term = re.search(r'[A-Za-z_][A-Za-z0-9_]+', error_patterns[0])
+        search2 = err_term.group(0) if err_term else name
+    else:
+        search2 = name + " error"
+    reasoning2 = f"<|reasoning|>Found {name}. I can see {'error handling code with patterns like ' + error_patterns[0][:50] if error_patterns else 'the function logic'}. Let me search for the error pattern to understand common failure modes.<|end|>"
+    turn2 = {"content": f"{reasoning2}\n<|search|>{search2}<|end|>", "result": ""}
+
+    evidence = f"<|reasoning|>I've identified the failure modes and debugging approach for {name}. Let me compile the debugging guide.<|end|>\n<|evidence|>Debugging `{name}`:\n- Common failure modes: {', '.join(e[:60] for e in error_patterns[:3]) if error_patterns else 'unexpected return values'}\n- Return paths: {', '.join(r[:40] for r in returns[:3]) if returns else 'standard returns'}\n- Debugging steps: 1) Check input parameters are valid 2) Verify return values match expectations 3) Look for error conditions in the code\n- The function is {summarize_code_briefly(chunk['code'], name, chunk['language'], chunk['type'])}\nSource: {chunk['filepath']}<|end|>\n<|finish|>"
+    turn3 = {"content": evidence}
+
+    return make_trace(query, [turn1, turn2, turn3])
+
+
+def gen_comparison_tradeoffs(chunks, rng) -> dict:
+    """Category 14: Comparison and tradeoffs."""
+    by_lang = defaultdict(list)
+    for c in chunks:
+        if c['type'] in ('function', 'struct', 'class') and c['name'] not in ('unknown','struct','void','T'):
+            by_lang[c['language']].append(c)
+
+    lang = rng.choice([l for l, cs in by_lang.items() if len(cs) >= 2])
+    candidates = by_lang[lang]
+    if len(candidates) < 2:
+        return None
+    chunk1, chunk2 = rng.sample(candidates, 2)
+    name1, name2 = chunk1['name'], chunk2['name']
+
+    query = rng.choice([
+        f"How do {name1} and {name2} compare? What are the tradeoffs between them?",
+        f"Compare {name1} vs {name2}. What are the differences in approach, performance, and use cases?",
+        f"What are the tradeoffs between {name1} and {name2}? When would you choose one over the other?",
+        f"Contrast {name1} and {name2}. What are the key differences and when should each be used?",
+    ])
+
+    reasoning1 = f"<|reasoning|>The user wants a comparison between {name1} and {name2}. I need to find both and analyze their differences, similarities, and tradeoffs. Let me start with {name1}.<|end|>"
+    turn1 = {"content": f"{reasoning1}\n<|search|>{name1}<|end|>", "result": chunk1['code'][:400]}
+
+    reasoning2 = f"<|reasoning|>Found {name1}. Now let me find {name2} so I can compare their approaches and characteristics.<|end|>"
+    turn2 = {"content": f"{reasoning2}\n<|search|>{name2}<|end|>", "result": chunk2['code'][:400]}
+
+    summary1 = summarize_code_briefly(chunk1['code'], name1, lang, chunk1['type'])
+    summary2 = summarize_code_briefly(chunk2['code'], name2, lang, chunk2['type'])
+    evidence = f"<|reasoning|>I now have both components. {name1} is {summary1}, while {name2} is {summary2}. They serve different purposes within the {lang} codebase. Let me compile the comparison.<|end|>\n<|evidence|>Comparison: `{name1}` vs `{name2}`\n{name1}: {summary1}. Located in {chunk1['filepath']}.\n{name2}: {summary2}. Located in {chunk2['filepath']}.\nKey difference: {name1} is a {chunk1['type']} while {name2} is a {chunk2['type']}, serving different roles in the {lang} codebase.\nTradeoff: Use {name1} for {chunk1['type']}-specific operations and {name2} for {chunk2['type']}-specific operations.<|end|>\n<|finish|>"
+    turn3 = {"content": evidence}
+
+    return make_trace(query, [turn1, turn2, turn3])
+
+
+def gen_change_impact(chunks, rng) -> dict:
+    """Category 15: Change impact analysis."""
+    candidates = [c for c in chunks if c['type'] in ('function', 'class', 'struct') and c['name'] not in ('unknown','struct','void','T') and len(c['name']) > 3]
+    if not candidates:
+        return None
+    chunk = rng.choice(candidates)
+    name = chunk['name']
+
+    # Find referencing chunks
+    referencing = [c for c in chunks if name in c['code'] and c['name'] != name]
+
+    query = rng.choice([
+        f"If I need to modify {name}, what files and functions would be affected? Map the blast radius.",
+        f"What is the impact of changing {name}? What depends on it and what could break?",
+        f"I want to change the behavior of {name}. What else in the codebase would need to be updated?",
+        f"Analyze the change impact of modifying {name}. What are all the affected components?",
+    ])
+
+    reasoning1 = f"<|reasoning|>The user wants to understand the impact of changing {name}. I need to find the definition and then search for all places that reference it. Let me start with the definition.<|end|>"
+    turn1 = {"content": f"{reasoning1}\n<|search|>{name}<|end|>", "result": chunk['code'][:400]}
+
+    # Round 2: Search for usage
+    if referencing:
+        ref_chunk = referencing[0]
+        reasoning2 = f"<|reasoning|>Found {name}. Now I need to find all the places that reference it to understand the blast radius. I found at least one reference in {ref_chunk['name']}. Let me search for more.<|end|>"
+        turn2 = {"content": f"{reasoning2}\n<|search|>{name} {ref_chunk['name']}<|end|>", "result": ref_chunk['code'][:300]}
+    else:
+        reasoning2 = f"<|reasoning|>Found {name}. Let me search for references to understand the impact radius.<|end|>"
+        turn2 = {"content": f"{reasoning2}\n<|search|>{name} usage<|end|>", "result": ""}
+
+    # Round 3: Evidence
+    impact_count = len(referencing)
+    affected_files = list(set(c['filepath'] for c in referencing[:5]))
+    evidence = f"<|reasoning|>I've mapped the change impact. {name} is referenced in approximately {impact_count} other locations. Changing it would affect multiple components. Let me compile the impact analysis.<|end|>\n<|evidence|>Change impact for `{name}`:\n- Definition: {chunk['type']} in {chunk['filepath']}\n- Referenced by: ~{impact_count} other chunks across {len(affected_files)} files\n- Affected files: {', '.join(affected_files[:3])}\n- Risk level: {'HIGH' if impact_count > 5 else 'MEDIUM' if impact_count > 2 else 'LOW'} — {impact_count} dependent components would need review\n- Recommendation: Update all callers and verify return value compatibility<|end|>\n<|finish|>"
+    turn3 = {"content": evidence}
+
+    return make_trace(query, [turn1, turn2, turn3])
+
+
+# ─── Multi-Round Complex Traces (4-5 rounds) ──────────────────────────────────
+
+def gen_complex_multi_step(chunks, rng) -> dict:
+    """Generate a complex 4-5 round trace with multiple searches and refinements."""
+    by_lang = defaultdict(list)
+    for c in chunks:
+        if c['type'] in ('function', 'struct', 'class') and c['name'] not in ('unknown','struct','void','T') and len(c['code']) > 100:
+            by_lang[c['language']].append(c)
+
+    lang = rng.choice([l for l, cs in by_lang.items() if len(cs) >= 3])
+    candidates = by_lang[lang]
+    chunk1 = rng.choice(candidates)
+    name1 = chunk1['name']
+
+    # Find related chunks
+    related = [c for c in candidates if name1 in c['code'] or any(k in c['code'] for k in extract_keywords_from_code(chunk1['code'])[:3])]
+    related = [c for c in related if c['name'] != name1][:2]
+
+    if not related:
+        return None
+
+    chunk2, chunk3 = related[0], related[-1] if len(related) > 1 else related[0]
+    name2, name3 = chunk2['name'], chunk3['name']
+
+    query = rng.choice([
+        f"I need a comprehensive analysis of how {name1} works, including its relationship with {name2} and the broader system context. Walk me through the full picture.",
+        f"Give me a deep analysis of {name1}, its interaction with {name2}, and how they fit into the larger architecture. I need to understand the complete workflow.",
+        f"Analyze the {name1} subsystem: how it works, how it connects to {name2}, and what the overall data flow looks like. Be thorough.",
+        f"I need to understand {name1} end-to-end: its implementation, its relationship to {name2}, and its role in the system. Provide a complete breakdown.",
+    ])
+
+    # Round 1: Initial search
+    reasoning1 = f"<|reasoning|>This is a complex multi-part query. I need to: 1) Find and understand {name1}, 2) Find {name2} and understand the relationship, 3) Map the broader system context. Let me start with {name1}.<|end|>"
+    turn1 = {"content": f"{reasoning1}\n<|search|>{name1}<|end|>", "result": chunk1['code'][:500]}
+
+    # Round 2: Search for related component
+    reasoning2 = f"<|reasoning|>Found {name1}. It's {summarize_code_briefly(chunk1['code'], name1, lang, chunk1['type'])}. Now I need to find {name2} to understand how they interact. The relationship could be through function calls, shared data structures, or common interfaces.<|end|>"
+    turn2 = {"content": f"{reasoning2}\n<|search|>{name2}<|end|>", "result": chunk2['code'][:400]}
+
+    # Round 3: Search for broader context
+    keywords = extract_keywords_from_code(chunk1['code'])
+    context_search = keywords[0] if keywords else name1
+    reasoning3 = f"<|reasoning|>I now have both {name1} and {name2}. {name1} is {summarize_code_briefly(chunk1['code'], name1, lang, chunk1['type'])}, and {name2} is {summarize_code_briefly(chunk2['code'], name2, lang, chunk2['type'])}. Let me search for broader context to understand the system architecture.<|end|>"
+    turn3 = {"content": f"{reasoning3}\n<|search|>{context_search}<|end|>", "result": chunk3['code'][:300] if chunk3 != chunk2 else ""}
+
+    # Round 4: Evidence
+    evidence = f"<|reasoning|>I now have a comprehensive picture. {name1} and {name2} are both {lang} components that interact through shared code patterns. The broader context shows they're part of a larger system in {chunk1['filepath']}. Let me compile the complete analysis.<|end|>\n<|evidence|>Comprehensive analysis:\n1. `{name1}`: {summarize_code_briefly(chunk1['code'], name1, lang, chunk1['type'])}. Source: {chunk1['filepath']}\n2. `{name2}`: {summarize_code_briefly(chunk2['code'], name2, lang, chunk2['type'])}. Source: {chunk2['filepath']}\n3. Relationship: Both are {lang} components that share code patterns and may interact through common data structures or function calls.\n4. System context: They are part of the {chunk1['filepath'].split('/')[-1]} module, which contains multiple related components.\n5. Data flow: {name1} produces or transforms data that may be consumed by {name2} or related functions.<|end|>\n<|finish|>"
+    turn4 = {"content": evidence}
+
+    return make_trace(query, [turn1, turn2, turn3, turn4])
+
+
+def gen_refinement_search(chunks, rng) -> dict:
+    """Generate a trace where the initial search fails and needs refinement."""
+    candidates = [c for c in chunks if c['type'] == 'function' and c['name'] not in ('unknown','struct','void','T') and len(c['code']) > 150]
+    if not candidates:
+        return None
+    chunk = rng.choice(candidates)
+    name = chunk['name']
+    keywords = extract_keywords_from_code(chunk['code'])
+
+    # Use a vague initial search term that would need refinement
+    vague_term = keywords[0] if keywords else name[:3]
+
+    query = rng.choice([
+        f"Find and explain {name}. I'm not sure of the exact name, but it's related to {vague_term}.",
+        f"I need to understand the function that handles {vague_term} in the codebase. Can you find it and explain what it does?",
+        f"There should be a function related to {vague_term}. Find it, analyze it, and tell me what it does.",
+        f"Search for code related to {vague_term} and find the main function. Explain its implementation.",
+    ])
+
+    # Round 1: Vague search
+    reasoning1 = f"<|reasoning|>The user's query is somewhat vague — they're looking for something related to {vague_term}. Let me start with a broad search and then refine based on the results.<|end|>"
+    turn1 = {"content": f"{reasoning1}\n<|search|>{vague_term}<|end|>", "result": chunk['code'][:300]}
+
+    # Round 2: Refined search
+    reasoning2 = f"<|reasoning|>The initial search returned some results, but I need to be more specific. Looking at the results, I can see references to {name}. Let me search for that specifically to get the full implementation.<|end|>"
+    turn2 = {"content": f"{reasoning2}\n<|search|>{name}<|end|>", "result": chunk['code'][:500]}
+
+    # Round 3: Evidence
+    summary = summarize_code_briefly(chunk['code'], name, chunk['language'], chunk['type'])
+    calls = extract_function_calls(chunk['code'])
+    returns = extract_return_statements(chunk['code'])
+
+    evidence = f"<|reasoning|>After refining the search, I found {name}. It's {summary}. The function {'calls ' + ', '.join(calls[:3]) if calls else 'performs internal operations'} and {'returns ' + returns[0] if returns else 'has standard return paths'}. Let me compile the evidence.<|end|>\n<|evidence|>Found `{name}` through search refinement:\n- Initial search for '{vague_term}' returned partial matches\n- Refined search for '{name}' found the full implementation\n- {name} is {summary}\n- Key operations: {', '.join(calls[:4]) if calls else 'internal processing'}\n- Return paths: {'; '.join(returns[:2]) if returns else 'standard'}\n- Source: {chunk['filepath']}<|end|>\n<|finish|>"
+    turn3 = {"content": evidence}
+
+    return make_trace(query, [turn1, turn2, turn3])
+
+
+# ─── Category Registry ────────────────────────────────────────────────────────
+
+CATEGORIES = {
+    "implementation": ("Implementation Deep Dive", gen_implementation_deep_dive),
+    "cross_file": ("Cross-File Data Flow Tracing", gen_cross_file_tracing),
+    "architecture": ("Architecture & Component Mapping", gen_architecture_mapping),
+    "usage": ("Usage Pattern Analysis", gen_usage_analysis),
+    "error_handling": ("Error Handling & Failure Paths", gen_error_handling),
+    "api_contract": ("API Contract & Interface", gen_api_contract),
+    "dependency": ("Dependency Graph Mapping", gen_dependency_graph),
+    "data_structure": ("Data Structure Analysis", gen_data_structure_analysis),
+    "configuration": ("Configuration & Parameters", gen_configuration_params),
+    "performance": ("Performance-Critical Paths", gen_performance_paths),
+    "security": ("Security & Validation", gen_security_validation),
+    "design_patterns": ("Design Pattern Recognition", gen_design_patterns),
+    "debugging": ("Debugging Assistance", gen_debugging_assistance),
+    "comparison": ("Comparison & Tradeoffs", gen_comparison_tradeoffs),
+    "change_impact": ("Change Impact Analysis", gen_change_impact),
+    "complex_multi": ("Complex Multi-Step Analysis", gen_complex_multi_step),
+    "refinement": ("Search Refinement", gen_refinement_search),
+}
+
+# ─── Main Generation Loop ─────────────────────────────────────────────────────
+
+def generate_traces(category: str, count: int, seed: int, chunks: list) -> list[dict]:
+    """Generate traces for a specific category."""
+    if category not in CATEGORIES:
+        print(f"Unknown category: {category}")
+        return []
+
+    cat_name, gen_func = CATEGORIES[category]
+    rng = random.Random(seed)
+    traces = []
+    attempts = 0
+    max_attempts = count * 5  # Allow retries for None returns
+
+    while len(traces) < count and attempts < max_attempts:
+        attempts += 1
+        try:
+            trace = gen_func(chunks, rng)
+            if trace is not None and trace.get("query") and len(trace.get("trace", [])) >= 4:
+                traces.append(trace)
+        except Exception as e:
+            # Silently skip failed generations
+            pass
+
+    return traces
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate high-quality SFT traces")
+    parser.add_argument("--category", type=str, default="all",
+                       help="Category name or 'all'")
+    parser.add_argument("--count", type=int, default=500,
+                       help="Number of traces per category")
+    parser.add_argument("--seed", type=int, default=42,
+                       help="Random seed")
+    parser.add_argument("--output", type=str, default=OUTPUT_DIR,
+                       help="Output directory")
+    args = parser.parse_args()
+
+    # Load chunks
+    print(f"Loading chunks from {CHUNKS_PATH}...")
+    chunks = [json.loads(l) for l in open(CHUNKS_PATH, encoding="utf-8")]
+    print(f"  Loaded {len(chunks):,} chunks")
+
+    # Create output directory
+    os.makedirs(args.output, exist_ok=True)
+
+    # Determine categories
+    if args.category == "all":
+        cats = list(CATEGORIES.keys())
+    else:
+        cats = [args.category]
+
+    total_traces = 0
+    for i, cat in enumerate(cats):
+        cat_name, _ = CATEGORIES[cat]
+        seed = args.seed + i * 1000
+        print(f"\n[{i+1}/{len(cats)}] Generating {args.count} traces for '{cat_name}' (seed={seed})...")
+
+        traces = generate_traces(cat, args.count, seed, chunks)
+
+        output_file = os.path.join(args.output, f"traces_{cat}.jsonl")
+        with open(output_file, "w", encoding="utf-8") as f:
+            for trace in traces:
+                f.write(json.dumps(trace, ensure_ascii=False) + "\n")
+
+        print(f"  Generated {len(traces)} traces → {output_file}")
+        total_traces += len(traces)
+
+    print(f"\nTotal: {total_traces} traces generated across {len(cats)} categories")
+
+
+if __name__ == "__main__":
+    main()

src/generate_traces_v2.py

@@ -0,0 +1,587 @@
+"""
+Diverse search-agent trace generator v2.
+
+Generates 20-50k unique, high-quality search agent traces using real code
+chunks from the deduplicated corpus. Each trace follows the Toast 1-style
+agentic search loop: reasoning -> search -> results -> analysis -> evidence -> finish.
+
+Improvements over generate_traces.py:
+  - Combinatorial query phrasings (hundreds of unique templates)
+  - Smarter evidence extraction (analyzes actual code structure)
+  - More trace types: simple, noise, multi-step, decomposition,
+    not-found, struct inspection, behavior analysis, cross-file,
+    usage tracking, error-handling patterns
+  - Trace-level dedup (no two traces share the same query+search combo)
+  - Varied reasoning text (not random.choice from 5 options)
+  - Realistic evidence packages that extract key facts, not dump code
+
+Output: data/traces/sft_traces_v2.jsonl
+"""
+
+import hashlib
+import json
+import os
+import random
+import re
+
+PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+DATA_DIR = os.path.join(PROJECT_DIR, "data")
+CHUNKS_PATH = os.path.join(DATA_DIR, "chunks", "chunks.jsonl")
+TRACES_DIR = os.path.join(DATA_DIR, "traces")
+TRACES_PATH = os.path.join(TRACES_DIR, "sft_traces_v2.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(2026)
+
+# ─── Query phrasing components (combinatorial) ───────────────────────────────
+QUERY_VERBS = ["What does", "How does", "Where is", "Explain", "Find",
+               "Describe", "What is", "How is", "What are", "Show me",
+               "Locate", "Identify", "Detail"]
+QUERY_ACTIONS = {
+    "function": ["do", "work", "behave", "function", "operate", "process",
+                 "handle", "compute", "return", "accept"],
+    "struct": ["contain", "have", "store", "represent", "look like",
+               "be structured", "be defined"],
+    "class": ["do", "work", "be structured", "be organized", "be designed",
+              "function", "be implemented"],
+    "enum": ["have", "contain", "represent", "be defined", "include"],
+    "macro": ["do", "define", "expand to", "be used for"],
+    "trait": ["define", "require", "provide", "specify"],
+    "interface": ["define", "require", "specify", "declare"],
+    "impl": ["implement", "provide", "define"],
+    "module": ["contain", "provide", "export", "define"],
+    "typedef": ["define", "represent", "be"],
+    "block": ["do", "work", "be used for"],
+}
+QUERY_SUFFIXES = ["?", "in the codebase?", "in this project?",
+                  "? Explain briefly.", "? Give me the details.",
+                  "", "? I need specifics.", "? Summarize the key points."]
+
+# ─── Reasoning components ────────────────────────────────────────────────────
+REASONING_OPENERS = [
+    "I need to find {desc}.", "Looking for {desc}.",
+    "The query is about {desc}.", "This requires locating {desc}.",
+    "Let me search for {desc}.", "I should find {desc} first.",
+    "Starting by searching for {desc}.", "The user wants to know about {desc}.",
+    "I'll look up {desc}.", "First, I need to find {desc}.",
+]
+REASONING_PLANS = [
+    "Let me search using '{term}'.",
+    "I'll query the index for '{term}'.",
+    "Searching for '{term}' should find it.",
+    "Let me issue a search for '{term}'.",
+    "I'll search with the term '{term}'.",
+    "A search for '{term}' should surface the relevant code.",
+    "Let me look for '{term}' in the codebase.",
+]
+ANALYSIS_FOUND = [
+    "Found the relevant code.", "This result contains what I need.",
+    "This is the right code.", "I found the target.",
+    "This matches the query.", "The search returned the relevant implementation.",
+    "This is exactly what I was looking for.",
+]
+ANALYSIS_NOISE = [
+    "The first result is relevant; the others are not directly related.",
+    "Result 1 matches. The remaining results appear to be unrelated code.",
+    "I found the target in the first result. Ignoring the noise.",
+    "The relevant code is in the first result; the rest don't match.",
+    "Only the first result is pertinent to this query.",
+]
+ANALYSIS_REFINE = [
+    "Found the definition. Now let me find where it's used.",
+    "I have the implementation. Let me search for usages.",
+    "Got the definition. Let me get more context with another search.",
+    "Found it. Let me do one more search for completeness.",
+    "Now I need to see how this is called. Let me search for that.",
+]
+NOT_FOUND_REASONING = [
+    "No results found for '{term}'. Let me try a broader search.",
+    "The search returned nothing for '{term}'. Let me try different terms.",
+    "No matches for '{term}'. Let me verify with another query.",
+    "Empty results for '{term}'. Let me rephrase the search.",
+]
+
+
+def generate_query(chunk: dict) -> str:
+    """Generate a diverse query for a chunk."""
+    name = chunk["name"]
+    typ = chunk["type"]
+    if name == "unknown" or name == "block":
+        # Use a key identifier from the code
+        ids = re.findall(r"\b(ngx_\w+|def\s+\w+|fn\s+\w+|struct\s+\w+|class\s+\w+)\b", chunk["code"])
+        if ids:
+            name = ids[0].replace("def ", "").replace("fn ", "").replace("struct ", "").replace("class ", "")
+        else:
+            words = re.findall(r"\b[a-z_]{4,}\b", chunk["code"])
+            name = words[0] if words else "this code"
+
+    verb = random.choice(QUERY_VERBS)
+    actions = QUERY_ACTIONS.get(typ, QUERY_ACTIONS["block"])
+    action = random.choice(actions)
+    suffix = random.choice(QUERY_SUFFIXES)
+
+    if verb in ("Where is", "Locate", "Find", "Show me", "Identify"):
+        return f"{verb} {name} {action}{suffix}".replace(f" {action}", "")
+    return f"{verb} {name} {action}{suffix}"
+
+
+def generate_description(chunk: dict) -> str:
+    name = chunk["name"]
+    typ = chunk["type"]
+    lang = chunk["language"]
+    type_words = {
+        "function": f"the {name} function", "struct": f"the {name} structure",
+        "class": f"the {name} class", "enum": f"the {name} enum",
+        "macro": f"the {name} macro", "trait": f"the {name} trait",
+        "interface": f"the {name} interface", "impl": f"the {name} impl block",
+        "module": f"the {name} module", "typedef": f"the {name} typedef",
+        "block": f"code related to {name}",
+    }
+    return type_words.get(typ, f"the {name} definition")
+
+
+def generate_search_term(chunk: dict) -> str:
+    name = chunk["name"]
+    if name != "unknown" and name != "block":
+        return name
+    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 ", "")
+    words = re.findall(r"\b[a-z_]{4,}\b", code)
+    return words[0] if words else "code"
+
+
+# ─── Evidence extraction (smarter) ───────────────────────────────────────────
+
+def extract_function_info(code: str, name: str) -> dict:
+    info = {"name": name, "params": [], "returns": "", "key_lines": [], "calls": []}
+    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()][:6]
+    ret_match = re.search(r"->\s*([^{]+)", code)
+    if ret_match:
+        info["returns"] = ret_match.group(1).strip()
+    calls = re.findall(r"\b(\w+)\s*\(", code)
+    info["calls"] = [c for c in calls if c != name and c not in ("if", "for", "while", "return", "print")][:5]
+    for line in code.split("\n"):
+        s = line.strip()
+        if not s or s.startswith("//") or s.startswith("#"):
+            continue
+        if any(kw in s for kw in ["return ", "if ", "for ", "while ", "raise ", "throw ", "Err(", "Ok(", "assert", "break", "continue"]):
+            info["key_lines"].append(s)
+    info["key_lines"] = info["key_lines"][:4]
+    return info
+
+
+def extract_struct_info(code: str, name: str) -> dict:
+    fields = []
+    for pat in [re.compile(r"^\s+(\w+)\s+(\w+);", re.M),
+                re.compile(r"^\s+(\w+):\s+([^,;]+)", re.M),
+                re.compile(r"^\s+(self\.)?(\w+)\s*[:=]", re.M),
+                re.compile(r"^\s+(\w+)\s+\w+", re.M)]:
+        for m in pat.findall(code):
+            fn = m[-1] if isinstance(m, tuple) else m
+            if fn and fn not in ("self", "pub", "fn", "def", "class", "struct", "const", "var"):
+                fields.append(fn)
+    return {"name": name, "fields": fields[:12]}
+
+
+def generate_evidence(chunk: dict) -> str:
+    code = chunk["code"]
+    name = chunk["name"]
+    typ = chunk["type"]
+    lang = chunk["language"]
+    filepath = chunk["filepath"]
+    repo = chunk.get("repo", "")
+
+    if typ == "function":
+        info = extract_function_info(code, name)
+        parts = [f"`{name}` is a {lang} function."]
+        if info["params"]:
+            params_str = ", ".join(info["params"])
+            parts.append(f"- Parameters: {params_str}")
+        if info["returns"]:
+            parts.append(f"- Returns: {info['returns']}")
+        if info["key_lines"]:
+            parts.append("- Key behavior:")
+            for line in info["key_lines"][:3]:
+                parts.append(f"  - `{line}`")
+        if info["calls"]:
+            parts.append(f"- Calls: {', '.join(info['calls'])}")
+        parts.append(f"- Source: {filepath}")
+        return "\n".join(parts)
+
+    elif typ == "struct":
+        info = extract_struct_info(code, name)
+        parts = [f"`{name}` is a {lang} structure with fields:"]
+        for field in info["fields"]:
+            parts.append(f"- `{field}`")
+        parts.append(f"- Source: {filepath}")
+        return "\n".join(parts)
+
+    elif typ == "class":
+        methods = re.findall(r"(?:def |fn |func |function )\s*(\w+)", code)
+        parts = [f"`{name}` is a {lang} class."]
+        if methods:
+            parts.append(f"- Methods: {', '.join(methods[:8])}")
+        base = re.search(r"class\s+\w+\s*(?:\(|:\s*|extends\s+|<\s*)(\w+)", code)
+        if base:
+            parts.append(f"- Inherits from: {base.group(1)}")
+        parts.append(f"- Source: {filepath}")
+        return "\n".join(parts)
+
+    elif typ == "enum":
+        variants = re.findall(r"^\s+(\w+)[,\s]*$", code, re.M)
+        parts = [f"`{name}` is a {lang} enum with variants:"]
+        for v in variants[:10]:
+            parts.append(f"- `{v}`")
+        parts.append(f"- Source: {filepath}")
+        return "\n".join(parts)
+
+    elif typ == "macro":
+        parts = [f"`{name}` is a {lang} macro/preprocessor definition:"]
+        for line in code.strip().split("\n")[:5]:
+            parts.append(f"  `{line}`")
+        parts.append(f"- Source: {filepath}")
+        return "\n".join(parts)
+
+    elif typ in ("trait", "interface"):
+        methods = re.findall(r"(?:fn |def |function )\s*(\w+)", code)
+        parts = [f"`{name}` is a {lang} {typ}."]
+        if methods:
+            parts.append(f"- Required methods: {', '.join(methods[:8])}")
+        parts.append(f"- Source: {filepath}")
+        return "\n".join(parts)
+
+    else:
+        parts = [f"Found code related to `{name}` ({lang}):"]
+        for line in code.strip().split("\n")[:8]:
+            parts.append(f"  `{line}`")
+        parts.append(f"- Source: {filepath}")
+        return "\n".join(parts)
+
+
+# ─── Trace builders ──────────────────────────────────────────────────────────
+
+def make_reasoning(desc: str, term: str) -> str:
+    opener = random.choice(REASONING_OPENERS).format(desc=desc)
+    plan = random.choice(REASONING_PLANS).format(term=term)
+    return f"{opener} {plan}"
+
+
+def trace_key(query: str, searches: list[str]) -> str:
+    """Dedup key for a trace."""
+    return hashlib.md5((query + "|" + "|".join(searches)).encode()).hexdigest()
+
+
+def make_trace_simple(chunk: dict) -> dict:
+    query = generate_query(chunk)
+    term = generate_search_term(chunk)
+    desc = generate_description(chunk)
+    evidence = generate_evidence(chunk)
+    reasoning1 = make_reasoning(desc, term)
+    analysis = random.choice(ANALYSIS_FOUND)
+    return {
+        "query": query,
+        "trace": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": query},
+            {"role": "assistant", "content": f"<|reasoning|>{reasoning1}<|end|>\n<|search|>{term}<|end|>"},
+            {"role": "result", "content": chunk["code"]},
+            {"role": "assistant", "content": f"<|reasoning|>{analysis} {desc.replace('the ', 'This is ', 1)}.<|end|>\n<|evidence|>{evidence}<|end|>\n<|finish|>"},
+        ],
+    }
+
+
+def make_trace_with_noise(chunk: dict, noise_chunks: list) -> dict:
+    query = generate_query(chunk)
+    term = generate_search_term(chunk)
+    desc = generate_description(chunk)
+    evidence = generate_evidence(chunk)
+    reasoning1 = make_reasoning(desc, term)
+    n_noise = random.choice([1, 2])
+    noise = random.sample(noise_chunks, min(n_noise, len(noise_chunks)))
+    analysis = random.choice(ANALYSIS_NOISE)
+    trace = [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {"role": "user", "content": query},
+        {"role": "assistant", "content": f"<|reasoning|>{reasoning1}<|end|>\n<|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}<|end|>\n<|evidence|>{evidence}<|end|>\n<|finish|>"})
+    return {"query": query, "trace": trace}
+
+
+def make_trace_multi_step(chunk: dict, related: dict) -> dict:
+    name = chunk["name"]
+    rname = related["name"]
+    query = f"What does {name} do and where is it used?"
+    term1 = generate_search_term(chunk)
+    term2 = generate_search_term(related)
+    desc = generate_description(chunk)
+    evidence1 = generate_evidence(chunk)
+    evidence2 = generate_evidence(related)
+    reasoning1 = make_reasoning(desc, term1)
+    reasoning2 = random.choice(ANALYSIS_REFINE)
+    return {
+        "query": query,
+        "trace": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": query},
+            {"role": "assistant", "content": f"<|reasoning|>{reasoning1}<|end|>\n<|search|>{term1}<|end|>"},
+            {"role": "result", "content": chunk["code"]},
+            {"role": "assistant", "content": f"<|reasoning|>Found {desc}. {reasoning2}<|end|>\n<|search|>{term2}<|end|>"},
+            {"role": "result", "content": related["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_decomposition(chunk: dict, related: dict) -> dict:
+    name = chunk["name"]
+    rname = related["name"]
+    query = f"How does {name} work and how does it relate to {rname}?"
+    term1 = generate_search_term(chunk)
+    term2 = generate_search_term(related)
+    evidence1 = generate_evidence(chunk)
+    evidence2 = generate_evidence(related)
+    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 {rname}. Let me search for {name} first.<|end|>\n<|search|>{term1}<|end|>"},
+            {"role": "result", "content": chunk["code"]},
+            {"role": "assistant", "content": f"<|reasoning|>Found {name}. Now let me search for {rname} to understand the relationship.<|end|>\n<|search|>{term2}<|end|>"},
+            {"role": "result", "content": related["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_not_found() -> dict:
+    fake_names = ["database_pool", "auth_middleware", "cache_manager", "session_store",
+                  "websocket_handler", "graphql_resolver", "redis_client", "email_service",
+                  "rate_limiter", "circuit_breaker", "event_bus", "message_queue",
+                  "load_balancer", "api_gateway", "service_mesh", "config_center"]
+    name = random.choice(fake_names)
+    query = f"Where is the {name} implemented?"
+    search1 = name
+    search2 = name.replace("_", " ")
+    reasoning1 = make_reasoning(f"the {name}", search1)
+    reasoning2 = random.choice(NOT_FOUND_REASONING).format(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_struct_inspection(chunk: dict) -> dict:
+    if chunk["type"] != "struct":
+        return make_trace_simple(chunk)
+    name = chunk["name"]
+    query = f"What fields does {name} have?"
+    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|>{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(chunk: dict) -> dict:
+    if chunk["type"] != "function":
+        return make_trace_simple(chunk)
+    name = chunk["name"]
+    query = f"What does {name} do? Explain its behavior."
+    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|>{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|>"},
+        ],
+    }
+
+
+def make_trace_cross_file(chunk: dict, related: dict) -> dict:
+    """Cross-file: definition in one file, usage in another."""
+    name = chunk["name"]
+    rname = related["name"]
+    query = f"How is {name} used across the codebase?"
+    term1 = generate_search_term(chunk)
+    term2 = generate_search_term(related)
+    evidence1 = generate_evidence(chunk)
+    evidence2 = generate_evidence(related)
+    return {
+        "query": query,
+        "trace": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": query},
+            {"role": "assistant", "content": f"<|reasoning|>I need to find the definition of {name} and then trace its usages across files.<|end|>\n<|search|>{term1}<|end|>"},
+            {"role": "result", "content": chunk["code"]},
+            {"role": "assistant", "content": f"<|reasoning|>Found the definition in {chunk['filepath']}. Now let me search for where it's called or referenced.<|end|>\n<|search|>{term2}<|end|>"},
+            {"role": "result", "content": related["code"]},
+            {"role": "assistant", "content": f"<|reasoning|>Found a usage in {related['filepath']}. I now have both the definition and a call site.<|end|>\n<|evidence|>Definition ({chunk['filepath']}):\n{evidence1}\n\nUsage ({related['filepath']}):\n{evidence2}<|end|>\n<|finish|>"},
+        ],
+    }
+
+
+# ─── Main ────────────────────────────────────────────────────────────────────
+
+def main():
+    os.makedirs(TRACES_DIR, exist_ok=True)
+
+    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")
+
+    good_chunks = [c for c in chunks if c["name"] != "unknown" and c["name"] != "block" and 50 < len(c["code"]) < 4000]
+    print(f"  Named chunks (good): {len(good_chunks):,}")
+
+    by_lang = {}
+    for c in good_chunks:
+        by_lang.setdefault(c["language"], []).append(c)
+
+    traces = []
+    seen_keys = set()
+
+    def add_trace(trace):
+        searches = [m["content"] for m in trace["trace"] if m["role"] == "assistant" and "<|search|>" in m["content"]]
+        # Extract search terms
+        search_terms = re.findall(r"<\|search\|>(.*?)<\|end\|>", "".join(m["content"] for m in trace["trace"] if m["role"] == "assistant"))
+        key = trace_key(trace["query"], search_terms)
+        if key in seen_keys:
+            return False
+        seen_keys.add(key)
+        traces.append(trace)
+        return True
+
+    def get_related(chunk):
+        lang_chunks = by_lang.get(chunk["language"], good_chunks)
+        pool = [c for c in lang_chunks if c["id"] != chunk["id"]]
+        return random.choice(pool) if pool else chunk
+
+    # Type 1: Simple lookup (25%)
+    n_simple = min(8000, len(good_chunks))
+    for chunk in random.sample(good_chunks, n_simple):
+        add_trace(make_trace_simple(chunk))
+    print(f"  Simple lookup: {len(traces)}")
+
+    # Type 2: With noise (20%)
+    n_noise = min(6000, len(good_chunks))
+    for chunk in random.sample(good_chunks, n_noise):
+        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:
+            add_trace(make_trace_with_noise(chunk, noise_pool))
+    print(f"  With noise: {len(traces)}")
+
+    # Type 3: Multi-step (15%)
+    n_multi = min(4000, len(good_chunks) // 2)
+    for chunk in random.sample(good_chunks, n_multi):
+        add_trace(make_trace_multi_step(chunk, get_related(chunk)))
+    print(f"  Multi-step: {len(traces)}")
+
+    # Type 4: Decomposition (10%)
+    n_decomp = min(3000, len(good_chunks) // 3)
+    for chunk in random.sample(good_chunks, n_decomp):
+        add_trace(make_trace_decomposition(chunk, get_related(chunk)))
+    print(f"  Decomposition: {len(traces)}")
+
+    # Type 5: Not found (5%)
+    for _ in range(1500):
+        add_trace(make_trace_not_found())
+    print(f"  Not found: {len(traces)}")
+
+    # Type 6: Struct inspection (10%)
+    struct_chunks = [c for c in good_chunks if c["type"] == "struct"]
+    for chunk in random.sample(struct_chunks, min(2000, len(struct_chunks))):
+        add_trace(make_trace_struct_inspection(chunk))
+    print(f"  Struct inspection: {len(traces)}")
+
+    # Type 7: Behavior analysis (10%)
+    func_chunks = [c for c in good_chunks if c["type"] == "function"]
+    for chunk in random.sample(func_chunks, min(2000, len(func_chunks))):
+        add_trace(make_trace_behavior(chunk))
+    print(f"  Behavior analysis: {len(traces)}")
+
+    # Type 8: Cross-file (5%)
+    n_cross = min(1500, len(good_chunks) // 4)
+    for chunk in random.sample(good_chunks, n_cross):
+        add_trace(make_trace_cross_file(chunk, get_related(chunk)))
+    print(f"  Cross-file: {len(traces)}")
+
+    random.shuffle(traces)
+
+    # Write
+    print(f"\nTotal unique traces: {len(traces):,}")
+    total_size = 0
+    with open(TRACES_PATH, "w", encoding="utf-8") as f:
+        for trace in traces:
+            line = json.dumps(trace, ensure_ascii=False)
+            f.write(line + "\n")
+            total_size += len(line)
+    print(f"Traces written to {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"]:
+        content = msg["content"]
+        if len(content) > 200:
+            content = content[:200] + "..."
+        print(f"\n[{msg['role']}]")
+        print(content)
+
+    # Stats
+    stats = {
+        "total_traces": len(traces),
+        "file_size_mb": total_size / 1e6,
+        "unique_queries": len(set(t["query"] for t in traces)),
+    }
+    stats_path = os.path.join(TRACES_DIR, "trace_stats.json")
+    with open(stats_path, "w") as f:
+        json.dump(stats, f, indent=2)
+    print(f"\nStats: {stats}")
+
+
+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/quality_filter.py

@@ -0,0 +1,291 @@
+"""
+Quality filtering, contamination control, and token-density validation.
+
+Takes the deduplicated documents and applies final quality gates:
+
+  QUALITY FILTERS:
+    - Real source code (language signatures at line starts)
+    - Balanced delimiters (braces/parens/brackets roughly match)
+    - Reasonable length (200–100k chars)
+    - Low repetition (unique line ratio)
+    - Clean ASCII (low non-ASCII ratio)
+    - High code-to-prose ratio
+    - Has structure (function/class/struct definitions)
+
+  CONTAMINATION CONTROL:
+    - No test files (already filtered in download, double-check here)
+    - No auto-generated code markers
+    - No license-only files
+    - No binary/garbage content
+    - No files with extremely high repetition (copy-paste blocks)
+
+  TOKEN DENSITY:
+    - Every kept document must be "dense in tokens" — meaning the
+      content tokenizes to a meaningful number of tokens relative to
+      its character length (no whitespace-padding, no huge comment blocks).
+    - Reports token density stats using the project tokenizer.
+
+Input:  data/dedup/documents_dedup.jsonl
+Output: data/quality/documents_quality.jsonl  +  quality_stats.json
+"""
+
+import json
+import os
+import re
+import sys
+
+PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+DATA_DIR = os.path.join(PROJECT_DIR, "data")
+DEDUP_DOCS_PATH = os.path.join(DATA_DIR, "dedup", "documents_dedup.jsonl")
+QUALITY_DIR = os.path.join(DATA_DIR, "quality")
+QUALITY_DOCS_PATH = os.path.join(QUALITY_DIR, "documents_quality.jsonl")
+STATS_PATH = os.path.join(QUALITY_DIR, "quality_stats.json")
+
+# ─── Code structure patterns (must be at line start) ─────────────────────────
+CODE_SIGNATURES = {
+    "python":   re.compile(r"^(def |class |import |from \S+ import |if __name__|@|async def )", re.M),
+    "js_ts":    re.compile(r"^(function |const |let |var |class |export |import |async function|interface |type \w+ =|enum )", re.M),
+    "rust":     re.compile(r"^(fn |pub fn |impl |struct |enum |trait |mod |use |pub struct|pub enum|macro_rules!|pub trait)", re.M),
+    "go":       re.compile(r"^(func |package |import |type \w+ struct|var |const )", re.M),
+    "c":        re.compile(r"^(#include|#define|#ifndef|#ifdef|#if |#endif|typedef |struct \w+|static |void |int |char )", re.M),
+    "cpp":      re.compile(r"^(#include|#define|#ifndef|#ifdef|#if |#endif|template |class \w+|namespace |struct \w+|void |int )", re.M),
+    "java":     re.compile(r"^(public |private |protected |class \w+|import |package |interface )", re.M),
+    "csharp":   re.compile(r"^(public |private |protected |internal |class \w+|using |namespace |interface )", re.M),
+    "ruby":     re.compile(r"^(def |class |module |require |require_relative |attr_|include )", re.M),
+}
+
+# ─── Contamination / exclusion patterns ──────────────────────────────────────
+AUTO_GEN = re.compile(
+    r"(?:auto[- ]generated|do not edit|generated by|code generated|"
+    r"DO NOT MODIFY|@generated|automatically generated|"
+    r"this file was generated)",
+    re.IGNORECASE,
+)
+LICENSE_ONLY = re.compile(r"^(?:/\*|//|#)\s*(?:copyright|licensed|mit license|apache license|bsd license|gnu|gpl)", re.I)
+PROSE_LINE = re.compile(r"^[A-Z][a-z]+ .* [a-z]+\.$", re.M)
+
+# ─── Quality scoring ─────────────────────────────────────────────────────────
+
+def detect_language_signatures(doc: dict) -> int:
+    """Count code-structure keywords at line starts. Returns count."""
+    content = doc["content"]
+    lang = doc["language"]
+    pat = CODE_SIGNATURES.get(lang)
+    if pat is None:
+        # Try all
+        return sum(len(p.findall(content)) for p in CODE_SIGNATURES.values())
+    return len(pat.findall(content))
+
+
+def score_document(doc: dict) -> tuple[float, str | None]:
+    """Score a document 0.0–1.0 on quality. Returns (score, reject_reason)."""
+    content = doc["content"]
+    lines = content.split("\n")
+    n_lines = len(lines)
+    length = len(content)
+
+    # ─── Hard rejects (contamination) ──────────────────────────────────────
+    if length < 200:
+        return 0.0, "too_short"
+    if length > 200_000:
+        return 0.0, "too_long"
+    if n_lines < 5:
+        return 0.0, "too_few_lines"
+
+    # Auto-generated code
+    if AUTO_GEN.search(content[:2000]):
+        return 0.0, "auto_generated"
+
+    # License-only files (first 10 lines are all license comments)
+    first_lines = "\n".join(lines[:10])
+    if LICENSE_ONLY.search(first_lines) and n_lines < 30:
+        return 0.0, "license_only"
+
+    # High non-ASCII (garbage/encoding issues)
+    non_ascii = sum(1 for c in content if ord(c) > 127)
+    if length > 0 and non_ascii / length > 0.03:
+        return 0.0, "high_non_ascii"
+
+    # Extremely high repetition (copy-paste blocks)
+    unique_lines = len(set(lines))
+    unique_ratio = unique_lines / max(n_lines, 1)
+    if unique_ratio < 0.20:
+        return 0.0, "high_repetition"
+
+    # ─── Soft scoring ──────────────────────────────────────────────────────
+    score = 0.0
+
+    # Base: passes hard filters
+    score += 0.15
+
+    # Length quality (sweet spot: 500–30000 chars)
+    if 500 <= length <= 30000:
+        score += 0.15
+    elif 200 <= length <= 80000:
+        score += 0.08
+
+    # Delimiter balance
+    braces = content.count("{") - content.count("}")
+    parens = content.count("(") - content.count(")")
+    brackets = content.count("[") - content.count("]")
+    total_delims = content.count("{") + content.count("(") + content.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 = detect_language_signatures(doc)
+    struct_density = min(struct_count / max(n_lines, 1) * 10, 1.0)
+    score += 0.15 * struct_density
+
+    # Low repetition (unique line ratio)
+    if unique_ratio > 0.7:
+        score += 0.10
+    elif unique_ratio > 0.5:
+        score += 0.05
+    else:
+        score -= 0.05
+
+    # Clean ASCII
+    if length > 0 and non_ascii / length < 0.005:
+        score += 0.05
+
+    # Indentation quality (indented lines indicate real code structure)
+    indented = sum(1 for l in lines if l.startswith("    ") or l.startswith("\t"))
+    if indented > 0 and indented / max(n_lines, 1) > 0.15:
+        score += 0.05
+
+    # Penalize high prose ratio (documentation, not code)
+    prose_lines = len(PROSE_LINE.findall(content))
+    prose_ratio = prose_lines / max(n_lines, 1)
+    if prose_ratio > 0.20:
+        score -= 0.15
+
+    # Comment density (sweet spot: 3–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.03 <= comment_ratio <= 0.30:
+        score += 0.10
+    elif comment_ratio > 0.50:
+        score -= 0.10  # too many comments = doc, not code
+
+    return min(max(score, 0.0), 1.0), None
+
+
+def check_token_density(doc: dict) -> tuple[bool, float]:
+    """Check that a document is dense in tokens (not whitespace-padded).
+
+    Returns (passes, chars_per_token_ratio).
+    A good code document should have ~2.5-5 chars per token.
+    If the ratio is very high (>15), it's likely whitespace/garbage.
+    If very low (<1.5), it may be all symbols.
+    """
+    content = doc["content"]
+    # Rough estimate: count non-whitespace characters as a proxy
+    # Real tokenization happens with the tokenizer, but this is a fast filter
+    non_ws = len(content) - content.count(" ") - content.count("\n") - content.count("\t") - content.count("\r")
+    # Approximate token count: split on whitespace + common code delimiters
+    approx_tokens = len(re.findall(r"\w+|[^\w\s]", content))
+    if approx_tokens == 0:
+        return False, 0.0
+    chars_per_token = len(content) / approx_tokens
+    # Good density: 2.0 - 8.0 chars per token
+    passes = 2.0 <= chars_per_token <= 12.0
+    return passes, chars_per_token
+
+
+def main():
+    os.makedirs(QUALITY_DIR, exist_ok=True)
+
+    print(f"Loading deduped documents from {DEDUP_DOCS_PATH}...")
+    docs = []
+    with open(DEDUP_DOCS_PATH, "r", encoding="utf-8") as f:
+        for line in f:
+            docs.append(json.loads(line))
+    print(f"  Loaded {len(docs):,} documents")
+
+    kept = []
+    reject_reasons = {}
+    scores = []
+    token_densities = []
+
+    for i, doc in enumerate(docs):
+        score, reject = score_document(doc)
+
+        if reject:
+            reject_reasons[reject] = reject_reasons.get(reject, 0) + 1
+            continue
+
+        if score < 0.40:
+            reject_reasons["low_score"] = reject_reasons.get("low_score", 0) + 1
+            continue
+
+        # Token density check
+        dense, cpt = check_token_density(doc)
+        token_densities.append(cpt)
+        if not dense:
+            reject_reasons["low_token_density"] = reject_reasons.get("low_token_density", 0) + 1
+            continue
+
+        doc = dict(doc)
+        doc["quality_score"] = round(score, 4)
+        kept.append(doc)
+        scores.append(score)
+
+        if (i + 1) % 5000 == 0:
+            print(f"  Processed {i+1}/{len(docs)} | kept {len(kept)} | "
+                  f"rejected {i+1 - len(kept)}")
+
+    # Stats
+    final_size = sum(len(d["content"]) for d in kept)
+    avg_score = sum(scores) / len(scores) if scores else 0
+    avg_cpt = sum(token_densities) / len(token_densities) if token_densities else 0
+
+    lang_counts = {}
+    for d in kept:
+        lang_counts[d["language"]] = lang_counts.get(d["language"], 0) + 1
+
+    stats = {
+        "input": len(docs),
+        "kept": len(kept),
+        "rejected": len(docs) - len(kept),
+        "reject_reasons": reject_reasons,
+        "final_size_mb": final_size / 1e6,
+        "avg_quality_score": round(avg_score, 4),
+        "avg_chars_per_token": round(avg_cpt, 2),
+        "language_distribution": lang_counts,
+        "min_score_threshold": 0.40,
+    }
+
+    print("\n" + "=" * 60)
+    print("QUALITY FILTERING COMPLETE")
+    print("=" * 60)
+    print(f"  Input:    {len(docs):,}")
+    print(f"  Kept:     {len(kept):,}")
+    print(f"  Rejected: {len(docs) - len(kept):,}")
+    print(f"  Reject reasons: {reject_reasons}")
+    print(f"  Final size: {final_size / 1e6:.1f} MB")
+    print(f"  Avg quality score: {avg_score:.3f}")
+    print(f"  Avg chars/token: {avg_cpt:.2f}")
+    print(f"  Languages: {lang_counts}")
+
+    # Write quality-filtered documents
+    with open(QUALITY_DOCS_PATH, "w", encoding="utf-8") as f:
+        for doc in kept:
+            f.write(json.dumps(doc, ensure_ascii=False) + "\n")
+    print(f"\nQuality documents written to {QUALITY_DOCS_PATH}")
+
+    with open(STATS_PATH, "w") as f:
+        json.dump(stats, f, indent=2)
+    print(f"Stats written to {STATS_PATH}")
+
+
+if __name__ == "__main__":
+    main()

src/search_agent.py

@@ -0,0 +1,385 @@
+"""
+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 = 512,
+    ):
+        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.
+
+        Stops at <|finish|> (end of agent turn) or after a <|search|>query<|end|>
+        sequence (so the harness can inject results before continuing).
+        """
+        input_len = input_ids.size(1)
+        recent_tokens = []
+        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, :]
+
+                # Greedy decoding — more reliable for a 500M model
+                next_token = next_logits.argmax(dim=-1, keepdim=True)
+                tok_id = next_token.item()
+                input_ids = torch.cat([input_ids, next_token], dim=1)
+                recent_tokens.append(tok_id)
+
+                # Stop on <|finish|>
+                if tok_id == self.finish_id:
+                    break
+
+                # Stop after <|search|>...<|end|> — harness needs to inject results
+                # Detect: we've seen <|search|> (32003) and then <|end|> (32008)
+                if tok_id == self.end_id and self.search_id in recent_tokens[-20:]:
+                    # Check the search was recent (within last 20 tokens)
+                    search_pos = None
+                    for i in range(len(recent_tokens) - 1, max(0, len(recent_tokens) - 20), -1):
+                        if recent_tokens[i] == self.search_id:
+                            search_pos = i
+                            break
+                    if search_pos is not None:
+                        # Only stop if <|end|> comes right after the search query
+                        # (not an <|end|> from a <|reasoning|> section that happened
+                        # to be within 20 tokens)
+                        # Check: is there a <|reasoning|> after the <|search|>?
+                        after_search = recent_tokens[search_pos+1:]
+                        if self.reasoning_id not in after_search and self.evidence_id not in after_search:
+                            break
+
+        # Decode only the newly generated tokens
+        generated_ids = input_ids[0, input_len:].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})
+
+            # Parse search query FIRST (before checking finish)
+            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})
+
+                # Also check if the response contains evidence (model generated
+                # everything in one go but we intercepted the search)
+                if self._has_finish(response):
+                    evidence = self._parse_evidence(response)
+                    if evidence:
+                        print(f"\n[EVIDENCE]: {evidence}")
+                        break
+            elif self._has_finish(response):
+                # No search but has finish — extract evidence
+                evidence = self._parse_evidence(response)
+                print(f"\n[EVIDENCE]: {evidence}")
+                break
+            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_modal.py

@@ -0,0 +1,840 @@
+"""
+Modal training script for Retriever500M.
+
+Runs pretraining and/or SFT on Modal cloud GPUs instead of the local laptop.
+Supports resuming from the existing bf16 checkpoint.
+
+Usage:
+  # Step 1: Seed the volume with the initial checkpoint
+  modal run src/train_modal.py --command seed
+
+  # Step 2: Run pretraining (continues from latest checkpoint)
+  modal run src/train_modal.py --command pretrain --steps 1000
+
+  # Step 3: Run SFT (starts from pretrained checkpoint)
+  modal run src/train_modal.py --command sft --steps 500
+
+  # Step 4: Download checkpoints back to local
+  modal run src/train_modal.py --command download
+
+  # Or do everything in one shot:
+  modal run src/train_modal.py --command pipeline --pretrain-steps 1000 --sft-steps 500
+"""
+
+import os
+import sys
+import json
+import time
+import argparse
+from dataclasses import asdict
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+
+import modal
+
+# ─── Modal Setup ─────────────────────────────────────────────────────────────
+
+APP_NAME = "retriever500m"
+VOLUME_NAME = "retriever500m-data"
+
+PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+SRC_DIR = os.path.join(PROJECT_DIR, "src")
+DATA_DIR = os.path.join(PROJECT_DIR, "data")
+TOKENIZER_DIR = os.path.join(PROJECT_DIR, "tokenizer")
+CHECKPOINT_DIR = os.path.join(PROJECT_DIR, "checkpoints")
+
+volume = modal.Volume.from_name(VOLUME_NAME, create_if_missing=True)
+
+# Image: PyTorch with CUDA, plus bitsandbytes for 8-bit optimizer.
+# Data, tokenizer, src, and initial checkpoint are baked into the image.
+# The volume is used only for saving new checkpoints and logs.
+image = (
+    modal.Image.from_registry(
+        "pytorch/pytorch:2.4.1-cuda12.1-cudnn9-runtime",
+        add_python="3.11",
+    )
+    .pip_install("tokenizers>=0.15,<0.21", "tqdm", "numpy")
+    .pip_install("bitsandbytes>=0.43,<0.45")
+    .add_local_dir(SRC_DIR, "/root/src")
+    .add_local_dir(TOKENIZER_DIR, "/root/tokenizer")
+    .add_local_dir(DATA_DIR, "/root/data", ignore=["raw_large/", "raw/", "dedup/"])
+    .add_local_dir(CHECKPOINT_DIR, "/root/seed_checkpoints")
+)
+
+app = modal.App(APP_NAME, image=image)
+
+# Remote paths inside the container.
+# Data/tokenizer/src are baked into the image.
+# The volume is mounted at /root/vol for checkpoints and logs.
+VOL_MOUNT = "/root/vol"
+REMOTE_DATA = "/root/data"              # baked into image
+REMOTE_CKPT = "/root/vol/checkpoints"   # on volume (for new checkpoints)
+REMOTE_LOGS = "/root/vol/logs"          # on volume
+REMOTE_TOKENIZER = "/root/tokenizer"    # baked into image
+REMOTE_SRC = "/root/src"                # baked into image
+REMOTE_SEED_CKPT = "/root/seed_checkpoints"  # baked into image (initial checkpoint)
+
+
+# ─── Seed command ────────────────────────────────────────────────────────────
+
+@app.function(volumes={VOL_MOUNT: volume})
+def seed():
+    """Copy the initial checkpoint from the image to the volume."""
+    import shutil
+
+    os.makedirs(REMOTE_CKPT, exist_ok=True)
+    os.makedirs(REMOTE_LOGS, exist_ok=True)
+
+    # Copy seed checkpoints from image to volume
+    if os.path.exists(REMOTE_SEED_CKPT):
+        for fname in os.listdir(REMOTE_SEED_CKPT):
+            src = os.path.join(REMOTE_SEED_CKPT, fname)
+            dst = os.path.join(REMOTE_CKPT, fname)
+            if os.path.isfile(src) and not os.path.exists(dst):
+                shutil.copy2(src, dst)
+                print(f"  Seeded {fname} ({os.path.getsize(src) / 1e6:.1f} MB)")
+            elif os.path.exists(dst):
+                print(f"  SKIP {fname} (already on volume)")
+
+    volume.commit()
+    print("Seed complete.")
+
+
+# ─── Download command ────────────────────────────────────────────────────────
+
+@app.function(volumes={VOL_MOUNT: volume})
+def download():
+    """Download checkpoints and logs from the Modal volume to local."""
+    import shutil
+
+    os.makedirs(CHECKPOINT_DIR, exist_ok=True)
+    os.makedirs(os.path.join(PROJECT_DIR, "logs"), exist_ok=True)
+
+    # Download checkpoints
+    for fname in os.listdir(REMOTE_CKPT):
+        src = os.path.join(REMOTE_CKPT, fname)
+        if os.path.isfile(src):
+            dst = os.path.join(CHECKPOINT_DIR, fname)
+            shutil.copy2(src, dst)
+            print(f"  Downloaded {fname} ({os.path.getsize(src) / 1e6:.1f} MB)")
+
+    # Download logs
+    for fname in os.listdir(REMOTE_LOGS):
+        src = os.path.join(REMOTE_LOGS, fname)
+        if os.path.isfile(src):
+            dst = os.path.join(PROJECT_DIR, "logs", fname)
+            shutil.copy2(src, dst)
+            print(f"  Downloaded log {fname}")
+
+    print("Download complete.")
+
+
+# ─── Training logic (shared) ─────────────────────────────────────────────────
+
+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 setup_optimizer(model, lr, use_8bit=True):
+    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
+
+
+# ─── Pretraining ─────────────────────────────────────────────────────────────
+
+def load_and_tokenize(corpus_path, tokenizer):
+    """Load corpus, tokenize, return 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")
+
+    chunk_size = 1_000_000
+    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, batch_size, seq_len, device):
+    max_start = len(tokens) - seq_len - 1
+    indices = np.random.randint(0, max_start, size=batch_size)
+    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
+
+
+def run_pretrain(steps, batch_size, grad_accum, seq_len, lr, warmup,
+                 save_every, log_every, corpus, resume, use_8bit_adam):
+    sys.path.insert(0, REMOTE_SRC)
+    from model import ModelConfig, Retriever500M
+    from tokenizers import Tokenizer
+
+    device = torch.device("cuda")
+    print(f"Device: {device}")
+    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(REMOTE_CKPT, exist_ok=True)
+    os.makedirs(REMOTE_LOGS, exist_ok=True)
+
+    # Tokenizer
+    tokenizer_path = os.path.join(REMOTE_TOKENIZER, "tokenizer.json")
+    tokenizer = Tokenizer.from_file(tokenizer_path)
+    vocab_size = tokenizer.get_vocab_size()
+    print(f"Vocab size: {vocab_size}")
+
+    # Data
+    if corpus == "curated":
+        corpus_path = os.path.join(REMOTE_DATA, "corpus_curated.txt")
+    elif corpus == "default":
+        corpus_path = os.path.join(REMOTE_DATA, "corpus.txt")
+    else:
+        corpus_path = 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=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 = setup_optimizer(model, lr, use_8bit=use_8bit_adam)
+
+    # Resume
+    start_step = 0
+    best_loss = float("inf")
+    accum_loss = 0.0
+    prev_log_steps = []
+
+    if resume:
+        resume_path = os.path.join(REMOTE_CKPT, "latest.pt")
+        if not os.path.exists(resume_path):
+            resume_path = os.path.join(REMOTE_CKPT, "latest_bf16.pt")
+        if not os.path.exists(resume_path):
+            # Fall back to seed checkpoint in image
+            resume_path = os.path.join(REMOTE_SEED_CKPT, "latest_bf16.pt")
+        if os.path.exists(resume_path):
+            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"])
+            start_step = int(ckpt.get("step", 0))
+            best_loss = float(ckpt.get("loss", float("inf")))
+            accum_loss = best_loss
+            print(f"  Resuming at step {start_step} (best_loss={best_loss:.4f})")
+
+            prev_log_path = os.path.join(REMOTE_LOGS, "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:
+                    pass
+        else:
+            print("WARNING: No checkpoint found, starting from scratch!")
+
+    # Training loop
+    effective_batch = batch_size * grad_accum
+    max_steps_total = start_step + steps
+    print(f"\nTraining configuration:")
+    print(f"  Micro batch size:  {batch_size}")
+    print(f"  Gradient accum:    {grad_accum}")
+    print(f"  Effective batch:   {effective_batch}")
+    print(f"  Sequence length:   {seq_len}")
+    print(f"  Learning rate:     {lr}")
+    print(f"  Steps this run:    {steps}")
+    print(f"  Start step:        {start_step}")
+    print(f"  Target step:       {max_steps_total}")
+    print()
+
+    log = {
+        "config": asdict(config),
+        "train_args": {"steps": steps, "batch_size": batch_size, "grad_accum": grad_accum,
+                       "seq_len": seq_len, "lr": lr, "warmup": warmup,
+                       "save_every": save_every, "log_every": log_every,
+                       "corpus": corpus, "resume": resume, "use_8bit_adam": use_8bit_adam},
+        "total_params": total_params,
+        "steps": list(prev_log_steps),
+    }
+
+    model.train()
+    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:
+        lr_now = get_lr(step, warmup, max_steps_total, lr, lr * 0.1)
+        for pg in optimizer.param_groups:
+            pg["lr"] = lr_now
+
+        optimizer.zero_grad(set_to_none=True)
+
+        total_loss = 0.0
+        for _ in range(grad_accum):
+            input_ids, targets = get_batch(tokens, batch_size, seq_len, device)
+            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+                out = model(input_ids, targets=targets, use_checkpoint=False)
+                loss = out["loss"] / 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 % 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
+
+            log_entry = {
+                "step": step, "loss": avg_loss, "ema_loss": accum_loss,
+                "lr": lr_now, "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_now:.2e}", "vram": f"{vram_used:.1f}G",
+            })
+
+        if (step + 1) % save_every == 0 or step == max_steps_total - 1:
+            ckpt_path = os.path.join(REMOTE_CKPT, 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}")
+
+            latest_path = os.path.join(REMOTE_CKPT, "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(REMOTE_CKPT, "best.pt")
+                torch.save({
+                    "model_state_dict": model.state_dict(),
+                    "config": asdict(config),
+                    "step": step + 1, "loss": accum_loss,
+                }, best_path)
+
+            # Save log
+            log_path = os.path.join(REMOTE_LOGS, "training_log.json")
+            with open(log_path, "w") as f:
+                json.dump(log, f, indent=2)
+
+            # Commit volume so checkpoints persist
+            volume.commit()
+
+        if step % 50 == 0:
+            torch.cuda.reset_peak_memory_stats()
+
+    # Final log save
+    log_path = os.path.join(REMOTE_LOGS, "training_log.json")
+    with open(log_path, "w") as f:
+        json.dump(log, f, indent=2)
+    volume.commit()
+
+    total_time = time.time() - start_time
+    print(f"\nPretraining 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}")
+    return accum_loss
+
+
+# ─── SFT ─────────────────────────────────────────────────────────────────────
+
+# Special token IDs
+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():
+    global SYSTEM_ID, USER_ID, ASSISTANT_ID, SEARCH_ID, RESULT_ID
+    global EVIDENCE_ID, REASONING_ID, FINISH_ID, END_ID
+    path = os.path.join(REMOTE_TOKENIZER, "special_tokens.json")
+    if os.path.exists(path):
+        with open(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)
+
+
+def format_trace_to_tokens(trace, tokenizer, max_seq_len=768):
+    messages = trace["trace"]
+    all_tokens = []
+    loss_mask = []
+
+    for msg in messages:
+        role = msg["role"]
+        content = msg["content"]
+
+        if role == "system":
+            tokens = [SYSTEM_ID] + tokenizer.encode(content).ids + [END_ID]
+            all_tokens.extend(tokens)
+            loss_mask.extend([0] * len(tokens))
+        elif role == "user":
+            tokens = [USER_ID] + tokenizer.encode(content).ids + [END_ID]
+            all_tokens.extend(tokens)
+            loss_mask.extend([0] * len(tokens))
+        elif role == "assistant":
+            tokens = [ASSISTANT_ID] + tokenizer.encode(content).ids + [END_ID]
+            all_tokens.extend(tokens)
+            loss_mask.extend([1] * len(tokens))
+        elif role == "result":
+            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))
+
+    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, tokenizer, max_seq_len=768):
+    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:
+                dataset.append((ids, mask))
+    print(f"  Loaded {len(dataset):,} traces")
+    return dataset
+
+
+def get_sft_batch(dataset, batch_size, seq_len, device):
+    indices = np.random.randint(0, len(dataset), size=batch_size)
+    input_ids_list = []
+    loss_mask_list = []
+
+    for idx in indices:
+        ids, mask = dataset[idx]
+        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 = torch.cat([input_ids[:, 1:], torch.zeros_like(input_ids[:, :1])], dim=1)
+    return input_ids, targets, loss_mask
+
+
+def run_sft(steps, batch_size, grad_accum, seq_len, lr, warmup,
+            save_every, log_every, use_8bit_adam):
+    from torch import nn
+    sys.path.insert(0, REMOTE_SRC)
+    from model import ModelConfig, Retriever500M
+    from tokenizers import Tokenizer
+
+    load_special_tokens()
+    device = torch.device("cuda")
+    print(f"Device: {device}")
+    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(REMOTE_CKPT, exist_ok=True)
+    os.makedirs(REMOTE_LOGS, exist_ok=True)
+
+    # Tokenizer (agent tokenizer with special tokens)
+    tokenizer_path = os.path.join(REMOTE_TOKENIZER, "tokenizer_agent.json")
+    tokenizer = Tokenizer.from_file(tokenizer_path)
+    vocab_size = tokenizer.get_vocab_size()
+    print(f"Vocab size: {vocab_size}")
+
+    # Data
+    traces_path = os.path.join(REMOTE_DATA, "sft_traces.jsonl")
+    gold_path = os.path.join(REMOTE_DATA, "gold_traces.jsonl")
+    dataset = load_sft_dataset(traces_path, tokenizer, seq_len)
+    gold_dataset = load_sft_dataset(gold_path, tokenizer, seq_len)
+    dataset.extend(gold_dataset)
+    print(f"  Total (with gold): {len(dataset):,}")
+
+    # Model
+    config = ModelConfig(
+        vocab_size=vocab_size,
+        d_model=1_280, n_layers=23, n_heads=20, d_ff=3_456,
+        max_seq_len=seq_len, dropout=0.0, tie_embeddings=True,
+    )
+    model = Retriever500M(config).to(device)
+
+    # Load pretrained checkpoint
+    ckpt_path = os.path.join(REMOTE_CKPT, "latest.pt")
+    if not os.path.exists(ckpt_path):
+        ckpt_path = os.path.join(REMOTE_CKPT, "latest_bf16.pt")
+    if not os.path.exists(ckpt_path):
+        ckpt_path = os.path.join(REMOTE_SEED_CKPT, "latest_bf16.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"])
+        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}")
+            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("WARNING: No checkpoint found, starting from scratch!")
+
+    total_params = model.count_parameters()
+    print(f"Model parameters: {total_params:,} ({total_params / 1e6:.1f}M)")
+
+    optimizer = setup_optimizer(model, lr, use_8bit=use_8bit_adam)
+
+    effective_batch = batch_size * grad_accum
+    print(f"\nSFT configuration:")
+    print(f"  Batch size:       {batch_size}")
+    print(f"  Grad accum:       {grad_accum}")
+    print(f"  Effective batch:  {effective_batch}")
+    print(f"  Sequence length:  {seq_len}")
+    print(f"  Learning rate:    {lr}")
+    print(f"  Steps:            {steps}")
+    print(f"  Warmup:           {warmup}")
+    print()
+
+    log = {
+        "config": asdict(config),
+        "train_args": {"steps": steps, "batch_size": batch_size, "grad_accum": grad_accum,
+                       "seq_len": seq_len, "lr": lr, "warmup": warmup,
+                       "save_every": save_every, "log_every": log_every,
+                       "use_8bit_adam": use_8bit_adam},
+        "total_params": total_params,
+        "steps": [],
+    }
+
+    model.train()
+    start_time = time.time()
+    accum_loss = 0.0
+    best_loss = float("inf")
+
+    pbar = tqdm(range(steps), desc="SFT")
+    for step in pbar:
+        lr_now = get_lr(step, warmup, steps, lr, lr * 0.1)
+        for pg in optimizer.param_groups:
+            pg["lr"] = lr_now
+
+        optimizer.zero_grad(set_to_none=True)
+
+        total_loss = 0.0
+        for _ in range(grad_accum):
+            input_ids, targets, loss_mask = get_sft_batch(
+                dataset, batch_size, seq_len, device
+            )
+            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+                out = model(input_ids, targets=targets, use_checkpoint=True)
+                logits = out["logits"]
+
+                if loss_mask.sum() > 0:
+                    shifted_mask = loss_mask[:, 1:].contiguous()
+                    masked_logits = logits[:, :-1, :].contiguous()
+                    masked_targets = targets[:, :-1].contiguous()
+
+                    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 / grad_accum
+                else:
+                    loss = out["loss"] / 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 % log_every == 0 or step == steps - 1:
+            elapsed = time.time() - start_time
+            steps_per_sec = (step + 1) / elapsed
+            vram_used = torch.cuda.max_memory_allocated() / 1e9
+
+            log_entry = {
+                "step": step, "loss": avg_loss, "ema_loss": accum_loss,
+                "lr": lr_now, "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_now:.2e}", "vram": f"{vram_used:.1f}G",
+            })
+
+        if (step + 1) % save_every == 0 or step == steps - 1:
+            ckpt_path = os.path.join(REMOTE_CKPT, 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(REMOTE_CKPT, "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(REMOTE_CKPT, "sft_best.pt")
+                torch.save({
+                    "model_state_dict": model.state_dict(),
+                    "config": asdict(config),
+                    "step": step + 1, "loss": accum_loss,
+                }, best_path)
+
+            log_path = os.path.join(REMOTE_LOGS, "sft_log.json")
+            with open(log_path, "w") as f:
+                json.dump(log, f, indent=2)
+
+            volume.commit()
+
+        if step % 50 == 0:
+            torch.cuda.reset_peak_memory_stats()
+
+    log_path = os.path.join(REMOTE_LOGS, "sft_log.json")
+    with open(log_path, "w") as f:
+        json.dump(log, f, indent=2)
+    volume.commit()
+
+    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}")
+    return accum_loss
+
+
+# ─── Modal entry points ──────────────────────────────────────────────────────
+
+GPU_CHOICES = {"a10g": "A10G", "a100": "A100", "h100": "H100"}
+
+
+@app.function(
+    volumes={VOL_MOUNT: volume},
+    gpu="A10G",
+    timeout=3600,
+)
+def pretrain(
+    steps=1000,
+    batch_size=8,
+    grad_accum=4,
+    seq_len=512,
+    lr=3e-4,
+    warmup=100,
+    save_every=200,
+    log_every=10,
+    corpus="curated",
+    resume=True,
+    use_8bit_adam=True,
+    gpu="a10g",
+):
+    """Run pretraining on Modal GPU."""
+    return run_pretrain(steps, batch_size, grad_accum, seq_len, lr, warmup,
+                        save_every, log_every, corpus, resume, use_8bit_adam)
+
+
+@app.function(
+    volumes={VOL_MOUNT: volume},
+    gpu="A10G",
+    timeout=3600,
+    env={"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"},
+)
+def sft(
+    steps=500,
+    batch_size=4,
+    grad_accum=8,
+    seq_len=768,
+    lr=5e-5,
+    warmup=20,
+    save_every=100,
+    log_every=10,
+    use_8bit_adam=True,
+    gpu="a10g",
+):
+    """Run SFT on Modal GPU."""
+    return run_sft(steps, batch_size, grad_accum, seq_len, lr, warmup,
+                   save_every, log_every, use_8bit_adam)
+
+
+@app.function(
+    volumes={VOL_MOUNT: volume},
+    gpu="A10G",
+    timeout=7200,
+)
+def pipeline(pretrain_steps=1000, sft_steps=500, gpu="a10g"):
+    """Run pretraining then SFT in one go."""
+    print("=" * 60)
+    print("PHASE 1: PRETRAINING")
+    print("=" * 60)
+    run_pretrain(
+        steps=pretrain_steps, batch_size=8, grad_accum=4, seq_len=512,
+        lr=3e-4, warmup=100, save_every=200, log_every=10,
+        corpus="curated", resume=True, use_8bit_adam=True,
+    )
+
+    print("\n" + "=" * 60)
+    print("PHASE 2: SFT")
+    print("=" * 60)
+    run_sft(
+        steps=sft_steps, batch_size=8, grad_accum=4, seq_len=768,
+        lr=5e-5, warmup=20, save_every=100, log_every=10,
+        use_8bit_adam=True,
+    )
+
+
+# ─── Local entry point for `modal run` ───────────────────────────────────────
+
+@app.local_entrypoint()
+def main(
+    command: str = "pipeline",
+    steps: int = 1000,
+    batch_size: int = 8,
+    grad_accum: int = 4,
+    seq_len: int = 512,
+    lr: float = 3e-4,
+    warmup: int = 100,
+    save_every: int = 200,
+    log_every: int = 10,
+    corpus: str = "curated",
+    resume: bool = True,
+    use_8bit_adam: bool = True,
+    gpu: str = "a10g",
+    pretrain_steps: int = 1000,
+    sft_steps: int = 500,
+):
+    if command == "seed":
+        seed.remote()
+    elif command == "pretrain":
+        result = pretrain.remote(
+            steps=steps, batch_size=batch_size, grad_accum=grad_accum,
+            seq_len=seq_len, lr=lr, warmup=warmup,
+            save_every=save_every, log_every=log_every,
+            corpus=corpus, resume=resume, use_8bit_adam=use_8bit_adam,
+            gpu=gpu,
+        )
+        print(f"Pretraining final loss: {result}")
+    elif command == "sft":
+        result = sft.remote(
+            steps=steps, batch_size=batch_size, grad_accum=grad_accum,
+            seq_len=seq_len, lr=lr, warmup=warmup,
+            save_every=save_every, log_every=log_every,
+            use_8bit_adam=use_8bit_adam, gpu=gpu,
+        )
+        print(f"SFT final loss: {result}")
+    elif command == "download":
+        download.remote()
+    elif command == "pipeline":
+        pipeline.remote(pretrain_steps=pretrain_steps, sft_steps=sft_steps, gpu=gpu)
+    else:
+        print(f"Unknown command: {command}")
+        print("Available: upload, pretrain, sft, download, pipeline")

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()