| 16 | |
| 17 | |
| 18 | def encode_file( |
| 19 | tokenizer, |
| 20 | data_path, |
| 21 | max_length, |
| 22 | pad_to_max_length=True, |
| 23 | return_tensors="pt", |
| 24 | overwrite_cache=False, |
| 25 | prefix="", |
| 26 | tok_name="", |
| 27 | ): |
| 28 | cache_path = Path(f"{data_path}_{tok_name}{max_length}.pt") |
| 29 | if not overwrite_cache and cache_path.exists(): |
| 30 | try: |
| 31 | examples = torch.load(cache_path) |
| 32 | assert isinstance(examples, list) |
| 33 | return examples |
| 34 | |
| 35 | except Exception: |
| 36 | print(f"failed to load from {cache_path}, retokenizing {data_path}") |
| 37 | data_path = Path(data_path) |
| 38 | |
| 39 | lns = lmap(str.strip, data_path.open().readlines()) |
| 40 | lns = [prefix + text for text in lns] |
| 41 | assert lns, f"found empty file at {data_path}" |
| 42 | examples = [] |
| 43 | for text in tqdm(lns, desc=f"Tokenizing {data_path.name}"): |
| 44 | tokenized = tokenizer( |
| 45 | [text], |
| 46 | max_length=max_length, |
| 47 | padding="max_length" if pad_to_max_length else None, |
| 48 | truncation=True, |
| 49 | add_prefix_space=True, |
| 50 | return_tensors=return_tensors, |
| 51 | ) |
| 52 | assert tokenized.input_ids.shape[1] == max_length |
| 53 | examples.append(tokenized) |
| 54 | torch.save(lmap(dict, examples), cache_path.open("wb")) |
| 55 | return examples |
| 56 | |
| 57 | |
| 58 | def lmap(f: Callable, x: Iterable) -> List: |