pixi.toml

[workspace]
authors = ["Michael Booth <michael@databooth.com.au>"]
channels = ["conda-forge", "https://conda.modular.com/max-nightly"]
name = "max-learning"
platforms = ["osx-arm64", "linux-64"]
version = "0.3.0"

[activation]
env = { MOJO_PATH = "$PIXI_PROJECT_ROOT/src:$HOME/code/github/databooth/mojo-toml/src:$HOME/code/github/databooth/mojo-dotenv/src" }

[dependencies]
modular = "==26.2.0.dev2026012505"  # Locked to prevent breaking changes
transformers = ">=4.57.3,<5"
pytorch = ">=2.9.1,<3"
torchvision = ">=0.22.1,<1"
optimum = ">=2.1.0,<3"
onnxruntime = ">=1.22.2,<2"
onnxscript = ">=0.5.6,<0.6"
numpy = ">=2.4.0,<3"
pytest = ">=9.0.2,<10"
tqdm = ">=4.67.1,<5"
scikit-learn = ">=1.6.1,<2"
pre-commit = ">=4.5.1,<5"
ruff = ">=0.8.6,<1"

[pypi-dependencies]
max-learning-python = { path = "src/python", editable = true }

[tasks]
# Setup
download-models = "bash models/download_models.sh"

# Python examples (MAX Graph) - numbered progression
example-elementwise-cpu = "python examples/python/01_elementwise/elementwise.py --device cpu"
example-elementwise-gpu = "python examples/python/01_elementwise/elementwise.py --device gpu"
example-linear = "python examples/python/02_linear_layer/linear_layer.py"
example-mlp = "python examples/python/03_mlp_regression/mlp_regression.py"
example-cnn = "python examples/python/04_cnn_mnist/cnn_mnist.py"
example-distilbert = { depends-on = ["download-models"], cmd = "python examples/python/03_distilbert_sentiment/distilbert_sentiment.py" }

# Benchmarking
benchmark-elementwise = "python benchmarks/01_elementwise/cpu_vs_gpu.py"
benchmark-elementwise-sizes = "python benchmarks/01_elementwise/cpu_vs_gpu_scaling.py"
benchmark-linear = "python benchmarks/02_linear_layer/cpu_vs_gpu.py"
benchmark-mlp = "python benchmarks/03_mlp/max_vs_pytorch.py"
benchmark-cnn = "python benchmarks/04_cnn/max_vs_pytorch.py"
benchmark-distilbert = { depends-on = ["download-models"], cmd = "python benchmarks/03_distilbert/max_vs_pytorch.py" }
benchmark-all = { depends-on = ["benchmark-elementwise", "benchmark-linear", "benchmark-mlp", "benchmark-cnn", "benchmark-distilbert"] }

# Mojo lexicon baseline (v0.1.0, non-MAX Graph)
mojo-build = "mojo build -I $HOME/code/github/databooth/mojo-toml/src -I $HOME/code/github/databooth/mojo-dotenv/src examples/mojo/lexicon_baseline/main.mojo -o lexicon-classifier"
mojo-run = "mojo run -I $HOME/code/github/databooth/mojo-toml/src -I $HOME/code/github/databooth/mojo-dotenv/src examples/mojo/lexicon_baseline/main.mojo"

# Testing
# test-mojo = "mojo run tests/mojo/test_*.mojo"  # Disabled: no mojo tests yet (uses TestSuite)
test-python = { depends-on = ["download-models"], cmd = "pytest tests/python/" }
test-all = { depends-on = ["test-python"] }  # Add test-mojo when implemented

# Development tools
mojo-version = "mojo --version"
mojo-format = "mojo format src/mojo/ tests/mojo/"
mojo-check = "mojo check src/mojo/ tests/mojo/"

# Code quality tasks
pre-commit = "pre-commit run --all-files"
pre-commit-install = "pre-commit install"
ruff-check = "ruff check ."
ruff-fix = "ruff check . --fix"
ruff-format = "ruff format ."

# Cleanup
clean-reports-md = "bash benchmarks/clean_reports.sh md"
clean-reports-json = "bash benchmarks/clean_reports.sh json"
clean-reports-csv = "bash benchmarks/clean_reports.sh csv"
clean-reports-all = "bash benchmarks/clean_reports.sh all"