EfficientIDS - MaxText Implementation

Transformer-based Sequential Recommendation System with Hierarchical Item Clustering

Migration from PAXml to MaxText

This is a port of EfficientIDS from PAXml (legacy) to MaxText for:

• ✅ JAX 0.6+ compatibility
• ✅ TPU v6e (Trillium) support
• ✅ Cleaner, more maintainable codebase
• ✅ Better performance and debugging

Directory Structure

efficientids_maxtext/
├── data/              # Data processing and loading
├── models/            # Model architectures
├── configs/           # Training configurations
├── metrics/           # Evaluation metrics
├── utils/             # Utilities and helpers
├── scripts/           # Training/eval scripts
├── tests/             # Unit tests
└── requirements.txt   # Dependencies

Quick Start

# Install dependencies
pip install -r requirements.txt

# Process MovieLens data
python data/process_movielens.py --output_dir ./data/ml1m_processed

# Train model
python scripts/train.py --config configs/qwen_config.py

# Evaluate
python scripts/eval.py --config configs/qwen_config.py --checkpoint_dir ./checkpoints

Requirements

• Python >= 3.10
• JAX >= 0.6.0
• Flax >= 0.10.0
• TPU v4/v5e/v6e

Original PAXml Implementation

See /Users/ritik.r/work/efficientids for the original PAXml version.