2 |
3 | # **GPT2** for Multiple Languages
4 |
5 | [](https://colab.research.google.com/github/imcaspar/gpt2-ml/blob/master/pretrained_model_demo.ipynb)
6 | [](https://github.com/imcaspar/gpt2-ml)
7 | [](https://github.com/imcaspar/gpt2-ml/releases)
8 | [](https://github.com/imcaspar/gpt2-ml/issues)
9 | [](https://github.com/imcaspar/gpt2-ml)
10 |
11 | [**中文说明**](./README_CN.md) | [**English**](./README.md)
12 |
13 | - [x] Simplifed GPT2 train scripts（based on Grover, supporting TPUs）
14 | - [x] Ported bert tokenizer, multilingual corpus compatible
15 | - [x] 1.5B GPT2 pretrained Chinese model ( ~15G corpus, 10w steps )
16 | - [x] Batteries-included Colab demo [#](https://github.com/imcaspar/gpt2-ml#google-colab)
17 | - [x] 1.5B GPT2 pretrained Chinese model ( ~30G corpus, 22w steps )
18 |
19 |
20 | ## Pretrained Model
21 | | Size | Language | Corpus | Vocab | Link1 | Link2 | SHA256 |
22 | | ---- | -------- | ------ | ----- | ----- | ----- | ------ |
23 | | 1.5B Params | Chinese | ~30G | CLUE ( 8021 tokens ) | [**Google Drive**](https://drive.google.com/file/d/1mT_qCQg4AWnAXTwKfsyyRWCRpgPrBJS3) | [**Baidu Pan (ffz6)**](https://pan.baidu.com/s/1yiuTHXUr2DpyBqmFYLJH6A) | e698cc97a7f5f706f84f58bb469d614e51d3c0ce5f9ab9bf77e01e3fcb41d482 | 24 | | 1.5B Params | Chinese | ~15G | Bert ( 21128 tokens ) | [**Google Drive**](https://drive.google.com/file/d/1IzWpQ6I2IgfV7CldZvFJnZ9byNDZdO4n) | [**Baidu Pan (q9vr)**](https://pan.baidu.com/s/1TA_3e-u2bXg_hcx_NwVbGw) | 4a6e5124df8db7ac2bdd902e6191b807
a6983a7f5d09fb10ce011f9a073b183e | 25 | 26 | Corpus from [THUCNews](http://thuctc.thunlp.org/#%E4%B8%AD%E6%96%87%E6%96%87%E6%9C%AC%E5%88%86%E7%B1%BB%E6%95%B0%E6%8D%AE%E9%9B%86THUCNews) and [nlp_chinese_corpus](https://github.com/brightmart/nlp_chinese_corpus) 27 | 28 | Using [Cloud TPU Pod v3-256](https://cloud.google.com/tpu/docs/types-zones#types) to train 22w steps 29 | 30 |  31 | 32 | 33 | ## Google Colab 34 | With just 2 clicks (not including Colab auth process), the 1.5B pretrained Chinese model demo is ready to go: 35 | 36 | [**[Colab Notebook]**](https://colab.research.google.com/github/imcaspar/gpt2-ml/blob/master/pretrained_model_demo.ipynb) 37 | 38 |
39 |
40 | ## Train
41 |
42 | ## Disclaimer
43 | The contents in this repository are for academic research purpose, and we do not provide any conclusive remarks.
44 |
45 | ## Citation
46 |
47 | ```
48 | @misc{GPT2-ML,
49 | author = {Zhibo Zhang},
50 | title = {GPT2-ML: GPT-2 for Multiple Languages},
51 | year = {2019},
52 | publisher = {GitHub},
53 | journal = {GitHub repository},
54 | howpublished = {\url{https://github.com/imcaspar/gpt2-ml}},
55 | }
56 | ```
57 |
58 | ## Reference
59 | https://github.com/google-research/bert
60 |
61 | https://github.com/rowanz/grover
62 |
63 | Research supported with Cloud TPUs from Google's TensorFlow Research Cloud (TFRC)
64 |
65 | ## Press
66 | [[机器之心] 只需单击三次，让中文GPT-2为你生成定制故事](https://mp.weixin.qq.com/s/FpoSNNKZSQOE2diPvJDHog)
67 |
68 | [[科学空间] 现在可以用Keras玩中文GPT2了](https://kexue.fm/archives/7292)
69 |
--------------------------------------------------------------------------------
/README_CN.md:
--------------------------------------------------------------------------------
1 |
2 |
3 | # **GPT2** for Multiple Languages
4 |
5 | [](https://colab.research.google.com/github/imcaspar/gpt2-ml/blob/master/pretrained_model_demo.ipynb)
6 | [](https://github.com/imcaspar/gpt2-ml)
7 | [](https://github.com/imcaspar/gpt2-ml/releases)
8 | [](https://github.com/imcaspar/gpt2-ml/issues)
9 | [](https://github.com/imcaspar/gpt2-ml)
10 |
11 | [**中文说明**](./README_CN.md) | [**English**](./README.md)
12 |
13 | - [x] 简化整理 GPT2 训练代码（based on Grover, supporting TPUs）
14 | - [x] 移植 bert tokenizer，添加多语言支持
15 | - [x] 15亿参数 GPT2 中文预训练模型( 15G 语料，训练 10w 步 )
16 | - [x] 开箱即用的模型生成效果 demo [#](https://github.com/imcaspar/gpt2-ml#google-colab)
17 | - [x] 15亿参数 GPT2 中文预训练模型( 30G 语料，训练 22w 步 )
18 |
19 |
20 | ## 预训练模型
21 | | Size | Language | Corpus | Vocab | Link1 | Link2 | SHA256 |
22 | | ---- | -------- | ------ | ----- | ----- | ----- | ------ |
23 | | 1.5B Params | Chinese | ~30G | CLUE ( 8021 tokens ) | [**Google Drive**](https://drive.google.com/file/d/1mT_qCQg4AWnAXTwKfsyyRWCRpgPrBJS3) | [**Baidu Pan (ffz6)**](https://pan.baidu.com/s/1yiuTHXUr2DpyBqmFYLJH6A) | e698cc97a7f5f706f84f58bb469d614e51d3c0ce5f9ab9bf77e01e3fcb41d482 | 24 | | 1.5B Params | Chinese | ~15G | Bert ( 21128 tokens ) | [**Google Drive**](https://drive.google.com/file/d/1IzWpQ6I2IgfV7CldZvFJnZ9byNDZdO4n) | [**Baidu Pan (q9vr)**](https://pan.baidu.com/s/1TA_3e-u2bXg_hcx_NwVbGw) | 4a6e5124df8db7ac2bdd902e6191b807
a6983a7f5d09fb10ce011f9a073b183e | 25 | 26 | 训练语料来自 [THUCNews](http://thuctc.thunlp.org/#%E4%B8%AD%E6%96%87%E6%96%87%E6%9C%AC%E5%88%86%E7%B1%BB%E6%95%B0%E6%8D%AE%E9%9B%86THUCNews) 以及 [nlp_chinese_corpus](https://github.com/brightmart/nlp_chinese_corpus)，清洗后总文本量约 15G 27 | 28 | 使用 [Cloud TPU Pod v3-256](https://cloud.google.com/tpu/docs/types-zones#types) 训练 22w 步 29 | 30 |  31 | 32 | 33 | ## Google Colab 34 | 只需两次鼠标点击（不包括 Colab 授权流程），体验 15 亿参数中文预训练模型生成效果： 35 | 36 | [**[Colab Notebook]**](https://colab.research.google.com/github/imcaspar/gpt2-ml/blob/master/pretrained_model_demo.ipynb) 37 | 38 |
39 |
40 | ## 训练
41 |
42 | ## 免责声明
43 | 该项目中的内容仅供技术研究参考，不作为任何结论性依据。
44 |
45 | ## Citation
46 |
47 | ```
48 | @misc{GPT2-ML,
49 | author = {Zhibo Zhang},
50 | title = {GPT2-ML: GPT-2 for Multiple Languages},
51 | year = {2019},
52 | publisher = {GitHub},
53 | journal = {GitHub repository},
54 | howpublished = {\url{https://github.com/imcaspar/gpt2-ml}},
55 | }
56 | ```
57 |
58 | ## Reference
59 | https://github.com/google-research/bert
60 |
61 | https://github.com/rowanz/grover
62 |
63 | Research supported with Cloud TPUs from Google's TensorFlow Research Cloud (TFRC)
64 |
65 | ## Press
66 | [[机器之心] 只需单击三次，让中文GPT-2为你生成定制故事](https://mp.weixin.qq.com/s/FpoSNNKZSQOE2diPvJDHog)
67 |
68 | [[科学空间] 现在可以用Keras玩中文GPT2了](https://kexue.fm/archives/7292)
--------------------------------------------------------------------------------
/configs/base.json:
--------------------------------------------------------------------------------
1 | {
2 | "vocab_size": 50270,
3 | "hidden_size": 768,
4 | "attention_probs_dropout_prob": 0.1,
5 | "hidden_dropout_prob": 0.1,
6 | "hidden_act": "gelu",
7 | "initializer_range": 0.02,
8 | "intermediate_size": 3072,
9 | "max_position_embeddings": 1024,
10 | "num_attention_heads": 12,
11 | "num_hidden_layers": 12
12 | }
--------------------------------------------------------------------------------
/configs/large.json:
--------------------------------------------------------------------------------
1 | {
2 | "vocab_size": 8021,
3 | "hidden_size": 1024,
4 | "attention_probs_dropout_prob": 0.1,
5 | "hidden_dropout_prob": 0.1,
6 | "hidden_act": "gelu",
7 | "initializer_range": 0.02,
8 | "intermediate_size": 4096,
9 | "max_position_embeddings": 1024,
10 | "num_attention_heads": 16,
11 | "num_hidden_layers": 24
12 | }
--------------------------------------------------------------------------------
/configs/mega.json:
--------------------------------------------------------------------------------
1 | {
2 | "vocab_size": 8021,
3 | "hidden_size": 1536,
4 | "attention_probs_dropout_prob": 0.1,
5 | "hidden_dropout_prob": 0.1,
6 | "hidden_act": "gelu",
7 | "initializer_range": 0.014142135623731,
8 | "intermediate_size": 6144,
9 | "max_position_embeddings": 1024,
10 | "num_attention_heads": 24,
11 | "num_hidden_layers": 48
12 | }
13 |
--------------------------------------------------------------------------------
/dataset/README.md:
--------------------------------------------------------------------------------
1 | curl -o ~/miniconda.sh -O https://repo.anaconda.com/miniconda/Miniconda2-latest-Linux-x86_64.sh
2 | chmod +x ~/miniconda.sh
3 | ~/miniconda.sh -b -p ~/conda
4 | rm ~/miniconda.sh
5 | ~/conda/bin/conda install -y python=3.7
6 | ~/conda/bin/conda init // exit shell
7 |
8 | conda create -n gpt3 python=3.7
9 |
10 |
11 | sudo apt install parallel
12 | pip install ujson==2.0.3
13 |
14 | export PYTHONPATH=$(pwd) //project path
--------------------------------------------------------------------------------
/dataset/prepare_data.py:
--------------------------------------------------------------------------------
1 | """
2 | Turn a merged corpus into tfrecord files.
3 |
4 | NOTE: You will want to do this using several processes. I did this on an AWS machine with 72 CPUs using GNU parallel
5 | as that's where I had the deduplicated RealNews dataset.
6 | """
7 | import argparse
8 | import ujson as json
9 | # from sample.encoder import get_encoder, tokenize_for_grover_training, detokenize, sliding_window, create_int_feature
10 | import random
11 | import tensorflow.compat.v1 as tf
12 | import collections
13 | import os
14 | from tempfile import TemporaryDirectory
15 |
16 | from tokenization import tokenization
17 |
18 | parser = argparse.ArgumentParser(description='SCRAPE!')
19 | parser.add_argument(
20 | '-fold',
21 | dest='fold',
22 | default=0,
23 | type=int,
24 | help='which fold we are on'
25 | )
26 | parser.add_argument(
27 | '-num_folds',
28 | dest='num_folds',
29 | default=1,
30 | type=int,
31 | help='Number of folds (corresponding to both the number of training files and the number of testing files)',
32 | )
33 | parser.add_argument(
34 | '-seed',
35 | dest='seed',
36 | default=1337,
37 | type=int,
38 | help='which seed to use'
39 | )
40 | parser.add_argument(
41 | '-base_fn',
42 | dest='base_fn',
43 | default='news2016zh_',
44 | type=str,
45 | help='We will output files that are like {base_fn}_{n}.tfrecord for n in 0, ..., 1023'
46 | )
47 |
48 | parser.add_argument(
49 | '-input_fn',
50 | dest='input_fn',
51 | default='realnews.jsonl',
52 | type=str,
53 | help='Base filename to use. THIS MUST BE A LOCAL FILE.'
54 | )
55 | parser.add_argument(
56 | '-max_seq_length',
57 | dest='max_seq_length',
58 | default=1024,
59 | type=int,
60 | help='Max sequence length',
61 | )
62 |
63 |
64 | args = parser.parse_args()
65 | random.seed(args.seed + args.fold)
66 |
67 | #encoder = get_encoder()
68 | tokenizer = tokenization.FullTokenizer(
69 | vocab_file="clue-vocab.txt", do_lower_case=True)
70 |
71 |
72 | class TFRecordWriter(object):
73 | def __init__(self, fn):
74 | self.fn = fn
75 | if fn.startswith('gs://'):
76 | from google.cloud import storage
77 | self.s3client = None
78 | self.gclient = storage.Client()
79 | self.storage_dir = TemporaryDirectory()
80 | self.writer = tf.python_io.TFRecordWriter(
81 | os.path.join(self.storage_dir.name, 'temp.tfrecord'))
82 | self.bucket_name, self.file_name = self.fn.split(
83 | 'gs://', 1)[1].split('/', 1)
84 |
85 | else:
86 | self.s3client = None
87 | self.gclient = None
88 | self.bucket_name = None
89 | self.file_name = None
90 | self.storage_dir = None
91 | self.writer = tf.python_io.TFRecordWriter(fn)
92 |
93 | def write(self, x):
94 | self.writer.write(x)
95 |
96 | def close(self):
97 | self.writer.close()
98 |
99 | if self.gclient is not None:
100 | bucket = self.gclient.get_bucket(self.bucket_name)
101 | blob = bucket.blob(self.file_name)
102 | blob.upload_from_filename(os.path.join(
103 | self.storage_dir.name, 'temp.tfrecord'))
104 | self.storage_dir.cleanup()
105 |
106 | def __enter__(self):
107 | # Called when entering "with" context.
108 | return self
109 |
110 | def __exit__(self, *_):
111 | # Called when exiting "with" context.
112 | # Upload shit
113 | print("CALLING CLOSE")
114 | self.close()
115 |
116 |
117 | def article_iterator(tokenizer):
118 | """ Iterate through the provided filename + tokenize"""
119 | assert os.path.exists(args.input_fn)
120 | for (dirpath, dirnames, filenames) in os.walk(args.input_fn):
121 | for filename in filenames:
122 | with open(os.path.join(dirpath, filename), 'r') as f:
123 | for l_no, l in enumerate(f):
124 | if l_no % args.num_folds == args.fold:
125 | article = json.loads(l)
126 |
127 | line = tokenization.convert_to_unicode(
128 | article['text']) # for news2016zh text body
129 | tokens = tokenizer.tokenize(line)
130 | input_ids = tokenizer.convert_tokens_to_ids(tokens)
131 |
132 | article['input_ids'] = input_ids
133 |
134 | article['inst_index'] = (l_no // args.num_folds)
135 | if article['inst_index'] < 100:
136 | print('---\nINPUT{}. {}\n---\nTokens: {}\n'.format(article['inst_index'],
137 | tokens,
138 | input_ids
139 | ), flush=True)
140 | if len(article['input_ids']) <= 64: # min size of article
141 | continue
142 | yield article
143 |
144 |
145 | def create_int_feature(values):
146 | feature = tf.train.Feature(
147 | int64_list=tf.train.Int64List(value=list(values)))
148 | return feature
149 |
150 |
151 | def buffered_and_sliding_window_article_iterator(tokenizer, final_desired_size=1025):
152 | """ We apply a sliding window to fix long sequences, and use a buffer that combines short sequences."""
153 | for article in article_iterator(tokenizer):
154 | if len(article['input_ids']) >= final_desired_size:
155 | article['input_ids'] = article['input_ids'][0:final_desired_size-1]
156 | while len(article['input_ids']) < final_desired_size:
157 | article['input_ids'].append(0)
158 | yield article
159 |
160 |
161 | # OK now write the tfrecord file
162 | total_written = 0
163 | train_file = args.base_fn + 'train_wiki19_{:04d}.tfrecord'.format(args.fold)
164 | with TFRecordWriter(train_file) as train_writer:
165 | for article in buffered_and_sliding_window_article_iterator(tokenizer,
166 | final_desired_size=args.max_seq_length + 1):
167 | writer2use = train_writer
168 | assert len(article['input_ids']) == (args.max_seq_length + 1)
169 |
170 | features = collections.OrderedDict()
171 | features["input_ids"] = create_int_feature(article['input_ids'])
172 | tf_example = tf.train.Example(
173 | features=tf.train.Features(feature=features))
174 |
175 | writer2use.write(tf_example.SerializeToString())
176 | total_written += 1
177 |
178 | # DEBUG
179 | if article['inst_index'] < 5:
180 | print("~~~\nIndex {}. ARTICLE: {}\n---\nTokens: {}\n\n".format(article['inst_index'],
181 | tokenizer.convert_ids_to_tokens(
182 | article['input_ids']),
183 | article['input_ids']
184 | ), flush=True)
185 | if article['inst_index'] % 1000 == 0:
186 | print("{} articles, {} written".format(
187 | article['inst_index'], total_written), flush=True)
188 | print("DONE UPLOADING", flush=True)
189 |
--------------------------------------------------------------------------------
/dataset/prepare_data.sh:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env bash
2 |
3 | NUM_FOLDS=1024
4 | MAX_SEQ_LENGTH=1024
5 | FN=${1}
6 | OUT_BUCKET=${2}
7 |
8 | rm -rf logs_${MAX_SEQ_LENGTH}
9 | mkdir logs_${MAX_SEQ_LENGTH}
10 | parallel -j $(nproc --all) --will-cite "python prepare_data.py -fold {1} -num_folds ${NUM_FOLDS} -base_fn gs://${OUT_BUCKET}/data_${MAX_SEQ_LENGTH}/ -input_fn ${FN} -max_seq_length ${MAX_SEQ_LENGTH} > logs_${MAX_SEQ_LENGTH}/log{1}.txt" ::: $(seq 0 $((${NUM_FOLDS}-1)))
11 |
--------------------------------------------------------------------------------
/dockerfiles/gpu-jupyter.Dockerfile:
--------------------------------------------------------------------------------
1 | FROM tensorflow/tensorflow:1.15.0-gpu-py3-jupyter
2 |
3 | RUN apt update && apt install -y --no-install-recommends git
4 | RUN git clone -q https://github.com/imcaspar/gpt2-ml && mkdir -p gpt2-ml/models/mega
5 |
6 | WORKDIR /gpt2-ml
7 |
8 | RUN perl 3rd/gdown.pl/gdown.pl https://drive.google.com/open?id=1n_5-tgPpQ1gqbyLPbP1PwiFi2eo7SWw_ models/mega/model.ckpt-100000.data-00000-of-00001
9 | RUN wget -q --show-progress https://github.com/imcaspar/gpt2-ml/releases/download/v0.5/model.ckpt-100000.index -P models/mega
10 | RUN wget -q --show-progress https://github.com/imcaspar/gpt2-ml/releases/download/v0.5/model.ckpt-100000.meta -P models/mega
11 |
12 | CMD ["bash", "-c", "jupyter notebook --ip 0.0.0.0 --no-browser --allow-root pretrained_model_demo.ipynb"]
--------------------------------------------------------------------------------
/pretrained_model_demo.ipynb:
--------------------------------------------------------------------------------
1 | {
2 | "cells": [
3 | {
4 | "cell_type": "markdown",
5 | "metadata": {},
6 | "source": [
7 | "[](https://github.com/imcaspar/gpt2-ml)\n",
8 | "[](https://github.com/imcaspar/gpt2-ml)\n",
9 | "[](https://twitter.com/intent/tweet?text=Wow:&url=https://github.com/imcaspar/gpt2-ml)\n",
10 | "### Instructions for running:\n",
11 | "\n",
12 | "* Press the ▶️button on the left of each of the cells\n",
13 | "* View the code: Double click any of the cells\n",
14 | "* Hide the code: Double click the right side of the cell"
15 | ]
16 | },
17 | {
18 | "cell_type": "code",
19 | "execution_count": null,
20 | "metadata": {},
21 | "outputs": [],
22 | "source": [
23 | "#@title #Prerequisites\n",
24 | "#%tensorflow_version 1.x\n",
25 | "!pip install -I tensorflow-gpu==1.15.4 &> tmp.log\n",
26 | "!git clone -q https://github.com/imcaspar/gpt2-ml\n",
27 | "%cd /content/gpt2-ml\n",
28 | "!mkdir -p /content/gpt2-ml/models/mega\n",
29 | "\n",
30 | "!perl 3rd/gdown.pl/gdown.pl https://drive.google.com/open?id=1mT_qCQg4AWnAXTwKfsyyRWCRpgPrBJS3 models/mega/model.ckpt-220000.data-00000-of-00001\n",
31 | "!wget -q --show-progress https://github.com/imcaspar/gpt2-ml/releases/download/v1.0/model.ckpt-220000.index -P models/mega\n",
32 | "!wget -q --show-progress https://github.com/imcaspar/gpt2-ml/releases/download/v1.0/model.ckpt-220000.meta -P models/mega\n",
33 | "!echo 'Download finished.🍺'\n",
34 | "# If gdown.pl failed, please uncomment following code & exec\n",
35 | "# !python scripts/down_gdrive_file.py -file_id='1mT_qCQg4AWnAXTwKfsyyRWCRpgPrBJS3' -file_path='models/mega/model.ckpt-220000.data-00000-of-00001'"
36 | ]
37 | },
38 | {
39 | "cell_type": "code",
40 | "execution_count": null,
41 | "metadata": {},
42 | "outputs": [],
43 | "source": [
44 | "#@title #Inference\n",
45 | "min_len = 150#@param {type:\"number\", min:5, max:1024, step:1}\n",
46 | "sp_num = 5#@param {type:\"number\", min:1, max:50, step:1}\n",
47 | "!PYTHONPATH=$(pwd) python scripts/demo.py -ckpt_fn models/mega/model.ckpt-220000 -min_len $min_len -samples $sp_num"
48 | ]
49 | }
50 | ],
51 | "metadata": {
52 | "colab": {
53 | "name": "15 亿参数 GPT2 中文预训练模型 | 1.5B GPT2 Pretrained Chinese Model",
54 | "provenance": [],
55 | "collapsed_sections": []
56 | },
57 | "kernelspec": {
58 | "name": "python3",
59 | "display_name": "Python 3"
60 | },
61 | "accelerator": "GPU"
62 | },
63 | "nbformat": 4,
64 | "nbformat_minor": 0
65 | }
--------------------------------------------------------------------------------
/requirements-gpu.txt:
--------------------------------------------------------------------------------
1 | pandas==0.24.2
2 | regex==2019.4.14
3 | h5py==2.10.0
4 | numpy==1.18.4
5 | tensorboard==1.15.0
6 | tensorflow-gpu==1.15.4
7 | tensorflow-estimator==1.15.1
8 | tqdm==4.31.1
9 | requests==2.22.0
10 | ujson==2.0.3
--------------------------------------------------------------------------------
/requirements-tpu.txt:
--------------------------------------------------------------------------------
1 | pandas==0.24.2
2 | regex==2019.4.14
3 | h5py==2.10.0
4 | numpy==1.18.4
5 | tensorboard==1.15.0
6 | tensorflow==1.15.4
7 | tensorflow-estimator==1.15.1
8 | tqdm==4.31.1
9 | requests==2.22.0
10 | ujson==2.0.3
--------------------------------------------------------------------------------
/scripts/demo.py:
--------------------------------------------------------------------------------
1 | import sys
2 | import os
3 | import argparse
4 | import json
5 | import re
6 |
7 | import tensorflow.compat.v1 as tf
8 | import numpy as np
9 |
10 | from train.modeling import GroverModel, GroverConfig, sample
11 | from tokenization import tokenization
12 |
13 | ##### ignore tf deprecated warning temporarily
14 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
15 | tf.logging.set_verbosity(tf.logging.DEBUG)
16 | from tensorflow.python.util import deprecation
17 | deprecation._PRINT_DEPRECATION_WARNINGS = False
18 | try:
19 | from tensorflow.python.util import module_wrapper as deprecation
20 | except ImportError:
21 | from tensorflow.python.util import deprecation_wrapper as deprecation
22 | deprecation._PER_MODULE_WARNING_LIMIT = 0
23 | #####
24 |
25 | parser = argparse.ArgumentParser(description='Contextual generation (aka given some metadata we will generate articles')
26 | parser.add_argument(
27 | '-metadata_fn',
28 | dest='metadata_fn',
29 | type=str,
30 | help='Path to a JSONL containing metadata',
31 | )
32 | parser.add_argument(
33 | '-out_fn',
34 | dest='out_fn',
35 | type=str,
36 | help='Out jsonl, which will contain the completed jsons',
37 | )
38 | parser.add_argument(
39 | '-input',
40 | dest='input',
41 | type=str,
42 | help='Text to complete',
43 | )
44 | parser.add_argument(
45 | '-config_fn',
46 | dest='config_fn',
47 | default='configs/mega.json',
48 | type=str,
49 | help='Configuration JSON for the model',
50 | )
51 | parser.add_argument(
52 | '-ckpt_fn',
53 | dest='ckpt_fn',
54 | default='../models/mega/model.ckpt',
55 | type=str,
56 | help='checkpoint file for the model',
57 | )
58 | parser.add_argument(
59 | '-target',
60 | dest='target',
61 | default='article',
62 | type=str,
63 | help='What to generate for each item in metadata_fn. can be article (body), title, etc.',
64 | )
65 | parser.add_argument(
66 | '-batch_size',
67 | dest='batch_size',
68 | default=1,
69 | type=int,
70 | help='How many things to generate per context. will split into chunks if need be',
71 | )
72 | parser.add_argument(
73 | '-num_folds',
74 | dest='num_folds',
75 | default=1,
76 | type=int,
77 | help='Number of folds. useful if we want to split up a big file into multiple jobs.',
78 | )
79 | parser.add_argument(
80 | '-fold',
81 | dest='fold',
82 | default=0,
83 | type=int,
84 | help='which fold we are on. useful if we want to split up a big file into multiple jobs.'
85 | )
86 | parser.add_argument(
87 | '-max_batch_size',
88 | dest='max_batch_size',
89 | default=None,
90 | type=int,
91 | help='max batch size. You can leave this out and we will infer one based on the number of hidden layers',
92 | )
93 | parser.add_argument(
94 | '-top_p',
95 | dest='top_p',
96 | default=0.95,
97 | type=float,
98 | help='p to use for top p sampling. if this isn\'t none, use this for everthing'
99 | )
100 | parser.add_argument(
101 | '-min_len',
102 | dest='min_len',
103 | default=1024,
104 | type=int,
105 | help='min length of sample',
106 | )
107 | parser.add_argument(
108 | '-eos_token',
109 | dest='eos_token',
110 | default=102,
111 | type=int,
112 | help='eos token id',
113 | )
114 | parser.add_argument(
115 | '-samples',
116 | dest='samples',
117 | default=5,
118 | type=int,
119 | help='num_samples',
120 | )
121 |
122 | def extract_generated_target(output_tokens, tokenizer):
123 | """
124 | Given some tokens that were generated, extract the target
125 | :param output_tokens: [num_tokens] thing that was generated
126 | :param encoder: how they were encoded
127 | :param target: the piece of metadata we wanted to generate!
128 | :return:
129 | """
130 | # Filter out first instance of start token
131 | assert output_tokens.ndim == 1
132 |
133 | start_ind = 0
134 | end_ind = output_tokens.shape[0]
135 |
136 | return {
137 | 'extraction': tokenization.printable_text(''.join(tokenizer.convert_ids_to_tokens(output_tokens))),
138 | 'start_ind': start_ind,
139 | 'end_ind': end_ind,
140 | }
141 |
142 | args = parser.parse_args()
143 | proj_root_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
144 | vocab_file_path = os.path.join(proj_root_path, "tokenization/clue-vocab.txt")
145 | tokenizer = tokenization.FullTokenizer(vocab_file=vocab_file_path , do_lower_case=True)
146 | news_config = GroverConfig.from_json_file(args.config_fn)
147 |
148 | # We might have to split the batch into multiple chunks if the batch size is too large
149 | default_mbs = {12: 32, 24: 16, 48: 3}
150 | max_batch_size = args.max_batch_size if args.max_batch_size is not None else default_mbs[news_config.num_hidden_layers]
151 |
152 | # factorize args.batch_size = (num_chunks * batch_size_per_chunk) s.t. batch_size_per_chunk < max_batch_size
153 | num_chunks = int(np.ceil(args.batch_size / max_batch_size))
154 | batch_size_per_chunk = int(np.ceil(args.batch_size / num_chunks))
155 |
156 | # This controls the top p for each generation.
157 | top_p = np.ones((num_chunks, batch_size_per_chunk), dtype=np.float32) * args.top_p
158 |
159 | tf_config = tf.ConfigProto(allow_soft_placement=True)
160 |
161 | with tf.Session(config=tf_config, graph=tf.Graph()) as sess:
162 | initial_context = tf.placeholder(tf.int32, [batch_size_per_chunk, None])
163 | p_for_topp = tf.placeholder(tf.float32, [batch_size_per_chunk])
164 | eos_token = tf.placeholder(tf.int32, [])
165 | min_len = tf.placeholder(tf.int32, [])
166 | tokens, probs = sample(news_config=news_config, initial_context=initial_context,
167 | eos_token=eos_token, min_len=min_len, ignore_ids=None, p_for_topp=p_for_topp,
168 | do_topk=False)
169 |
170 | saver = tf.train.Saver()
171 | saver.restore(sess, args.ckpt_fn)
172 | print('🍺Model loaded. \nInput something please:⬇️')
173 | text = input()
174 | while text != "":
175 | for i in range(args.samples):
176 | print("Sample,", i + 1, " of ", args.samples)
177 | line = tokenization.convert_to_unicode(text)
178 | bert_tokens = tokenizer.tokenize(line)
179 | encoded = tokenizer.convert_tokens_to_ids(bert_tokens)
180 | context_formatted = []
181 | context_formatted.extend(encoded)
182 | # Format context end
183 |
184 | gens = []
185 | gens_raw = []
186 | gen_probs = []
187 |
188 | for chunk_i in range(num_chunks):
189 | tokens_out, probs_out = sess.run([tokens, probs],
190 | feed_dict={initial_context: [context_formatted] * batch_size_per_chunk,
191 | eos_token: args.eos_token, min_len: args.min_len,
192 | p_for_topp: top_p[chunk_i]})
193 |
194 | for t_i, p_i in zip(tokens_out, probs_out):
195 | extraction = extract_generated_target(output_tokens=t_i, tokenizer=tokenizer)
196 | gens.append(extraction['extraction'])
197 |
198 | l = re.findall('.{1,70}', gens[0].replace('[UNK]', '').replace('##', ''))
199 | print("\n".join(l))
200 | print('Next try:⬇️')
201 | text = input()
202 |
--------------------------------------------------------------------------------
/scripts/down_gdrive_file.py:
--------------------------------------------------------------------------------
1 | import argparse
2 |
3 | from google.colab import auth
4 | from googleapiclient.discovery import build
5 | from apiclient.http import MediaIoBaseDownload
6 | from tqdm import tqdm
7 |
8 | parser = argparse.ArgumentParser(description='Simple file download script for Google Drive')
9 | parser.add_argument(
10 | '-file_id',
11 | dest='file_id',
12 | type=str,
13 | help='File id in Google Drive URL',
14 | )
15 | parser.add_argument(
16 | '-file_path',
17 | dest='file_path',
18 | type=str,
19 | help='Output file path',
20 | )
21 |
22 | args = parser.parse_args()
23 |
24 | auth.authenticate_user()
25 | drive_service = build('drive', 'v3')
26 |
27 | # file_id, file_ext = ('1n_5-tgPpQ1gqbyLPbP1PwiFi2eo7SWw_', '.data-00000-of-00001')
28 | # filename = '%s/model.ckpt-%d%s' % (model_dir, 100000, file_ext)
29 | req = drive_service.files().get_media(fileId=args.file_id)
30 | with open(args.file_path, 'wb') as f:
31 | downloader = MediaIoBaseDownload(f, req, chunksize=100*1024*1024)
32 | done = False
33 | pbar = tqdm(total=100, desc='%s' % args.file_path)
34 | progress = 0
35 | while done is False:
36 | status, done = downloader.next_chunk()
37 | new_progress = int(status.progress() * 100)
38 | pbar.update(new_progress - progress)
39 | progress = new_progress
40 | pbar.close()
41 |
--------------------------------------------------------------------------------
/tokenization/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/imcaspar/gpt2-ml/f6286b16cbbee6dddbe1ba557fffb47eaf998cd1/tokenization/__init__.py
--------------------------------------------------------------------------------
/tokenization/tokenization.py:
--------------------------------------------------------------------------------
1 | # coding=utf-8
2 | # Copyright 2018 The Google AI Language Team Authors.
3 | #
4 | # Licensed under the Apache License, Version 2.0 (the "License");
5 | # you may not use this file except in compliance with the License.
6 | # You may obtain a copy of the License at
7 | #
8 | # http://www.apache.org/licenses/LICENSE-2.0
9 | #
10 | # Unless required by applicable law or agreed to in writing, software
11 | # distributed under the License is distributed on an "AS IS" BASIS,
12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 | # See the License for the specific language governing permissions and
14 | # limitations under the License.
15 | """Tokenization classes."""
16 |
17 | from __future__ import absolute_import
18 | from __future__ import division
19 | from __future__ import print_function
20 |
21 | import collections
22 | import re
23 | import unicodedata
24 | import six
25 | import tensorflow.compat.v1 as tf
26 |
27 |
28 | def validate_case_matches_checkpoint(do_lower_case, init_checkpoint):
29 | """Checks whether the casing config is consistent with the checkpoint name."""
30 |
31 | # The casing has to be passed in by the user and there is no explicit check
32 | # as to whether it matches the checkpoint. The casing information probably
33 | # should have been stored in the bert_config.json file, but it's not, so
34 | # we have to heuristically detect it to validate.
35 |
36 | if not init_checkpoint:
37 | return
38 |
39 | m = re.match("^.*?([A-Za-z0-9_-]+)/bert_model.ckpt", init_checkpoint)
40 | if m is None:
41 | return
42 |
43 | model_name = m.group(1)
44 |
45 | lower_models = [
46 | "uncased_L-24_H-1024_A-16", "uncased_L-12_H-768_A-12",
47 | "multilingual_L-12_H-768_A-12", "chinese_L-12_H-768_A-12"
48 | ]
49 |
50 | cased_models = [
51 | "cased_L-12_H-768_A-12", "cased_L-24_H-1024_A-16",
52 | "multi_cased_L-12_H-768_A-12"
53 | ]
54 |
55 | is_bad_config = False
56 | if model_name in lower_models and not do_lower_case:
57 | is_bad_config = True
58 | actual_flag = "False"
59 | case_name = "lowercased"
60 | opposite_flag = "True"
61 |
62 | if model_name in cased_models and do_lower_case:
63 | is_bad_config = True
64 | actual_flag = "True"
65 | case_name = "cased"
66 | opposite_flag = "False"
67 |
68 | if is_bad_config:
69 | raise ValueError(
70 | "You passed in `--do_lower_case=%s` with `--init_checkpoint=%s`. "
71 | "However, `%s` seems to be a %s model, so you "
72 | "should pass in `--do_lower_case=%s` so that the fine-tuning matches "
73 | "how the model was pre-training. If this error is wrong, please "
74 | "just comment out this check." % (actual_flag, init_checkpoint,
75 | model_name, case_name, opposite_flag))
76 |
77 |
78 | def convert_to_unicode(text):
79 | """Converts `text` to Unicode (if it's not already), assuming utf-8 input."""
80 | if six.PY3:
81 | if isinstance(text, str):
82 | return text
83 | elif isinstance(text, bytes):
84 | return text.decode("utf-8", "ignore")
85 | else:
86 | raise ValueError("Unsupported string type: %s" % (type(text)))
87 | elif six.PY2:
88 | if isinstance(text, str):
89 | return text.decode("utf-8", "ignore")
90 | elif isinstance(text, unicode):
91 | return text
92 | else:
93 | raise ValueError("Unsupported string type: %s" % (type(text)))
94 | else:
95 | raise ValueError("Not running on Python2 or Python 3?")
96 |
97 |
98 | def printable_text(text):
99 | """Returns text encoded in a way suitable for print or `tf.logging`."""
100 |
101 | # These functions want `str` for both Python2 and Python3, but in one case
102 | # it's a Unicode string and in the other it's a byte string.
103 | if six.PY3:
104 | if isinstance(text, str):
105 | return text
106 | elif isinstance(text, bytes):
107 | return text.decode("utf-8", "ignore")
108 | else:
109 | raise ValueError("Unsupported string type: %s" % (type(text)))
110 | elif six.PY2:
111 | if isinstance(text, str):
112 | return text
113 | elif isinstance(text, unicode):
114 | return text.encode("utf-8")
115 | else:
116 | raise ValueError("Unsupported string type: %s" % (type(text)))
117 | else:
118 | raise ValueError("Not running on Python2 or Python 3?")
119 |
120 |
121 | def load_vocab(vocab_file):
122 | """Loads a vocabulary file into a dictionary."""
123 | vocab = collections.OrderedDict()
124 | index = 0
125 | with tf.gfile.GFile(vocab_file, "r") as reader:
126 | while True:
127 | token = convert_to_unicode(reader.readline())
128 | if not token:
129 | break
130 | token = token.strip()
131 | vocab[token] = index
132 | index += 1
133 | return vocab
134 |
135 |
136 | def convert_by_vocab(vocab, items):
137 | """Converts a sequence of [tokens|ids] using the vocab."""
138 | output = []
139 | for item in items:
140 | output.append(vocab[item])
141 | return output
142 |
143 |
144 | def convert_tokens_to_ids(vocab, tokens):
145 | return convert_by_vocab(vocab, tokens)
146 |
147 |
148 | def convert_ids_to_tokens(inv_vocab, ids):
149 | return convert_by_vocab(inv_vocab, ids)
150 |
151 |
152 | def whitespace_tokenize(text):
153 | """Runs basic whitespace cleaning and splitting on a piece of text."""
154 | text = text.strip()
155 | if not text:
156 | return []
157 | tokens = text.split()
158 | return tokens
159 |
160 |
161 | class FullTokenizer(object):
162 | """Runs end-to-end tokenziation."""
163 |
164 | def __init__(self, vocab_file, do_lower_case=True):
165 | self.vocab = load_vocab(vocab_file)
166 | self.inv_vocab = {v: k for k, v in self.vocab.items()}
167 | self.basic_tokenizer = BasicTokenizer(do_lower_case=do_lower_case)
168 | self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.vocab)
169 |
170 | def tokenize(self, text):
171 | split_tokens = []
172 | for token in self.basic_tokenizer.tokenize(text):
173 | for sub_token in self.wordpiece_tokenizer.tokenize(token):
174 | split_tokens.append(sub_token)
175 |
176 | return split_tokens
177 |
178 | def convert_tokens_to_ids(self, tokens):
179 | return convert_by_vocab(self.vocab, tokens)
180 |
181 | def convert_ids_to_tokens(self, ids):
182 | return convert_by_vocab(self.inv_vocab, ids)
183 |
184 |
185 | class BasicTokenizer(object):
186 | """Runs basic tokenization (punctuation splitting, lower casing, etc.)."""
187 |
188 | def __init__(self, do_lower_case=True):
189 | """Constructs a BasicTokenizer.
190 |
191 | Args:
192 | do_lower_case: Whether to lower case the input.
193 | """
194 | self.do_lower_case = do_lower_case
195 |
196 | def tokenize(self, text):
197 | """Tokenizes a piece of text."""
198 | text = convert_to_unicode(text)
199 | text = self._clean_text(text)
200 |
201 | # This was added on November 1st, 2018 for the multilingual and Chinese
202 | # models. This is also applied to the English models now, but it doesn't
203 | # matter since the English models were not trained on any Chinese data
204 | # and generally don't have any Chinese data in them (there are Chinese
205 | # characters in the vocabulary because Wikipedia does have some Chinese
206 | # words in the English Wikipedia.).
207 | text = self._tokenize_chinese_chars(text)
208 |
209 | orig_tokens = whitespace_tokenize(text)
210 | split_tokens = []
211 | for token in orig_tokens:
212 | if self.do_lower_case:
213 | token = token.lower()
214 | token = self._run_strip_accents(token)
215 | split_tokens.extend(self._run_split_on_punc(token))
216 |
217 | output_tokens = whitespace_tokenize(" ".join(split_tokens))
218 | return output_tokens
219 |
220 | def _run_strip_accents(self, text):
221 | """Strips accents from a piece of text."""
222 | text = unicodedata.normalize("NFD", text)
223 | output = []
224 | for char in text:
225 | cat = unicodedata.category(char)
226 | if cat == "Mn":
227 | continue
228 | output.append(char)
229 | return "".join(output)
230 |
231 | def _run_split_on_punc(self, text):
232 | """Splits punctuation on a piece of text."""
233 | chars = list(text)
234 | i = 0
235 | start_new_word = True
236 | output = []
237 | while i < len(chars):
238 | char = chars[i]
239 | if _is_punctuation(char):
240 | output.append([char])
241 | start_new_word = True
242 | else:
243 | if start_new_word:
244 | output.append([])
245 | start_new_word = False
246 | output[-1].append(char)
247 | i += 1
248 |
249 | return ["".join(x) for x in output]
250 |
251 | def _tokenize_chinese_chars(self, text):
252 | """Adds whitespace around any CJK character."""
253 | output = []
254 | for char in text:
255 | cp = ord(char)
256 | if self._is_chinese_char(cp):
257 | output.append(" ")
258 | output.append(char)
259 | output.append(" ")
260 | else:
261 | output.append(char)
262 | return "".join(output)
263 |
264 | def _is_chinese_char(self, cp):
265 | """Checks whether CP is the codepoint of a CJK character."""
266 | # This defines a "chinese character" as anything in the CJK Unicode block:
267 | # https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
268 | #
269 | # Note that the CJK Unicode block is NOT all Japanese and Korean characters,
270 | # despite its name. The modern Korean Hangul alphabet is a different block,
271 | # as is Japanese Hiragana and Katakana. Those alphabets are used to write
272 | # space-separated words, so they are not treated specially and handled
273 | # like the all of the other languages.
274 | if ((cp >= 0x4E00 and cp <= 0x9FFF) or #
275 | (cp >= 0x3400 and cp <= 0x4DBF) or #
276 | (cp >= 0x20000 and cp <= 0x2A6DF) or #
277 | (cp >= 0x2A700 and cp <= 0x2B73F) or #
278 | (cp >= 0x2B740 and cp <= 0x2B81F) or #
279 | (cp >= 0x2B820 and cp <= 0x2CEAF) or
280 | (cp >= 0xF900 and cp <= 0xFAFF) or #
281 | (cp >= 0x2F800 and cp <= 0x2FA1F)): #
282 | return True
283 |
284 | return False
285 |
286 | def _clean_text(self, text):
287 | """Performs invalid character removal and whitespace cleanup on text."""
288 | output = []
289 | for char in text:
290 | cp = ord(char)
291 | if cp == 0 or cp == 0xfffd or _is_control(char):
292 | continue
293 | if _is_whitespace(char):
294 | output.append(" ")
295 | else:
296 | output.append(char)
297 | return "".join(output)
298 |
299 |
300 | class WordpieceTokenizer(object):
301 | """Runs WordPiece tokenziation."""
302 |
303 | def __init__(self, vocab, unk_token="[UNK]", max_input_chars_per_word=200):
304 | self.vocab = vocab
305 | self.unk_token = unk_token
306 | self.max_input_chars_per_word = max_input_chars_per_word
307 |
308 | def tokenize(self, text):
309 | """Tokenizes a piece of text into its word pieces.
310 |
311 | This uses a greedy longest-match-first algorithm to perform tokenization
312 | using the given vocabulary.
313 |
314 | For example:
315 | input = "unaffable"
316 | output = ["un", "##aff", "##able"]
317 |
318 | Args:
319 | text: A single token or whitespace separated tokens. This should have
320 | already been passed through `BasicTokenizer.
321 |
322 | Returns:
323 | A list of wordpiece tokens.
324 | """
325 |
326 | text = convert_to_unicode(text)
327 |
328 | output_tokens = []
329 | for token in whitespace_tokenize(text):
330 | chars = list(token)
331 | if len(chars) > self.max_input_chars_per_word:
332 | output_tokens.append(self.unk_token)
333 | continue
334 |
335 | is_bad = False
336 | start = 0
337 | sub_tokens = []
338 | while start < len(chars):
339 | end = len(chars)
340 | cur_substr = None
341 | while start < end:
342 | substr = "".join(chars[start:end])
343 | if start > 0:
344 | substr = "##" + substr
345 | if substr in self.vocab:
346 | cur_substr = substr
347 | break
348 | end -= 1
349 | if cur_substr is None:
350 | is_bad = True
351 | break
352 | sub_tokens.append(cur_substr)
353 | start = end
354 |
355 | if is_bad:
356 | output_tokens.append(self.unk_token)
357 | else:
358 | output_tokens.extend(sub_tokens)
359 | return output_tokens
360 |
361 |
362 | def _is_whitespace(char):
363 | """Checks whether `chars` is a whitespace character."""
364 | # \t, \n, and \r are technically contorl characters but we treat them
365 | # as whitespace since they are generally considered as such.
366 | if char == " " or char == "\t" or char == "\n" or char == "\r":
367 | return True
368 | cat = unicodedata.category(char)
369 | if cat == "Zs":
370 | return True
371 | return False
372 |
373 |
374 | def _is_control(char):
375 | """Checks whether `chars` is a control character."""
376 | # These are technically control characters but we count them as whitespace
377 | # characters.
378 | if char == "\t" or char == "\n" or char == "\r":
379 | return False
380 | cat = unicodedata.category(char)
381 | if cat in ("Cc", "Cf"):
382 | return True
383 | return False
384 |
385 |
386 | def _is_punctuation(char):
387 | """Checks whether `chars` is a punctuation character."""
388 | cp = ord(char)
389 | # We treat all non-letter/number ASCII as punctuation.
390 | # Characters such as "^", "$", and "`" are not in the Unicode
391 | # Punctuation class but we treat them as punctuation anyways, for
392 | # consistency.
393 | if ((cp >= 33 and cp <= 47) or (cp >= 58 and cp <= 64) or
394 | (cp >= 91 and cp <= 96) or (cp >= 123 and cp <= 126)):
395 | return True
396 | cat = unicodedata.category(char)
397 | if cat.startswith("P"):
398 | return True
399 | return False
400 |
--------------------------------------------------------------------------------
/train/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/imcaspar/gpt2-ml/f6286b16cbbee6dddbe1ba557fffb47eaf998cd1/train/__init__.py
--------------------------------------------------------------------------------
/train/dataloader.py:
--------------------------------------------------------------------------------
1 | # Original work Copyright 2018 The Google AI Language Team Authors.
2 | # Modified work Copyright 2019 Rowan Zellers
3 | #
4 | # Licensed under the Apache License, Version 2.0 (the "License");
5 | # you may not use this file except in compliance with the License.
6 | # You may obtain a copy of the License at
7 | #
8 | # http://www.apache.org/licenses/LICENSE-2.0
9 | #
10 | # Unless required by applicable law or agreed to in writing, software
11 | # distributed under the License is distributed on an "AS IS" BASIS,
12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 | # See the License for the specific language governing permissions and
14 | # limitations under the License.
15 |
16 | import collections
17 | import tensorflow.compat.v1 as tf
18 |
19 |
20 | def _decode_record(record, name_to_features):
21 | """Decodes a record to a TensorFlow example."""
22 | example = tf.parse_single_example(record, name_to_features)
23 |
24 | # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
25 | # So cast all int64 to int32.
26 | for name in list(example.keys()):
27 | t = example[name]
28 | if t.dtype == tf.int64:
29 | t = tf.cast(t, tf.int32)
30 | example[name] = t
31 | return example
32 |
33 |
34 | def input_fn_builder(input_files,
35 | seq_length,
36 | is_training,
37 | num_cpu_threads=4,
38 | evaluate_for_fixed_number_of_steps=True):
39 | """Creates an `input_fn` closure to be passed to TPUEstimator."""
40 |
41 | def input_fn(params):
42 | """The actual input function."""
43 | batch_size = params["batch_size"]
44 | name_to_features = {
45 | "input_ids": tf.FixedLenFeature([seq_length + 1], tf.int64),
46 | }
47 |
48 | # For training, we want a lot of parallel reading and shuffling.
49 | # For eval, we want no shuffling and parallel reading doesn't matter.
50 | if is_training:
51 | d = tf.data.Dataset.from_tensor_slices(tf.constant(input_files))
52 | d = d.repeat()
53 | d = d.shuffle(buffer_size=len(input_files))
54 |
55 | # `cycle_length` is the number of parallel files that get read.
56 | cycle_length = min(num_cpu_threads, len(input_files))
57 |
58 | # `sloppy` mode means that the interleaving is not exact. This adds
59 | # even more randomness to the training pipeline.
60 | d = d.apply(
61 | tf.data.experimental.parallel_interleave(
62 | tf.data.TFRecordDataset,
63 | sloppy=is_training,
64 | cycle_length=cycle_length))
65 | d = d.shuffle(buffer_size=100)
66 | else:
67 | d = tf.data.TFRecordDataset(input_files)
68 | # If we evaluate for a fixed number of steps we don't want to encounter
69 | # out-of-range exceptions.
70 | if evaluate_for_fixed_number_of_steps:
71 | d = d.repeat()
72 |
73 | # We must `drop_remainder` on training because the TPU requires fixed
74 | # size dimensions. For eval, we assume we are evaluating on the CPU or GPU
75 | # and we *don't* want to drop the remainder, otherwise we wont cover
76 | # every sample.
77 | d = d.apply(
78 | tf.data.experimental.map_and_batch(
79 | lambda record: _decode_record(record, name_to_features),
80 | batch_size=batch_size,
81 | num_parallel_batches=num_cpu_threads,
82 | drop_remainder=True))
83 | return d
84 |
85 | return input_fn
86 |
87 |
88 | # ~~~~~~~~~~~~~~ This is for classification / AF ~~~~~~~~~~~~~~~~~~
89 | def classification_convert_examples_to_features(
90 | examples, max_seq_length, batch_size, encoder, output_file, labels, pad_extra_examples=False,
91 | chop_from_front_if_needed=True):
92 | """Convert a set of `InputExample`s to a TFRecord file."""
93 |
94 | writer = tf.python_io.TFRecordWriter(output_file)
95 |
96 | label_map = {label: i for i, label in enumerate(labels)}
97 |
98 | for (ex_index, example) in enumerate(examples):
99 | if ex_index % 10000 == 0:
100 | tf.logging.info("Writing example %d of %d" % (ex_index, len(examples)))
101 |
102 | # begin_summary is our [CLS] token
103 | tokens = example['ids'] + [encoder.begin_summary]
104 |
105 | if len(tokens) > max_seq_length:
106 | if chop_from_front_if_needed:
107 | tokens = tokens[-max_seq_length:]
108 | else:
109 | tokens = example['ids'][:(max_seq_length-1)] + [encoder.begin_summary]
110 | elif len(tokens) < max_seq_length:
111 | tokens.extend([encoder.padding] * (max_seq_length - len(tokens)))
112 |
113 | features = collections.OrderedDict()
114 | features['input_ids'] = tf.train.Feature(int64_list=tf.train.Int64List(value=tokens))
115 | features['label_ids'] = tf.train.Feature(int64_list=tf.train.Int64List(value=[label_map[example['label']]]))
116 | features['is_real_example'] = tf.train.Feature(int64_list=tf.train.Int64List(value=[1]))
117 | tf_example = tf.train.Example(features=tf.train.Features(feature=features))
118 | writer.write(tf_example.SerializeToString())
119 |
120 | if pad_extra_examples:
121 | for x in range(len(examples) % batch_size):
122 | features = collections.OrderedDict()
123 | features['input_ids'] = tf.train.Feature(int64_list=tf.train.Int64List(value=[0]*max_seq_length))
124 | features['label_ids'] = tf.train.Feature(int64_list=tf.train.Int64List(value=[0]))
125 | features['is_real_example'] = tf.train.Feature(int64_list=tf.train.Int64List(value=[0]))
126 | tf_example = tf.train.Example(features=tf.train.Features(feature=features))
127 | writer.write(tf_example.SerializeToString())
128 | writer.close()
129 |
130 |
131 | def classification_input_fn_builder(input_file, seq_length, is_training,
132 | drop_remainder,
133 | buffer_size=100):
134 | """Creates an `input_fn` closure to be passed to TPUEstimator."""
135 |
136 | name_to_features = {
137 | "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
138 | "label_ids": tf.FixedLenFeature([], tf.int64),
139 | "is_real_example": tf.FixedLenFeature([], tf.int64),
140 | }
141 |
142 | def input_fn(params):
143 | """The actual input function."""
144 | batch_size = params["batch_size"]
145 |
146 | # For training, we want a lot of parallel reading and shuffling.
147 | # For eval, we want no shuffling and parallel reading doesn't matter.
148 | d = tf.data.TFRecordDataset(input_file)
149 | if is_training:
150 | d = d.repeat()
151 | d = d.shuffle(buffer_size=buffer_size)
152 |
153 | d = d.apply(
154 | tf.data.experimental.map_and_batch(
155 | lambda record: _decode_record(record, name_to_features),
156 | batch_size=batch_size,
157 | drop_remainder=drop_remainder))
158 |
159 | return d
160 |
161 | return input_fn
162 |
--------------------------------------------------------------------------------
/train/modeling.py:
--------------------------------------------------------------------------------
1 | # Original work Copyright 2018 The Google AI Language Team Authors.
2 | # Modified work Copyright 2019 Rowan Zellers
3 | #
4 | # Licensed under the Apache License, Version 2.0 (the "License");
5 | # you may not use this file except in compliance with the License.
6 | # You may obtain a copy of the License at
7 | #
8 | # http://www.apache.org/licenses/LICENSE-2.0
9 | #
10 | # Unless required by applicable law or agreed to in writing, software
11 | # distributed under the License is distributed on an "AS IS" BASIS,
12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 | # See the License for the specific language governing permissions and
14 | # limitations under the License.
15 |
16 | import copy
17 | import json
18 | import math
19 |
20 | import six
21 | import tensorflow.compat.v1 as tf
22 |
23 | from train import optimization_adafactor
24 | from train.utils import get_assignment_map_from_checkpoint, get_shape_list, get_attention_mask, gelu, layer_norm, dropout, \
25 | construct_scalar_host_call
26 |
27 | class GroverConfig(object):
28 | """Configuration for `GroverModel`"""
29 |
30 | def __init__(self,
31 | vocab_size,
32 | hidden_size=768,
33 | num_hidden_layers=12,
34 | num_attention_heads=12,
35 | intermediate_size=3072,
36 | hidden_act="gelu",
37 | hidden_dropout_prob=0.1,
38 | attention_probs_dropout_prob=0.1,
39 | max_position_embeddings=512,
40 | initializer_range=0.02):
41 | """Constructs NewsConfig.
42 |
43 | Args:
44 | vocab_size: Vocabulary size of `inputs_ids` in `GroverModel`.
45 | hidden_size: Size of the layers
46 | num_hidden_layers: Number of hidden layers in the Transformer encoder.
47 | num_attention_heads: Number of attention heads for each attention layer in
48 | the Transformer encoder.
49 | intermediate_size: The size of the "intermediate" (i.e., feed-forward)
50 | layer in the Transformer encoder.
51 | hidden_act: The non-linear activation function (function or string) in the
52 | encoder and pooler.
53 | hidden_dropout_prob: The dropout probability for all fully connected
54 | layers in the embeddings, encoder, and pooler.
55 | attention_probs_dropout_prob: The dropout ratio for the attention
56 | probabilities.
57 | max_position_embeddings: The maximum sequence length that this model might
58 | ever be used with. Typically set this to something large just in case
59 | (e.g., 512 or 1024 or 2048).
60 | initializer_range: The stdev of the truncated_normal_initializer for
61 | initializing all weight matrices.
62 | """
63 | self.vocab_size = vocab_size
64 | self.hidden_size = hidden_size
65 | self.num_hidden_layers = num_hidden_layers
66 | self.num_attention_heads = num_attention_heads
67 | self.hidden_act = hidden_act
68 | self.intermediate_size = intermediate_size
69 | self.hidden_dropout_prob = hidden_dropout_prob
70 | self.attention_probs_dropout_prob = attention_probs_dropout_prob
71 | self.max_position_embeddings = max_position_embeddings
72 | self.initializer_range = initializer_range
73 | self.pad_token_id = 0
74 |
75 | @classmethod
76 | def from_dict(cls, json_object):
77 | """Constructs a `NewsConfig` from a Python dictionary of parameters."""
78 | config = GroverConfig(vocab_size=None)
79 | for (key, value) in six.iteritems(json_object):
80 | config.__dict__[key] = value
81 | return config
82 |
83 | @classmethod
84 | def from_json_file(cls, json_file):
85 | """Constructs a `NewsConfig` from a json file of parameters."""
86 | with tf.gfile.GFile(json_file, "r") as reader:
87 | text = reader.read()
88 | return cls.from_dict(json.loads(text))
89 |
90 | def to_dict(self):
91 | """Serializes this instance to a Python dictionary."""
92 | output = copy.deepcopy(self.__dict__)
93 | return output
94 |
95 | def to_json_string(self):
96 | """Serializes this instance to a JSON string."""
97 | return json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"
98 |
99 |
100 | def mask_attention_for_ltr(attention_scores, attention_mask):
101 | """
102 | Mask attention so that we're only predicting going forward
103 | :param attention_scores: [batch, heads, dst_sequence, src_sequence], where information flows from src to dst.
104 | :param attention_mask [query_length, key_length]
105 | :return: masked attention
106 | """
107 | # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
108 | # masked positions, this operation will create a tensor which is 0.0 for
109 | # positions we want to attend and -10000.0 for masked positions.
110 | mask = attention_mask[None, None]
111 | return attention_scores * mask - tf.cast(1e10, attention_scores.dtype) * (1 - mask)
112 |
113 |
114 | def create_initializer(initializer_range=0.02):
115 | """Creates a `truncated_normal_initializer` with the given range."""
116 | return tf.truncated_normal_initializer(stddev=initializer_range)
117 |
118 |
119 | def _attention_projection_and_transpose(x_flat, batch_size, seq_length, num_attention_heads, size_per_head,
120 | name, initializer_range=0.02):
121 | """
122 | :param x_flat: [batch_size*seq_length, width]
123 | :return: A fixed up tensor of size [batch_size, num_attention_heads, seq_length, size_per_head]
124 | """
125 | batch_size_seq_length, dim = get_shape_list(x_flat, expected_rank=2)
126 |
127 | if dim != size_per_head * num_attention_heads:
128 | raise ValueError("passed in a tensor of shape {} when size_per_head={} and num_attention_heads={}".format(
129 | (batch_size_seq_length, dim), size_per_head, num_attention_heads
130 | ))
131 |
132 | projected = tf.layers.dense(
133 | x_flat,
134 | num_attention_heads * size_per_head,
135 | name=name,
136 | kernel_initializer=create_initializer(initializer_range))
137 |
138 | projected = tf.reshape(
139 | projected, [batch_size, seq_length, num_attention_heads, size_per_head])
140 | output_tensor = tf.transpose(projected, [0, 2, 1, 3])
141 | return output_tensor
142 |
143 |
144 | def attention_layer(x_flat, attention_mask, batch_size, seq_length, size_per_head=512, num_attention_heads=1, *,
145 | cache=None,
146 | initializer_range=0.02, hidden_dropout_prob=0.1,
147 | attention_probs_dropout_prob=0.1, do_cache=False):
148 | """
149 |
150 | :param x_flat: Tensor input, should be [batch_size*seq_length, dim]
151 | :param attention_mask: Attention mask to use of size [seq_length, seq_length+cached_length]
152 | :param size_per_head: dim = size_per_head * num_attention_heads
153 | :param num_attention_heads: dim = size_per_head * num_attention_heads
154 | :param cache: Optionally some past (cached) things of size
155 | [batch, 2, heads, sequence, features], where 2 is [k, v]
156 | :param do_cache: True if we should return cache
157 | :return: A new tensor of shape [batch_size, seq_length, dim]
158 | as well as a new cache "cached_keys_and_values" that will be of size
159 | [batch_size, 2, num_attention_heads, seq_length, dim]
160 | """
161 | batch_size_seq_length, dim = get_shape_list(x_flat, expected_rank=2)
162 |
163 | if dim != size_per_head * num_attention_heads:
164 | raise ValueError("passed in a tensor of shape {} when size_per_head={} and num_attention_heads={}".format(
165 | (batch_size_seq_length, dim), size_per_head, num_attention_heads
166 | ))
167 |
168 | query = _attention_projection_and_transpose(x_flat, batch_size=batch_size, seq_length=seq_length,
169 | num_attention_heads=num_attention_heads, size_per_head=size_per_head,
170 | name='query_layer',
171 | initializer_range=initializer_range)
172 | key = _attention_projection_and_transpose(x_flat, batch_size=batch_size, seq_length=seq_length,
173 | num_attention_heads=num_attention_heads, size_per_head=size_per_head,
174 | name='key_layer',
175 | initializer_range=initializer_range)
176 |
177 | value = _attention_projection_and_transpose(x_flat, batch_size=batch_size, seq_length=seq_length,
178 | num_attention_heads=num_attention_heads, size_per_head=size_per_head,
179 | name='value_layer',
180 | initializer_range=initializer_range)
181 |
182 | # Add to cache
183 | cached_keys_and_values = tf.stack([key, value], axis=1) if do_cache else None
184 |
185 | # Things that were relevant from the cache
186 | if cache is not None:
187 | pk, pv = tf.unstack(cache, axis=1)
188 | key = tf.concat([pk, key], axis=-2)
189 | value = tf.concat([pv, value], axis=-2)
190 |
191 | # Multiply [batch_size, num_attention_heads, seq_length, size_per_head] with
192 | # [batch_size, num_attention_heads, size_per_head, seq_length+cached_length] ->
193 | # [batch_size, num_attention_heads, seq_length, seq_length+cached_length]
194 | attention_scores = tf.matmul(query, key, transpose_b=True)
195 | attention_scores = tf.multiply(attention_scores,
196 | 1.0 / math.sqrt(float(size_per_head)))
197 | attention_scores = mask_attention_for_ltr(attention_scores, attention_mask)
198 | attention_probs = tf.nn.softmax(attention_scores)
199 |
200 | # This is actually dropping out entire tokens to attend to, which might
201 | # seem a bit unusual, but is taken from the original Transformer paper.
202 | # NOPENOPENOPENOPE
203 | # attention_probs = factoreddropout(attention_probs, attention_probs_dropout_prob)
204 |
205 | # Multiply [batch_size, num_attention_heads, seq_length, seq_length+cached_length] with
206 | # [batch_size, num_attention_heads, seq_length+cached_length, size_per_head] ->
207 | # [batch_size, num_attention_heads, seq_length, size_per_head] ->
208 | context_layer = tf.matmul(attention_probs, value)
209 |
210 | # `context_layer` = [batch_size, seq_length, num_attention_heads, size_per_head]
211 | context_layer = tf.transpose(context_layer, [0, 2, 1, 3])
212 | context_layer = tf.reshape(context_layer, [batch_size * seq_length, num_attention_heads * size_per_head])
213 |
214 | context_layer_projected = tf.layers.dense(
215 | context_layer,
216 | num_attention_heads * size_per_head,
217 | kernel_initializer=create_initializer(initializer_range),
218 | name='context_projection_layer'
219 | )
220 | context_layer_projected = dropout(context_layer_projected, hidden_dropout_prob)
221 |
222 | return context_layer_projected, cached_keys_and_values
223 |
224 |
225 | def residual_mlp_layer(x_flat, intermediate_size, initializer_range=0.02, hidden_dropout_prob=0.1):
226 | """
227 | :param x: The attention output. It should be [batch_size*seq_length, dim]
228 | :param intermediate_size: the hidden projection. By default this is the input_dim * 4.
229 |
230 | in the original GPT we would return layer_norm(x_norm + h1) rather than layer_norm(x + h1)
231 |
232 | :return:
233 | """
234 | batch_size_seq_length, hidden_size = get_shape_list(x_flat, expected_rank=2)
235 | x_norm = layer_norm(x_flat, name='mlp_ln0')
236 |
237 | intermediate_output = tf.layers.dense(
238 | x_norm,
239 | intermediate_size,
240 | activation=gelu,
241 | kernel_initializer=create_initializer(initializer_range),
242 | name='intermediate',
243 | )
244 |
245 | output_for_residual = tf.layers.dense(
246 | intermediate_output,
247 | hidden_size,
248 | name='output',
249 | kernel_initializer=create_initializer(initializer_range))
250 | output_for_residual = dropout(output_for_residual, hidden_dropout_prob)
251 |
252 | layer_output = layer_norm(x_flat + output_for_residual, name='mlp_ln1')
253 | return layer_output
254 |
255 |
256 | def embed(input_ids,
257 | vocab_size,
258 | embedding_size,
259 | position_offset=0,
260 | initializer_range=0.02,
261 | max_position_embeddings=512,
262 | use_one_hot_embeddings=True):
263 | """reur and position embeddings
264 | :param input_ids: int Tensor of shape [batch_size, seq_length].
265 | :param vocab_size: number of words in vocab
266 | :param embedding_size: dimensionality of the embedding
267 | :param position_offset: aka number of cached tokens.
268 | :param initializer_range: float. Range of the weight initialization.
269 | :param max_position_embeddings: int. Maximum sequence length.
270 | :param use_one_hot_embeddings: probably want this to be true
271 | :return: [batch_size, seq_length, embedding_size] embedded tensor
272 | """
273 | (batch_size, seq_length) = get_shape_list(input_ids, expected_rank=2)
274 |
275 | embedding_table = tf.get_variable(
276 | name='word_embed',
277 | shape=[vocab_size, embedding_size],
278 | initializer=create_initializer(initializer_range),
279 | )
280 |
281 | assert_op = tf.assert_less_equal(tf.reduce_max(input_ids), vocab_size - 1)
282 | with tf.control_dependencies([assert_op]):
283 | if use_one_hot_embeddings:
284 | flat_input_ids = tf.reshape(input_ids, [-1])
285 | one_hot_input_ids = tf.one_hot(flat_input_ids, depth=vocab_size)
286 | output_flat = tf.matmul(one_hot_input_ids, embedding_table)
287 | else:
288 | output_flat = tf.nn.embedding_lookup(embedding_table, input_ids)
289 |
290 | embedded_input = tf.reshape(output_flat, [batch_size, seq_length, embedding_size])
291 |
292 | assert_op = tf.assert_less_equal(seq_length, max_position_embeddings)
293 |
294 | with tf.control_dependencies([assert_op]):
295 | full_position_embeddings = tf.get_variable(
296 | name='pos_embed',
297 | shape=[max_position_embeddings, embedding_size],
298 | initializer=create_initializer(initializer_range),
299 | )
300 | # Since the position embedding table is a learned variable, we create it
301 | # using a (long) sequence length `max_position_embeddings`. The actual
302 | # sequence length might be shorter than this, for faster training of
303 | # tasks that do not have long sequences.
304 | #
305 | # So `full_position_embeddings` is effectively an embedding table
306 | # for position [0, 1, 2, ..., max_position_embeddings-1], and the current
307 | # sequence has positions [0, 1, 2, ... seq_length-1], so we can just
308 | # perform a slice.
309 | if position_offset == 0:
310 | embedded_input += tf.slice(full_position_embeddings, [0, 0], [seq_length, -1])[None]
311 | else:
312 | # Tensorflow is too stupid to allow slicing
313 | flat_pos_ids = (tf.range(seq_length, dtype=tf.int32) + position_offset)
314 | one_hot_pos_ids = tf.one_hot(flat_pos_ids, depth=max_position_embeddings)
315 |
316 | # [seq_length, full_position_embeddings], [full_position_embeddings, dim]
317 | seq_embeds = tf.matmul(one_hot_pos_ids, full_position_embeddings)
318 | embedded_input += seq_embeds[None]
319 |
320 | # embedded_input += tf.slice(full_position_embeddings[position_offset:], [0, 0], [seq_length, -1])[None]
321 |
322 | return layer_norm(embedded_input, name='embed_norm'), embedding_table
323 |
324 |
325 | def _top_p_sample(logits, ignore_ids=None, num_samples=1, p=0.9):
326 | """
327 | Does top-p sampling. if ignore_ids is on, then we will zero out those logits.
328 | :param logits: [batch_size, vocab_size] tensor
329 | :param ignore_ids: [vocab_size] one-hot representation of the indices we'd like to ignore and never predict,
330 | like padding maybe
331 | :param p: topp threshold to use, either a float or a [batch_size] vector
332 | :return: [batch_size, num_samples] samples
333 |
334 | # TODO FIGURE OUT HOW TO DO THIS ON TPUS. IT'S HELLA SLOW RIGHT NOW, DUE TO ARGSORT I THINK
335 | """
336 | with tf.variable_scope('top_p_sample'):
337 | batch_size, vocab_size = get_shape_list(logits, expected_rank=2)
338 |
339 | probs = tf.nn.softmax(logits if ignore_ids is None else logits - tf.cast(ignore_ids[None], tf.float32) * 1e10,
340 | axis=-1)
341 |
342 | if isinstance(p, float) and p > 0.999999:
343 | # Don't do top-p sampling in this case
344 | print("Top-p sampling DISABLED", flush=True)
345 | return {
346 | 'probs': probs,
347 | 'sample': tf.random.categorical(
348 | logits=logits if ignore_ids is None else logits - tf.cast(ignore_ids[None], tf.float32) * 1e10,
349 | num_samples=num_samples, dtype=tf.int32),
350 | }
351 |
352 | # [batch_size, vocab_perm]
353 | indices = tf.argsort(probs, direction='DESCENDING')
354 | cumulative_probabilities = tf.math.cumsum(tf.batch_gather(probs, indices), axis=-1, exclusive=False)
355 |
356 | # find the top pth index to cut off. careful we don't want to cutoff everything!
357 | # result will be [batch_size, vocab_perm]
358 | p_expanded = p if isinstance(p, float) else p[:, None]
359 | exclude_mask = tf.logical_not(
360 | tf.logical_or(cumulative_probabilities < p_expanded, tf.range(vocab_size)[None] < 1))
361 |
362 | # OPTION A - sample in the sorted space, then unsort.
363 | logits_to_use = tf.batch_gather(logits, indices) - tf.cast(exclude_mask, tf.float32) * 1e10
364 | sample_perm = tf.random.categorical(logits=logits_to_use, num_samples=num_samples)
365 | sample = tf.batch_gather(indices, sample_perm)
366 |
367 | # OPTION B - unsort first - Indices need to go back to 0 -> N-1 -- then sample
368 | # unperm_indices = tf.argsort(indices, direction='ASCENDING')
369 | # include_mask_unperm = tf.batch_gather(include_mask, unperm_indices)
370 | # logits_to_use = logits - (1 - tf.cast(include_mask_unperm, tf.float32)) * 1e10
371 | # sample = tf.random.categorical(logits=logits_to_use, num_samples=num_samples, dtype=tf.int32)
372 |
373 | return {
374 | 'probs': probs,
375 | 'sample': sample,
376 | }
377 |
378 |
379 | def _top_k_sample(logits, ignore_ids=None, num_samples=1, k=10):
380 | """
381 | Does top-k sampling. if ignore_ids is on, then we will zero out those logits.
382 | :param logits: [batch_size, vocab_size] tensor
383 | :param ignore_ids: [vocab_size] one-hot representation of the indices we'd like to ignore and never predict,
384 | like padding maybe
385 | :param p: topp threshold to use, either a float or a [batch_size] vector
386 | :return: [batch_size, num_samples] samples
387 |
388 | # TODO FIGURE OUT HOW TO DO THIS ON TPUS. IT'S HELLA SLOW RIGHT NOW, DUE TO ARGSORT I THINK
389 | """
390 | with tf.variable_scope('top_p_sample'):
391 | batch_size, vocab_size = get_shape_list(logits, expected_rank=2)
392 |
393 | probs = tf.nn.softmax(logits if ignore_ids is None else logits - tf.cast(ignore_ids[None], tf.float32) * 1e10,
394 | axis=-1)
395 | # [batch_size, vocab_perm]
396 | indices = tf.argsort(probs, direction='DESCENDING')
397 |
398 | # find the top pth index to cut off. careful we don't want to cutoff everything!
399 | # result will be [batch_size, vocab_perm]
400 | k_expanded = k if isinstance(k, int) else k[:, None]
401 | exclude_mask = tf.range(vocab_size)[None] >= k_expanded
402 |
403 | # OPTION A - sample in the sorted space, then unsort.
404 | logits_to_use = tf.batch_gather(logits, indices) - tf.cast(exclude_mask, tf.float32) * 1e10
405 | sample_perm = tf.random.categorical(logits=logits_to_use, num_samples=num_samples)
406 | sample = tf.batch_gather(indices, sample_perm)
407 |
408 | return {
409 | 'probs': probs,
410 | 'sample': sample,
411 | }
412 |
413 |
414 | class GroverModel(object):
415 | def __init__(self,
416 | config: GroverConfig,
417 | is_training,
418 | input_ids,
419 | cache=None,
420 | do_cache=False,
421 | pad_token_id=0,
422 | chop_off_last_token=True,
423 | scope=None,
424 | reuse=False):
425 | """
426 | :param config:
427 | :param is_training:
428 | :param input_ids: Tensor thats of size [batch_size, seq_length]
429 | :param cache: Optionally, a tensor to use that will contain cached information of the size
430 | [batch_size, num_layers, 2, num_heads, cache_length, features]
431 | :param do_cache: Whether to cache again.
432 | :param pad_token_id: Which token will be used for padding (probably 0.)
433 | :param chop_off_last_token: True if we will end up using this for TRAINING only. False if we want to generate.
434 | it means the last token in input_ids will not be processed by the model as input
435 | :param scope: scope to run this on
436 | """
437 | self.config = copy.deepcopy(config)
438 | self.is_training = is_training
439 | self.pad_token_id = pad_token_id
440 |
441 | if not is_training:
442 | self.config.hidden_dropout_prob = 0.0
443 | self.config.attention_probs_dropout_prob = 0.0
444 |
445 | if chop_off_last_token:
446 | self.target_ids = input_ids[:, 1:]
447 | self.input_ids = input_ids[:, :-1]
448 | else:
449 | self.input_ids = input_ids
450 | self.target_ids = tf.concat((input_ids[:, 1:],
451 | tf.constant(self.pad_token_id, dtype=self.input_ids.dtype,
452 | shape=[get_shape_list(self.input_ids, 2)[0], 1])), 1)
453 |
454 | self.batch_size, self.seq_length = get_shape_list(self.input_ids, 2)
455 |
456 | if cache is None:
457 | caches = [None] * config.num_hidden_layers
458 | self.cache_length = 0
459 | else:
460 | batch_size_, num_layers_, two_, num_heads_, self.cache_length, features_ = get_shape_list(
461 | cache, expected_rank=6)
462 | assert batch_size_ == self.batch_size
463 | assert num_layers_ == config.num_hidden_layers
464 | assert two_ == 2
465 | assert num_heads_ == config.num_attention_heads
466 | assert features_ == (config.hidden_size // config.num_attention_heads)
467 | caches = tf.unstack(cache, axis=1)
468 |
469 | with tf.variable_scope(scope, default_name='newslm', reuse=reuse):
470 | with tf.variable_scope("embeddings"):
471 | embeddings, self.embedding_table = embed(self.input_ids, config.vocab_size,
472 | config.hidden_size,
473 | position_offset=self.cache_length,
474 | initializer_range=config.initializer_range,
475 | max_position_embeddings=config.max_position_embeddings,
476 | use_one_hot_embeddings=True)
477 |
478 | mask = get_attention_mask(self.seq_length, self.seq_length + self.cache_length, dtype=embeddings.dtype)
479 |
480 | # We keep the representation as a 2D tensor to avoid re-shaping it back and
481 | # forth from a 3D tensor to a 2D tensor. Re-shapes are normally free on
482 | # the GPU/CPU but may not be free on the TPU, so we want to minimize them to
483 | # help the optimizer.
484 | hidden_state = tf.reshape(embeddings, [self.batch_size * self.seq_length, self.config.hidden_size])
485 | new_kvs = []
486 | for layer_idx, layer_cache in enumerate(caches):
487 | with tf.variable_scope('layer{:02d}'.format(layer_idx)):
488 | # [batch_size * seq_length, hidden_size]
489 | attention_output, new_kv = attention_layer(
490 | hidden_state,
491 | mask,
492 | batch_size=self.batch_size,
493 | seq_length=self.seq_length,
494 | size_per_head=config.hidden_size // config.num_attention_heads,
495 | num_attention_heads=config.num_attention_heads,
496 | initializer_range=config.initializer_range,
497 | hidden_dropout_prob=self.config.hidden_dropout_prob,
498 | attention_probs_dropout_prob=self.config.attention_probs_dropout_prob,
499 | do_cache=do_cache,
500 | cache=layer_cache,
501 | )
502 | new_kvs.append(new_kv)
503 |
504 | # [batch_size * seq_length, hidden_size]
505 | hidden_state = residual_mlp_layer(hidden_state + attention_output,
506 | intermediate_size=config.intermediate_size,
507 | hidden_dropout_prob=self.config.hidden_dropout_prob)
508 | self.hidden_state = hidden_state
509 |
510 | self.new_kvs = tf.stack(new_kvs, axis=1) if do_cache else None
511 |
512 | # Note that the hidden state is still flat (batch_size*hidden_size)
513 | self.logits_flat = tf.matmul(self.hidden_state, self.embedding_table, transpose_b=True)
514 |
515 | # THE OUTPUT BIAS DOES NOT SPARK JOY
516 | # output_bias = tf.get_variable('output_bias', shape=[config.vocab_size], initializer=tf.zeros_initializer())
517 | # self.logits_flat = tf.nn.bias_add(self.logits_flat, output_bias)
518 |
519 | @property
520 | def log_probs(self):
521 | logprobs_flat = tf.nn.log_softmax(self.logits_flat, axis=-1)
522 | return tf.reshape(logprobs_flat, [self.batch_size, self.seq_length, -1])
523 |
524 | def lm_loss(self):
525 | """
526 | :return: stuff
527 | """
528 | target_ids_flat = tf.reshape(self.target_ids, [-1])
529 |
530 | # 1 if it's valid and 0 otherwise.
531 | label_weights = tf.cast(tf.not_equal(target_ids_flat, self.pad_token_id), dtype=self.logits_flat.dtype)
532 |
533 | # [batch_size * seq_length, vocab_size]
534 | one_hot_labels = tf.one_hot(target_ids_flat,
535 | depth=self.config.vocab_size,
536 | dtype=self.logits_flat.dtype)
537 |
538 | # [batch_size * seq_length, vocab_size]
539 | logprobs_flat = tf.nn.log_softmax(self.logits_flat, axis=-1)
540 |
541 | per_example_loss = -tf.reduce_sum(logprobs_flat * one_hot_labels, axis=[-1])
542 |
543 | # per_example_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=self.logits_flat, labels=target_ids_flat)
544 |
545 | numerator = tf.reduce_sum(label_weights * per_example_loss)
546 | denominator = tf.reduce_sum(label_weights) + 1e-5
547 | loss = numerator / denominator
548 | return loss
549 |
550 | def pooled_output(self, clf_token):
551 | """
552 | Extract pooled output given a token that says where we should look
553 | :param clf_token:
554 | :return:
555 | """
556 | pool_idx = tf.cast(tf.argmax(tf.cast(tf.equal(self.input_ids, clf_token), tf.float32), 1), tf.int32)
557 | return tf.gather(self.hidden_state, tf.range(self.batch_size, dtype=tf.int32) * self.seq_length + pool_idx)
558 |
559 |
560 | def model_fn_builder(config: GroverConfig, init_checkpoint, learning_rate,
561 | num_train_steps, num_warmup_steps, use_tpu):
562 | """Returns `model_fn` closure for TPUEstimator."""
563 |
564 | def model_fn(features, labels, mode, params): # pylint: disable=unused-argument
565 | """The `model_fn` for TPUEstimator."""
566 |
567 | tf.logging.info("*** Features ***")
568 | for name in sorted(features.keys()):
569 | tf.logging.info(" name = %s, shape = %s" % (name, features[name].shape))
570 |
571 | input_ids = features["input_ids"]
572 |
573 | is_training = (mode == tf.estimator.ModeKeys.TRAIN)
574 |
575 | model = GroverModel(
576 | config=config,
577 | is_training=is_training,
578 | input_ids=input_ids,
579 | pad_token_id=config.pad_token_id,
580 | chop_off_last_token=True,
581 | )
582 |
583 | total_loss = model.lm_loss()
584 |
585 | if is_training:
586 | train_op, train_metrics = optimization_adafactor.create_optimizer(
587 | total_loss, learning_rate, num_train_steps, num_warmup_steps, use_tpu)
588 | tvars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
589 | else:
590 | train_op = None
591 | train_metrics = {}
592 | tvars = tf.trainable_variables()
593 |
594 | initialized_variable_names = {}
595 | scaffold_fn = None
596 | if init_checkpoint:
597 | (assignment_map, initialized_variable_names
598 | ) = get_assignment_map_from_checkpoint(tvars, init_checkpoint)
599 | if use_tpu:
600 | def tpu_scaffold():
601 | tf.train.init_from_checkpoint(init_checkpoint, assignment_map)
602 | return tf.train.Scaffold()
603 |
604 | scaffold_fn = tpu_scaffold
605 | else:
606 | tf.train.init_from_checkpoint(init_checkpoint, assignment_map)
607 |
608 | tf.logging.info("**** Trainable Variables ****")
609 | for var in tvars:
610 | init_string = ""
611 | if var.name in initialized_variable_names:
612 | init_string = ", *INIT_FROM_CKPT*"
613 | tf.logging.info(" name = %s, shape = %s%s", var.name, var.shape,
614 | init_string)
615 |
616 | output_spec = None
617 | if mode == tf.estimator.ModeKeys.TRAIN:
618 | if use_tpu:
619 | output_spec = tf.contrib.tpu.TPUEstimatorSpec(
620 | mode=mode,
621 | loss=total_loss,
622 | train_op=train_op,
623 | host_call=construct_scalar_host_call(metric_dict=train_metrics, model_dir=params['model_dir'],
624 | prefix='training/'),
625 | scaffold_fn=scaffold_fn)
626 | else:
627 | output_spec = tf.contrib.tpu.TPUEstimatorSpec(
628 | mode=mode,
629 | loss=total_loss,
630 | train_op=train_op,
631 | training_hooks=[
632 | tf.train.LoggingTensorHook({'loss': tf.metrics.mean(total_loss)[1]}, every_n_iter=100)],
633 | scaffold_fn=scaffold_fn)
634 |
635 | elif mode == tf.estimator.ModeKeys.EVAL:
636 | def metric_fn(total_loss):
637 | loss = tf.metrics.mean(values=total_loss)
638 | return {
639 | "eval_loss": loss,
640 | }
641 |
642 | eval_metrics = (metric_fn,
643 | [total_loss])
644 | output_spec = tf.contrib.tpu.TPUEstimatorSpec(
645 | mode=mode,
646 | loss=total_loss,
647 | eval_metrics=eval_metrics,
648 | scaffold_fn=scaffold_fn)
649 | else:
650 | gt_logprobs = tf.squeeze(tf.batch_gather(model.log_probs, model.target_ids[:, :, None]), axis=2)
651 |
652 | # Need top-p required under topp sampling!
653 | better_than_gt = model.log_probs > gt_logprobs[:, :, None]
654 | top_p_required = tf.reduce_sum(tf.cast(better_than_gt, tf.float32) * tf.exp(model.log_probs), axis=2)
655 |
656 | # No top-p sampling for now, since this seems to be too slow on TPUs
657 | if use_tpu:
658 | predictions = tf.reshape(
659 | tf.random.categorical(logits=model.logits_flat, num_samples=1),
660 | get_shape_list(model.target_ids),
661 | )
662 | else:
663 | # Argmax
664 | # predictions = tf.math.argmax(model.log_probs, axis=-1, output_type=tf.int32)
665 | predictions = tf.reshape(
666 | _top_p_sample(model.logits_flat, num_samples=1, p=0.99)['sample'],
667 | get_shape_list(model.target_ids),
668 | )
669 | pred_logprobs = tf.squeeze(tf.batch_gather(model.log_probs, predictions[:, :, None]), axis=2)
670 |
671 | output_spec = tf.contrib.tpu.TPUEstimatorSpec(
672 | mode=mode,
673 | predictions={'gt_logprobs': gt_logprobs,
674 | 'top_p_required': top_p_required,
675 | 'predictions': predictions,
676 | 'pred_logprobs': pred_logprobs,
677 | 'labels': input_ids},
678 | scaffold_fn=scaffold_fn)
679 | return output_spec
680 |
681 | return model_fn
682 |
683 |
684 | def sample_step(tokens, ignore_ids, news_config, batch_size=1, p_for_topp=0.95, cache=None, do_topk=False):
685 | """
686 | Helper function that samples from grover for a single step
687 | :param tokens: [batch_size, n_ctx_b] tokens that we will predict from
688 | :param ignore_ids: [n_vocab] mask of the tokens we don't want to predict
689 | :param news_config: config for the GroverModel
690 | :param batch_size: batch size to use
691 | :param p_for_topp: top-p or top-k threshold
692 | :param cache: [batch_size, news_config.num_hidden_layers, 2,
693 | news_config.num_attention_heads, n_ctx_a,
694 | news_config.hidden_size // news_config.num_attention_heads] OR, None
695 | :return: new_tokens, size [batch_size]
696 | new_probs, also size [batch_size]
697 | new_cache, size [batch_size, news_config.num_hidden_layers, 2, n_ctx_b,
698 | news_config.num_attention_heads, news_config.hidden_size // news_config.num_attention_heads]
699 | """
700 | model = GroverModel(
701 | config=news_config,
702 | is_training=False,
703 | input_ids=tokens,
704 | reuse=tf.AUTO_REUSE,
705 | scope='newslm',
706 | chop_off_last_token=False,
707 | do_cache=True,
708 | cache=cache,
709 | )
710 |
711 | # Extract the FINAL SEQ LENGTH
712 | batch_size_times_seq_length, vocab_size = get_shape_list(model.logits_flat, expected_rank=2)
713 | next_logits = tf.reshape(model.logits_flat, [batch_size, -1, vocab_size])[:, -1]
714 |
715 | if do_topk:
716 | sample_info = _top_k_sample(next_logits, num_samples=1, k=tf.cast(p_for_topp, dtype=tf.int32))
717 | else:
718 | sample_info = _top_p_sample(next_logits, ignore_ids=ignore_ids, num_samples=1, p=p_for_topp)
719 |
720 | new_tokens = tf.squeeze(sample_info['sample'], 1)
721 | new_probs = tf.squeeze(tf.batch_gather(sample_info['probs'], sample_info['sample']), 1)
722 | return {
723 | 'new_tokens': new_tokens,
724 | 'new_probs': new_probs,
725 | 'new_cache': model.new_kvs,
726 | }
727 |
728 |
729 | def initialize_from_context(initial_context, ignore_ids, news_config, p_for_topp=0.95, do_topk=False):
730 | """ same signature as sample_step"""
731 | batch_size, _ = get_shape_list(initial_context, expected_rank=2)
732 |
733 | context_output = sample_step(tokens=initial_context, ignore_ids=ignore_ids, news_config=news_config,
734 | batch_size=batch_size, p_for_topp=p_for_topp, cache=None, do_topk=do_topk)
735 | return {
736 | 'tokens': tf.concat([initial_context, context_output['new_tokens'][:, None]], 1),
737 | 'cache': context_output['new_cache'],
738 | 'probs': context_output['new_probs'][:, None]
739 | }
740 |
741 |
742 | def sample(news_config: GroverConfig, initial_context, eos_token, min_len, ignore_ids=None, p_for_topp=0.95,
743 | do_topk=False):
744 | """
745 | V1 version of: sample outputs from a model, and do it all at once
746 | :param news_config: Configuration used to construct the model
747 | :param initial_context: [batch_size, seq_length] that we'll start generating with
748 | :param eos_token: Stop generating if you see this (tf scalar)
749 | :param min_len: min length of sample
750 | :param ignore_ids: NEVER GENERATE THESE [vocab_size]
751 | :return:
752 | """
753 | batch_size, _ = get_shape_list(initial_context, expected_rank=2)
754 |
755 | if ignore_ids is None:
756 | ignore_ids = tf.constant([x == 0 for x in range(news_config.vocab_size)], dtype=tf.bool)
757 |
758 | with tf.name_scope('sample_sequence'):
759 | # Initial call to get cache
760 | context_output = initialize_from_context(initial_context, ignore_ids=ignore_ids, news_config=news_config,
761 | p_for_topp=p_for_topp,
762 | do_topk=do_topk)
763 | ctx = context_output['tokens']
764 | cache = context_output['cache']
765 | probs = context_output['probs']
766 |
767 | def body(ctx, cache, probs):
768 | """ for whatever reason this didn't work when I ran it on more than one at once... ugh."""
769 | next_outputs = sample_step(ctx[:, -1][:, None], ignore_ids=ignore_ids, news_config=news_config,
770 | batch_size=batch_size, p_for_topp=p_for_topp, cache=cache,
771 | do_topk=do_topk)
772 |
773 | # Update everything
774 | new_cache = tf.concat([cache, next_outputs['new_cache']], axis=-2)
775 | new_ids = tf.concat([ctx, next_outputs['new_tokens'][:, None]], axis=1)
776 | new_probs = tf.concat([probs, next_outputs['new_probs'][:, None]], axis=1)
777 | return [new_ids, new_cache, new_probs]
778 |
779 | def cond(ctx, cache, probs):
780 | # ctx = tf.Print(ctx,[tf.shape(ctx)])
781 | is_eos = tf.reduce_all(tf.reduce_any(tf.equal(ctx[:,-1:], eos_token), axis=1))
782 | is_len = tf.greater(get_shape_list(ctx)[1], min_len)
783 | return tf.logical_not(tf.logical_and(is_eos, is_len))
784 |
785 | tokens, cache, probs = tf.while_loop(
786 | cond=cond, body=body, maximum_iterations=1025 - get_shape_list(ctx)[1],
787 | loop_vars=[ctx, cache, probs],
788 | shape_invariants=[tf.TensorShape([batch_size, None]),
789 | tf.TensorShape(
790 | [batch_size, news_config.num_hidden_layers, 2,
791 | news_config.num_attention_heads,
792 | None, news_config.hidden_size // news_config.num_attention_heads]),
793 | tf.TensorShape([batch_size, None]),
794 | ],
795 | back_prop=False,
796 | )
797 | return tokens, probs
798 |
--------------------------------------------------------------------------------
/train/optimization_adafactor.py:
--------------------------------------------------------------------------------
1 | # Original work Copyright 2018 The Google AI Language Team Authors.
2 | # Modified work Copyright 2019 Rowan Zellers
3 | #
4 | # Licensed under the Apache License, Version 2.0 (the "License");
5 | # you may not use this file except in compliance with the License.
6 | # You may obtain a copy of the License at
7 | #
8 | # http://www.apache.org/licenses/LICENSE-2.0
9 | #
10 | # Unless required by applicable law or agreed to in writing, software
11 | # distributed under the License is distributed on an "AS IS" BASIS,
12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 | # See the License for the specific language governing permissions and
14 | # limitations under the License.
15 | import re
16 | import tensorflow as tf
17 | from train.utils import get_shape_list
18 |
19 |
20 | def create_optimizer(loss, init_lr, num_train_steps, num_warmup_steps, use_tpu):
21 | """Creates an optimizer training op."""
22 | global_step = tf.train.get_or_create_global_step()
23 |
24 | learning_rate = tf.constant(value=init_lr, shape=[], dtype=tf.float32)
25 |
26 | # Implements linear decay of the learning rate.
27 | learning_rate = tf.train.polynomial_decay(
28 | learning_rate,
29 | global_step,
30 | num_train_steps,
31 | end_learning_rate=0.0,
32 | power=1.0,
33 | cycle=False)
34 |
35 | # Implements linear warmup. I.e., if global_step < num_warmup_steps, the
36 | # learning rate will be `global_step/num_warmup_steps * init_lr`.
37 | if num_warmup_steps:
38 | global_steps_int = tf.cast(global_step, tf.int32)
39 | warmup_steps_int = tf.constant(num_warmup_steps, dtype=tf.int32)
40 |
41 | global_steps_float = tf.cast(global_steps_int, tf.float32)
42 | warmup_steps_float = tf.cast(warmup_steps_int, tf.float32)
43 |
44 | warmup_percent_done = global_steps_float / warmup_steps_float
45 | warmup_learning_rate = init_lr * warmup_percent_done
46 |
47 | is_warmup = tf.cast(global_steps_int < warmup_steps_int, tf.float32)
48 | learning_rate = (
49 | (1.0 - is_warmup) * learning_rate + is_warmup * warmup_learning_rate)
50 |
51 | # It is recommended that you use this optimizer for fine tuning, since this
52 | # is how the model was trained (note that the Adam m/v variables are NOT
53 | # loaded from init_checkpoint.)
54 | optimizer = AdaFactorOptimizer(
55 | learning_rate=learning_rate,
56 | weight_decay_rate=0.01,
57 | beta_1=0.9,
58 | beta_2=0.999,
59 | epsilon=1e-6,
60 | exclude_from_weight_decay=["LayerNorm", "layer_norm", "bias"])
61 |
62 | if use_tpu:
63 | optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer)
64 |
65 | tvars = tf.trainable_variables()
66 | grads = tf.gradients(loss, tvars)
67 |
68 | # You could do this, but instead we don't because a) it's slow and b) we already did the 'update clipping'
69 | # (grads, _) = tf.clip_by_global_norm(grads, clip_norm=1.0)
70 |
71 | train_op = optimizer.apply_gradients(
72 | zip(grads, tvars), global_step=global_step)
73 |
74 | # Normally the global step update is done inside of `apply_gradients`.
75 | # However, `AdaFactorOptimizer` doesn't do this. But if you use
76 | # a different optimizer, you should probably take this line out.
77 | new_global_step = global_step + 1
78 | train_op = tf.group(train_op, [global_step.assign(new_global_step)])
79 |
80 | train_metrics = {
81 | 'learning_rate': learning_rate,
82 | 'minibatch_loss': loss,
83 | # 'minibatch_ppl': tf.math.exp(loss),
84 | }
85 | return train_op, train_metrics
86 |
87 |
88 | class AdaFactorOptimizer(tf.compat.v1.train.Optimizer):
89 | """here's the optimizer we'll use"""
90 |
91 | def __init__(self,
92 | learning_rate,
93 | weight_decay_rate=0.0,
94 | beta_1=0.9,
95 | beta_2=0.999,
96 | epsilon=1e-6,
97 | exclude_from_weight_decay=None,
98 | clipping_rate=1.0,
99 | name="AdaFactorOptimizer"):
100 | """Constructs a AdaFactorOptimizer."""
101 | super(AdaFactorOptimizer, self).__init__(False, name)
102 |
103 | self.learning_rate = learning_rate
104 | self.weight_decay_rate = weight_decay_rate
105 | self.beta_1 = beta_1
106 | self.beta_2 = beta_2
107 | self.epsilon = epsilon
108 | self.epsilon1 = 1e-30
109 | self.epsilon2 = 0.001
110 | self.clipping_rate = clipping_rate
111 | self.exclude_from_weight_decay = exclude_from_weight_decay
112 | self.use_locking = False
113 |
114 | def _use_factored(self, shape):
115 | return len(shape) >= 2
116 |
117 | def _parameter_scale(self, var):
118 | """Estimate the scale of the parameters from the current values.
119 | We include a minimum value of 0.001 to give it a chance to escape 0
120 | if it was zero-initialized.
121 | Instead of using the value, we could impute the scale from the shape,
122 | as initializers do.
123 | Args:
124 | var: a variable or Tensor.
125 | Returns:
126 | a Scalar
127 | """
128 | return tf.maximum(reduce_rms(var), self.epsilon2)
129 |
130 | def apply_gradients(self, grads_and_vars, global_step=None, name=None):
131 | """See base class."""
132 | assignments = []
133 | for (grad, param) in grads_and_vars:
134 | if grad is None or param is None:
135 | continue
136 |
137 | param_name = self._get_variable_name(param.name)
138 | shape_list = get_shape_list(param, expected_rank=[1, 2])
139 |
140 | # decay_rate = 1 - tf.pow(tf.cast(tf.train.get_or_create_global_step(), tf.float32) + 1.0, -0.8)
141 | decay_rate = self.beta_2
142 | grad_squared = tf.square(grad) + self.epsilon1
143 |
144 | update_scale = self.learning_rate
145 | # update_scale = self.learning_rate * tf.cast(self._parameter_scale(param), dtype=tf.float32)
146 |
147 | # HACK: Make things dependent on grad.
148 | # This confounds the XLA rewriter and keeps it from fusing computations
149 | # across different variables. This fusion is a bad for HBM usage, since
150 | # it causes the gradients to persist in memory.
151 | grad_squared_mean = tf.reduce_mean(grad_squared)
152 | decay_rate += grad_squared_mean * 1e-30
153 | update_scale += grad_squared_mean * 1e-30
154 |
155 | # END HACK
156 |
157 | if self._use_factored(shape_list):
158 | num_rows, num_columns = shape_list
159 |
160 | vr = tf.get_variable(
161 | name=param_name + "/adafactor_vr",
162 | shape=[num_rows],
163 | dtype=tf.float32,
164 | trainable=False,
165 | initializer=tf.zeros_initializer())
166 | vc = tf.get_variable(
167 | name=param_name + "/adafactor_vc",
168 | shape=[num_columns],
169 | dtype=tf.float32,
170 | trainable=False,
171 | initializer=tf.zeros_initializer())
172 |
173 | next_vr = decay_rate * vr + (1 - decay_rate) * tf.reduce_mean(grad_squared, 1)
174 | next_vc = decay_rate * vc + (1 - decay_rate) * tf.reduce_mean(grad_squared, 0)
175 |
176 | long_term_mean = tf.reduce_mean(next_vr, -1, keepdims=True)
177 | r_factor = tf.rsqrt(next_vr / long_term_mean + self.epsilon1)
178 | c_factor = tf.rsqrt(next_vc + self.epsilon1)
179 | update = grad * tf.expand_dims(r_factor, -1) * tf.expand_dims(c_factor, -2)
180 |
181 | assignments.append(vr.assign(next_vr, use_locking=self.use_locking))
182 | assignments.append(vc.assign(next_vc, use_locking=self.use_locking))
183 | else:
184 | v = tf.get_variable(
185 | name=param_name + "/adafactor_v",
186 | shape=shape_list,
187 | dtype=tf.float32,
188 | trainable=False,
189 | initializer=tf.zeros_initializer())
190 | next_v = decay_rate * v + (1 - decay_rate) * grad_squared
191 |
192 | assignments.append(v.assign(next_v, use_locking=self.use_locking))
193 | update = grad * tf.rsqrt(next_v + self.epsilon1)
194 |
195 | clipping_denom = tf.maximum(1.0, reduce_rms(update) / self.clipping_rate)
196 | update /= clipping_denom
197 |
198 | # Do weight decay
199 | # Just adding the square of the weights to the loss function is *not*
200 | # the correct way of using L2 regularization/weight decay with Adam,
201 | # since that will interact with the m and v parameters in strange ways.
202 | #
203 | # Instead we want ot decay the weights in a manner that doesn't interact
204 | # with the m/v parameters. This is equivalent to adding the square
205 | # # of the weights to the loss with plain (non-momentum) SGD.
206 | if self._do_use_weight_decay(param_name):
207 | update += self.weight_decay_rate * param
208 |
209 | update_with_lr = update_scale * update
210 | next_param = param - update_with_lr
211 |
212 | assignments.append(param.assign(next_param, use_locking=self.use_locking))
213 | return tf.group(*assignments, name=name)
214 |
215 | def _do_use_weight_decay(self, param_name):
216 | """Whether to use L2 weight decay for `param_name`."""
217 | if not self.weight_decay_rate:
218 | return False
219 | if self.exclude_from_weight_decay:
220 | for r in self.exclude_from_weight_decay:
221 | if re.search(r, param_name) is not None:
222 | return False
223 | return True
224 |
225 | def _get_variable_name(self, param_name):
226 | """Get the variable name from the tensor name."""
227 | m = re.match("^(.*):\\d+$", param_name)
228 | if m is not None:
229 | param_name = m.group(1)
230 | return param_name
231 |
232 |
233 | def reduce_rms(x):
234 | return tf.sqrt(tf.reduce_mean(tf.square(x)))
235 |
--------------------------------------------------------------------------------
/train/train_tpu.py:
--------------------------------------------------------------------------------
1 | # Original work Copyright 2018 The Google AI Language Team Authors.
2 | # Modified work Copyright 2019 Rowan Zellers
3 | #
4 | # Licensed under the Apache License, Version 2.0 (the "License");
5 | # you may not use this file except in compliance with the License.
6 | # You may obtain a copy of the License at
7 | #
8 | # http://www.apache.org/licenses/LICENSE-2.0
9 | #
10 | # Unless required by applicable law or agreed to in writing, software
11 | # distributed under the License is distributed on an "AS IS" BASIS,
12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 | # See the License for the specific language governing permissions and
14 | # limitations under the License.
15 |
16 | """ Training script! """
17 |
18 | import tensorflow.compat.v1 as tf
19 |
20 | from train.dataloader import input_fn_builder
21 | from train.modeling import model_fn_builder, GroverConfig
22 |
23 | flags = tf.flags
24 |
25 | FLAGS = flags.FLAGS
26 |
27 | ## Required parameters
28 | flags.DEFINE_string(
29 | "config_file", 'configs/base.json',
30 | "The config json file corresponding to the pre-trained news model. "
31 | "This specifies the model architecture.")
32 |
33 | flags.DEFINE_string(
34 | "input_file", None,
35 | "Input TF example files (can be a glob or comma separated).")
36 |
37 | flags.DEFINE_string(
38 | "output_dir", None,
39 | "The output directory where the model checkpoints will be written.")
40 |
41 | ## Other parameters
42 | flags.DEFINE_string(
43 | "init_checkpoint", None,
44 | "Initial checkpoint (usually from a pre-trained model).")
45 |
46 | flags.DEFINE_integer(
47 | "max_seq_length", 1024,
48 | "The maximum total input sequence length after BPE tokenization. "
49 | "Sequences longer than this will be truncated, and sequences shorter "
50 | "than this will be padded. Must match data generation.")
51 |
52 | flags.DEFINE_integer("train_batch_size", 32, "Total batch size for training.")
53 |
54 | flags.DEFINE_float("learning_rate", 5e-5, "The initial learning rate for adafactor.")
55 |
56 | flags.DEFINE_integer("num_train_steps", 100000, "Number of training steps.")
57 |
58 | flags.DEFINE_integer("num_warmup_steps", 10000, "Number of warmup steps.")
59 |
60 | flags.DEFINE_integer("save_checkpoints_steps", 1000,
61 | "How often to save the model checkpoint.")
62 |
63 | flags.DEFINE_integer("iterations_per_loop", 1000,
64 | "How many steps to make in each estimator call.")
65 |
66 | flags.DEFINE_integer("max_eval_steps", 100, "Maximum number of eval steps.")
67 |
68 | flags.DEFINE_bool("use_tpu", False, "Whether to use TPU or GPU/CPU.")
69 |
70 | flags.DEFINE_string(
71 | "tpu_name", None,
72 | "The Cloud TPU to use for training. This should be either the name "
73 | "used when creating the Cloud TPU, or a grpc://ip.address.of.tpu:8470 "
74 | "url.")
75 |
76 | flags.DEFINE_string(
77 | "tpu_zone", None,
78 | "[Optional] GCE zone where the Cloud TPU is located in. If not "
79 | "specified, we will attempt to automatically detect the GCE project from "
80 | "metadata.")
81 |
82 | flags.DEFINE_string(
83 | "gcp_project", None,
84 | "[Optional] Project name for the Cloud TPU-enabled project. If not "
85 | "specified, we will attempt to automatically detect the GCE project from "
86 | "metadata.")
87 |
88 | flags.DEFINE_string("master", None, "[Optional] TensorFlow master URL.")
89 |
90 | flags.DEFINE_integer(
91 | "num_tpu_cores", 8,
92 | "Only used if `use_tpu` is True. Total number of TPU cores to use.")
93 |
94 |
95 | def main(_):
96 | tf.logging.set_verbosity(tf.logging.INFO)
97 |
98 | news_config = GroverConfig.from_json_file(FLAGS.config_file)
99 |
100 | tf.gfile.MakeDirs(FLAGS.output_dir)
101 |
102 | input_files = []
103 | for input_pattern in FLAGS.input_file.split(","):
104 | input_files.extend(tf.gfile.Glob(input_pattern))
105 |
106 | tf.logging.info("*** Input Files ***")
107 | for input_file in input_files:
108 | tf.logging.info(" %s" % input_file)
109 |
110 | tpu_cluster_resolver = None
111 | if FLAGS.use_tpu and FLAGS.tpu_name:
112 | tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(
113 | FLAGS.tpu_name, zone=FLAGS.tpu_zone, project=FLAGS.gcp_project)
114 |
115 | is_per_host = tf.contrib.tpu.InputPipelineConfig.PER_HOST_V2
116 | run_config = tf.contrib.tpu.RunConfig(
117 | cluster=tpu_cluster_resolver,
118 | master=FLAGS.master,
119 | model_dir=FLAGS.output_dir,
120 | save_checkpoints_steps=FLAGS.save_checkpoints_steps,
121 | keep_checkpoint_max=None,
122 | tpu_config=tf.contrib.tpu.TPUConfig(
123 | iterations_per_loop=FLAGS.iterations_per_loop,
124 | num_shards=FLAGS.num_tpu_cores,
125 | per_host_input_for_training=is_per_host))
126 |
127 | model_fn = model_fn_builder(news_config, init_checkpoint=FLAGS.init_checkpoint,
128 | learning_rate=FLAGS.learning_rate,
129 | num_train_steps=FLAGS.num_train_steps,
130 | num_warmup_steps=FLAGS.num_warmup_steps,
131 | use_tpu=FLAGS.use_tpu,
132 | )
133 |
134 | # If TPU is not available, this will fall back to normal Estimator on CPU
135 | # or GPU.
136 | estimator = tf.contrib.tpu.TPUEstimator(
137 | use_tpu=FLAGS.use_tpu,
138 | model_fn=model_fn,
139 | config=run_config,
140 | train_batch_size=FLAGS.train_batch_size,
141 | eval_batch_size=FLAGS.train_batch_size,
142 | params={'model_dir': FLAGS.output_dir}
143 | )
144 |
145 | tf.logging.info("***** Running training *****")
146 | tf.logging.info(" Batch size = %d", FLAGS.train_batch_size)
147 | train_input_fn = input_fn_builder(
148 | input_files=input_files,
149 | seq_length=FLAGS.max_seq_length,
150 | is_training=True)
151 |
152 | estimator.train(input_fn=train_input_fn, max_steps=FLAGS.num_train_steps)
153 |
154 | if __name__ == "__main__":
155 | flags.mark_flag_as_required("input_file")
156 | flags.mark_flag_as_required("output_dir")
157 | tf.app.run()
158 |
--------------------------------------------------------------------------------
/train/train_tpu_adafactor.sh:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env bash
2 |
3 | export PYTHONPATH=../
4 |
5 | learning_rate=1e-4
6 | init_checkpoint=""
7 | max_seq_length=1024
8 | save_checkpoint_steps=1000
9 |
10 | # You can customize the training here
11 | # mega, medium, or base
12 | model_type="mega"
13 | OUTPUT_DIR="gs://" # put your output directory here
14 | input_file="gs://" # put your input files here, it can also be something like "*.tfrecord"
15 |
16 | if [ ${model_type} == "base" ]; then
17 | num_tpu_cores=32
18 | batch_size_per_core=16
19 | elif [ ${model_type} == "medium" ]; then
20 | num_tpu_cores=128
21 | batch_size_per_core=4
22 | elif [ ${model_type} == "mega" ]; then
23 | num_tpu_cores=256
24 | batch_size_per_core=2
25 | fi
26 |
27 |
28 | # there are 20k * 1024 examples so this translates to 20 epochs. seems ok and i can run for more if needed
29 | num_train_steps=800000
30 |
31 | # Make sure batch size scales.
32 | let batch_size="$batch_size_per_core * $num_tpu_cores"
33 |
34 | python train_tpu.py \
35 | --config_file=configs/${model_type}.json \
36 | --input_file=${input_file} \
37 | --output_dir=${OUTPUT_DIR} \
38 | --max_seq_length=${max_seq_length} \
39 | --train_batch_size=${batch_size} \
40 | --learning_rate=${learning_rate} \
41 | --num_train_steps=${num_train_steps} \
42 | --num_warmup_steps=10000 \
43 | --save_checkpoints_steps=${save_checkpoint_steps} \
44 | --iterations_per_loop=${save_checkpoint_steps} \
45 | --use_tpu=True \
46 | --tpu_name=$(hostname) \
47 | --num_tpu_cores=$num_tpu_cores \
48 | --init_checkpoint=${init_checkpoint}
49 |
--------------------------------------------------------------------------------
/train/utils.py:
--------------------------------------------------------------------------------
1 | # Original work Copyright 2018 The Google AI Language Team Authors.
2 | # Modified work Copyright 2019 Rowan Zellers
3 | #
4 | # Licensed under the Apache License, Version 2.0 (the "License");
5 | # you may not use this file except in compliance with the License.
6 | # You may obtain a copy of the License at
7 | #
8 | # http://www.apache.org/licenses/LICENSE-2.0
9 | #
10 | # Unless required by applicable law or agreed to in writing, software
11 | # distributed under the License is distributed on an "AS IS" BASIS,
12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 | # See the License for the specific language governing permissions and
14 | # limitations under the License.
15 |
16 | import collections
17 | import re
18 |
19 | import six
20 | import tensorflow.compat.v1 as tf
21 | import numpy as np
22 | from tensorflow.python.lib.io import file_io
23 |
24 |
25 | def _save_np(absolute_fn, array):
26 | if absolute_fn.startswith('gs://'):
27 | with file_io.FileIO(absolute_fn, 'w') as f:
28 | np.save(f, array)
29 | else:
30 | np.save(absolute_fn, array)
31 |
32 |
33 | def assert_rank(tensor, expected_rank, name=None):
34 | """Raises an exception if the tensor rank is not of the expected rank.
35 |
36 | Args:
37 | tensor: A tf.Tensor to check the rank of.
38 | expected_rank: Python integer or list of integers, expected rank.
39 | name: Optional name of the tensor for the error message.
40 |
41 | Raises:
42 | ValueError: If the expected shape doesn't match the actual shape.
43 | """
44 | if name is None:
45 | name = tensor.name
46 |
47 | expected_rank_dict = {}
48 | if isinstance(expected_rank, six.integer_types):
49 | expected_rank_dict[expected_rank] = True
50 | else:
51 | for x in expected_rank:
52 | expected_rank_dict[x] = True
53 |
54 | actual_rank = tensor.shape.ndims
55 | if actual_rank not in expected_rank_dict:
56 | scope_name = tf.get_variable_scope().name
57 | raise ValueError(
58 | "For the tensor `%s` in scope `%s`, the actual rank "
59 | "`%d` (shape = %s) is not equal to the expected rank `%s`" %
60 | (name, scope_name, actual_rank, str(tensor.shape), str(expected_rank)))
61 |
62 |
63 | def get_shape_list(tensor, expected_rank=None, name=None):
64 | """Returns a list of the shape of tensor, preferring static dimensions.
65 |
66 | Args:
67 | tensor: A tf.Tensor object to find the shape of.
68 | expected_rank: (optional) int. The expected rank of `tensor`. If this is
69 | specified and the `tensor` has a different rank, and exception will be
70 | thrown.
71 | name: Optional name of the tensor for the error message.
72 |
73 | Returns:
74 | A list of dimensions of the shape of tensor. All static dimensions will
75 | be returned as python integers, and dynamic dimensions will be returned
76 | as tf.Tensor scalars.
77 | """
78 | if name is None:
79 | name = tensor.name
80 |
81 | if expected_rank is not None:
82 | assert_rank(tensor, expected_rank, name)
83 |
84 | shape = tensor.shape.as_list()
85 |
86 | non_static_indexes = []
87 | for (index, dim) in enumerate(shape):
88 | if dim is None:
89 | non_static_indexes.append(index)
90 |
91 | if not non_static_indexes:
92 | return shape
93 |
94 | dyn_shape = tf.shape(tensor)
95 | for index in non_static_indexes:
96 | shape[index] = dyn_shape[index]
97 | return shape
98 |
99 |
100 | def gelu(input_tensor):
101 | """Gaussian Error Linear Unit.
102 |
103 | This is a smoother version of the RELU.
104 | Original paper: https://arxiv.org/abs/1606.08415
105 |
106 | Args:
107 | input_tensor: float Tensor to perform activation.
108 |
109 | Returns:
110 | `input_tensor` with the GELU activation applied.
111 | """
112 | cdf = 0.5 * (1.0 + tf.erf(input_tensor / tf.sqrt(2.0)))
113 | return input_tensor * cdf
114 |
115 |
116 | def layer_norm(input_tensor, name=None, epsilon=1e-5):
117 | """Run layer normalization on the last dimension of the tensor."""
118 | name2use = f'LayerNorm_{name}' if name is not None else name
119 | with tf.variable_scope(name2use, default_name='LayerNorm'):
120 | dim = input_tensor.shape[-1].value
121 | gamma = tf.get_variable('gamma', [dim], initializer=tf.constant_initializer(1))
122 | beta = tf.get_variable('beta', [dim], initializer=tf.constant_initializer(0))
123 | mean = tf.reduce_mean(input_tensor, axis=-1, keepdims=True)
124 | std = tf.reduce_mean(tf.square(input_tensor - mean), axis=-1, keepdims=True)
125 | input_tensor = (input_tensor - mean) * tf.rsqrt(std + epsilon)
126 | input_tensor = input_tensor * gamma + beta
127 | return input_tensor
128 |
129 |
130 | def dropout(input_tensor, dropout_prob):
131 | """Perform dropout.
132 |
133 | Args:
134 | input_tensor: float Tensor.
135 | dropout_prob: Python float. The probability of dropping out a value (NOT of
136 | *keeping* a dimension as in `tf.nn.dropout`).
137 |
138 | Returns:
139 | A version of `input_tensor` with dropout applied.
140 | """
141 | if dropout_prob is None or dropout_prob == 0.0:
142 | return input_tensor
143 | output = tf.nn.dropout(input_tensor, rate=dropout_prob)
144 | return output
145 |
146 |
147 | def get_attention_mask(nd, ns, *, dtype):
148 | """
149 | this is a TPU compatible version of tf.matrix_band_part(tf.ones([nd, ns]), -1, ns-nd)
150 | where the lower right triangle contains 1s
151 | """
152 | i = tf.range(nd)[:, None]
153 | j = tf.range(ns)
154 | m = i >= j - ns + nd
155 | return tf.cast(m, dtype)
156 |
157 |
158 | def get_assignment_map_from_checkpoint(tvars, init_checkpoint):
159 | """Compute the union of the current variables and checkpoint variables."""
160 | assignment_map = {}
161 | initialized_variable_names = {}
162 |
163 | name_to_variable = collections.OrderedDict()
164 | for var in tvars:
165 | name = var.name
166 | m = re.match("^(.*):\\d+$", name)
167 | if m is not None:
168 | name = m.group(1)
169 | name_to_variable[name] = var
170 |
171 | init_vars = tf.train.list_variables(init_checkpoint)
172 |
173 | assignment_map = collections.OrderedDict()
174 | for x in init_vars:
175 | (name, var) = (x[0], x[1])
176 | if name not in name_to_variable:
177 | continue
178 | assignment_map[name] = name
179 | initialized_variable_names[name] = 1
180 | initialized_variable_names[name + ":0"] = 1
181 | return (assignment_map, initialized_variable_names)
182 |
183 |
184 | def construct_scalar_host_call(metric_dict, model_dir, prefix=""):
185 | """Construct a host call to log scalars when training on TPU.
186 |
187 | Args:
188 | metric_dict: A dict of the tensors to be logged.
189 | model_dir: The location to write the summary.
190 | prefix: The prefix (if any) to prepend to the metric names.
191 |
192 | Returns:
193 | A tuple of (function, args_to_be_passed_to_said_function)
194 | """
195 | metric_names = list(metric_dict.keys())
196 |
197 | def host_call_fn(global_step, *args):
198 | """Training host call. Creates scalar summaries for training metrics.
199 |
200 | This function is executed on the CPU and should not directly reference
201 | any Tensors in the rest of the `model_fn`. To pass Tensors from the
202 | model to the `metric_fn`, provide as part of the `host_call`. See
203 | https://www.tensorflow.org/api_docs/python/tf/contrib/tpu/TPUEstimatorSpec
204 | for more information.
205 |
206 | Arguments should match the list of `Tensor` objects passed as the second
207 | element in the tuple passed to `host_call`.
208 |
209 | Args:
210 | global_step: `Tensor with shape `[batch]` for the global_step
211 | *args: Remaining tensors to log.
212 |
213 | Returns:
214 | List of summary ops to run on the CPU host.
215 | """
216 | step = global_step[0]
217 | with tf.contrib.summary.create_file_writer(
218 | logdir=model_dir, filename_suffix=".host_call").as_default():
219 | with tf.contrib.summary.always_record_summaries():
220 | for i, name in enumerate(metric_names):
221 | tf.contrib.summary.scalar(prefix + name, args[i][0], step=step)
222 |
223 | return tf.contrib.summary.all_summary_ops()
224 |
225 | # To log the current learning rate, and gradient norm for Tensorboard, the
226 | # summary op needs to be run on the host CPU via host_call. host_call
227 | # expects [batch_size, ...] Tensors, thus reshape to introduce a batch
228 | # dimension. These Tensors are implicitly concatenated to
229 | # [params['batch_size']].
230 | global_step_tensor = tf.reshape(
231 | tf.compat.v1.train.get_or_create_global_step(), [1])
232 | other_tensors = [tf.reshape(metric_dict[key], [1]) for key in metric_names]
233 |
234 | return host_call_fn, [global_step_tensor] + other_tensors
235 |
--------------------------------------------------------------------------------