gpt-neox

1
# Copyright (c) 2024, EleutherAI
2
#
3
# Licensed under the Apache License, Version 2.0 (the "License");
4
# you may not use this file except in compliance with the License.
5
# You may obtain a copy of the License at
6
#
7
#     http://www.apache.org/licenses/LICENSE-2.0
8
#
9
# Unless required by applicable law or agreed to in writing, software
10
# distributed under the License is distributed on an "AS IS" BASIS,
11
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
# See the License for the specific language governing permissions and
13
# limitations under the License.
14

15
from tools.datasets.corpora import prepare_dataset, DATA_DOWNLOADERS
16
import argparse
17

18
TOKENIZER_CHOICES = [
19
    "HFGPT2Tokenizer",
20
    "HFTokenizer",
21
    "GPT2BPETokenizer",
22
    "CharLevelTokenizer",
23
    "TiktokenTokenizer",
24
    "SPMTokenizer",
25
]
26
DATASET_CHOICES = [i for i in DATA_DOWNLOADERS.keys() if i != "pass"]
27

28

29
def get_args():
30
    parser = argparse.ArgumentParser(description="Download & preprocess neox datasets")
31
    parser.add_argument(
32
        "dataset",
33
        nargs="?",
34
        default="enwik8",
35
        help="name of dataset to download.",
36
        choices=DATASET_CHOICES,
37
    )
38
    parser.add_argument(
39
        "-t",
40
        "--tokenizer",
41
        default="GPT2BPETokenizer",
42
        choices=TOKENIZER_CHOICES,
43
        help=f'Type of tokenizer to use - choose from {", ".join(TOKENIZER_CHOICES)}',
44
    )
45
    parser.add_argument(
46
        "-d",
47
        "--data-dir",
48
        default=None,
49
        help=f"Directory to which to download datasets / tokenizer "
50
        f"files - defaults to ./data",
51
    )
52
    parser.add_argument(
53
        "-v", "--vocab-file", default=None, help=f"Tokenizer vocab file (if required)"
54
    )
55
    parser.add_argument(
56
        "-m", "--merge-file", default=None, help=f"Tokenizer merge file (if required)"
57
    )
58
    parser.add_argument(
59
        "-f",
60
        "--force-redownload",
61
        dest="force_redownload",
62
        default=False,
63
        action="store_true",
64
    )
65
    return parser.parse_args()
66

67

68
if __name__ == "__main__":
69
    args = get_args()
70
    prepare_dataset(
71
        dataset_name=args.dataset,
72
        tokenizer_type=args.tokenizer,
73
        data_dir=args.data_dir,
74
        vocab_file=args.vocab_file,
75
        merge_file=args.merge_file,
76
        force_redownload=args.force_redownload,
77
    )
78