# Difficulty-Aware Semantic-ID Optimization for Generative Recommendation

Official code release for the paper **Difficulty-Aware Semantic-ID Optimization for Generative Recommendation (DASO)**.

This repository is a focused DASO paper repo. It is not a general-purpose `MiniOneRec` framework dump. The codebase is intentionally narrowed to the training and evaluation path used by the paper:

1. build semantic IDs
2. run supervised fine-tuning (SFT)
3. run post-SFT reinforcement learning with either `GRPO` or `DASO`
4. evaluate with constrained semantic-ID decoding

## What This Repo Contains

- semantic-ID construction code under `rq/`
- dataset conversion and preprocessing utilities under `data/` and `convert_dataset.py`
- SFT training code in `sft.py`
- RL training code in `rl.py`
- the customized trainer in `minionerec_trainer.py`
- constrained decoding and metric computation in `evaluate.py`, `calc.py`, and `calc_level.py`

## What This Repo Does Not Try To Be

This release removes unrelated experimental branches, internal notes, self-distillation code, GPR variants, reward ablations, and old exploratory scripts. The public post-SFT methods exposed by this repo are:

- `GRPO`: the main baseline after SFT
- `DASO`: the difficulty-aware method in the paper

## Repository Structure

| Path | Description |
| --- | --- |
| `rq/` | semantic-ID construction: text embedding, RQ-VAE training, index generation |
| `data/` | processed Amazon data, index files, and preprocessing scripts |
| `config/zero2_opt.yaml` | Accelerate + DeepSpeed config used by RL training |
| `convert_dataset.py` | converts processed interaction data into the DASO training format |
| `convert_dataset.sh` | lightweight wrapper around `convert_dataset.py` |
| `sft.py` | supervised fine-tuning on semantic IDs |
| `rl.py` | RL fine-tuning entry for `GRPO` and `DASO` |
| `minionerec_trainer.py` | customized GRPO trainer with difficulty-aware sibling-relative advantages |
| `evaluate.py` | constrained decoding evaluation |
| `calc.py` | ranking metric computation |
| `calc_level.py` | per-level semantic-ID analysis |
| `split.py` / `merge.py` | shard evaluation inputs and merge outputs across GPUs |

## Data Included In The Repo

The repo currently contains processed public Amazon splits and semantic-ID resources for:

- `Industrial_and_Scientific`
- `Office_Products`

Included files cover:

- `train / valid / test` CSV splits
- `info` files mapping semantic ID, item title, and item ID
- `index` files mapping item IDs to semantic IDs
- precomputed item embedding arrays used during semantic-ID construction

If you only want to reproduce the paper pipeline, you can start directly from the included processed data.

## Environment Setup

Recommended environment:

- Python `3.11`
- PyTorch `2.6.0`
- `accelerate==1.10.1`
- `transformers==4.57.1`
- `trl==0.24.0`
- bf16-capable NVIDIA GPUs

Create the environment:

```bash
conda create -n daso python=3.11 -y
conda activate daso
pip install -r requirements.txt
```

## Hardware

The default scripts are written for multi-GPU training.

- `SFT`: default scripts assume 4 GPUs
- `GRPO`: default scripts assume 4 GPUs
- `DASO`: default scripts assume 8 GPUs

You can override GPU selection with `CUDA_VISIBLE_DEVICES`, and several scripts also expose `MODEL_PATH`, `OUTPUT_DIR`, `MASTER_ADDR`, `MASTER_PORT`, or `CUDA_LIST` through environment variables.

## Quick Start

If you want the shortest path to running the repo:

### 1. SFT

Industrial:

```bash
bash run_sft_industrial.sh
```

Office:

```bash
bash run_sft_office.sh
```

### 2. GRPO baseline

Industrial:

```bash
bash run_grpo_industrial.sh
```

Office:

```bash
bash run_grpo_office.sh
```

### 3. DASO

Industrial:

```bash
bash run_daso_industrial.sh
```

Office:

```bash
bash run_daso_office.sh
```

### 4. Evaluation

Industrial:

```bash
bash run_eval_industrial.sh /path/to/checkpoint
```

Office:

```bash
bash run_eval_office.sh /path/to/checkpoint
```

## Paper-Facing Entrypoints

These are the main public scripts:

- `run_sft_industrial.sh`
- `run_sft_office.sh`
- `run_grpo_industrial.sh`
- `run_grpo_office.sh`
- `run_daso_industrial.sh`
- `run_daso_office.sh`
- `run_eval_industrial.sh`
- `run_eval_office.sh`

## Semantic-ID Construction

The semantic-ID pipeline is retained in this repo because it is part of the full paper workflow.

The high-level flow is:

1. preprocess raw Amazon data
2. encode item text into dense embeddings
3. train RQ-VAE on item embeddings
4. generate semantic-ID indices
5. convert the interaction data into the DASO training format

Relevant files:

- `data/amazon18_data_process.py`
- `data/amazon18_data_process.sh`
- `data/amazon23_data_process.py`
- `data/amazon23_data_process.sh`
- `rq/text2emb/amazon_text2emb.py`
- `rq/text2emb/amazon_text2emb.sh`
- `rq/rqvae.py`
- `rq/rqvae.sh`
- `rq/generate_indices.py`
- `convert_dataset.py`
- `convert_dataset.sh`

If you need to rebuild the whole pipeline from raw data, the rough command sequence is:

```bash
bash data/amazon18_data_process.sh
bash rq/text2emb/amazon_text2emb.sh
bash rq/rqvae.sh
python rq/generate_indices.py ...
bash convert_dataset.sh
```

The exact arguments depend on dataset paths and model checkpoints.

The helper wrapper can also be driven by environment variables:

```bash
DASO_DATASET_NAME=Office_Products \
DASO_INPUT_DIR=data/Amazon18/Office_Products \
DASO_OUTPUT_DIR=data/Amazon18 \
bash convert_dataset.sh
```

## Training Details

### SFT

`sft.py` trains the language model to predict the next item semantic ID from user history, together with auxiliary semantic-ID alignment data already built into the training dataset construction.

Default launch scripts:

- `run_sft_industrial.sh`
- `run_sft_office.sh`

Common script knobs:

- `BASE_MODEL`
- `OUTPUT_DIR`
- `CUDA_VISIBLE_DEVICES`
- `MASTER_ADDR`
- `MASTER_PORT`

### RL

`rl.py` is the public post-SFT RL entry.

Supported reward configurations in the public release:

- `rule`
- `ranking`

`GRPO` uses the standard group-relative policy optimization path.

`DASO` is exposed through:

```bash
--daso True --daso_alpha 0.5
```

In the current implementation, DASO is the paper-facing name for the difficulty-aware sibling-relative advantage training path built on top of the GRPO loop. The training logic is implemented in `minionerec_trainer.py`.

## Evaluation

Evaluation uses constrained decoding so generated outputs must stay within the valid semantic-ID space.

The evaluation path is:

1. split test data into per-GPU shards
2. decode with valid semantic-ID prefix constraints
3. merge shard outputs
4. compute ranking metrics
5. compute per-level semantic-ID accuracy statistics

Main files:

- `run_eval_industrial.sh`
- `run_eval_office.sh`
- `evaluate.py`
- `split.py`
- `merge.py`
- `calc.py`
- `calc_level.py`

By default, evaluation writes merged outputs under `./results/<checkpoint_name>/`.

## Important Notes

- This repo currently ships processed data for the two public Amazon subsets used in the paper workflow.
- The repository keeps semantic-ID construction code, SFT code, the GRPO baseline, and DASO.
- The repo has been cleaned to remove unrelated experimental branches so the public release reflects the paper more closely.
- Some internal class and file names still reflect the codebase history, but the exposed scripts and README are now aligned with the DASO release path.
- The top-level repo is now intentionally small: if a file is not directly useful for semantic-ID construction, SFT, GRPO, DASO, or constrained evaluation, it has likely been removed from the public release path.

## License

This repository keeps the original [LICENSE](./LICENSE). Before pushing the repo publicly, it is still worth doing one final manual pass to confirm that every included file is intended for redistribution.