OneRetrieval (ICDE 2027)

Reference implementation for the paper OneRetrieval: Unifying Multi-Branch E-commerce Retrieval with an Editable Generative Model (ICDE 2027).

The repository provides an end-to-end, self-contained pipeline that produces all training data (Stage 0 → Stage 3) required by the paper, from a small synthetic sample dataset. Every script is aligned with a specific section of the paper (see the mapping table below).

---

1. Environment

• Python ≥ 3.9

• Install dependencies:

  pip install -r requirements.txt

• Optional: GPU + vLLM is only needed by update_2_3_llm_attr_important.py (the LLM primary-subject precedence step). For a quick end-to-end smoke test you can skip it and feed an empty dictionary (see Step 3 below).

• BGE embedding: sample_data/gen_samples.py uses the BAAI/bge-small-zh-v1.5 model (via sentence-transformers) to encode attribute words into real semantic embeddings — no random noise. On first run it will download the model (~90 MB). You can override the model with:

  BGE_MODEL_NAME=BAAI/bge-large-zh-v1.5 python sample_data/gen_samples.py

---

2. Repository Layout

oneretrieval_code/
├── README.md
├── requirements.txt
├── .gitignore
├── oneretrieval_icde_2027.pdf           # paper (for reference)
├── update_0_attribute_merge_analysis.py # Step 1: 18-cat → 6-group merge
├── update_1_get_kmeans_model.py         # Step 2: codebook construction
├── update_2_3_llm_attr_important.py     # Step 3: LLM primary-subject pick
├── update_4_get_sid.py                  # Step 4: item / query → 6-token SID
├── update_5_process_train_dataset.py    # Step 5: build Stage 0/1/2/3 data
└── sample_data/
    └── gen_samples.py                   # one-shot synthetic data generator

Running python sample_data/gen_samples.py will (re)create the following files under sample_data/ (they are not committed to git):

File	Rows	Columns
attr.txt	345	attr\tcate
attr_emb.txt	345	attr\tcate\tpv\temb(BGE-encoded, \x02-sep)
item.txt	300	item_id\ttitle
triple.txt	500	19 columns of q ↔ i interaction

---

3. Quick Start (walk through Stage 0 – Stage 3)

# 0) Generate the synthetic sample dataset
python sample_data/gen_samples.py

# 1) Information-theoretic attribute-category merging (18 → 6 groups)
python update_0_attribute_merge_analysis.py \
    --attr_file  sample_data/attr.txt \
    --item_file  sample_data/item.txt \
    --output_dir results/merge_analysis_demo

# 2) Build 6-group codebook with four-block layout (empty+cluster+solo+reserved)
python update_1_get_kmeans_model.py \
    --input      sample_data/attr_emb.txt \
    --output_dir results/kmeans_demo \
    --n_reserved 1                       # paper recommends 10 for production

# 3) LLM primary-subject precedence
#    (Needs vLLM + GPU. For a local smoke test, skip and use an empty dict:)
mkdir -p results/llm_demo
echo '{}' > results/llm_demo/attr_important_dict.json
# Real run example (produces results/llm_demo/attr_important_dict.json):
# python update_2_3_llm_attr_important.py \
#     --attr_file  sample_data/attr_emb.txt \
#     --item_file  sample_data/item.txt \
#     --output_dir results/llm_demo \
#     --model_path <path-or-hub-id-of-your-LLM> \
#     --tensor_parallel_size 2 \
#     --batch_size 1024

# 4) Encode items and queries into 6-token SIDs
python update_4_get_sid.py --mode item \
    --codebook   results/kmeans_demo/codebook_encoded.txt \
    --attr_pv    sample_data/attr_emb.txt \
    --important  results/llm_demo/attr_important_dict.json \
    --input      sample_data/item.txt \
    --output     results/sid_demo/item_sid.txt

python update_4_get_sid.py --mode query \
    --codebook   results/kmeans_demo/codebook_encoded.txt \
    --attr_pv    sample_data/attr_emb.txt \
    --important  results/llm_demo/attr_important_dict.json \
    --input      sample_data/triple.txt \
    --output     results/sid_demo/query_sid.txt

# 5) Build the four-stage SFT training data
python update_5_process_train_dataset.py stage0 \
    --codebook results/kmeans_demo/codebook_encoded.txt \
    --output   results/stage0.txt

python update_5_process_train_dataset.py stage1 \
    --query results/sid_demo/query_sid.txt \
    --item  results/sid_demo/item_sid.txt \
    --output results/stage1.txt

python update_5_process_train_dataset.py stage2 \
    --pair   sample_data/triple.txt \
    --output results/stage2.txt

python update_5_process_train_dataset.py stage3 \
    --triple sample_data/triple.txt \
    --output results/stage3.txt

After a successful run you should see (on the provided sample data):

Output	Rows	Content
results/sid_demo/item_sid.txt	300	item_id\tsid\ttitle
results/sid_demo/query_sid.txt	500	query\tsid
results/stage0.txt	~5 900	attr↔SID bidirectional + zero-sid
results/stage1.txt	600	6 tasks: q↔SID, title↔SID, q→cate, title→cate
results/stage2.txt	1 000	4 tasks: q↔title, q_sid↔i_sid
results/stage3.txt	500	q + q_sid + HIST_q + HIST_s → i_sid

---

4. Paper Mapping

Script	Paper section
update_0_attribute_merge_analysis.py	§III-C Information-theoretic attribute-category merging (18 → 6)
update_1_get_kmeans_model.py	§III-D Codebook Construction (non-uniform capacity, four-block layout)
update_2_3_llm_attr_important.py	§III-D3 Primary-subject precedence (LLM pairwise resolution)
update_4_get_sid.py	§III-D3 Item / query record → 6-token SID
update_5_process_train_dataset.py stage0	§III-E Stage 0 Attribute–SID alignment
update_5_process_train_dataset.py stage1	§III-E Stage 1 Content alignment (4 bidirectional + 2 category tasks)
update_5_process_train_dataset.py stage2	§III-E Stage 2 Collaborative co-occurrence (surface + SID level)
update_5_process_train_dataset.py stage3	§III-E Stage 3 Personalized retrieval + reserved-slot self-routing

---

5. Model Training

This repository only produces the training data (Stage 0 – Stage 3) for OneRetrieval. For the actual model training (BART-base backbone, SFT loop, DPO refinement, etc.), we recommend reusing the training pipeline in

benchen4395/onesearch-family

Specifically, the outputs of this repo (results/stage0.txt ~ results/stage3.txt) can be fed directly into the SFT with Self-Distillation track of that repository — that track is what the OneRetrieval paper uses for the multi-stage curriculum learning.

Each line in our produced files is in prompt\tresponse format, which is directly compatible with the SFT data loader in onesearch-family.

---

6. Notes

• The four sample_data/*.txt files are re-generatable and therefore git-ignored. Run python sample_data/gen_samples.py once after cloning.
• results/ is generated at run time and is git-ignored as well.
• All produced training files use TAB as the prompt/response separator so they can be loaded by any standard seq2seq trainer.

7. Citation

If you find this code useful, please cite:

@article{zhang2026oneretrieval,
  title={OneRetrieval: Unifying Multi-Branch E-commerce Retrieval with an Editable Generative Model},
  author={Zhang, Xuxin and Chen, Ben and Lv, Yue and Wang, Siyuan and Li, Yupeng and Ma, Yufei and Liang, Zihan and Zhao, Tong and Yang, Ying and Dai, Huangyu and others},
  journal={arXiv preprint arXiv:2606.13533},
  year={2026}
}