google-deepmind / google-deepmind/gemma

Refactor DPO dataset construction to remove hardcoded parameters in and TODOs in dpo.py in examples

Open
#498 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
5.7k
Forks
1k
Avg merge
10h 33m
Merged PRs (30d)
2

Description

**Description**
In **examples/dpo.py**, the `_make_dataset` function currently relies on hardcoded values for `max_length` (512) and `batch_size` (16). There is also an outstanding TODO comment (`# TODO(epot): !!!!`) explicitly flagging this section for cleanup.

This differs from other examples like `seq2seq.py` and `lora.py`, where these parameters are defined in `get_config()` and passed down as arguments. This inconsistency makes it harder to configure the DPO training run without modifying the internal logic of the dataset builder.

**Solution**
I propose refactoring `examples/dpo.py` to match the pattern used in the other example scripts:
1. Move the `batch_size` and `max_length` definitions to the top of `get_config()`.
2. Update `_make_dataset` to accept these values as keyword arguments (`*`, `training`, `batch_size`, `max_length`).
3. Remove the placeholder TODO comment.

**Alternatives I've considered**
Leaving it as-is works functionality-wise, but it leaves technical debt in the codebase and makes the examples inconsistent for new users trying to learn the library.

**Additional context**
I am happy to open a PR to standardize this example.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.