oneretrieval

KO | EN

gitlite — search

#javascript #python #hacktoberfest #react #ai #typescript #llm #go #golang #android #machine-learning #rust #deep-learning #linux

★ 12 Open GitHub ↗

No description available.

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# 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:

```bash
  pip install -r requirements.txt
  ```

- **Optional**: GPU + [vLLM](https://github.com/vllm-project/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`](https://huggingface.co/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:

```bash
  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)

```bash
# 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](https://github.com/benchen4395/onesearch-family/tree/main)**

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:

```bibtex
@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}
}
```

oneretrieval

Explore Similar Repositories

Was this content helpful?

Recent Feedback

Download README