Qwen example notebook

[peregilk]

I am trying out your Qwen example notebook. It links to Kaggle but as far I can see, there are no JAX/Flax checkpoints there. How do I convert/find these?

[vfdev-5]

You can download the model and extract to a folder using these commands in a separate cell:

!curl -L -o model.tar.gz https://www.kaggle.com/api/v1/models/qwen-lm/qwen-3/transformers/0.6b/1/download
!mkdir -p model_cp_path && tar -xvzf model.tar.gz -C model_cp_path

and change MODEL_CP_PATH = 'model_cp_path'

EDIT:
In case you see an error when using the sampler like:

ValueError: Incorrect number of leaves, expected 311 leaves, but got 1.

we may need to install the latest dependencies:

!pip install -U jax[cuda12]
!pip install -U git+https://github.com/google/flax
!pip install git+https://github.com/google/tunix