Error while running main.py - Most likely missing chkpoint location

[shivajid]

I ran into couple of issues:

• Testing this on v6e-8

It requires two packages perhaps we should have requirements.txt
Ran

"python3 main.py"

Got the following error. Do I need to put the checkpoint in certain location

Traceback (most recent call last):
  File "/home/shivajid/jax-llm-examples/deepseek_r1_jax/main.py", line 42, in <module>
    tokenizer = dsjax.load_tokenizer(ckpt_path / "tokenizer.json", ckpt_path / "tokenizer_config.json")
  File "/home/shivajid/jax-llm-examples/deepseek_r1_jax/deepseek_r1_jax/model.py", line 192, in load_tokenizer
    config = json.loads(tokenizer_config_path.expanduser().resolve().read_text())
  File "/home/shivajid/miniconda3/envs/vllm/lib/python3.10/site-packages/etils/epath/abstract_path.py", line 150, in read_text
    with self.open('r', encoding=encoding) as f:
  File "/home/shivajid/miniconda3/envs/vllm/lib/python3.10/site-packages/etils/epath/gpath.py", line 257, in open
    gfile = self._backend.open(self._path_str, mode)
  File "/home/shivajid/miniconda3/envs/vllm/lib/python3.10/site-packages/etils/epath/backend.py", line 133, in open
    return open(path, mode, encoding=encoding)
FileNotFoundError: [Errno 2] No such file or directory: '/home/shivajid/bucket/deepseek-r1-jax-chkpt/tokenizer_config.json'

[rdyro]

That's right, you need to convert the checkpoint from huggingface first since we're not distributing the JAX checkpoint at the moment.

Take a look at: https://github.com/jax-ml/jax-llm-examples/blob/main/deepseek_r1_jax/scripts/convert_hf_r1_checkpoint.py

You can modify the root_path to point to the huggingface checkpoint directory and dest_path to where you want to output the converted JAX checkpoint.

Let me know if you run into any issues!