birgermoell commited on
Commit
1d2d61e
1 Parent(s): 003c665

Uploaded model

Browse files
README.md ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ language: da
3
+ widget:
4
+ - text: "Jeg elsker livet"
5
+ ---
6
+
7
+ # GPT2-svenska-wikipedia
8
+ A Danish GPT2 style model trained using Flax CLM pipeline on the Danish
9
+ part of the wiki40b dataset.
10
+
11
+ https://huggingface.co/datasets/wiki40b
12
+
13
+
14
+ ## Data cleaning and preprocessing
15
+ The data was cleaned and preprocessed using the following script. Make sure to install depencies for beam_runner to make the dataset work.
16
+
17
+ ```python
18
+ from datasets import load_dataset
19
+ def load_and_clean_wiki():
20
+ dataset = load_dataset('wiki40b', 'da', beam_runner='DirectRunner', split="train")
21
+ #dataset = load_dataset('wiki40b', 'sv', beam_runner='DirectRunner')
22
+ dataset = dataset.remove_columns(['wikidata_id', 'version_id'])
23
+ filtered_dataset = dataset.map(filter_wikipedia)
24
+ # filtered_dataset[:3]
25
+ # print(filtered_dataset[:3])
26
+ return filtered_dataset
27
+
28
+ def filter_wikipedia(batch):
29
+ batch["text"] = " ".join(batch["text"].split("\n_START_SECTION_\n"))
30
+ batch["text"] = " ".join(batch["text"].split("\n_START_ARTICLE_\n"))
31
+ batch["text"] = " ".join(batch["text"].split("\n_START_ARTICLE_\n"))
32
+ batch["text"] = " ".join(batch["text"].split("\n_START_PARAGRAPH_\n"))
33
+ batch["text"] = " ".join(batch["text"].split("_NEWLINE_"))
34
+ batch["text"] = " ".join(batch["text"].split("\xa0"))
35
+ return batch
36
+ ```
37
+
38
+ ## Training script
39
+ The following training script was used to train the model.
40
+ ```bash
41
+ ./run_clm_flax.py --output_dir="${MODEL_DIR}" --model_type="gpt2" --config_name="${MODEL_DIR}" --tokenizer_name="${MODEL_DIR}" --dataset_name="wiki40b" --dataset_config_name="da" --do_train --do_eval --block_size="512" --per_device_train_batch_size="64" --per_device_eval_batch_size="64" --learning_rate="5e-3" --warmup_steps="1000" --adam_beta1="0.9" --adam_beta2="0.98" --weight_decay="0.01" --overwrite_output_dir --num_train_epochs="20" --logging_steps="500" --save_steps="1000" --eval_steps="2500" --push_to_hub
42
+ ```
43
+
added_tokens.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {"<|endoftext|>": 50265}
config.json ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": ".",
3
+ "activation_function": "gelu_new",
4
+ "architectures": [
5
+ "GPT2LMHeadModel"
6
+ ],
7
+ "attn_pdrop": 0.0,
8
+ "bos_token_id": 50256,
9
+ "embd_pdrop": 0.0,
10
+ "eos_token_id": 50256,
11
+ "gradient_checkpointing": false,
12
+ "initializer_range": 0.02,
13
+ "layer_norm_epsilon": 1e-05,
14
+ "model_type": "gpt2",
15
+ "n_ctx": 1024,
16
+ "n_embd": 768,
17
+ "n_head": 12,
18
+ "n_inner": null,
19
+ "n_layer": 12,
20
+ "n_positions": 1024,
21
+ "resid_pdrop": 0.0,
22
+ "scale_attn_weights": true,
23
+ "summary_activation": null,
24
+ "summary_first_dropout": 0.1,
25
+ "summary_proj_to_labels": true,
26
+ "summary_type": "cls_index",
27
+ "summary_use_proj": true,
28
+ "task_specific_params": {
29
+ "text-generation": {
30
+ "do_sample": true,
31
+ "max_length": 50
32
+ }
33
+ },
34
+ "transformers_version": "4.8.2",
35
+ "use_cache": true,
36
+ "vocab_size": 50257
37
+ }
evaluate.py ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import GPT2Tokenizer, GPT2Model, FlaxGPT2LMHeadModel, GPT2LMHeadModel, pipeline, set_seed
2
+
3
+ tokenizer = GPT2Tokenizer.from_pretrained("flax-community/swe-gpt-wiki")
4
+ model = GPT2LMHeadModel.from_pretrained("flax-community/swe-gpt-wiki")
5
+
6
+
7
+ generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
8
+ set_seed(42)
9
+ result = generator("Jeg elsker livet", max_length=150, num_return_sequences=5)
10
+ print(result)
events.out.tfevents.1626210352.t1v-n-905ccd3e-w-0.81578.3.v2 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:76d130c4a36ddf3c7806834b1585c3ce4c5576656b7a4c5ecfb81f8f24b2a0ff
3
+ size 441341
flax_model.msgpack ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:2e1c774950aca8cbe48b10b23e878b0ab53f74c157abad541731f0d805d36b44
3
+ size 497764120
make_config.py ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ from transformers import GPT2Config
2
+
3
+ model_dir = "./swe-gpt-wiki" # ${MODEL_DIR}
4
+
5
+ config = GPT2Config.from_pretrained("gpt2", resid_pdrop=0.0, embd_pdrop=0.0, attn_pdrop=0.0)
6
+ config.save_pretrained(model_dir)
merges.txt ADDED
The diff for this file is too large to render. See raw diff
 
pytorch_model.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c483c8ff7b5ee38b196395c81fc9bdcb80c93762a1b9affc4d62509844ac9f5f
3
+ size 510401385
run_clm_flax.py ADDED
@@ -0,0 +1,645 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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=causal-lm
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 logging
25
+ import math
26
+ import os
27
+ import sys
28
+ import time
29
+ from dataclasses import dataclass, field
30
+ from pathlib import Path
31
+ from typing import Callable, Optional
32
+
33
+ import datasets
34
+ from datasets import Dataset, load_dataset
35
+ from tqdm import tqdm
36
+
37
+ import jax
38
+ import jax.numpy as jnp
39
+ import optax
40
+ import transformers
41
+ from load_from_hf import load_and_clean_wiki
42
+ from flax import jax_utils, traverse_util
43
+ from flax.jax_utils import unreplicate
44
+ from flax.training import train_state
45
+ from flax.training.common_utils import get_metrics, onehot, shard, shard_prng_key
46
+ from transformers import (
47
+ CONFIG_MAPPING,
48
+ FLAX_MODEL_FOR_CAUSAL_LM_MAPPING,
49
+ AutoConfig,
50
+ AutoTokenizer,
51
+ FlaxAutoModelForCausalLM,
52
+ HfArgumentParser,
53
+ TrainingArguments,
54
+ is_tensorboard_available,
55
+ )
56
+ from transformers.testing_utils import CaptureLogger
57
+
58
+
59
+ logger = logging.getLogger(__name__)
60
+
61
+ MODEL_CONFIG_CLASSES = list(FLAX_MODEL_FOR_CAUSAL_LM_MAPPING.keys())
62
+ MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
63
+
64
+
65
+ @dataclass
66
+ class ModelArguments:
67
+ """
68
+ Arguments pertaining to which model/config/tokenizer we are going to fine-tune, or train from scratch.
69
+ """
70
+
71
+ model_name_or_path: Optional[str] = field(
72
+ default=None,
73
+ metadata={
74
+ "help": "The model checkpoint for weights initialization."
75
+ "Don't set if you want to train a model from scratch."
76
+ },
77
+ )
78
+ model_type: Optional[str] = field(
79
+ default=None,
80
+ metadata={"help": "If training from scratch, pass a model type from the list: " + ", ".join(MODEL_TYPES)},
81
+ )
82
+ config_name: Optional[str] = field(
83
+ default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
84
+ )
85
+ tokenizer_name: Optional[str] = field(
86
+ default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
87
+ )
88
+ cache_dir: Optional[str] = field(
89
+ default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
90
+ )
91
+ use_fast_tokenizer: bool = field(
92
+ default=True,
93
+ metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
94
+ )
95
+ dtype: Optional[str] = field(
96
+ default="float32",
97
+ metadata={
98
+ "help": "Floating-point format in which the model weights should be initialized and trained. Choose one of `[float32, float16, bfloat16]`."
99
+ },
100
+ )
101
+
102
+
103
+ @dataclass
104
+ class DataTrainingArguments:
105
+ """
106
+ Arguments pertaining to what data we are going to input our model for training and eval.
107
+ """
108
+
109
+ dataset_name: Optional[str] = field(
110
+ default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
111
+ )
112
+ dataset_config_name: Optional[str] = field(
113
+ default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
114
+ )
115
+ train_file: Optional[str] = field(default=None, metadata={"help": "The input training data file (a text file)."})
116
+ validation_file: Optional[str] = field(
117
+ default=None,
118
+ metadata={"help": "An optional input evaluation data file to evaluate the perplexity on (a text file)."},
119
+ )
120
+ max_train_samples: Optional[int] = field(
121
+ default=None,
122
+ metadata={
123
+ "help": "For debugging purposes or quicker training, truncate the number of training examples to this "
124
+ "value if set."
125
+ },
126
+ )
127
+ max_eval_samples: Optional[int] = field(
128
+ default=None,
129
+ metadata={
130
+ "help": "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
131
+ "value if set."
132
+ },
133
+ )
134
+ overwrite_cache: bool = field(
135
+ default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
136
+ )
137
+ validation_split_percentage: Optional[int] = field(
138
+ default=5,
139
+ metadata={
140
+ "help": "The percentage of the train set used as validation set in case there's no validation split"
141
+ },
142
+ )
143
+ block_size: Optional[int] = field(
144
+ default=None,
145
+ metadata={
146
+ "help": "Optional input sequence length after tokenization. "
147
+ "The training dataset will be truncated in block of this size for training. "
148
+ "Default to the model max input length for single sentence inputs (take into account special tokens)."
149
+ },
150
+ )
151
+ overwrite_cache: bool = field(
152
+ default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
153
+ )
154
+ preprocessing_num_workers: Optional[int] = field(
155
+ default=None,
156
+ metadata={"help": "The number of processes to use for the preprocessing."},
157
+ )
158
+
159
+ def __post_init__(self):
160
+ if self.dataset_name is None and self.train_file is None and self.validation_file is None:
161
+ raise ValueError("Need either a dataset name or a training/validation file.")
162
+ else:
163
+ if self.train_file is not None:
164
+ extension = self.train_file.split(".")[-1]
165
+ assert extension in ["csv", "json", "txt"], "`train_file` should be a csv, a json or a txt file."
166
+ if self.validation_file is not None:
167
+ extension = self.validation_file.split(".")[-1]
168
+ assert extension in ["csv", "json", "txt"], "`validation_file` should be a csv, a json or a txt file."
169
+
170
+
171
+ class TrainState(train_state.TrainState):
172
+ dropout_rng: jnp.ndarray
173
+
174
+ def replicate(self):
175
+ return jax_utils.replicate(self).replace(dropout_rng=shard_prng_key(self.dropout_rng))
176
+
177
+
178
+ def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuffle: bool = False):
179
+ """
180
+ Returns batches of size `batch_size` from truncated `dataset`, sharded over all local devices.
181
+ Shuffle batches if `shuffle` is `True`.
182
+ """
183
+ steps_per_epoch = len(dataset) // batch_size
184
+
185
+ if shuffle:
186
+ batch_idx = jax.random.permutation(rng, len(dataset))
187
+ else:
188
+ batch_idx = jnp.arange(len(dataset))
189
+
190
+ batch_idx = batch_idx[: steps_per_epoch * batch_size] # Skip incomplete batch.
191
+ batch_idx = batch_idx.reshape((steps_per_epoch, batch_size))
192
+
193
+ for idx in batch_idx:
194
+ batch = dataset[idx]
195
+ batch = {k: jnp.array(v) for k, v in batch.items()}
196
+
197
+ batch = shard(batch)
198
+
199
+ yield batch
200
+
201
+
202
+ def write_train_metric(summary_writer, train_metrics, train_time, step):
203
+ summary_writer.scalar("train_time", train_time, step)
204
+
205
+ train_metrics = get_metrics(train_metrics)
206
+ for key, vals in train_metrics.items():
207
+ tag = f"train_{key}"
208
+ for i, val in enumerate(vals):
209
+ summary_writer.scalar(tag, val, step - len(vals) + i + 1)
210
+
211
+
212
+ def write_eval_metric(summary_writer, eval_metrics, step):
213
+ for metric_name, value in eval_metrics.items():
214
+ summary_writer.scalar(f"eval_{metric_name}", value, step)
215
+
216
+
217
+ def create_learning_rate_fn(
218
+ train_ds_size: int, train_batch_size: int, num_train_epochs: int, num_warmup_steps: int, learning_rate: float
219
+ ) -> Callable[[int], jnp.array]:
220
+ """Returns a linear warmup, linear_decay learning rate function."""
221
+ steps_per_epoch = train_ds_size // train_batch_size
222
+ num_train_steps = steps_per_epoch * num_train_epochs
223
+ warmup_fn = optax.linear_schedule(init_value=0.0, end_value=learning_rate, transition_steps=num_warmup_steps)
224
+ decay_fn = optax.linear_schedule(
225
+ init_value=learning_rate, end_value=0, transition_steps=num_train_steps - num_warmup_steps
226
+ )
227
+ schedule_fn = optax.join_schedules(schedules=[warmup_fn, decay_fn], boundaries=[num_warmup_steps])
228
+ return schedule_fn
229
+
230
+
231
+ def main():
232
+ # See all possible arguments in src/transformers/training_args.py
233
+ # or by passing the --help flag to this script.
234
+ # We now keep distinct sets of args, for a cleaner separation of concerns.
235
+
236
+ parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
237
+ if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
238
+ # If we pass only one argument to the script and it's the path to a json file,
239
+ # let's parse it to get our arguments.
240
+ model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
241
+ else:
242
+ model_args, data_args, training_args = parser.parse_args_into_dataclasses()
243
+
244
+ if (
245
+ os.path.exists(training_args.output_dir)
246
+ and os.listdir(training_args.output_dir)
247
+ and training_args.do_train
248
+ and not training_args.overwrite_output_dir
249
+ ):
250
+ raise ValueError(
251
+ f"Output directory ({training_args.output_dir}) already exists and is not empty."
252
+ "Use --overwrite_output_dir to overcome."
253
+ )
254
+
255
+ # Make one log on every process with the configuration for debugging.
256
+ logging.basicConfig(
257
+ format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
258
+ datefmt="%m/%d/%Y %H:%M:%S",
259
+ level=logging.INFO,
260
+ )
261
+ # Setup logging, we only want one process per machine to log things on the screen.
262
+ logger.setLevel(logging.INFO if jax.process_index() == 0 else logging.ERROR)
263
+ if jax.process_index() == 0:
264
+ datasets.utils.logging.set_verbosity_warning()
265
+ transformers.utils.logging.set_verbosity_info()
266
+ else:
267
+ datasets.utils.logging.set_verbosity_error()
268
+ transformers.utils.logging.set_verbosity_error()
269
+
270
+ # Set the verbosity to info of the Transformers logger (on main process only):
271
+ logger.info(f"Training/evaluation parameters {training_args}")
272
+
273
+ # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
274
+ # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
275
+ # (the dataset will be downloaded automatically from the datasets Hub).
276
+ #
277
+ # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
278
+ # 'text' is found. You can easily tweak this behavior (see below).
279
+ #
280
+ # In distributed training, the load_dataset function guarantees that only one local process can concurrently
281
+ # download the dataset.
282
+ if data_args.dataset_name is not None:
283
+
284
+ # loading the wiki data from the load and clean file
285
+ dataset = load_and_clean_wiki("da")
286
+ print("the dataset is", dataset)
287
+
288
+ # if "validation" not in dataset.keys():
289
+ # dataset["validation"] = load_dataset(
290
+ # data_args.dataset_name,
291
+ # data_args.dataset_config_name,
292
+ # split=f"train[:{data_args.validation_split_percentage}%]",
293
+ # cache_dir=model_args.cache_dir,
294
+ # )
295
+ # dataset["train"] = load_dataset(
296
+ # data_args.dataset_name,
297
+ # data_args.dataset_config_name,
298
+ # split=f"train[{data_args.validation_split_percentage}%:]",
299
+ # cache_dir=model_args.cache_dir,
300
+ # )
301
+ else:
302
+ data_files = {}
303
+ if data_args.train_file is not None:
304
+ data_files["train"] = data_args.train_file
305
+ if data_args.validation_file is not None:
306
+ data_files["validation"] = data_args.validation_file
307
+ extension = data_args.train_file.split(".")[-1]
308
+ if extension == "txt":
309
+ extension = "text"
310
+ dataset = load_dataset(extension, data_files=data_files, cache_dir=model_args.cache_dir)
311
+ # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
312
+ # https://huggingface.co/docs/datasets/loading_datasets.html.
313
+
314
+ # Load pretrained model and tokenizer
315
+
316
+ # Distributed training:
317
+ # The .from_pretrained methods guarantee that only one local process can concurrently
318
+ # download model & vocab.
319
+ if model_args.config_name:
320
+ config = AutoConfig.from_pretrained(model_args.config_name, cache_dir=model_args.cache_dir)
321
+ elif model_args.model_name_or_path:
322
+ config = AutoConfig.from_pretrained(model_args.model_name_or_path, cache_dir=model_args.cache_dir)
323
+ else:
324
+ config = CONFIG_MAPPING[model_args.model_type]()
325
+ logger.warning("You are instantiating a new config instance from scratch.")
326
+
327
+ if model_args.tokenizer_name:
328
+ tokenizer = AutoTokenizer.from_pretrained(
329
+ model_args.tokenizer_name, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
330
+ )
331
+ elif model_args.model_name_or_path:
332
+ tokenizer = AutoTokenizer.from_pretrained(
333
+ model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
334
+ )
335
+ else:
336
+ raise ValueError(
337
+ "You are instantiating a new tokenizer from scratch. This is not supported by this script."
338
+ "You can do it from another script, save it, and load it from here, using --tokenizer_name."
339
+ )
340
+
341
+ if model_args.model_name_or_path:
342
+ model = FlaxAutoModelForCausalLM.from_pretrained(
343
+ model_args.model_name_or_path, config=config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
344
+ )
345
+ else:
346
+ model = FlaxAutoModelForCausalLM.from_config(
347
+ config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
348
+ )
349
+
350
+ # Preprocessing the datasets.
351
+ # First we tokenize all the texts.
352
+ # if training_args.do_train:
353
+ # column_names = dataset["train"].column_names
354
+ # else:
355
+ # column_names = dataset["validation"].column_names
356
+ text_column_name = "text"
357
+
358
+ # since this will be pickled to avoid _LazyModule error in Hasher force logger loading before tokenize_function
359
+ tok_logger = transformers.utils.logging.get_logger("transformers.tokenization_utils_base")
360
+
361
+ def tokenize_function(examples):
362
+ with CaptureLogger(tok_logger) as cl:
363
+ output = tokenizer(examples[text_column_name])
364
+ # clm input could be much much longer than block_size
365
+ if "Token indices sequence length is longer than the" in cl.out:
366
+ tok_logger.warning(
367
+ "^^^^^^^^^^^^^^^^ Please ignore the warning above - this long input will be chunked into smaller bits before being passed to the model."
368
+ )
369
+ return output
370
+
371
+ tokenized_datasets = dataset.map(
372
+ tokenize_function,
373
+ batched=True,
374
+ num_proc=data_args.preprocessing_num_workers,
375
+ remove_columns=dataset.column_names,
376
+ load_from_cache_file=not data_args.overwrite_cache,
377
+ )
378
+
379
+ if data_args.block_size is None:
380
+ block_size = tokenizer.model_max_length
381
+ if block_size > config.max_position_embeddings:
382
+ logger.warning(
383
+ f"The tokenizer picked seems to have a very large `model_max_length` ({tokenizer.model_max_length}). "
384
+ "Picking 1024 instead. You can change that default value by passing --block_size xxx."
385
+ )
386
+ block_size = 1024
387
+ else:
388
+ if data_args.block_size > tokenizer.model_max_length:
389
+ logger.warning(
390
+ f"The block_size passed ({data_args.block_size}) is larger than the maximum length for the model"
391
+ f"({tokenizer.model_max_length}). Using block_size={tokenizer.model_max_length}."
392
+ )
393
+ block_size = min(data_args.block_size, tokenizer.model_max_length)
394
+
395
+ # Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size.
396
+ def group_texts(examples):
397
+ # Concatenate all texts.
398
+ # print("the examples are", examples)
399
+ # import pdb
400
+ # pdb.set_trace()
401
+ concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
402
+ total_length = len(concatenated_examples[list(examples.keys())[0]])
403
+ # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
404
+ # customize this part to your needs.
405
+ if total_length >= block_size:
406
+ total_length = (total_length // block_size) * block_size
407
+ # Split by chunks of max_len.
408
+ result = {
409
+ k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
410
+ for k, t in concatenated_examples.items()
411
+ }
412
+ result["labels"] = result["input_ids"].copy()
413
+ return result
414
+
415
+ # Note that with `batched=True`, this map processes 1,000 texts together, so group_texts throws away a remainder
416
+ # for each of those groups of 1,000 texts. You can adjust that batch_size here but a higher value might be slower
417
+ # to preprocess.
418
+ #
419
+ # To speed up this part, we use multiprocessing. See the documentation of the map method for more information:
420
+ # https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.map
421
+
422
+ lm_datasets = tokenized_datasets.map(
423
+ group_texts,
424
+ batched=True,
425
+ num_proc=data_args.preprocessing_num_workers,
426
+ load_from_cache_file=not data_args.overwrite_cache,
427
+ )
428
+
429
+ if training_args.do_train:
430
+ # if "train" not in tokenized_datasets:
431
+ # raise ValueError("--do_train requires a train dataset")
432
+ # train_dataset = lm_datasets["train"]
433
+ train_dataset = lm_datasets
434
+ if data_args.max_train_samples is not None:
435
+ train_dataset = train_dataset.select(range(data_args.max_train_samples))
436
+
437
+ # if training_args.do_eval:
438
+ # if "validation" not in tokenized_datasets:
439
+ # raise ValueError("--do_eval requires a validation dataset")
440
+ # eval_dataset = lm_datasets["validation"]
441
+ # if data_args.max_eval_samples is not None:
442
+ # eval_dataset = eval_dataset.select(range(data_args.max_eval_samples))
443
+
444
+ # Enable tensorboard only on the master node
445
+ has_tensorboard = is_tensorboard_available()
446
+ if has_tensorboard and jax.process_index() == 0:
447
+ try:
448
+ from flax.metrics.tensorboard import SummaryWriter
449
+
450
+ summary_writer = SummaryWriter(log_dir=Path(training_args.output_dir))
451
+ except ImportError as ie:
452
+ has_tensorboard = False
453
+ logger.warning(
454
+ f"Unable to display metrics through TensorBoard because some package are not installed: {ie}"
455
+ )
456
+ else:
457
+ logger.warning(
458
+ "Unable to display metrics through TensorBoard because the package is not installed: "
459
+ "Please run pip install tensorboard to enable."
460
+ )
461
+
462
+ # Initialize our training
463
+ rng = jax.random.PRNGKey(training_args.seed)
464
+ rng, dropout_rng = jax.random.split(rng)
465
+
466
+ # Store some constant
467
+ num_epochs = int(training_args.num_train_epochs)
468
+ train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()
469
+ eval_batch_size = int(training_args.per_device_eval_batch_size) * jax.device_count()
470
+ steps_per_epoch = len(train_dataset) // train_batch_size
471
+ total_train_steps = steps_per_epoch * num_epochs
472
+
473
+ # Create learning rate schedule
474
+ linear_decay_lr_schedule_fn = create_learning_rate_fn(
475
+ len(train_dataset),
476
+ train_batch_size,
477
+ training_args.num_train_epochs,
478
+ training_args.warmup_steps,
479
+ training_args.learning_rate,
480
+ )
481
+
482
+ # We use Optax's "masking" functionality to not apply weight decay
483
+ # to bias and LayerNorm scale parameters. decay_mask_fn returns a
484
+ # mask boolean with the same structure as the parameters.
485
+ # The mask is True for parameters that should be decayed.
486
+ # Note that this mask is specifically adapted for FlaxGPT2.
487
+ # For other models, one should correct the layer norm parameter naming
488
+ # accordingly.
489
+ def decay_mask_fn(params):
490
+ flat_params = traverse_util.flatten_dict(params)
491
+ flat_mask = {
492
+ path: (path[-1] != "bias" and path[-2:] not in [("ln_1", "scale"), ("ln_2", "scale"), ("ln_f", "scale")])
493
+ for path in flat_params
494
+ }
495
+ return traverse_util.unflatten_dict(flat_mask)
496
+
497
+ # create adam optimizer
498
+ if training_args.adafactor:
499
+ # We use the default parameters here to initialize adafactor,
500
+ # For more details about the parameters please check https://github.com/deepmind/optax/blob/ed02befef9bf81cbbf236be3d2b0e032e9ed4a40/optax/_src/alias.py#L74
501
+ optimizer = optax.adafactor(
502
+ learning_rate=linear_decay_lr_schedule_fn,
503
+ )
504
+ else:
505
+ optimizer = optax.adamw(
506
+ learning_rate=linear_decay_lr_schedule_fn,
507
+ b1=training_args.adam_beta1,
508
+ b2=training_args.adam_beta2,
509
+ eps=training_args.adam_epsilon,
510
+ weight_decay=training_args.weight_decay,
511
+ mask=decay_mask_fn,
512
+ )
513
+
514
+ # Setup train state
515
+ state = TrainState.create(apply_fn=model.__call__, params=model.params, tx=optimizer, dropout_rng=dropout_rng)
516
+
517
+ def loss_fn(logits, labels):
518
+ shift_logits = logits[..., :-1, :]
519
+ shift_labels = labels[..., 1:]
520
+ loss = optax.softmax_cross_entropy(shift_logits, onehot(shift_labels, shift_logits.shape[-1]))
521
+ return loss.mean()
522
+
523
+ # Define gradient update step fn
524
+ def train_step(state, batch):
525
+ dropout_rng, new_dropout_rng = jax.random.split(state.dropout_rng)
526
+
527
+ def compute_loss(params):
528
+ labels = batch.pop("labels")
529
+ logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
530
+ loss = loss_fn(logits, labels)
531
+ return loss
532
+
533
+ grad_fn = jax.value_and_grad(compute_loss)
534
+ loss, grad = grad_fn(state.params)
535
+ grad = jax.lax.pmean(grad, "batch")
536
+
537
+ new_state = state.apply_gradients(grads=grad, dropout_rng=new_dropout_rng)
538
+
539
+ metrics = {"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(state.step)}
540
+ metrics = jax.lax.pmean(metrics, axis_name="batch")
541
+
542
+ return new_state, metrics
543
+
544
+ # Define eval fn
545
+ def eval_step(params, batch):
546
+ labels = batch.pop("labels")
547
+ logits = model(**batch, params=params, train=False)[0]
548
+ loss = loss_fn(logits, labels)
549
+
550
+ # summarize metrics
551
+ metrics = {"loss": loss}
552
+ metrics = jax.lax.pmean(metrics, axis_name="batch")
553
+ return metrics
554
+
555
+ # Create parallel version of the train and eval step
556
+ p_train_step = jax.pmap(train_step, "batch", donate_argnums=(0,))
557
+ p_eval_step = jax.pmap(eval_step, "batch")
558
+
559
+ # Replicate the train state on each device
560
+ state = state.replicate()
561
+
562
+ logger.info("***** Running training *****")
563
+ logger.info(f" Num examples = {len(train_dataset)}")
564
+ logger.info(f" Num Epochs = {num_epochs}")
565
+ logger.info(f" Instantaneous batch size per device = {training_args.per_device_train_batch_size}")
566
+ logger.info(f" Total train batch size (w. parallel & distributed) = {train_batch_size}")
567
+ logger.info(f" Total optimization steps = {total_train_steps}")
568
+
569
+ train_time = 0
570
+ train_metrics = []
571
+ epochs = tqdm(range(num_epochs), desc=f"Epoch ... (1/{num_epochs})", position=0)
572
+ for epoch in epochs:
573
+ # ======================== Training ================================
574
+ train_start = time.time()
575
+
576
+ # Create sampling rng
577
+ rng, input_rng = jax.random.split(rng)
578
+
579
+ # Generate an epoch by shuffling sampling indices from the train dataset
580
+ train_loader = data_loader(input_rng, train_dataset, train_batch_size, shuffle=True)
581
+ steps_per_epoch = len(train_dataset) // train_batch_size
582
+ # train
583
+ for step in tqdm(range(steps_per_epoch), desc="Training...", position=1, leave=False):
584
+ batch = next(train_loader)
585
+ state, train_metric = p_train_step(state, batch)
586
+ train_metrics.append(train_metric)
587
+
588
+ cur_step = epoch * (len(train_dataset) // train_batch_size) + step
589
+
590
+ if cur_step % training_args.logging_steps == 0 and cur_step > 0:
591
+ # Save metrics
592
+ train_metric = unreplicate(train_metric)
593
+ train_time += time.time() - train_start
594
+ if has_tensorboard and jax.process_index() == 0:
595
+ write_train_metric(summary_writer, train_metrics, train_time, cur_step)
596
+
597
+ epochs.write(
598
+ f"Step... ({cur_step} | Loss: {train_metric['loss'].mean()}, Learning Rate: {train_metric['learning_rate'].mean()})"
599
+ )
600
+
601
+ train_metrics = []
602
+
603
+ # if cur_step % training_args.eval_steps == 0 and cur_step > 0:
604
+ # # ======================== Evaluating ==============================
605
+ # eval_metrics = []
606
+ # eval_loader = data_loader(input_rng, eval_dataset, eval_batch_size)
607
+ # eval_steps = len(eval_dataset) // eval_batch_size
608
+ # for _ in tqdm(range(eval_steps), desc="Evaluating...", position=2, leave=False):
609
+ # # Model forward
610
+ # batch = next(eval_loader)
611
+ # metrics = p_eval_step(state.params, batch)
612
+ # eval_metrics.append(metrics)
613
+
614
+ # # normalize eval metrics
615
+ # eval_metrics = get_metrics(eval_metrics)
616
+ # eval_metrics = jax.tree_map(jnp.mean, eval_metrics)
617
+
618
+ # try:
619
+ # eval_metrics["perplexity"] = math.exp(eval_metrics["loss"])
620
+ # except OverflowError:
621
+ # eval_metrics["perplexity"] = float("inf")
622
+
623
+ # # Print metrics and update progress bar
624
+ # desc = f"Step... ({cur_step} | Eval Loss: {eval_metrics['loss']} | Eval Perplexity: {eval_metrics['perplexity']})"
625
+ # epochs.write(desc)
626
+ # epochs.desc = desc
627
+
628
+ # # Save metrics
629
+ # if has_tensorboard and jax.process_index() == 0:
630
+ # write_eval_metric(summary_writer, eval_metrics, cur_step)
631
+
632
+ if cur_step % training_args.save_steps == 0 and cur_step > 0:
633
+ # save checkpoint after each epoch and push checkpoint to the hub
634
+ if jax.process_index() == 0:
635
+ params = jax.device_get(unreplicate(state.params))
636
+ model.save_pretrained(
637
+ training_args.output_dir,
638
+ params=params,
639
+ push_to_hub=training_args.push_to_hub,
640
+ commit_message=f"Saving weights and logs of step {cur_step}",
641
+ )
642
+
643
+
644
+ if __name__ == "__main__":
645
+ main()
save_model.py ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ from transformers import AutoTokenizer, GPT2LMHeadModel
2
+ '''
3
+ This is a script to convert the Jax model and the tokenizer to Pytorch model
4
+ '''
5
+ model = GPT2LMHeadModel.from_pretrained(".", from_flax=True)
6
+ model.save_pretrained(".")
7
+ tokenizer = AutoTokenizer.from_pretrained(".")
8
+ tokenizer.save_pretrained(".")
special_tokens_map.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {"bos_token": "<|endoftext|>", "eos_token": "<|endoftext|>", "unk_token": "<|endoftext|>"}
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"}
vocab.json ADDED
The diff for this file is too large to render. See raw diff