Enrich model card with examples and training details

README.md

@@ -16,82 +16,144 @@ library_name: transformers
 pipeline_tag: text-generation
 ---
 
-# SearchLM NL2BM25 — GRPO v2 Shaped Reward (Qwen2.5-3B-Instruct)
+# SearchLM NL2BM25 — GRPO v2 Shaped Reward ✅ (Qwen2.5-3B-Instruct)
 
-The best-performing SearchLM checkpoint. Trained via GRPO with a shaped reward
-that eliminates the keyword-bag gaming observed in [GRPO v1](Supreeth/searchlm-nl2bm25-grpo).
+**Part of the [SearchLM collection](https://huggingface.co/collections/Supreeth/searchlm) · [GitHub](https://github.com/SupreethRao99/searchLM)**
 
-See the [SearchLM collection](https://huggingface.co/collections/Supreeth/searchlm) for all checkpoints
-and the [v2 analysis report](https://github.com/SupreethRao99/searchLM/blob/main/REWARD_HACKING_REPORT_V2.md)
-for full benchmark results and behavioral analysis.
+The best-performing SearchLM checkpoint. Trained via GRPO with a shaped reward that
+eliminated the specification gaming found in [GRPO v1](Supreeth/searchlm-nl2bm25-grpo) while simultaneously
+improving retrieval quality. Achieves **NDCG@10 = 0.577** on NFCorpus and **0.657** on SciFact.
+
+> **Pipeline position:** `base → SFT v1 → GRPO v1 (⚠️) → SFT v2 → `**`GRPO v2 ✅`**
+
+---
 
 ## What it does
 
-Given a natural language query, the model reasons step-by-step about key concepts,
-synonyms, and boolean structure, then emits a Tantivy boolean search query:
+The model reasons step-by-step about key concepts, synonym expansion, and boolean structure,
+then emits a Tantivy-compatible boolean search query:
 
+**Input:** `Do Cholesterol Statin Drugs Cause Breast Cancer?`
+
+**Output:**
+```
+<reasoning>
+Key concepts:
+1. Statin drugs — synonyms: statin, "HMG-CoA reductase inhibitor", simvastatin,
+   atorvastatin, lovastatin, pravastatin
+2. Causal relationship — cause, risk, association, induce, "increase risk"
+3. Breast cancer — "breast cancer", "breast carcinoma", "breast neoplasm"
+
+Strategy: AND the three concept groups; use OR to expand synonyms within each.
+Phrase-quote multi-word terms; keep AND chains short to avoid zero-result queries.
+</reasoning>
+<query>(statin OR "HMG-CoA reductase inhibitor" OR simvastatin OR atorvastatin OR lovastatin)
+AND (cause OR risk OR association OR induce)
+AND ("breast cancer" OR "breast carcinoma" OR "breast neoplasm")</query>
 ```
-Input: Do Cholesterol Statin Drugs Cause Breast Cancer?
 
-Output:
+Compare to [GRPO v1](Supreeth/searchlm-nl2bm25-grpo)'s output for the same query:
+```
 <reasoning>
-Key concepts: (1) cholesterol-lowering statin drugs — synonyms: statin, HMG-CoA reductase
-inhibitor, simvastatin, atorvastatin; (2) causal relationship — cause, risk, association,
-induce; (3) breast cancer — "breast cancer", "breast carcinoma", "breast neoplasm".
-Connect with AND; expand synonyms with OR.
 </reasoning>
-<query>(statin OR "HMG-CoA reductase inhibitor" OR simvastatin OR atorvastatin) AND
-(cause OR risk OR association OR induce) AND ("breast cancer" OR "breast carcinoma")</query>
+<query>Cholesterol Statin Breast Cancer</query>
 ```
 
+GRPO v2 generates 147-token completions with substantive reasoning; GRPO v1 generated 5-token
+keyword bags with empty reasoning blocks.
+
+---
+
 ## How v2 eliminated reward hacking
 
-The v1 reward (`0.6 × NDCG@10 + 0.4 × MRR`) was gameable with keyword bags on small corpora.
-v2 uses a shaped reward with three anti-gaming mechanisms:
+The v1 reward (`0.6 × NDCG@10 + 0.4 × MRR`) was gameable with keyword bags on small corpora
+because BM25 recall on 3–5K doc indexes is high for distinctive nouns. Three mechanisms
+closed this gap in v2:
 
+```python
+# v2 reward function
+base   = 0.6 * max(0, ndcg_at_10 - keyword_baseline_ndcg)  # must beat noun-extraction
+       + 0.4 * mrr
+shaped = complexity_mult * base                              # 1.0 with boolean ops, 0.5 without
+       + 0.15 * min(reasoning_tokens / 100, 1.0)            # up to +0.15 reasoning bonus
+reward = 0.0 if len(query.split()) < 3 else shaped          # hard gate: ≥3 tokens required
 ```
-base   = 0.6 × max(0, NDCG@10 − keyword_baseline) + 0.4 × MRR
-shaped = complexity_mult × base + 0.15 × min(reasoning_tokens / 100, 1.0)
-reward = 0.0  if  len(query.split()) < 3  else  shaped
 
-where:
-  keyword_baseline = NDCG@10 of naive noun-extraction (precomputed per query)
-  complexity_mult  = 1.0 with boolean operators, 0.5 without
-```
+| Mechanism | Effect |
+|-----------|--------|
+| Keyword baseline delta | Model earns zero NDCG credit for matching naive noun-extraction |
+| Hard length gate | Single/double-word queries unconditionally return 0.0 |
+| Reasoning depth bonus | Up to +0.15 reward for ≥100-token reasoning blocks |
+| Complexity multiplier | Queries without boolean operators earn half credit |
 
-1. **Keyword baseline delta** — must beat noun-extraction to earn NDCG credit
-2. **Hard length gate** — queries < 3 tokens → reward = 0.0
-3. **Reasoning depth bonus** — up to +0.15 for substantive reasoning blocks
+---
+
+## All SearchLM checkpoints
 
-## Benchmark (test split, NDCG@10)
+| Model | NFCorpus NDCG@10 | SciFact NDCG@10 | Mean tokens | Boolean ops |
+|-------|-----------------|----------------|-------------|-------------|
+| base (Qwen2.5-3B-Instruct) | 0.455 | 0.386 | 120 | ~20% |
+| [SFT v1](https://huggingface.co/Supreeth/searchlm-nl2bm25-sft) | 0.441 | 0.273 | 95 | ~80% |
+| [GRPO v1](https://huggingface.co/Supreeth/searchlm-nl2bm25-grpo) ⚠️ | 0.556 | 0.608 | **5–7** | **0%** |
+| [SFT v2](https://huggingface.co/Supreeth/searchlm-nl2bm25-sft-v2) | 0.466 | 0.358 | 109 | ~65% |
+| [**GRPO v2**](https://huggingface.co/Supreeth/searchlm-nl2bm25-grpo-v2) ✅ | **0.577** | **0.657** | 147 | ~35% |
 
-| Dataset | Base | SFT v1 | GRPO v1 | SFT v2 | **GRPO v2** |
-|---------|------|--------|--------|--------|------------|
-| NFCorpus | 0.455 | 0.441 | 0.556 | 0.466 | **0.577** |
-| SciFact  | 0.386 | 0.273 | 0.608 | 0.358 | **0.657** |
+Evaluated on BEIR test splits (NFCorpus: 323 queries, SciFact: 300 queries).
+
+---
 
-## Behavioral comparison vs GRPO v1
+## Behavioral comparison (GRPO v1 vs GRPO v2)
 
-| Dimension | GRPO v1 | GRPO v2 |
+| Dimension | [GRPO v1](Supreeth/searchlm-nl2bm25-grpo) ⚠️ | **GRPO v2** ✅ |
 |-----------|---------|---------|
-| Mean completion length | 6 tokens | **147 tokens** |
-| Boolean operator usage | 0% | ~35% |
-| `frac_reward_zero_std` (early training) | 90–96% | **0%** |
-| Reasoning block | empty | substantive |
+| NFCorpus NDCG@10 | 0.556 | **0.577** (+0.021) |
+| SciFact NDCG@10 | 0.608 | **0.657** (+0.049) |
+| Mean completion length | **5–7 tokens** | 147 tokens |
+| Boolean operator usage | **0%** | ~35% |
+| Phrase usage | **0%** | ~25% |
+| `frac_reward_zero_std` (step 1) | **90–96%** | 0.0% |
+| `frac_reward_zero_std` (final) | **90–96%** | ~61% |
+| Reasoning block | **empty** | substantive |
+
+The shaped reward did not sacrifice performance to eliminate gaming — it improved both.
+
+---
 
 ## Training Details
 
 | Setting | Value |
 |---------|-------|
 | Base model | [searchlm-nl2bm25-sft-v2](Supreeth/searchlm-nl2bm25-sft-v2) |
-| Method | GRPO (TRL + vLLM colocate) |
-| Reward | Shaped (baseline delta + complexity mult + reasoning bonus) |
-| Training datasets | NFCorpus + SciFact + FiQA (3K queries, 57K docs) |
+| Method | GRPO (TRL GRPOTrainer + vLLM colocate, single H100) |
+| Reward | Shaped: `complexity_mult × (0.6 × ΔNDCG + 0.4 × MRR) + 0.15 × reasoning_depth` |
+| Training datasets | NFCorpus + SciFact + FiQA-2018 (3K queries, 57,638 docs) |
+| `num_generations` | 8 (was 2 in v1) |
 | Epochs | 1 |
-| `num_generations` | 8 |
-| Steps | 2,879 |
-| Hardware | NVIDIA H100 80 GB (~3h 3m) |
+| Steps | 2,879 (~3.3s/step) |
+| Batch size | 2 (+ 8 grad accum = effective 16) |
+| Learning rate | 1e-6 |
+| vLLM GPU utilisation | 0.30 (24 GB KV cache) |
+| Max new tokens | 1,024 |
+| Gradient checkpointing | yes |
+| Hardware | NVIDIA H100 80 GB |
+| Training time | ~3h 3m |
+| Final train loss | 0.0012 |
 | Final mean reward | ~0.29 |
+| W&B run | `supreethrao/searchlm/runs/9x1tg52j` |
+
+### Why these hyperparameters
+
+**`num_generations=8`**: v1 used 2, leading to 90-96% of groups having zero within-group
+reward variance (no gradient signal). With 8 completions, variance emerged from step 1.
+
+**`vllm_gpu_memory_utilization=0.30`**: On H100 80GB, Adam fp32 optimizer states for a 3B
+model require ~24 GB. At 0.45 utilisation, vLLM reserved 36 GB and Adam states OOM'd. 0.30
+leaves 56 GB for the training stack.
+
+**`torch_compile=False`**: Compiled backward pass materialised fp32 FFN intermediate buffers
+(~90 MB each) that eager + gradient checkpointing avoids, causing OOM at batch_size=4.
+
+---
 
 ## Usage
 
@@ -105,7 +167,8 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 tokenizer = AutoTokenizer.from_pretrained("Supreeth/searchlm-nl2bm25-grpo-v2")
 
-SYSTEM_PROMPT = """You are an expert information retrieval specialist. Convert the natural language query into a Tantivy boolean search query.
+SYSTEM_PROMPT = """You are an expert information retrieval specialist. Convert the \
+natural language query into a Tantivy boolean search query.
 
 Output format (strictly follow this):
 <reasoning>
@@ -124,18 +187,32 @@ outputs = model.generate(**inputs, max_new_tokens=512, temperature=0.7, do_sampl
 print(tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True))
 ```
 
+---
+
 ## Tantivy Boolean Syntax
 
-| Construct | Syntax |
-|-----------|--------|
-| Single term | `cancer` |
-| Exact phrase | `"bone density"` |
-| AND | `vitamin AND calcium` |
-| OR | `cancer OR tumor OR malignancy` |
-| NOT | `NOT review` |
-| Grouping | `(cat OR feline) AND behavior` |
-| Field scope | `title:"machine learning"` |
-| Boost | `cancer^2 OR tumor` |
+[Tantivy](https://github.com/quickwit-oss/tantivy) is a full-text search engine library.
+The model targets its query language:
+
+| Construct | Syntax | Example |
+|-----------|--------|---------|
+| Single term | `word` | `cancer` |
+| Exact phrase | `"phrase"` | `"bone density"` |
+| AND | `A AND B` | `vitamin AND calcium` |
+| OR | `A OR B` | `cancer OR tumor OR malignancy` |
+| NOT | `NOT A` | `NOT review` |
+| Grouping | `(A OR B)` | `(cat OR feline) AND behavior` |
+| Field scope | `field:term` | `title:"machine learning"` |
+| Boost | `term^N` | `cancer^2 OR tumor` |
+
+---
+
+## Related resources
+
+- **Code:** [SupreethRao99/searchLM](https://github.com/SupreethRao99/searchLM)
+- **Analysis:** [Full five-checkpoint comparison report](https://github.com/SupreethRao99/searchLM/blob/main/REWARD_HACKING_REPORT_V2.md)
+- **Dataset:** [Supreeth/nl2bm25-sft](https://huggingface.co/datasets/Supreeth/nl2bm25-sft)
+- **Collection:** [SearchLM collection](https://huggingface.co/collections/Supreeth/searchlm)
 
 ## Citation