Add dataset_processor_path CLI knob for custom datasets

[py4]

Add a dataset_processor_path CLI/yaml knob that lets users plug in a custom process_data(dataset_name, model_tokenizer, template_config, tmvp_config, x) -> dict function from a user-provided Python file, instead of editing maxtext to support a new dataset shape.

Why: the built-in utils_rl.process_data is hardcoded for a small set of dataset schemas (GSM8K, etc.). For users running RL on custom datasets with different answer columns / cleaning rules, the alternative was either (1) edit maxtext source (fork divergence) or (2) reformat the dataset to look like GSM8K (lossy). This knob gives a clean third option: ship your dataset processor as a Python file and point maxtext at it.

Changes (2 files, +41/-16 lines):

• src/maxtext/trainers/post_train/rl/train_rl.py:
  • New _load_custom_callable(module_path, function_name) helper that uses importlib.util.spec_from_file_location to load a function from an arbitrary .py file (without adding to sys.path).
  • prepare_datasets checks trainer_config.dataset_processor_path; if set, loads process_data from that file and substitutes for utils_rl.process_data in the dataset pipeline.
• src/maxtext/configs/post_train/rl.yml: new top-level knob dataset_processor_path: '' with comment documenting the signature contract.

Backward compatible: default empty string falls back to utils_rl.process_data (identical to old behavior). The _load_custom_callable helper is only invoked when the user explicitly sets the path.

User-facing contract:

# user_process_data.py
def process_data(dataset_name, model_tokenizer, template_config, tmvp_config, x):
    return {"prompts": ..., "question": ..., "answer": ...}

python3 -m maxtext.trainers.post_train.rl.train_rl rl.yml \
  dataset_processor_path=/path/to/user_process_data.py \
  ...

Checklist

• Tested locally with a custom processor file (VTC-style raw-text prompt template); produced expected outputs
• Backward compatible: default empty string preserves utils_rl.process_data behavior
• No effect on non-RL paths (only prepare_datasets in the RL trainer touched)
• _load_custom_callable doesn't pollute sys.path (uses spec_from_file_location)

[codecov]

Codecov Report

❌ Patch coverage is 73.68421% with 5 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/trainers/post_train/rl/train_rl.py	40.00%	2 Missing and 1 partial ⚠️
src/maxtext/trainers/post_train/rl/utils_rl.py	85.71%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[khatwanimohit]

LGTM

[A9isha]

Just one request for the refactor - the rest look good, thank you Pooya!