(max_sample)
| 95 | return sources, targets |
| 96 | |
| 97 | def get_wiki_dataset(max_sample): |
| 98 | wiki_dataset = load_dataset("wikitext", 'wikitext-2-raw-v1', split='train') |
| 99 | # wiki_dataset = load_dataset("/root/model/datasets/wikitext/wikitext", 'wikitext-2-raw-v1', split='train') |
| 100 | |
| 101 | wiki_long = [] |
| 102 | for text in wiki_dataset['text']: |
| 103 | if len(text) > 512: |
| 104 | wiki_long.append(text) |
| 105 | wiki_front = [long[:128] for long in wiki_long] |
| 106 | |
| 107 | targets = sources = wiki_front |
| 108 | |
| 109 | return extract_random_dataset(sources, targets, max_sample) |
| 110 | |
| 111 | def get_redpajama_dataset(max_sample): |
| 112 | # wiki_dataset = load_dataset("wikitext", 'wikitext-2-raw-v1', split='train') |
no test coverage detected