Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add option to check dataset labels in SFTTrainer #1414

Closed
wants to merge 8 commits into from
Closed

Add option to check dataset labels in SFTTrainer #1414

Changes from 2 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
3d1f68a
add option check_dataset_labels and log tokens and labels of first pr…
geronimi73 Mar 10, 2024
8dc5390
removed nonsense logger call
geronimi73 Mar 10, 2024
3853037
add docstring, default check_dataset_labels=False
geronimi73 Mar 11, 2024
9632fee
Update trl/trainer/sft_trainer.py
geronimi73 Mar 12, 2024
d435e41
spaces for ruff
geronimi73 Mar 12, 2024
a760a17
add logger
geronimi73 Mar 13, 2024
24a1d3e
style
geronimi73 Mar 14, 2024
cdb865c
Update trl/trainer/sft_trainer.py
geronimi73 Mar 15, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions trl/trainer/sft_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -149,6 +149,7 @@ def __init__(
model_init_kwargs: Optional[Dict] = None,
dataset_kwargs: Optional[Dict] = None,
eval_packing: Optional[bool] = None,
check_dataset_labels: Optional[bool] = None,
):
if model_init_kwargs is None:
model_init_kwargs = {}
Expand Down Expand Up @@ -302,6 +303,15 @@ def make_inputs_require_grad(module, input, output):
"overflow issues when training a model in half-precision. You might consider adding `tokenizer.padding_side = 'right'` to your code."
)

if check_dataset_labels:
if train_dataset is not None and len(train_dataset) > 0:
input_ids, attention_mask, labels = data_collator([train_dataset[0]]).values()
# print is obviously the wrong choice but no logger
print(f"check_dataset_labels:")
print(tokenizer.decode(input_ids[0]))
for token, label in zip(input_ids[0], labels[0]):
print(token.item(), f"'{tokenizer.decode(token)}'", label.item())

super().__init__(
model=model,
args=args,
Expand Down