Skip to content

fixed bug in run_mlm_flax_stream.py #17203

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

Merged
merged 8 commits into from
May 16, 2022
Merged

fixed bug in run_mlm_flax_stream.py #17203

Changes from 5 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
92e120f
fixed bug run_mlm_flax_stream.py
KennethEnevoldsen May 12, 2022
06c0fcd
Update run_mlm_flax_stream.py
KennethEnevoldsen May 12, 2022
63cf612
adding missing paranthesis
KennethEnevoldsen May 12, 2022
147eb67
formatted to black
May 12, 2022
a5aa1e8
remove cols from dataset instead
May 12, 2022
4ea3407
reformat to black
May 12, 2022
0339a10
moved rem. columns to map
May 15, 2022
d50700e
formatted to black
May 15, 2022
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
8 changes: 6 additions & 2 deletions examples/research_projects/jax-projects/dataset-streaming/run_mlm_flax_stream.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -280,8 +280,10 @@ def advance_iter_and_group_samples(train_iterator, num_samples, max_seq_length):
tokenized_samples = next(train_iterator)
i += len(tokenized_samples["input_ids"])

# concatenate tokenized samples to list
samples = {k: samples[k] + tokenized_samples[k] for k in tokenized_samples.keys()}
# concatenate tokenized samples to list (excluding "id" and "text")
samples = {
k: samples[k] + tokenized_samples[k] for k in ["input_ids", "attention_mask", "special_tokens_mask"]
}

# Concatenated tokens are split to lists of length `max_seq_length`.
# Note that remainedr of % max_seq_length are thrown away.
Expand Down Expand Up @@ -399,10 +401,12 @@ def write_eval_metric(summary_writer, eval_metrics, step):
def tokenize_function(examples):
return tokenizer(examples[data_args.text_column_name], return_special_tokens_mask=True)


tokenized_datasets = dataset.map(
tokenize_function,
batched=True,
)
tokenized_datasets.remove_columns(dataset.features.keys())

shuffle_seed = training_args.seed
tokenized_datasets = tokenized_datasets.shuffle(buffer_size=data_args.shuffle_buffer_size, seed=shuffle_seed)
Expand Down