Skip to content
BranchesTags
 
 

Repository files navigation

OneRec: Unified Recommendation Framework

A modular framework for generative recommendation with 4 orthogonal dimensions:

Dimension Options Description
Backbone gemma, qwen, t5, gpt2 Transformer model for encoding
Head embedding, sid Prediction strategy
Reasoning on, off Whether to generate reasoning
Trainer supervised, recpo Training algorithm

Quick Start

# Install dependencies
pip install -r requirements.txt

# Train with different configurations:

# RRec original: Gemma + Embedding + Reasoning ON + RecPO
python scripts/train.py \
    --backbone gemma \
    --head embedding \
    --reasoning on \
    --trainer recpo \
    --dataset_dir data/Musical_Instruments

# TIGER style: T5 + SID + Reasoning OFF + Supervised  
python scripts/train.py \
    --backbone t5 \
    --head sid \
    --reasoning off \
    --trainer supervised \
    --sid_codebook_path data/MI/codebook.pt

# New combination: Gemma + SID + Reasoning ON + RecPO
python scripts/train.py \
    --backbone gemma \
    --head sid \
    --reasoning on \
    --trainer recpo \
    --sid_codebook_path data/MI/codebook.pt

Architecture

┌─────────────────────────────────────────────────────────────┐
│                           Pipeline                          │
│                                                             │
│  ┌──────────────┐   ┌──────────────┐   ┌──────────────┐    │
│  │  Backbone    │   │  Reasoning   │   │    Head      │    │
│  │ (Dimension 1)│──▶│ (Dimension 3)│──▶│ (Dimension 2)│    │
│  │              │   │  (Optional)  │   │              │    │
│  │ gemma/qwen/  │   │ <think>...   │   │ embedding/   │    │
│  │ t5/gpt2      │   │ </think>     │   │ sid          │    │
│  └──────────────┘   └──────────────┘   └──────────────┘    │
│                           │                                 │
│                           ▼                                 │
│                 ┌──────────────────┐                        │
│                 │     Trainer      │                        │
│                 │   (Dimension 4)  │                        │
│                 │ supervised/recpo │                        │
│                 └──────────────────┘                        │
└─────────────────────────────────────────────────────────────┘

Directory Structure

RRec/
├── src/
│   ├── backbones/          # Dimension 1: Transformer models
│   │   ├── gemma.py        # Gemma-2-2b-it
│   │   ├── qwen.py         # Qwen2.5-3B-Instruct
│   │   ├── t5.py           # T5 Encoder
│   │   └── gpt2.py         # GPT-2
│   │
│   ├── heads/              # Dimension 2: Prediction heads
│   │   ├── embedding_head.py  # Dot-product similarity
│   │   └── sid_head.py        # Semantic ID generation
│   │
│   ├── reasoning/          # Dimension 3: Reasoning module
│   │   ├── generator.py    # vLLM/HF generation
│   │   └── prompt_builder.py
│   │
│   ├── trainers/           # Dimension 4: Training algorithms
│   │   ├── base_trainer.py
│   │   ├── supervised_trainer.py
│   │   └── recpo_trainer.py
│   │
│   ├── sid/                # SID learning (separate stage)
│   │   ├── residual_quantization.py
│   │   └── mini_batch_kmeans.py
│   │
│   ├── data/               # Data loading
│   │   ├── collators/
│   │   └── prompters/
│   │
│   └── utils/              # Utilities
│
├── scripts/
│   └── train.py            # Main entry point
│
└── requirements.txt

Combination Matrix

All 16 combinations are valid:

Backbone Head Reasoning Trainer Notes
gemma embedding on recpo RRec original
gemma embedding off supervised Baseline
gemma sid on recpo New: LLM + SID + reasoning
t5 sid off supervised ≈ GRID TIGER
qwen embedding on recpo RRec with Qwen
... ... ... ... Any combination

Key Arguments

# Backbone
--backbone gemma          # gemma | qwen | t5 | gpt2
--model_path /path/to/model
--use_lora True
--lora_r 4

# Head
--head embedding          # embedding | sid
--sid_codebook_path path  # Required for SID head
--similarity_type dot     # dot | cosine | L2

# Reasoning
--reasoning on            # on | off
--max_new_tokens 300
--group_size 4            # Samples per input (RecPO)
--use_vllm True

# Trainer
--trainer recpo           # supervised | recpo
--epsilon_low 0.2         # PPO clipping (RecPO)
--advantage_type gaussian

# General
--dataset_dir data/MI
--dataset_category Musical_Instruments
--num_train_epochs 10
--learning_rate 1e-5

Results

All experiments use Gemma-2B backbone on Musical Instruments dataset. Metrics are reported on the test set.

SID versus Single Embedding

  • Supervised -- No thinking, directly generate the answer
  • RecPO -- Thinking
Method NDCG@5 NDCG@10 Recall@5 Recall@10
Single Embedding + Supervised 0.0412 0.0538 0.0587 0.0891
SID + Supervised 0.0083 0.0112 0.0121 0.0195
Single Embedding + RecPO 0.0518 0.0659 0.0721 0.1073
SID + RecPO 0.0097 0.0134 0.0148 0.0226

Thinking versus No Thinking (Single Embedding)

Method NDCG@5 NDCG@10 Recall@5 Recall@10
No Thinking + Supervised 0.0412 0.0538 0.0587 0.0891
Thinking + RecPO 0.0518 0.0659 0.0721 0.1073

Takeaways and Next Steps

  • Reasoning leads to performance improvement --> Build distillation functions
  • SID performs bad --> Compare with Clark's implementation to identify the cause

About

No description, website, or topics provided.

Resources

Readme
License
Activity

Stars

0 stars

Watchers

0 watching

Forks

0 forks
Report repository

Releases

Packages

Contributors

Languages