p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] train_tokenizer.py (python)
        print(f"Restored model state from step {step}")
        
        # Restore dataloader state if available
        try:
            dataloader_state_path = os.path.join(args.ckpt_dir, f"dataloader_state_{step}")
            if os.path.exists(dataloader_state_path):
                dataloader_iterator.load(dataloader_state_path)
                print(f"Restored dataloader state from step {step}")
        except Exception as e:
            print(f"Could not restore dataloader state: {e}")
    
    dataloader = (jax.make_array_from_process_local_data(videos_sharding, elem) for elem in dataloader_iterator) # type: ignore
[content] train_tokenizer.py (python)
        print(f"Restored dataloader and model state from step {step}")

    # --- TRAIN LOOP ---
    dataloader = iter(grain_dataloader)
    dataloader = (jax.make_array_from_process_local_data(videos_sharding, elem) for elem in dataloader) # type: ignore
[content] train_tokenizer.py (python)
                # Save dataloader state
                try:
                    dataloader_state_path = os.path.join(args.ckpt_dir, f"dataloader_state_{step}")
                    dataloader_iterator.save(dataloader_state_path)
                except Exception as e:
                    print(f"Could not save dataloader state: {e}")
[content] train_tokenizer.py (python)
    # --- Create iterator from dataloader ---
    dataloader = iter(grain_dataloader)
    dataloader = (jax.make_array_from_process_local_data(videos_sharding, elem) for elem in dataloader) # type: ignore
[content] train_tokenizer.py (python)
Create iterator from dataloader ---
    dataloader = iter(grain_dataloader)
    dataloader = (jax.make_array_from_process_local_data(videos_sharding, elem) for elem in dataloader) # type: ignore
    
    # ---
[content] train_tokenizer.py (python)

    # --- Create iterator from dataloader ---
    dataloader = iter(grain_dataloader)
    dataloader = (jax.make_array_from_process_local_data(videos_sharding, elem) for elem in dataloader) # type: ignore
[content] train_tokenizer.py (python)
    handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
    handler_registry.add('model_state', ocp.args.StandardSave)
    handler_registry.add('model_state', ocp.args.StandardRestore)
    handler_registry.add('dataloader_state', ocp.args.StandardSave, grain.checkpoint.CheckpointHandler())
    handler_registry.add('dataloader_state', ocp.args.StandardRestore, grain.checkpoint.CheckpointHandler())
[content] train_tokenizer.py (python)
    handler_registry.add('dataloader_state', ocp.args.StandardSave, grain.checkpoint.CheckpointHandler())
    handler_registry.add('dataloader_state', ocp.args.StandardRestore, grain.checkpoint.CheckpointHandler())
[content] train_tokenizer.py (python)
    handler_registry.add('model_state', ocp.args.StandardSave, ocp.handlers.StandardCheckpointHandler())
    handler_registry.add('model_state', ocp.args.StandardRestore, ocp.handlers.StandardCheckpointHandler())
    handler_registry.add('dataloader_state', grain.checkpoint.CheckpointSave, grain.checkpoint.CheckpointHandler())
    handler_registry.add('dataloader_state', grain.checkpoint.CheckpointRestore, grain.checkpoint.CheckpointHandler())
[content] train_tokenizer.py (python)
    handler_registry.add('model_state', ocp.args.StandardSave)
    handler_registry.add('model_state', ocp.args.StandardRestore)
    handler_registry.add('dataloader_state', ocp.args.StandardSave, grain.checkpoint.CheckpointHandler())
    handler_registry.add('dataloader_state', ocp.args.StandardRestore, grain.checkpoint.CheckpointHandler())
[content] train_tokenizer.py (python)
    # --- Create DataLoaderIterator from dataloader ---
    initial_state = grain_dataloader._create_initial_state()
    grain_iterator = grain.DataLoaderIterator(grain_dataloader, initial_state)
[content] train_tokenizer.py (python)
            args=ocp.args.Composite(
                model_state=train_state_restore_args,
                dataloader_state=grain_iterator,
            )
        )
        train_state = restored.model_state
        grain_dataloader = restored.dataloader_state
[content] train_tokenizer.py (python)
    handler_registry.add('model_state', ocp.args.StandardSave, ocp.handlers.StandardCheckpointHandler())
    handler_registry.add('model_state', ocp.args.StandardRestore, ocp.handlers.StandardCheckpointHandler())
    handler_registry.add('dataloader_state', grain.checkpoint.PyGrainCheckpointHandler())
[content] train_tokenizer.py (python)
    handler_registry.add('dataloader_state', grain.checkpoint.CheckpointSave, grain.checkpoint.CheckpointHandler())
    handler_registry.add('dataloader_state', grain.checkpoint.CheckpointRestore, grain.checkpoint.CheckpointHandler())
[content] train_tokenizer.py (python)
    handler_registry.add('dataloader_state', grain.checkpoint.CheckpointSave, grain.checkpoint.CheckpointHandler)
    handler_registry.add('dataloader_state', grain.checkpoint.CheckpointRestore, grain.checkpoint.CheckpointHandler)
[content] train_tokenizer.py (python)
    
    # --- Create DataLoaderIterator from dataloader ---
    initial_state = grain_dataloader._create_initial_state()
    grain_iterator = grain.DataLoaderIterator(grain_dataloader, initial_state)
    
    # --- Restore checkpoint ---
    if args.restore_ckpt:
        # Find the latest checkpoint
        checkpoint_files = [f for f in os.listdir(args.ckpt_dir) if f.startswith("tokenizer_") and f.endswith(".ckpt")]
        if checkpoint_files:
            latest_checkpoint = max(checkpoint_files, key=lambda x: int(x.split("_")[-1].split(".")[0]))
            checkpoint_path = os.path.join(args.ckpt_dir, latest_checkpoint)
            
            restore_target = {"model": train_state}
            restore_args = orbax_utils.restore_args_from_target(restore_target)
            restored = PyTreeCheckpointer().restore(
                checkpoint_path, 
                item=restore_target, 
                restore_args=restore_args
            )
            train_state = restored["model"]
            step = int(latest_checkpoint.split("_")[-1].split(".")[0])
            print(f"Restored model state from step {step}")
[content] train_tokenizer.py (python)
                ckpt = {"model": train_state}
                orbax_checkpointer = PyTreeCheckpointer()
                save_args = orbax_utils.save_args_from_target(ckpt)
                orbax_checkpointer.save(
                    os.path.join(args.ckpt_dir, f"tokenizer_{ts}_{step}.ckpt"),
                    ckpt,
                    save_args=save_args,
[content] train_tokenizer.py (python)
                train_state_save_args = orbax_utils.save_args_from_target(train_state)
                checkpoint_manager.save(
                    step,
                    args=ocp.args.Composite(
                        model_state=train_state_save_args,
                        dataloader_state=grain.checkpoint.CheckpointSave(grain_iterator),
                    )
[content] train_tokenizer.py (python)
        )
        train_state = restored["model_state"]
        grain_iterator = restored["dataloader_state"]
        step = checkpoint_manager.latest_step() or 0
        print(f"Restored dataloader and model state from step {step}")
[content] train_tokenizer.py (python)
        train_state_restore_args = orbax_utils.restore_args_from_target(train_state)
        
        restored = checkpoint_manager.restore(
            checkpoint_manager.latest_step(),
            args=ocp.args.Composite(
                model_state=train_state_restore_args,
                dataloader_state=grain.checkpoint.CheckpointRestore(grain_iterator),
[content] train_tokenizer.py (python)
    handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
    handler_registry.add('model_state', ocp.args.StandardSave, ocp.handlers.StandardCheckpointHandler())
    handler_registry.add('model_state', ocp.args.StandardRestore, ocp.handlers.StandardCheckpointHandler())
    handler_registry.add('dataloader_state', grain.checkpoint.CheckpointSave, grain.checkpoint.CheckpointHandler)
    handler_registry.add('dataloader_state', grain.checkpoint.CheckpointRestore, grain.checkpoint.CheckpointHandler)
    
    checkpoint_options = ocp.CheckpointManagerOptions(
        save_interval_steps=args.log_checkpoint_interval,
        max_to_keep=3,
        step_format_fixed_length=6,
        cleanup_tmp_directories=True,
    )
    
    checkpoint_manager = ocp.CheckpointManager(
        os.path.abspath(args.ckpt_dir),
        options=checkpoint_options,
        handler_registry=handler_registry,
    )
[content] train_tokenizer.py (python)
        # Find the latest checkpoint
        checkpoint_files = [f for f in os.listdir(args.ckpt_dir) if f.startswith("tokenizer_") and f.endswith(".ckpt")]
        if checkpoint_files:
            latest_checkpoint = max(checkpoint_files, key=lambda x: int(x.split("_")[-1].split(".")[0]))
            checkpoint_path = os.path.join(args.ckpt_dir, latest_checkpoint)
            
            restore_target = {"model": train_state}
            restore_args = orbax_utils.restore_args_from_target(restore_target)
            restored = PyTreeCheckpointer().restore(
                checkpoint_path, 
                item=restore_target, 
                restore_args=restore_args
[content] train_tokenizer.py (python)
                checkpoint_manager.save(
                    step,
                    args=ocp.args.Composite(
                        model_state=ocp.args.StandardSave(train_state),
                        dataloader_state=grain.checkpoint.CheckpointSave(grain_iterator),
                    )
[content] train_tokenizer.py (python)
    handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
    handler_registry.add('model_state', ocp.args.StandardSave, ocp.handlers.StandardCheckpointHandler())
    handler_registry.add('model_state', ocp.args.StandardRestore, ocp.handlers.StandardCheckpointHandler())
    handler_registry.add('dataloader_state', grain.checkpoint.CheckpointSave, grain.checkpoint.CheckpointHandler())
    handler_registry.add('dataloader_state', grain.checkpoint.CheckpointRestore, grain.checkpoint.CheckpointHandler())
    
    checkpoint_options = ocp.CheckpointManagerOptions(
        save_interval_steps=args.log_checkpoint_interval,
        max_to_keep=3,
        step_format_fixed_length=6,
        cleanup_tmp_directories=True,
    )
    
    checkpoint_manager = ocp.CheckpointManager(
        os.path.abspath(args.ckpt_dir),
        options=checkpoint_options,
        handler_registry=handler_registry,
    )
[content] train_tokenizer.py (python)
    
    # --- Restore checkpoint ---
    if args.restore_ckpt:
        train_state_restore_args = orbax_utils.restore_args_from_target(train_state)
        
        restored = checkpoint_manager.restore(
            checkpoint_manager.latest_step(),
            args=ocp.args.Composite(
                model_state=train_state_restore_args,
                dataloader_state=grain.checkpoint.CheckpointRestore(grain_iterator),
            )
        )
        train_state = restored["model_state"]
        grain_iterator = restored["dataloader_state"]
        step = checkpoint_manager.latest_step() or 0
        print(f"Restored dataloader and model state from step {step}")
[content] train_tokenizer.py (python)
    handler_registry.add('model_state', ocp.args.StandardSave, ocp.handlers.StandardCheckpointHandler)
    handler_registry.add('model_state', ocp.args.StandardRestore, ocp.handlers.StandardCheckpointHandler)
    handler_registry.add('dataloader_state', grain.checkpoint.CheckpointSave, grain.checkpoint.CheckpointHandler)
    handler_registry.add('dataloader_state', grain.checkpoint.CheckpointRestore, grain.checkpoint.CheckpointHandler)