This load_with_retry function is duplicated across multiple notebooks. To improve maintainability, consider moving this and other shared utility functions into a common Python script that can be imported by all notebooks.

Additionally, the caching logic is not robust. It uses a fixed DATASET_CACHE path regardless of the dataset path argument. This can lead to cache collisions if different datasets are used across notebooks in the same environment. A better approach would be to create a unique cache directory for each dataset, for example by using a hash of the dataset path.

Example for a more robust caching mechanism:

import hashlib

def get_cache_path(dataset_path):
    hash_object = hashlib.md5(dataset_path.encode())
    return f"/kaggle/working/_dataset_cache_{hash_object.hexdigest()}"

# In load_with_retry:
dataset_cache_path = get_cache_path(path)
if os.path.exists(dataset_cache_path):
    # ... load from disk
...
ds.save_to_disk(dataset_cache_path)

The use of global variables (PRINTED_TIMES, PRINT_EVERY_STEPS) makes this function stateful and harder to test and reuse. This can lead to unexpected behavior if the function is called multiple times or from different places.

A better approach would be to manage this state within a class or by using a closure. For example, you could wrap the reward function in a class:

class RewardLogger:
    def __init__(self, print_every_steps=5):
        self.printed_times = 0
        self.print_every_steps = print_every_steps

    def check_numbers(self, prompts, completions, answer, **kwargs):
        # ... (rest of the function logic)
        if self.printed_times % self.print_every_steps == 0:
            # ... print statement
        self.printed_times += 1
        # ...
        return scores

reward_logger = RewardLogger()
# in GRPOTrainer:
# reward_funcs = [..., reward_logger.check_numbers]

This encapsulates the state and makes the code cleaner and more maintainable.

Patching trainer.args after the GRPOTrainer has been initialized can be confusing and error-prone. The arguments are used during the trainer's __init__. Modifying them afterwards might not have the intended effect for all parameters.

It would be clearer and safer to configure all training arguments directly in the GRPOConfig object before passing it to the trainer. This makes the configuration explicit and centralized.

For example, move the patched arguments to the GRPOConfig in the preceding cell:

training_args = GRPOConfig(
    learning_rate = 5e-6,
    # ... other args
    per_device_train_batch_size = 1,
    gradient_accumulation_steps = 32,
    gradient_checkpointing = False,
    fp16 = True,
    bf16 = False,
    dataloader_pin_memory = True,
    dataloader_num_workers = 2,
    ddp_find_unused_parameters = False,
    max_grad_norm = 0.3,
    # ...
)

The regex used to parse the torch version is fragile and might fail on development or more complex version strings (e.g., 2.1.0.dev20230310+cu118). A more robust way to get the major.minor version would be to use string splitting, which is also more readable.

For example:

import torch
v = '.'.join(torch.__version__.split('.')[:2])

This Jinja template has a syntax error. The {% endif %} on the third line of the template string does not have a corresponding {% if ... %} block. This will raise a TemplateSyntaxError if this code block is executed. Although it's currently under an if False block, it's best to fix it for future use. It seems to be a copy-paste error and should likely be removed.

Defining a large and complex class like DeepSeekOCR2DataCollator inside a notebook cell harms readability and reusability. This class is also duplicated in other notebooks, which will make maintenance difficult.

Consider moving this class to a separate utils.py file and importing it into the notebooks where it's needed. This will make the notebooks cleaner and easier to follow, and any changes to the data collator will only need to be made in one place.