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

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-4B"

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=64,#[32,64]
            do_sample=False,
            num_beams=6, #[3, 6]
            num_beam_groups=3,
            diversity_penalty=1.0,
            num_return_sequences=3, #[1, 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)