storm-qwen3-1.7B/README.md

---
base_model: Qwen/Qwen3-1.7B
pipeline_tag: text-generation
tags:
- qwen
- qwen3
- lora
- storm
- information-retrieval
- keyword-generation
- query-expansion
---

# STORM Qwen3 Keyword Generator

Preprint for EMNLP 2026

## Abstract
Modern retrieval increasingly relies on dense and learned-sparse neural models that are effective but require encoding the 
entire corpus into a specialized index, which must be rebuilt whenever the model changes. Lexical retrievers like BM25 stay 
efficient, transparent, and run on a standard inverted index that need not change as models evolve, but suffer from vocabulary
mismatch. LLM query rewriting can help, yet prompted rewriters emit well-formed but retrieval-ineffective---or harmful---terms,
and training against a retrieval reward gives only delayed, sequence-level supervision that obscures which terms helped. 
We introduce STORM (Stepwise Token Optimization with Reward-guided beaM search), a self-supervised framework for lexical 
query expansion. STORM trains the rewriter through generation guided by retrieval metrics: at each step, a beam of candidate
expansions is scored against the BM25 index and low-reward continuations are pruned, turning the retrieval reward into a 
token-level signal that concentrates exploration on retrieval-effective vocabulary. Across TREC DL and BEIR, STORM lets 
0.6B--8B backbones match or surpass competitive LLM rewriters and far larger proprietary models while retrieving as fast 
as plain BM25, and transfers zero-shot to 18 languages (MIRACL), beating dedicated multilingual dense retrievers. 
STORM is thus a competitive, infrastructure-light alternative to dense neural retrieval.


## Install

```bash
pip install -U torch transformers accelerate safetensors
```

## Load and run
Note: Group beam search is optional; standard beam search (num_beams=5, num_return_sequences=1 or 3) or 
decoding without beam search also works, although our slightly best-performing configuration used group beam search with num_beams=6 and num_return_sequences=3.


```python
import torch
from transformers import AutoModelForCausalLM
from transformers import AutoTokenizer

model_id = "Arthur-75/storm-qwen3-1.7B"

system_prompt = (
    "From the query generate new semantic related keywords.\n"
    "Output the result strictly as a single comma-separated line."
)

if torch.cuda.is_available():
    device = "cuda"
elif torch.backends.mps.is_available():
    device = "mps"
else:
    device = "cpu"

tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
#tokenizer.padding_side = "left"

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    dtype=torch.float32,
    #device_map=None,
    #low_cpu_mem_usage=True,
    trust_remote_code=True,
)

model = model.to(device)
model.eval()



def generate_keywords(query: str):
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": f"[QUERY]: {query.strip()}\n[KEYWORDS]: "},
    ]

    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        add_special_tokens=False,
    ).to(device)

    prompt_len = inputs["input_ids"].shape[1]

    with torch.inference_mode():
        output = model.generate(
            **inputs,
            max_new_tokens=32,#[32,64]
            do_sample=False,
            num_beams=6,
            num_beam_groups=3,
            diversity_penalty=1.0,
            num_return_sequences=3,
            #repetition_penalty=1.0,
            custom_generate="transformers-community/group-beam-search"
           # pad_token_id=tokenizer.pad_token_id,
           # eos_token_id=tokenizer.eos_token_id,
        )

    decoded = tokenizer.batch_decode(
        output[:, prompt_len:],
        skip_special_tokens=True,
    )

    return decoded


query = "What are the symptoms of vitamin D deficiency?"
outputs = generate_keywords(query)

print(outputs)
```