![\"Open](\"https://colab.research.google.com/assets/colab-badge.svg\"){kind=icon}
"
11 | ]
12 | },
13 | {
14 | "cell_type": "markdown",
15 | "metadata": {
16 | "colab_type": "text",
17 | "id": "fnOHnctkG6kW"
18 | },
19 | "source": [
20 | "# XLNet IMDB movie review classification project\n",
21 | "\n",
22 | "This notebook is for classifying the [imdb sentiment dataset](https://ai.stanford.edu/~amaas/data/sentiment/). It will be easy to edit this notebook in order to run all of the classification tasks referenced in the [XLNet paper](https://arxiv.org/abs/1906.08237). Whilst you cannot expect to obtain the state-of-the-art results in the paper on a GPU, this model will still score very highly. "
23 | ]
24 | },
25 | {
26 | "cell_type": "markdown",
27 | "metadata": {
28 | "colab_type": "text",
29 | "id": "2mBzLdrdzodb"
30 | },
31 | "source": [
32 | "## Setup\n",
33 | "Install dependencies"
34 | ]
35 | },
36 | {
37 | "cell_type": "code",
38 | "execution_count": 0,
39 | "metadata": {
40 | "colab": {},
41 | "colab_type": "code",
42 | "id": "hRHRPImGUth7"
43 | },
44 | "outputs": [],
45 | "source": [
46 | "! pip install sentencepiece"
47 | ]
48 | },
49 | {
50 | "cell_type": "markdown",
51 | "metadata": {
52 | "colab_type": "text",
53 | "id": "jy8gUsPuJNyw"
54 | },
55 | "source": [
56 | "Download the pretrained XLNet model and unzip"
57 | ]
58 | },
59 | {
60 | "cell_type": "code",
61 | "execution_count": 0,
62 | "metadata": {
63 | "colab": {},
64 | "colab_type": "code",
65 | "id": "HfPDGsUtHKG0"
66 | },
67 | "outputs": [],
68 | "source": [
69 | "# only needs to be done once\n",
70 | "! wget https://storage.googleapis.com/xlnet/released_models/cased_L-24_H-1024_A-16.zip\n",
71 | "! unzip cased_L-24_H-1024_A-16.zip "
72 | ]
73 | },
74 | {
75 | "cell_type": "markdown",
76 | "metadata": {
77 | "colab_type": "text",
78 | "id": "4uUwjq3BJRbu"
79 | },
80 | "source": [
81 | "Download extract the imdb dataset - surpessing output"
82 | ]
83 | },
84 | {
85 | "cell_type": "code",
86 | "execution_count": 0,
87 | "metadata": {
88 | "colab": {},
89 | "colab_type": "code",
90 | "id": "QOGRICbOIsU8"
91 | },
92 | "outputs": [],
93 | "source": [
94 | "! wget http://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz\n",
95 | "! tar zxf aclImdb_v1.tar.gz"
96 | ]
97 | },
98 | {
99 | "cell_type": "markdown",
100 | "metadata": {
101 | "colab_type": "text",
102 | "id": "yGY_ggUUMrwU"
103 | },
104 | "source": [
105 | "Git clone XLNet repo for access to run_classifier and the rest of the xlnet module"
106 | ]
107 | },
108 | {
109 | "cell_type": "code",
110 | "execution_count": 0,
111 | "metadata": {
112 | "colab": {},
113 | "colab_type": "code",
114 | "id": "-r190eYVMpiG"
115 | },
116 | "outputs": [],
117 | "source": [
118 | "! git clone https://github.com/zihangdai/xlnet.git"
119 | ]
120 | },
121 | {
122 | "cell_type": "markdown",
123 | "metadata": {
124 | "colab_type": "text",
125 | "id": "jDP-IaVuPC-z"
126 | },
127 | "source": [
128 | "## Define Variables\n",
129 | "Define all the dirs: data, xlnet scripts & pretrained model. \n",
130 | "If you would like to save models then you can authenticate a GCP account and use that for the OUTPUT_DIR & CHECKPOINT_DIR - you will need a large amount storage to fix these models. \n",
131 | "\n",
132 | "Alternatively it is easy to integrate a google drive account, checkout this guide for [I/O in colab](https://colab.research.google.com/notebooks/io.ipynb) but rememeber these will take up a large amount of storage. \n"
133 | ]
134 | },
135 | {
136 | "cell_type": "code",
137 | "execution_count": 0,
138 | "metadata": {
139 | "colab": {},
140 | "colab_type": "code",
141 | "id": "y7N_xVwavQlV"
142 | },
143 | "outputs": [],
144 | "source": [
145 | "SCRIPTS_DIR = 'xlnet' #@param {type:\"string\"}\n",
146 | "DATA_DIR = 'aclImdb' #@param {type:\"string\"}\n",
147 | "OUTPUT_DIR = 'proc_data/imdb' #@param {type:\"string\"}\n",
148 | "PRETRAINED_MODEL_DIR = 'xlnet_cased_L-24_H-1024_A-16' #@param {type:\"string\"}\n",
149 | "CHECKPOINT_DIR = 'exp/imdb' #@param {type:\"string\"}"
150 | ]
151 | },
152 | {
153 | "cell_type": "markdown",
154 | "metadata": {
155 | "colab_type": "text",
156 | "id": "jR6euqwL1KBV"
157 | },
158 | "source": [
159 | "## Run Model\n",
160 | "This will set off the fine tuning of XLNet. There are a few things to note here:\n",
161 | "\n",
162 | "\n",
163 | "1. This script will train and evaluate the model\n",
164 | "2. This will store the results locally on colab and will be lost when you are disconnected from the runtime\n",
165 | "3. This uses the large version of the model (base not released presently)\n",
166 | "4. We are using a max seq length of 128 with a batch size of 8 please refer to the [README](https://github.com/zihangdai/xlnet#memory-issue-during-finetuning) for why this is.\n",
167 | "5. This will take approx 4hrs to run on GPU.\n",
168 | "\n"
169 | ]
170 | },
171 | {
172 | "cell_type": "code",
173 | "execution_count": 0,
174 | "metadata": {
175 | "colab": {},
176 | "colab_type": "code",
177 | "id": "CEMuT6LU0avg"
178 | },
179 | "outputs": [],
180 | "source": [
181 | "train_command = \"python xlnet/run_classifier.py \\\n",
182 | " --do_train=True \\\n",
183 | " --do_eval=True \\\n",
184 | " --eval_all_ckpt=True \\\n",
185 | " --task_name=imdb \\\n",
186 | " --data_dir=\"+DATA_DIR+\" \\\n",
187 | " --output_dir=\"+OUTPUT_DIR+\" \\\n",
188 | " --model_dir=\"+CHECKPOINT_DIR+\" \\\n",
189 | " --uncased=False \\\n",
190 | " --spiece_model_file=\"+PRETRAINED_MODEL_DIR+\"/spiece.model \\\n",
191 | " --model_config_path=\"+PRETRAINED_MODEL_DIR+\"/xlnet_config.json \\\n",
192 | " --init_checkpoint=\"+PRETRAINED_MODEL_DIR+\"/xlnet_model.ckpt \\\n",
193 | " --max_seq_length=128 \\\n",
194 | " --train_batch_size=8 \\\n",
195 | " --eval_batch_size=8 \\\n",
196 | " --num_hosts=1 \\\n",
197 | " --num_core_per_host=1 \\\n",
198 | " --learning_rate=2e-5 \\\n",
199 | " --train_steps=4000 \\\n",
200 | " --warmup_steps=500 \\\n",
201 | " --save_steps=500 \\\n",
202 | " --iterations=500\"\n",
203 | "\n",
204 | "! {train_command}\n"
205 | ]
206 | },
207 | {
208 | "cell_type": "markdown",
209 | "metadata": {
210 | "colab_type": "text",
211 | "id": "VvhqD-sO0Kyh"
212 | },
213 | "source": [
214 | "## Running & Results\n",
215 | "These are the results that I got from running this experiment\n",
216 | "### Params\n",
217 | "* --max_seq_length=128 \\\n",
218 | "* --train_batch_size= 8 \n",
219 | "\n",
220 | "### Times\n",
221 | "* Training: 1hr 11mins\n",
222 | "* Evaluation: 2.5hr\n",
223 | "\n",
224 | "### Results\n",
225 | "* Most accurate model on final step\n",
226 | "* Accuracy: 0.92416, eval_loss: 0.31708\n"
227 | ]
228 | },
229 | {
230 | "cell_type": "markdown",
231 | "metadata": {
232 | "colab_type": "text",
233 | "id": "XUW2avFM_fi_"
234 | },
235 | "source": [
236 | "### Model\n",
237 | "\n",
238 | "* The trained model checkpoints can be found in 'exp/imdb'\n",
239 | "\n"
240 | ]
241 | }
242 | ],
243 | "metadata": {
244 | "accelerator": "GPU",
245 | "colab": {
246 | "collapsed_sections": [],
247 | "include_colab_link": true,
248 | "name": "XLNet-imdb-GPU.ipynb",
249 | "provenance": [],
250 | "toc_visible": true,
251 | "version": "0.3.2"
252 | },
253 | "kernelspec": {
254 | "display_name": "Python 3",
255 | "language": "python",
256 | "name": "python3"
257 | }
258 | },
259 | "nbformat": 4,
260 | "nbformat_minor": 1
261 | }
262 |
--------------------------------------------------------------------------------
/prepro_utils.py:
--------------------------------------------------------------------------------
1 | # coding=utf-8
2 | from __future__ import absolute_import
3 | from __future__ import division
4 | from __future__ import print_function
5 |
6 | import unicodedata
7 | import six
8 | from functools import partial
9 |
10 |
11 | SPIECE_UNDERLINE = '▁'
12 |
13 |
14 | def printable_text(text):
15 | """Returns text encoded in a way suitable for print or `tf.logging`."""
16 |
17 | # These functions want `str` for both Python2 and Python3, but in one case
18 | # it's a Unicode string and in the other it's a byte string.
19 | if six.PY3:
20 | if isinstance(text, str):
21 | return text
22 | elif isinstance(text, bytes):
23 | return text.decode("utf-8", "ignore")
24 | else:
25 | raise ValueError("Unsupported string type: %s" % (type(text)))
26 | elif six.PY2:
27 | if isinstance(text, str):
28 | return text
29 | elif isinstance(text, unicode):
30 | return text.encode("utf-8")
31 | else:
32 | raise ValueError("Unsupported string type: %s" % (type(text)))
33 | else:
34 | raise ValueError("Not running on Python2 or Python 3?")
35 |
36 |
37 | def print_(*args):
38 | new_args = []
39 | for arg in args:
40 | if isinstance(arg, list):
41 | s = [printable_text(i) for i in arg]
42 | s = ' '.join(s)
43 | new_args.append(s)
44 | else:
45 | new_args.append(printable_text(arg))
46 | print(*new_args)
47 |
48 |
49 | def preprocess_text(inputs, lower=False, remove_space=True, keep_accents=False):
50 | if remove_space:
51 | outputs = ' '.join(inputs.strip().split())
52 | else:
53 | outputs = inputs
54 | outputs = outputs.replace("``", '"').replace("''", '"')
55 |
56 | if six.PY2 and isinstance(outputs, str):
57 | outputs = outputs.decode('utf-8')
58 |
59 | if not keep_accents:
60 | outputs = unicodedata.normalize('NFKD', outputs)
61 | outputs = ''.join([c for c in outputs if not unicodedata.combining(c)])
62 | if lower:
63 | outputs = outputs.lower()
64 |
65 | return outputs
66 |
67 |
68 | def encode_pieces(sp_model, text, return_unicode=True, sample=False):
69 | # return_unicode is used only for py2
70 |
71 | # note(zhiliny): in some systems, sentencepiece only accepts str for py2
72 | if six.PY2 and isinstance(text, unicode):
73 | text = text.encode('utf-8')
74 |
75 | if not sample:
76 | pieces = sp_model.EncodeAsPieces(text)
77 | else:
78 | pieces = sp_model.SampleEncodeAsPieces(text, 64, 0.1)
79 | new_pieces = []
80 | for piece in pieces:
81 | if len(piece) > 1 and piece[-1] == ',' and piece[-2].isdigit():
82 | cur_pieces = sp_model.EncodeAsPieces(
83 | piece[:-1].replace(SPIECE_UNDERLINE, ''))
84 | if piece[0] != SPIECE_UNDERLINE and cur_pieces[0][0] == SPIECE_UNDERLINE:
85 | if len(cur_pieces[0]) == 1:
86 | cur_pieces = cur_pieces[1:]
87 | else:
88 | cur_pieces[0] = cur_pieces[0][1:]
89 | cur_pieces.append(piece[-1])
90 | new_pieces.extend(cur_pieces)
91 | else:
92 | new_pieces.append(piece)
93 |
94 | # note(zhiliny): convert back to unicode for py2
95 | if six.PY2 and return_unicode:
96 | ret_pieces = []
97 | for piece in new_pieces:
98 | if isinstance(piece, str):
99 | piece = piece.decode('utf-8')
100 | ret_pieces.append(piece)
101 | new_pieces = ret_pieces
102 |
103 | return new_pieces
104 |
105 |
106 | def encode_ids(sp_model, text, sample=False):
107 | pieces = encode_pieces(sp_model, text, return_unicode=False, sample=sample)
108 | ids = [sp_model.PieceToId(piece) for piece in pieces]
109 | return ids
110 |
111 |
112 | if __name__ == '__main__':
113 | import sentencepiece as spm
114 |
115 | sp = spm.SentencePieceProcessor()
116 | sp.load('sp10m.uncased.v3.model')
117 |
118 | print_(u'I was born in 2000, and this is falsé.')
119 | print_(u'ORIGINAL', sp.EncodeAsPieces(u'I was born in 2000, and this is falsé.'))
120 | print_(u'OURS', encode_pieces(sp, u'I was born in 2000, and this is falsé.'))
121 | print(encode_ids(sp, u'I was born in 2000, and this is falsé.'))
122 | print_('')
123 | prepro_func = partial(preprocess_text, lower=True)
124 | print_(prepro_func('I was born in 2000, and this is falsé.'))
125 | print_('ORIGINAL', sp.EncodeAsPieces(prepro_func('I was born in 2000, and this is falsé.')))
126 | print_('OURS', encode_pieces(sp, prepro_func('I was born in 2000, and this is falsé.')))
127 | print(encode_ids(sp, prepro_func('I was born in 2000, and this is falsé.')))
128 | print_('')
129 | print_('I was born in 2000, and this is falsé.')
130 | print_('ORIGINAL', sp.EncodeAsPieces('I was born in 2000, and this is falsé.'))
131 | print_('OURS', encode_pieces(sp, 'I was born in 2000, and this is falsé.'))
132 | print(encode_ids(sp, 'I was born in 2000, and this is falsé.'))
133 | print_('')
134 | print_('I was born in 92000, and this is falsé.')
135 | print_('ORIGINAL', sp.EncodeAsPieces('I was born in 92000, and this is falsé.'))
136 | print_('OURS', encode_pieces(sp, 'I was born in 92000, and this is falsé.'))
137 | print(encode_ids(sp, 'I was born in 92000, and this is falsé.'))
138 |
139 |
--------------------------------------------------------------------------------
/run_classifier.py:
--------------------------------------------------------------------------------
1 | from __future__ import absolute_import
2 | from __future__ import division
3 | from __future__ import print_function
4 |
5 | from os.path import join
6 | from absl import flags
7 | import os
8 | import sys
9 | import csv
10 | import collections
11 | import numpy as np
12 | import time
13 | import math
14 | import json
15 | import random
16 | from copy import copy
17 | from collections import defaultdict as dd
18 |
19 | import absl.logging as _logging # pylint: disable=unused-import
20 | import tensorflow as tf
21 |
22 | import sentencepiece as spm
23 |
24 | from data_utils import SEP_ID, VOCAB_SIZE, CLS_ID
25 | import model_utils
26 | import function_builder
27 | from classifier_utils import PaddingInputExample
28 | from classifier_utils import convert_single_example
29 | from prepro_utils import preprocess_text, encode_ids
30 |
31 |
32 | # Model
33 | flags.DEFINE_string("model_config_path", default=None,
34 | help="Model config path.")
35 | flags.DEFINE_float("dropout", default=0.1,
36 | help="Dropout rate.")
37 | flags.DEFINE_float("dropatt", default=0.1,
38 | help="Attention dropout rate.")
39 | flags.DEFINE_integer("clamp_len", default=-1,
40 | help="Clamp length")
41 | flags.DEFINE_string("summary_type", default="last",
42 | help="Method used to summarize a sequence into a compact vector.")
43 | flags.DEFINE_bool("use_summ_proj", default=True,
44 | help="Whether to use projection for summarizing sequences.")
45 | flags.DEFINE_bool("use_bfloat16", False,
46 | help="Whether to use bfloat16.")
47 |
48 | # Parameter initialization
49 | flags.DEFINE_enum("init", default="normal",
50 | enum_values=["normal", "uniform"],
51 | help="Initialization method.")
52 | flags.DEFINE_float("init_std", default=0.02,
53 | help="Initialization std when init is normal.")
54 | flags.DEFINE_float("init_range", default=0.1,
55 | help="Initialization std when init is uniform.")
56 |
57 | # I/O paths
58 | flags.DEFINE_bool("overwrite_data", default=False,
59 | help="If False, will use cached data if available.")
60 | flags.DEFINE_string("init_checkpoint", default=None,
61 | help="checkpoint path for initializing the model. "
62 | "Could be a pretrained model or a finetuned model.")
63 | flags.DEFINE_string("output_dir", default="",
64 | help="Output dir for TF records.")
65 | flags.DEFINE_string("spiece_model_file", default="",
66 | help="Sentence Piece model path.")
67 | flags.DEFINE_string("model_dir", default="",
68 | help="Directory for saving the finetuned model.")
69 | flags.DEFINE_string("data_dir", default="",
70 | help="Directory for input data.")
71 |
72 | # TPUs and machines
73 | flags.DEFINE_bool("use_tpu", default=False, help="whether to use TPU.")
74 | flags.DEFINE_integer("num_hosts", default=1, help="How many TPU hosts.")
75 | flags.DEFINE_integer("num_core_per_host", default=8,
76 | help="8 for TPU v2 and v3-8, 16 for larger TPU v3 pod. In the context "
77 | "of GPU training, it refers to the number of GPUs used.")
78 | flags.DEFINE_string("tpu_job_name", default=None, help="TPU worker job name.")
79 | flags.DEFINE_string("tpu", default=None, help="TPU name.")
80 | flags.DEFINE_string("tpu_zone", default=None, help="TPU zone.")
81 | flags.DEFINE_string("gcp_project", default=None, help="gcp project.")
82 | flags.DEFINE_string("master", default=None, help="master")
83 | flags.DEFINE_integer("iterations", default=1000,
84 | help="number of iterations per TPU training loop.")
85 |
86 | # training
87 | flags.DEFINE_bool("do_train", default=False, help="whether to do training")
88 | flags.DEFINE_integer("train_steps", default=1000,
89 | help="Number of training steps")
90 | flags.DEFINE_integer("warmup_steps", default=0, help="number of warmup steps")
91 | flags.DEFINE_float("learning_rate", default=1e-5, help="initial learning rate")
92 | flags.DEFINE_float("lr_layer_decay_rate", 1.0,
93 | "Top layer: lr[L] = FLAGS.learning_rate."
94 | "Low layer: lr[l-1] = lr[l] * lr_layer_decay_rate.")
95 | flags.DEFINE_float("min_lr_ratio", default=0.0,
96 | help="min lr ratio for cos decay.")
97 | flags.DEFINE_float("clip", default=1.0, help="Gradient clipping")
98 | flags.DEFINE_integer("max_save", default=0,
99 | help="Max number of checkpoints to save. Use 0 to save all.")
100 | flags.DEFINE_integer("save_steps", default=None,
101 | help="Save the model for every save_steps. "
102 | "If None, not to save any model.")
103 | flags.DEFINE_integer("train_batch_size", default=8,
104 | help="Batch size for training")
105 | flags.DEFINE_float("weight_decay", default=0.00, help="Weight decay rate")
106 | flags.DEFINE_float("adam_epsilon", default=1e-8, help="Adam epsilon")
107 | flags.DEFINE_string("decay_method", default="poly", help="poly or cos")
108 |
109 | # evaluation
110 | flags.DEFINE_bool("do_eval", default=False, help="whether to do eval")
111 | flags.DEFINE_bool("do_predict", default=False, help="whether to do prediction")
112 | flags.DEFINE_float("predict_threshold", default=0,
113 | help="Threshold for binary prediction.")
114 | flags.DEFINE_string("eval_split", default="dev", help="could be dev or test")
115 | flags.DEFINE_integer("eval_batch_size", default=128,
116 | help="batch size for evaluation")
117 | flags.DEFINE_integer("predict_batch_size", default=128,
118 | help="batch size for prediction.")
119 | flags.DEFINE_string("predict_dir", default=None,
120 | help="Dir for saving prediction files.")
121 | flags.DEFINE_bool("eval_all_ckpt", default=False,
122 | help="Eval all ckpts. If False, only evaluate the last one.")
123 | flags.DEFINE_string("predict_ckpt", default=None,
124 | help="Ckpt path for do_predict. If None, use the last one.")
125 |
126 | # task specific
127 | flags.DEFINE_string("task_name", default=None, help="Task name")
128 | flags.DEFINE_integer("max_seq_length", default=128, help="Max sequence length")
129 | flags.DEFINE_integer("shuffle_buffer", default=2048,
130 | help="Buffer size used for shuffle.")
131 | flags.DEFINE_integer("num_passes", default=1,
132 | help="Num passes for processing training data. "
133 | "This is use to batch data without loss for TPUs.")
134 | flags.DEFINE_bool("uncased", default=False,
135 | help="Use uncased.")
136 | flags.DEFINE_string("cls_scope", default=None,
137 | help="Classifier layer scope.")
138 | flags.DEFINE_bool("is_regression", default=False,
139 | help="Whether it's a regression task.")
140 |
141 | FLAGS = flags.FLAGS
142 |
143 |
144 | class InputExample(object):
145 | """A single training/test example for simple sequence classification."""
146 |
147 | def __init__(self, guid, text_a, text_b=None, label=None):
148 | """Constructs a InputExample.
149 | Args:
150 | guid: Unique id for the example.
151 | text_a: string. The untokenized text of the first sequence. For single
152 | sequence tasks, only this sequence must be specified.
153 | text_b: (Optional) string. The untokenized text of the second sequence.
154 | Only must be specified for sequence pair tasks.
155 | label: (Optional) string. The label of the example. This should be
156 | specified for train and dev examples, but not for test examples.
157 | """
158 | self.guid = guid
159 | self.text_a = text_a
160 | self.text_b = text_b
161 | self.label = label
162 |
163 |
164 | class DataProcessor(object):
165 | """Base class for data converters for sequence classification data sets."""
166 |
167 | def get_train_examples(self, data_dir):
168 | """Gets a collection of `InputExample`s for the train set."""
169 | raise NotImplementedError()
170 |
171 | def get_dev_examples(self, data_dir):
172 | """Gets a collection of `InputExample`s for the dev set."""
173 | raise NotImplementedError()
174 |
175 | def get_test_examples(self, data_dir):
176 | """Gets a collection of `InputExample`s for prediction."""
177 | raise NotImplementedError()
178 |
179 | def get_labels(self):
180 | """Gets the list of labels for this data set."""
181 | raise NotImplementedError()
182 |
183 | @classmethod
184 | def _read_tsv(cls, input_file, quotechar=None):
185 | """Reads a tab separated value file."""
186 | with tf.gfile.Open(input_file, "r") as f:
187 | reader = csv.reader(f, delimiter="\t", quotechar=quotechar)
188 | lines = []
189 | for line in reader:
190 | if len(line) == 0: continue
191 | lines.append(line)
192 | return lines
193 |
194 |
195 | class GLUEProcessor(DataProcessor):
196 | def __init__(self):
197 | self.train_file = "train.tsv"
198 | self.dev_file = "dev.tsv"
199 | self.test_file = "test.tsv"
200 | self.label_column = None
201 | self.text_a_column = None
202 | self.text_b_column = None
203 | self.contains_header = True
204 | self.test_text_a_column = None
205 | self.test_text_b_column = None
206 | self.test_contains_header = True
207 |
208 | def get_train_examples(self, data_dir):
209 | """See base class."""
210 | return self._create_examples(
211 | self._read_tsv(os.path.join(data_dir, self.train_file)), "train")
212 |
213 | def get_dev_examples(self, data_dir):
214 | """See base class."""
215 | return self._create_examples(
216 | self._read_tsv(os.path.join(data_dir, self.dev_file)), "dev")
217 |
218 | def get_test_examples(self, data_dir):
219 | """See base class."""
220 | if self.test_text_a_column is None:
221 | self.test_text_a_column = self.text_a_column
222 | if self.test_text_b_column is None:
223 | self.test_text_b_column = self.text_b_column
224 |
225 | return self._create_examples(
226 | self._read_tsv(os.path.join(data_dir, self.test_file)), "test")
227 |
228 | def get_labels(self):
229 | """See base class."""
230 | return ["0", "1"]
231 |
232 | def _create_examples(self, lines, set_type):
233 | """Creates examples for the training and dev sets."""
234 | examples = []
235 | for (i, line) in enumerate(lines):
236 | if i == 0 and self.contains_header and set_type != "test":
237 | continue
238 | if i == 0 and self.test_contains_header and set_type == "test":
239 | continue
240 | guid = "%s-%s" % (set_type, i)
241 |
242 | a_column = (self.text_a_column if set_type != "test" else
243 | self.test_text_a_column)
244 | b_column = (self.text_b_column if set_type != "test" else
245 | self.test_text_b_column)
246 |
247 | # there are some incomplete lines in QNLI
248 | if len(line) <= a_column:
249 | tf.logging.warning('Incomplete line, ignored.')
250 | continue
251 | text_a = line[a_column]
252 |
253 | if b_column is not None:
254 | if len(line) <= b_column:
255 | tf.logging.warning('Incomplete line, ignored.')
256 | continue
257 | text_b = line[b_column]
258 | else:
259 | text_b = None
260 |
261 | if set_type == "test":
262 | label = self.get_labels()[0]
263 | else:
264 | if len(line) <= self.label_column:
265 | tf.logging.warning('Incomplete line, ignored.')
266 | continue
267 | label = line[self.label_column]
268 | examples.append(
269 | InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
270 | return examples
271 |
272 |
273 | class Yelp5Processor(DataProcessor):
274 | def get_train_examples(self, data_dir):
275 | return self._create_examples(os.path.join(data_dir, "train.csv"))
276 |
277 | def get_dev_examples(self, data_dir):
278 | return self._create_examples(os.path.join(data_dir, "test.csv"))
279 |
280 | def get_labels(self):
281 | """See base class."""
282 | return ["1", "2", "3", "4", "5"]
283 |
284 | def _create_examples(self, input_file):
285 | """Creates examples for the training and dev sets."""
286 | examples = []
287 | with tf.gfile.Open(input_file) as f:
288 | reader = csv.reader(f)
289 | for i, line in enumerate(reader):
290 |
291 | label = line[0]
292 | text_a = line[1].replace('""', '"').replace('\\"', '"')
293 | examples.append(
294 | InputExample(guid=str(i), text_a=text_a, text_b=None, label=label))
295 | return examples
296 |
297 |
298 | class ImdbProcessor(DataProcessor):
299 | def get_labels(self):
300 | return ["neg", "pos"]
301 |
302 | def get_train_examples(self, data_dir):
303 | return self._create_examples(os.path.join(data_dir, "train"))
304 |
305 | def get_dev_examples(self, data_dir):
306 | return self._create_examples(os.path.join(data_dir, "test"))
307 |
308 | def _create_examples(self, data_dir):
309 | examples = []
310 | for label in ["neg", "pos"]:
311 | cur_dir = os.path.join(data_dir, label)
312 | for filename in tf.gfile.ListDirectory(cur_dir):
313 | if not filename.endswith("txt"): continue
314 |
315 | path = os.path.join(cur_dir, filename)
316 | with tf.gfile.Open(path) as f:
317 | text = f.read().strip().replace("", " ") 318 | examples.append(InputExample( 319 | guid="unused_id", text_a=text, text_b=None, label=label)) 320 | return examples 321 | 322 | 323 | class MnliMatchedProcessor(GLUEProcessor): 324 | def __init__(self): 325 | super(MnliMatchedProcessor, self).__init__() 326 | self.dev_file = "dev_matched.tsv" 327 | self.test_file = "test_matched.tsv" 328 | self.label_column = -1 329 | self.text_a_column = 8 330 | self.text_b_column = 9 331 | 332 | def get_labels(self): 333 | return ["contradiction", "entailment", "neutral"] 334 | 335 | 336 | class MnliMismatchedProcessor(MnliMatchedProcessor): 337 | def __init__(self): 338 | super(MnliMismatchedProcessor, self).__init__() 339 | self.dev_file = "dev_mismatched.tsv" 340 | self.test_file = "test_mismatched.tsv" 341 | 342 | 343 | class StsbProcessor(GLUEProcessor): 344 | def __init__(self): 345 | super(StsbProcessor, self).__init__() 346 | self.label_column = 9 347 | self.text_a_column = 7 348 | self.text_b_column = 8 349 | 350 | def get_labels(self): 351 | return [0.0] 352 | 353 | def _create_examples(self, lines, set_type): 354 | """Creates examples for the training and dev sets.""" 355 | examples = [] 356 | for (i, line) in enumerate(lines): 357 | if i == 0 and self.contains_header and set_type != "test": 358 | continue 359 | if i == 0 and self.test_contains_header and set_type == "test": 360 | continue 361 | guid = "%s-%s" % (set_type, i) 362 | 363 | a_column = (self.text_a_column if set_type != "test" else 364 | self.test_text_a_column) 365 | b_column = (self.text_b_column if set_type != "test" else 366 | self.test_text_b_column) 367 | 368 | # there are some incomplete lines in QNLI 369 | if len(line) <= a_column: 370 | tf.logging.warning('Incomplete line, ignored.') 371 | continue 372 | text_a = line[a_column] 373 | 374 | if b_column is not None: 375 | if len(line) <= b_column: 376 | tf.logging.warning('Incomplete line, ignored.') 377 | continue 378 | text_b = line[b_column] 379 | else: 380 | text_b = None 381 | 382 | if set_type == "test": 383 | label = self.get_labels()[0] 384 | else: 385 | if len(line) <= self.label_column: 386 | tf.logging.warning('Incomplete line, ignored.') 387 | continue 388 | label = float(line[self.label_column]) 389 | examples.append( 390 | InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label)) 391 | 392 | return examples 393 | 394 | 395 | def file_based_convert_examples_to_features( 396 | examples, label_list, max_seq_length, tokenize_fn, output_file, 397 | num_passes=1): 398 | """Convert a set of `InputExample`s to a TFRecord file.""" 399 | 400 | # do not create duplicated records 401 | if tf.gfile.Exists(output_file) and not FLAGS.overwrite_data: 402 | tf.logging.info("Do not overwrite tfrecord {} exists.".format(output_file)) 403 | return 404 | 405 | tf.logging.info("Create new tfrecord {}.".format(output_file)) 406 | 407 | writer = tf.python_io.TFRecordWriter(output_file) 408 | 409 | if num_passes > 1: 410 | examples *= num_passes 411 | 412 | for (ex_index, example) in enumerate(examples): 413 | if ex_index % 10000 == 0: 414 | tf.logging.info("Writing example {} of {}".format(ex_index, 415 | len(examples))) 416 | 417 | feature = convert_single_example(ex_index, example, label_list, 418 | max_seq_length, tokenize_fn) 419 | 420 | def create_int_feature(values): 421 | f = tf.train.Feature(int64_list=tf.train.Int64List(value=list(values))) 422 | return f 423 | 424 | def create_float_feature(values): 425 | f = tf.train.Feature(float_list=tf.train.FloatList(value=list(values))) 426 | return f 427 | 428 | features = collections.OrderedDict() 429 | features["input_ids"] = create_int_feature(feature.input_ids) 430 | features["input_mask"] = create_float_feature(feature.input_mask) 431 | features["segment_ids"] = create_int_feature(feature.segment_ids) 432 | if label_list is not None: 433 | features["label_ids"] = create_int_feature([feature.label_id]) 434 | else: 435 | features["label_ids"] = create_float_feature([float(feature.label_id)]) 436 | features["is_real_example"] = create_int_feature( 437 | [int(feature.is_real_example)]) 438 | 439 | tf_example = tf.train.Example(features=tf.train.Features(feature=features)) 440 | writer.write(tf_example.SerializeToString()) 441 | writer.close() 442 | 443 | 444 | def file_based_input_fn_builder(input_file, seq_length, is_training, 445 | drop_remainder): 446 | """Creates an `input_fn` closure to be passed to TPUEstimator.""" 447 | 448 | 449 | name_to_features = { 450 | "input_ids": tf.FixedLenFeature([seq_length], tf.int64), 451 | "input_mask": tf.FixedLenFeature([seq_length], tf.float32), 452 | "segment_ids": tf.FixedLenFeature([seq_length], tf.int64), 453 | "label_ids": tf.FixedLenFeature([], tf.int64), 454 | "is_real_example": tf.FixedLenFeature([], tf.int64), 455 | } 456 | if FLAGS.is_regression: 457 | name_to_features["label_ids"] = tf.FixedLenFeature([], tf.float32) 458 | 459 | tf.logging.info("Input tfrecord file {}".format(input_file)) 460 | 461 | def _decode_record(record, name_to_features): 462 | """Decodes a record to a TensorFlow example.""" 463 | example = tf.parse_single_example(record, name_to_features) 464 | 465 | # tf.Example only supports tf.int64, but the TPU only supports tf.int32. 466 | # So cast all int64 to int32. 467 | for name in list(example.keys()): 468 | t = example[name] 469 | if t.dtype == tf.int64: 470 | t = tf.cast(t, tf.int32) 471 | example[name] = t 472 | 473 | return example 474 | 475 | def input_fn(params, input_context=None): 476 | """The actual input function.""" 477 | if FLAGS.use_tpu: 478 | batch_size = params["batch_size"] 479 | elif is_training: 480 | batch_size = FLAGS.train_batch_size 481 | elif FLAGS.do_eval: 482 | batch_size = FLAGS.eval_batch_size 483 | else: 484 | batch_size = FLAGS.predict_batch_size 485 | 486 | d = tf.data.TFRecordDataset(input_file) 487 | # Shard the dataset to difference devices 488 | if input_context is not None: 489 | tf.logging.info("Input pipeline id %d out of %d", 490 | input_context.input_pipeline_id, input_context.num_replicas_in_sync) 491 | d = d.shard(input_context.num_input_pipelines, 492 | input_context.input_pipeline_id) 493 | 494 | # For training, we want a lot of parallel reading and shuffling. 495 | # For eval, we want no shuffling and parallel reading doesn't matter. 496 | if is_training: 497 | d = d.shuffle(buffer_size=FLAGS.shuffle_buffer) 498 | d = d.repeat() 499 | 500 | d = d.apply( 501 | tf.contrib.data.map_and_batch( 502 | lambda record: _decode_record(record, name_to_features), 503 | batch_size=batch_size, 504 | drop_remainder=drop_remainder)) 505 | 506 | return d 507 | 508 | return input_fn 509 | 510 | 511 | def get_model_fn(n_class): 512 | def model_fn(features, labels, mode, params): 513 | #### Training or Evaluation 514 | is_training = (mode == tf.estimator.ModeKeys.TRAIN) 515 | 516 | #### Get loss from inputs 517 | if FLAGS.is_regression: 518 | (total_loss, per_example_loss, logits 519 | ) = function_builder.get_regression_loss(FLAGS, features, is_training) 520 | else: 521 | (total_loss, per_example_loss, logits 522 | ) = function_builder.get_classification_loss( 523 | FLAGS, features, n_class, is_training) 524 | 525 | #### Check model parameters 526 | num_params = sum([np.prod(v.shape) for v in tf.trainable_variables()]) 527 | tf.logging.info('#params: {}'.format(num_params)) 528 | 529 | #### load pretrained models 530 | scaffold_fn = model_utils.init_from_checkpoint(FLAGS) 531 | 532 | #### Evaluation mode 533 | if mode == tf.estimator.ModeKeys.EVAL: 534 | assert FLAGS.num_hosts == 1 535 | 536 | def metric_fn(per_example_loss, label_ids, logits, is_real_example): 537 | predictions = tf.argmax(logits, axis=-1, output_type=tf.int32) 538 | eval_input_dict = { 539 | 'labels': label_ids, 540 | 'predictions': predictions, 541 | 'weights': is_real_example 542 | } 543 | accuracy = tf.metrics.accuracy(**eval_input_dict) 544 | 545 | loss = tf.metrics.mean(values=per_example_loss, weights=is_real_example) 546 | return { 547 | 'eval_accuracy': accuracy, 548 | 'eval_loss': loss} 549 | 550 | def regression_metric_fn( 551 | per_example_loss, label_ids, logits, is_real_example): 552 | loss = tf.metrics.mean(values=per_example_loss, weights=is_real_example) 553 | pearsonr = tf.contrib.metrics.streaming_pearson_correlation( 554 | logits, label_ids, weights=is_real_example) 555 | return {'eval_loss': loss, 'eval_pearsonr': pearsonr} 556 | 557 | is_real_example = tf.cast(features["is_real_example"], dtype=tf.float32) 558 | 559 | #### Constucting evaluation TPUEstimatorSpec with new cache. 560 | label_ids = tf.reshape(features['label_ids'], [-1]) 561 | 562 | if FLAGS.is_regression: 563 | metric_fn = regression_metric_fn 564 | else: 565 | metric_fn = metric_fn 566 | metric_args = [per_example_loss, label_ids, logits, is_real_example] 567 | 568 | if FLAGS.use_tpu: 569 | eval_spec = tf.contrib.tpu.TPUEstimatorSpec( 570 | mode=mode, 571 | loss=total_loss, 572 | eval_metrics=(metric_fn, metric_args), 573 | scaffold_fn=scaffold_fn) 574 | else: 575 | eval_spec = tf.estimator.EstimatorSpec( 576 | mode=mode, 577 | loss=total_loss, 578 | eval_metric_ops=metric_fn(*metric_args)) 579 | 580 | return eval_spec 581 | 582 | elif mode == tf.estimator.ModeKeys.PREDICT: 583 | label_ids = tf.reshape(features["label_ids"], [-1]) 584 | 585 | predictions = { 586 | "logits": logits, 587 | "labels": label_ids, 588 | "is_real": features["is_real_example"] 589 | } 590 | 591 | if FLAGS.use_tpu: 592 | output_spec = tf.contrib.tpu.TPUEstimatorSpec( 593 | mode=mode, predictions=predictions, scaffold_fn=scaffold_fn) 594 | else: 595 | output_spec = tf.estimator.EstimatorSpec( 596 | mode=mode, predictions=predictions) 597 | return output_spec 598 | 599 | #### Configuring the optimizer 600 | train_op, learning_rate, _ = model_utils.get_train_op(FLAGS, total_loss) 601 | 602 | monitor_dict = {} 603 | monitor_dict["lr"] = learning_rate 604 | 605 | #### Constucting training TPUEstimatorSpec with new cache. 606 | if FLAGS.use_tpu: 607 | #### Creating host calls 608 | if not FLAGS.is_regression: 609 | label_ids = tf.reshape(features['label_ids'], [-1]) 610 | predictions = tf.argmax(logits, axis=-1, output_type=label_ids.dtype) 611 | is_correct = tf.equal(predictions, label_ids) 612 | accuracy = tf.reduce_mean(tf.cast(is_correct, tf.float32)) 613 | 614 | monitor_dict["accuracy"] = accuracy 615 | 616 | host_call = function_builder.construct_scalar_host_call( 617 | monitor_dict=monitor_dict, 618 | model_dir=FLAGS.model_dir, 619 | prefix="train/", 620 | reduce_fn=tf.reduce_mean) 621 | else: 622 | host_call = None 623 | 624 | train_spec = tf.contrib.tpu.TPUEstimatorSpec( 625 | mode=mode, loss=total_loss, train_op=train_op, host_call=host_call, 626 | scaffold_fn=scaffold_fn) 627 | else: 628 | train_spec = tf.estimator.EstimatorSpec( 629 | mode=mode, loss=total_loss, train_op=train_op) 630 | 631 | return train_spec 632 | 633 | return model_fn 634 | 635 | 636 | def main(_): 637 | tf.logging.set_verbosity(tf.logging.INFO) 638 | 639 | #### Validate flags 640 | if FLAGS.save_steps is not None: 641 | FLAGS.iterations = min(FLAGS.iterations, FLAGS.save_steps) 642 | 643 | if FLAGS.do_predict: 644 | predict_dir = FLAGS.predict_dir 645 | if not tf.gfile.Exists(predict_dir): 646 | tf.gfile.MakeDirs(predict_dir) 647 | 648 | processors = { 649 | "mnli_matched": MnliMatchedProcessor, 650 | "mnli_mismatched": MnliMismatchedProcessor, 651 | 'sts-b': StsbProcessor, 652 | 'imdb': ImdbProcessor, 653 | "yelp5": Yelp5Processor 654 | } 655 | 656 | if not FLAGS.do_train and not FLAGS.do_eval and not FLAGS.do_predict: 657 | raise ValueError( 658 | "At least one of `do_train`, `do_eval, `do_predict` or " 659 | "`do_submit` must be True.") 660 | 661 | if not tf.gfile.Exists(FLAGS.output_dir): 662 | tf.gfile.MakeDirs(FLAGS.output_dir) 663 | 664 | task_name = FLAGS.task_name.lower() 665 | 666 | if task_name not in processors: 667 | raise ValueError("Task not found: %s" % (task_name)) 668 | 669 | processor = processors[task_name]() 670 | label_list = processor.get_labels() if not FLAGS.is_regression else None 671 | 672 | sp = spm.SentencePieceProcessor() 673 | sp.Load(FLAGS.spiece_model_file) 674 | def tokenize_fn(text): 675 | text = preprocess_text(text, lower=FLAGS.uncased) 676 | return encode_ids(sp, text) 677 | 678 | run_config = model_utils.configure_tpu(FLAGS) 679 | 680 | model_fn = get_model_fn(len(label_list) if label_list is not None else None) 681 | 682 | spm_basename = os.path.basename(FLAGS.spiece_model_file) 683 | 684 | # If TPU is not available, this will fall back to normal Estimator on CPU 685 | # or GPU. 686 | if FLAGS.use_tpu: 687 | estimator = tf.contrib.tpu.TPUEstimator( 688 | use_tpu=FLAGS.use_tpu, 689 | model_fn=model_fn, 690 | config=run_config, 691 | train_batch_size=FLAGS.train_batch_size, 692 | predict_batch_size=FLAGS.predict_batch_size, 693 | eval_batch_size=FLAGS.eval_batch_size) 694 | else: 695 | estimator = tf.estimator.Estimator( 696 | model_fn=model_fn, 697 | config=run_config) 698 | 699 | if FLAGS.do_train: 700 | train_file_base = "{}.len-{}.train.tf_record".format( 701 | spm_basename, FLAGS.max_seq_length) 702 | train_file = os.path.join(FLAGS.output_dir, train_file_base) 703 | tf.logging.info("Use tfrecord file {}".format(train_file)) 704 | 705 | train_examples = processor.get_train_examples(FLAGS.data_dir) 706 | np.random.shuffle(train_examples) 707 | tf.logging.info("Num of train samples: {}".format(len(train_examples))) 708 | 709 | file_based_convert_examples_to_features( 710 | train_examples, label_list, FLAGS.max_seq_length, tokenize_fn, 711 | train_file, FLAGS.num_passes) 712 | 713 | train_input_fn = file_based_input_fn_builder( 714 | input_file=train_file, 715 | seq_length=FLAGS.max_seq_length, 716 | is_training=True, 717 | drop_remainder=True) 718 | 719 | estimator.train(input_fn=train_input_fn, max_steps=FLAGS.train_steps) 720 | 721 | if FLAGS.do_eval or FLAGS.do_predict: 722 | if FLAGS.eval_split == "dev": 723 | eval_examples = processor.get_dev_examples(FLAGS.data_dir) 724 | else: 725 | eval_examples = processor.get_test_examples(FLAGS.data_dir) 726 | 727 | tf.logging.info("Num of eval samples: {}".format(len(eval_examples))) 728 | 729 | if FLAGS.do_eval: 730 | # TPU requires a fixed batch size for all batches, therefore the number 731 | # of examples must be a multiple of the batch size, or else examples 732 | # will get dropped. So we pad with fake examples which are ignored 733 | # later on. These do NOT count towards the metric (all tf.metrics 734 | # support a per-instance weight, and these get a weight of 0.0). 735 | # 736 | # Modified in XL: We also adopt the same mechanism for GPUs. 737 | while len(eval_examples) % FLAGS.eval_batch_size != 0: 738 | eval_examples.append(PaddingInputExample()) 739 | 740 | eval_file_base = "{}.len-{}.{}.eval.tf_record".format( 741 | spm_basename, FLAGS.max_seq_length, FLAGS.eval_split) 742 | eval_file = os.path.join(FLAGS.output_dir, eval_file_base) 743 | 744 | file_based_convert_examples_to_features( 745 | eval_examples, label_list, FLAGS.max_seq_length, tokenize_fn, 746 | eval_file) 747 | 748 | assert len(eval_examples) % FLAGS.eval_batch_size == 0 749 | eval_steps = int(len(eval_examples) // FLAGS.eval_batch_size) 750 | 751 | eval_input_fn = file_based_input_fn_builder( 752 | input_file=eval_file, 753 | seq_length=FLAGS.max_seq_length, 754 | is_training=False, 755 | drop_remainder=True) 756 | 757 | # Filter out all checkpoints in the directory 758 | steps_and_files = [] 759 | filenames = tf.gfile.ListDirectory(FLAGS.model_dir) 760 | 761 | for filename in filenames: 762 | if filename.endswith(".index"): 763 | ckpt_name = filename[:-6] 764 | cur_filename = join(FLAGS.model_dir, ckpt_name) 765 | global_step = int(cur_filename.split("-")[-1]) 766 | tf.logging.info("Add {} to eval list.".format(cur_filename)) 767 | steps_and_files.append([global_step, cur_filename]) 768 | steps_and_files = sorted(steps_and_files, key=lambda x: x[0]) 769 | 770 | # Decide whether to evaluate all ckpts 771 | if not FLAGS.eval_all_ckpt: 772 | steps_and_files = steps_and_files[-1:] 773 | 774 | eval_results = [] 775 | for global_step, filename in sorted(steps_and_files, key=lambda x: x[0]): 776 | ret = estimator.evaluate( 777 | input_fn=eval_input_fn, 778 | steps=eval_steps, 779 | checkpoint_path=filename) 780 | 781 | ret["step"] = global_step 782 | ret["path"] = filename 783 | 784 | eval_results.append(ret) 785 | 786 | tf.logging.info("=" * 80) 787 | log_str = "Eval result | " 788 | for key, val in sorted(ret.items(), key=lambda x: x[0]): 789 | log_str += "{} {} | ".format(key, val) 790 | tf.logging.info(log_str) 791 | 792 | key_name = "eval_pearsonr" if FLAGS.is_regression else "eval_accuracy" 793 | eval_results.sort(key=lambda x: x[key_name], reverse=True) 794 | 795 | tf.logging.info("=" * 80) 796 | log_str = "Best result | " 797 | for key, val in sorted(eval_results[0].items(), key=lambda x: x[0]): 798 | log_str += "{} {} | ".format(key, val) 799 | tf.logging.info(log_str) 800 | 801 | if FLAGS.do_predict: 802 | eval_file_base = "{}.len-{}.{}.predict.tf_record".format( 803 | spm_basename, FLAGS.max_seq_length, FLAGS.eval_split) 804 | eval_file = os.path.join(FLAGS.output_dir, eval_file_base) 805 | 806 | file_based_convert_examples_to_features( 807 | eval_examples, label_list, FLAGS.max_seq_length, tokenize_fn, 808 | eval_file) 809 | 810 | pred_input_fn = file_based_input_fn_builder( 811 | input_file=eval_file, 812 | seq_length=FLAGS.max_seq_length, 813 | is_training=False, 814 | drop_remainder=False) 815 | 816 | predict_results = [] 817 | with tf.gfile.Open(os.path.join(predict_dir, "{}.tsv".format( 818 | task_name)), "w") as fout: 819 | fout.write("index\tprediction\n") 820 | 821 | for pred_cnt, result in enumerate(estimator.predict( 822 | input_fn=pred_input_fn, 823 | yield_single_examples=True, 824 | checkpoint_path=FLAGS.predict_ckpt)): 825 | if pred_cnt % 1000 == 0: 826 | tf.logging.info("Predicting submission for example: {}".format( 827 | pred_cnt)) 828 | 829 | logits = [float(x) for x in result["logits"].flat] 830 | predict_results.append(logits) 831 | 832 | if len(logits) == 1: 833 | label_out = logits[0] 834 | elif len(logits) == 2: 835 | if logits[1] - logits[0] > FLAGS.predict_threshold: 836 | label_out = label_list[1] 837 | else: 838 | label_out = label_list[0] 839 | elif len(logits) > 2: 840 | max_index = np.argmax(np.array(logits, dtype=np.float32)) 841 | label_out = label_list[max_index] 842 | else: 843 | raise NotImplementedError 844 | 845 | fout.write("{}\t{}\n".format(pred_cnt, label_out)) 846 | 847 | predict_json_path = os.path.join(predict_dir, "{}.logits.json".format( 848 | task_name)) 849 | 850 | with tf.gfile.Open(predict_json_path, "w") as fp: 851 | json.dump(predict_results, fp, indent=4) 852 | 853 | 854 | if __name__ == "__main__": 855 | tf.app.run() 856 | -------------------------------------------------------------------------------- /run_race.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | 5 | from os.path import join 6 | from absl import flags 7 | import os 8 | import csv 9 | import collections 10 | import numpy as np 11 | import time 12 | import math 13 | import json 14 | import random 15 | from copy import copy 16 | from collections import defaultdict as dd 17 | 18 | from scipy.stats import pearsonr, spearmanr 19 | from sklearn.metrics import matthews_corrcoef, f1_score 20 | 21 | import absl.logging as _logging # pylint: disable=unused-import 22 | 23 | import tensorflow as tf 24 | import sentencepiece as spm 25 | 26 | from data_utils import SEP_ID, VOCAB_SIZE, CLS_ID 27 | import model_utils 28 | import function_builder 29 | from classifier_utils import PaddingInputExample 30 | from classifier_utils import convert_single_example 31 | from prepro_utils import preprocess_text, encode_ids 32 | 33 | # Model 34 | flags.DEFINE_string("model_config_path", default=None, 35 | help="Model config path.") 36 | flags.DEFINE_float("dropout", default=0.1, 37 | help="Dropout rate.") 38 | flags.DEFINE_float("dropatt", default=0.1, 39 | help="Attention dropout rate.") 40 | flags.DEFINE_integer("clamp_len", default=-1, 41 | help="Clamp length") 42 | flags.DEFINE_string("summary_type", default="last", 43 | help="Method used to summarize a sequence into a compact vector.") 44 | flags.DEFINE_bool("use_summ_proj", default=True, 45 | help="Whether to use projection for summarizing sequences.") 46 | flags.DEFINE_bool("use_bfloat16", default=False, 47 | help="Whether to use bfloat16.") 48 | 49 | # Parameter initialization 50 | flags.DEFINE_enum("init", default="normal", 51 | enum_values=["normal", "uniform"], 52 | help="Initialization method.") 53 | flags.DEFINE_float("init_std", default=0.02, 54 | help="Initialization std when init is normal.") 55 | flags.DEFINE_float("init_range", default=0.1, 56 | help="Initialization std when init is uniform.") 57 | 58 | # I/O paths 59 | flags.DEFINE_bool("overwrite_data", default=False, 60 | help="If False, will use cached data if available.") 61 | flags.DEFINE_string("init_checkpoint", default=None, 62 | help="checkpoint path for initializing the model. " 63 | "Could be a pretrained model or a finetuned model.") 64 | flags.DEFINE_string("output_dir", default="", 65 | help="Output dir for TF records.") 66 | flags.DEFINE_string("spiece_model_file", default="", 67 | help="Sentence Piece model path.") 68 | flags.DEFINE_string("model_dir", default="", 69 | help="Directory for saving the finetuned model.") 70 | flags.DEFINE_string("data_dir", default="", 71 | help="Directory for input data.") 72 | 73 | # TPUs and machines 74 | flags.DEFINE_bool("use_tpu", default=False, help="whether to use TPU.") 75 | flags.DEFINE_integer("num_hosts", default=1, help="How many TPU hosts.") 76 | flags.DEFINE_integer("num_core_per_host", default=8, 77 | help="8 for TPU v2 and v3-8, 16 for larger TPU v3 pod. In the context " 78 | "of GPU training, it refers to the number of GPUs used.") 79 | flags.DEFINE_string("tpu_job_name", default=None, help="TPU worker job name.") 80 | flags.DEFINE_string("tpu", default=None, help="TPU name.") 81 | flags.DEFINE_string("tpu_zone", default=None, help="TPU zone.") 82 | flags.DEFINE_string("gcp_project", default=None, help="gcp project.") 83 | flags.DEFINE_string("master", default=None, help="master") 84 | flags.DEFINE_integer("iterations", default=1000, 85 | help="number of iterations per TPU training loop.") 86 | 87 | # Training 88 | flags.DEFINE_bool("do_train", default=False, help="whether to do training") 89 | flags.DEFINE_integer("train_steps", default=12000, 90 | help="Number of training steps") 91 | flags.DEFINE_integer("warmup_steps", default=0, help="number of warmup steps") 92 | flags.DEFINE_float("learning_rate", default=2e-5, help="initial learning rate") 93 | flags.DEFINE_float("lr_layer_decay_rate", 1.0, 94 | "Top layer: lr[L] = FLAGS.learning_rate." 95 | "Low layer: lr[l-1] = lr[l] * lr_layer_decay_rate.") 96 | flags.DEFINE_float("min_lr_ratio", default=0.0, 97 | help="min lr ratio for cos decay.") 98 | flags.DEFINE_float("clip", default=1.0, help="Gradient clipping") 99 | flags.DEFINE_integer("max_save", default=0, 100 | help="Max number of checkpoints to save. Use 0 to save all.") 101 | flags.DEFINE_integer("save_steps", default=None, 102 | help="Save the model for every save_steps. " 103 | "If None, not to save any model.") 104 | flags.DEFINE_integer("train_batch_size", default=8, 105 | help="Batch size for training. Note that batch size 1 corresponds to " 106 | "4 sequences: one paragraph + one quesetion + 4 candidate answers.") 107 | flags.DEFINE_float("weight_decay", default=0.00, help="weight decay rate") 108 | flags.DEFINE_float("adam_epsilon", default=1e-6, help="adam epsilon") 109 | flags.DEFINE_string("decay_method", default="poly", help="poly or cos") 110 | 111 | # Evaluation 112 | flags.DEFINE_bool("do_eval", default=False, help="whether to do eval") 113 | flags.DEFINE_string("eval_split", default="dev", 114 | help="could be dev or test") 115 | flags.DEFINE_integer("eval_batch_size", default=32, 116 | help="Batch size for evaluation.") 117 | 118 | # Data config 119 | flags.DEFINE_integer("max_seq_length", default=512, 120 | help="Max length for the paragraph.") 121 | flags.DEFINE_integer("max_qa_length", default=128, 122 | help="Max length for the concatenated question and answer.") 123 | flags.DEFINE_integer("shuffle_buffer", default=2048, 124 | help="Buffer size used for shuffle.") 125 | flags.DEFINE_bool("uncased", default=False, 126 | help="Use uncased.") 127 | flags.DEFINE_bool("high_only", default=False, 128 | help="Evaluate on high school only.") 129 | flags.DEFINE_bool("middle_only", default=False, 130 | help="Evaluate on middle school only.") 131 | 132 | FLAGS = flags.FLAGS 133 | 134 | SEG_ID_A = 0 135 | SEG_ID_B = 1 136 | SEG_ID_CLS = 2 137 | SEG_ID_SEP = 3 138 | SEG_ID_PAD = 4 139 | 140 | 141 | class PaddingInputExample(object): 142 | """Fake example so the num input examples is a multiple of the batch size. 143 | When running eval/predict on the TPU, we need to pad the number of examples 144 | to be a multiple of the batch size, because the TPU requires a fixed batch 145 | size. The alternative is to drop the last batch, which is bad because it means 146 | the entire output data won't be generated. 147 | We use this class instead of `None` because treating `None` as padding 148 | battches could cause silent errors. 149 | """ 150 | 151 | 152 | class InputFeatures(object): 153 | """A single set of features of data.""" 154 | 155 | def __init__(self, 156 | input_ids, 157 | input_mask, 158 | segment_ids, 159 | label_id, 160 | is_real_example=True): 161 | self.input_ids = input_ids 162 | self.input_mask = input_mask 163 | self.segment_ids = segment_ids 164 | self.label_id = label_id 165 | self.is_real_example = is_real_example 166 | 167 | 168 | def convert_single_example(example, tokenize_fn): 169 | """Converts a single `InputExample` into a single `InputFeatures`.""" 170 | 171 | if isinstance(example, PaddingInputExample): 172 | return InputFeatures( 173 | input_ids=[0] * FLAGS.max_seq_length * 4, 174 | input_mask=[1] * FLAGS.max_seq_length * 4, 175 | segment_ids=[0] * FLAGS.max_seq_length * 4, 176 | label_id=0, 177 | is_real_example=False) 178 | 179 | input_ids, input_mask, all_seg_ids = [], [], [] 180 | tokens_context = tokenize_fn(example.context) 181 | for i in range(len(example.qa_list)): 182 | tokens_qa = tokenize_fn(example.qa_list[i]) 183 | if len(tokens_qa) > FLAGS.max_qa_length: 184 | tokens_qa = tokens_qa[- FLAGS.max_qa_length:] 185 | 186 | if len(tokens_context) + len(tokens_qa) > FLAGS.max_seq_length - 3: 187 | tokens = tokens_context[: FLAGS.max_seq_length - 3 - len(tokens_qa)] 188 | else: 189 | tokens = tokens_context 190 | 191 | segment_ids = [SEG_ID_A] * len(tokens) 192 | 193 | tokens.append(SEP_ID) 194 | segment_ids.append(SEG_ID_A) 195 | 196 | tokens.extend(tokens_qa) 197 | segment_ids.extend([SEG_ID_B] * len(tokens_qa)) 198 | 199 | tokens.append(SEP_ID) 200 | segment_ids.append(SEG_ID_B) 201 | 202 | tokens.append(CLS_ID) 203 | segment_ids.append(SEG_ID_CLS) 204 | 205 | cur_input_ids = tokens 206 | cur_input_mask = [0] * len(cur_input_ids) 207 | 208 | if len(cur_input_ids) < FLAGS.max_seq_length: 209 | delta_len = FLAGS.max_seq_length - len(cur_input_ids) 210 | cur_input_ids = [0] * delta_len + cur_input_ids 211 | cur_input_mask = [1] * delta_len + cur_input_mask 212 | segment_ids = [SEG_ID_PAD] * delta_len + segment_ids 213 | 214 | assert len(cur_input_ids) == FLAGS.max_seq_length 215 | assert len(cur_input_mask) == FLAGS.max_seq_length 216 | assert len(segment_ids) == FLAGS.max_seq_length 217 | 218 | input_ids.extend(cur_input_ids) 219 | input_mask.extend(cur_input_mask) 220 | all_seg_ids.extend(segment_ids) 221 | 222 | label_id = example.label 223 | 224 | feature = InputFeatures( 225 | input_ids=input_ids, 226 | input_mask=input_mask, 227 | segment_ids=all_seg_ids, 228 | label_id=label_id) 229 | return feature 230 | 231 | 232 | class InputExample(object): 233 | def __init__(self, context, qa_list, label, level): 234 | self.context = context 235 | self.qa_list = qa_list 236 | self.label = label 237 | self.level = level 238 | 239 | 240 | def get_examples(data_dir, set_type): 241 | examples = [] 242 | 243 | for level in ["middle", "high"]: 244 | if level == "middle" and FLAGS.high_only: continue 245 | if level == "high" and FLAGS.middle_only: continue 246 | 247 | cur_dir = os.path.join(data_dir, set_type, level) 248 | for filename in tf.gfile.ListDirectory(cur_dir): 249 | cur_path = os.path.join(cur_dir, filename) 250 | with tf.gfile.Open(cur_path) as f: 251 | cur_data = json.load(f) 252 | 253 | answers = cur_data["answers"] 254 | options = cur_data["options"] 255 | questions = cur_data["questions"] 256 | context = cur_data["article"] 257 | 258 | for i in range(len(answers)): 259 | label = ord(answers[i]) - ord("A") 260 | qa_list = [] 261 | 262 | question = questions[i] 263 | for j in range(4): 264 | option = options[i][j] 265 | 266 | if "_" in question: 267 | qa_cat = question.replace("_", option) 268 | else: 269 | qa_cat = " ".join([question, option]) 270 | 271 | qa_list.append(qa_cat) 272 | 273 | examples.append(InputExample(context, qa_list, label, level)) 274 | 275 | return examples 276 | 277 | 278 | def file_based_convert_examples_to_features(examples, tokenize_fn, output_file): 279 | if tf.gfile.Exists(output_file) and not FLAGS.overwrite_data: 280 | return 281 | 282 | tf.logging.info("Start writing tfrecord %s.", output_file) 283 | writer = tf.python_io.TFRecordWriter(output_file) 284 | 285 | for ex_index, example in enumerate(examples): 286 | if ex_index % 10000 == 0: 287 | tf.logging.info("Writing example %d of %d" % (ex_index, len(examples))) 288 | 289 | feature = convert_single_example(example, tokenize_fn) 290 | 291 | def create_int_feature(values): 292 | f = tf.train.Feature(int64_list=tf.train.Int64List(value=list(values))) 293 | return f 294 | 295 | def create_float_feature(values): 296 | f = tf.train.Feature(float_list=tf.train.FloatList(value=list(values))) 297 | return f 298 | 299 | features = collections.OrderedDict() 300 | features["input_ids"] = create_int_feature(feature.input_ids) 301 | features["input_mask"] = create_float_feature(feature.input_mask) 302 | features["segment_ids"] = create_int_feature(feature.segment_ids) 303 | features["label_ids"] = create_int_feature([feature.label_id]) 304 | features["is_real_example"] = create_int_feature( 305 | [int(feature.is_real_example)]) 306 | 307 | tf_example = tf.train.Example(features=tf.train.Features(feature=features)) 308 | writer.write(tf_example.SerializeToString()) 309 | writer.close() 310 | 311 | 312 | def file_based_input_fn_builder(input_file, seq_length, is_training, 313 | drop_remainder): 314 | """Creates an `input_fn` closure to be passed to TPUEstimator.""" 315 | 316 | name_to_features = { 317 | "input_ids": tf.FixedLenFeature([seq_length * 4], tf.int64), 318 | "input_mask": tf.FixedLenFeature([seq_length * 4], tf.float32), 319 | "segment_ids": tf.FixedLenFeature([seq_length * 4], tf.int64), 320 | "label_ids": tf.FixedLenFeature([], tf.int64), 321 | "is_real_example": tf.FixedLenFeature([], tf.int64), 322 | } 323 | 324 | tf.logging.info("Input tfrecord file {}".format(input_file)) 325 | 326 | def _decode_record(record, name_to_features): 327 | """Decodes a record to a TensorFlow example.""" 328 | example = tf.parse_single_example(record, name_to_features) 329 | 330 | # tf.Example only supports tf.int64, but the TPU only supports tf.int32. 331 | # So cast all int64 to int32. 332 | for name in list(example.keys()): 333 | t = example[name] 334 | if t.dtype == tf.int64: 335 | t = tf.cast(t, tf.int32) 336 | example[name] = t 337 | 338 | return example 339 | 340 | def input_fn(params): 341 | """The actual input function.""" 342 | if FLAGS.use_tpu: 343 | batch_size = params["batch_size"] 344 | elif is_training: 345 | batch_size = FLAGS.train_batch_size 346 | elif FLAGS.do_eval: 347 | batch_size = FLAGS.eval_batch_size 348 | 349 | # For training, we want a lot of parallel reading and shuffling. 350 | # For eval, we want no shuffling and parallel reading doesn't matter. 351 | d = tf.data.TFRecordDataset(input_file) 352 | if is_training: 353 | d = d.shuffle(buffer_size=FLAGS.shuffle_buffer) 354 | d = d.repeat() 355 | # d = d.shuffle(buffer_size=100) 356 | 357 | d = d.apply( 358 | tf.contrib.data.map_and_batch( 359 | lambda record: _decode_record(record, name_to_features), 360 | batch_size=batch_size, 361 | drop_remainder=drop_remainder)) 362 | 363 | return d 364 | 365 | return input_fn 366 | 367 | 368 | def get_model_fn(): 369 | def model_fn(features, labels, mode, params): 370 | #### Training or Evaluation 371 | is_training = (mode == tf.estimator.ModeKeys.TRAIN) 372 | 373 | total_loss, per_example_loss, logits = function_builder.get_race_loss( 374 | FLAGS, features, is_training) 375 | 376 | #### Check model parameters 377 | num_params = sum([np.prod(v.shape) for v in tf.trainable_variables()]) 378 | tf.logging.info('#params: {}'.format(num_params)) 379 | 380 | #### load pretrained models 381 | scaffold_fn = model_utils.init_from_checkpoint(FLAGS) 382 | 383 | #### Evaluation mode 384 | if mode == tf.estimator.ModeKeys.EVAL: 385 | assert FLAGS.num_hosts == 1 386 | 387 | def metric_fn(per_example_loss, label_ids, logits, is_real_example): 388 | predictions = tf.argmax(logits, axis=-1, output_type=tf.int32) 389 | eval_input_dict = { 390 | 'labels': label_ids, 391 | 'predictions': predictions, 392 | 'weights': is_real_example 393 | } 394 | accuracy = tf.metrics.accuracy(**eval_input_dict) 395 | 396 | loss = tf.metrics.mean(values=per_example_loss, weights=is_real_example) 397 | return { 398 | 'eval_accuracy': accuracy, 399 | 'eval_loss': loss} 400 | 401 | is_real_example = tf.cast(features["is_real_example"], dtype=tf.float32) 402 | 403 | #### Constucting evaluation TPUEstimatorSpec with new cache. 404 | label_ids = tf.reshape(features['label_ids'], [-1]) 405 | metric_args = [per_example_loss, label_ids, logits, is_real_example] 406 | 407 | if FLAGS.use_tpu: 408 | eval_spec = tf.contrib.tpu.TPUEstimatorSpec( 409 | mode=mode, 410 | loss=total_loss, 411 | eval_metrics=(metric_fn, metric_args), 412 | scaffold_fn=scaffold_fn) 413 | else: 414 | eval_spec = tf.estimator.EstimatorSpec( 415 | mode=mode, 416 | loss=total_loss, 417 | eval_metric_ops=metric_fn(*metric_args)) 418 | 419 | return eval_spec 420 | 421 | 422 | #### Configuring the optimizer 423 | train_op, learning_rate, _ = model_utils.get_train_op(FLAGS, total_loss) 424 | 425 | monitor_dict = {} 426 | monitor_dict["lr"] = learning_rate 427 | 428 | #### Constucting training TPUEstimatorSpec with new cache. 429 | if FLAGS.use_tpu: 430 | #### Creating host calls 431 | host_call = None 432 | 433 | train_spec = tf.contrib.tpu.TPUEstimatorSpec( 434 | mode=mode, loss=total_loss, train_op=train_op, host_call=host_call, 435 | scaffold_fn=scaffold_fn) 436 | else: 437 | train_spec = tf.estimator.EstimatorSpec( 438 | mode=mode, loss=total_loss, train_op=train_op) 439 | 440 | return train_spec 441 | 442 | return model_fn 443 | 444 | 445 | def main(_): 446 | tf.logging.set_verbosity(tf.logging.INFO) 447 | 448 | #### Validate flags 449 | if FLAGS.save_steps is not None: 450 | FLAGS.iterations = min(FLAGS.iterations, FLAGS.save_steps) 451 | 452 | if not FLAGS.do_train and not FLAGS.do_eval: 453 | raise ValueError( 454 | "At least one of `do_train` or `do_eval` must be True.") 455 | 456 | if not tf.gfile.Exists(FLAGS.output_dir): 457 | tf.gfile.MakeDirs(FLAGS.output_dir) 458 | 459 | sp = spm.SentencePieceProcessor() 460 | sp.Load(FLAGS.spiece_model_file) 461 | def tokenize_fn(text): 462 | text = preprocess_text(text, lower=FLAGS.uncased) 463 | return encode_ids(sp, text) 464 | 465 | # TPU Configuration 466 | run_config = model_utils.configure_tpu(FLAGS) 467 | 468 | model_fn = get_model_fn() 469 | 470 | spm_basename = os.path.basename(FLAGS.spiece_model_file) 471 | 472 | # If TPU is not available, this will fall back to normal Estimator on CPU 473 | # or GPU. 474 | if FLAGS.use_tpu: 475 | estimator = tf.contrib.tpu.TPUEstimator( 476 | use_tpu=FLAGS.use_tpu, 477 | model_fn=model_fn, 478 | config=run_config, 479 | train_batch_size=FLAGS.train_batch_size, 480 | eval_batch_size=FLAGS.eval_batch_size) 481 | else: 482 | estimator = tf.estimator.Estimator( 483 | model_fn=model_fn, 484 | config=run_config) 485 | 486 | if FLAGS.do_train: 487 | train_file_base = "{}.len-{}.train.tf_record".format( 488 | spm_basename, FLAGS.max_seq_length) 489 | train_file = os.path.join(FLAGS.output_dir, train_file_base) 490 | 491 | if not tf.gfile.Exists(train_file) or FLAGS.overwrite_data: 492 | train_examples = get_examples(FLAGS.data_dir, "train") 493 | random.shuffle(train_examples) 494 | file_based_convert_examples_to_features( 495 | train_examples, tokenize_fn, train_file) 496 | 497 | train_input_fn = file_based_input_fn_builder( 498 | input_file=train_file, 499 | seq_length=FLAGS.max_seq_length, 500 | is_training=True, 501 | drop_remainder=True) 502 | estimator.train(input_fn=train_input_fn, max_steps=FLAGS.train_steps) 503 | 504 | if FLAGS.do_eval: 505 | eval_examples = get_examples(FLAGS.data_dir, FLAGS.eval_split) 506 | tf.logging.info("Num of eval samples: {}".format(len(eval_examples))) 507 | 508 | # TPU requires a fixed batch size for all batches, therefore the number 509 | # of examples must be a multiple of the batch size, or else examples 510 | # will get dropped. So we pad with fake examples which are ignored 511 | # later on. These do NOT count towards the metric (all tf.metrics 512 | # support a per-instance weight, and these get a weight of 0.0). 513 | # 514 | # Modified in XL: We also adopt the same mechanism for GPUs. 515 | 516 | while len(eval_examples) % FLAGS.eval_batch_size != 0: 517 | eval_examples.append(PaddingInputExample()) 518 | 519 | eval_file_base = "{}.len-{}.{}.tf_record".format( 520 | spm_basename, FLAGS.max_seq_length, FLAGS.eval_split) 521 | 522 | if FLAGS.high_only: 523 | eval_file_base = "high." + eval_file_base 524 | elif FLAGS.middle_only: 525 | eval_file_base = "middle." + eval_file_base 526 | 527 | eval_file = os.path.join(FLAGS.output_dir, eval_file_base) 528 | file_based_convert_examples_to_features( 529 | eval_examples, tokenize_fn, eval_file) 530 | 531 | assert len(eval_examples) % FLAGS.eval_batch_size == 0 532 | eval_steps = int(len(eval_examples) // FLAGS.eval_batch_size) 533 | 534 | eval_input_fn = file_based_input_fn_builder( 535 | input_file=eval_file, 536 | seq_length=FLAGS.max_seq_length, 537 | is_training=False, 538 | drop_remainder=True) 539 | 540 | ret = estimator.evaluate( 541 | input_fn=eval_input_fn, 542 | steps=eval_steps) 543 | 544 | # Log current result 545 | tf.logging.info("=" * 80) 546 | log_str = "Eval | " 547 | for key, val in ret.items(): 548 | log_str += "{} {} | ".format(key, val) 549 | tf.logging.info(log_str) 550 | tf.logging.info("=" * 80) 551 | 552 | 553 | if __name__ == "__main__": 554 | tf.app.run() 555 | -------------------------------------------------------------------------------- /scripts/gpu_squad_base.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | #### local path 4 | SQUAD_DIR=data/squad 5 | INIT_CKPT_DIR=xlnet_cased_L-12_H-768_A-12 6 | PROC_DATA_DIR=proc_data/squad 7 | MODEL_DIR=experiment/squad 8 | 9 | #### Use 3 GPUs, each with 8 seqlen-512 samples 10 | 11 | python run_squad.py \ 12 | --use_tpu=False \ 13 | --num_hosts=1 \ 14 | --num_core_per_host=3 \ 15 | --model_config_path=${INIT_CKPT_DIR}/xlnet_config.json \ 16 | --spiece_model_file=${INIT_CKPT_DIR}/spiece.model \ 17 | --output_dir=${PROC_DATA_DIR} \ 18 | --init_checkpoint=${INIT_CKPT_DIR}/xlnet_model.ckpt \ 19 | --model_dir=${MODEL_DIR} \ 20 | --train_file=${SQUAD_DIR}/train-v2.0.json \ 21 | --predict_file=${SQUAD_DIR}/dev-v2.0.json \ 22 | --uncased=False \ 23 | --max_seq_length=512 \ 24 | --do_train=True \ 25 | --train_batch_size=8 \ 26 | --do_predict=True \ 27 | --predict_batch_size=32 \ 28 | --learning_rate=2e-5 \ 29 | --adam_epsilon=1e-6 \ 30 | --iterations=1000 \ 31 | --save_steps=1000 \ 32 | --train_steps=12000 \ 33 | --warmup_steps=1000 \ 34 | $@ 35 | -------------------------------------------------------------------------------- /scripts/prepro_squad.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | #### local path 4 | SQUAD_DIR=data/squad 5 | INIT_CKPT_DIR=xlnet_cased_L-24_H-1024_A-16 6 | 7 | #### google storage path 8 | GS_ROOT= 9 | GS_PROC_DATA_DIR=${GS_ROOT}/proc_data/squad 10 | 11 | python run_squad.py \ 12 | --use_tpu=False \ 13 | --do_prepro=True \ 14 | --spiece_model_file=${INIT_CKPT_DIR}/spiece.model \ 15 | --train_file=${SQUAD_DIR}/train-v2.0.json \ 16 | --output_dir=${GS_PROC_DATA_DIR} \ 17 | --uncased=False \ 18 | --max_seq_length=512 \ 19 | $@ 20 | 21 | #### Potential multi-processing version 22 | # NUM_PROC=8 23 | # for i in `seq 0 $((NUM_PROC - 1))`; do 24 | # python run_squad.py \ 25 | # --use_tpu=False \ 26 | # --do_prepro=True \ 27 | # --spiece_model_file=${INIT_CKPT_DIR}/spiece.model \ 28 | # --train_file=${SQUAD_DIR}/train-v2.0.json \ 29 | # --output_dir=${GS_PROC_DATA_DIR} \ 30 | # --uncased=False \ 31 | # --max_seq_length=512 \ 32 | # --num_proc=${NUM_PROC} \ 33 | # --proc_id=${i} \ 34 | # $@ & 35 | # done 36 | -------------------------------------------------------------------------------- /scripts/tpu_race_large_bsz32.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | #### local path 4 | RACE_DIR=data/RACE 5 | INIT_CKPT_DIR=xlnet_cased_L-24_H-1024_A-16 6 | 7 | #### google storage path 8 | GS_ROOT= 9 | GS_INIT_CKPT_DIR=${GS_ROOT}/${INIT_CKPT_DIR} 10 | GS_PROC_DATA_DIR=${GS_ROOT}/proc_data/race 11 | GS_MODEL_DIR=${GS_ROOT}/experiment/race 12 | 13 | # TPU name in google cloud 14 | TPU_NAME= 15 | 16 | python run_race.py \ 17 | --use_tpu=True \ 18 | --tpu=${TPU_NAME} \ 19 | --num_hosts=4 \ 20 | --num_core_per_host=8 \ 21 | --model_config_path=${INIT_CKPT_DIR}/xlnet_config.json \ 22 | --spiece_model_file=${INIT_CKPT_DIR}/spiece.model \ 23 | --output_dir=${GS_PROC_DATA_DIR} \ 24 | --init_checkpoint=${GS_INIT_CKPT_DIR}/xlnet_model.ckpt \ 25 | --model_dir=${GS_MODEL_DIR} \ 26 | --data_dir=${RACE_DIR} \ 27 | --max_seq_length=512 \ 28 | --max_qa_length=128 \ 29 | --uncased=False \ 30 | --do_train=True \ 31 | --train_batch_size=32 \ 32 | --do_eval=True \ 33 | --eval_batch_size=32 \ 34 | --train_steps=12000 \ 35 | --save_steps=1000 \ 36 | --iterations=1000 \ 37 | --warmup_steps=1000 \ 38 | --learning_rate=2e-5 \ 39 | --weight_decay=0 \ 40 | --adam_epsilon=1e-6 \ 41 | $@ 42 | -------------------------------------------------------------------------------- /scripts/tpu_race_large_bsz8.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | #### local path 4 | RACE_DIR=data/RACE 5 | INIT_CKPT_DIR=xlnet_cased_L-24_H-1024_A-16 6 | 7 | #### google storage path 8 | GS_ROOT= 9 | GS_INIT_CKPT_DIR=${GS_ROOT}/${INIT_CKPT_DIR} 10 | GS_PROC_DATA_DIR=${GS_ROOT}/proc_data/race 11 | GS_MODEL_DIR=${GS_ROOT}/experiment/race 12 | 13 | # TPU name in google cloud 14 | TPU_NAME= 15 | 16 | python run_race.py \ 17 | --use_tpu=True \ 18 | --tpu=${TPU_NAME} \ 19 | --num_hosts=1 \ 20 | --num_core_per_host=8 \ 21 | --model_config_path=${INIT_CKPT_DIR}/xlnet_config.json \ 22 | --spiece_model_file=${INIT_CKPT_DIR}/spiece.model \ 23 | --output_dir=${GS_PROC_DATA_DIR} \ 24 | --init_checkpoint=${GS_INIT_CKPT_DIR}/xlnet_model.ckpt \ 25 | --model_dir=${GS_MODEL_DIR} \ 26 | --data_dir=${RACE_DIR} \ 27 | --max_seq_length=512 \ 28 | --max_qa_length=128 \ 29 | --uncased=False \ 30 | --do_train=True \ 31 | --train_batch_size=8 \ 32 | --do_eval=True \ 33 | --eval_batch_size=32 \ 34 | --train_steps=12000 \ 35 | --save_steps=1000 \ 36 | --iterations=1000 \ 37 | --warmup_steps=1000 \ 38 | --learning_rate=2e-5 \ 39 | --weight_decay=0 \ 40 | --adam_epsilon=1e-6 \ 41 | $@ 42 | -------------------------------------------------------------------------------- /scripts/tpu_squad_large.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | #### local path 4 | SQUAD_DIR=data/squad 5 | INIT_CKPT_DIR=xlnet_cased_L-24_H-1024_A-16 6 | 7 | #### google storage path 8 | GS_ROOT= 9 | GS_INIT_CKPT_DIR=${GS_ROOT}/${INIT_CKPT_DIR} 10 | GS_PROC_DATA_DIR=${GS_ROOT}/proc_data/squad 11 | GS_MODEL_DIR=${GS_ROOT}/experiment/squad 12 | 13 | # TPU name in google cloud 14 | TPU_NAME= 15 | 16 | python run_squad.py \ 17 | --use_tpu=True \ 18 | --tpu=${TPU_NAME} \ 19 | --num_hosts=1 \ 20 | --num_core_per_host=8 \ 21 | --model_config_path=${INIT_CKPT_DIR}/xlnet_config.json \ 22 | --spiece_model_file=${INIT_CKPT_DIR}/spiece.model \ 23 | --output_dir=${GS_PROC_DATA_DIR} \ 24 | --init_checkpoint=${GS_INIT_CKPT_DIR}/xlnet_model.ckpt \ 25 | --model_dir=${GS_MODEL_DIR} \ 26 | --train_file=${SQUAD_DIR}/train-v2.0.json \ 27 | --predict_file=${SQUAD_DIR}/dev-v2.0.json \ 28 | --uncased=False \ 29 | --max_seq_length=512 \ 30 | --do_train=True \ 31 | --train_batch_size=48 \ 32 | --do_predict=True \ 33 | --predict_batch_size=32 \ 34 | --learning_rate=3e-5 \ 35 | --adam_epsilon=1e-6 \ 36 | --iterations=1000 \ 37 | --save_steps=1000 \ 38 | --train_steps=8000 \ 39 | --warmup_steps=1000 \ 40 | $@ 41 | -------------------------------------------------------------------------------- /squad_utils.py: -------------------------------------------------------------------------------- 1 | """Official evaluation script for SQuAD version 2.0. 2 | 3 | In addition to basic functionality, we also compute additional statistics and 4 | plot precision-recall curves if an additional na_prob.json file is provided. 5 | This file is expected to map question ID's to the model's predicted probability 6 | that a question is unanswerable. 7 | """ 8 | import argparse 9 | import collections 10 | import json 11 | import numpy as np 12 | import os 13 | import re 14 | import string 15 | import sys 16 | 17 | OPTS = None 18 | 19 | def parse_args(): 20 | parser = argparse.ArgumentParser('Official evaluation script for SQuAD version 2.0.') 21 | parser.add_argument('data_file', metavar='data.json', help='Input data JSON file.') 22 | parser.add_argument('pred_file', metavar='pred.json', help='Model predictions.') 23 | parser.add_argument('--out-file', '-o', metavar='eval.json', 24 | help='Write accuracy metrics to file (default is stdout).') 25 | parser.add_argument('--na-prob-file', '-n', metavar='na_prob.json', 26 | help='Model estimates of probability of no answer.') 27 | parser.add_argument('--na-prob-thresh', '-t', type=float, default=1.0, 28 | help='Predict "" if no-answer probability exceeds this (default = 1.0).') 29 | parser.add_argument('--out-image-dir', '-p', metavar='out_images', default=None, 30 | help='Save precision-recall curves to directory.') 31 | parser.add_argument('--verbose', '-v', action='store_true') 32 | if len(sys.argv) == 1: 33 | parser.print_help() 34 | sys.exit(1) 35 | return parser.parse_args() 36 | 37 | def make_qid_to_has_ans(dataset): 38 | qid_to_has_ans = {} 39 | for article in dataset: 40 | for p in article['paragraphs']: 41 | for qa in p['qas']: 42 | qid_to_has_ans[qa['id']] = bool(qa['answers']) 43 | return qid_to_has_ans 44 | 45 | def normalize_answer(s): 46 | """Lower text and remove punctuation, articles and extra whitespace.""" 47 | def remove_articles(text): 48 | regex = re.compile(r'\b(a|an|the)\b', re.UNICODE) 49 | return re.sub(regex, ' ', text) 50 | def white_space_fix(text): 51 | return ' '.join(text.split()) 52 | def remove_punc(text): 53 | exclude = set(string.punctuation) 54 | return ''.join(ch for ch in text if ch not in exclude) 55 | def lower(text): 56 | return text.lower() 57 | return white_space_fix(remove_articles(remove_punc(lower(s)))) 58 | 59 | def get_tokens(s): 60 | if not s: return [] 61 | return normalize_answer(s).split() 62 | 63 | def compute_exact(a_gold, a_pred): 64 | return int(normalize_answer(a_gold) == normalize_answer(a_pred)) 65 | 66 | def compute_f1(a_gold, a_pred): 67 | gold_toks = get_tokens(a_gold) 68 | pred_toks = get_tokens(a_pred) 69 | common = collections.Counter(gold_toks) & collections.Counter(pred_toks) 70 | num_same = sum(common.values()) 71 | if len(gold_toks) == 0 or len(pred_toks) == 0: 72 | # If either is no-answer, then F1 is 1 if they agree, 0 otherwise 73 | return int(gold_toks == pred_toks) 74 | if num_same == 0: 75 | return 0 76 | precision = 1.0 * num_same / len(pred_toks) 77 | recall = 1.0 * num_same / len(gold_toks) 78 | f1 = (2 * precision * recall) / (precision + recall) 79 | return f1 80 | 81 | def get_raw_scores(dataset, preds): 82 | exact_scores = {} 83 | f1_scores = {} 84 | for article in dataset: 85 | for p in article['paragraphs']: 86 | for qa in p['qas']: 87 | qid = qa['id'] 88 | gold_answers = [a['text'] for a in qa['answers'] 89 | if normalize_answer(a['text'])] 90 | if not gold_answers: 91 | # For unanswerable questions, only correct answer is empty string 92 | gold_answers = [''] 93 | if qid not in preds: 94 | print('Missing prediction for %s' % qid) 95 | continue 96 | a_pred = preds[qid] 97 | # Take max over all gold answers 98 | exact_scores[qid] = max(compute_exact(a, a_pred) for a in gold_answers) 99 | f1_scores[qid] = max(compute_f1(a, a_pred) for a in gold_answers) 100 | return exact_scores, f1_scores 101 | 102 | def apply_no_ans_threshold(scores, na_probs, qid_to_has_ans, na_prob_thresh): 103 | new_scores = {} 104 | for qid, s in scores.items(): 105 | pred_na = na_probs[qid] > na_prob_thresh 106 | if pred_na: 107 | new_scores[qid] = float(not qid_to_has_ans[qid]) 108 | else: 109 | new_scores[qid] = s 110 | return new_scores 111 | 112 | def make_eval_dict(exact_scores, f1_scores, qid_list=None): 113 | if not qid_list: 114 | total = len(exact_scores) 115 | return collections.OrderedDict([ 116 | ('exact', 100.0 * sum(exact_scores.values()) / total), 117 | ('f1', 100.0 * sum(f1_scores.values()) / total), 118 | ('total', total), 119 | ]) 120 | else: 121 | total = len(qid_list) 122 | return collections.OrderedDict([ 123 | ('exact', 100.0 * sum(exact_scores[k] for k in qid_list) / total), 124 | ('f1', 100.0 * sum(f1_scores[k] for k in qid_list) / total), 125 | ('total', total), 126 | ]) 127 | 128 | def merge_eval(main_eval, new_eval, prefix): 129 | for k in new_eval: 130 | main_eval['%s_%s' % (prefix, k)] = new_eval[k] 131 | 132 | def plot_pr_curve(precisions, recalls, out_image, title): 133 | plt.step(recalls, precisions, color='b', alpha=0.2, where='post') 134 | plt.fill_between(recalls, precisions, step='post', alpha=0.2, color='b') 135 | plt.xlabel('Recall') 136 | plt.ylabel('Precision') 137 | plt.xlim([0.0, 1.05]) 138 | plt.ylim([0.0, 1.05]) 139 | plt.title(title) 140 | plt.savefig(out_image) 141 | plt.clf() 142 | 143 | def make_precision_recall_eval(scores, na_probs, num_true_pos, qid_to_has_ans, 144 | out_image=None, title=None): 145 | qid_list = sorted(na_probs, key=lambda k: na_probs[k]) 146 | true_pos = 0.0 147 | cur_p = 1.0 148 | cur_r = 0.0 149 | precisions = [1.0] 150 | recalls = [0.0] 151 | avg_prec = 0.0 152 | for i, qid in enumerate(qid_list): 153 | if qid_to_has_ans[qid]: 154 | true_pos += scores[qid] 155 | cur_p = true_pos / float(i+1) 156 | cur_r = true_pos / float(num_true_pos) 157 | if i == len(qid_list) - 1 or na_probs[qid] != na_probs[qid_list[i+1]]: 158 | # i.e., if we can put a threshold after this point 159 | avg_prec += cur_p * (cur_r - recalls[-1]) 160 | precisions.append(cur_p) 161 | recalls.append(cur_r) 162 | if out_image: 163 | plot_pr_curve(precisions, recalls, out_image, title) 164 | return {'ap': 100.0 * avg_prec} 165 | 166 | def run_precision_recall_analysis(main_eval, exact_raw, f1_raw, na_probs, 167 | qid_to_has_ans, out_image_dir): 168 | if out_image_dir and not os.path.exists(out_image_dir): 169 | os.makedirs(out_image_dir) 170 | num_true_pos = sum(1 for v in qid_to_has_ans.values() if v) 171 | if num_true_pos == 0: 172 | return 173 | pr_exact = make_precision_recall_eval( 174 | exact_raw, na_probs, num_true_pos, qid_to_has_ans, 175 | out_image=os.path.join(out_image_dir, 'pr_exact.png'), 176 | title='Precision-Recall curve for Exact Match score') 177 | pr_f1 = make_precision_recall_eval( 178 | f1_raw, na_probs, num_true_pos, qid_to_has_ans, 179 | out_image=os.path.join(out_image_dir, 'pr_f1.png'), 180 | title='Precision-Recall curve for F1 score') 181 | oracle_scores = {k: float(v) for k, v in qid_to_has_ans.items()} 182 | pr_oracle = make_precision_recall_eval( 183 | oracle_scores, na_probs, num_true_pos, qid_to_has_ans, 184 | out_image=os.path.join(out_image_dir, 'pr_oracle.png'), 185 | title='Oracle Precision-Recall curve (binary task of HasAns vs. NoAns)') 186 | merge_eval(main_eval, pr_exact, 'pr_exact') 187 | merge_eval(main_eval, pr_f1, 'pr_f1') 188 | merge_eval(main_eval, pr_oracle, 'pr_oracle') 189 | 190 | def histogram_na_prob(na_probs, qid_list, image_dir, name): 191 | if not qid_list: 192 | return 193 | x = [na_probs[k] for k in qid_list] 194 | weights = np.ones_like(x) / float(len(x)) 195 | plt.hist(x, weights=weights, bins=20, range=(0.0, 1.0)) 196 | plt.xlabel('Model probability of no-answer') 197 | plt.ylabel('Proportion of dataset') 198 | plt.title('Histogram of no-answer probability: %s' % name) 199 | plt.savefig(os.path.join(image_dir, 'na_prob_hist_%s.png' % name)) 200 | plt.clf() 201 | 202 | def find_best_thresh(preds, scores, na_probs, qid_to_has_ans): 203 | num_no_ans = sum(1 for k in qid_to_has_ans if not qid_to_has_ans[k]) 204 | cur_score = num_no_ans 205 | best_score = cur_score 206 | best_thresh = 0.0 207 | qid_list = sorted(na_probs, key=lambda k: na_probs[k]) 208 | for i, qid in enumerate(qid_list): 209 | if qid not in scores: continue 210 | if qid_to_has_ans[qid]: 211 | diff = scores[qid] 212 | else: 213 | if preds[qid]: 214 | diff = -1 215 | else: 216 | diff = 0 217 | cur_score += diff 218 | if cur_score > best_score: 219 | best_score = cur_score 220 | best_thresh = na_probs[qid] 221 | return 100.0 * best_score / len(scores), best_thresh 222 | 223 | def find_best_thresh_v2(preds, scores, na_probs, qid_to_has_ans): 224 | num_no_ans = sum(1 for k in qid_to_has_ans if not qid_to_has_ans[k]) 225 | cur_score = num_no_ans 226 | best_score = cur_score 227 | best_thresh = 0.0 228 | qid_list = sorted(na_probs, key=lambda k: na_probs[k]) 229 | for i, qid in enumerate(qid_list): 230 | if qid not in scores: continue 231 | if qid_to_has_ans[qid]: 232 | diff = scores[qid] 233 | else: 234 | if preds[qid]: 235 | diff = -1 236 | else: 237 | diff = 0 238 | cur_score += diff 239 | if cur_score > best_score: 240 | best_score = cur_score 241 | best_thresh = na_probs[qid] 242 | 243 | has_ans_score, has_ans_cnt = 0, 0 244 | for qid in qid_list: 245 | if not qid_to_has_ans[qid]: continue 246 | has_ans_cnt += 1 247 | 248 | if qid not in scores: continue 249 | has_ans_score += scores[qid] 250 | 251 | return 100.0 * best_score / len(scores), best_thresh, 1.0 * has_ans_score / has_ans_cnt 252 | 253 | def find_all_best_thresh(main_eval, preds, exact_raw, f1_raw, na_probs, qid_to_has_ans): 254 | best_exact, exact_thresh = find_best_thresh(preds, exact_raw, na_probs, qid_to_has_ans) 255 | best_f1, f1_thresh = find_best_thresh(preds, f1_raw, na_probs, qid_to_has_ans) 256 | main_eval['best_exact'] = best_exact 257 | main_eval['best_exact_thresh'] = exact_thresh 258 | main_eval['best_f1'] = best_f1 259 | main_eval['best_f1_thresh'] = f1_thresh 260 | 261 | def find_all_best_thresh_v2(main_eval, preds, exact_raw, f1_raw, na_probs, qid_to_has_ans): 262 | best_exact, exact_thresh, has_ans_exact = find_best_thresh_v2(preds, exact_raw, na_probs, qid_to_has_ans) 263 | best_f1, f1_thresh, has_ans_f1 = find_best_thresh_v2(preds, f1_raw, na_probs, qid_to_has_ans) 264 | main_eval['best_exact'] = best_exact 265 | main_eval['best_exact_thresh'] = exact_thresh 266 | main_eval['best_f1'] = best_f1 267 | main_eval['best_f1_thresh'] = f1_thresh 268 | main_eval['has_ans_exact'] = has_ans_exact 269 | main_eval['has_ans_f1'] = has_ans_f1 270 | 271 | def main(): 272 | with open(OPTS.data_file) as f: 273 | dataset_json = json.load(f) 274 | dataset = dataset_json['data'] 275 | with open(OPTS.pred_file) as f: 276 | preds = json.load(f) 277 | 278 | new_orig_data = [] 279 | for article in dataset: 280 | for p in article['paragraphs']: 281 | for qa in p['qas']: 282 | if qa['id'] in preds: 283 | new_para = {'qas': [qa]} 284 | new_article = {'paragraphs': [new_para]} 285 | new_orig_data.append(new_article) 286 | dataset = new_orig_data 287 | 288 | if OPTS.na_prob_file: 289 | with open(OPTS.na_prob_file) as f: 290 | na_probs = json.load(f) 291 | else: 292 | na_probs = {k: 0.0 for k in preds} 293 | qid_to_has_ans = make_qid_to_has_ans(dataset) # maps qid to True/False 294 | has_ans_qids = [k for k, v in qid_to_has_ans.items() if v] 295 | no_ans_qids = [k for k, v in qid_to_has_ans.items() if not v] 296 | exact_raw, f1_raw = get_raw_scores(dataset, preds) 297 | exact_thresh = apply_no_ans_threshold(exact_raw, na_probs, qid_to_has_ans, 298 | OPTS.na_prob_thresh) 299 | f1_thresh = apply_no_ans_threshold(f1_raw, na_probs, qid_to_has_ans, 300 | OPTS.na_prob_thresh) 301 | out_eval = make_eval_dict(exact_thresh, f1_thresh) 302 | if has_ans_qids: 303 | has_ans_eval = make_eval_dict(exact_thresh, f1_thresh, qid_list=has_ans_qids) 304 | merge_eval(out_eval, has_ans_eval, 'HasAns') 305 | if no_ans_qids: 306 | no_ans_eval = make_eval_dict(exact_thresh, f1_thresh, qid_list=no_ans_qids) 307 | merge_eval(out_eval, no_ans_eval, 'NoAns') 308 | if OPTS.na_prob_file: 309 | find_all_best_thresh(out_eval, preds, exact_raw, f1_raw, na_probs, qid_to_has_ans) 310 | if OPTS.na_prob_file and OPTS.out_image_dir: 311 | run_precision_recall_analysis(out_eval, exact_raw, f1_raw, na_probs, 312 | qid_to_has_ans, OPTS.out_image_dir) 313 | histogram_na_prob(na_probs, has_ans_qids, OPTS.out_image_dir, 'hasAns') 314 | histogram_na_prob(na_probs, no_ans_qids, OPTS.out_image_dir, 'noAns') 315 | if OPTS.out_file: 316 | with open(OPTS.out_file, 'w') as f: 317 | json.dump(out_eval, f) 318 | else: 319 | print(json.dumps(out_eval, indent=2)) 320 | 321 | if __name__ == '__main__': 322 | OPTS = parse_args() 323 | if OPTS.out_image_dir: 324 | import matplotlib 325 | matplotlib.use('Agg') 326 | import matplotlib.pyplot as plt 327 | main() 328 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | """Pretraining on TPUs.""" 2 | from __future__ import absolute_import 3 | from __future__ import division 4 | from __future__ import print_function 5 | 6 | import os 7 | 8 | from absl import app 9 | from absl import flags 10 | import absl.logging as _logging # pylint: disable=unused-import 11 | 12 | import numpy as np 13 | 14 | import tensorflow as tf 15 | import model_utils 16 | import tpu_estimator 17 | import function_builder 18 | import data_utils 19 | 20 | # TPU parameters 21 | flags.DEFINE_string("master", default=None, 22 | help="master") 23 | flags.DEFINE_string("tpu", default=None, 24 | help="The Cloud TPU to use for training. This should be either the name " 25 | "used when creating the Cloud TPU, or a grpc://ip.address.of.tpu:8470 url.") 26 | flags.DEFINE_string("gcp_project", default=None, 27 | help="Project name for the Cloud TPU-enabled project. If not specified, " 28 | "we will attempt to automatically detect the GCE project from metadata.") 29 | flags.DEFINE_string("tpu_zone",default=None, 30 | help="GCE zone where the Cloud TPU is located in. If not specified, we " 31 | "will attempt to automatically detect the GCE project from metadata.") 32 | flags.DEFINE_bool("use_tpu", default=True, 33 | help="Use TPUs rather than plain CPUs.") 34 | flags.DEFINE_integer("num_hosts", default=1, 35 | help="number of TPU hosts") 36 | flags.DEFINE_integer("num_core_per_host", default=8, 37 | help="number of cores per host") 38 | flags.DEFINE_bool("track_mean", default=False, 39 | help="Whether to track mean loss.") 40 | 41 | # Experiment (data/checkpoint/directory) config 42 | flags.DEFINE_integer("num_passes", default=1, 43 | help="Number of passed used for training.") 44 | flags.DEFINE_string("record_info_dir", default=None, 45 | help="Path to local directory containing `record_info-lm.json`.") 46 | flags.DEFINE_string("model_dir", default=None, 47 | help="Estimator model_dir.") 48 | flags.DEFINE_string("init_checkpoint", default=None, 49 | help="Checkpoint path for initializing the model.") 50 | 51 | # Optimization config 52 | flags.DEFINE_float("learning_rate", default=1e-4, 53 | help="Maximum learning rate.") 54 | flags.DEFINE_float("clip", default=1.0, 55 | help="Gradient clipping value.") 56 | # lr decay 57 | flags.DEFINE_float("min_lr_ratio", default=0.001, 58 | help="Minimum ratio learning rate.") 59 | flags.DEFINE_integer("warmup_steps", default=0, 60 | help="Number of steps for linear lr warmup.") 61 | flags.DEFINE_float("adam_epsilon", default=1e-8, 62 | help="Adam epsilon.") 63 | flags.DEFINE_string("decay_method", default="poly", 64 | help="Poly or cos.") 65 | flags.DEFINE_float("weight_decay", default=0.0, 66 | help="Weight decay rate.") 67 | 68 | # Training config 69 | flags.DEFINE_integer("train_batch_size", default=16, 70 | help="Size of the train batch across all hosts.") 71 | flags.DEFINE_integer("train_steps", default=100000, 72 | help="Total number of training steps.") 73 | flags.DEFINE_integer("iterations", default=1000, 74 | help="Number of iterations per repeat loop.") 75 | flags.DEFINE_integer("save_steps", default=None, 76 | help="Number of steps for model checkpointing. " 77 | "None for not saving checkpoints") 78 | flags.DEFINE_integer("max_save", default=100000, 79 | help="Maximum number of checkpoints to save.") 80 | 81 | # Data config 82 | flags.DEFINE_integer("seq_len", default=0, 83 | help="Sequence length for pretraining.") 84 | flags.DEFINE_integer("reuse_len", default=0, 85 | help="How many tokens to be reused in the next batch. " 86 | "Could be half of `seq_len`.") 87 | flags.DEFINE_bool("uncased", False, 88 | help="Use uncased inputs or not.") 89 | flags.DEFINE_integer("perm_size", 0, 90 | help="Window size of permutation.") 91 | flags.DEFINE_bool("bi_data", default=True, 92 | help="Use bidirectional data streams, i.e., forward & backward.") 93 | flags.DEFINE_integer("mask_alpha", default=6, 94 | help="How many tokens to form a group.") 95 | flags.DEFINE_integer("mask_beta", default=1, 96 | help="How many tokens to mask within each group.") 97 | flags.DEFINE_integer("num_predict", default=None, 98 | help="Number of tokens to predict in partial prediction.") 99 | flags.DEFINE_integer("n_token", 32000, help="Vocab size") 100 | 101 | # Model config 102 | flags.DEFINE_integer("mem_len", default=0, 103 | help="Number of steps to cache") 104 | flags.DEFINE_bool("same_length", default=False, 105 | help="Same length attention") 106 | flags.DEFINE_integer("clamp_len", default=-1, 107 | help="Clamp length") 108 | 109 | flags.DEFINE_integer("n_layer", default=6, 110 | help="Number of layers.") 111 | flags.DEFINE_integer("d_model", default=32, 112 | help="Dimension of the model.") 113 | flags.DEFINE_integer("d_embed", default=32, 114 | help="Dimension of the embeddings.") 115 | flags.DEFINE_integer("n_head", default=4, 116 | help="Number of attention heads.") 117 | flags.DEFINE_integer("d_head", default=8, 118 | help="Dimension of each attention head.") 119 | flags.DEFINE_integer("d_inner", default=32, 120 | help="Dimension of inner hidden size in positionwise feed-forward.") 121 | flags.DEFINE_float("dropout", default=0.0, 122 | help="Dropout rate.") 123 | flags.DEFINE_float("dropatt", default=0.0, 124 | help="Attention dropout rate.") 125 | flags.DEFINE_bool("untie_r", default=False, 126 | help="Untie r_w_bias and r_r_bias") 127 | flags.DEFINE_string("summary_type", default="last", 128 | help="Method used to summarize a sequence into a compact vector.") 129 | flags.DEFINE_string("ff_activation", default="relu", 130 | help="Activation type used in position-wise feed-forward.") 131 | flags.DEFINE_bool("use_bfloat16", False, 132 | help="Whether to use bfloat16.") 133 | 134 | # Parameter initialization 135 | flags.DEFINE_enum("init", default="normal", 136 | enum_values=["normal", "uniform"], 137 | help="Initialization method.") 138 | flags.DEFINE_float("init_std", default=0.02, 139 | help="Initialization std when init is normal.") 140 | flags.DEFINE_float("init_range", default=0.1, 141 | help="Initialization std when init is uniform.") 142 | 143 | FLAGS = flags.FLAGS 144 | 145 | 146 | def get_model_fn(): 147 | """doc.""" 148 | def model_fn(features, labels, mode, params): 149 | """doc.""" 150 | #### Training or Evaluation 151 | is_training = (mode == tf.estimator.ModeKeys.TRAIN) 152 | assert is_training 153 | 154 | #### Retrieve `mems` from `params["cache"]` 155 | mems = {} 156 | idx = 0 157 | if FLAGS.mem_len > 0: 158 | mems["mems"] = params["cache"] 159 | 160 | #### Get loss from inputs 161 | total_loss, new_mems, monitor_dict = function_builder.get_loss( 162 | FLAGS, features, labels, mems, is_training) 163 | 164 | #### Turn `new_mems` into `new_cache` 165 | new_cache = [] 166 | if FLAGS.mem_len > 0: 167 | new_cache += new_mems["mems"] 168 | 169 | #### Check model parameters 170 | num_params = sum([np.prod(v.shape) for v in tf.trainable_variables()]) 171 | tf.logging.info("#params: {}".format(num_params)) 172 | 173 | #### Configuring the optimizer 174 | train_op, learning_rate, gnorm = model_utils.get_train_op( 175 | FLAGS, total_loss) 176 | monitor_dict["lr"] = learning_rate 177 | monitor_dict["gnorm"] = gnorm 178 | 179 | #### Customized initial checkpoint 180 | scaffold_fn = model_utils.init_from_checkpoint(FLAGS, global_vars=True) 181 | 182 | #### Creating host calls 183 | host_call = function_builder.construct_scalar_host_call( 184 | monitor_dict=monitor_dict, 185 | model_dir=FLAGS.model_dir, 186 | prefix="train/", 187 | reduce_fn=tf.reduce_mean) 188 | 189 | #### Constucting training TPUEstimatorSpec with new cache. 190 | train_spec = tf.contrib.tpu.TPUEstimatorSpec( 191 | mode=mode, loss=total_loss, train_op=train_op, host_call=host_call, 192 | scaffold_fn=scaffold_fn) 193 | 194 | train_spec.cache = new_cache 195 | 196 | return train_spec 197 | 198 | return model_fn 199 | 200 | 201 | def get_cache_fn(mem_len): 202 | """doc.""" 203 | tf_float = tf.bfloat16 if FLAGS.use_bfloat16 else tf.float32 204 | def cache_fn(batch_size): 205 | mems = [] 206 | if FLAGS.mem_len > 0: 207 | for _ in range(FLAGS.n_layer): 208 | zeros = tf.zeros( 209 | [mem_len, batch_size, FLAGS.d_model], 210 | dtype=tf_float) 211 | mems.append(zeros) 212 | 213 | return mems 214 | 215 | if mem_len > 0: 216 | return cache_fn 217 | else: 218 | return None 219 | 220 | 221 | def get_input_fn(split): 222 | """doc.""" 223 | assert split == "train" 224 | batch_size = FLAGS.train_batch_size 225 | 226 | input_fn, record_info_dict = data_utils.get_input_fn( 227 | tfrecord_dir=FLAGS.record_info_dir, 228 | split=split, 229 | bsz_per_host=batch_size // FLAGS.num_hosts, 230 | seq_len=FLAGS.seq_len, 231 | reuse_len=FLAGS.reuse_len, 232 | bi_data=FLAGS.bi_data, 233 | num_hosts=FLAGS.num_hosts, 234 | num_core_per_host=FLAGS.num_core_per_host, 235 | perm_size=FLAGS.perm_size, 236 | mask_alpha=FLAGS.mask_alpha, 237 | mask_beta=FLAGS.mask_beta, 238 | uncased=FLAGS.uncased, 239 | num_passes=FLAGS.num_passes, 240 | use_bfloat16=FLAGS.use_bfloat16, 241 | num_predict=FLAGS.num_predict) 242 | 243 | return input_fn, record_info_dict 244 | 245 | 246 | def main(unused_argv): 247 | del unused_argv # Unused 248 | 249 | tf.logging.set_verbosity(tf.logging.INFO) 250 | 251 | assert FLAGS.seq_len > 0 252 | assert FLAGS.perm_size > 0 253 | 254 | FLAGS.n_token = data_utils.VOCAB_SIZE 255 | tf.logging.info("n_token {}".format(FLAGS.n_token)) 256 | 257 | if not tf.gfile.Exists(FLAGS.model_dir): 258 | tf.gfile.MakeDirs(FLAGS.model_dir) 259 | 260 | # Get train input function 261 | train_input_fn, train_record_info_dict = get_input_fn("train") 262 | 263 | tf.logging.info("num of batches {}".format( 264 | train_record_info_dict["num_batch"])) 265 | 266 | # Get train cache function 267 | train_cache_fn = get_cache_fn(FLAGS.mem_len) 268 | 269 | ##### Get model function 270 | model_fn = get_model_fn() 271 | 272 | ##### Create TPUEstimator 273 | # TPU Configuration 274 | run_config = model_utils.configure_tpu(FLAGS) 275 | 276 | # TPU Estimator 277 | estimator = tpu_estimator.TPUEstimator( 278 | model_fn=model_fn, 279 | train_cache_fn=train_cache_fn, 280 | use_tpu=FLAGS.use_tpu, 281 | config=run_config, 282 | params={"track_mean": FLAGS.track_mean}, 283 | train_batch_size=FLAGS.train_batch_size, 284 | eval_on_tpu=FLAGS.use_tpu) 285 | 286 | #### Training 287 | estimator.train(input_fn=train_input_fn, max_steps=FLAGS.train_steps) 288 | 289 | 290 | if __name__ == "__main__": 291 | app.run(main) 292 | -------------------------------------------------------------------------------- /train_gpu.py: -------------------------------------------------------------------------------- 1 | """Pretraining on GPUs.""" 2 | from __future__ import absolute_import 3 | from __future__ import division 4 | from __future__ import print_function 5 | 6 | import os, sys 7 | import math 8 | import json 9 | import time 10 | import numpy as np 11 | 12 | from absl import flags 13 | import absl.logging as _logging # pylint: disable=unused-import 14 | 15 | import tensorflow as tf 16 | 17 | import data_utils 18 | import model_utils 19 | from gpu_utils import assign_to_gpu, average_grads_and_vars 20 | import function_builder 21 | 22 | 23 | # GPU config 24 | flags.DEFINE_integer("num_hosts", default=1, 25 | help="Number of hosts") 26 | flags.DEFINE_integer("num_core_per_host", default=8, 27 | help="Number of cores per host") 28 | flags.DEFINE_bool("use_tpu", default=False, 29 | help="Whether to use TPUs for training.") 30 | 31 | # Experiment (data/checkpoint/directory) config 32 | flags.DEFINE_integer("num_passes", default=1, 33 | help="Number of passed used for training.") 34 | flags.DEFINE_string("record_info_dir", default=None, 35 | help="Path to local directory containing `record_info-lm.json`.") 36 | flags.DEFINE_string("model_dir", default=None, 37 | help="Estimator model_dir.") 38 | flags.DEFINE_string("init_checkpoint", default=None, 39 | help="checkpoint path for initializing the model.") 40 | 41 | # Optimization config 42 | flags.DEFINE_float("learning_rate", default=1e-4, 43 | help="Maximum learning rate.") 44 | flags.DEFINE_float("clip", default=1.0, 45 | help="Gradient clipping value.") 46 | # for cosine decay 47 | flags.DEFINE_float("min_lr_ratio", default=0.001, 48 | help="Minimum ratio learning rate.") 49 | flags.DEFINE_integer("warmup_steps", default=0, 50 | help="Number of steps for linear lr warmup.") 51 | flags.DEFINE_float("adam_epsilon", default=1e-8, 52 | help="Adam epsilon") 53 | flags.DEFINE_string("decay_method", default="poly", 54 | help="poly or cos") 55 | flags.DEFINE_float("weight_decay", default=0.0, 56 | help="weight decay") 57 | 58 | # Training config 59 | flags.DEFINE_integer("train_batch_size", default=16, 60 | help="Size of train batch.") 61 | flags.DEFINE_integer("train_steps", default=100000, 62 | help="Total number of training steps.") 63 | flags.DEFINE_integer("iterations", default=1000, 64 | help="Number of iterations per repeat loop.") 65 | flags.DEFINE_integer("save_steps", default=None, 66 | help="number of steps for model checkpointing.") 67 | 68 | # Data config 69 | flags.DEFINE_integer('seq_len', default=0, 70 | help='Sequence length for pretraining.') 71 | flags.DEFINE_integer('reuse_len', default=0, 72 | help="How many tokens to be reused in the next batch. " 73 | "Could be half of seq_len") 74 | flags.DEFINE_bool("bi_data", default=True, 75 | help="Use bidirectional data streams, i.e., forward & backward.") 76 | flags.DEFINE_integer("mask_alpha", default=6, 77 | help="How many tokens to form a group.") 78 | flags.DEFINE_integer("mask_beta", default=1, 79 | help="How many tokens to mask within each group.") 80 | flags.DEFINE_integer("num_predict", default=None, 81 | help="Number of tokens to predict in partial prediction.") 82 | flags.DEFINE_integer('perm_size', default=None, 83 | help='perm size.') 84 | flags.DEFINE_bool("uncased", False, 85 | help="Use uncased inputs or not.") 86 | flags.DEFINE_integer("n_token", 32000, help="Vocab size") 87 | 88 | # Model config 89 | flags.DEFINE_integer("mem_len", default=0, 90 | help="Number of steps to cache") 91 | flags.DEFINE_bool("same_length", default=False, 92 | help="Same length attention") 93 | flags.DEFINE_integer("clamp_len", default=-1, 94 | help="Clamp length") 95 | 96 | flags.DEFINE_integer("n_layer", default=6, 97 | help="Number of layers.") 98 | flags.DEFINE_integer("d_model", default=32, 99 | help="Dimension of the model.") 100 | flags.DEFINE_integer("d_embed", default=32, 101 | help="Dimension of the embeddings.") 102 | flags.DEFINE_integer("n_head", default=4, 103 | help="Number of attention heads.") 104 | flags.DEFINE_integer("d_head", default=8, 105 | help="Dimension of each attention head.") 106 | flags.DEFINE_integer("d_inner", default=32, 107 | help="Dimension of inner hidden size in positionwise feed-forward.") 108 | flags.DEFINE_float("dropout", default=0.0, 109 | help="Dropout rate.") 110 | flags.DEFINE_float("dropatt", default=0.0, 111 | help="Attention dropout rate.") 112 | flags.DEFINE_bool("untie_r", default=False, 113 | help="Untie r_w_bias and r_r_bias") 114 | flags.DEFINE_string("summary_type", default="last", 115 | help="Method used to summarize a sequence into a compact vector.") 116 | flags.DEFINE_string("ff_activation", default="relu", 117 | help="Activation type used in position-wise feed-forward.") 118 | flags.DEFINE_bool("use_bfloat16", False, 119 | help="Whether to use bfloat16.") 120 | 121 | # Parameter initialization 122 | flags.DEFINE_enum("init", default="normal", 123 | enum_values=["normal", "uniform"], 124 | help="Initialization method.") 125 | flags.DEFINE_float("init_std", default=0.02, 126 | help="Initialization std when init is normal.") 127 | flags.DEFINE_float("init_range", default=0.1, 128 | help="Initialization std when init is uniform.") 129 | 130 | 131 | FLAGS = flags.FLAGS 132 | 133 | 134 | def get_model_fn(): 135 | def model_fn(features, labels, mems, is_training): 136 | #### Get loss from inputs 137 | total_loss, new_mems, monitor_dict = function_builder.get_loss( 138 | FLAGS, features, labels, mems, is_training) 139 | 140 | #### Check model parameters 141 | num_params = sum([np.prod(v.shape) for v in tf.trainable_variables()]) 142 | tf.logging.info('#params: {}'.format(num_params)) 143 | 144 | # GPU 145 | assert is_training 146 | all_vars = tf.trainable_variables() 147 | grads = tf.gradients(total_loss, all_vars) 148 | grads_and_vars = list(zip(grads, all_vars)) 149 | 150 | return total_loss, new_mems, grads_and_vars 151 | 152 | return model_fn 153 | 154 | 155 | def single_core_graph(is_training, features, mems): 156 | model_fn = get_model_fn() 157 | 158 | model_ret = model_fn( 159 | features=features, 160 | labels=None, 161 | mems=mems, 162 | is_training=is_training) 163 | 164 | return model_ret 165 | 166 | 167 | def create_mems_tf(bsz_per_core): 168 | mems = [tf.placeholder(dtype=tf.float32, 169 | shape=[FLAGS.mem_len, bsz_per_core, FLAGS.d_model]) 170 | for layer in range(FLAGS.n_layer)] 171 | 172 | return mems 173 | 174 | 175 | def initialize_mems_np(bsz_per_core): 176 | mems_np = [np.zeros(shape=[FLAGS.mem_len, bsz_per_core, FLAGS.d_model], 177 | dtype=np.float32) 178 | for layer in range(FLAGS.n_layer)] 179 | 180 | return mems_np 181 | 182 | 183 | def train(ps_device): 184 | ##### Get input function and model function 185 | 186 | train_input_fn, record_info_dict = data_utils.get_input_fn( 187 | tfrecord_dir=FLAGS.record_info_dir, 188 | split="train", 189 | bsz_per_host=FLAGS.train_batch_size, 190 | seq_len=FLAGS.seq_len, 191 | reuse_len=FLAGS.reuse_len, 192 | bi_data=FLAGS.bi_data, 193 | num_hosts=1, 194 | num_core_per_host=1, # set to one no matter how many GPUs 195 | perm_size=FLAGS.perm_size, 196 | mask_alpha=FLAGS.mask_alpha, 197 | mask_beta=FLAGS.mask_beta, 198 | uncased=FLAGS.uncased, 199 | num_passes=FLAGS.num_passes, 200 | use_bfloat16=FLAGS.use_bfloat16, 201 | num_predict=FLAGS.num_predict) 202 | 203 | # for key, info in record_info_dict.items(): 204 | tf.logging.info("num of batches {}".format(record_info_dict["num_batch"])) 205 | 206 | ##### Create input tensors / placeholders 207 | bsz_per_core = FLAGS.train_batch_size // FLAGS.num_core_per_host 208 | 209 | params = { 210 | "batch_size": FLAGS.train_batch_size # the whole batch 211 | } 212 | train_set = train_input_fn(params) 213 | 214 | example = train_set.make_one_shot_iterator().get_next() 215 | 216 | if FLAGS.num_core_per_host > 1: 217 | examples = [{} for _ in range(FLAGS.num_core_per_host)] 218 | for key in example.keys(): 219 | vals = tf.split(example[key], FLAGS.num_core_per_host, 0) 220 | for device_id in range(FLAGS.num_core_per_host): 221 | examples[device_id][key] = vals[device_id] 222 | else: 223 | examples = [example] 224 | 225 | ##### Create computational graph 226 | tower_mems, tower_losses, tower_new_mems, tower_grads_and_vars = [], [], [], [] 227 | 228 | for i in range(FLAGS.num_core_per_host): 229 | reuse = True if i > 0 else None 230 | with tf.device(assign_to_gpu(i, ps_device)), \ 231 | tf.variable_scope(tf.get_variable_scope(), reuse=reuse): 232 | 233 | # The mems for each tower is a dictionary 234 | mems_i = {} 235 | if FLAGS.mem_len: 236 | mems_i["mems"] = create_mems_tf(bsz_per_core) 237 | 238 | loss_i, new_mems_i, grads_and_vars_i = single_core_graph( 239 | is_training=True, 240 | features=examples[i], 241 | mems=mems_i) 242 | 243 | tower_mems.append(mems_i) 244 | tower_losses.append(loss_i) 245 | tower_new_mems.append(new_mems_i) 246 | tower_grads_and_vars.append(grads_and_vars_i) 247 | 248 | ## average losses and gradients across towers 249 | if len(tower_losses) > 1: 250 | loss = tf.add_n(tower_losses) / len(tower_losses) 251 | grads_and_vars = average_grads_and_vars(tower_grads_and_vars) 252 | else: 253 | loss = tower_losses[0] 254 | grads_and_vars = tower_grads_and_vars[0] 255 | 256 | ## get train op 257 | train_op, learning_rate, gnorm = model_utils.get_train_op(FLAGS, None, 258 | grads_and_vars=grads_and_vars) 259 | global_step = tf.train.get_global_step() 260 | 261 | ##### Training loop 262 | # initialize mems 263 | tower_mems_np = [] 264 | for i in range(FLAGS.num_core_per_host): 265 | mems_i_np = {} 266 | for key in tower_mems[i].keys(): 267 | mems_i_np[key] = initialize_mems_np(bsz_per_core) 268 | tower_mems_np.append(mems_i_np) 269 | 270 | saver = tf.train.Saver() 271 | 272 | gpu_options = tf.GPUOptions(allow_growth=True) 273 | 274 | model_utils.init_from_checkpoint(FLAGS, global_vars=True) 275 | 276 | with tf.Session(config=tf.ConfigProto(allow_soft_placement=True, 277 | gpu_options=gpu_options)) as sess: 278 | sess.run(tf.global_variables_initializer()) 279 | 280 | fetches = [loss, tower_new_mems, global_step, gnorm, learning_rate, train_op] 281 | 282 | total_loss, prev_step = 0., -1 283 | while True: 284 | feed_dict = {} 285 | for i in range(FLAGS.num_core_per_host): 286 | for key in tower_mems_np[i].keys(): 287 | for m, m_np in zip(tower_mems[i][key], tower_mems_np[i][key]): 288 | feed_dict[m] = m_np 289 | 290 | fetched = sess.run(fetches, feed_dict=feed_dict) 291 | 292 | loss_np, tower_mems_np, curr_step = fetched[:3] 293 | total_loss += loss_np 294 | 295 | if curr_step > 0 and curr_step % FLAGS.iterations == 0: 296 | curr_loss = total_loss / (curr_step - prev_step) 297 | tf.logging.info("[{}] | gnorm {:.2f} lr {:8.6f} " 298 | "| loss {:.2f} | pplx {:>7.2f}, bpc {:>7.4f}".format( 299 | curr_step, fetched[-3], fetched[-2], 300 | curr_loss, math.exp(curr_loss), curr_loss / math.log(2))) 301 | total_loss, prev_step = 0., curr_step 302 | 303 | if curr_step > 0 and curr_step % FLAGS.save_steps == 0: 304 | save_path = os.path.join(FLAGS.model_dir, "model.ckpt") 305 | saver.save(sess, save_path) 306 | tf.logging.info("Model saved in path: {}".format(save_path)) 307 | 308 | if curr_step >= FLAGS.train_steps: 309 | break 310 | 311 | 312 | def main(unused_argv): 313 | del unused_argv # Unused 314 | 315 | tf.logging.set_verbosity(tf.logging.INFO) 316 | 317 | # Get corpus info 318 | FLAGS.n_token = data_utils.VOCAB_SIZE 319 | tf.logging.info("n_token {}".format(FLAGS.n_token)) 320 | 321 | if not tf.gfile.Exists(FLAGS.model_dir): 322 | tf.gfile.MakeDirs(FLAGS.model_dir) 323 | 324 | train("/gpu:0") 325 | 326 | 327 | if __name__ == "__main__": 328 | tf.app.run() 329 | -------------------------------------------------------------------------------- /xlnet.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | 5 | import json 6 | import os 7 | import tensorflow as tf 8 | import modeling 9 | 10 | 11 | def _get_initializer(FLAGS): 12 | """Get variable intializer.""" 13 | if FLAGS.init == "uniform": 14 | initializer = tf.initializers.random_uniform( 15 | minval=-FLAGS.init_range, 16 | maxval=FLAGS.init_range, 17 | seed=None) 18 | elif FLAGS.init == "normal": 19 | initializer = tf.initializers.random_normal( 20 | stddev=FLAGS.init_std, 21 | seed=None) 22 | else: 23 | raise ValueError("Initializer {} not supported".format(FLAGS.init)) 24 | return initializer 25 | 26 | 27 | class XLNetConfig(object): 28 | """XLNetConfig contains hyperparameters that are specific to a model checkpoint; 29 | i.e., these hyperparameters should be the same between 30 | pretraining and finetuning. 31 | 32 | The following hyperparameters are defined: 33 | n_layer: int, the number of layers. 34 | d_model: int, the hidden size. 35 | n_head: int, the number of attention heads. 36 | d_head: int, the dimension size of each attention head. 37 | d_inner: int, the hidden size in feed-forward layers. 38 | ff_activation: str, "relu" or "gelu". 39 | untie_r: bool, whether to untie the biases in attention. 40 | n_token: int, the vocab size. 41 | """ 42 | 43 | def __init__(self, FLAGS=None, json_path=None): 44 | """Constructing an XLNetConfig. 45 | One of FLAGS or json_path should be provided.""" 46 | 47 | assert FLAGS is not None or json_path is not None 48 | 49 | self.keys = ["n_layer", "d_model", "n_head", "d_head", "d_inner", 50 | "ff_activation", "untie_r", "n_token"] 51 | 52 | if FLAGS is not None: 53 | self.init_from_flags(FLAGS) 54 | 55 | if json_path is not None: 56 | self.init_from_json(json_path) 57 | 58 | def init_from_flags(self, FLAGS): 59 | for key in self.keys: 60 | setattr(self, key, getattr(FLAGS, key)) 61 | 62 | def init_from_json(self, json_path): 63 | with tf.gfile.Open(json_path) as f: 64 | json_data = json.load(f) 65 | for key in self.keys: 66 | setattr(self, key, json_data[key]) 67 | 68 | def to_json(self, json_path): 69 | """Save XLNetConfig to a json file.""" 70 | json_data = {} 71 | for key in self.keys: 72 | json_data[key] = getattr(self, key) 73 | 74 | json_dir = os.path.dirname(json_path) 75 | if not tf.gfile.Exists(json_dir): 76 | tf.gfile.MakeDirs(json_dir) 77 | with tf.gfile.Open(json_path, "w") as f: 78 | json.dump(json_data, f, indent=4, sort_keys=True) 79 | 80 | 81 | def create_run_config(is_training, is_finetune, FLAGS): 82 | kwargs = dict( 83 | is_training=is_training, 84 | use_tpu=FLAGS.use_tpu, 85 | use_bfloat16=FLAGS.use_bfloat16, 86 | dropout=FLAGS.dropout, 87 | dropatt=FLAGS.dropatt, 88 | init=FLAGS.init, 89 | init_range=FLAGS.init_range, 90 | init_std=FLAGS.init_std, 91 | clamp_len=FLAGS.clamp_len) 92 | 93 | if not is_finetune: 94 | kwargs.update(dict( 95 | mem_len=FLAGS.mem_len, 96 | reuse_len=FLAGS.reuse_len, 97 | bi_data=FLAGS.bi_data, 98 | clamp_len=FLAGS.clamp_len, 99 | same_length=FLAGS.same_length)) 100 | 101 | return RunConfig(**kwargs) 102 | 103 | 104 | class RunConfig(object): 105 | """RunConfig contains hyperparameters that could be different 106 | between pretraining and finetuning. 107 | These hyperparameters can also be changed from run to run. 108 | We store them separately from XLNetConfig for flexibility. 109 | """ 110 | 111 | def __init__(self, is_training, use_tpu, use_bfloat16, dropout, dropatt, 112 | init="normal", init_range=0.1, init_std=0.02, mem_len=None, 113 | reuse_len=None, bi_data=False, clamp_len=-1, same_length=False): 114 | """ 115 | Args: 116 | is_training: bool, whether in training mode. 117 | use_tpu: bool, whether TPUs are used. 118 | use_bfloat16: bool, use bfloat16 instead of float32. 119 | dropout: float, dropout rate. 120 | dropatt: float, dropout rate on attention probabilities. 121 | init: str, the initialization scheme, either "normal" or "uniform". 122 | init_range: float, initialize the parameters with a uniform distribution 123 | in [-init_range, init_range]. Only effective when init="uniform". 124 | init_std: float, initialize the parameters with a normal distribution 125 | with mean 0 and stddev init_std. Only effective when init="normal". 126 | mem_len: int, the number of tokens to cache. 127 | reuse_len: int, the number of tokens in the currect batch to be cached 128 | and reused in the future. 129 | bi_data: bool, whether to use bidirectional input pipeline. 130 | Usually set to True during pretraining and False during finetuning. 131 | clamp_len: int, clamp all relative distances larger than clamp_len. 132 | -1 means no clamping. 133 | same_length: bool, whether to use the same attention length for each token. 134 | """ 135 | 136 | self.init = init 137 | self.init_range = init_range 138 | self.init_std = init_std 139 | self.is_training = is_training 140 | self.dropout = dropout 141 | self.dropatt = dropatt 142 | self.use_tpu = use_tpu 143 | self.use_bfloat16 = use_bfloat16 144 | self.mem_len = mem_len 145 | self.reuse_len = reuse_len 146 | self.bi_data = bi_data 147 | self.clamp_len = clamp_len 148 | self.same_length = same_length 149 | 150 | 151 | class XLNetModel(object): 152 | """A wrapper of the XLNet model used during both pretraining and finetuning.""" 153 | 154 | def __init__(self, xlnet_config, run_config, input_ids, seg_ids, input_mask, 155 | mems=None, perm_mask=None, target_mapping=None, inp_q=None, 156 | **kwargs): 157 | """ 158 | Args: 159 | xlnet_config: XLNetConfig, 160 | run_config: RunConfig, 161 | input_ids: int32 Tensor in shape [len, bsz], the input token IDs. 162 | seg_ids: int32 Tensor in shape [len, bsz], the input segment IDs. 163 | input_mask: float32 Tensor in shape [len, bsz], the input mask. 164 | 0 for real tokens and 1 for padding. 165 | mems: a list of float32 Tensors in shape [mem_len, bsz, d_model], memory 166 | from previous batches. The length of the list equals n_layer. 167 | If None, no memory is used. 168 | perm_mask: float32 Tensor in shape [len, len, bsz]. 169 | If perm_mask[i, j, k] = 0, i attend to j in batch k; 170 | if perm_mask[i, j, k] = 1, i does not attend to j in batch k. 171 | If None, each position attends to all the others. 172 | target_mapping: float32 Tensor in shape [num_predict, len, bsz]. 173 | If target_mapping[i, j, k] = 1, the i-th predict in batch k is 174 | on the j-th token. 175 | Only used during pretraining for partial prediction. 176 | Set to None during finetuning. 177 | inp_q: float32 Tensor in shape [len, bsz]. 178 | 1 for tokens with losses and 0 for tokens without losses. 179 | Only used during pretraining for two-stream attention. 180 | Set to None during finetuning. 181 | """ 182 | 183 | initializer = _get_initializer(run_config) 184 | 185 | tfm_args = dict( 186 | n_token=xlnet_config.n_token, 187 | initializer=initializer, 188 | attn_type="bi", 189 | n_layer=xlnet_config.n_layer, 190 | d_model=xlnet_config.d_model, 191 | n_head=xlnet_config.n_head, 192 | d_head=xlnet_config.d_head, 193 | d_inner=xlnet_config.d_inner, 194 | ff_activation=xlnet_config.ff_activation, 195 | untie_r=xlnet_config.untie_r, 196 | 197 | is_training=run_config.is_training, 198 | use_bfloat16=run_config.use_bfloat16, 199 | use_tpu=run_config.use_tpu, 200 | dropout=run_config.dropout, 201 | dropatt=run_config.dropatt, 202 | 203 | mem_len=run_config.mem_len, 204 | reuse_len=run_config.reuse_len, 205 | bi_data=run_config.bi_data, 206 | clamp_len=run_config.clamp_len, 207 | same_length=run_config.same_length 208 | ) 209 | 210 | input_args = dict( 211 | inp_k=input_ids, 212 | seg_id=seg_ids, 213 | input_mask=input_mask, 214 | mems=mems, 215 | perm_mask=perm_mask, 216 | target_mapping=target_mapping, 217 | inp_q=inp_q) 218 | tfm_args.update(input_args) 219 | 220 | with tf.variable_scope("model", reuse=tf.AUTO_REUSE): 221 | (self.output, self.new_mems, self.lookup_table 222 | ) = modeling.transformer_xl(**tfm_args) 223 | 224 | self.input_mask = input_mask 225 | self.initializer = initializer 226 | self.xlnet_config = xlnet_config 227 | self.run_config = run_config 228 | 229 | def get_pooled_out(self, summary_type, use_summ_proj=True): 230 | """ 231 | Args: 232 | summary_type: str, "last", "first", "mean", or "attn". The method 233 | to pool the input to get a vector representation. 234 | use_summ_proj: bool, whether to use a linear projection during pooling. 235 | 236 | Returns: 237 | float32 Tensor in shape [bsz, d_model], the pooled representation. 238 | """ 239 | 240 | xlnet_config = self.xlnet_config 241 | run_config = self.run_config 242 | 243 | with tf.variable_scope("model", reuse=tf.AUTO_REUSE): 244 | summary = modeling.summarize_sequence( 245 | summary_type=summary_type, 246 | hidden=self.output, 247 | d_model=xlnet_config.d_model, 248 | n_head=xlnet_config.n_head, 249 | d_head=xlnet_config.d_head, 250 | dropout=run_config.dropout, 251 | dropatt=run_config.dropatt, 252 | is_training=run_config.is_training, 253 | input_mask=self.input_mask, 254 | initializer=self.initializer, 255 | use_proj=use_summ_proj) 256 | 257 | return summary 258 | 259 | def get_sequence_output(self): 260 | """ 261 | Returns: 262 | float32 Tensor in shape [len, bsz, d_model]. The last layer hidden 263 | representation of XLNet. 264 | """ 265 | 266 | return self.output 267 | 268 | def get_new_memory(self): 269 | """ 270 | Returns: 271 | list of float32 Tensors in shape [mem_len, bsz, d_model], the new 272 | memory that concatenates the previous memory with the current input 273 | representations. 274 | The length of the list equals n_layer. 275 | """ 276 | return self.new_mems 277 | 278 | def get_embedding_table(self): 279 | """ 280 | Returns: 281 | float32 Tensor in shape [n_token, d_model]. The embedding lookup table. 282 | Used for tying embeddings between input and output layers. 283 | """ 284 | return self.lookup_table 285 | 286 | def get_initializer(self): 287 | """ 288 | Returns: 289 | A tf initializer. Used to initialize variables in layers on top of XLNet. 290 | """ 291 | return self.initializer 292 | 293 | --------------------------------------------------------------------------------