split completion text to sequence_len (#616)

This commit is contained in:
Wing Lian
2023-09-21 21:51:25 -04:00
committed by GitHub
parent 2844eb22b6
commit 97d3776ce6
4 changed files with 87 additions and 59 deletions

View File

@@ -38,10 +38,15 @@ class TokenizedPromptDataset(Dataset):
def process(self, dataset):
features = dataset.features.keys()
num_proc = min(64, os.cpu_count())
map_kwargs = {}
if self.prompt_tokenizer.supports_batched:
map_kwargs["batched"] = True
map_kwargs["batch_size"] = 100
return dataset.map(
self.prompt_tokenizer.tokenize_prompt,
num_proc=num_proc,
remove_columns=features,
**map_kwargs,
)