(dataset_path: Union[str, List[str]])
| 30 | |
| 31 | |
| 32 | def load_pretrain_dataset(dataset_path: Union[str, List[str]]) -> Dict: |
| 33 | if type(dataset_path) is str: |
| 34 | dataset_path = [dataset_path] |
| 35 | |
| 36 | for p in dataset_path: |
| 37 | if not os.path.isdir(p): |
| 38 | if p.endswith(".gz") or p.endswith(".jsonl"): |
| 39 | print(f"loading from {p}") |
| 40 | yield from stream_jsonl(p) |
| 41 | else: |
| 42 | p_list = glob.glob(p + "/*") |
| 43 | for p_ in p_list: |
| 44 | if p_.endswith(".gz") or p_.endswith(".jsonl"): |
| 45 | print(f"loading from {p_}") |
| 46 | yield from stream_jsonl(p_) |
| 47 | |
| 48 | |
| 49 | def process_sample( |