├── tmp_data └── sample.tfrecords ├── requirements.txt ├── bert_config.json ├── vocab.py ├── __init__.py ├── CONTRIBUTING.md ├── .gitignore ├── README.md ├── sample_text.txt ├── vocab.txt ├── optimization_gpu.py ├── LICENSE ├── tokenization.py ├── create_pretraining_data.py ├── run_pretraining_gpu_v2.py ├── run_pretraining_gpu.py └── modeling.py /tmp_data/sample.tfrecords: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/guotong1988/BERT-pre-training/HEAD/tmp_data/sample.tfrecords -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | tensorflow >= 1.11.0 # CPU Version of TensorFlow. 2 | # tensorflow-gpu >= 1.11.0 # GPU version of TensorFlow. 3 | -------------------------------------------------------------------------------- /bert_config.json: -------------------------------------------------------------------------------- 1 | { 2 | "attention_probs_dropout_prob": 0.1, 3 | "hidden_act": "gelu", 4 | "hidden_dropout_prob": 0.1, 5 | "hidden_size": 128, 6 | "initializer_range": 0.02, 7 | "intermediate_size": 256, 8 | "max_position_embeddings": 512, 9 | "num_attention_heads": 8, 10 | "num_hidden_layers": 6, 11 | "type_vocab_size": 2, 12 | "vocab_size": 453 13 | } 14 | -------------------------------------------------------------------------------- /vocab.py: -------------------------------------------------------------------------------- 1 | from nltk.tokenize import wordpunct_tokenize 2 | f = open("sample_text.txt",mode="r",encoding="utf-8") 3 | f2 = open("vocab.txt",mode="w",encoding="utf-8") 4 | lines = f.readlines() 5 | vocab = set() 6 | for line in lines: 7 | word_list = wordpunct_tokenize(line) 8 | for word in word_list: 9 | vocab.add(word) 10 | 11 | f2.write("[PAD]\n") 12 | f2.write("[UNK]\n") 13 | f2.write("[CLS]\n") 14 | f2.write("[SEP]\n") 15 | f2.write("[MASK]\n") 16 | for word in list(vocab): 17 | f2.write(word) 18 | f2.write("\n") 19 | 20 | f2.close() 21 | f.close() -------------------------------------------------------------------------------- /__init__.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # Copyright 2018 The Google AI Language Team Authors. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # How to Contribute 2 | 3 | BERT needs to maintain permanent compatibility with the pre-trained model files, 4 | so we do not plan to make any major changes to this library (other than what was 5 | promised in the README). However, we can accept small patches related to 6 | re-factoring and documentation. To submit contributes, there are just a few 7 | small guidelines you need to follow. 8 | 9 | ## Contributor License Agreement 10 | 11 | Contributions to this project must be accompanied by a Contributor License 12 | Agreement. You (or your employer) retain the copyright to your contribution; 13 | this simply gives us permission to use and redistribute your contributions as 14 | part of the project. Head over to to see 15 | your current agreements on file or to sign a new one. 16 | 17 | You generally only need to submit a CLA once, so if you've already submitted one 18 | (even if it was for a different project), you probably don't need to do it 19 | again. 20 | 21 | ## Code reviews 22 | 23 | All submissions, including submissions by project members, require review. We 24 | use GitHub pull requests for this purpose. Consult 25 | [GitHub Help](https://help.github.com/articles/about-pull-requests/) for more 26 | information on using pull requests. 27 | 28 | ## Community Guidelines 29 | 30 | This project follows 31 | [Google's Open Source Community Guidelines](https://opensource.google.com/conduct/). 32 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Initially taken from Github's Python gitignore file 2 | 3 | # Byte-compiled / optimized / DLL files 4 | __pycache__/ 5 | *.py[cod] 6 | *$py.class 7 | 8 | # C extensions 9 | *.so 10 | 11 | # Distribution / packaging 12 | .Python 13 | build/ 14 | develop-eggs/ 15 | dist/ 16 | downloads/ 17 | eggs/ 18 | .eggs/ 19 | lib/ 20 | lib64/ 21 | parts/ 22 | sdist/ 23 | var/ 24 | wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | .hypothesis/ 51 | .pytest_cache/ 52 | 53 | # Translations 54 | *.mo 55 | *.pot 56 | 57 | # Django stuff: 58 | *.log 59 | local_settings.py 60 | db.sqlite3 61 | 62 | # Flask stuff: 63 | instance/ 64 | .webassets-cache 65 | 66 | # Scrapy stuff: 67 | .scrapy 68 | 69 | # Sphinx documentation 70 | docs/_build/ 71 | 72 | # PyBuilder 73 | target/ 74 | 75 | # Jupyter Notebook 76 | .ipynb_checkpoints 77 | 78 | # IPython 79 | profile_default/ 80 | ipython_config.py 81 | 82 | # pyenv 83 | .python-version 84 | 85 | # celery beat schedule file 86 | celerybeat-schedule 87 | 88 | # SageMath parsed files 89 | *.sage.py 90 | 91 | # Environments 92 | .env 93 | .venv 94 | env/ 95 | venv/ 96 | ENV/ 97 | env.bak/ 98 | venv.bak/ 99 | 100 | # Spyder project settings 101 | .spyderproject 102 | .spyproject 103 | 104 | # Rope project settings 105 | .ropeproject 106 | 107 | # mkdocs documentation 108 | /site 109 | 110 | # mypy 111 | .mypy_cache/ 112 | .dmypy.json 113 | dmypy.json 114 | 115 | # Pyre type checker 116 | .pyre/ 117 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # BERT MULTI-GPU PRE-TRAIN ON ONE MACHINE WITHOUT HOROVOD (Data Parallelism) 2 | 3 | [![LICENSE](https://img.shields.io/badge/license-Anti%20996-blue.svg)](https://github.com/996icu/996.ICU/blob/master/LICENSE) 4 | [![996.icu](https://img.shields.io/badge/link-996.icu-red.svg)](https://996.icu) 5 | 6 | BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding 7 | 8 | # REASONABLE / PRINCIPLE 9 | 10 | More gpu means more data in a batch (, batch size is larger). And the gradients of a batch data is averaged for back-propagation. 11 | 12 | If the sum learning rate of one batch is fixed, then the learning rate of one data is smaller, when batch size is larger. 13 | 14 | If the learning rate of one data is fixed, then the sum learning rate of one batch is larger, when batch size is larger. 15 | 16 | **It seems:** More gpu --> Larger sum learning rate of one batch --> Faster training. 17 | 18 | But as the learning rate can only be in the range of 1e-4 to 3e-5, [**train_samples_per_second**](https://github.com/guotong1988/LLM-post-training/issues/1) becomes the evaluation metric for speed. 19 | 20 | # WHATS NEW 21 | 22 | Using 1-GPU (100 batch size) vs using 4-GPU (400 batch size) for the same learning rate (0.00001) and same pre-training steps (1,000,000) will be no difference of 0.1% in downstream task accuracy. 23 | 24 | # REQUIREMENT 25 | 26 | python 3 27 | 28 | tensorflow 1.14 - 1.15 29 | 30 | # TRAINING 31 | 32 | 0, edit the input and output file name in `create_pretraining_data.py` and `run_pretraining_gpu.py` 33 | 34 | 1, run `create_pretraining_data.py` 35 | 36 | 2, run `run_pretraining_gpu_v2.py` 37 | 38 | # PARAMETERS 39 | 40 | Edit `n_gpus` in `run_pretraining_gpu_v2.py` 41 | 42 | 43 | # DATA 44 | 45 | In `sample_text.txt`, sentence is end by `\n`, paragraph is splitted by empty line. 46 | 47 | # EXPERIMENT RESULT ON DOWNSTREAM TASKS 48 | 49 | Quora question pairs English dataset, 50 | 51 | Official BERT: ACC 91.2, AUC 96.9 52 | 53 | This BERT with pretrain loss 2.05: ACC 90.1, AUC 96.3 54 | 55 | # NOTE 56 | 57 | ### 1) 58 | For `HierarchicalCopyAllReduce` `MirroredStrategy`, `global_step/sec` shows the sum of multi gpus' steps. 59 | ### 2) 60 | `batch_size` is the `batch_size` per GPU, not the `global_batch_size` 61 | 62 | 63 | -------------------------------------------------------------------------------- /sample_text.txt: -------------------------------------------------------------------------------- 1 | This text is included to make sure Unicode is handled properly: 力加勝北区ᴵᴺᵀᵃছজটডণত 2 | Text should be one-sentence-per-line, with empty lines between documents. 3 | This sample text is public domain and was randomly selected from Project Guttenberg. 4 | 5 | The rain had only ceased with the gray streaks of morning at Blazing Star, and the settlement awoke to a moral sense of cleanliness, and the finding of forgotten knives, tin cups, and smaller camp utensils, where the heavy showers had washed away the debris and dust heaps before the cabin doors. 6 | Indeed, it was recorded in Blazing Star that a fortunate early riser had once picked up on the highway a solid chunk of gold quartz which the rain had freed from its incumbering soil, and washed into immediate and glittering popularity. 7 | Possibly this may have been the reason why early risers in that locality, during the rainy season, adopted a thoughtful habit of body, and seldom lifted their eyes to the rifted or india-ink washed skies above them. 8 | "Cass" Beard had risen early that morning, but not with a view to discovery. 9 | A leak in his cabin roof,--quite consistent with his careless, improvident habits,--had roused him at 4 A. M., with a flooded "bunk" and wet blankets. 10 | The chips from his wood pile refused to kindle a fire to dry his bed-clothes, and he had recourse to a more provident neighbor's to supply the deficiency. 11 | This was nearly opposite. 12 | Mr. Cassius crossed the highway, and stopped suddenly. 13 | Something glittered in the nearest red pool before him. 14 | Gold, surely! 15 | But, wonderful to relate, not an irregular, shapeless fragment of crude ore, fresh from Nature's crucible, but a bit of jeweler's handicraft in the form of a plain gold ring. 16 | Looking at it more attentively, he saw that it bore the inscription, "May to Cass." 17 | Like most of his fellow gold-seekers, Cass was superstitious. 18 | 19 | The fountain of classic wisdom, Hypatia herself. 20 | As the ancient sage--the name is unimportant to a monk--pumped water nightly that he might study by day, so I, the guardian of cloaks and parasols, at the sacred doors of her lecture-room, imbibe celestial knowledge. 21 | From my youth I felt in me a soul above the matter-entangled herd. 22 | She revealed to me the glorious fact, that I am a spark of Divinity itself. 23 | A fallen star, I am, sir!' continued he, pensively, stroking his lean stomach--'a fallen star!--fallen, if the dignity of philosophy will allow of the simile, among the hogs of the lower world--indeed, even into the hog-bucket itself. Well, after all, I will show you the way to the Archbishop's. 24 | There is a philosophic pleasure in opening one's treasures to the modest young. 25 | Perhaps you will assist me by carrying this basket of fruit?' And the little man jumped up, put his basket on Philammon's head, and trotted off up a neighbouring street. 26 | Philammon followed, half contemptuous, half wondering at what this philosophy might be, which could feed the self-conceit of anything so abject as his ragged little apish guide; 27 | but the novel roar and whirl of the street, the perpetual stream of busy faces, the line of curricles, palanquins, laden asses, camels, elephants, which met and passed him, and squeezed him up steps and into doorways, as they threaded their way through the great Moon-gate into the ample street beyond, drove everything from his mind but wondering curiosity, and a vague, helpless dread of that great living wilderness, more terrible than any dead wilderness of sand which he had left behind. 28 | Already he longed for the repose, the silence of the Laura--for faces which knew him and smiled upon him; but it was too late to turn back now. 29 | His guide held on for more than a mile up the great main street, crossed in the centre of the city, at right angles, by one equally magnificent, at each end of which, miles away, appeared, dim and distant over the heads of the living stream of passengers, the yellow sand-hills of the desert; 30 | while at the end of the vista in front of them gleamed the blue harbour, through a network of countless masts. 31 | At last they reached the quay at the opposite end of the street; 32 | and there burst on Philammon's astonished eyes a vast semicircle of blue sea, ringed with palaces and towers. 33 | He stopped involuntarily; and his little guide stopped also, and looked askance at the young monk, to watch the effect which that grand panorama should produce on him. 34 | -------------------------------------------------------------------------------- /vocab.txt: -------------------------------------------------------------------------------- 1 | [PAD] 2 | [UNK] 3 | [CLS] 4 | [SEP] 5 | [MASK] 6 | locality 7 | lean 8 | this 9 | great 10 | during 11 | street 12 | bit 13 | herd 14 | treasures 15 | upon 16 | washed 17 | attentively 18 | palanquins 19 | fragment 20 | refused 21 | morning 22 | dead 23 | reached 24 | handicraft 25 | line 26 | Possibly 27 | had 28 | living 29 | semicircle 30 | repose 31 | smiled 32 | surely 33 | my 34 | deficiency 35 | way 36 | but 37 | beyond 38 | philosophy 39 | thoughtful 40 | faces 41 | burst 42 | passed 43 | as 44 | season 45 | helpless 46 | unimportant 47 | chips 48 | shapeless 49 | asses 50 | astonished 51 | rainy 52 | panorama 53 | skies 54 | dread 55 | anything 56 | ragged 57 | From 58 | provident 59 | popularity 60 | itself 61 | dust 62 | public 63 | crude 64 | Star 65 | Something 66 | matter 67 | wonderful 68 | dignity 69 | flooded 70 | self 71 | fact 72 | fruit 73 | Indeed 74 | last 75 | cabin 76 | was 77 | empty 78 | Project 79 | reason 80 | moral 81 | conceit 82 | recorded 83 | end 84 | may 85 | suddenly 86 | view 87 | There 88 | - 89 | of 90 | over 91 | involuntarily 92 | As 93 | riser 94 | pile 95 | before 96 | body 97 | to 98 | angles 99 | star 100 | ringed 101 | highway 102 | head 103 | steps 104 | that 105 | followed 106 | more 107 | and 108 | drove 109 | careless 110 | fortunate 111 | it 112 | Perhaps 113 | or 114 | ?' 115 | glittered 116 | the 117 | quite 118 | philosophic 119 | above 120 | bed 121 | turn 122 | lifted 123 | busy 124 | than 125 | seekers 126 | appeared 127 | room 128 | should 129 | classic 130 | ink 131 | effect 132 | mile 133 | glittering 134 | he 135 | sacred 136 | assist 137 | I 138 | And 139 | make 140 | for 141 | dim 142 | towers 143 | a 144 | grand 145 | passengers 146 | consistent 147 | Nature 148 | met 149 | study 150 | countless 151 | He 152 | on 153 | stomach 154 | left 155 | from 156 | Already 157 | : 158 | main 159 | pensively 160 | text 161 | risen 162 | ! 163 | nearest 164 | finding 165 | discovery 166 | !-- 167 | Cassius 168 | wood 169 | imbibe 170 | its 171 | leak 172 | palaces 173 | A 174 | ring 175 | kindle 176 | once 177 | Hypatia 178 | clothes 179 | sand 180 | equally 181 | knowledge 182 | form 183 | handled 184 | entangled 185 | celestial 186 | trotted 187 | hills 188 | distant 189 | which 190 | incumbering 191 | stopped 192 | abject 193 | doorways 194 | awoke 195 | any 196 | ore 197 | doors 198 | now 199 | between 200 | solid 201 | right 202 | modest 203 | soul 204 | an 205 | rain 206 | ancient 207 | streaks 208 | miles 209 | also 210 | curiosity 211 | held 212 | neighbor 213 | wondering 214 | ; 215 | pumped 216 | put 217 | sure 218 | late 219 | felt 220 | everything 221 | This 222 | domain 223 | red 224 | wisdom 225 | produce 226 | forgotten 227 | masts 228 | cloaks 229 | after 230 | young 231 | elephants 232 | heaps 233 | basket 234 | nightly 235 | have 236 | while 237 | Gold 238 | sentence 239 | harbour 240 | novel 241 | stroking 242 | spark 243 | gold 244 | !' 245 | front 246 | At 247 | day 248 | pleasure 249 | inscription 250 | Looking 251 | been 252 | most 253 | lines 254 | lecture 255 | ample 256 | revealed 257 | looked 258 | tin 259 | camels 260 | 4 261 | stream 262 | askance 263 | Unicode 264 | Laura 265 | M 266 | magnificent 267 | opposite 268 | am 269 | me 270 | lower 271 | city 272 | bore 273 | half 274 | vast 275 | quartz 276 | roof 277 | her 278 | crucible 279 | allow 280 | little 281 | even 282 | indeed 283 | Well 284 | mind 285 | sea 286 | gate 287 | silence 288 | vague 289 | heavy 290 | supply 291 | bunk 292 | bucket 293 | " 294 | fresh 295 | perpetual 296 | ' 297 | continued 298 | vista 299 | crossed 300 | will 301 | 力加勝北区ᴵᴺᵀᵃছজটডণত 302 | only 303 | through 304 | fellow 305 | off 306 | freed 307 | nearly 308 | man 309 | But 310 | network 311 | why 312 | Divinity 313 | hogs 314 | contemptuous 315 | longed 316 | knew 317 | his 318 | at 319 | where 320 | Moon 321 | simile 322 | threaded 323 | jeweler 324 | feed 325 | May 326 | smaller 327 | pool 328 | glorious 329 | recourse 330 | carrying 331 | with 332 | ., 333 | The 334 | whirl 335 | s 336 | ceased 337 | superstitious 338 | them 339 | relate 340 | cups 341 | plain 342 | they 343 | knives 344 | roused 345 | in 346 | away 347 | him 348 | Like 349 | terrible 350 | immediate 351 | showers 352 | you 353 | squeezed 354 | be 355 | dry 356 | rifted 357 | india 358 | back 359 | Text 360 | blue 361 | all 362 | wet 363 | irregular 364 | wilderness 365 | ." 366 | Philammon 367 | gray 368 | their 369 | sense 370 | quay 371 | Beard 372 | too 373 | adopted 374 | Blazing 375 | roar 376 | among 377 | blankets 378 | . 379 | up 380 | camp 381 | eyes 382 | guide 383 | by 384 | curricles 385 | debris 386 | might 387 | parasols 388 | -- 389 | habits 390 | herself 391 | saw 392 | , 393 | fire 394 | included 395 | gleamed 396 | sage 397 | Mr 398 | settlement 399 | is 400 | fallen 401 | His 402 | world 403 | centre 404 | neighbouring 405 | yellow 406 | sample 407 | into 408 | show 409 | seldom 410 | utensils 411 | desert 412 | picked 413 | name 414 | per 415 | watch 416 | selected 417 | guardian 418 | habit 419 | properly 420 | ,-- 421 | randomly 422 | fountain 423 | if 424 | what 425 | opening 426 | Archbishop 427 | one 428 | soil 429 | cleanliness 430 | Cass 431 | risers 432 | documents 433 | improvident 434 | --' 435 | hog 436 | Guttenberg 437 | heads 438 | not 439 | water 440 | apish 441 | so 442 | sir 443 | jumped 444 | could 445 | behind 446 | laden 447 | each 448 | chunk 449 | early 450 | there 451 | monk 452 | She 453 | youth -------------------------------------------------------------------------------- /optimization_gpu.py: -------------------------------------------------------------------------------- 1 | """Functions and classes related to optimization (weight updates).""" 2 | 3 | from __future__ import absolute_import 4 | from __future__ import division 5 | from __future__ import print_function 6 | 7 | import re 8 | import tensorflow as tf 9 | from tensorflow.python.framework import ops 10 | from tensorflow.python.ops import control_flow_ops 11 | from tensorflow.python.ops import math_ops 12 | from tensorflow.python.training import optimizer 13 | from tensorflow.python.ops import state_ops 14 | from tensorflow.python.ops import resource_variable_ops 15 | 16 | 17 | def create_optimizer(loss, init_lr, num_train_steps, num_warmup_steps, use_tpu): 18 | """Creates an optimizer training op.""" 19 | global_step = tf.train.get_or_create_global_step() 20 | 21 | learning_rate = tf.constant(value=init_lr, shape=[], dtype=tf.float32) 22 | 23 | # Implements linear decay of the learning rate. 24 | learning_rate = tf.train.polynomial_decay( 25 | learning_rate, 26 | global_step, 27 | num_train_steps, 28 | end_learning_rate=0.0, 29 | power=1.0, 30 | cycle=False) 31 | 32 | # Implements linear warmup. I.e., if global_step < num_warmup_steps, the 33 | # learning rate will be `global_step/num_warmup_steps * init_lr`. 34 | if num_warmup_steps: 35 | global_steps_int = tf.cast(global_step, tf.int32) 36 | warmup_steps_int = tf.constant(num_warmup_steps, dtype=tf.int32) 37 | 38 | global_steps_float = tf.cast(global_steps_int, tf.float32) 39 | warmup_steps_float = tf.cast(warmup_steps_int, tf.float32) 40 | 41 | warmup_percent_done = global_steps_float / warmup_steps_float 42 | warmup_learning_rate = init_lr * warmup_percent_done 43 | 44 | is_warmup = tf.cast(global_steps_int < warmup_steps_int, tf.float32) 45 | learning_rate = ( 46 | (1.0 - is_warmup) * learning_rate + is_warmup * warmup_learning_rate) 47 | 48 | # It is recommended that you use this optimizer for fine tuning, since this 49 | # is how the model was trained (note that the Adam m/v variables are NOT 50 | # loaded from init_checkpoint.) 51 | optimizer = AdamWeightDecayOptimizer( 52 | learning_rate=learning_rate, 53 | weight_decay_rate=0.01, 54 | beta_1=0.9, 55 | beta_2=0.999, 56 | epsilon=1e-6, 57 | exclude_from_weight_decay=["LayerNorm", "layer_norm", "bias"]) 58 | 59 | if use_tpu: 60 | optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer) 61 | 62 | tvars = tf.trainable_variables() 63 | grads = tf.gradients(loss, tvars) 64 | 65 | # This is how the model was pre-trained. 66 | (grads, _) = tf.clip_by_global_norm(grads, clip_norm=1.0) 67 | 68 | train_op = optimizer.apply_gradients( 69 | zip(grads, tvars), global_step=global_step) 70 | 71 | # Normally the global step update is done inside of `apply_gradients`. 72 | # However, `AdamWeightDecayOptimizer` doesn't do this. But if you use 73 | # a different optimizer, you should probably take this line out. 74 | #new_global_step = global_step + 1 75 | #train_op = tf.group(train_op, [global_step.assign(new_global_step)]) 76 | return train_op 77 | 78 | class AdamWeightDecayOptimizer(optimizer.Optimizer): 79 | """A basic Adam optimizer that includes "correct" L2 weight decay.""" 80 | 81 | def __init__(self, 82 | learning_rate, 83 | weight_decay_rate=0.0, 84 | beta_1=0.9, 85 | beta_2=0.999, 86 | epsilon=1e-6, 87 | exclude_from_weight_decay=None, 88 | name="AdamWeightDecayOptimizer"): 89 | """Constructs a AdamWeightDecayOptimizer.""" 90 | super(AdamWeightDecayOptimizer, self).__init__(False, name) 91 | 92 | self.learning_rate = learning_rate 93 | self.weight_decay_rate = weight_decay_rate 94 | self.beta_1 = beta_1 95 | self.beta_2 = beta_2 96 | self.epsilon = epsilon 97 | self.exclude_from_weight_decay = exclude_from_weight_decay 98 | 99 | def _prepare(self): 100 | self.learning_rate_t = ops.convert_to_tensor( 101 | self.learning_rate, name='learning_rate') 102 | self.weight_decay_rate_t = ops.convert_to_tensor( 103 | self.weight_decay_rate, name='weight_decay_rate') 104 | self.beta_1_t = ops.convert_to_tensor(self.beta_1, name='beta_1') 105 | self.beta_2_t = ops.convert_to_tensor(self.beta_2, name='beta_2') 106 | self.epsilon_t = ops.convert_to_tensor(self.epsilon, name='epsilon') 107 | 108 | def _create_slots(self, var_list): 109 | for v in var_list: 110 | self._zeros_slot(v, 'm', self._name) 111 | self._zeros_slot(v, 'v', self._name) 112 | 113 | def _apply_dense(self, grad, var): 114 | learning_rate_t = math_ops.cast( 115 | self.learning_rate_t, var.dtype.base_dtype) 116 | beta_1_t = math_ops.cast(self.beta_1_t, var.dtype.base_dtype) 117 | beta_2_t = math_ops.cast(self.beta_2_t, var.dtype.base_dtype) 118 | epsilon_t = math_ops.cast(self.epsilon_t, var.dtype.base_dtype) 119 | weight_decay_rate_t = math_ops.cast( 120 | self.weight_decay_rate_t, var.dtype.base_dtype) 121 | 122 | m = self.get_slot(var, 'm') 123 | v = self.get_slot(var, 'v') 124 | 125 | # Standard Adam update. 126 | next_m = ( 127 | tf.multiply(beta_1_t, m) + 128 | tf.multiply(1.0 - beta_1_t, grad)) 129 | next_v = ( 130 | tf.multiply(beta_2_t, v) + tf.multiply(1.0 - beta_2_t, 131 | tf.square(grad))) 132 | 133 | update = next_m / (tf.sqrt(next_v) + epsilon_t) 134 | 135 | if self._do_use_weight_decay(var.name): 136 | update += weight_decay_rate_t * var 137 | 138 | update_with_lr = learning_rate_t * update 139 | 140 | next_param = var - update_with_lr 141 | 142 | return control_flow_ops.group(*[var.assign(next_param), 143 | m.assign(next_m), 144 | v.assign(next_v)]) 145 | 146 | def _resource_apply_dense(self, grad, var): 147 | learning_rate_t = math_ops.cast( 148 | self.learning_rate_t, var.dtype.base_dtype) 149 | beta_1_t = math_ops.cast(self.beta_1_t, var.dtype.base_dtype) 150 | beta_2_t = math_ops.cast(self.beta_2_t, var.dtype.base_dtype) 151 | epsilon_t = math_ops.cast(self.epsilon_t, var.dtype.base_dtype) 152 | weight_decay_rate_t = math_ops.cast( 153 | self.weight_decay_rate_t, var.dtype.base_dtype) 154 | 155 | m = self.get_slot(var, 'm') 156 | v = self.get_slot(var, 'v') 157 | 158 | # Standard Adam update. 159 | next_m = ( 160 | tf.multiply(beta_1_t, m) + 161 | tf.multiply(1.0 - beta_1_t, grad)) 162 | next_v = ( 163 | tf.multiply(beta_2_t, v) + tf.multiply(1.0 - beta_2_t, 164 | tf.square(grad))) 165 | 166 | update = next_m / (tf.sqrt(next_v) + epsilon_t) 167 | 168 | if self._do_use_weight_decay(var.name): 169 | update += weight_decay_rate_t * var 170 | 171 | update_with_lr = learning_rate_t * update 172 | 173 | next_param = var - update_with_lr 174 | 175 | return control_flow_ops.group(*[var.assign(next_param), 176 | m.assign(next_m), 177 | v.assign(next_v)]) 178 | 179 | def _apply_sparse_shared(self, grad, var, indices, scatter_add): 180 | learning_rate_t = math_ops.cast( 181 | self.learning_rate_t, var.dtype.base_dtype) 182 | beta_1_t = math_ops.cast(self.beta_1_t, var.dtype.base_dtype) 183 | beta_2_t = math_ops.cast(self.beta_2_t, var.dtype.base_dtype) 184 | epsilon_t = math_ops.cast(self.epsilon_t, var.dtype.base_dtype) 185 | weight_decay_rate_t = math_ops.cast( 186 | self.weight_decay_rate_t, var.dtype.base_dtype) 187 | 188 | m = self.get_slot(var, 'm') 189 | v = self.get_slot(var, 'v') 190 | 191 | m_t = state_ops.assign(m, m * beta_1_t, 192 | use_locking=self._use_locking) 193 | 194 | m_scaled_g_values = grad * (1 - beta_1_t) 195 | with ops.control_dependencies([m_t]): 196 | m_t = scatter_add(m, indices, m_scaled_g_values) 197 | 198 | v_scaled_g_values = (grad * grad) * (1 - beta_2_t) 199 | v_t = state_ops.assign(v, v * beta_2_t, use_locking=self._use_locking) 200 | with ops.control_dependencies([v_t]): 201 | v_t = scatter_add(v, indices, v_scaled_g_values) 202 | 203 | update = m_t / (math_ops.sqrt(v_t) + epsilon_t) 204 | 205 | if self._do_use_weight_decay(var.name): 206 | update += weight_decay_rate_t * var 207 | 208 | update_with_lr = learning_rate_t * update 209 | 210 | var_update = state_ops.assign_sub(var, 211 | update_with_lr, 212 | use_locking=self._use_locking) 213 | return control_flow_ops.group(*[var_update, m_t, v_t]) 214 | 215 | def _apply_sparse(self, grad, var): 216 | return self._apply_sparse_shared( 217 | grad.values, var, grad.indices, 218 | lambda x, i, v: state_ops.scatter_add( # pylint: disable=g-long-lambda 219 | x, i, v, use_locking=self._use_locking)) 220 | 221 | def _resource_scatter_add(self, x, i, v): 222 | with ops.control_dependencies( 223 | [resource_variable_ops.resource_scatter_add( 224 | x.handle, i, v)]): 225 | return x.value() 226 | 227 | def _resource_apply_sparse(self, grad, var, indices): 228 | return self._apply_sparse_shared( 229 | grad, var, indices, self._resource_scatter_add) 230 | 231 | def _do_use_weight_decay(self, param_name): 232 | """Whether to use L2 weight decay for `param_name`.""" 233 | if not self.weight_decay_rate: 234 | return False 235 | if self.exclude_from_weight_decay: 236 | for r in self.exclude_from_weight_decay: 237 | if re.search(r, param_name) is not None: 238 | return False 239 | return True 240 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | 2 | Apache License 3 | Version 2.0, January 2004 4 | http://www.apache.org/licenses/ 5 | 6 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 7 | 8 | 1. Definitions. 9 | 10 | "License" shall mean the terms and conditions for use, reproduction, 11 | and distribution as defined by Sections 1 through 9 of this document. 12 | 13 | "Licensor" shall mean the copyright owner or entity authorized by 14 | the copyright owner that is granting the License. 15 | 16 | "Legal Entity" shall mean the union of the acting entity and all 17 | other entities that control, are controlled by, or are under common 18 | control with that entity. For the purposes of this definition, 19 | "control" means (i) the power, direct or indirect, to cause the 20 | direction or management of such entity, whether by contract or 21 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 22 | outstanding shares, or (iii) beneficial ownership of such entity. 23 | 24 | "You" (or "Your") shall mean an individual or Legal Entity 25 | exercising permissions granted by this License. 26 | 27 | "Source" form shall mean the preferred form for making modifications, 28 | including but not limited to software source code, documentation 29 | source, and configuration files. 30 | 31 | "Object" form shall mean any form resulting from mechanical 32 | transformation or translation of a Source form, including but 33 | not limited to compiled object code, generated documentation, 34 | and conversions to other media types. 35 | 36 | "Work" shall mean the work of authorship, whether in Source or 37 | Object form, made available under the License, as indicated by a 38 | copyright notice that is included in or attached to the work 39 | (an example is provided in the Appendix below). 40 | 41 | "Derivative Works" shall mean any work, whether in Source or Object 42 | form, that is based on (or derived from) the Work and for which the 43 | editorial revisions, annotations, elaborations, or other modifications 44 | represent, as a whole, an original work of authorship. For the purposes 45 | of this License, Derivative Works shall not include works that remain 46 | separable from, or merely link (or bind by name) to the interfaces of, 47 | the Work and Derivative Works thereof. 48 | 49 | "Contribution" shall mean any work of authorship, including 50 | the original version of the Work and any modifications or additions 51 | to that Work or Derivative Works thereof, that is intentionally 52 | submitted to Licensor for inclusion in the Work by the copyright owner 53 | or by an individual or Legal Entity authorized to submit on behalf of 54 | the copyright owner. For the purposes of this definition, "submitted" 55 | means any form of electronic, verbal, or written communication sent 56 | to the Licensor or its representatives, including but not limited to 57 | communication on electronic mailing lists, source code control systems, 58 | and issue tracking systems that are managed by, or on behalf of, the 59 | Licensor for the purpose of discussing and improving the Work, but 60 | excluding communication that is conspicuously marked or otherwise 61 | designated in writing by the copyright owner as "Not a Contribution." 62 | 63 | "Contributor" shall mean Licensor and any individual or Legal Entity 64 | on behalf of whom a Contribution has been received by Licensor and 65 | subsequently incorporated within the Work. 66 | 67 | 2. Grant of Copyright License. Subject to the terms and conditions of 68 | this License, each Contributor hereby grants to You a perpetual, 69 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 70 | copyright license to reproduce, prepare Derivative Works of, 71 | publicly display, publicly perform, sublicense, and distribute the 72 | Work and such Derivative Works in Source or Object form. 73 | 74 | 3. Grant of Patent License. Subject to the terms and conditions of 75 | this License, each Contributor hereby grants to You a perpetual, 76 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 77 | (except as stated in this section) patent license to make, have made, 78 | use, offer to sell, sell, import, and otherwise transfer the Work, 79 | where such license applies only to those patent claims licensable 80 | by such Contributor that are necessarily infringed by their 81 | Contribution(s) alone or by combination of their Contribution(s) 82 | with the Work to which such Contribution(s) was submitted. If You 83 | institute patent litigation against any entity (including a 84 | cross-claim or counterclaim in a lawsuit) alleging that the Work 85 | or a Contribution incorporated within the Work constitutes direct 86 | or contributory patent infringement, then any patent licenses 87 | granted to You under this License for that Work shall terminate 88 | as of the date such litigation is filed. 89 | 90 | 4. Redistribution. You may reproduce and distribute copies of the 91 | Work or Derivative Works thereof in any medium, with or without 92 | modifications, and in Source or Object form, provided that You 93 | meet the following conditions: 94 | 95 | (a) You must give any other recipients of the Work or 96 | Derivative Works a copy of this License; and 97 | 98 | (b) You must cause any modified files to carry prominent notices 99 | stating that You changed the files; and 100 | 101 | (c) You must retain, in the Source form of any Derivative Works 102 | that You distribute, all copyright, patent, trademark, and 103 | attribution notices from the Source form of the Work, 104 | excluding those notices that do not pertain to any part of 105 | the Derivative Works; and 106 | 107 | (d) If the Work includes a "NOTICE" text file as part of its 108 | distribution, then any Derivative Works that You distribute must 109 | include a readable copy of the attribution notices contained 110 | within such NOTICE file, excluding those notices that do not 111 | pertain to any part of the Derivative Works, in at least one 112 | of the following places: within a NOTICE text file distributed 113 | as part of the Derivative Works; within the Source form or 114 | documentation, if provided along with the Derivative Works; or, 115 | within a display generated by the Derivative Works, if and 116 | wherever such third-party notices normally appear. The contents 117 | of the NOTICE file are for informational purposes only and 118 | do not modify the License. You may add Your own attribution 119 | notices within Derivative Works that You distribute, alongside 120 | or as an addendum to the NOTICE text from the Work, provided 121 | that such additional attribution notices cannot be construed 122 | as modifying the License. 123 | 124 | You may add Your own copyright statement to Your modifications and 125 | may provide additional or different license terms and conditions 126 | for use, reproduction, or distribution of Your modifications, or 127 | for any such Derivative Works as a whole, provided Your use, 128 | reproduction, and distribution of the Work otherwise complies with 129 | the conditions stated in this License. 130 | 131 | 5. Submission of Contributions. Unless You explicitly state otherwise, 132 | any Contribution intentionally submitted for inclusion in the Work 133 | by You to the Licensor shall be under the terms and conditions of 134 | this License, without any additional terms or conditions. 135 | Notwithstanding the above, nothing herein shall supersede or modify 136 | the terms of any separate license agreement you may have executed 137 | with Licensor regarding such Contributions. 138 | 139 | 6. Trademarks. This License does not grant permission to use the trade 140 | names, trademarks, service marks, or product names of the Licensor, 141 | except as required for reasonable and customary use in describing the 142 | origin of the Work and reproducing the content of the NOTICE file. 143 | 144 | 7. Disclaimer of Warranty. Unless required by applicable law or 145 | agreed to in writing, Licensor provides the Work (and each 146 | Contributor provides its Contributions) on an "AS IS" BASIS, 147 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 148 | implied, including, without limitation, any warranties or conditions 149 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 150 | PARTICULAR PURPOSE. You are solely responsible for determining the 151 | appropriateness of using or redistributing the Work and assume any 152 | risks associated with Your exercise of permissions under this License. 153 | 154 | 8. Limitation of Liability. In no event and under no legal theory, 155 | whether in tort (including negligence), contract, or otherwise, 156 | unless required by applicable law (such as deliberate and grossly 157 | negligent acts) or agreed to in writing, shall any Contributor be 158 | liable to You for damages, including any direct, indirect, special, 159 | incidental, or consequential damages of any character arising as a 160 | result of this License or out of the use or inability to use the 161 | Work (including but not limited to damages for loss of goodwill, 162 | work stoppage, computer failure or malfunction, or any and all 163 | other commercial damages or losses), even if such Contributor 164 | has been advised of the possibility of such damages. 165 | 166 | 9. Accepting Warranty or Additional Liability. While redistributing 167 | the Work or Derivative Works thereof, You may choose to offer, 168 | and charge a fee for, acceptance of support, warranty, indemnity, 169 | or other liability obligations and/or rights consistent with this 170 | License. However, in accepting such obligations, You may act only 171 | on Your own behalf and on Your sole responsibility, not on behalf 172 | of any other Contributor, and only if You agree to indemnify, 173 | defend, and hold each Contributor harmless for any liability 174 | incurred by, or claims asserted against, such Contributor by reason 175 | of your accepting any such warranty or additional liability. 176 | 177 | END OF TERMS AND CONDITIONS 178 | 179 | APPENDIX: How to apply the Apache License to your work. 180 | 181 | To apply the Apache License to your work, attach the following 182 | boilerplate notice, with the fields enclosed by brackets "[]" 183 | replaced with your own identifying information. (Don't include 184 | the brackets!) The text should be enclosed in the appropriate 185 | comment syntax for the file format. We also recommend that a 186 | file or class name and description of purpose be included on the 187 | same "printed page" as the copyright notice for easier 188 | identification within third-party archives. 189 | 190 | Copyright [yyyy] [name of copyright owner] 191 | 192 | Licensed under the Apache License, Version 2.0 (the "License"); 193 | you may not use this file except in compliance with the License. 194 | You may obtain a copy of the License at 195 | 196 | http://www.apache.org/licenses/LICENSE-2.0 197 | 198 | Unless required by applicable law or agreed to in writing, software 199 | distributed under the License is distributed on an "AS IS" BASIS, 200 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 201 | See the License for the specific language governing permissions and 202 | limitations under the License. 203 | -------------------------------------------------------------------------------- /tokenization.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # Copyright 2018 The Google AI Language Team Authors. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | """Tokenization classes.""" 16 | 17 | from __future__ import absolute_import 18 | from __future__ import division 19 | from __future__ import print_function 20 | 21 | import collections 22 | import re 23 | import unicodedata 24 | import six 25 | import tensorflow as tf 26 | 27 | 28 | def validate_case_matches_checkpoint(do_lower_case, init_checkpoint): 29 | """Checks whether the casing config is consistent with the checkpoint name.""" 30 | 31 | # The casing has to be passed in by the user and there is no explicit check 32 | # as to whether it matches the checkpoint. The casing information probably 33 | # should have been stored in the bert_config.json file, but it's not, so 34 | # we have to heuristically detect it to validate. 35 | 36 | if not init_checkpoint: 37 | return 38 | 39 | m = re.match("^.*?([A-Za-z0-9_-]+)/bert_model.ckpt", init_checkpoint) 40 | if m is None: 41 | return 42 | 43 | model_name = m.group(1) 44 | 45 | lower_models = [ 46 | "uncased_L-24_H-1024_A-16", "uncased_L-12_H-768_A-12", 47 | "multilingual_L-12_H-768_A-12", "chinese_L-12_H-768_A-12" 48 | ] 49 | 50 | cased_models = [ 51 | "cased_L-12_H-768_A-12", "cased_L-24_H-1024_A-16", 52 | "multi_cased_L-12_H-768_A-12" 53 | ] 54 | 55 | is_bad_config = False 56 | if model_name in lower_models and not do_lower_case: 57 | is_bad_config = True 58 | actual_flag = "False" 59 | case_name = "lowercased" 60 | opposite_flag = "True" 61 | 62 | if model_name in cased_models and do_lower_case: 63 | is_bad_config = True 64 | actual_flag = "True" 65 | case_name = "cased" 66 | opposite_flag = "False" 67 | 68 | if is_bad_config: 69 | raise ValueError( 70 | "You passed in `--do_lower_case=%s` with `--init_checkpoint=%s`. " 71 | "However, `%s` seems to be a %s model, so you " 72 | "should pass in `--do_lower_case=%s` so that the fine-tuning matches " 73 | "how the model was pre-training. If this error is wrong, please " 74 | "just comment out this check." % (actual_flag, init_checkpoint, 75 | model_name, case_name, opposite_flag)) 76 | 77 | 78 | def convert_to_unicode(text): 79 | """Converts `text` to Unicode (if it's not already), assuming utf-8 input.""" 80 | if six.PY3: 81 | if isinstance(text, str): 82 | return text 83 | elif isinstance(text, bytes): 84 | return text.decode("utf-8", "ignore") 85 | else: 86 | raise ValueError("Unsupported string type: %s" % (type(text))) 87 | elif six.PY2: 88 | if isinstance(text, str): 89 | return text.decode("utf-8", "ignore") 90 | elif isinstance(text, unicode): 91 | return text 92 | else: 93 | raise ValueError("Unsupported string type: %s" % (type(text))) 94 | else: 95 | raise ValueError("Not running on Python2 or Python 3?") 96 | 97 | 98 | def printable_text(text): 99 | """Returns text encoded in a way suitable for print or `tf.logging`.""" 100 | 101 | # These functions want `str` for both Python2 and Python3, but in one case 102 | # it's a Unicode string and in the other it's a byte string. 103 | if six.PY3: 104 | if isinstance(text, str): 105 | return text 106 | elif isinstance(text, bytes): 107 | return text.decode("utf-8", "ignore") 108 | else: 109 | raise ValueError("Unsupported string type: %s" % (type(text))) 110 | elif six.PY2: 111 | if isinstance(text, str): 112 | return text 113 | elif isinstance(text, unicode): 114 | return text.encode("utf-8") 115 | else: 116 | raise ValueError("Unsupported string type: %s" % (type(text))) 117 | else: 118 | raise ValueError("Not running on Python2 or Python 3?") 119 | 120 | 121 | def load_vocab(vocab_file): 122 | """Loads a vocabulary file into a dictionary.""" 123 | vocab = collections.OrderedDict() 124 | index = 0 125 | with tf.gfile.GFile(vocab_file, "r") as reader: 126 | while True: 127 | token = convert_to_unicode(reader.readline()) 128 | if not token: 129 | break 130 | token = token.strip() 131 | vocab[token] = index 132 | index += 1 133 | return vocab 134 | 135 | 136 | def convert_by_vocab(vocab, items): 137 | """Converts a sequence of [tokens|ids] using the vocab.""" 138 | output = [] 139 | for item in items: 140 | output.append(vocab[item]) 141 | return output 142 | 143 | 144 | def convert_tokens_to_ids(vocab, tokens): 145 | return convert_by_vocab(vocab, tokens) 146 | 147 | 148 | def convert_ids_to_tokens(inv_vocab, ids): 149 | return convert_by_vocab(inv_vocab, ids) 150 | 151 | 152 | def whitespace_tokenize(text): 153 | """Runs basic whitespace cleaning and splitting on a piece of text.""" 154 | text = text.strip() 155 | if not text: 156 | return [] 157 | tokens = text.split() 158 | return tokens 159 | 160 | 161 | class FullTokenizer(object): 162 | """Runs end-to-end tokenziation.""" 163 | 164 | def __init__(self, vocab_file, do_lower_case=True): 165 | self.vocab = load_vocab(vocab_file) 166 | self.inv_vocab = {v: k for k, v in self.vocab.items()} 167 | self.basic_tokenizer = BasicTokenizer(do_lower_case=do_lower_case) 168 | self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.vocab) 169 | 170 | def tokenize(self, text): 171 | split_tokens = [] 172 | for token in self.basic_tokenizer.tokenize(text): 173 | for sub_token in self.wordpiece_tokenizer.tokenize(token): 174 | split_tokens.append(sub_token) 175 | 176 | return split_tokens 177 | 178 | def convert_tokens_to_ids(self, tokens): 179 | return convert_by_vocab(self.vocab, tokens) 180 | 181 | def convert_ids_to_tokens(self, ids): 182 | return convert_by_vocab(self.inv_vocab, ids) 183 | 184 | 185 | class BasicTokenizer(object): 186 | """Runs basic tokenization (punctuation splitting, lower casing, etc.).""" 187 | 188 | def __init__(self, do_lower_case=True): 189 | """Constructs a BasicTokenizer. 190 | 191 | Args: 192 | do_lower_case: Whether to lower case the input. 193 | """ 194 | self.do_lower_case = do_lower_case 195 | 196 | def tokenize(self, text): 197 | """Tokenizes a piece of text.""" 198 | text = convert_to_unicode(text) 199 | text = self._clean_text(text) 200 | 201 | # This was added on November 1st, 2018 for the multilingual and Chinese 202 | # models. This is also applied to the English models now, but it doesn't 203 | # matter since the English models were not trained on any Chinese data 204 | # and generally don't have any Chinese data in them (there are Chinese 205 | # characters in the vocabulary because Wikipedia does have some Chinese 206 | # words in the English Wikipedia.). 207 | text = self._tokenize_chinese_chars(text) 208 | 209 | orig_tokens = whitespace_tokenize(text) 210 | split_tokens = [] 211 | for token in orig_tokens: 212 | if self.do_lower_case: 213 | token = token.lower() 214 | token = self._run_strip_accents(token) 215 | split_tokens.extend(self._run_split_on_punc(token)) 216 | 217 | output_tokens = whitespace_tokenize(" ".join(split_tokens)) 218 | return output_tokens 219 | 220 | def _run_strip_accents(self, text): 221 | """Strips accents from a piece of text.""" 222 | text = unicodedata.normalize("NFD", text) 223 | output = [] 224 | for char in text: 225 | cat = unicodedata.category(char) 226 | if cat == "Mn": 227 | continue 228 | output.append(char) 229 | return "".join(output) 230 | 231 | def _run_split_on_punc(self, text): 232 | """Splits punctuation on a piece of text.""" 233 | chars = list(text) 234 | i = 0 235 | start_new_word = True 236 | output = [] 237 | while i < len(chars): 238 | char = chars[i] 239 | if _is_punctuation(char): 240 | output.append([char]) 241 | start_new_word = True 242 | else: 243 | if start_new_word: 244 | output.append([]) 245 | start_new_word = False 246 | output[-1].append(char) 247 | i += 1 248 | 249 | return ["".join(x) for x in output] 250 | 251 | def _tokenize_chinese_chars(self, text): 252 | """Adds whitespace around any CJK character.""" 253 | output = [] 254 | for char in text: 255 | cp = ord(char) 256 | if self._is_chinese_char(cp): 257 | output.append(" ") 258 | output.append(char) 259 | output.append(" ") 260 | else: 261 | output.append(char) 262 | return "".join(output) 263 | 264 | def _is_chinese_char(self, cp): 265 | """Checks whether CP is the codepoint of a CJK character.""" 266 | # This defines a "chinese character" as anything in the CJK Unicode block: 267 | # https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block) 268 | # 269 | # Note that the CJK Unicode block is NOT all Japanese and Korean characters, 270 | # despite its name. The modern Korean Hangul alphabet is a different block, 271 | # as is Japanese Hiragana and Katakana. Those alphabets are used to write 272 | # space-separated words, so they are not treated specially and handled 273 | # like the all of the other languages. 274 | if ((cp >= 0x4E00 and cp <= 0x9FFF) or # 275 | (cp >= 0x3400 and cp <= 0x4DBF) or # 276 | (cp >= 0x20000 and cp <= 0x2A6DF) or # 277 | (cp >= 0x2A700 and cp <= 0x2B73F) or # 278 | (cp >= 0x2B740 and cp <= 0x2B81F) or # 279 | (cp >= 0x2B820 and cp <= 0x2CEAF) or 280 | (cp >= 0xF900 and cp <= 0xFAFF) or # 281 | (cp >= 0x2F800 and cp <= 0x2FA1F)): # 282 | return True 283 | 284 | return False 285 | 286 | def _clean_text(self, text): 287 | """Performs invalid character removal and whitespace cleanup on text.""" 288 | output = [] 289 | for char in text: 290 | cp = ord(char) 291 | if cp == 0 or cp == 0xfffd or _is_control(char): 292 | continue 293 | if _is_whitespace(char): 294 | output.append(" ") 295 | else: 296 | output.append(char) 297 | return "".join(output) 298 | 299 | 300 | class WordpieceTokenizer(object): 301 | """Runs WordPiece tokenziation.""" 302 | 303 | def __init__(self, vocab, unk_token="[UNK]", max_input_chars_per_word=200): 304 | self.vocab = vocab 305 | self.unk_token = unk_token 306 | self.max_input_chars_per_word = max_input_chars_per_word 307 | 308 | def tokenize(self, text): 309 | """Tokenizes a piece of text into its word pieces. 310 | 311 | This uses a greedy longest-match-first algorithm to perform tokenization 312 | using the given vocabulary. 313 | 314 | For example: 315 | input = "unaffable" 316 | output = ["un", "##aff", "##able"] 317 | 318 | Args: 319 | text: A single token or whitespace separated tokens. This should have 320 | already been passed through `BasicTokenizer. 321 | 322 | Returns: 323 | A list of wordpiece tokens. 324 | """ 325 | 326 | text = convert_to_unicode(text) 327 | 328 | output_tokens = [] 329 | for token in whitespace_tokenize(text): 330 | chars = list(token) 331 | if len(chars) > self.max_input_chars_per_word: 332 | output_tokens.append(self.unk_token) 333 | continue 334 | 335 | is_bad = False 336 | start = 0 337 | sub_tokens = [] 338 | while start < len(chars): 339 | end = len(chars) 340 | cur_substr = None 341 | while start < end: 342 | substr = "".join(chars[start:end]) 343 | if start > 0: 344 | substr = "##" + substr 345 | if substr in self.vocab: 346 | cur_substr = substr 347 | break 348 | end -= 1 349 | if cur_substr is None: 350 | is_bad = True 351 | break 352 | sub_tokens.append(cur_substr) 353 | start = end 354 | 355 | if is_bad: 356 | output_tokens.append(self.unk_token) 357 | else: 358 | output_tokens.extend(sub_tokens) 359 | return output_tokens 360 | 361 | 362 | def _is_whitespace(char): 363 | """Checks whether `chars` is a whitespace character.""" 364 | # \t, \n, and \r are technically contorl characters but we treat them 365 | # as whitespace since they are generally considered as such. 366 | if char == " " or char == "\t" or char == "\n" or char == "\r": 367 | return True 368 | cat = unicodedata.category(char) 369 | if cat == "Zs": 370 | return True 371 | return False 372 | 373 | 374 | def _is_control(char): 375 | """Checks whether `chars` is a control character.""" 376 | # These are technically control characters but we count them as whitespace 377 | # characters. 378 | if char == "\t" or char == "\n" or char == "\r": 379 | return False 380 | cat = unicodedata.category(char) 381 | if cat.startswith("C"): 382 | return True 383 | return False 384 | 385 | 386 | def _is_punctuation(char): 387 | """Checks whether `chars` is a punctuation character.""" 388 | cp = ord(char) 389 | # We treat all non-letter/number ASCII as punctuation. 390 | # Characters such as "^", "$", and "`" are not in the Unicode 391 | # Punctuation class but we treat them as punctuation anyways, for 392 | # consistency. 393 | if ((cp >= 33 and cp <= 47) or (cp >= 58 and cp <= 64) or 394 | (cp >= 91 and cp <= 96) or (cp >= 123 and cp <= 126)): 395 | return True 396 | cat = unicodedata.category(char) 397 | if cat.startswith("P"): 398 | return True 399 | return False 400 | -------------------------------------------------------------------------------- /create_pretraining_data.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # Copyright 2018 The Google AI Language Team Authors. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | """Create masked LM/next sentence masked_lm TF examples for BERT.""" 16 | 17 | from __future__ import absolute_import 18 | from __future__ import division 19 | from __future__ import print_function 20 | 21 | import collections 22 | import random 23 | import tensorflow as tf 24 | import tokenization 25 | 26 | flags = tf.flags 27 | 28 | FLAGS = flags.FLAGS 29 | 30 | flags.DEFINE_string("input_dir", "/data1/tong.guo/1-billion-word-language-modeling-benchmark-r13output/training-monolingual.tokenized.shuffled", 31 | "Input raw text file (or comma-separated list of files).") 32 | 33 | flags.DEFINE_string("input_file", "data/xab", 34 | "Input raw text file (or comma-separated list of files).") 35 | 36 | flags.DEFINE_string( 37 | "output_file", "tmp_data/sample2.tfrecords", 38 | "Output TF example file (or comma-separated list of files).") 39 | 40 | flags.DEFINE_string("vocab_file", "vocab/vocab.txt", 41 | "The vocabulary file that the BERT model was trained on.") 42 | 43 | flags.DEFINE_bool( 44 | "do_lower_case", True, 45 | "Whether to lower case the input text. Should be True for uncased " 46 | "models and False for cased models.") 47 | 48 | flags.DEFINE_integer("max_seq_length", 64, "Maximum sequence length.") 49 | 50 | flags.DEFINE_integer("max_predictions_per_seq", 20, 51 | "Maximum number of masked LM predictions per sequence.") 52 | 53 | flags.DEFINE_integer("random_seed", 12345, "Random seed for data generation.") 54 | 55 | flags.DEFINE_integer( 56 | "dupe_factor", 10, 57 | "Number of times to duplicate the input data (with different masks).") 58 | 59 | flags.DEFINE_float("masked_lm_prob", 0.15, "Masked LM probability.") 60 | 61 | flags.DEFINE_float( 62 | "short_seq_prob", 0.1, 63 | "Probability of creating sequences which are shorter than the " 64 | "maximum length.") 65 | 66 | 67 | class TrainingInstance(object): 68 | """A single training instance (sentence pair).""" 69 | 70 | def __init__(self, tokens, segment_ids, masked_lm_positions, masked_lm_labels, 71 | is_random_next): 72 | self.tokens = tokens 73 | self.segment_ids = segment_ids 74 | self.is_random_next = is_random_next 75 | self.masked_lm_positions = masked_lm_positions 76 | self.masked_lm_labels = masked_lm_labels 77 | 78 | def __str__(self): 79 | s = "" 80 | s += "tokens: %s\n" % (" ".join( 81 | [tokenization.printable_text(x) for x in self.tokens])) 82 | s += "segment_ids: %s\n" % (" ".join([str(x) for x in self.segment_ids])) 83 | s += "is_random_next: %s\n" % self.is_random_next 84 | s += "masked_lm_positions: %s\n" % (" ".join( 85 | [str(x) for x in self.masked_lm_positions])) 86 | s += "masked_lm_labels: %s\n" % (" ".join( 87 | [tokenization.printable_text(x) for x in self.masked_lm_labels])) 88 | s += "\n" 89 | return s 90 | 91 | def __repr__(self): 92 | return self.__str__() 93 | 94 | 95 | def write_instance_to_example_files(instances, tokenizer, max_seq_length, 96 | max_predictions_per_seq, output_files): 97 | """Create TF example files from `TrainingInstance`s.""" 98 | writers = [] 99 | for output_file in output_files: 100 | writers.append(tf.python_io.TFRecordWriter(output_file)) 101 | 102 | writer_index = 0 103 | 104 | total_written = 0 105 | for (inst_index, instance) in enumerate(instances): 106 | input_ids = tokenizer.convert_tokens_to_ids(instance.tokens) 107 | input_mask = [1] * len(input_ids) 108 | segment_ids = list(instance.segment_ids) 109 | assert len(input_ids) <= max_seq_length 110 | 111 | while len(input_ids) < max_seq_length: 112 | input_ids.append(0) 113 | input_mask.append(0) 114 | segment_ids.append(0) 115 | 116 | assert len(input_ids) == max_seq_length 117 | assert len(input_mask) == max_seq_length 118 | assert len(segment_ids) == max_seq_length 119 | 120 | masked_lm_positions = list(instance.masked_lm_positions) 121 | masked_lm_ids = tokenizer.convert_tokens_to_ids(instance.masked_lm_labels) 122 | masked_lm_weights = [1.0] * len(masked_lm_ids) 123 | 124 | while len(masked_lm_positions) < max_predictions_per_seq: 125 | masked_lm_positions.append(0) 126 | masked_lm_ids.append(0) 127 | masked_lm_weights.append(0.0) 128 | 129 | next_sentence_label = 1 if instance.is_random_next else 0 130 | 131 | features = collections.OrderedDict() 132 | features["input_ids"] = create_int_feature(input_ids) 133 | features["input_mask"] = create_int_feature(input_mask) 134 | features["segment_ids"] = create_int_feature(segment_ids) 135 | features["masked_lm_positions"] = create_int_feature(masked_lm_positions) 136 | features["masked_lm_ids"] = create_int_feature(masked_lm_ids) 137 | features["masked_lm_weights"] = create_float_feature(masked_lm_weights) 138 | features["next_sentence_labels"] = create_int_feature([next_sentence_label]) 139 | 140 | tf_example = tf.train.Example(features=tf.train.Features(feature=features)) 141 | 142 | writers[writer_index].write(tf_example.SerializeToString()) 143 | writer_index = (writer_index + 1) % len(writers) 144 | 145 | total_written += 1 146 | 147 | if inst_index < 20: 148 | tf.logging.info("*** Example ***") 149 | tf.logging.info("tokens: %s" % " ".join( 150 | [tokenization.printable_text(x) for x in instance.tokens])) 151 | 152 | for feature_name in features.keys(): 153 | feature = features[feature_name] 154 | values = [] 155 | if feature.int64_list.value: 156 | values = feature.int64_list.value 157 | elif feature.float_list.value: 158 | values = feature.float_list.value 159 | tf.logging.info( 160 | "%s: %s" % (feature_name, " ".join([str(x) for x in values]))) 161 | 162 | for writer in writers: 163 | writer.close() 164 | 165 | tf.logging.info("Wrote %d total instances", total_written) 166 | 167 | 168 | def create_int_feature(values): 169 | feature = tf.train.Feature(int64_list=tf.train.Int64List(value=list(values))) 170 | return feature 171 | 172 | 173 | def create_float_feature(values): 174 | feature = tf.train.Feature(float_list=tf.train.FloatList(value=list(values))) 175 | return feature 176 | 177 | 178 | def create_training_instances(input_files, tokenizer, max_seq_length, 179 | dupe_factor, short_seq_prob, masked_lm_prob, 180 | max_predictions_per_seq, rng): 181 | """Create `TrainingInstance`s from raw text.""" 182 | all_documents = [[]] 183 | index=0 184 | # Input file format: 185 | # (1) One sentence per line. These should ideally be actual sentences, not 186 | # entire paragraphs or arbitrary spans of text. (Because we use the 187 | # sentence boundaries for the "next sentence prediction" task). 188 | # (2) Blank lines between documents. Document boundaries are needed so 189 | # that the "next sentence prediction" task doesn't span between documents. 190 | for input_file in input_files: 191 | with tf.gfile.GFile(input_file, "r") as reader: 192 | while True: 193 | index+=1 194 | if index%1000==0: 195 | print(index) 196 | line = tokenization.convert_to_unicode(reader.readline()) 197 | if not line: 198 | break 199 | line = line.strip() 200 | 201 | # Empty lines are used as document delimiters 202 | if not line: 203 | all_documents.append([]) 204 | tokens = tokenizer.tokenize(line) 205 | if tokens: 206 | all_documents[-1].append(tokens) 207 | 208 | # Remove empty documents 209 | all_documents = [x for x in all_documents if x] 210 | rng.shuffle(all_documents) 211 | 212 | vocab_words = list(tokenizer.vocab.keys()) 213 | instances = [] 214 | for _ in range(dupe_factor): 215 | for document_index in range(len(all_documents)): 216 | if document_index%1000==0: 217 | print(document_index) 218 | instances.extend( 219 | create_instances_from_document( 220 | all_documents, document_index, max_seq_length, short_seq_prob, 221 | masked_lm_prob, max_predictions_per_seq, vocab_words, rng)) 222 | 223 | rng.shuffle(instances) 224 | return instances 225 | 226 | 227 | def create_instances_from_document( 228 | all_documents, document_index, max_seq_length, short_seq_prob, 229 | masked_lm_prob, max_predictions_per_seq, vocab_words, rng): 230 | """Creates `TrainingInstance`s for a single document.""" 231 | document = all_documents[document_index] 232 | 233 | # Account for [CLS], [SEP], [SEP] 234 | max_num_tokens = max_seq_length - 3 235 | 236 | # We *usually* want to fill up the entire sequence since we are padding 237 | # to `max_seq_length` anyways, so short sequences are generally wasted 238 | # computation. However, we *sometimes* 239 | # (i.e., short_seq_prob == 0.1 == 10% of the time) want to use shorter 240 | # sequences to minimize the mismatch between pre-training and fine-tuning. 241 | # The `target_seq_length` is just a rough target however, whereas 242 | # `max_seq_length` is a hard limit. 243 | target_seq_length = max_num_tokens 244 | if rng.random() < short_seq_prob: 245 | target_seq_length = rng.randint(2, max_num_tokens) 246 | 247 | # We DON'T just concatenate all of the tokens from a document into a long 248 | # sequence and choose an arbitrary split point because this would make the 249 | # next sentence prediction task too easy. Instead, we split the input into 250 | # segments "A" and "B" based on the actual "sentences" provided by the user 251 | # input. 252 | instances = [] 253 | current_chunk = [] 254 | current_length = 0 255 | i = 0 256 | while i < len(document): 257 | segment = document[i] 258 | current_chunk.append(segment) 259 | current_length += len(segment) 260 | if i == len(document) - 1 or current_length >= target_seq_length: 261 | if current_chunk: 262 | # `a_end` is how many segments from `current_chunk` go into the `A` 263 | # (first) sentence. 264 | a_end = 1 265 | if len(current_chunk) >= 2: 266 | a_end = rng.randint(1, len(current_chunk) - 1) 267 | 268 | tokens_a = [] 269 | for j in range(a_end): 270 | tokens_a.extend(current_chunk[j]) 271 | 272 | tokens_b = [] 273 | # Random next 274 | is_random_next = False 275 | if len(current_chunk) == 1 or rng.random() < 0.5: 276 | is_random_next = True 277 | target_b_length = target_seq_length - len(tokens_a) 278 | 279 | # This should rarely go for more than one iteration for large 280 | # corpora. However, just to be careful, we try to make sure that 281 | # the random document is not the same as the document 282 | # we're processing. 283 | for _ in range(10): 284 | random_document_index = rng.randint(0, len(all_documents) - 1) 285 | if random_document_index != document_index: 286 | break 287 | 288 | random_document = all_documents[random_document_index] 289 | random_start = rng.randint(0, len(random_document) - 1) 290 | for j in range(random_start, len(random_document)): 291 | tokens_b.extend(random_document[j]) 292 | if len(tokens_b) >= target_b_length: 293 | break 294 | # We didn't actually use these segments so we "put them back" so 295 | # they don't go to waste. 296 | num_unused_segments = len(current_chunk) - a_end 297 | i -= num_unused_segments 298 | # Actual next 299 | else: 300 | is_random_next = False 301 | for j in range(a_end, len(current_chunk)): 302 | tokens_b.extend(current_chunk[j]) 303 | truncate_seq_pair(tokens_a, tokens_b, max_num_tokens, rng) 304 | 305 | assert len(tokens_a) >= 1 306 | assert len(tokens_b) >= 1 307 | 308 | tokens = [] 309 | segment_ids = [] 310 | tokens.append("[CLS]") 311 | segment_ids.append(0) 312 | for token in tokens_a: 313 | tokens.append(token) 314 | segment_ids.append(0) 315 | 316 | tokens.append("[SEP]") 317 | segment_ids.append(0) 318 | 319 | for token in tokens_b: 320 | tokens.append(token) 321 | segment_ids.append(1) 322 | tokens.append("[SEP]") 323 | segment_ids.append(1) 324 | 325 | (tokens, masked_lm_positions, 326 | masked_lm_labels) = create_masked_lm_predictions( 327 | tokens, masked_lm_prob, max_predictions_per_seq, vocab_words, rng) 328 | instance = TrainingInstance( 329 | tokens=tokens, 330 | segment_ids=segment_ids, 331 | is_random_next=is_random_next, 332 | masked_lm_positions=masked_lm_positions, 333 | masked_lm_labels=masked_lm_labels) 334 | instances.append(instance) 335 | current_chunk = [] 336 | current_length = 0 337 | i += 1 338 | 339 | return instances 340 | 341 | 342 | MaskedLmInstance = collections.namedtuple("MaskedLmInstance", 343 | ["index", "label"]) 344 | 345 | 346 | def create_masked_lm_predictions(tokens, masked_lm_prob, 347 | max_predictions_per_seq, vocab_words, rng): 348 | """Creates the predictions for the masked LM objective.""" 349 | 350 | cand_indexes = [] 351 | for (i, token) in enumerate(tokens): 352 | if token == "[CLS]" or token == "[SEP]": 353 | continue 354 | cand_indexes.append(i) 355 | 356 | rng.shuffle(cand_indexes) 357 | 358 | output_tokens = list(tokens) 359 | 360 | num_to_predict = min(max_predictions_per_seq, 361 | max(1, int(round(len(tokens) * masked_lm_prob)))) 362 | 363 | masked_lms = [] 364 | covered_indexes = set() 365 | for index in cand_indexes: 366 | if len(masked_lms) >= num_to_predict: 367 | break 368 | if index in covered_indexes: 369 | continue 370 | covered_indexes.add(index) 371 | 372 | masked_token = None 373 | # 80% of the time, replace with [MASK] 374 | if rng.random() < 0.8: 375 | masked_token = "[MASK]" 376 | else: 377 | # 10% of the time, keep original 378 | if rng.random() < 0.5: 379 | masked_token = tokens[index] 380 | # 10% of the time, replace with random word 381 | else: 382 | masked_token = vocab_words[rng.randint(0, len(vocab_words) - 1)] 383 | 384 | output_tokens[index] = masked_token 385 | 386 | masked_lms.append(MaskedLmInstance(index=index, label=tokens[index])) 387 | 388 | masked_lms = sorted(masked_lms, key=lambda x: x.index) 389 | 390 | masked_lm_positions = [] 391 | masked_lm_labels = [] 392 | for p in masked_lms: 393 | masked_lm_positions.append(p.index) 394 | masked_lm_labels.append(p.label) 395 | 396 | return (output_tokens, masked_lm_positions, masked_lm_labels) 397 | 398 | 399 | def truncate_seq_pair(tokens_a, tokens_b, max_num_tokens, rng): 400 | """Truncates a pair of sequences to a maximum sequence length.""" 401 | while True: 402 | total_length = len(tokens_a) + len(tokens_b) 403 | if total_length <= max_num_tokens: 404 | break 405 | 406 | trunc_tokens = tokens_a if len(tokens_a) > len(tokens_b) else tokens_b 407 | assert len(trunc_tokens) >= 1 408 | 409 | # We want to sometimes truncate from the front and sometimes from the 410 | # back to add more randomness and avoid biases. 411 | if rng.random() < 0.5: 412 | del trunc_tokens[0] 413 | else: 414 | trunc_tokens.pop() 415 | 416 | import os 417 | 418 | def main(_): 419 | tf.logging.set_verbosity(tf.logging.INFO) 420 | 421 | tokenizer = tokenization.FullTokenizer( 422 | vocab_file=FLAGS.vocab_file, do_lower_case=FLAGS.do_lower_case) 423 | 424 | input_files = []#[FLAGS.input_dir+"/"+f for f in os.listdir(FLAGS.input_dir)] 425 | for input_pattern in FLAGS.input_file.split(","): 426 | input_files.extend(tf.gfile.Glob(input_pattern)) 427 | 428 | tf.logging.info("*** Reading from input files ***") 429 | for input_file in input_files: 430 | tf.logging.info(" %s", input_file) 431 | 432 | rng = random.Random(FLAGS.random_seed) 433 | instances = create_training_instances( 434 | input_files, tokenizer, FLAGS.max_seq_length, FLAGS.dupe_factor, 435 | FLAGS.short_seq_prob, FLAGS.masked_lm_prob, FLAGS.max_predictions_per_seq, 436 | rng) 437 | 438 | output_files = FLAGS.output_file.split(",") 439 | tf.logging.info("*** Writing to output files ***") 440 | for output_file in output_files: 441 | tf.logging.info(" %s", output_file) 442 | 443 | write_instance_to_example_files(instances, tokenizer, FLAGS.max_seq_length, 444 | FLAGS.max_predictions_per_seq, output_files) 445 | 446 | 447 | if __name__ == "__main__": 448 | flags.mark_flag_as_required("input_file") 449 | flags.mark_flag_as_required("output_file") 450 | flags.mark_flag_as_required("vocab_file") 451 | tf.app.run() 452 | -------------------------------------------------------------------------------- /run_pretraining_gpu_v2.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # Copyright 2018 The Google AI Language Team Authors. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | """Run masked LM/next sentence masked_lm pre-training for BERT.""" 16 | 17 | from __future__ import absolute_import 18 | from __future__ import division 19 | from __future__ import print_function 20 | 21 | import os 22 | import modeling 23 | import optimization_gpu as optimization 24 | import tensorflow as tf 25 | from tensorflow.python.estimator.run_config import RunConfig 26 | from tensorflow.python.estimator.estimator import Estimator 27 | from tensorflow.python.distribute.cross_device_ops import AllReduceCrossDeviceOps 28 | 29 | flags = tf.flags 30 | 31 | FLAGS = flags.FLAGS 32 | 33 | flags.DEFINE_integer( 34 | "n_gpus", 6, 35 | "GPU number") 36 | 37 | ## Required parameters 38 | flags.DEFINE_string( 39 | "bert_config_file", "bert_config.json", 40 | "The config json file corresponding to the pre-trained BERT model. " 41 | "This specifies the model architecture.") 42 | 43 | flags.DEFINE_string( 44 | "input_file", "tmp_data_128/sample1.tfrecords,tmp_data_128/sample2.tfrecords,tmp_data_128/sample3.tfrecords,tmp_data_128/sample4.tfrecords,tmp_data_128/sample5.tfrecords,tmp_data_128/sample6.tfrecords,tmp_data_128/sample7.tfrecords,tmp_data_128/sample8.tfrecords,tmp_data_128/sample9.tfrecords,tmp_data_128/sample10.tfrecords", 45 | "Input TF example files (can be a glob or comma separated).") 46 | 47 | flags.DEFINE_string( 48 | "output_dir", "output2", 49 | "The output directory where the model checkpoints will be written.") 50 | 51 | ## Other parameters 52 | flags.DEFINE_string( 53 | "init_checkpoint", None, 54 | "Initial checkpoint (usually from a pre-trained BERT model).") 55 | 56 | flags.DEFINE_integer( 57 | "max_seq_length", 128, 58 | "The maximum total input sequence length after WordPiece tokenization. " 59 | "Sequences longer than this will be truncated, and sequences shorter " 60 | "than this will be padded. Must match data generation.") 61 | 62 | flags.DEFINE_integer( 63 | "max_predictions_per_seq", 20, 64 | "Maximum number of masked LM predictions per sequence. " 65 | "Must match data generation.") 66 | 67 | flags.DEFINE_bool("do_train", True, "Whether to run training.") 68 | 69 | flags.DEFINE_bool("do_eval", False, "Whether to run eval on the dev set.") 70 | 71 | flags.DEFINE_integer("train_batch_size", 32, "Total batch size for training.") # 32 or 8? 72 | 73 | flags.DEFINE_integer("eval_batch_size", 8, "Total batch size for eval.") 74 | 75 | flags.DEFINE_float("learning_rate", 5e-5, "The initial learning rate for Adam.") 76 | 77 | flags.DEFINE_integer("num_train_steps", 300000, "Number of training steps.") 78 | 79 | flags.DEFINE_integer("num_warmup_steps", 10000, "Number of warmup steps.") 80 | 81 | flags.DEFINE_integer("save_checkpoints_steps", 1000, 82 | "How often to save the model checkpoint.") 83 | 84 | flags.DEFINE_integer("iterations_per_loop", 1000, 85 | "How many steps to make in each estimator call.") 86 | 87 | flags.DEFINE_integer("max_eval_steps", 100, "Maximum number of eval steps.") 88 | 89 | flags.DEFINE_bool("use_tpu", False, "Whether to use TPU or GPU/CPU.") 90 | 91 | tf.flags.DEFINE_string( 92 | "tpu_name", None, 93 | "The Cloud TPU to use for training. This should be either the name " 94 | "used when creating the Cloud TPU, or a grpc://ip.address.of.tpu:8470 " 95 | "url.") 96 | 97 | tf.flags.DEFINE_string( 98 | "tpu_zone", None, 99 | "[Optional] GCE zone where the Cloud TPU is located in. If not " 100 | "specified, we will attempt to automatically detect the GCE project from " 101 | "metadata.") 102 | 103 | tf.flags.DEFINE_string( 104 | "gcp_project", None, 105 | "[Optional] Project name for the Cloud TPU-enabled project. If not " 106 | "specified, we will attempt to automatically detect the GCE project from " 107 | "metadata.") 108 | 109 | tf.flags.DEFINE_string("master", None, "[Optional] TensorFlow master URL.") 110 | 111 | flags.DEFINE_integer( 112 | "num_tpu_cores", 8, 113 | "Only used if `use_tpu` is True. Total number of TPU cores to use.") 114 | 115 | 116 | def model_fn_builder(bert_config, init_checkpoint, learning_rate, 117 | num_train_steps, num_warmup_steps, use_tpu, 118 | use_one_hot_embeddings): 119 | """Returns `model_fn` closure for TPUEstimator.""" 120 | 121 | def model_fn(features, labels, mode, params): # pylint: disable=unused-argument 122 | """The `model_fn` for TPUEstimator.""" 123 | 124 | tf.logging.info("*** Features ***") 125 | for name in sorted(features.keys()): 126 | tf.logging.info(" name = %s, shape = %s" % (name, features[name].shape)) 127 | 128 | input_ids = features["input_ids"] 129 | input_mask = features["input_mask"] 130 | segment_ids = features["segment_ids"] 131 | masked_lm_positions = features["masked_lm_positions"] 132 | masked_lm_ids = features["masked_lm_ids"] 133 | masked_lm_weights = features["masked_lm_weights"] 134 | next_sentence_labels = features["next_sentence_labels"] 135 | 136 | is_training = (mode == tf.estimator.ModeKeys.TRAIN) 137 | 138 | model = modeling.BertModel( 139 | config=bert_config, 140 | is_training=is_training, 141 | input_ids=input_ids, 142 | input_mask=input_mask, 143 | token_type_ids=segment_ids, 144 | use_one_hot_embeddings=use_one_hot_embeddings) 145 | 146 | (masked_lm_loss, 147 | masked_lm_example_loss, masked_lm_log_probs) = get_masked_lm_output( 148 | bert_config, model.get_sequence_output(), model.get_embedding_table(), 149 | masked_lm_positions, masked_lm_ids, masked_lm_weights) 150 | 151 | (next_sentence_loss, next_sentence_example_loss, 152 | next_sentence_log_probs) = get_next_sentence_output( 153 | bert_config, model.get_pooled_output(), next_sentence_labels) 154 | 155 | total_loss = masked_lm_loss + next_sentence_loss 156 | 157 | tvars = tf.trainable_variables() 158 | 159 | initialized_variable_names = {} 160 | scaffold_fn = None 161 | if init_checkpoint: 162 | (assignment_map, initialized_variable_names 163 | ) = modeling.get_assignment_map_from_checkpoint(tvars, init_checkpoint) 164 | if use_tpu: 165 | 166 | def tpu_scaffold(): 167 | tf.train.init_from_checkpoint(init_checkpoint, assignment_map) 168 | return tf.train.Scaffold() 169 | 170 | scaffold_fn = tpu_scaffold 171 | else: 172 | tf.train.init_from_checkpoint(init_checkpoint, assignment_map) 173 | 174 | tf.logging.info("**** Trainable Variables ****") 175 | for var in tvars: 176 | init_string = "" 177 | if var.name in initialized_variable_names: 178 | init_string = ", *INIT_FROM_CKPT*" 179 | tf.logging.info(" name = %s, shape = %s%s", var.name, var.shape, 180 | init_string) 181 | 182 | output_spec = None 183 | if mode == tf.estimator.ModeKeys.TRAIN: 184 | train_op = optimization.create_optimizer( 185 | total_loss, learning_rate, num_train_steps, num_warmup_steps, False) 186 | output_spec = tf.estimator.EstimatorSpec( 187 | mode=mode, 188 | loss=total_loss, 189 | train_op=train_op, 190 | scaffold=scaffold_fn) 191 | elif mode == tf.estimator.ModeKeys.EVAL: 192 | 193 | def metric_fn(masked_lm_example_loss, masked_lm_log_probs, masked_lm_ids, 194 | masked_lm_weights, next_sentence_example_loss, 195 | next_sentence_log_probs, next_sentence_labels): 196 | """Computes the loss and accuracy of the model.""" 197 | masked_lm_log_probs = tf.reshape(masked_lm_log_probs, 198 | [-1, masked_lm_log_probs.shape[-1]]) 199 | masked_lm_predictions = tf.argmax( 200 | masked_lm_log_probs, axis=-1, output_type=tf.int32) 201 | masked_lm_example_loss = tf.reshape(masked_lm_example_loss, [-1]) 202 | masked_lm_ids = tf.reshape(masked_lm_ids, [-1]) 203 | masked_lm_weights = tf.reshape(masked_lm_weights, [-1]) 204 | masked_lm_accuracy = tf.metrics.accuracy( 205 | labels=masked_lm_ids, 206 | predictions=masked_lm_predictions, 207 | weights=masked_lm_weights) 208 | masked_lm_mean_loss = tf.metrics.mean( 209 | values=masked_lm_example_loss, weights=masked_lm_weights) 210 | 211 | next_sentence_log_probs = tf.reshape( 212 | next_sentence_log_probs, [-1, next_sentence_log_probs.shape[-1]]) 213 | next_sentence_predictions = tf.argmax( 214 | next_sentence_log_probs, axis=-1, output_type=tf.int32) 215 | next_sentence_labels = tf.reshape(next_sentence_labels, [-1]) 216 | next_sentence_accuracy = tf.metrics.accuracy( 217 | labels=next_sentence_labels, predictions=next_sentence_predictions) 218 | next_sentence_mean_loss = tf.metrics.mean( 219 | values=next_sentence_example_loss) 220 | 221 | return { 222 | "masked_lm_accuracy": masked_lm_accuracy, 223 | "masked_lm_loss": masked_lm_mean_loss, 224 | "next_sentence_accuracy": next_sentence_accuracy, 225 | "next_sentence_loss": next_sentence_mean_loss, 226 | } 227 | 228 | eval_metrics = (metric_fn, [ 229 | masked_lm_example_loss, masked_lm_log_probs, masked_lm_ids, 230 | masked_lm_weights, next_sentence_example_loss, 231 | next_sentence_log_probs, next_sentence_labels 232 | ]) 233 | output_spec = tf.contrib.tpu.TPUEstimatorSpec( 234 | mode=mode, 235 | loss=total_loss, 236 | eval_metrics=eval_metrics, 237 | scaffold_fn=scaffold_fn) 238 | else: 239 | raise ValueError("Only TRAIN and EVAL modes are supported: %s" % (mode)) 240 | 241 | return output_spec 242 | 243 | return model_fn 244 | 245 | 246 | def get_masked_lm_output(bert_config, input_tensor, output_weights, positions, 247 | label_ids, label_weights): 248 | """Get loss and log probs for the masked LM.""" 249 | input_tensor = gather_indexes(input_tensor, positions) 250 | 251 | with tf.variable_scope("cls/predictions"): 252 | # We apply one more non-linear transformation before the output layer. 253 | # This matrix is not used after pre-training. 254 | with tf.variable_scope("transform"): 255 | input_tensor = tf.layers.dense( 256 | input_tensor, 257 | units=bert_config.hidden_size, 258 | activation=modeling.get_activation(bert_config.hidden_act), 259 | kernel_initializer=modeling.create_initializer( 260 | bert_config.initializer_range)) 261 | input_tensor = modeling.layer_norm(input_tensor) 262 | 263 | # The output weights are the same as the input embeddings, but there is 264 | # an output-only bias for each token. 265 | output_bias = tf.get_variable( 266 | "output_bias", 267 | shape=[bert_config.vocab_size], 268 | initializer=tf.zeros_initializer()) 269 | logits = tf.matmul(input_tensor, output_weights, transpose_b=True) 270 | logits = tf.nn.bias_add(logits, output_bias) 271 | log_probs = tf.nn.log_softmax(logits, axis=-1) 272 | 273 | label_ids = tf.reshape(label_ids, [-1]) 274 | label_weights = tf.reshape(label_weights, [-1]) 275 | 276 | one_hot_labels = tf.one_hot( 277 | label_ids, depth=bert_config.vocab_size, dtype=tf.float32) 278 | 279 | # The `positions` tensor might be zero-padded (if the sequence is too 280 | # short to have the maximum number of predictions). The `label_weights` 281 | # tensor has a value of 1.0 for every real prediction and 0.0 for the 282 | # padding predictions. 283 | per_example_loss = -tf.reduce_sum(log_probs * one_hot_labels, axis=[-1]) 284 | numerator = tf.reduce_sum(label_weights * per_example_loss) 285 | denominator = tf.reduce_sum(label_weights) + 1e-5 286 | loss = numerator / denominator 287 | 288 | return (loss, per_example_loss, log_probs) 289 | 290 | 291 | def get_next_sentence_output(bert_config, input_tensor, labels): 292 | """Get loss and log probs for the next sentence prediction.""" 293 | 294 | # Simple binary classification. Note that 0 is "next sentence" and 1 is 295 | # "random sentence". This weight matrix is not used after pre-training. 296 | with tf.variable_scope("cls/seq_relationship"): 297 | output_weights = tf.get_variable( 298 | "output_weights", 299 | shape=[2, bert_config.hidden_size], 300 | initializer=modeling.create_initializer(bert_config.initializer_range)) 301 | output_bias = tf.get_variable( 302 | "output_bias", shape=[2], initializer=tf.zeros_initializer()) 303 | 304 | logits = tf.matmul(input_tensor, output_weights, transpose_b=True) 305 | logits = tf.nn.bias_add(logits, output_bias) 306 | log_probs = tf.nn.log_softmax(logits, axis=-1) 307 | labels = tf.reshape(labels, [-1]) 308 | one_hot_labels = tf.one_hot(labels, depth=2, dtype=tf.float32) 309 | per_example_loss = -tf.reduce_sum(one_hot_labels * log_probs, axis=-1) 310 | loss = tf.reduce_mean(per_example_loss) 311 | return (loss, per_example_loss, log_probs) 312 | 313 | 314 | def gather_indexes(sequence_tensor, positions): 315 | """Gathers the vectors at the specific positions over a minibatch.""" 316 | sequence_shape = modeling.get_shape_list(sequence_tensor, expected_rank=3) 317 | batch_size = sequence_shape[0] 318 | seq_length = sequence_shape[1] 319 | width = sequence_shape[2] 320 | 321 | flat_offsets = tf.reshape( 322 | tf.range(0, batch_size, dtype=tf.int32) * seq_length, [-1, 1]) 323 | flat_positions = tf.reshape(positions + flat_offsets, [-1]) 324 | flat_sequence_tensor = tf.reshape(sequence_tensor, 325 | [batch_size * seq_length, width]) 326 | output_tensor = tf.gather(flat_sequence_tensor, flat_positions) 327 | return output_tensor 328 | 329 | 330 | def input_fn_builder(input_files, 331 | max_seq_length, 332 | max_predictions_per_seq, 333 | is_training, 334 | num_cpu_threads=4): 335 | """Creates an `input_fn` closure to be passed to TPUEstimator.""" 336 | 337 | def input_fn(params): 338 | """The actual input function.""" 339 | batch_size = FLAGS.train_batch_size 340 | 341 | name_to_features = { 342 | "input_ids": 343 | tf.FixedLenFeature([max_seq_length], tf.int64), 344 | "input_mask": 345 | tf.FixedLenFeature([max_seq_length], tf.int64), 346 | "segment_ids": 347 | tf.FixedLenFeature([max_seq_length], tf.int64), 348 | "masked_lm_positions": 349 | tf.FixedLenFeature([max_predictions_per_seq], tf.int64), 350 | "masked_lm_ids": 351 | tf.FixedLenFeature([max_predictions_per_seq], tf.int64), 352 | "masked_lm_weights": 353 | tf.FixedLenFeature([max_predictions_per_seq], tf.float32), 354 | "next_sentence_labels": 355 | tf.FixedLenFeature([1], tf.int64), 356 | } 357 | 358 | # For training, we want a lot of parallel reading and shuffling. 359 | # For eval, we want no shuffling and parallel reading doesn't matter. 360 | if is_training: 361 | d = tf.data.Dataset.from_tensor_slices(tf.constant(input_files)) 362 | d = d.repeat() 363 | d = d.shuffle(buffer_size=len(input_files)) 364 | 365 | # `cycle_length` is the number of parallel files that get read. 366 | cycle_length = min(num_cpu_threads, len(input_files)) 367 | 368 | # `sloppy` mode means that the interleaving is not exact. This adds 369 | # even more randomness to the training pipeline. 370 | d = d.apply( 371 | tf.contrib.data.parallel_interleave( 372 | tf.data.TFRecordDataset, 373 | sloppy=is_training, 374 | cycle_length=cycle_length)) 375 | d = d.shuffle(buffer_size=100) 376 | else: 377 | d = tf.data.TFRecordDataset(input_files) 378 | # Since we evaluate for a fixed number of steps we don't want to encounter 379 | # out-of-range exceptions. 380 | d = d.repeat() 381 | 382 | # We must `drop_remainder` on training because the TPU requires fixed 383 | # size dimensions. For eval, we assume we are evaluating on the CPU or GPU 384 | # and we *don't* want to drop the remainder, otherwise we wont cover 385 | # every sample. 386 | d = d.apply( 387 | tf.contrib.data.map_and_batch( 388 | lambda record: _decode_record(record, name_to_features), 389 | batch_size=batch_size, 390 | num_parallel_batches=num_cpu_threads, 391 | drop_remainder=True)) 392 | return d 393 | 394 | return input_fn 395 | 396 | 397 | def _decode_record(record, name_to_features): 398 | """Decodes a record to a TensorFlow example.""" 399 | example = tf.parse_single_example(record, name_to_features) 400 | 401 | # tf.Example only supports tf.int64, but the TPU only supports tf.int32. 402 | # So cast all int64 to int32. 403 | for name in list(example.keys()): 404 | t = example[name] 405 | if t.dtype == tf.int64: 406 | t = tf.to_int32(t) 407 | example[name] = t 408 | 409 | return example 410 | 411 | 412 | def main(_): 413 | tf.logging.set_verbosity(tf.logging.INFO) 414 | 415 | if not FLAGS.do_train and not FLAGS.do_eval: 416 | raise ValueError("At least one of `do_train` or `do_eval` must be True.") 417 | 418 | bert_config = modeling.BertConfig.from_json_file(FLAGS.bert_config_file) 419 | 420 | tf.gfile.MakeDirs(FLAGS.output_dir) 421 | 422 | input_files = [] 423 | for input_pattern in FLAGS.input_file.split(","): 424 | input_files.extend(tf.gfile.Glob(input_pattern)) 425 | 426 | tf.logging.info("*** Input Files ***") 427 | for input_file in input_files: 428 | tf.logging.info(" %s" % input_file) 429 | 430 | tpu_cluster_resolver = None 431 | if FLAGS.use_tpu and FLAGS.tpu_name: 432 | tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver( 433 | FLAGS.tpu_name, zone=FLAGS.tpu_zone, project=FLAGS.gcp_project) 434 | 435 | is_per_host = tf.contrib.tpu.InputPipelineConfig.PER_HOST_V2 436 | 437 | dist_strategy = tf.contrib.distribute.MirroredStrategy( 438 | num_gpus=FLAGS.n_gpus, 439 | cross_device_ops=AllReduceCrossDeviceOps('nccl', num_packs=FLAGS.n_gpus), 440 | ) 441 | 442 | ''' IF ERROR COULD TRY 443 | dist_strategy = tf.contrib.distribute.MirroredStrategy( 444 | devices=["device:GPU:%d" % i for i in range(FLAGS.n_gpus)], 445 | cross_tower_ops=tf.distribute.HierarchicalCopyAllReduce()) 446 | ''' 447 | 448 | log_every_n_steps = 8 449 | run_config = RunConfig( 450 | train_distribute=dist_strategy, 451 | eval_distribute=dist_strategy, 452 | log_step_count_steps=log_every_n_steps, 453 | model_dir=FLAGS.output_dir, 454 | save_checkpoints_steps=FLAGS.save_checkpoints_steps) 455 | 456 | 457 | model_fn = model_fn_builder( 458 | bert_config=bert_config, 459 | init_checkpoint=FLAGS.init_checkpoint, 460 | learning_rate=FLAGS.learning_rate, 461 | num_train_steps=FLAGS.num_train_steps, 462 | num_warmup_steps=FLAGS.num_warmup_steps, 463 | use_tpu=FLAGS.use_tpu, 464 | use_one_hot_embeddings=FLAGS.use_tpu) 465 | 466 | # If TPU is not available, this will fall back to normal Estimator on CPU 467 | # or GPU. 468 | estimator = Estimator( 469 | model_fn=model_fn, 470 | params={}, 471 | config=run_config) 472 | 473 | if FLAGS.do_train: 474 | tf.logging.info("***** Running training *****") 475 | tf.logging.info(" Batch size = %d", FLAGS.train_batch_size) 476 | train_input_fn = input_fn_builder( 477 | input_files=input_files, 478 | max_seq_length=FLAGS.max_seq_length, 479 | max_predictions_per_seq=FLAGS.max_predictions_per_seq, 480 | is_training=True) 481 | estimator.train(input_fn=train_input_fn, max_steps=FLAGS.num_train_steps) 482 | 483 | if FLAGS.do_eval: 484 | tf.logging.info("***** Running evaluation *****") 485 | tf.logging.info(" Batch size = %d", FLAGS.eval_batch_size) 486 | 487 | eval_input_fn = input_fn_builder( 488 | input_files=input_files, 489 | max_seq_length=FLAGS.max_seq_length, 490 | max_predictions_per_seq=FLAGS.max_predictions_per_seq, 491 | is_training=False) 492 | 493 | result = estimator.evaluate( 494 | input_fn=eval_input_fn, steps=FLAGS.max_eval_steps) 495 | 496 | output_eval_file = os.path.join(FLAGS.output_dir, "eval_results.txt") 497 | with tf.gfile.GFile(output_eval_file, "w") as writer: 498 | tf.logging.info("***** Eval results *****") 499 | for key in sorted(result.keys()): 500 | tf.logging.info(" %s = %s", key, str(result[key])) 501 | writer.write("%s = %s\n" % (key, str(result[key]))) 502 | 503 | 504 | if __name__ == "__main__": 505 | flags.mark_flag_as_required("input_file") 506 | flags.mark_flag_as_required("bert_config_file") 507 | flags.mark_flag_as_required("output_dir") 508 | tf.app.run() 509 | -------------------------------------------------------------------------------- /run_pretraining_gpu.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # Copyright 2018 The Google AI Language Team Authors. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | """Run masked LM/next sentence masked_lm pre-training for BERT.""" 16 | 17 | from __future__ import absolute_import 18 | from __future__ import division 19 | from __future__ import print_function 20 | import time 21 | import os 22 | import modeling 23 | # import optimization_gpu 24 | import tensorflow as tf 25 | 26 | flags = tf.flags 27 | 28 | FLAGS = flags.FLAGS 29 | 30 | ## Required parameters 31 | flags.DEFINE_integer( 32 | "n_gpus", 10, 33 | "GPU number") 34 | 35 | flags.DEFINE_string( 36 | "bert_config_file", "bert_config.json", 37 | "The config json file corresponding to the pre-trained BERT model. " 38 | "This specifies the model architecture.") 39 | 40 | flags.DEFINE_string( 41 | "input_file", "tmp_data_128/sample1.tfrecords,tmp_data_128/sample2.tfrecords,tmp_data_128/sample3.tfrecords,tmp_data_128/sample4.tfrecords,tmp_data_128/sample5.tfrecords,tmp_data_128/sample6.tfrecords,tmp_data_128/sample7.tfrecords,tmp_data_128/sample8.tfrecords,tmp_data_128/sample9.tfrecords,tmp_data_128/sample10.tfrecords", 42 | "Input TF example files (can be a glob or comma separated).") 43 | 44 | flags.DEFINE_string( 45 | "output_dir", "output", 46 | "The output directory where the model checkpoints will be written.") 47 | 48 | ## Other parameters 49 | flags.DEFINE_string( 50 | "init_checkpoint", None, 51 | "Initial checkpoint (usually from a pre-trained BERT model).") 52 | 53 | flags.DEFINE_integer( 54 | "max_seq_length", 128, 55 | "The maximum total input sequence length after WordPiece tokenization. " 56 | "Sequences longer than this will be truncated, and sequences shorter " 57 | "than this will be padded. Must match data generation.") 58 | 59 | flags.DEFINE_integer( 60 | "max_predictions_per_seq", 20, 61 | "Maximum number of masked LM predictions per sequence. " 62 | "Must match data generation.") 63 | 64 | flags.DEFINE_bool("do_train", True, "Whether to run training.") 65 | 66 | flags.DEFINE_bool("do_eval", False, "Whether to run eval on the dev set.") 67 | 68 | flags.DEFINE_integer("train_batch_size", 32, "Total batch size for training.") 69 | 70 | flags.DEFINE_integer("eval_batch_size", 8, "Total batch size for eval.") 71 | 72 | flags.DEFINE_float("learning_rate", 5e-5, "The initial learning rate for Adam.") 73 | 74 | flags.DEFINE_integer("num_train_steps", 100000, "Number of training steps.") 75 | 76 | flags.DEFINE_integer("num_warmup_steps", 10000, "Number of warmup steps.") 77 | 78 | flags.DEFINE_integer("save_checkpoints_steps", 1000, 79 | "How often to save the model checkpoint.") 80 | 81 | flags.DEFINE_integer("iterations_per_loop", 1000, 82 | "How many steps to make in each estimator call.") 83 | 84 | flags.DEFINE_integer("max_eval_steps", 100, "Maximum number of eval steps.") 85 | 86 | flags.DEFINE_bool("use_tpu", False, "Whether to use TPU or GPU/CPU.") 87 | 88 | tf.flags.DEFINE_string( 89 | "tpu_name", None, 90 | "The Cloud TPU to use for training. This should be either the name " 91 | "used when creating the Cloud TPU, or a grpc://ip.address.of.tpu:8470 " 92 | "url.") 93 | 94 | tf.flags.DEFINE_string( 95 | "tpu_zone", None, 96 | "[Optional] GCE zone where the Cloud TPU is located in. If not " 97 | "specified, we will attempt to automatically detect the GCE project from " 98 | "metadata.") 99 | 100 | tf.flags.DEFINE_string( 101 | "gcp_project", None, 102 | "[Optional] Project name for the Cloud TPU-enabled project. If not " 103 | "specified, we will attempt to automatically detect the GCE project from " 104 | "metadata.") 105 | 106 | tf.flags.DEFINE_string("master", None, "[Optional] TensorFlow master URL.") 107 | 108 | flags.DEFINE_integer( 109 | "num_tpu_cores", 8, 110 | "Only used if `use_tpu` is True. Total number of TPU cores to use.") 111 | 112 | 113 | 114 | def _deduplicate_indexed_slices(values, indices): 115 | """Sums `values` associated with any non-unique `indices`. 116 | Args: 117 | values: A `Tensor` with rank >= 1. 118 | indices: A one-dimensional integer `Tensor`, indexing into the first 119 | dimension of `values` (as in an IndexedSlices object). 120 | Returns: 121 | A tuple of (`summed_values`, `unique_indices`) where `unique_indices` is a 122 | de-duplicated version of `indices` and `summed_values` contains the sum of 123 | `values` slices associated with each unique index. 124 | """ 125 | unique_indices, new_index_positions = tf.unique(indices) 126 | summed_values = tf.unsorted_segment_sum( 127 | values, new_index_positions, 128 | tf.shape(unique_indices)[0]) 129 | return (summed_values, unique_indices) 130 | 131 | 132 | def average_gradients(tower_grads, batch_size, options): 133 | # calculate average gradient for each shared variable across all GPUs 134 | average_grads = [] 135 | count = 0 136 | for grad_and_vars in zip(*tower_grads): 137 | # Note that each grad_and_vars looks like the following: 138 | # ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN)) 139 | # We need to average the gradients across each GPU. 140 | count += 1 141 | g0, v0 = grad_and_vars[0] 142 | 143 | if g0 is None: 144 | # no gradient for this variable, skip it 145 | average_grads.append((g0, v0)) 146 | continue 147 | 148 | if isinstance(g0, tf.IndexedSlices): 149 | # If the gradient is type IndexedSlices then this is a sparse 150 | # gradient with attributes indices and values. 151 | # To average, need to concat them individually then create 152 | # a new IndexedSlices object. 153 | indices = [] 154 | values = [] 155 | for g, v in grad_and_vars: 156 | indices.append(g.indices) 157 | values.append(g.values) 158 | all_indices = tf.concat(indices, 0) 159 | avg_values = tf.concat(values, 0) / len(grad_and_vars) 160 | # deduplicate across indices 161 | av, ai = _deduplicate_indexed_slices(avg_values, all_indices) 162 | grad = tf.IndexedSlices(av, ai, dense_shape=g0.dense_shape) 163 | 164 | else: 165 | # a normal tensor can just do a simple average 166 | grads = [] 167 | for g, v in grad_and_vars: 168 | # Add 0 dimension to the gradients to represent the tower. 169 | expanded_g = tf.expand_dims(g, 0) 170 | # Append on a 'tower' dimension which we will average over 171 | grads.append(expanded_g) 172 | 173 | # Average over the 'tower' dimension. 174 | grad = tf.concat(grads, 0) 175 | grad = tf.reduce_mean(grad, 0) 176 | 177 | # the Variables are redundant because they are shared 178 | # across towers. So.. just return the first tower's pointer to 179 | # the Variable. 180 | v = grad_and_vars[0][1] 181 | grad_and_var = (grad, v) 182 | 183 | average_grads.append(grad_and_var) 184 | 185 | assert len(average_grads) == len(list(zip(*tower_grads))) 186 | 187 | return average_grads 188 | 189 | 190 | def clip_by_global_norm_summary(t_list, clip_norm, norm_name, variables): 191 | # wrapper around tf.clip_by_global_norm that also does summary ops of norms 192 | 193 | # compute norms 194 | # use global_norm with one element to handle IndexedSlices vs dense 195 | norms = [tf.global_norm([t]) for t in t_list] 196 | 197 | # summary ops before clipping 198 | summary_ops = [] 199 | for ns, v in zip(norms, variables): 200 | name = 'norm_pre_clip/' + v.name.replace(":", "_") 201 | summary_ops.append(tf.summary.scalar(name, ns)) 202 | 203 | # clip 204 | clipped_t_list, tf_norm = tf.clip_by_global_norm(t_list, clip_norm) 205 | 206 | # summary ops after clipping 207 | norms_post = [tf.global_norm([t]) for t in clipped_t_list] 208 | for ns, v in zip(norms_post, variables): 209 | name = 'norm_post_clip/' + v.name.replace(":", "_") 210 | summary_ops.append(tf.summary.scalar(name, ns)) 211 | 212 | summary_ops.append(tf.summary.scalar(norm_name, tf_norm)) 213 | 214 | return clipped_t_list, tf_norm, summary_ops 215 | 216 | 217 | def clip_grads(grads, all_clip_norm_val, do_summaries): 218 | # grads = [(grad1, var1), (grad2, var2), ...] 219 | def _clip_norms(grad_and_vars, val, name): 220 | # grad_and_vars is a list of (g, v) pairs 221 | grad_tensors = [g for g, v in grad_and_vars] 222 | vv = [v for g, v in grad_and_vars] 223 | scaled_val = val 224 | if do_summaries: 225 | clipped_tensors, g_norm, so = clip_by_global_norm_summary( 226 | grad_tensors, scaled_val, name, vv) 227 | else: 228 | so = [] 229 | clipped_tensors, g_norm = tf.clip_by_global_norm( 230 | grad_tensors, scaled_val) 231 | 232 | ret = [] 233 | for t, (g, v) in zip(clipped_tensors, grad_and_vars): 234 | ret.append((t, v)) 235 | 236 | return ret, so 237 | 238 | ret, summary_ops = _clip_norms(grads, all_clip_norm_val, 'norm_grad') 239 | 240 | assert len(ret) == len(grads) 241 | 242 | return ret, summary_ops 243 | 244 | def model_fn_builder(bert_config, init_checkpoint, learning_rate, 245 | num_train_steps, num_warmup_steps, use_tpu, 246 | use_one_hot_embeddings): 247 | """Returns `model_fn` closure for TPUEstimator.""" 248 | 249 | def model_fn(features, labels, mode, params): # pylint: disable=unused-argument 250 | """The `model_fn` for TPUEstimator.""" 251 | 252 | tf.logging.info("*** Features ***") 253 | for name in sorted(features.keys()): 254 | tf.logging.info(" name = %s, shape = %s" % (name, features[name].shape)) 255 | 256 | input_ids = features["input_ids"] 257 | input_mask = features["input_mask"] 258 | segment_ids = features["segment_ids"] 259 | masked_lm_positions = features["masked_lm_positions"] 260 | masked_lm_ids = features["masked_lm_ids"] 261 | masked_lm_weights = features["masked_lm_weights"] 262 | next_sentence_labels = features["next_sentence_labels"] 263 | 264 | is_training = (mode == tf.estimator.ModeKeys.TRAIN) 265 | 266 | model = modeling.BertModel( 267 | config=bert_config, 268 | is_training=is_training, 269 | input_ids=input_ids, 270 | input_mask=input_mask, 271 | token_type_ids=segment_ids, 272 | use_one_hot_embeddings=use_one_hot_embeddings) 273 | 274 | (masked_lm_loss, 275 | masked_lm_example_loss, masked_lm_log_probs) = get_masked_lm_output( 276 | bert_config, model.get_sequence_output(), model.get_embedding_table(), 277 | masked_lm_positions, masked_lm_ids, masked_lm_weights) 278 | 279 | (next_sentence_loss, next_sentence_example_loss, 280 | next_sentence_log_probs) = get_next_sentence_output( 281 | bert_config, model.get_pooled_output(), next_sentence_labels) 282 | 283 | total_loss = masked_lm_loss + next_sentence_loss 284 | 285 | tvars = tf.trainable_variables() 286 | 287 | initialized_variable_names = {} 288 | scaffold_fn = None 289 | if init_checkpoint: 290 | (assignment_map, initialized_variable_names 291 | ) = modeling.get_assignment_map_from_checkpoint(tvars, init_checkpoint) 292 | if use_tpu: 293 | 294 | def tpu_scaffold(): 295 | tf.train.init_from_checkpoint(init_checkpoint, assignment_map) 296 | return tf.train.Scaffold() 297 | 298 | scaffold_fn = tpu_scaffold 299 | else: 300 | tf.train.init_from_checkpoint(init_checkpoint, assignment_map) 301 | 302 | tf.logging.info("**** Trainable Variables ****") 303 | for var in tvars: 304 | init_string = "" 305 | if var.name in initialized_variable_names: 306 | init_string = ", *INIT_FROM_CKPT*" 307 | tf.logging.info(" name = %s, shape = %s%s", var.name, var.shape, 308 | init_string) 309 | 310 | output_spec = None 311 | if mode == tf.estimator.ModeKeys.TRAIN: 312 | train_op = optimization.create_optimizer( 313 | total_loss, learning_rate, num_train_steps, num_warmup_steps, use_tpu) 314 | 315 | output_spec = tf.contrib.tpu.TPUEstimatorSpec( 316 | mode=mode, 317 | loss=total_loss, 318 | train_op=train_op, 319 | scaffold_fn=scaffold_fn) 320 | elif mode == tf.estimator.ModeKeys.EVAL: 321 | 322 | def metric_fn(masked_lm_example_loss, masked_lm_log_probs, masked_lm_ids, 323 | masked_lm_weights, next_sentence_example_loss, 324 | next_sentence_log_probs, next_sentence_labels): 325 | """Computes the loss and accuracy of the model.""" 326 | masked_lm_log_probs = tf.reshape(masked_lm_log_probs, 327 | [-1, masked_lm_log_probs.shape[-1]]) 328 | masked_lm_predictions = tf.argmax( 329 | masked_lm_log_probs, axis=-1, output_type=tf.int32) 330 | masked_lm_example_loss = tf.reshape(masked_lm_example_loss, [-1]) 331 | masked_lm_ids = tf.reshape(masked_lm_ids, [-1]) 332 | masked_lm_weights = tf.reshape(masked_lm_weights, [-1]) 333 | masked_lm_accuracy = tf.metrics.accuracy( 334 | labels=masked_lm_ids, 335 | predictions=masked_lm_predictions, 336 | weights=masked_lm_weights) 337 | masked_lm_mean_loss = tf.metrics.mean( 338 | values=masked_lm_example_loss, weights=masked_lm_weights) 339 | 340 | next_sentence_log_probs = tf.reshape( 341 | next_sentence_log_probs, [-1, next_sentence_log_probs.shape[-1]]) 342 | next_sentence_predictions = tf.argmax( 343 | next_sentence_log_probs, axis=-1, output_type=tf.int32) 344 | next_sentence_labels = tf.reshape(next_sentence_labels, [-1]) 345 | next_sentence_accuracy = tf.metrics.accuracy( 346 | labels=next_sentence_labels, predictions=next_sentence_predictions) 347 | next_sentence_mean_loss = tf.metrics.mean( 348 | values=next_sentence_example_loss) 349 | 350 | return { 351 | "masked_lm_accuracy": masked_lm_accuracy, 352 | "masked_lm_loss": masked_lm_mean_loss, 353 | "next_sentence_accuracy": next_sentence_accuracy, 354 | "next_sentence_loss": next_sentence_mean_loss, 355 | } 356 | 357 | eval_metrics = (metric_fn, [ 358 | masked_lm_example_loss, masked_lm_log_probs, masked_lm_ids, 359 | masked_lm_weights, next_sentence_example_loss, 360 | next_sentence_log_probs, next_sentence_labels 361 | ]) 362 | output_spec = tf.contrib.tpu.TPUEstimatorSpec( 363 | mode=mode, 364 | loss=total_loss, 365 | eval_metrics=eval_metrics, 366 | scaffold_fn=scaffold_fn) 367 | else: 368 | raise ValueError("Only TRAIN and EVAL modes are supported: %s" % (mode)) 369 | 370 | return output_spec 371 | 372 | return model_fn 373 | 374 | 375 | def get_masked_lm_output(bert_config, input_tensor, output_weights, positions, 376 | label_ids, label_weights): 377 | """Get loss and log probs for the masked LM.""" 378 | input_tensor = gather_indexes(input_tensor, positions) 379 | 380 | with tf.variable_scope("cls/predictions"): 381 | # We apply one more non-linear transformation before the output layer. 382 | # This matrix is not used after pre-training. 383 | with tf.variable_scope("transform"): 384 | input_tensor = tf.layers.dense( 385 | input_tensor, 386 | units=bert_config.hidden_size, 387 | activation=modeling.get_activation(bert_config.hidden_act), 388 | kernel_initializer=modeling.create_initializer( 389 | bert_config.initializer_range)) 390 | input_tensor = modeling.layer_norm(input_tensor) 391 | 392 | # The output weights are the same as the input embeddings, but there is 393 | # an output-only bias for each token. 394 | output_bias = tf.get_variable( 395 | "output_bias", 396 | shape=[bert_config.vocab_size], 397 | initializer=tf.zeros_initializer()) 398 | logits = tf.matmul(input_tensor, output_weights, transpose_b=True) 399 | logits = tf.nn.bias_add(logits, output_bias) 400 | log_probs = tf.nn.log_softmax(logits, axis=-1) 401 | 402 | label_ids = tf.reshape(label_ids, [-1]) 403 | label_weights = tf.reshape(label_weights, [-1]) 404 | 405 | one_hot_labels = tf.one_hot( 406 | label_ids, depth=bert_config.vocab_size, dtype=tf.float32) 407 | 408 | # The `positions` tensor might be zero-padded (if the sequence is too 409 | # short to have the maximum number of predictions). The `label_weights` 410 | # tensor has a value of 1.0 for every real prediction and 0.0 for the 411 | # padding predictions. 412 | per_example_loss = -tf.reduce_sum(log_probs * one_hot_labels, axis=[-1]) 413 | numerator = tf.reduce_sum(label_weights * per_example_loss) 414 | denominator = tf.reduce_sum(label_weights) + 1e-5 415 | loss = numerator / denominator 416 | 417 | return (loss, per_example_loss, log_probs) 418 | 419 | 420 | def get_next_sentence_output(bert_config, input_tensor, labels): 421 | """Get loss and log probs for the next sentence prediction.""" 422 | 423 | # Simple binary classification. Note that 0 is "next sentence" and 1 is 424 | # "random sentence". This weight matrix is not used after pre-training. 425 | with tf.variable_scope("cls/seq_relationship"): 426 | output_weights = tf.get_variable( 427 | "output_weights", 428 | shape=[2, bert_config.hidden_size], 429 | initializer=modeling.create_initializer(bert_config.initializer_range)) 430 | output_bias = tf.get_variable( 431 | "output_bias", shape=[2], initializer=tf.zeros_initializer()) 432 | 433 | logits = tf.matmul(input_tensor, output_weights, transpose_b=True) 434 | logits = tf.nn.bias_add(logits, output_bias) 435 | log_probs = tf.nn.log_softmax(logits, axis=-1) 436 | labels = tf.reshape(labels, [-1]) 437 | one_hot_labels = tf.one_hot(labels, depth=2, dtype=tf.float32) 438 | per_example_loss = -tf.reduce_sum(one_hot_labels * log_probs, axis=-1) 439 | loss = tf.reduce_mean(per_example_loss) 440 | return (loss, per_example_loss, log_probs) 441 | 442 | 443 | def gather_indexes(sequence_tensor, positions): 444 | """Gathers the vectors at the specific positions over a minibatch.""" 445 | sequence_shape = modeling.get_shape_list(sequence_tensor, expected_rank=3) 446 | batch_size = sequence_shape[0] 447 | seq_length = sequence_shape[1] 448 | width = sequence_shape[2] 449 | 450 | flat_offsets = tf.reshape( 451 | tf.range(0, batch_size, dtype=tf.int32) * seq_length, [-1, 1]) 452 | flat_positions = tf.reshape(positions + flat_offsets, [-1]) 453 | flat_sequence_tensor = tf.reshape(sequence_tensor, 454 | [batch_size * seq_length, width]) 455 | output_tensor = tf.gather(flat_sequence_tensor, flat_positions) 456 | return output_tensor 457 | 458 | 459 | 460 | def parse_input_fn_result(result): 461 | """Gets features, labels, and hooks from the result of an Estimator input_fn. 462 | 463 | Args: 464 | result: output of an input_fn to an estimator, which should be one of: 465 | 466 | * A 'tf.data.Dataset' object: Outputs of `Dataset` object must be a 467 | tuple (features, labels) with same constraints as below. 468 | * A tuple (features, labels): Where `features` is a `Tensor` or a 469 | dictionary of string feature name to `Tensor` and `labels` is a 470 | `Tensor` or a dictionary of string label name to `Tensor`. Both 471 | `features` and `labels` are consumed by `model_fn`. They should 472 | satisfy the expectation of `model_fn` from inputs. 473 | 474 | Returns: 475 | Tuple of features, labels, and input_hooks, where features are as described 476 | above, labels are as described above or None, and input_hooks are a list 477 | of SessionRunHooks to be included when running. 478 | 479 | Raises: 480 | ValueError: if the result is a list or tuple of length != 2. 481 | """ 482 | try: 483 | # We can't just check whether this is a tf.data.Dataset instance here, 484 | # as this is plausibly a PerDeviceDataset. Try treating as a dataset first. 485 | iterator = result.make_initializable_iterator() 486 | except AttributeError: 487 | # Not a dataset or dataset-like-object. Move along. 488 | pass 489 | else: 490 | result = iterator.get_next() 491 | return result,iterator 492 | 493 | 494 | def input_fn(input_files, 495 | batch_size, 496 | max_seq_length, 497 | max_predictions_per_seq, 498 | is_training, 499 | num_cpu_threads=4): 500 | """The actual input function.""" 501 | # batch_size = params["batch_size"] 502 | 503 | name_to_features = { 504 | "input_ids": 505 | tf.FixedLenFeature([max_seq_length], tf.int64), 506 | "input_mask": 507 | tf.FixedLenFeature([max_seq_length], tf.int64), 508 | "segment_ids": 509 | tf.FixedLenFeature([max_seq_length], tf.int64), 510 | "masked_lm_positions": 511 | tf.FixedLenFeature([max_predictions_per_seq], tf.int64), 512 | "masked_lm_ids": 513 | tf.FixedLenFeature([max_predictions_per_seq], tf.int64), 514 | "masked_lm_weights": 515 | tf.FixedLenFeature([max_predictions_per_seq], tf.float32), 516 | "next_sentence_labels": 517 | tf.FixedLenFeature([1], tf.int64), 518 | } 519 | 520 | # For training, we want a lot of parallel reading and shuffling. 521 | # For eval, we want no shuffling and parallel reading doesn't matter. 522 | if is_training: 523 | d = tf.data.Dataset.from_tensor_slices(tf.constant(input_files)) 524 | d = d.repeat() 525 | d = d.shuffle(buffer_size=len(input_files)) 526 | 527 | # `cycle_length` is the number of parallel files that get read. 528 | cycle_length = min(num_cpu_threads, len(input_files)) 529 | 530 | # `sloppy` mode means that the interleaving is not exact. This adds 531 | # even more randomness to the training pipeline. 532 | d = d.apply( 533 | tf.contrib.data.parallel_interleave( 534 | tf.data.TFRecordDataset, 535 | sloppy=is_training, 536 | cycle_length=cycle_length)) 537 | d = d.shuffle(buffer_size=100) 538 | else: 539 | d = tf.data.TFRecordDataset(input_files) 540 | # Since we evaluate for a fixed number of steps we don't want to encounter 541 | # out-of-range exceptions. 542 | d = d.repeat() 543 | 544 | # We must `drop_remainder` on training because the TPU requires fixed 545 | # size dimensions. For eval, we assume we are evaluating on the CPU or GPU 546 | # and we *don't* want to drop the remainder, otherwise we wont cover 547 | # every sample. 548 | d = d.apply( 549 | tf.contrib.data.map_and_batch( 550 | lambda record: _decode_record(record, name_to_features), 551 | batch_size=batch_size, 552 | num_parallel_batches=num_cpu_threads, 553 | drop_remainder=True)) 554 | return d 555 | 556 | 557 | def _decode_record(record, name_to_features): 558 | """Decodes a record to a TensorFlow example.""" 559 | example = tf.parse_single_example(record, name_to_features) 560 | 561 | # tf.Example only supports tf.int64, but the TPU only supports tf.int32. 562 | # So cast all int64 to int32. 563 | for name in list(example.keys()): 564 | t = example[name] 565 | if t.dtype == tf.int64: 566 | t = tf.to_int32(t) 567 | example[name] = t 568 | 569 | return example 570 | 571 | 572 | def main(_): 573 | mode = tf.estimator.ModeKeys.TRAIN 574 | use_one_hot_embeddings = FLAGS.use_tpu 575 | 576 | tf.logging.set_verbosity(tf.logging.INFO) 577 | 578 | if not FLAGS.do_train and not FLAGS.do_eval: 579 | raise ValueError("At least one of `do_train` or `do_eval` must be True.") 580 | 581 | bert_config = modeling.BertConfig.from_json_file(FLAGS.bert_config_file) 582 | 583 | #tf.gfile.MakeDirs(FLAGS.output_dir) 584 | 585 | input_files = [] 586 | for input_pattern in FLAGS.input_file.split(","): 587 | input_files.extend(tf.gfile.Glob(input_pattern)) 588 | 589 | tf.logging.info("*** Input Files ***") 590 | for input_file in input_files: 591 | tf.logging.info(" %s" % input_file) 592 | 593 | if FLAGS.do_train: 594 | tf.logging.info("***** Running training *****") 595 | tf.logging.info(" Batch size = %d", FLAGS.train_batch_size) 596 | n_gpus = FLAGS.n_gpus 597 | batch_size = FLAGS.train_batch_size 598 | d = input_fn(input_files,FLAGS.train_batch_size*n_gpus,FLAGS.max_seq_length, 599 | FLAGS.max_predictions_per_seq,True) 600 | features,iterator = parse_input_fn_result(d) 601 | # train_input_fn = input_fn_builder( 602 | # input_files=input_files, 603 | # max_seq_length=FLAGS.max_seq_length, 604 | # max_predictions_per_seq=FLAGS.max_predictions_per_seq, 605 | # is_training=True) 606 | # estimator.train(input_fn=train_input_fn, max_steps=FLAGS.num_train_steps) 607 | 608 | input_ids_list = tf.split(features["input_ids"], n_gpus, axis=0) 609 | input_mask_list = tf.split(features["input_mask"], n_gpus, axis=0) 610 | segment_ids_list = tf.split(features["segment_ids"], n_gpus, axis=0) 611 | masked_lm_positions_list = tf.split(features["masked_lm_positions"], n_gpus, axis=0) 612 | masked_lm_ids_list = tf.split(features["masked_lm_ids"], n_gpus, axis=0) 613 | masked_lm_weights_list = tf.split(features["masked_lm_weights"], n_gpus, axis=0) 614 | next_sentence_labels_list = tf.split(features["next_sentence_labels"], n_gpus, axis=0) 615 | 616 | # multi-gpu train 617 | 618 | # optimizer = optimization_gpu.create_optimizer( 619 | # None, FLAGS.learning_rate, FLAGS.num_train_steps, FLAGS.num_warmup_steps, False) 620 | optimizer = tf.train.AdamOptimizer(learning_rate=FLAGS.learning_rate) 621 | 622 | # calculate the gradients on each GPU 623 | tower_grads = [] 624 | models = [] 625 | loss_print = tf.get_variable( 626 | 'train_perplexity', [], 627 | initializer=tf.constant_initializer(0.0), trainable=False) 628 | for k in range(n_gpus): 629 | with tf.device('/gpu:%d' % k): 630 | with tf.variable_scope('lm', reuse=k > 0): 631 | # calculate the loss for one model replica and get 632 | # lstm states 633 | 634 | input_ids = input_ids_list[k] 635 | input_mask = input_mask_list[k] 636 | segment_ids = segment_ids_list[k] 637 | masked_lm_positions = masked_lm_positions_list[k] 638 | masked_lm_ids = masked_lm_ids_list[k] 639 | masked_lm_weights = masked_lm_weights_list[k] 640 | next_sentence_labels = next_sentence_labels_list[k] 641 | 642 | is_training = (mode == tf.estimator.ModeKeys.TRAIN) 643 | 644 | model = modeling.BertModel( 645 | config=bert_config, 646 | is_training=is_training, 647 | input_ids=input_ids, 648 | input_mask=input_mask, 649 | token_type_ids=segment_ids, 650 | use_one_hot_embeddings=use_one_hot_embeddings) 651 | 652 | (masked_lm_loss, 653 | masked_lm_example_loss, masked_lm_log_probs) = get_masked_lm_output( 654 | bert_config, model.get_sequence_output(), model.get_embedding_table(), 655 | masked_lm_positions, masked_lm_ids, masked_lm_weights) 656 | 657 | (next_sentence_loss, next_sentence_example_loss, 658 | next_sentence_log_probs) = get_next_sentence_output( 659 | bert_config, model.get_pooled_output(), next_sentence_labels) 660 | 661 | total_loss = masked_lm_loss + next_sentence_loss 662 | 663 | 664 | loss = total_loss 665 | models.append(model) 666 | # get gradients 667 | grads = optimizer.compute_gradients( 668 | loss, 669 | aggregation_method=tf.AggregationMethod.EXPERIMENTAL_TREE, 670 | ) 671 | tower_grads.append(grads) 672 | # keep track of loss across all GPUs 673 | loss_print += loss 674 | 675 | average_grads = average_gradients(tower_grads, None, None) 676 | average_grads, norm_summary_ops = clip_grads(average_grads, 10.0, True) 677 | loss_print = loss_print / n_gpus 678 | train_op = optimizer.apply_gradients(average_grads) 679 | init = tf.global_variables_initializer() 680 | saver = tf.train.Saver(tf.global_variables(), max_to_keep=2) 681 | with tf.Session(config=tf.ConfigProto( 682 | allow_soft_placement=True)) as sess: 683 | sess.run(init) 684 | sess.run(iterator.initializer) 685 | checkpoint_path = os.path.join(FLAGS.output_dir, 'model.ckpt') 686 | if os.path.exists(FLAGS.output_dir): 687 | saver.restore(sess, checkpoint_path) 688 | 689 | count = 0 690 | t0 = time.time() 691 | sum = 0 692 | while True: 693 | 694 | _, loss_print_ = sess.run([train_op, loss_print]) 695 | # optimistic_restore(sess, checkpoint_path + "-0") 696 | # loss_print_2 = sess.run([loss_print]) 697 | sum+=loss_print_ 698 | count += 1 699 | if count%300==0: 700 | print("------------") 701 | print(time.time() - t0," s") 702 | t0 = time.time() 703 | print("loss ",sum/count) 704 | sum=0 705 | count=0 706 | checkpoint_path = os.path.join(FLAGS.output_dir, 'model.ckpt') 707 | saver.save(sess, checkpoint_path) 708 | 709 | 710 | if FLAGS.do_eval: 711 | tf.logging.info("***** Running evaluation *****") 712 | tf.logging.info(" Batch size = %d", FLAGS.eval_batch_size) 713 | 714 | eval_input_fn = input_fn_builder( 715 | input_files=input_files, 716 | max_seq_length=FLAGS.max_seq_length, 717 | max_predictions_per_seq=FLAGS.max_predictions_per_seq, 718 | is_training=False) 719 | 720 | result = estimator.evaluate( 721 | input_fn=eval_input_fn, steps=FLAGS.max_eval_steps) 722 | 723 | output_eval_file = os.path.join(FLAGS.output_dir, "eval_results.txt") 724 | with tf.gfile.GFile(output_eval_file, "w") as writer: 725 | tf.logging.info("***** Eval results *****") 726 | for key in sorted(result.keys()): 727 | tf.logging.info(" %s = %s", key, str(result[key])) 728 | writer.write("%s = %s\n" % (key, str(result[key]))) 729 | 730 | 731 | if __name__ == "__main__": 732 | flags.mark_flag_as_required("input_file") 733 | flags.mark_flag_as_required("bert_config_file") 734 | flags.mark_flag_as_required("output_dir") 735 | tf.app.run() 736 | -------------------------------------------------------------------------------- /modeling.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # Copyright 2018 The Google AI Language Team Authors. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | """The main BERT model and related functions.""" 16 | 17 | from __future__ import absolute_import 18 | from __future__ import division 19 | from __future__ import print_function 20 | 21 | import collections 22 | import copy 23 | import json 24 | import math 25 | import re 26 | import numpy as np 27 | import six 28 | import tensorflow as tf 29 | 30 | 31 | class BertConfig(object): 32 | """Configuration for `BertModel`.""" 33 | 34 | def __init__(self, 35 | vocab_size, 36 | hidden_size=768, 37 | num_hidden_layers=12, 38 | num_attention_heads=12, 39 | intermediate_size=3072, 40 | hidden_act="gelu", 41 | hidden_dropout_prob=0.1, 42 | attention_probs_dropout_prob=0.1, 43 | max_position_embeddings=512, 44 | type_vocab_size=16, 45 | initializer_range=0.02): 46 | """Constructs BertConfig. 47 | Args: 48 | vocab_size: Vocabulary size of `inputs_ids` in `BertModel`. 49 | hidden_size: Size of the encoder layers and the pooler layer. 50 | num_hidden_layers: Number of hidden layers in the Transformer encoder. 51 | num_attention_heads: Number of attention heads for each attention layer in 52 | the Transformer encoder. 53 | intermediate_size: The size of the "intermediate" (i.e., feed-forward) 54 | layer in the Transformer encoder. 55 | hidden_act: The non-linear activation function (function or string) in the 56 | encoder and pooler. 57 | hidden_dropout_prob: The dropout probability for all fully connected 58 | layers in the embeddings, encoder, and pooler. 59 | attention_probs_dropout_prob: The dropout ratio for the attention 60 | probabilities. 61 | max_position_embeddings: The maximum sequence length that this model might 62 | ever be used with. Typically set this to something large just in case 63 | (e.g., 512 or 1024 or 2048). 64 | type_vocab_size: The vocabulary size of the `token_type_ids` passed into 65 | `BertModel`. 66 | initializer_range: The stdev of the truncated_normal_initializer for 67 | initializing all weight matrices. 68 | """ 69 | self.vocab_size = vocab_size 70 | self.hidden_size = hidden_size 71 | self.num_hidden_layers = num_hidden_layers 72 | self.num_attention_heads = num_attention_heads 73 | self.hidden_act = hidden_act 74 | self.intermediate_size = intermediate_size 75 | self.hidden_dropout_prob = hidden_dropout_prob 76 | self.attention_probs_dropout_prob = attention_probs_dropout_prob 77 | self.max_position_embeddings = max_position_embeddings 78 | self.type_vocab_size = type_vocab_size 79 | self.initializer_range = initializer_range 80 | 81 | @classmethod 82 | def from_dict(cls, json_object): 83 | """Constructs a `BertConfig` from a Python dictionary of parameters.""" 84 | config = BertConfig(vocab_size=None) 85 | for (key, value) in six.iteritems(json_object): 86 | config.__dict__[key] = value 87 | return config 88 | 89 | @classmethod 90 | def from_json_file(cls, json_file): 91 | """Constructs a `BertConfig` from a json file of parameters.""" 92 | with tf.gfile.GFile(json_file, "r") as reader: 93 | text = reader.read() 94 | return cls.from_dict(json.loads(text)) 95 | 96 | def to_dict(self): 97 | """Serializes this instance to a Python dictionary.""" 98 | output = copy.deepcopy(self.__dict__) 99 | return output 100 | 101 | def to_json_string(self): 102 | """Serializes this instance to a JSON string.""" 103 | return json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n" 104 | 105 | 106 | class BertModel(object): 107 | """BERT model ("Bidirectional Encoder Representations from Transformers"). 108 | Example usage: 109 | ```python 110 | # Already been converted into WordPiece token ids 111 | input_ids = tf.constant([[31, 51, 99], [15, 5, 0]]) 112 | input_mask = tf.constant([[1, 1, 1], [1, 1, 0]]) 113 | token_type_ids = tf.constant([[0, 0, 1], [0, 2, 0]]) 114 | config = modeling.BertConfig(vocab_size=32000, hidden_size=512, 115 | num_hidden_layers=8, num_attention_heads=6, intermediate_size=1024) 116 | model = modeling.BertModel(config=config, is_training=True, 117 | input_ids=input_ids, input_mask=input_mask, token_type_ids=token_type_ids) 118 | label_embeddings = tf.get_variable(...) 119 | pooled_output = model.get_pooled_output() 120 | logits = tf.matmul(pooled_output, label_embeddings) 121 | ... 122 | ``` 123 | """ 124 | 125 | def __init__(self, 126 | config, 127 | is_training, 128 | input_ids, 129 | input_mask=None, 130 | token_type_ids=None, 131 | use_one_hot_embeddings=False, 132 | scope=None): 133 | """Constructor for BertModel. 134 | Args: 135 | config: `BertConfig` instance. 136 | is_training: bool. true for training model, false for eval model. Controls 137 | whether dropout will be applied. 138 | input_ids: int32 Tensor of shape [batch_size, seq_length]. 139 | input_mask: (optional) int32 Tensor of shape [batch_size, seq_length]. 140 | token_type_ids: (optional) int32 Tensor of shape [batch_size, seq_length]. 141 | use_one_hot_embeddings: (optional) bool. Whether to use one-hot word 142 | embeddings or tf.embedding_lookup() for the word embeddings. 143 | scope: (optional) variable scope. Defaults to "bert". 144 | Raises: 145 | ValueError: The config is invalid or one of the input tensor shapes 146 | is invalid. 147 | """ 148 | config = copy.deepcopy(config) 149 | if not is_training: 150 | config.hidden_dropout_prob = 0.0 151 | config.attention_probs_dropout_prob = 0.0 152 | 153 | input_shape = get_shape_list(input_ids, expected_rank=2) 154 | batch_size = input_shape[0] 155 | seq_length = input_shape[1] 156 | 157 | if input_mask is None: 158 | input_mask = tf.ones(shape=[batch_size, seq_length], dtype=tf.int32) 159 | 160 | if token_type_ids is None: 161 | token_type_ids = tf.zeros(shape=[batch_size, seq_length], dtype=tf.int32) 162 | 163 | with tf.variable_scope(scope, default_name="bert"): 164 | with tf.variable_scope("embeddings"): 165 | # Perform embedding lookup on the word ids. 166 | (self.embedding_output, self.embedding_table) = embedding_lookup( 167 | input_ids=input_ids, 168 | vocab_size=config.vocab_size, 169 | embedding_size=config.hidden_size, 170 | initializer_range=config.initializer_range, 171 | word_embedding_name="word_embeddings", 172 | use_one_hot_embeddings=use_one_hot_embeddings) 173 | 174 | # Add positional embeddings and token type embeddings, then layer 175 | # normalize and perform dropout. 176 | self.embedding_output = embedding_postprocessor( 177 | input_tensor=self.embedding_output, 178 | use_token_type=True, 179 | token_type_ids=token_type_ids, 180 | token_type_vocab_size=config.type_vocab_size, 181 | token_type_embedding_name="token_type_embeddings", 182 | use_position_embeddings=True, 183 | position_embedding_name="position_embeddings", 184 | initializer_range=config.initializer_range, 185 | max_position_embeddings=config.max_position_embeddings, 186 | dropout_prob=config.hidden_dropout_prob) 187 | 188 | with tf.variable_scope("encoder"): 189 | # This converts a 2D mask of shape [batch_size, seq_length] to a 3D 190 | # mask of shape [batch_size, seq_length, seq_length] which is used 191 | # for the attention scores. 192 | attention_mask = create_attention_mask_from_input_mask( 193 | input_ids, input_mask) 194 | 195 | # Run the stacked transformer. 196 | # `sequence_output` shape = [batch_size, seq_length, hidden_size]. 197 | self.all_encoder_layers = transformer_model( 198 | input_tensor=self.embedding_output, 199 | attention_mask=attention_mask, 200 | hidden_size=config.hidden_size, 201 | num_hidden_layers=config.num_hidden_layers, 202 | num_attention_heads=config.num_attention_heads, 203 | intermediate_size=config.intermediate_size, 204 | intermediate_act_fn=get_activation(config.hidden_act), 205 | hidden_dropout_prob=config.hidden_dropout_prob, 206 | attention_probs_dropout_prob=config.attention_probs_dropout_prob, 207 | initializer_range=config.initializer_range, 208 | do_return_all_layers=True) 209 | 210 | self.sequence_output = self.all_encoder_layers[-1] 211 | # The "pooler" converts the encoded sequence tensor of shape 212 | # [batch_size, seq_length, hidden_size] to a tensor of shape 213 | # [batch_size, hidden_size]. This is necessary for segment-level 214 | # (or segment-pair-level) classification tasks where we need a fixed 215 | # dimensional representation of the segment. 216 | with tf.variable_scope("pooler"): 217 | # We "pool" the model by simply taking the hidden state corresponding 218 | # to the first token. We assume that this has been pre-trained 219 | first_token_tensor = tf.squeeze(self.sequence_output[:, 0:1, :], axis=1) 220 | self.pooled_output = tf.layers.dense( 221 | first_token_tensor, 222 | config.hidden_size, 223 | activation=tf.tanh, 224 | kernel_initializer=create_initializer(config.initializer_range)) 225 | 226 | def get_pooled_output(self): 227 | return self.pooled_output 228 | 229 | def get_sequence_output(self): 230 | """Gets final hidden layer of encoder. 231 | Returns: 232 | float Tensor of shape [batch_size, seq_length, hidden_size] corresponding 233 | to the final hidden of the transformer encoder. 234 | """ 235 | return self.sequence_output 236 | 237 | def get_all_encoder_layers(self): 238 | return self.all_encoder_layers 239 | 240 | def get_embedding_output(self): 241 | """Gets output of the embedding lookup (i.e., input to the transformer). 242 | Returns: 243 | float Tensor of shape [batch_size, seq_length, hidden_size] corresponding 244 | to the output of the embedding layer, after summing the word 245 | embeddings with the positional embeddings and the token type embeddings, 246 | then performing layer normalization. This is the input to the transformer. 247 | """ 248 | return self.embedding_output 249 | 250 | def get_embedding_table(self): 251 | return self.embedding_table 252 | 253 | 254 | def gelu(x): 255 | """Gaussian Error Linear Unit. 256 | This is a smoother version of the RELU. 257 | Original paper: https://arxiv.org/abs/1606.08415 258 | Args: 259 | x: float Tensor to perform activation. 260 | Returns: 261 | `x` with the GELU activation applied. 262 | """ 263 | cdf = 0.5 * (1.0 + tf.tanh( 264 | (np.sqrt(2 / np.pi) * (x + 0.044715 * tf.pow(x, 3))))) 265 | return x * cdf 266 | 267 | 268 | def get_activation(activation_string): 269 | """Maps a string to a Python function, e.g., "relu" => `tf.nn.relu`. 270 | Args: 271 | activation_string: String name of the activation function. 272 | Returns: 273 | A Python function corresponding to the activation function. If 274 | `activation_string` is None, empty, or "linear", this will return None. 275 | If `activation_string` is not a string, it will return `activation_string`. 276 | Raises: 277 | ValueError: The `activation_string` does not correspond to a known 278 | activation. 279 | """ 280 | 281 | # We assume that anything that"s not a string is already an activation 282 | # function, so we just return it. 283 | if not isinstance(activation_string, six.string_types): 284 | return activation_string 285 | 286 | if not activation_string: 287 | return None 288 | 289 | act = activation_string.lower() 290 | if act == "linear": 291 | return None 292 | elif act == "relu": 293 | return tf.nn.relu 294 | elif act == "gelu": 295 | return gelu 296 | elif act == "tanh": 297 | return tf.tanh 298 | else: 299 | raise ValueError("Unsupported activation: %s" % act) 300 | 301 | 302 | def get_assignment_map_from_checkpoint(tvars, init_checkpoint): 303 | """Compute the union of the current variables and checkpoint variables.""" 304 | assignment_map = {} 305 | initialized_variable_names = {} 306 | 307 | name_to_variable = collections.OrderedDict() 308 | for var in tvars: 309 | name = var.name 310 | m = re.match("^(.*):\\d+$", name) 311 | if m is not None: 312 | name = m.group(1) 313 | name_to_variable[name] = var 314 | 315 | init_vars = tf.train.list_variables(init_checkpoint) 316 | 317 | assignment_map = collections.OrderedDict() 318 | for x in init_vars: 319 | (name, var) = (x[0], x[1]) 320 | if name not in name_to_variable: 321 | continue 322 | assignment_map[name] = name 323 | initialized_variable_names[name] = 1 324 | initialized_variable_names[name + ":0"] = 1 325 | 326 | return (assignment_map, initialized_variable_names) 327 | 328 | 329 | def dropout(input_tensor, dropout_prob): 330 | """Perform dropout. 331 | Args: 332 | input_tensor: float Tensor. 333 | dropout_prob: Python float. The probability of dropping out a value (NOT of 334 | *keeping* a dimension as in `tf.nn.dropout`). 335 | Returns: 336 | A version of `input_tensor` with dropout applied. 337 | """ 338 | if dropout_prob is None or dropout_prob == 0.0: 339 | return input_tensor 340 | 341 | output = tf.nn.dropout(input_tensor, 1.0 - dropout_prob) 342 | return output 343 | 344 | 345 | def layer_norm(input_tensor, name=None): 346 | """Run layer normalization on the last dimension of the tensor.""" 347 | return tf.contrib.layers.layer_norm( 348 | inputs=input_tensor, begin_norm_axis=-1, begin_params_axis=-1, scope=name) 349 | 350 | 351 | def layer_norm_and_dropout(input_tensor, dropout_prob, name=None): 352 | """Runs layer normalization followed by dropout.""" 353 | output_tensor = layer_norm(input_tensor, name) 354 | output_tensor = dropout(output_tensor, dropout_prob) 355 | return output_tensor 356 | 357 | 358 | def create_initializer(initializer_range=0.02): 359 | """Creates a `truncated_normal_initializer` with the given range.""" 360 | return tf.truncated_normal_initializer(stddev=initializer_range) 361 | 362 | 363 | def embedding_lookup(input_ids, 364 | vocab_size, 365 | embedding_size=128, 366 | initializer_range=0.02, 367 | word_embedding_name="word_embeddings", 368 | use_one_hot_embeddings=False): 369 | """Looks up words embeddings for id tensor. 370 | Args: 371 | input_ids: int32 Tensor of shape [batch_size, seq_length] containing word 372 | ids. 373 | vocab_size: int. Size of the embedding vocabulary. 374 | embedding_size: int. Width of the word embeddings. 375 | initializer_range: float. Embedding initialization range. 376 | word_embedding_name: string. Name of the embedding table. 377 | use_one_hot_embeddings: bool. If True, use one-hot method for word 378 | embeddings. If False, use `tf.gather()`. 379 | Returns: 380 | float Tensor of shape [batch_size, seq_length, embedding_size]. 381 | """ 382 | # This function assumes that the input is of shape [batch_size, seq_length, 383 | # num_inputs]. 384 | # 385 | # If the input is a 2D tensor of shape [batch_size, seq_length], we 386 | # reshape to [batch_size, seq_length, 1]. 387 | if input_ids.shape.ndims == 2: 388 | input_ids = tf.expand_dims(input_ids, axis=[-1]) 389 | 390 | embedding_table = tf.get_variable( 391 | name=word_embedding_name, 392 | shape=[vocab_size, embedding_size], 393 | initializer=create_initializer(initializer_range)) 394 | 395 | flat_input_ids = tf.reshape(input_ids, [-1]) 396 | if use_one_hot_embeddings: 397 | one_hot_input_ids = tf.one_hot(flat_input_ids, depth=vocab_size) 398 | output = tf.matmul(one_hot_input_ids, embedding_table) 399 | else: 400 | output = tf.gather(embedding_table, flat_input_ids) 401 | 402 | input_shape = get_shape_list(input_ids) 403 | 404 | output = tf.reshape(output, 405 | input_shape[0:-1] + [input_shape[-1] * embedding_size]) 406 | return (output, embedding_table) 407 | 408 | 409 | def embedding_postprocessor(input_tensor, 410 | use_token_type=False, 411 | token_type_ids=None, 412 | token_type_vocab_size=16, 413 | token_type_embedding_name="token_type_embeddings", 414 | use_position_embeddings=True, 415 | position_embedding_name="position_embeddings", 416 | initializer_range=0.02, 417 | max_position_embeddings=512, 418 | dropout_prob=0.1): 419 | """Performs various post-processing on a word embedding tensor. 420 | Args: 421 | input_tensor: float Tensor of shape [batch_size, seq_length, 422 | embedding_size]. 423 | use_token_type: bool. Whether to add embeddings for `token_type_ids`. 424 | token_type_ids: (optional) int32 Tensor of shape [batch_size, seq_length]. 425 | Must be specified if `use_token_type` is True. 426 | token_type_vocab_size: int. The vocabulary size of `token_type_ids`. 427 | token_type_embedding_name: string. The name of the embedding table variable 428 | for token type ids. 429 | use_position_embeddings: bool. Whether to add position embeddings for the 430 | position of each token in the sequence. 431 | position_embedding_name: string. The name of the embedding table variable 432 | for positional embeddings. 433 | initializer_range: float. Range of the weight initialization. 434 | max_position_embeddings: int. Maximum sequence length that might ever be 435 | used with this model. This can be longer than the sequence length of 436 | input_tensor, but cannot be shorter. 437 | dropout_prob: float. Dropout probability applied to the final output tensor. 438 | Returns: 439 | float tensor with same shape as `input_tensor`. 440 | Raises: 441 | ValueError: One of the tensor shapes or input values is invalid. 442 | """ 443 | input_shape = get_shape_list(input_tensor, expected_rank=3) 444 | batch_size = input_shape[0] 445 | seq_length = input_shape[1] 446 | width = input_shape[2] 447 | 448 | output = input_tensor 449 | 450 | if use_token_type: 451 | if token_type_ids is None: 452 | raise ValueError("`token_type_ids` must be specified if" 453 | "`use_token_type` is True.") 454 | token_type_table = tf.get_variable( 455 | name=token_type_embedding_name, 456 | shape=[token_type_vocab_size, width], 457 | initializer=create_initializer(initializer_range)) 458 | # This vocab will be small so we always do one-hot here, since it is always 459 | # faster for a small vocabulary. 460 | flat_token_type_ids = tf.reshape(token_type_ids, [-1]) 461 | one_hot_ids = tf.one_hot(flat_token_type_ids, depth=token_type_vocab_size) 462 | token_type_embeddings = tf.matmul(one_hot_ids, token_type_table) 463 | token_type_embeddings = tf.reshape(token_type_embeddings, 464 | [batch_size, seq_length, width]) 465 | output += token_type_embeddings 466 | 467 | if use_position_embeddings: 468 | assert_op = tf.assert_less_equal(seq_length, max_position_embeddings) 469 | with tf.control_dependencies([assert_op]): 470 | full_position_embeddings = tf.get_variable( 471 | name=position_embedding_name, 472 | shape=[max_position_embeddings, width], 473 | initializer=create_initializer(initializer_range)) 474 | # Since the position embedding table is a learned variable, we create it 475 | # using a (long) sequence length `max_position_embeddings`. The actual 476 | # sequence length might be shorter than this, for faster training of 477 | # tasks that do not have long sequences. 478 | # 479 | # So `full_position_embeddings` is effectively an embedding table 480 | # for position [0, 1, 2, ..., max_position_embeddings-1], and the current 481 | # sequence has positions [0, 1, 2, ... seq_length-1], so we can just 482 | # perform a slice. 483 | position_embeddings = tf.slice(full_position_embeddings, [0, 0], 484 | [seq_length, -1]) 485 | num_dims = len(output.shape.as_list()) 486 | 487 | # Only the last two dimensions are relevant (`seq_length` and `width`), so 488 | # we broadcast among the first dimensions, which is typically just 489 | # the batch size. 490 | position_broadcast_shape = [] 491 | for _ in range(num_dims - 2): 492 | position_broadcast_shape.append(1) 493 | position_broadcast_shape.extend([seq_length, width]) 494 | position_embeddings = tf.reshape(position_embeddings, 495 | position_broadcast_shape) 496 | output += position_embeddings 497 | 498 | output = layer_norm_and_dropout(output, dropout_prob) 499 | return output 500 | 501 | 502 | def create_attention_mask_from_input_mask(from_tensor, to_mask): 503 | """Create 3D attention mask from a 2D tensor mask. 504 | Args: 505 | from_tensor: 2D or 3D Tensor of shape [batch_size, from_seq_length, ...]. 506 | to_mask: int32 Tensor of shape [batch_size, to_seq_length]. 507 | Returns: 508 | float Tensor of shape [batch_size, from_seq_length, to_seq_length]. 509 | """ 510 | from_shape = get_shape_list(from_tensor, expected_rank=[2, 3]) 511 | batch_size = from_shape[0] 512 | from_seq_length = from_shape[1] 513 | 514 | to_shape = get_shape_list(to_mask, expected_rank=2) 515 | to_seq_length = to_shape[1] 516 | 517 | to_mask = tf.cast( 518 | tf.reshape(to_mask, [batch_size, 1, to_seq_length]), tf.float32) 519 | 520 | # We don't assume that `from_tensor` is a mask (although it could be). We 521 | # don't actually care if we attend *from* padding tokens (only *to* padding) 522 | # tokens so we create a tensor of all ones. 523 | # 524 | # `broadcast_ones` = [batch_size, from_seq_length, 1] 525 | broadcast_ones = tf.ones( 526 | shape=[batch_size, from_seq_length, 1], dtype=tf.float32) 527 | 528 | # Here we broadcast along two dimensions to create the mask. 529 | mask = broadcast_ones * to_mask 530 | 531 | return mask 532 | 533 | 534 | def attention_layer(from_tensor, 535 | to_tensor, 536 | attention_mask=None, 537 | num_attention_heads=1, 538 | size_per_head=512, 539 | query_act=None, 540 | key_act=None, 541 | value_act=None, 542 | attention_probs_dropout_prob=0.0, 543 | initializer_range=0.02, 544 | do_return_2d_tensor=False, 545 | batch_size=None, 546 | from_seq_length=None, 547 | to_seq_length=None): 548 | """Performs multi-headed attention from `from_tensor` to `to_tensor`. 549 | This is an implementation of multi-headed attention based on "Attention 550 | is all you Need". If `from_tensor` and `to_tensor` are the same, then 551 | this is self-attention. Each timestep in `from_tensor` attends to the 552 | corresponding sequence in `to_tensor`, and returns a fixed-with vector. 553 | This function first projects `from_tensor` into a "query" tensor and 554 | `to_tensor` into "key" and "value" tensors. These are (effectively) a list 555 | of tensors of length `num_attention_heads`, where each tensor is of shape 556 | [batch_size, seq_length, size_per_head]. 557 | Then, the query and key tensors are dot-producted and scaled. These are 558 | softmaxed to obtain attention probabilities. The value tensors are then 559 | interpolated by these probabilities, then concatenated back to a single 560 | tensor and returned. 561 | In practice, the multi-headed attention are done with transposes and 562 | reshapes rather than actual separate tensors. 563 | Args: 564 | from_tensor: float Tensor of shape [batch_size, from_seq_length, 565 | from_width]. 566 | to_tensor: float Tensor of shape [batch_size, to_seq_length, to_width]. 567 | attention_mask: (optional) int32 Tensor of shape [batch_size, 568 | from_seq_length, to_seq_length]. The values should be 1 or 0. The 569 | attention scores will effectively be set to -infinity for any positions in 570 | the mask that are 0, and will be unchanged for positions that are 1. 571 | num_attention_heads: int. Number of attention heads. 572 | size_per_head: int. Size of each attention head. 573 | query_act: (optional) Activation function for the query transform. 574 | key_act: (optional) Activation function for the key transform. 575 | value_act: (optional) Activation function for the value transform. 576 | attention_probs_dropout_prob: (optional) float. Dropout probability of the 577 | attention probabilities. 578 | initializer_range: float. Range of the weight initializer. 579 | do_return_2d_tensor: bool. If True, the output will be of shape [batch_size 580 | * from_seq_length, num_attention_heads * size_per_head]. If False, the 581 | output will be of shape [batch_size, from_seq_length, num_attention_heads 582 | * size_per_head]. 583 | batch_size: (Optional) int. If the input is 2D, this might be the batch size 584 | of the 3D version of the `from_tensor` and `to_tensor`. 585 | from_seq_length: (Optional) If the input is 2D, this might be the seq length 586 | of the 3D version of the `from_tensor`. 587 | to_seq_length: (Optional) If the input is 2D, this might be the seq length 588 | of the 3D version of the `to_tensor`. 589 | Returns: 590 | float Tensor of shape [batch_size, from_seq_length, 591 | num_attention_heads * size_per_head]. (If `do_return_2d_tensor` is 592 | true, this will be of shape [batch_size * from_seq_length, 593 | num_attention_heads * size_per_head]). 594 | Raises: 595 | ValueError: Any of the arguments or tensor shapes are invalid. 596 | """ 597 | 598 | def transpose_for_scores(input_tensor, batch_size, num_attention_heads, 599 | seq_length, width): 600 | output_tensor = tf.reshape( 601 | input_tensor, [batch_size, seq_length, num_attention_heads, width]) 602 | 603 | output_tensor = tf.transpose(output_tensor, [0, 2, 1, 3]) 604 | return output_tensor 605 | 606 | from_shape = get_shape_list(from_tensor, expected_rank=[2, 3]) 607 | to_shape = get_shape_list(to_tensor, expected_rank=[2, 3]) 608 | 609 | if len(from_shape) != len(to_shape): 610 | raise ValueError( 611 | "The rank of `from_tensor` must match the rank of `to_tensor`.") 612 | 613 | if len(from_shape) == 3: 614 | batch_size = from_shape[0] 615 | from_seq_length = from_shape[1] 616 | to_seq_length = to_shape[1] 617 | elif len(from_shape) == 2: 618 | if (batch_size is None or from_seq_length is None or to_seq_length is None): 619 | raise ValueError( 620 | "When passing in rank 2 tensors to attention_layer, the values " 621 | "for `batch_size`, `from_seq_length`, and `to_seq_length` " 622 | "must all be specified.") 623 | 624 | # Scalar dimensions referenced here: 625 | # B = batch size (number of sequences) 626 | # F = `from_tensor` sequence length 627 | # T = `to_tensor` sequence length 628 | # N = `num_attention_heads` 629 | # H = `size_per_head` 630 | 631 | from_tensor_2d = reshape_to_matrix(from_tensor) 632 | to_tensor_2d = reshape_to_matrix(to_tensor) 633 | 634 | # `query_layer` = [B*F, N*H] 635 | query_layer = tf.layers.dense( 636 | from_tensor_2d, 637 | num_attention_heads * size_per_head, 638 | activation=query_act, 639 | name="query", 640 | kernel_initializer=create_initializer(initializer_range)) 641 | 642 | # `key_layer` = [B*T, N*H] 643 | key_layer = tf.layers.dense( 644 | to_tensor_2d, 645 | num_attention_heads * size_per_head, 646 | activation=key_act, 647 | name="key", 648 | kernel_initializer=create_initializer(initializer_range)) 649 | 650 | # `value_layer` = [B*T, N*H] 651 | value_layer = tf.layers.dense( 652 | to_tensor_2d, 653 | num_attention_heads * size_per_head, 654 | activation=value_act, 655 | name="value", 656 | kernel_initializer=create_initializer(initializer_range)) 657 | 658 | # `query_layer` = [B, N, F, H] 659 | query_layer = transpose_for_scores(query_layer, batch_size, 660 | num_attention_heads, from_seq_length, 661 | size_per_head) 662 | 663 | # `key_layer` = [B, N, T, H] 664 | key_layer = transpose_for_scores(key_layer, batch_size, num_attention_heads, 665 | to_seq_length, size_per_head) 666 | 667 | # Take the dot product between "query" and "key" to get the raw 668 | # attention scores. 669 | # `attention_scores` = [B, N, F, T] 670 | attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True) 671 | attention_scores = tf.multiply(attention_scores, 672 | 1.0 / math.sqrt(float(size_per_head))) 673 | 674 | if attention_mask is not None: 675 | # `attention_mask` = [B, 1, F, T] 676 | attention_mask = tf.expand_dims(attention_mask, axis=[1]) 677 | 678 | # Since attention_mask is 1.0 for positions we want to attend and 0.0 for 679 | # masked positions, this operation will create a tensor which is 0.0 for 680 | # positions we want to attend and -10000.0 for masked positions. 681 | adder = (1.0 - tf.cast(attention_mask, tf.float32)) * -10000.0 682 | 683 | # Since we are adding it to the raw scores before the softmax, this is 684 | # effectively the same as removing these entirely. 685 | attention_scores += adder 686 | 687 | # Normalize the attention scores to probabilities. 688 | # `attention_probs` = [B, N, F, T] 689 | attention_probs = tf.nn.softmax(attention_scores) 690 | 691 | # This is actually dropping out entire tokens to attend to, which might 692 | # seem a bit unusual, but is taken from the original Transformer paper. 693 | attention_probs = dropout(attention_probs, attention_probs_dropout_prob) 694 | 695 | # `value_layer` = [B, T, N, H] 696 | value_layer = tf.reshape( 697 | value_layer, 698 | [batch_size, to_seq_length, num_attention_heads, size_per_head]) 699 | 700 | # `value_layer` = [B, N, T, H] 701 | value_layer = tf.transpose(value_layer, [0, 2, 1, 3]) 702 | 703 | # `context_layer` = [B, N, F, H] 704 | context_layer = tf.matmul(attention_probs, value_layer) 705 | 706 | # `context_layer` = [B, F, N, H] 707 | context_layer = tf.transpose(context_layer, [0, 2, 1, 3]) 708 | 709 | if do_return_2d_tensor: 710 | # `context_layer` = [B*F, N*H] 711 | context_layer = tf.reshape( 712 | context_layer, 713 | [batch_size * from_seq_length, num_attention_heads * size_per_head]) 714 | else: 715 | # `context_layer` = [B, F, N*H] 716 | context_layer = tf.reshape( 717 | context_layer, 718 | [batch_size, from_seq_length, num_attention_heads * size_per_head]) 719 | 720 | return context_layer 721 | 722 | 723 | def transformer_model(input_tensor, 724 | attention_mask=None, 725 | hidden_size=768, 726 | num_hidden_layers=12, 727 | num_attention_heads=12, 728 | intermediate_size=3072, 729 | intermediate_act_fn=gelu, 730 | hidden_dropout_prob=0.1, 731 | attention_probs_dropout_prob=0.1, 732 | initializer_range=0.02, 733 | do_return_all_layers=False): 734 | """Multi-headed, multi-layer Transformer from "Attention is All You Need". 735 | This is almost an exact implementation of the original Transformer encoder. 736 | See the original paper: 737 | https://arxiv.org/abs/1706.03762 738 | Also see: 739 | https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/models/transformer.py 740 | Args: 741 | input_tensor: float Tensor of shape [batch_size, seq_length, hidden_size]. 742 | attention_mask: (optional) int32 Tensor of shape [batch_size, seq_length, 743 | seq_length], with 1 for positions that can be attended to and 0 in 744 | positions that should not be. 745 | hidden_size: int. Hidden size of the Transformer. 746 | num_hidden_layers: int. Number of layers (blocks) in the Transformer. 747 | num_attention_heads: int. Number of attention heads in the Transformer. 748 | intermediate_size: int. The size of the "intermediate" (a.k.a., feed 749 | forward) layer. 750 | intermediate_act_fn: function. The non-linear activation function to apply 751 | to the output of the intermediate/feed-forward layer. 752 | hidden_dropout_prob: float. Dropout probability for the hidden layers. 753 | attention_probs_dropout_prob: float. Dropout probability of the attention 754 | probabilities. 755 | initializer_range: float. Range of the initializer (stddev of truncated 756 | normal). 757 | do_return_all_layers: Whether to also return all layers or just the final 758 | layer. 759 | Returns: 760 | float Tensor of shape [batch_size, seq_length, hidden_size], the final 761 | hidden layer of the Transformer. 762 | Raises: 763 | ValueError: A Tensor shape or parameter is invalid. 764 | """ 765 | if hidden_size % num_attention_heads != 0: 766 | raise ValueError( 767 | "The hidden size (%d) is not a multiple of the number of attention " 768 | "heads (%d)" % (hidden_size, num_attention_heads)) 769 | 770 | attention_head_size = int(hidden_size / num_attention_heads) 771 | input_shape = get_shape_list(input_tensor, expected_rank=3) 772 | batch_size = input_shape[0] 773 | seq_length = input_shape[1] 774 | input_width = input_shape[2] 775 | 776 | # The Transformer performs sum residuals on all layers so the input needs 777 | # to be the same as the hidden size. 778 | if input_width != hidden_size: 779 | raise ValueError("The width of the input tensor (%d) != hidden size (%d)" % 780 | (input_width, hidden_size)) 781 | 782 | # We keep the representation as a 2D tensor to avoid re-shaping it back and 783 | # forth from a 3D tensor to a 2D tensor. Re-shapes are normally free on 784 | # the GPU/CPU but may not be free on the TPU, so we want to minimize them to 785 | # help the optimizer. 786 | prev_output = reshape_to_matrix(input_tensor) 787 | 788 | all_layer_outputs = [] 789 | for layer_idx in range(num_hidden_layers): 790 | with tf.variable_scope("layer_%d" % layer_idx): 791 | layer_input = prev_output 792 | 793 | with tf.variable_scope("attention"): 794 | attention_heads = [] 795 | with tf.variable_scope("self"): 796 | attention_head = attention_layer( 797 | from_tensor=layer_input, 798 | to_tensor=layer_input, 799 | attention_mask=attention_mask, 800 | num_attention_heads=num_attention_heads, 801 | size_per_head=attention_head_size, 802 | attention_probs_dropout_prob=attention_probs_dropout_prob, 803 | initializer_range=initializer_range, 804 | do_return_2d_tensor=True, 805 | batch_size=batch_size, 806 | from_seq_length=seq_length, 807 | to_seq_length=seq_length) 808 | attention_heads.append(attention_head) 809 | 810 | attention_output = None 811 | if len(attention_heads) == 1: 812 | attention_output = attention_heads[0] 813 | else: 814 | # In the case where we have other sequences, we just concatenate 815 | # them to the self-attention head before the projection. 816 | attention_output = tf.concat(attention_heads, axis=-1) 817 | 818 | # Run a linear projection of `hidden_size` then add a residual 819 | # with `layer_input`. 820 | with tf.variable_scope("output"): 821 | attention_output = tf.layers.dense( 822 | attention_output, 823 | hidden_size, 824 | kernel_initializer=create_initializer(initializer_range)) 825 | attention_output = dropout(attention_output, hidden_dropout_prob) 826 | attention_output = layer_norm(attention_output + layer_input) 827 | 828 | # The activation is only applied to the "intermediate" hidden layer. 829 | with tf.variable_scope("intermediate"): 830 | intermediate_output = tf.layers.dense( 831 | attention_output, 832 | intermediate_size, 833 | activation=intermediate_act_fn, 834 | kernel_initializer=create_initializer(initializer_range)) 835 | 836 | # Down-project back to `hidden_size` then add the residual. 837 | with tf.variable_scope("output"): 838 | layer_output = tf.layers.dense( 839 | intermediate_output, 840 | hidden_size, 841 | kernel_initializer=create_initializer(initializer_range)) 842 | layer_output = dropout(layer_output, hidden_dropout_prob) 843 | layer_output = layer_norm(layer_output + attention_output) 844 | prev_output = layer_output 845 | all_layer_outputs.append(layer_output) 846 | 847 | if do_return_all_layers: 848 | final_outputs = [] 849 | for layer_output in all_layer_outputs: 850 | final_output = reshape_from_matrix(layer_output, input_shape) 851 | final_outputs.append(final_output) 852 | return final_outputs 853 | else: 854 | final_output = reshape_from_matrix(prev_output, input_shape) 855 | return final_output 856 | 857 | 858 | def get_shape_list(tensor, expected_rank=None, name=None): 859 | """Returns a list of the shape of tensor, preferring static dimensions. 860 | Args: 861 | tensor: A tf.Tensor object to find the shape of. 862 | expected_rank: (optional) int. The expected rank of `tensor`. If this is 863 | specified and the `tensor` has a different rank, and exception will be 864 | thrown. 865 | name: Optional name of the tensor for the error message. 866 | Returns: 867 | A list of dimensions of the shape of tensor. All static dimensions will 868 | be returned as python integers, and dynamic dimensions will be returned 869 | as tf.Tensor scalars. 870 | """ 871 | if name is None: 872 | name = tensor.name 873 | 874 | if expected_rank is not None: 875 | assert_rank(tensor, expected_rank, name) 876 | 877 | shape = tensor.shape.as_list() 878 | 879 | non_static_indexes = [] 880 | for (index, dim) in enumerate(shape): 881 | if dim is None: 882 | non_static_indexes.append(index) 883 | 884 | if not non_static_indexes: 885 | return shape 886 | 887 | dyn_shape = tf.shape(tensor) 888 | for index in non_static_indexes: 889 | shape[index] = dyn_shape[index] 890 | return shape 891 | 892 | 893 | def reshape_to_matrix(input_tensor): 894 | """Reshapes a >= rank 2 tensor to a rank 2 tensor (i.e., a matrix).""" 895 | ndims = input_tensor.shape.ndims 896 | if ndims < 2: 897 | raise ValueError("Input tensor must have at least rank 2. Shape = %s" % 898 | (input_tensor.shape)) 899 | if ndims == 2: 900 | return input_tensor 901 | 902 | width = input_tensor.shape[-1] 903 | output_tensor = tf.reshape(input_tensor, [-1, width]) 904 | return output_tensor 905 | 906 | 907 | def reshape_from_matrix(output_tensor, orig_shape_list): 908 | """Reshapes a rank 2 tensor back to its original rank >= 2 tensor.""" 909 | if len(orig_shape_list) == 2: 910 | return output_tensor 911 | 912 | output_shape = get_shape_list(output_tensor) 913 | 914 | orig_dims = orig_shape_list[0:-1] 915 | width = output_shape[-1] 916 | 917 | return tf.reshape(output_tensor, orig_dims + [width]) 918 | 919 | 920 | def assert_rank(tensor, expected_rank, name=None): 921 | """Raises an exception if the tensor rank is not of the expected rank. 922 | Args: 923 | tensor: A tf.Tensor to check the rank of. 924 | expected_rank: Python integer or list of integers, expected rank. 925 | name: Optional name of the tensor for the error message. 926 | Raises: 927 | ValueError: If the expected shape doesn't match the actual shape. 928 | """ 929 | if name is None: 930 | name = tensor.name 931 | 932 | expected_rank_dict = {} 933 | if isinstance(expected_rank, six.integer_types): 934 | expected_rank_dict[expected_rank] = True 935 | else: 936 | for x in expected_rank: 937 | expected_rank_dict[x] = True 938 | 939 | actual_rank = tensor.shape.ndims 940 | if actual_rank not in expected_rank_dict: 941 | scope_name = tf.get_variable_scope().name 942 | raise ValueError( 943 | "For the tensor `%s` in scope `%s`, the actual rank " 944 | "`%d` (shape = %s) is not equal to the expected rank `%s`" % 945 | (name, scope_name, actual_rank, str(tensor.shape), str(expected_rank))) 946 | --------------------------------------------------------------------------------