aapot commited on
Commit
12280fc
1 Parent(s): 780cf63

Saving weights and logs of step 10000

Browse files
added_tokens.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {"<|endoftext|>": 50257}
config.json ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "./",
3
+ "activation_function": "gelu_new",
4
+ "architectures": [
5
+ "GPT2LMHeadModel"
6
+ ],
7
+ "attn_pdrop": 0.1,
8
+ "bos_token_id": 50256,
9
+ "embd_pdrop": 0.1,
10
+ "eos_token_id": 50256,
11
+ "initializer_range": 0.02,
12
+ "layer_norm_epsilon": 1e-05,
13
+ "model_type": "gpt2",
14
+ "n_ctx": 1024,
15
+ "n_embd": 768,
16
+ "n_head": 12,
17
+ "n_inner": null,
18
+ "n_layer": 12,
19
+ "n_positions": 1024,
20
+ "reorder_and_upcast_attn": false,
21
+ "resid_pdrop": 0.1,
22
+ "scale_attn_by_inverse_layer_idx": false,
23
+ "scale_attn_weights": true,
24
+ "summary_activation": null,
25
+ "summary_first_dropout": 0.1,
26
+ "summary_proj_to_labels": true,
27
+ "summary_type": "cls_index",
28
+ "summary_use_proj": true,
29
+ "task_specific_params": {
30
+ "text-generation": {
31
+ "do_sample": true,
32
+ "max_length": 50
33
+ }
34
+ },
35
+ "transformers_version": "4.16.0.dev0",
36
+ "use_cache": true,
37
+ "vocab_size": 50257
38
+ }
flax_model.msgpack ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:4bf5c41eb8e8bd43c6b368f4ae4f338e25ddf4fbf583ad377e3d85ecd6610d92
3
+ size 497764120
merges.txt ADDED
The diff for this file is too large to render. See raw diff
 
replace_token_script.py ADDED
@@ -0,0 +1,80 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ''''This script was used to replace the final index of tokenizer.json and vocab.json
2
+ with "<|endoftext|>" token. Also reassociate the corresponding merges'''
3
+
4
+ import json
5
+
6
+ tokenizer_path = 'tokenizer.json'
7
+ model_config_path = 'config.json'
8
+ vocab_path = 'vocab.json'
9
+
10
+ with open(vocab_path, "r") as f:
11
+ vocab_data = json.load(f)
12
+
13
+ with open(tokenizer_path, "r") as f:
14
+ tokenizer_data = json.load(f)
15
+
16
+ with open(model_config_path, "r") as f:
17
+ model_config = json.load(f)
18
+
19
+ model_vocab_size = model_config['vocab_size']
20
+ tokenizer_vocab = tokenizer_data['model']['vocab']
21
+
22
+ mergeslength = len(tokenizer_data['model']['merges'])
23
+
24
+ #readjust added_tokens 'id' to model_vocab_size - 1
25
+ tokenizer_data['added_tokens'][-1]['id'] = model_vocab_size - 1
26
+
27
+ final_index = model_vocab_size - 1
28
+ eos = '<|endoftext|>'
29
+
30
+ #retrieve the key of final index
31
+ old_key_final_index_tokenizer = list(tokenizer_data['model']['vocab'].keys())[final_index]
32
+ old_key_final_index_vocab = list(vocab_data.keys())[final_index]
33
+ old_key_final_index_vocab_min2 = list(vocab_data.keys())[final_index - 1]
34
+ old_key_final_index_tokenizer_merges = tokenizer_data['model']['merges'][mergeslength - 1]
35
+
36
+ print(f"old_key_final_index_tokenizer = {old_key_final_index_tokenizer}")
37
+ print(f"old_key_final_index_vocab = {old_key_final_index_vocab}")
38
+ print(f"old_key_final_index_vocab_min2 = {old_key_final_index_vocab_min2}")
39
+ print(f"old_key_final_index_tokenizer_merges = {old_key_final_index_tokenizer_merges}")
40
+
41
+ #replace old key with new key
42
+ tokenizer_data['model']['vocab']['<|endoftext|>'] = tokenizer_data['model']['vocab'][old_key_final_index_tokenizer]
43
+ vocab_data[eos] = vocab_data[old_key_final_index_vocab]
44
+
45
+ #replace the final merges idx with vocab_data - 1
46
+ tokenizer_data['model']['merges'] = tokenizer_data['model']['merges'][: mergeslength - 1]
47
+
48
+
49
+ #delete old key
50
+ del tokenizer_data['model']['vocab'][old_key_final_index_tokenizer]
51
+ del vocab_data[old_key_final_index_vocab]
52
+
53
+ #check updated key
54
+ old_key_final_index_tokenizer = list(tokenizer_data['model']['vocab'].keys())[final_index]
55
+ old_key_final_index_vocab = list(vocab_data.keys())[final_index]
56
+ old_key_final_index_tokenizer_merges = tokenizer_data['model']['merges'][mergeslength - 2]
57
+
58
+ print(len(tokenizer_data['model']['merges']))
59
+ print()
60
+ print(f"updated old_key_final_index_tokenizer = {old_key_final_index_tokenizer}")
61
+ print(f"updated old_key_final_index_vocab = {old_key_final_index_vocab}")
62
+ print(f"updated old_key_final_index_tokenizer_merges = {old_key_final_index_tokenizer_merges}")
63
+
64
+ with open(tokenizer_path, "w")as f:
65
+ json.dump(tokenizer_data, f)
66
+
67
+ with open(vocab_path, "w")as f:
68
+ json.dump(vocab_data, f)
69
+
70
+ with open('merges.txt') as f:
71
+ lines = f.readlines()
72
+
73
+ with open("merges.txt", "w") as f:
74
+ for i in range(len(lines) - 1):
75
+ f.write(lines[i])
76
+
77
+ with open('merges.txt') as f:
78
+ newlines = f.readlines()
79
+
80
+ print(f"newlines[len(newlines) - 1] = {newlines[len(newlines) - 1]}")
run_clm_flax.py ADDED
@@ -0,0 +1,839 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python
2
+ # coding=utf-8
3
+ # Copyright 2021 The HuggingFace Team All rights reserved.
4
+ #
5
+ # Licensed under the Apache License, Version 2.0 (the "License");
6
+ # you may not use this file except in compliance with the License.
7
+ # You may obtain a copy of the License at
8
+ #
9
+ # http://www.apache.org/licenses/LICENSE-2.0
10
+ #
11
+ # Unless required by applicable law or agreed to in writing, software
12
+ # distributed under the License is distributed on an "AS IS" BASIS,
13
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14
+ # See the License for the specific language governing permissions and
15
+ # limitations under the License.
16
+ """
17
+ Pre-training/Fine-tuning the library models for causal language modeling (GPT, GPT-2, CTRL, ...) on a text file or a dataset.
18
+
19
+ Here is the full list of checkpoints on the hub that can be fine-tuned by this script:
20
+ https://huggingface.co/models?filter=text-generation
21
+ """
22
+ # You can also adapt this script on your own causal language modeling task. Pointers for this are left as comments.
23
+
24
+ import json
25
+ import logging
26
+ import math
27
+ import os
28
+ import sys
29
+ import time
30
+ import gc
31
+ from dataclasses import asdict, dataclass, field
32
+ from enum import Enum
33
+ from itertools import chain
34
+ from pathlib import Path
35
+ from typing import Callable, Optional
36
+
37
+ import datasets
38
+ import numpy as np
39
+ from datasets import Dataset, load_dataset, load_from_disk
40
+ from tqdm import tqdm
41
+
42
+ import jax
43
+ import jax.numpy as jnp
44
+ import optax
45
+ import transformers
46
+ from flax import jax_utils, traverse_util
47
+ from flax.jax_utils import unreplicate
48
+ from flax.training import train_state
49
+ from flax.training.common_utils import get_metrics, onehot, shard, shard_prng_key
50
+ from huggingface_hub import Repository
51
+ from transformers import (
52
+ CONFIG_MAPPING,
53
+ FLAX_MODEL_FOR_CAUSAL_LM_MAPPING,
54
+ AutoConfig,
55
+ AutoTokenizer,
56
+ FlaxAutoModelForCausalLM,
57
+ HfArgumentParser,
58
+ is_tensorboard_available,
59
+ set_seed,
60
+ )
61
+ from transformers.file_utils import get_full_repo_name
62
+ from transformers.testing_utils import CaptureLogger
63
+
64
+
65
+ logger = logging.getLogger(__name__)
66
+
67
+ MODEL_CONFIG_CLASSES = list(FLAX_MODEL_FOR_CAUSAL_LM_MAPPING.keys())
68
+ MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
69
+
70
+
71
+ @dataclass
72
+ class TrainingArguments:
73
+ output_dir: str = field(
74
+ metadata={"help": "The output directory where the model predictions and checkpoints will be written."},
75
+ )
76
+ overwrite_output_dir: bool = field(
77
+ default=False,
78
+ metadata={
79
+ "help": (
80
+ "Overwrite the content of the output directory. "
81
+ "Use this to continue training if output_dir points to a checkpoint directory."
82
+ )
83
+ },
84
+ )
85
+ do_train: bool = field(default=False, metadata={"help": "Whether to run training."})
86
+ do_eval: bool = field(default=False, metadata={"help": "Whether to run eval on the dev set."})
87
+ per_device_train_batch_size: int = field(
88
+ default=8, metadata={"help": "Batch size per GPU/TPU core/CPU for training."}
89
+ )
90
+ per_device_eval_batch_size: int = field(
91
+ default=8, metadata={"help": "Batch size per GPU/TPU core/CPU for evaluation."}
92
+ )
93
+ learning_rate: float = field(default=5e-5, metadata={"help": "The initial learning rate for AdamW."})
94
+ weight_decay: float = field(default=0.0, metadata={"help": "Weight decay for AdamW if we apply some."})
95
+ adam_beta1: float = field(default=0.9, metadata={"help": "Beta1 for AdamW optimizer"})
96
+ adam_beta2: float = field(default=0.999, metadata={"help": "Beta2 for AdamW optimizer"})
97
+ adam_epsilon: float = field(default=1e-8, metadata={"help": "Epsilon for AdamW optimizer."})
98
+ adafactor: bool = field(default=False, metadata={"help": "Whether or not to replace AdamW by Adafactor."})
99
+ num_train_epochs: float = field(default=3.0, metadata={"help": "Total number of training epochs to perform."})
100
+ warmup_steps: int = field(default=0, metadata={"help": "Linear warmup over warmup_steps."})
101
+ warmup_ratio: float = field(default=0.0, metadata={"help": "Linear warmup ratio of total train steps."})
102
+ cosine_decay: bool = field(
103
+ default=False, metadata={"help": "Whether or not to use cosine decay instead of the basic linear decay schedule."}
104
+ )
105
+ gradient_clipping: bool = field(
106
+ default=False, metadata={"help": "Whether or not to use gradient clipping."}
107
+ )
108
+ logging_steps: int = field(default=500, metadata={"help": "Log every X updates steps."})
109
+ save_steps: int = field(default=500, metadata={"help": "Save checkpoint every X updates steps."})
110
+ eval_steps: int = field(default=None, metadata={"help": "Run an evaluation every X steps."})
111
+ seed: int = field(default=42, metadata={"help": "Random seed that will be set at the beginning of training."})
112
+ push_to_hub: bool = field(
113
+ default=False, metadata={"help": "Whether or not to upload the trained model to the model hub after training."}
114
+ )
115
+ hub_model_id: str = field(
116
+ default=None, metadata={"help": "The name of the repository to keep in sync with the local `output_dir`."}
117
+ )
118
+ hub_token: str = field(default=None, metadata={"help": "The token to use to push to the Model Hub."})
119
+
120
+ def __post_init__(self):
121
+ if self.output_dir is not None:
122
+ self.output_dir = os.path.expanduser(self.output_dir)
123
+
124
+ def to_dict(self):
125
+ """
126
+ Serializes this instance while replace `Enum` by their values (for JSON serialization support). It obfuscates
127
+ the token values by removing their value.
128
+ """
129
+ d = asdict(self)
130
+ for k, v in d.items():
131
+ if isinstance(v, Enum):
132
+ d[k] = v.value
133
+ if isinstance(v, list) and len(v) > 0 and isinstance(v[0], Enum):
134
+ d[k] = [x.value for x in v]
135
+ if k.endswith("_token"):
136
+ d[k] = f"<{k.upper()}>"
137
+ return d
138
+
139
+
140
+ @dataclass
141
+ class ModelArguments:
142
+ """
143
+ Arguments pertaining to which model/config/tokenizer we are going to fine-tune, or train from scratch.
144
+ """
145
+
146
+ model_name_or_path: Optional[str] = field(
147
+ default=None,
148
+ metadata={
149
+ "help": "The model checkpoint for weights initialization."
150
+ "Don't set if you want to train a model from scratch."
151
+ },
152
+ )
153
+ model_type: Optional[str] = field(
154
+ default=None,
155
+ metadata={"help": "If training from scratch, pass a model type from the list: " + ", ".join(MODEL_TYPES)},
156
+ )
157
+ config_name: Optional[str] = field(
158
+ default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
159
+ )
160
+ tokenizer_name: Optional[str] = field(
161
+ default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
162
+ )
163
+ cache_dir: Optional[str] = field(
164
+ default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
165
+ )
166
+ use_fast_tokenizer: bool = field(
167
+ default=True,
168
+ metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
169
+ )
170
+ dtype: Optional[str] = field(
171
+ default="float32",
172
+ metadata={
173
+ "help": "Floating-point format in which the model weights should be initialized and trained. Choose one of `[float32, float16, bfloat16]`."
174
+ },
175
+ )
176
+
177
+
178
+ @dataclass
179
+ class DataTrainingArguments:
180
+ """
181
+ Arguments pertaining to what data we are going to input our model for training and eval.
182
+ """
183
+
184
+ dataset_name: Optional[str] = field(
185
+ default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
186
+ )
187
+ dataset_config_name: Optional[str] = field(
188
+ default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
189
+ )
190
+ dataset_filepath: Optional[str] = field(
191
+ default=None, metadata={"help": "Filepath to locally saved HF Dataset (with 'dataset.save_to_disk' method) to use for training"}
192
+ )
193
+ train_file: Optional[str] = field(default=None, metadata={"help": "The input training data file (a text file)."})
194
+ validation_file: Optional[str] = field(
195
+ default=None,
196
+ metadata={"help": "An optional input evaluation data file to evaluate the perplexity on (a text file)."},
197
+ )
198
+ max_train_samples: Optional[int] = field(
199
+ default=None,
200
+ metadata={
201
+ "help": "For debugging purposes or quicker training, truncate the number of training examples to this "
202
+ "value if set."
203
+ },
204
+ )
205
+ max_eval_samples: Optional[int] = field(
206
+ default=None,
207
+ metadata={
208
+ "help": "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
209
+ "value if set."
210
+ },
211
+ )
212
+ overwrite_cache: bool = field(
213
+ default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
214
+ )
215
+ validation_split_percentage: Optional[int] = field(
216
+ default=5,
217
+ metadata={
218
+ "help": "The percentage of the train set used as validation set in case there's no validation split"
219
+ },
220
+ )
221
+ block_size: Optional[int] = field(
222
+ default=None,
223
+ metadata={
224
+ "help": "Optional input sequence length after tokenization. "
225
+ "The training dataset will be truncated in block of this size for training. "
226
+ "Default to the model max input length for single sentence inputs (take into account special tokens)."
227
+ },
228
+ )
229
+ overwrite_cache: bool = field(
230
+ default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
231
+ )
232
+ preprocessing_num_workers: Optional[int] = field(
233
+ default=None,
234
+ metadata={"help": "The number of processes to use for the preprocessing."},
235
+ )
236
+ keep_linebreaks: bool = field(
237
+ default=True, metadata={"help": "Whether to keep line breaks when using TXT files or not."}
238
+ )
239
+
240
+ def __post_init__(self):
241
+ if self.dataset_name is None and self.train_file is None and self.dataset_filepath is None and self.validation_file is None:
242
+ raise ValueError("Need either a dataset name or a training/validation file.")
243
+ else:
244
+ if self.train_file is not None:
245
+ extension = self.train_file.split(".")[-1]
246
+ assert extension in ["csv", "json", "txt"], "`train_file` should be a csv, a json or a txt file."
247
+ if self.validation_file is not None:
248
+ extension = self.validation_file.split(".")[-1]
249
+ assert extension in ["csv", "json", "txt"], "`validation_file` should be a csv, a json or a txt file."
250
+
251
+
252
+ class TrainState(train_state.TrainState):
253
+ dropout_rng: jnp.ndarray
254
+
255
+ def replicate(self):
256
+ return jax_utils.replicate(self).replace(dropout_rng=shard_prng_key(self.dropout_rng))
257
+
258
+
259
+ def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuffle: bool = False):
260
+ """
261
+ Returns batches of size `batch_size` from truncated `dataset`, sharded over all local devices.
262
+ Shuffle batches if `shuffle` is `True`.
263
+ """
264
+ steps_per_epoch = len(dataset) // batch_size
265
+
266
+ if shuffle:
267
+ batch_idx = jax.random.permutation(rng, len(dataset))
268
+ else:
269
+ batch_idx = jnp.arange(len(dataset))
270
+
271
+ batch_idx = batch_idx[: steps_per_epoch * batch_size] # Skip incomplete batch.
272
+ batch_idx = batch_idx.reshape((steps_per_epoch, batch_size))
273
+
274
+ for idx in batch_idx:
275
+ batch = dataset[idx]
276
+ batch = {k: np.array(v) for k, v in batch.items()}
277
+
278
+ yield batch
279
+
280
+
281
+ def write_train_metric(summary_writer, train_metrics, train_time, step):
282
+ summary_writer.scalar("train_time", train_time, step)
283
+
284
+ train_metrics = get_metrics(train_metrics)
285
+ for key, vals in train_metrics.items():
286
+ tag = f"train_{key}"
287
+ for i, val in enumerate(vals):
288
+ summary_writer.scalar(tag, val, step - len(vals) + i + 1)
289
+
290
+
291
+ def write_eval_metric(summary_writer, eval_metrics, step):
292
+ for metric_name, value in eval_metrics.items():
293
+ summary_writer.scalar(f"eval_{metric_name}", value, step)
294
+
295
+
296
+ def create_linear_learning_rate_fn(
297
+ train_ds_size: int, train_batch_size: int, num_train_epochs: int, num_warmup_steps: int, learning_rate: float
298
+ ) -> Callable[[int], jnp.array]:
299
+ """Returns a linear warmup, linear decay learning rate function."""
300
+ steps_per_epoch = train_ds_size // train_batch_size
301
+ num_train_steps = steps_per_epoch * num_train_epochs
302
+ warmup_fn = optax.linear_schedule(init_value=0.0, end_value=learning_rate, transition_steps=num_warmup_steps)
303
+ decay_fn = optax.linear_schedule(
304
+ init_value=learning_rate, end_value=0, transition_steps=num_train_steps - num_warmup_steps
305
+ )
306
+ schedule_fn = optax.join_schedules(schedules=[warmup_fn, decay_fn], boundaries=[num_warmup_steps])
307
+ return schedule_fn
308
+
309
+ def create_cosine_learning_rate_fn(
310
+ train_ds_size: int, train_batch_size: int, num_train_epochs: int, num_warmup_steps: int, learning_rate: float
311
+ ) -> Callable[[int], jnp.array]:
312
+ """Returns a linear warmup, cosine decay learning rate function."""
313
+ steps_per_epoch = train_ds_size // train_batch_size
314
+ num_train_steps = steps_per_epoch * num_train_epochs
315
+ warmup_fn = optax.linear_schedule(init_value=0.0, end_value=learning_rate, transition_steps=num_warmup_steps)
316
+ decay_fn = optax.cosine_decay_schedule(
317
+ init_value=learning_rate, decay_steps=num_train_steps - num_warmup_steps, alpha=0.1
318
+ )
319
+ schedule_fn = optax.join_schedules(schedules=[warmup_fn, decay_fn], boundaries=[num_warmup_steps])
320
+ return schedule_fn
321
+
322
+
323
+ def main():
324
+ # See all possible arguments in src/transformers/training_args.py
325
+ # or by passing the --help flag to this script.
326
+ # We now keep distinct sets of args, for a cleaner separation of concerns.
327
+
328
+ parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
329
+ if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
330
+ # If we pass only one argument to the script and it's the path to a json file,
331
+ # let's parse it to get our arguments.
332
+ model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
333
+ else:
334
+ model_args, data_args, training_args = parser.parse_args_into_dataclasses()
335
+
336
+ if (
337
+ os.path.exists(training_args.output_dir)
338
+ and os.listdir(training_args.output_dir)
339
+ and training_args.do_train
340
+ and not training_args.overwrite_output_dir
341
+ ):
342
+ raise ValueError(
343
+ f"Output directory ({training_args.output_dir}) already exists and is not empty."
344
+ "Use --overwrite_output_dir to overcome."
345
+ )
346
+
347
+ # Make one log on every process with the configuration for debugging.
348
+ logging.basicConfig(
349
+ format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
350
+ datefmt="%m/%d/%Y %H:%M:%S",
351
+ level=logging.INFO,
352
+ )
353
+ # Setup logging, we only want one process per machine to log things on the screen.
354
+ logger.setLevel(logging.INFO if jax.process_index() == 0 else logging.ERROR)
355
+ if jax.process_index() == 0:
356
+ datasets.utils.logging.set_verbosity_warning()
357
+ transformers.utils.logging.set_verbosity_info()
358
+ else:
359
+ datasets.utils.logging.set_verbosity_error()
360
+ transformers.utils.logging.set_verbosity_error()
361
+
362
+ # Set the verbosity to info of the Transformers logger (on main process only):
363
+ logger.info(f"Training/evaluation parameters {training_args}")
364
+
365
+ # Set seed before initializing model.
366
+ set_seed(training_args.seed)
367
+
368
+ # Handle the repository creation
369
+ if training_args.push_to_hub:
370
+ if training_args.hub_model_id is None:
371
+ repo_name = get_full_repo_name(
372
+ Path(training_args.output_dir).absolute().name, token=training_args.hub_token
373
+ )
374
+ else:
375
+ repo_name = training_args.hub_model_id
376
+ repo = Repository(training_args.output_dir, clone_from=repo_name)
377
+
378
+ # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
379
+ # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
380
+ # (the dataset will be downloaded automatically from the datasets Hub).
381
+ #
382
+ # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
383
+ # 'text' is found. You can easily tweak this behavior (see below).
384
+ #
385
+ # In distributed training, the load_dataset function guarantees that only one local process can concurrently
386
+ # download the dataset.
387
+ if data_args.dataset_name is not None:
388
+ # Downloading and loading a dataset from the hub.
389
+ dataset = load_dataset(
390
+ data_args.dataset_name, data_args.dataset_config_name, cache_dir=model_args.cache_dir, keep_in_memory=False
391
+ )
392
+
393
+ if "validation" not in dataset.keys():
394
+ dataset["validation"] = load_dataset(
395
+ data_args.dataset_name,
396
+ data_args.dataset_config_name,
397
+ split=f"train[:{data_args.validation_split_percentage}%]",
398
+ cache_dir=model_args.cache_dir,
399
+ )
400
+ dataset["train"] = load_dataset(
401
+ data_args.dataset_name,
402
+ data_args.dataset_config_name,
403
+ split=f"train[{data_args.validation_split_percentage}%:]",
404
+ cache_dir=model_args.cache_dir,
405
+ )
406
+
407
+ elif data_args.dataset_filepath is not None:
408
+ # Loading a dataset from local file.
409
+ dataset = load_from_disk(data_args.dataset_filepath)
410
+ if "validation" not in dataset.keys():
411
+ dataset = datasets.train_test_split(test_size=data_args.validation_split_percentage/100)
412
+ dataset["validation"] = dataset["test"]
413
+ del dataset["test"]
414
+
415
+ else:
416
+ data_files = {}
417
+ dataset_args = {}
418
+ if data_args.train_file is not None:
419
+ data_files["train"] = data_args.train_file
420
+ if data_args.validation_file is not None:
421
+ data_files["validation"] = data_args.validation_file
422
+ extension = data_args.train_file.split(".")[-1]
423
+ if extension == "txt":
424
+ extension = "text"
425
+ dataset_args["keep_linebreaks"] = data_args.keep_linebreaks
426
+ dataset = load_dataset(extension, data_files=data_files, cache_dir=model_args.cache_dir, **dataset_args)
427
+
428
+ if "validation" not in dataset.keys():
429
+ dataset["validation"] = load_dataset(
430
+ extension,
431
+ data_files=data_files,
432
+ split=f"train[:{data_args.validation_split_percentage}%]",
433
+ cache_dir=model_args.cache_dir,
434
+ **dataset_args,
435
+ )
436
+ dataset["train"] = load_dataset(
437
+ extension,
438
+ data_files=data_files,
439
+ split=f"train[{data_args.validation_split_percentage}%:]",
440
+ cache_dir=model_args.cache_dir,
441
+ **dataset_args,
442
+ )
443
+ # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
444
+ # https://huggingface.co/docs/datasets/loading_datasets.html.
445
+
446
+ # Load pretrained model and tokenizer
447
+
448
+ # Distributed training:
449
+ # The .from_pretrained methods guarantee that only one local process can concurrently
450
+ # download model & vocab.
451
+ if model_args.config_name:
452
+ config = AutoConfig.from_pretrained(model_args.config_name, cache_dir=model_args.cache_dir)
453
+ elif model_args.model_name_or_path:
454
+ config = AutoConfig.from_pretrained(model_args.model_name_or_path, cache_dir=model_args.cache_dir)
455
+ else:
456
+ config = CONFIG_MAPPING[model_args.model_type]()
457
+ logger.warning("You are instantiating a new config instance from scratch.")
458
+
459
+ if model_args.tokenizer_name:
460
+ tokenizer = AutoTokenizer.from_pretrained(
461
+ model_args.tokenizer_name, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
462
+ )
463
+ elif model_args.model_name_or_path:
464
+ tokenizer = AutoTokenizer.from_pretrained(
465
+ model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
466
+ )
467
+ else:
468
+ raise ValueError(
469
+ "You are instantiating a new tokenizer from scratch. This is not supported by this script."
470
+ "You can do it from another script, save it, and load it from here, using --tokenizer_name."
471
+ )
472
+
473
+ if tokenizer.pad_token is None:
474
+ tokenizer.pad_token = tokenizer.eos_token
475
+
476
+ if model_args.model_name_or_path:
477
+ model = FlaxAutoModelForCausalLM.from_pretrained(
478
+ model_args.model_name_or_path, config=config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
479
+ )
480
+ else:
481
+ model = FlaxAutoModelForCausalLM.from_config(
482
+ config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
483
+ )
484
+
485
+ # Preprocessing the datasets.
486
+ # First we tokenize all the texts.
487
+ if training_args.do_train:
488
+ column_names = dataset["train"].column_names
489
+ else:
490
+ column_names = dataset["validation"].column_names
491
+ text_column_name = "text" if "text" in column_names else column_names[0]
492
+
493
+ # since this will be pickled to avoid _LazyModule error in Hasher force logger loading before tokenize_function
494
+ tok_logger = transformers.utils.logging.get_logger("transformers.tokenization_utils_base")
495
+
496
+ def tokenize_function(examples):
497
+ with CaptureLogger(tok_logger) as cl:
498
+ output = tokenizer(examples[text_column_name])
499
+ # clm input could be much much longer than block_size
500
+ if "Token indices sequence length is longer than the" in cl.out:
501
+ tok_logger.warning(
502
+ "^^^^^^^^^^^^^^^^ Please ignore the warning above - this long input will be chunked into smaller bits before being passed to the model."
503
+ )
504
+ return output
505
+
506
+ tokenized_datasets = dataset.map(
507
+ tokenize_function,
508
+ batched=True,
509
+ num_proc=data_args.preprocessing_num_workers,
510
+ remove_columns=column_names,
511
+ load_from_cache_file=not data_args.overwrite_cache,
512
+ )
513
+
514
+ if data_args.block_size is None:
515
+ block_size = tokenizer.model_max_length
516
+ if block_size > config.max_position_embeddings:
517
+ logger.warning(
518
+ f"The tokenizer picked seems to have a very large `model_max_length` ({tokenizer.model_max_length}). "
519
+ "Picking 1024 instead. You can change that default value by passing --block_size xxx."
520
+ )
521
+ block_size = 1024
522
+ else:
523
+ if data_args.block_size > tokenizer.model_max_length:
524
+ logger.warning(
525
+ f"The block_size passed ({data_args.block_size}) is larger than the maximum length for the model"
526
+ f"({tokenizer.model_max_length}). Using block_size={tokenizer.model_max_length}."
527
+ )
528
+ block_size = min(data_args.block_size, tokenizer.model_max_length)
529
+
530
+ # Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size.
531
+ def group_texts(examples):
532
+ # Concatenate all texts.
533
+ concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
534
+ total_length = len(concatenated_examples[list(examples.keys())[0]])
535
+ # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
536
+ # customize this part to your needs.
537
+ if total_length >= block_size:
538
+ total_length = (total_length // block_size) * block_size
539
+ # Split by chunks of max_len.
540
+ result = {
541
+ k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
542
+ for k, t in concatenated_examples.items()
543
+ }
544
+ result["labels"] = result["input_ids"].copy()
545
+ return result
546
+
547
+ # Note that with `batched=True`, this map processes 1,000 texts together, so group_texts throws away a remainder
548
+ # for each of those groups of 1,000 texts. You can adjust that batch_size here but a higher value might be slower
549
+ # to preprocess.
550
+ #
551
+ # To speed up this part, we use multiprocessing. See the documentation of the map method for more information:
552
+ # https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.map
553
+
554
+ lm_datasets = tokenized_datasets.map(
555
+ group_texts,
556
+ batched=True,
557
+ num_proc=data_args.preprocessing_num_workers,
558
+ load_from_cache_file=not data_args.overwrite_cache,
559
+ )
560
+
561
+ if training_args.do_train:
562
+ if "train" not in tokenized_datasets:
563
+ raise ValueError("--do_train requires a train dataset")
564
+ train_dataset = lm_datasets["train"]
565
+ if data_args.max_train_samples is not None:
566
+ train_dataset = train_dataset.select(range(data_args.max_train_samples))
567
+
568
+ # test to see that tokenization worked
569
+ detokenized_example = tokenizer.decode(train_dataset[0]["input_ids"])
570
+ logger.info(f"Detokenized example: {detokenized_example}")
571
+ detokenized_example = tokenizer.decode(train_dataset[-1]["input_ids"])
572
+ logger.info(f"Detokenized example 2: {detokenized_example}")
573
+
574
+ if training_args.do_eval:
575
+ if "validation" not in tokenized_datasets:
576
+ raise ValueError("--do_eval requires a validation dataset")
577
+ eval_dataset = lm_datasets["validation"]
578
+ if data_args.max_eval_samples is not None:
579
+ eval_dataset = eval_dataset.select(range(data_args.max_eval_samples))
580
+
581
+ # Enable tensorboard only on the master node
582
+ has_tensorboard = is_tensorboard_available()
583
+ if has_tensorboard and jax.process_index() == 0:
584
+ try:
585
+ from flax.metrics.tensorboard import SummaryWriter
586
+
587
+ summary_writer = SummaryWriter(log_dir=Path(os.path.join(training_args.output_dir, "runs")))
588
+ except ImportError as ie:
589
+ has_tensorboard = False
590
+ logger.warning(
591
+ f"Unable to display metrics through TensorBoard because some package are not installed: {ie}"
592
+ )
593
+ else:
594
+ logger.warning(
595
+ "Unable to display metrics through TensorBoard because the package is not installed: "
596
+ "Please run pip install tensorboard to enable."
597
+ )
598
+
599
+ # Initialize our training
600
+ rng = jax.random.PRNGKey(training_args.seed)
601
+ rng, dropout_rng = jax.random.split(rng)
602
+
603
+ # Store some constant
604
+ num_epochs = int(training_args.num_train_epochs)
605
+ train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()
606
+ eval_batch_size = int(training_args.per_device_eval_batch_size) * jax.device_count()
607
+ steps_per_epoch = len(train_dataset) // train_batch_size
608
+ total_train_steps = steps_per_epoch * num_epochs
609
+
610
+ if training_args.warmup_ratio > 0:
611
+ warmup_steps = int(total_train_steps * training_args.warmup_ratio)
612
+ else:
613
+ warmup_steps = training_args.warmup_steps
614
+
615
+ # Create learning rate schedule
616
+ if training_args.cosine_decay:
617
+ lr_schedule_fn = create_cosine_learning_rate_fn(
618
+ len(train_dataset),
619
+ train_batch_size,
620
+ training_args.num_train_epochs,
621
+ warmup_steps,
622
+ training_args.learning_rate,
623
+ )
624
+ else:
625
+ lr_schedule_fn = create_linear_learning_rate_fn(
626
+ len(train_dataset),
627
+ train_batch_size,
628
+ training_args.num_train_epochs,
629
+ warmup_steps,
630
+ training_args.learning_rate,
631
+ )
632
+
633
+ # We use Optax's "masking" functionality to not apply weight decay
634
+ # to bias and LayerNorm scale parameters. decay_mask_fn returns a
635
+ # mask boolean with the same structure as the parameters.
636
+ # The mask is True for parameters that should be decayed.
637
+ # Note that this mask is specifically adapted for FlaxGPT2.
638
+ # For other models, one should correct the layer norm parameter naming
639
+ # accordingly.
640
+ def decay_mask_fn(params):
641
+ flat_params = traverse_util.flatten_dict(params)
642
+ flat_mask = {
643
+ path: (path[-1] != "bias" and path[-2:] not in [("ln_1", "scale"), ("ln_2", "scale"), ("ln_f", "scale")])
644
+ for path in flat_params
645
+ }
646
+ return traverse_util.unflatten_dict(flat_mask)
647
+
648
+ # create adam optimizer
649
+ if training_args.adafactor:
650
+ # We use the default parameters here to initialize adafactor,
651
+ # For more details about the parameters please check https://github.com/deepmind/optax/blob/ed02befef9bf81cbbf236be3d2b0e032e9ed4a40/optax/_src/alias.py#L74
652
+ optimizer = optax.adafactor(
653
+ learning_rate=lr_schedule_fn,
654
+ )
655
+ else:
656
+ optimizer = optax.adamw(
657
+ learning_rate=lr_schedule_fn,
658
+ b1=training_args.adam_beta1,
659
+ b2=training_args.adam_beta2,
660
+ eps=training_args.adam_epsilon,
661
+ weight_decay=training_args.weight_decay,
662
+ mask=decay_mask_fn,
663
+ )
664
+ if training_args.gradient_clipping:
665
+ optimizer = optax.chain(
666
+ optax.clip_by_global_norm(1.),
667
+ optimizer
668
+ )
669
+
670
+ # Setup train state
671
+ state = TrainState.create(apply_fn=model.__call__, params=model.params, tx=optimizer, dropout_rng=dropout_rng)
672
+
673
+ def loss_fn(logits, labels):
674
+ shift_logits = logits[..., :-1, :]
675
+ shift_labels = labels[..., 1:]
676
+ loss = optax.softmax_cross_entropy(shift_logits, onehot(shift_labels, shift_logits.shape[-1]))
677
+ return loss.mean()
678
+
679
+ # Define gradient update step fn
680
+ def train_step(state, batch):
681
+ dropout_rng, new_dropout_rng = jax.random.split(state.dropout_rng)
682
+
683
+ def compute_loss(params):
684
+ labels = batch.pop("labels")
685
+ logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
686
+ loss = loss_fn(logits, labels)
687
+ return loss
688
+
689
+ grad_fn = jax.value_and_grad(compute_loss)
690
+ loss, grad = grad_fn(state.params)
691
+ grad = jax.lax.pmean(grad, "batch")
692
+
693
+ new_state = state.apply_gradients(grads=grad, dropout_rng=new_dropout_rng)
694
+
695
+ metrics = {"loss": loss, "learning_rate": lr_schedule_fn(state.step)}
696
+ metrics = jax.lax.pmean(metrics, axis_name="batch")
697
+
698
+ return new_state, metrics
699
+
700
+ # Define eval fn
701
+ def eval_step(params, batch):
702
+ labels = batch.pop("labels")
703
+ logits = model(**batch, params=params, train=False)[0]
704
+ loss = loss_fn(logits, labels)
705
+
706
+ # summarize metrics
707
+ metrics = {"loss": loss}
708
+ metrics = jax.lax.pmean(metrics, axis_name="batch")
709
+ return metrics
710
+
711
+ # Create parallel version of the train and eval step
712
+ p_train_step = jax.pmap(train_step, "batch", donate_argnums=(0,))
713
+ p_eval_step = jax.pmap(eval_step, "batch")
714
+
715
+ # Replicate the train state on each device
716
+ state = state.replicate()
717
+
718
+ logger.info("***** Running training *****")
719
+ logger.info(f" Num examples = {len(train_dataset)}")
720
+ logger.info(f" Num Epochs = {num_epochs}")
721
+ logger.info(f" Instantaneous batch size per device = {training_args.per_device_train_batch_size}")
722
+ logger.info(f" Total train batch size (w. parallel & distributed) = {train_batch_size}")
723
+ logger.info(f" Total optimization steps = {total_train_steps}")
724
+
725
+ train_time = 0
726
+ train_metrics = []
727
+ epochs = tqdm(range(num_epochs), desc="Epoch ... ", position=0)
728
+ for epoch in epochs:
729
+ # ======================== Training ================================
730
+ train_start = time.time()
731
+
732
+ # Create sampling rng
733
+ rng, input_rng = jax.random.split(rng)
734
+
735
+ # Generate an epoch by shuffling sampling indices from the train dataset
736
+ train_loader = data_loader(input_rng, train_dataset, train_batch_size, shuffle=True)
737
+ steps_per_epoch = len(train_dataset) // train_batch_size
738
+ # train
739
+ for step in tqdm(range(steps_per_epoch), desc="Training...", position=1, leave=False):
740
+ batch = next(train_loader)
741
+ batch = shard(batch)
742
+ state, train_metric = p_train_step(state, batch)
743
+ train_metrics.append(train_metric)
744
+
745
+ cur_step = epoch * (len(train_dataset) // train_batch_size) + step
746
+
747
+ if cur_step % training_args.logging_steps == 0 and cur_step > 0:
748
+ # Save metrics
749
+ train_metric = unreplicate(train_metric)
750
+ train_time += time.time() - train_start
751
+ if has_tensorboard and jax.process_index() == 0:
752
+ write_train_metric(summary_writer, train_metrics, train_time, cur_step)
753
+
754
+ epochs.write(
755
+ f"Step... ({cur_step} | Loss: {train_metric['loss'].mean()}, Learning Rate: {train_metric['learning_rate'].mean()})"
756
+ )
757
+
758
+ train_metrics = []
759
+
760
+ if cur_step % training_args.eval_steps == 0 and cur_step > 0:
761
+ # ======================== Evaluating ==============================
762
+ eval_metrics = []
763
+ eval_loader = data_loader(input_rng, eval_dataset, eval_batch_size)
764
+ eval_steps = len(eval_dataset) // eval_batch_size
765
+ for _ in tqdm(range(eval_steps), desc="Evaluating...", position=2, leave=False):
766
+ # Model forward
767
+ batch = next(eval_loader)
768
+ batch = shard(batch)
769
+ metrics = p_eval_step(state.params, batch)
770
+ eval_metrics.append(metrics)
771
+
772
+ # normalize eval metrics
773
+ eval_metrics = get_metrics(eval_metrics)
774
+ eval_metrics = jax.tree_map(jnp.mean, eval_metrics)
775
+
776
+ try:
777
+ eval_metrics["perplexity"] = math.exp(eval_metrics["loss"])
778
+ except OverflowError:
779
+ eval_metrics["perplexity"] = float("inf")
780
+
781
+ # Print metrics and update progress bar
782
+ desc = f"Step... ({cur_step} | Eval Loss: {eval_metrics['loss']} | Eval Perplexity: {eval_metrics['perplexity']})"
783
+ epochs.write(desc)
784
+ epochs.desc = desc
785
+
786
+ # Save metrics
787
+ if has_tensorboard and jax.process_index() == 0:
788
+ write_eval_metric(summary_writer, eval_metrics, cur_step)
789
+
790
+ if cur_step % training_args.save_steps == 0 and cur_step > 0:
791
+ # save checkpoint after each epoch and push checkpoint to the hub
792
+ if jax.process_index() == 0:
793
+ params = jax.device_get(unreplicate(state.params))
794
+ model.save_pretrained(training_args.output_dir, params=params)
795
+ tokenizer.save_pretrained(training_args.output_dir)
796
+ if training_args.push_to_hub:
797
+ repo.push_to_hub(commit_message=f"Saving weights and logs of step {cur_step}", blocking=False)
798
+
799
+ # save also at the end of epoch
800
+ try:
801
+ if jax.process_index() == 0:
802
+ params = jax.device_get(unreplicate(state.params))
803
+ model.save_pretrained(training_args.output_dir, params=params)
804
+ tokenizer.save_pretrained(training_args.output_dir)
805
+ if training_args.push_to_hub:
806
+ repo.push_to_hub(commit_message=f"Saving weights and logs of epoch {epoch}", blocking=False)
807
+ except:
808
+ # push to hub fails the whole script if nothing new to commit
809
+ pass
810
+
811
+ # Eval after training
812
+ if training_args.do_eval:
813
+ eval_metrics = []
814
+ eval_loader = data_loader(input_rng, eval_dataset, eval_batch_size)
815
+ eval_steps = len(eval_dataset) // eval_batch_size
816
+ for _ in tqdm(range(eval_steps), desc="Evaluating...", position=2, leave=False):
817
+ # Model forward
818
+ batch = shard(next(eval_loader))
819
+ metrics = p_eval_step(state.params, batch)
820
+ eval_metrics.append(metrics)
821
+
822
+ # normalize eval metrics
823
+ eval_metrics = get_metrics(eval_metrics)
824
+ eval_metrics = jax.tree_map(lambda x: jnp.mean(x).item(), eval_metrics)
825
+
826
+ try:
827
+ eval_metrics["perplexity"] = math.exp(eval_metrics["loss"])
828
+ except OverflowError:
829
+ eval_metrics["perplexity"] = float("inf")
830
+
831
+ if jax.process_index() == 0:
832
+ eval_metrics = {f"eval_{metric_name}": value for metric_name, value in eval_metrics.items()}
833
+ path = os.path.join(training_args.output_dir, "eval_results.json")
834
+ with open(path, "w") as f:
835
+ json.dump(eval_metrics, f, indent=4, sort_keys=True)
836
+
837
+
838
+ if __name__ == "__main__":
839
+ main()
runs/events.out.tfevents.1642099734.t1v-n-42145f73-w-0.2317757.0.v2 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3310027638bb05951334514395d3e81a9d7eb8ce109681141532c4e59751cb88
3
+ size 1471449
special_tokens_map.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {"bos_token": "<|endoftext|>", "eos_token": "<|endoftext|>", "unk_token": "<|endoftext|>", "pad_token": "<|endoftext|>"}
start_train.sh ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # set train hyperparams
2
+ unset LD_PRELOAD
3
+ export HF_DATASETS_CACHE="/researchdisk/datasets_cache"
4
+ export USE_TORCH=0
5
+ python3 run_clm_flax.py \
6
+ --output_dir="./" \
7
+ --model_type="gpt2" \
8
+ --config_name="./" \
9
+ --tokenizer_name="./" \
10
+ --dataset_filepath="/researchdisk/training_dataset_full_deduplicated" \
11
+ --do_train --do_eval \
12
+ --block_size="512" \
13
+ --per_device_train_batch_size="32" \
14
+ --per_device_eval_batch_size="32" \
15
+ --preprocessing_num_workers="64" \
16
+ --learning_rate="5e-3" \
17
+ --warmup_ratio="0.01" \
18
+ --cosine_decay \
19
+ --adam_beta1="0.9" \
20
+ --adam_beta2="0.98" \
21
+ --adam_epsilon="1e-8" \
22
+ --weight_decay="0.01" \
23
+ --overwrite_output_dir \
24
+ --logging_steps="500" \
25
+ --eval_steps="10000" \
26
+ --save_steps="10000" \
27
+ --num_train_epochs="10" \
28
+ --dtype="bfloat16" \
29
+ --push_to_hub \
30
+ --hub_model_id="Finnish-NLP/gpt2-finnish"
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer_config.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {"unk_token": "<|endoftext|>", "bos_token": "<|endoftext|>", "eos_token": "<|endoftext|>", "add_prefix_space": false, "special_tokens_map_file": null, "name_or_path": "./", "tokenizer_class": "GPT2Tokenizer"}
train_tokenizer.py ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from datasets import load_from_disk
2
+ from tokenizers import trainers, Tokenizer, normalizers, ByteLevelBPETokenizer
3
+ from transformers import AutoTokenizer
4
+
5
+
6
+ model_dir = "./"
7
+
8
+ # load dataset
9
+ dataset = load_from_disk("/researchdisk/training_dataset_full_deduplicated")
10
+ dataset = dataset["train"]
11
+
12
+ # Instantiate tokenizer
13
+ tokenizer = ByteLevelBPETokenizer()
14
+ def batch_iterator(batch_size=1000):
15
+ for i in range(0, len(dataset), batch_size):
16
+ yield dataset[i: i + batch_size]["text"]
17
+
18
+ # Customized training
19
+ tokenizer.train_from_iterator(batch_iterator(), vocab_size=50257, min_frequency=2, special_tokens=[
20
+ "<s>",
21
+ "<pad>",
22
+ "</s>",
23
+ "<unk>",
24
+ "<mask>",
25
+ ])
26
+
27
+ # Save files to disk
28
+ tokenizer.save(f"{model_dir}/tokenizer.json")
29
+ tokenizer = AutoTokenizer.from_pretrained(model_dir)
30
+ tokenizer.save_pretrained(model_dir)
vocab.json ADDED
The diff for this file is too large to render. See raw diff