| from datasets import load_dataset | |
| from tokenizers import trainers, Tokenizer, normalizers, ByteLevelBPETokenizer | |
| import os | |
| # export DATASET="${HOME}/data/nedd_wiki_news/nedd_wiki_news.py" # Name of the dataset in the Huggingface Hub | |
| # export DATASET_CONFIG="nedd_nl" # Config of the dataset in the Huggingface Hub | |
| # export DATASET_SPLIT="train" # Split to use for training tokenizer and model | |
| # export VOCAB_SIZE="50257" | |
| # export MODEL_PATH="${HOME}/data/${HF_PROJECT}" # Path to the model, e.g. here inside the mount | |
| dataset_name = os.environ.get("DATASET") | |
| dataset_config = os.environ.get("DATASET_CONFIG") | |
| dataset_split = os.environ.get("DATASET_SPLIT") | |
| vocab_size = int(os.environ.get("VOCAB_SIZE")) | |
| model_path = os.environ.get("MODEL_PATH") | |
| # load dataset | |
| dataset = load_dataset(dataset_name, dataset_config, split=dataset_split) | |
| # Instantiate tokenizer | |
| tokenizer = ByteLevelBPETokenizer() | |
| def batch_iterator(batch_size=100): | |
| for i in range(0, len(dataset), batch_size): | |
| yield dataset[i: i + batch_size]["text"] | |
| # Customized training | |
| tokenizer.train_from_iterator(batch_iterator(), vocab_size=vocab_size, min_frequency=2, special_tokens=[ | |
| "<s>", | |
| "<pad>", | |
| "</s>", | |
| "<unk>", | |
| "<mask>", | |
| ]) | |
| # Save files to disk | |
| tokenizer.save("tokenizer.json") | |
| tokenizer.save_model(".") | |