# makeitwork1

Retriever500M — a 497M-parameter decoder-only transformer trained as a search agent.

## Architecture

Custom model defined in `src/model.py` (`Retriever500M`):
- d_model=1280, n_layers=23, n_heads=20, d_ff=3456
- RoPE positional embeddings
- Tied input/output embeddings

## Training

Two stages, both run on Modal (volume `retriever500m-data`):

1. **Pretraining** (1711 steps, corpus_curated.txt, lr=3e-4, seq_len=512):
   final EMA loss = 0.122
2. **SFT** (500 steps, sft_traces.jsonl + gold_traces.jsonl, lr=5e-5, seq_len=768):
   final EMA loss = 0.084

The checkpoint in this repo (`sft_latest.pt`) is the final SFT weights
(model_state_dict + config, step 500). Vocab size = 32009 (includes 9 special
agent tokens: system, user, assistant, search, result, evidence, reasoning,
finish, end).

## Loading

```python
import torch, sys
sys.path.insert(0, "src")
from model import ModelConfig, Retriever500M

ckpt = torch.load("sft_latest.pt", map_location="cuda", weights_only=False)
config = ModelConfig(**ckpt["config"])
model = Retriever500M(config).to("cuda")
model.load_state_dict(ckpt["model_state_dict"])
model.eval()
```

Tokenizer: `tokenizer/tokenizer_agent.json` (HuggingFace `tokenizers` library).