初始化项目,由ModelHub XC社区提供模型
Model: Supreeth/searchlm-nl2bm25-grpo-v2 Source: Original Platform
This commit is contained in:
226
README.md
Normal file
226
README.md
Normal file
@@ -0,0 +1,226 @@
|
||||
---
|
||||
language:
|
||||
- en
|
||||
license: apache-2.0
|
||||
base_model: Qwen/Qwen2.5-3B-Instruct
|
||||
tags:
|
||||
- information-retrieval
|
||||
- boolean-search
|
||||
- NL2BM25
|
||||
- GRPO
|
||||
- RLVR
|
||||
- tantivy
|
||||
- BEIR
|
||||
- searchlm
|
||||
library_name: transformers
|
||||
pipeline_tag: text-generation
|
||||
---
|
||||
|
||||
# SearchLM NL2BM25 — GRPO v2 Shaped Reward ✅ (Qwen2.5-3B-Instruct)
|
||||
|
||||
**Part of the [SearchLM collection](https://huggingface.co/collections/Supreeth/searchlm) · [GitHub](https://github.com/SupreethRao99/searchLM)**
|
||||
|
||||
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
|
||||
|
||||
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>
|
||||
```
|
||||
|
||||
Compare to [GRPO v1](Supreeth/searchlm-nl2bm25-grpo)'s output for the same query:
|
||||
```
|
||||
<reasoning>
|
||||
</reasoning>
|
||||
<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
|
||||
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
|
||||
```
|
||||
|
||||
| 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 |
|
||||
|
||||
---
|
||||
|
||||
## All SearchLM checkpoints
|
||||
|
||||
| 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% |
|
||||
|
||||
Evaluated on BEIR test splits (NFCorpus: 323 queries, SciFact: 300 queries).
|
||||
|
||||
---
|
||||
|
||||
## Behavioral comparison (GRPO v1 vs GRPO v2)
|
||||
|
||||
| Dimension | [GRPO v1](Supreeth/searchlm-nl2bm25-grpo) ⚠️ | **GRPO v2** ✅ |
|
||||
|-----------|---------|---------|
|
||||
| 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 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 |
|
||||
| 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
|
||||
|
||||
```python
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"Supreeth/searchlm-nl2bm25-grpo-v2",
|
||||
torch_dtype="auto",
|
||||
device_map="auto",
|
||||
)
|
||||
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.
|
||||
|
||||
Output format (strictly follow this):
|
||||
<reasoning>
|
||||
Step-by-step concept extraction and synonym expansion.
|
||||
</reasoning>
|
||||
<query>your boolean query here</query>"""
|
||||
|
||||
nl_query = "effects of climate change on coral reef ecosystems"
|
||||
messages = [
|
||||
{"role": "system", "content": SYSTEM_PROMPT},
|
||||
{"role": "user", "content": f"Convert to a Tantivy boolean search query:\n\n{nl_query}"},
|
||||
]
|
||||
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
|
||||
inputs = tokenizer(text, return_tensors="pt").to(model.device)
|
||||
outputs = model.generate(**inputs, max_new_tokens=512, temperature=0.7, do_sample=True)
|
||||
print(tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True))
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Tantivy Boolean Syntax
|
||||
|
||||
[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
|
||||
|
||||
```bibtex
|
||||
@misc{searchlm2026,
|
||||
title = {SearchLM: Training Small Language Models for Boolean Query Generation via RLVR},
|
||||
author = {Rao, Supreeth},
|
||||
year = {2026},
|
||||
url = {https://github.com/SupreethRao99/searchLM},
|
||||
}
|
||||
```
|
||||
Reference in New Issue
Block a user