added_training_stuff

generator/ctrl/LICENSE.txt

@@ -0,0 +1,12 @@
+Copyright (c) 2019, Salesforce.com, Inc.
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+
+* Neither the name of Salesforce.com nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

generator/ctrl/control_codes.txt

@@ -0,0 +1,55 @@
+0.000296793    Pregnancy
+0.000127197    Christianity
+0.003084531    Explain
+0.000180196    Fitness
+6.88985E-05    Saving
+0.000217295    Ask
+8.47981E-05    Ass
+0.000143097    Joke
+0.000196096    Questions
+0.000127197    Thoughts
+0.000169596    Retail
+0.000270294    Feminism
+0.000111298    Writing
+0.000402791    Atheism
+1.05998E-06    Netflix
+0.000365692    Computing
+0.000132497    Opinion
+0.000169596    Alone
+0.000323293    Funny
+0.000249094    Gaming
+0.000402791    Human
+0.000132497    India
+2.11995E-08    Joker
+0.000201395    Diet
+0.000238495    Legal
+6.35986E-06    Norman
+3.60392E-07    Tip
+0.000302093    Weight
+0.000132497    Movies
+0.000111298    Running
+7.41983E-05    Science
+0.00135147    Horror
+0.000291493    Confession
+0.000190796    Finance
+0.000413391    Politics
+7.41983E-05    Scary
+0.000206695    Support
+6.35986E-05    Technologies
+0.000243795    Teenage
+0.000217295    Event
+0.000206695    Learned
+0.000121897    Notion
+0.0847981    Wikipedia
+0.095927851    Books
+0.001176574    Extract
+0.000127197    Confessions
+0.000227895    Conspiracy
+0.365691808    Links
+0.000423991    Narcissus
+0.000280894    Relationship
+0.000922179    Relationships
+0.153696557    Reviews
+0.043877717    News
+0.129847091    Translation
+0.111297507    multilingual

generator/ctrl/ctrl_generator.py

@@ -219,7 +219,7 @@ class CTRLGenerator():
         forbidden_tokens = ['<unk>', 'Sco@@', "&amp@@", "1]@@", "2]@@", "3]@@", "4]@@", "https://www.@@", "[@@", ":@@",
                             "Edit", "&@@", "2:","1:", ":", "Edit@@", "EDI@@", "EDIT@@", "edit", "TL@@", "tl@@", ";@@",
                             '**', "http://@@", "Redd@@", "UP@@", "mom", "Up@@", "Me:", "Update", "mom@@", "Part",
-                            "http://www.@@"]
+                            "http://www.@@", "edit@@", "*@@"]
 
         if num_new_lines > self.max_new_lines:
             forbidden_tokens.append("\n")

generator/ctrl/estimator.patch

@@ -0,0 +1,24 @@
+47c47
+< 
+---
+> import tensorflow as tf 
+228c228
+< def _create_keras_model_fn(keras_model, custom_objects=None):
+---
+> def _create_keras_model_fn(keras_model, params=None, custom_objects=None):
+239c239
+<   def model_fn(features, labels, mode):
+---
+>   def model_fn(features, labels, mode, params=None):
+448c448
+<   if keras_model._is_graph_network:
+---
+>   if False:
+462,464c462,464
+<   estimator = estimator_lib.Estimator(keras_model_fn,
+<                                       config=config,
+<                                       warm_start_from=warm_start_path)
+---
+>   estimator = tf.contrib.tpu.TPUEstimator(keras_model_fn, use_tpu=False, train_batch_size=4, eval_batch_size=4,
+>                                           config=config,
+>                                           warm_start_from=warm_start_path)

generator/ctrl/generation.py

@@ -0,0 +1,288 @@
+from __future__ import division
+from __future__ import print_function
+import tensorflow as tf
+import os
+import numpy as np
+tf.enable_eager_execution()
+import transformer
+import argparse
+import pdb
+import sys
+import re
+from collections import Counter
+from tensorflow.python import debug as tf_debug
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import embedding_ops
+import fastBPE
+import platform
+
+use_py3 = platform.python_version()[0] == '3'
+
+parser = argparse.ArgumentParser(description='TensorFlow code for generating from CTRL')
+parser.add_argument('--model_dir', type=str, required=True,
+                                        help='location of model checkpoint')
+parser.add_argument('--seed', type=int, default=1337,
+                                        help='random seed for TensorFlow, numpy and PythonHash')
+parser.add_argument('--generate_num', type=int, default=256,
+                                        help='number of tokens to generate')
+parser.add_argument('--temperature', type=float, default=0,
+                                        help='temperature for sampling distribution; 0 means greedy')
+parser.add_argument('--nucleus', type=float, default=0.,
+                                        help='cumulative probability cutoff for nucleus sampling; 0 means no nucleus sampling')
+parser.add_argument('--topk', type=int, default=0,
+                                        help='topk value for sampling from the softmax distribution ; 0 means no topk preferred')
+parser.add_argument('--penalty', type=float, default=1.2,
+                                        help='repetition penalty for greedy sampling')
+parser.add_argument('--print_once', action='store_true',
+                                        help='the completion is printed only at the end; not every word')
+parser.add_argument('--topn', type=int, default=0,
+                                        help='print top-n candidates during generations; defaults to 0 which is no printing')
+
+args = parser.parse_args()
+tf.random.set_random_seed(args.seed)
+os.environ['PYTHONHASHSEED'] = str(args.seed)
+np.random.seed(args.seed)
+
+# load the vocabulary from file
+vocab = open('vocab').read().decode(encoding='utf-8').split('\n') if not use_py3 else open('vocab', encoding='utf-8').read().split('\n')
+vocab = list(map(lambda x: x.split(' ')[0], vocab)) + ['<unk>'] + ['\n']
+print ('{} unique words'.format(len(vocab)))
+
+# length of the vocabulary
+vocab_size = len(vocab)
+
+# define the numericalization map
+# idx2word maps the numericalized ID to the word
+# word2idx maps the word to the numericalized ID
+word2idx = {u:i for i, u in enumerate(vocab)}
+idx2word = np.array(vocab)
+
+
+
+# sequence length to use for the transformer
+# the model is trained with a seq_length of 512
+# so, any value <= 512 should work
+seq_length = min(args.generate_num, 256)
+
+
+
+
+# the dimension of the transformer
+embedding_dim = 1280
+
+
+# Now, we begin defining the model
+# we defer the transformer definition to transformer.py
+# here, we only define the tied softmax layer
+# this layer ties the softmax weights to the input embeddings
+class TiedEmbeddingSoftmax(tf.keras.layers.Layer):
+
+  def __init__(self, vocab_size=vocab_size, embedding_size=embedding_dim, **kwargs):
+    super(TiedEmbeddingSoftmax, self).__init__()
+    self.w = self.add_weight(name='w', shape=(vocab_size, embedding_size),
+                             initializer='random_normal',
+                             trainable=True)
+    self.b = self.add_weight(name='b', shape=(vocab_size,),
+                             initializer='zeros',
+                             trainable=True)
+
+  def call(self, inputs, embed=True):
+    if embed:
+      dtype = tf.keras.backend.dtype(inputs)
+      if dtype != 'int32' and dtype != 'int64':
+        inputs = math_ops.cast(inputs, 'int32')
+      return embedding_ops.embedding_lookup(self.w, inputs)
+    else:
+      return tf.tensordot(inputs, tf.transpose(self.w), 1) + self.b
+
+# input for the keras model
+tokens = tf.keras.layers.Input(shape=(seq_length,), dtype='int32')
+
+# instantiates a tied softmax class
+tied_embedding_softmax = TiedEmbeddingSoftmax()
+
+# embedded tokens, before passing it to the transformer
+embedded = tied_embedding_softmax(tokens, embed=True)
+
+# the activations after passing it from the transformer
+# for some odd reason, TPUs don't play well with specifying the arguments of the Encoder() function
+# so you have to leave them at their defaults
+transformed = transformer.Encoder()(embedded, training=False)
+
+
+# pass the activations from our tiedsoftmax class
+# this time with embed=False denoting that we are doing the softmax operation
+# and not a lookup
+logits = tied_embedding_softmax(transformed, embed=False)
+
+
+# finally, define the Keras model with inputs as tokens and outputs as the logits we just computed
+model = tf.keras.Model(inputs=tokens, outputs=logits)
+
+
+# the loss function is a simple categorical crossentropy between the logits and the labels
+def loss(labels, logits):
+    return tf.keras.losses.sparse_categorical_crossentropy(labels, logits, from_logits=True)
+
+# the optimizer is not used since this code only supports inference
+# however, to compile the model, we still define it
+optimizer = tf.contrib.tpu.CrossShardOptimizer(
+    tf.contrib.estimator.clip_gradients_by_norm(
+        tf.train.AdagradOptimizer(learning_rate=1e-2), 0.25)
+    )        
+
+# compile the model with the optimizer and loss            
+model.compile(optimizer=optimizer, loss=loss)
+print(model.summary())
+
+
+# IMPORTANT
+# this is where the saved model is presented to the code
+# the model directory should have the model checkpoint and
+# a checkpoint file
+run_config = tf.contrib.tpu.RunConfig(
+        model_dir=args.model_dir)
+
+
+# this converts the Keras model to a TensorFlow estimator
+# this step is critical
+# remember to patch the TF 1.14 file before running the code, else you're going to see errors here
+estimator_model = tf.keras.estimator.model_to_estimator(keras_model=model, config=run_config)
+
+# we now create a serving function from this estimator
+# this enables us to load the model once and easily query it multiple times
+def serving_input_fn():
+    inputs = {'input_1': tf.placeholder(tf.int32, [1,seq_length])}
+    return tf.estimator.export.ServingInputReceiver(inputs, inputs)
+predict_fn = tf.contrib.predictor.from_estimator(estimator_model, serving_input_fn)
+
+# almost there, we now take the user prompt and tokenize with BPE
+# load BPE codes
+bpe = fastBPE.fastBPE('codes', 'vocab')
+
+temperature = args.temperature
+nucleusprob = args.nucleus
+penalty = args.penalty
+topk = args.topk
+
+while True:
+    prompt = raw_input('ENTER PROMPT: ') if not use_py3 else input('ENTER PROMPT: ')
+
+    # tokenize provided prompt
+    split_prompt = bpe.apply([prompt])[0].split()
+    text = [word2idx[i] for i in split_prompt]
+
+    # pad with 0s and create a mini-batch of 2 (arbitrary, for ease of code)
+    padded_text = text + [0] * (args.generate_num - len(text))
+    tokens_generated = np.tile(padded_text, (1,1))
+    try:
+        for token in range(len(text)-1, args.generate_num-1):
+          # get the logits from the prediction function
+          # the logic here is a bit convoluted because we are allowing generation past 512 tokens
+          # this is done by sliding the window over (past 512 tokens) and continuing prediction
+          # I'm sure this can be simplified (TODO)
+          if token <= seq_length:
+            prompt_logits = predict_fn({'input_1':tokens_generated[:, :seq_length]})['tied_embedding_softmax'].squeeze() / (temperature if temperature>0 else 1.)
+            _token = token if token < seq_length else -1
+          else:
+            _token = -1
+            end = token + 1
+            start = token - seq_length + 2
+            prompt_logits = predict_fn({'input_1':np.hstack((tokens_generated[:,0:1], tokens_generated[:,start:end]))})['tied_embedding_softmax'].squeeze() / (temperature if temperature>0 else 1.)
+
+
+          # if penalty (for repetition) is non-zero,
+          # discount the logits from already generated tokens
+          if penalty>0:
+              penalized_so_far = set()
+              for _ in range(token+1):
+                 generated_token = tokens_generated[0][_]
+                 # don't penalize newlines
+                 # you could also choose not to penalize frequent words
+                 # (which incidentally are sorted in the vocab file)
+                 # but I don't do that
+                 # if it prints too many new lines instead of continuing generating text,
+                 # you might want to comment this out
+                 if idx2word[generated_token] == '\n':
+                     continue
+                 if generated_token in penalized_so_far:
+                     continue
+                 penalized_so_far.add(generated_token)
+                 prompt_logits[_token][generated_token] /= penalty
+
+          # disallow some tokens
+          prompt_logits[_token][word2idx['<unk>']] = -1e8
+
+          # sometimes, when generating from reddit,
+          # it tries to generate the Score (reddit Karma) immediately after generating the Title:
+          # to disallow this, we can just prevent it from generating Score
+          prompt_logits[_token][word2idx['Sco@@']] = -1e8
+
+
+          # compute probabilities from logits
+          prompt_probs = np.exp(prompt_logits[_token])
+          prompt_probs = prompt_probs / sum(prompt_probs)
+          pruned_list = np.argsort(prompt_probs)[::-1]
+          # if you are using nucleus prob, then compute the nucleus probability size
+          if nucleusprob > 0.:
+            minimum_topk = 1
+            nucleus = max(np.where(np.cumsum(np.sort(prompt_probs)[::-1])>nucleusprob)[0][0], minimum_topk)
+          elif topk > 0:
+            # we are over-loading notation here
+            # if you choose to specify a topk instead of a nucleus,
+            # we will hardcode the nucleus to be just that
+            nucleus = topk
+          else:
+            # if you specify neither nucleus or topk,
+            # then we will use the whole list
+            nucleus = len(pruned_list)
+            
+          pruned_list = pruned_list[:nucleus]  
+          # if you want to disallow more complex tokens, you can do so here
+          # for instance, if you want to disallow anything with the phrase `http`,
+          # you can delete theme from the pruned_list
+          # you can comment this out, I'm keeping it in for demonstration purpose
+          tokens_to_disallow = []
+          for _ in range(len(pruned_list)):
+              if 'http' in idx2word[pruned_list[_]]:
+                  tokens_to_disallow.append(_)
+          pruned_list = np.delete(pruned_list, tokens_to_disallow)
+
+          if args.topn > 0 :
+            print('TOPN :: top-n alternatives:', [idx2word[_] for _ in pruned_list[:args.topn]])
+
+          # if temperature is 0
+          # just pick the first (most probable) token
+          if temperature==0:
+              idx = pruned_list[0]
+          else:
+            # else,
+            # sample from the pruned_list with the logits
+            chosen_idx = int(tf.random.categorical(np.expand_dims(prompt_logits[_token][pruned_list],0), num_samples=1).numpy())
+            idx = pruned_list[chosen_idx]
+
+          if args.topn > 0 :
+            print('TOPN :: chosen word:', idx2word[idx])
+
+          # assign the token for generation
+          tokens_generated[0][token+1] = idx
+
+          # clear screen if you want to
+          # os.system("clear")
+
+
+          tokens_generated_so_far = ' '.join([idx2word[c] for c in tokens_generated[0].squeeze()[:token+2]])
+          tokens_generated_so_far = re.sub('(@@ )', '', string=tokens_generated_so_far)
+          tokens_generated_so_far = re.sub('(@@ ?$)', '', string=tokens_generated_so_far)              
+
+          if not args.print_once:
+            print('---------------------------------------')
+            print(tokens_generated_so_far)
+            print()
+        print('---------------------------------------')            
+        print(tokens_generated_so_far)
+        print()
+
+    except KeyboardInterrupt: #Exception as e:
+        print('Continuing')
+            

generator/ctrl/pytorch_generation.py

@@ -0,0 +1,283 @@
+from __future__ import print_function
+import torch
+import os
+import tqdm
+import pdb
+import numpy as np
+import platform
+import hashlib
+import pytorch_transformer
+import re
+import argparse
+import tensorflow as tf
+import fastBPE
+from tensorflow.python import pywrap_tensorflow
+
+use_py3 = platform.python_version()[0] == '3'
+
+parser = argparse.ArgumentParser(description='TensorFlow code for generating from CTRL')
+parser.add_argument('--model_path', type=str, required=True,
+                                        help='location of model *data* checkpoint; this is NOT the directory but rather the model checkpoint')
+parser.add_argument('--seed', type=int, default=1337,
+                                        help='random seed for TensorFlow, numpy and PythonHash')
+parser.add_argument('--generate_num', type=int, default=256,
+                                        help='number of tokens to generate')
+parser.add_argument('--temperature', type=float, default=0,
+                                        help='temperature for sampling distribution; 0 means greedy')
+parser.add_argument('--nucleus', type=float, default=0.,
+                                        help='cumulative probability cutoff for nucleus sampling; 0 means no nucleus sampling')
+parser.add_argument('--topk', type=int, default=0,
+                                        help='topk value for sampling from the softmax distribution ; 0 means no topk preferred')
+parser.add_argument('--penalty', type=float, default=1.2,
+                                        help='repetition penalty for greedy sampling')
+parser.add_argument('--print_once', action='store_true',
+                                        help='the completion is printed only at the end; not every word')
+parser.add_argument('--topn', type=int, default=0,
+                                        help='print top-n candidates during generations; defaults to 0 which is no printing')
+
+args = parser.parse_args()
+torch.manual_seed(args.seed)
+torch.cuda.manual_seed_all(args.seed)
+os.environ['PYTHONHASHSEED'] = str(args.seed)
+np.random.seed(args.seed)
+
+# load the vocabulary from file
+vocab = open('vocab').read().decode(encoding='utf-8').split('\n') if not use_py3 else open('vocab', encoding='utf-8').read().split('\n')
+vocab = list(map(lambda x: x.split(' ')[0], vocab)) + ['<unk>'] + ['\n']
+print ('{} unique words'.format(len(vocab)))
+
+# length of the vocabulary
+vocab_size = len(vocab)
+
+# define the numericalization map
+# idx2word maps the numericalized ID to the word
+# word2idx maps the word to the numericalized ID
+word2idx = {u:i for i, u in enumerate(vocab)}
+idx2word = np.array(vocab)
+embedding_dim = 1280
+
+class TiedEmbeddingSoftmax(torch.nn.Module):
+
+  def __init__(self, vocab_size=vocab_size, embedding_size=embedding_dim, **kwargs):
+    super(TiedEmbeddingSoftmax, self).__init__()
+    self.w = torch.nn.Parameter(torch.zeros(vocab_size, embedding_size))
+    self.b = torch.nn.Parameter(torch.zeros(vocab_size))
+
+  def forward(self, inputs, embed=True):
+    if embed:
+      return torch.nn.functional.embedding(inputs, self.w)
+    else:
+      return torch.tensordot(inputs, self.w.t(), 1) + self.b
+
+
+test_softmax = TiedEmbeddingSoftmax()
+test_encoder = pytorch_transformer.Encoder()
+
+def predict_fn(inputs):
+  with torch.no_grad():
+    embedded = torch.tensor(inputs['input_1']).cuda()
+    embedded = test_softmax(embedded, embed=True)
+    embedded = test_encoder(embedded)
+    embedded = test_softmax(embedded, embed=False)
+  return embedded
+
+
+
+
+
+
+bpe = fastBPE.fastBPE('codes', 'vocab')
+seq_length = min(args.generate_num, 256)
+pytorch_model_hash = hashlib.md5(args.model_path.encode('utf-8')).hexdigest()
+temperature = args.temperature
+nucleusprob = args.nucleus
+penalty = args.penalty
+topk = args.topk
+
+# try to load the model from a (cached) PyTorch checkpoint
+# if one is not available, then create it by converting the weights
+if os.path.exists(pytorch_model_hash):
+  print('Found PyTorch checkpoint @', pytorch_model_hash)
+  print('Loading instead of converting from TensorFlow')
+  checkpoint = torch.load(pytorch_model_hash)
+  test_softmax.load_state_dict(checkpoint['softmax'])
+  test_encoder.load_state_dict(checkpoint['encoder'])
+
+  test_softmax.to('cuda')
+  test_encoder.to('cuda')
+
+else:
+  print('Could not find PyTorch checkpoint')
+  print('Converting weights and will store the PyTorch checkpoint as ', pytorch_model_hash)
+  chkpt_for_reader = '.'.join(args.model_path.split('.')[:-1])
+  reader = pywrap_tensorflow.NewCheckpointReader(chkpt_for_reader)
+  test_softmax.w = torch.nn.Parameter(torch.tensor(reader.get_tensor('w')).to('cuda'))
+  test_softmax.b = torch.nn.Parameter(torch.tensor(reader.get_tensor('b')).to('cuda'))
+
+  list_of_variables = list(filter(lambda x: 'Adagrad' not in x, reader.get_variable_to_shape_map().keys()))
+
+  str2parameter = lambda x: torch.nn.Parameter(torch.tensor(reader.get_tensor(x)).t().to('cuda'))
+
+  test_encoder.layernorm.weight = str2parameter('encoder/layer_normalization_96/gamma')
+  test_encoder.layernorm.bias = str2parameter('encoder/layer_normalization_96/beta')
+  for i in tqdm.tqdm(range(48)):
+    if i==0:
+      layer_variables = sorted(filter(lambda x: 'layer/' in x, list_of_variables))
+    else:
+      layer_variables = sorted(filter(lambda x: 'layer_'+str(i)+'/' in x, list_of_variables))
+
+    current_layer = getattr(test_encoder, 'layer'+str(i))
+
+    current_layer.layernorm1.bias = str2parameter(layer_variables[0])
+    current_layer.layernorm1.weight = str2parameter(layer_variables[1])
+
+    current_layer.layernorm2.bias = str2parameter(layer_variables[2])
+    current_layer.layernorm2.weight = str2parameter(layer_variables[3])
+
+
+    current_layer.multi_head_attention.Wq.bias = str2parameter(layer_variables[4])
+    current_layer.multi_head_attention.Wq.weight = str2parameter(layer_variables[5])
+    current_layer.multi_head_attention.Wk.bias = str2parameter(layer_variables[6])
+    current_layer.multi_head_attention.Wk.weight = str2parameter(layer_variables[7])
+    current_layer.multi_head_attention.Wv.bias = str2parameter(layer_variables[8])
+    current_layer.multi_head_attention.Wv.weight = str2parameter(layer_variables[9])
+    current_layer.multi_head_attention.dense.bias = str2parameter(layer_variables[10])
+    current_layer.multi_head_attention.dense.weight = str2parameter(layer_variables[11])
+    current_layer.ffn[0].bias = str2parameter(layer_variables[12])
+    current_layer.ffn[0].weight = str2parameter(layer_variables[13])
+    current_layer.ffn[2].bias = str2parameter(layer_variables[14])
+    current_layer.ffn[2].weight = str2parameter(layer_variables[15])
+  torch.save({
+    'softmax': test_softmax.state_dict(),
+    'encoder': test_encoder.state_dict(),
+  }, pytorch_model_hash)
+
+test_softmax.eval()
+test_encoder.eval()
+
+
+
+
+while True:
+  prompt = raw_input('ENTER PROMPT: ') if not use_py3 else input('ENTER PROMPT: ')
+
+  # tokenize provided prompt
+  split_prompt = bpe.apply([prompt])[0].split()
+  text = [word2idx[i] for i in split_prompt]
+
+
+  # pad with 0s and create a mini-batch of 2 (arbitrary, for ease of code)
+  padded_text = text + [0] * (args.generate_num - len(text))
+  tokens_generated = np.tile(padded_text, (1,1))
+  try:
+    for token in range(len(text)-1, args.generate_num-1):
+      # get the logits from the prediction function
+      # the logic here is a bit convoluted because we are allowing generation past 512 tokens
+      # this is done by sliding the window over (past 512 tokens) and continuing prediction
+      # I'm sure this can be simplified (TODO)
+      if token <= seq_length:
+        prompt_logits = predict_fn({'input_1':tokens_generated[:, :seq_length]}).squeeze() / (temperature if temperature>0 else 1.)
+        _token = token if token < seq_length else -1
+      else:
+        _token = -1
+        end = token + 1
+        start = token - seq_length + 2
+        prompt_logits = predict_fn({'input_1':np.hstack((tokens_generated[:,0:1], tokens_generated[:,start:end]))}).squeeze() / (temperature if temperature>0 else 1.)
+
+      prompt_logits = prompt_logits.cpu().detach().numpy()
+
+      # if penalty (for repetition) is non-zero,
+      # discount the logits from already generated tokens
+      if penalty>0:
+          penalized_so_far = set()
+          for _ in range(token+1):
+             generated_token = tokens_generated[0][_]
+             # don't penalize newlines
+             # you could also choose not to penalize frequent words
+             # (which incidentally are sorted in the vocab file)
+             # but I don't do that
+             # if it prints too many new lines instead of continuing generating text,
+             # you might want to comment this out
+             #if idx2word[generated_token] == '\n':
+             #    continue
+             if generated_token in penalized_so_far:
+                 continue
+             penalized_so_far.add(generated_token)
+             prompt_logits[_token][generated_token] /= penalty
+
+      # disallow some tokens
+      prompt_logits[_token][word2idx['<unk>']] = -1e8
+
+      # sometimes, when generating from reddit,
+      # it tries to generate the Score (reddit Karma) immediately after generating the Title:
+      # to disallow this, we can just prevent it from generating Score
+      prompt_logits[_token][word2idx['Sco@@']] = -1e8
+
+
+      # compute probabilities from logits
+      prompt_probs = np.exp(prompt_logits[_token])
+      prompt_probs = prompt_probs / sum(prompt_probs)
+      pruned_list = np.argsort(prompt_probs)[::-1]
+      # if you are using nucleus prob, then compute the nucleus probability size
+      if nucleusprob > 0.:
+        minimum_topk = 1
+        nucleus = max(np.where(np.cumsum(np.sort(prompt_probs)[::-1])>nucleusprob)[0][0], minimum_topk)
+      elif topk > 0:
+        # we are over-loading notation here
+        # if you choose to specify a topk instead of a nucleus,
+        # we will hardcode the nucleus to be just that
+        nucleus = topk
+      else:
+        # if you specify neither nucleus or topk,
+        # then we will use the whole list
+        nucleus = len(pruned_list)
+
+      pruned_list = pruned_list[:nucleus]  
+      # if you want to disallow more complex tokens, you can do so here
+      # for instance, if you want to disallow anything with the phrase `http`,
+      # you can delete theme from the pruned_list
+      # you can comment this out, I'm keeping it in for demonstration purpose
+      tokens_to_disallow = []
+      for _ in range(len(pruned_list)):
+          if 'http' in idx2word[pruned_list[_]]:
+              tokens_to_disallow.append(_)
+      pruned_list = np.delete(pruned_list, tokens_to_disallow)
+
+      if args.topn > 0 :
+        print('TOPN :: top-n alternatives:', [idx2word[_] for _ in pruned_list[:args.topn]])
+
+      # if temperature is 0
+      # just pick the first (most probable) token
+      if temperature==0:
+          idx = pruned_list[0]
+      else:
+        # else,
+        # sample from the pruned_list with the logits
+        chosen_idx = torch.distributions.categorical.Categorical(torch.tensor(np.expand_dims(prompt_logits[_token][pruned_list],0))).sample().numpy()[0]
+        idx = pruned_list[chosen_idx]
+
+      if args.topn > 0 :
+        print('TOPN :: chosen word:', idx2word[idx])
+
+      # assign the token for generation
+      tokens_generated[0][token+1] = idx
+
+      # clear screen if you want to
+      # os.system("clear")
+
+
+      tokens_generated_so_far = ' '.join([idx2word[c] for c in tokens_generated[0].squeeze()[:token+2]])
+      tokens_generated_so_far = re.sub('(@@ )', '', string=tokens_generated_so_far)
+      tokens_generated_so_far = re.sub('(@@ ?$)', '', string=tokens_generated_so_far)              
+
+      if not args.print_once:
+        print('---------------------------------------')
+        print(tokens_generated_so_far)
+        print()
+    print('---------------------------------------')            
+    print(tokens_generated_so_far)
+    print()
+    
+  except KeyboardInterrupt: #Exception as e:
+    print('Continuing')
+

generator/ctrl/pytorch_transformer.py

@@ -0,0 +1,143 @@
+from __future__ import print_function
+import torch
+import os
+import tqdm
+import pdb
+import numpy as np
+import platform
+import re
+import argparse
+import tensorflow as tf
+from tensorflow.python import pywrap_tensorflow
+
+
+def angle_defn(pos, i, d_model_size):
+  angle_rates = 1 / np.power(10000, (2 * (i//2)) / np.float32(d_model_size))
+  return pos * angle_rates
+
+def positional_encoding(position, d_model_size):
+  # create the sinusoidal pattern for the positional encoding
+  angle_rads = angle_defn(np.arange(position)[:, np.newaxis], np.arange(d_model_size)[np.newaxis, :], d_model_size)
+  
+  sines = np.sin(angle_rads[:, 0::2])
+  cosines = np.cos(angle_rads[:, 1::2])
+  
+  pos_encoding = torch.tensor(np.concatenate([sines, cosines], axis=-1)[np.newaxis, ...], dtype=torch.float)
+  return pos_encoding
+
+def scaled_dot_product_attention(q, k, v, mask):
+  # calculate attention
+  matmul_qk = torch.matmul(q, k.permute(0,1,3,2))
+  
+  dk = k.shape[-1]
+  scaled_attention_logits = matmul_qk / np.sqrt(dk)
+
+  if mask is not None:
+    scaled_attention_logits += (mask * -1e9)
+    
+  attention_weights = torch.softmax(scaled_attention_logits, dim=-1) 
+  output = torch.matmul(attention_weights, v)
+  
+  return output
+
+
+class MultiHeadAttention(torch.nn.Module):
+  def __init__(self, d_model_size, num_heads):
+    super(MultiHeadAttention, self).__init__()
+    self.num_heads = num_heads
+    self.d_model_size = d_model_size
+    
+    self.depth = int(d_model_size / self.num_heads)
+    
+    self.Wq = torch.nn.Linear(d_model_size, d_model_size)
+    self.Wk = torch.nn.Linear(d_model_size, d_model_size)
+    self.Wv = torch.nn.Linear(d_model_size, d_model_size)
+    
+    self.dense = torch.nn.Linear(d_model_size, d_model_size)
+        
+  def split_into_heads(self, x, batch_size):
+    x = x.reshape(batch_size, -1, self.num_heads, self.depth)
+    return x.permute([0, 2, 1, 3])
+    
+  def forward(self, v, k, q, mask):
+    batch_size = q.shape[0]
+    
+    q = self.Wq(q)
+    k = self.Wk(k)
+    v = self.Wv(v)
+    
+    q = self.split_into_heads(q, batch_size)
+    k = self.split_into_heads(k, batch_size)
+    v = self.split_into_heads(v, batch_size)
+    
+    scaled_attention = scaled_dot_product_attention(q, k, v, mask).permute([0, 2, 1, 3])
+    original_size_attention = scaled_attention.reshape(batch_size, -1, self.d_model_size)
+    output = self.dense(original_size_attention)
+        
+    return output
+
+
+
+def point_wise_feed_forward_network(d_model_size, dff):
+  return torch.nn.Sequential(torch.nn.Linear(d_model_size, dff), torch.nn.ReLU(), torch.nn.Linear(dff, d_model_size))
+
+
+class EncoderLayer(torch.nn.Module):
+  def __init__(self, d_model_size, num_heads, dff, rate=0.1):
+    super(EncoderLayer, self).__init__()
+
+    self.multi_head_attention = MultiHeadAttention(d_model_size, num_heads)
+    self.ffn = point_wise_feed_forward_network(d_model_size, dff)
+
+    self.layernorm1 = torch.nn.LayerNorm(d_model_size, eps=1e-6)
+    self.layernorm2 = torch.nn.LayerNorm(d_model_size, eps=1e-6)
+    
+    self.dropout1 = torch.nn.Dropout(rate)
+    self.dropout2 = torch.nn.Dropout(rate)
+    
+  def forward(self, x, mask):
+    normed = self.layernorm1(x)
+    attn_output  = self.multi_head_attention(normed, normed, normed, mask)
+    attn_output = self.dropout1(attn_output)
+    out1 = x + attn_output
+
+    out2 = self.layernorm2(out1)
+    ffn_output = self.ffn(out2)
+    ffn_output = self.dropout2(ffn_output)
+    out2 = out1 + ffn_output
+    
+    return out2
+
+
+
+
+class Encoder(torch.nn.Module):
+  def __init__(self, num_layers=48, d_model_size=1280, num_heads=16, dff=8192, input_vocab_size=50000,
+               rate=0.1, **kwargs):
+    super(Encoder, self).__init__()
+
+    self.d_model_size = d_model_size
+    self.num_layers = num_layers
+    
+    self.pos_encoding = positional_encoding(input_vocab_size, self.d_model_size).to('cuda')
+
+    for i in range(num_layers):
+      setattr(self, "layer%i" % i, EncoderLayer(d_model_size, num_heads, dff, rate))
+    
+    self.layernorm = torch.nn.LayerNorm(d_model_size, eps=1e-6)  
+    self.dropout = torch.nn.Dropout(rate)
+
+  def forward(self, x):
+
+    seq_len = x.shape[1]
+    
+    mask = torch.triu(torch.ones(seq_len, seq_len), 1).to('cuda')
+    
+    x *= np.sqrt(self.d_model_size)
+    x += self.pos_encoding[:, :seq_len, :]
+
+    x = self.dropout(x)
+    
+    for i in range(self.num_layers):
+      x = getattr(self, "layer%i" % i)(x, mask)
+    return self.layernorm(x)

generator/ctrl/source_attribution.py

@@ -0,0 +1,192 @@
+from __future__ import division
+from __future__ import print_function
+import tensorflow as tf
+import os
+import numpy as np
+tf.enable_eager_execution()
+import transformer
+import argparse
+import pdb
+import sys
+import re
+from collections import Counter
+from tensorflow.python import debug as tf_debug
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import embedding_ops
+import fastBPE
+import platform
+
+use_py3 = platform.python_version()[0] == '3'
+
+parser = argparse.ArgumentParser(description='TensorFlow code for generating from CTRL')
+parser.add_argument('--model_dir', type=str, required=True,
+                                        help='location of model checkpoint')
+parser.add_argument('--seed', type=int, default=1337,
+                                        help='random seed for TensorFlow, numpy and PythonHash')
+
+args = parser.parse_args()
+tf.random.set_random_seed(args.seed)
+os.environ['PYTHONHASHSEED'] = str(args.seed)
+np.random.seed(args.seed)
+
+# load the vocabulary from file
+vocab = open('vocab').read().decode(encoding='utf-8').split('\n') if not use_py3 else open('vocab', encoding='utf-8').read().split('\n')
+vocab = list(map(lambda x: x.split(' ')[0], vocab)) + ['<unk>'] + ['\n']
+print ('{} unique words'.format(len(vocab)))
+
+# length of the vocabulary
+vocab_size = len(vocab)
+
+# define the numericalization map
+# idx2word maps the numericalized ID to the word
+# word2idx maps the word to the numericalized ID
+word2idx = {u:i for i, u in enumerate(vocab)}
+idx2word = np.array(vocab)
+
+
+
+# sequence length to use for the transformer
+# the model is trained with a seq_length of 512
+# so, any value <= 512 should work
+seq_length = 256
+
+
+
+
+# the dimension of the transformer
+embedding_dim = 1280
+
+
+# Now, we begin defining the model
+# we defer the transformer definition to transformer.py
+# here, we only define the tied softmax layer
+# this layer ties the softmax weights to the input embeddings
+class TiedEmbeddingSoftmax(tf.keras.layers.Layer):
+
+  def __init__(self, vocab_size=vocab_size, embedding_size=embedding_dim, **kwargs):
+    super(TiedEmbeddingSoftmax, self).__init__()
+    self.w = self.add_weight(name='w', shape=(vocab_size, embedding_size),
+                             initializer='random_normal',
+                             trainable=True)
+    self.b = self.add_weight(name='b', shape=(vocab_size,),
+                             initializer='zeros',
+                             trainable=True)
+
+  def call(self, inputs, embed=True):
+    if embed:
+      dtype = tf.keras.backend.dtype(inputs)
+      if dtype != 'int32' and dtype != 'int64':
+        inputs = math_ops.cast(inputs, 'int32')
+      return embedding_ops.embedding_lookup(self.w, inputs)
+    else:
+      return tf.tensordot(inputs, tf.transpose(self.w), 1) + self.b
+
+# input for the keras model
+tokens = tf.keras.layers.Input(shape=(seq_length,), dtype='int32')
+
+# instantiates a tied softmax class
+tied_embedding_softmax = TiedEmbeddingSoftmax()
+
+# embedded tokens, before passing it to the transformer
+embedded = tied_embedding_softmax(tokens, embed=True)
+
+# the activations after passing it from the transformer
+# for some odd reason, TPUs don't play well with specifying the arguments of the Encoder() function
+# so you have to leave them at their defaults
+transformed = transformer.Encoder()(embedded, training=False)
+
+
+# pass the activations from our tiedsoftmax class
+# this time with embed=False denoting that we are doing the softmax operation
+# and not a lookup
+logits = tied_embedding_softmax(transformed, embed=False)
+
+
+# finally, define the Keras model with inputs as tokens and outputs as the logits we just computed
+model = tf.keras.Model(inputs=tokens, outputs=logits)
+
+
+# the loss function is a simple categorical crossentropy between the logits and the labels
+def loss(labels, logits):
+    return tf.keras.losses.sparse_categorical_crossentropy(labels, logits, from_logits=True)
+
+# the optimizer is not used since this code only supports inference
+# however, to compile the model, we still define it
+optimizer = tf.contrib.tpu.CrossShardOptimizer(
+    tf.contrib.estimator.clip_gradients_by_norm(
+        tf.train.AdagradOptimizer(learning_rate=1e-2), 0.25)
+    )        
+
+# compile the model with the optimizer and loss            
+model.compile(optimizer=optimizer, loss=loss)
+print(model.summary())
+
+
+# IMPORTANT
+# this is where the saved model is presented to the code
+# the model directory should have the model checkpoint and
+# a checkpoint file
+run_config = tf.contrib.tpu.RunConfig(
+        model_dir=args.model_dir)
+
+
+# this converts the Keras model to a TensorFlow estimator
+# this step is critical
+# remember to patch the TF 1.14 file before running the code, else you're going to see errors here
+estimator_model = tf.keras.estimator.model_to_estimator(keras_model=model, config=run_config)
+
+# we now create a serving function from this estimator
+# this enables us to load the model once and easily query it multiple times
+def serving_input_fn():
+    inputs = {'input_1': tf.placeholder(tf.int32, [2,seq_length])}
+    return tf.estimator.export.ServingInputReceiver(inputs, inputs)
+predict_fn = tf.contrib.predictor.from_estimator(estimator_model, serving_input_fn)
+
+# almost there, we now take the user prompt and tokenize with BPE
+# load BPE codes
+bpe = fastBPE.fastBPE('codes', 'vocab')
+
+domains = []
+with open('control_codes.txt', 'r') as f:
+    domains = [line.split() for line in f.readlines()]
+    domains = [(t[1], float(t[0])) for t in domains]
+
+    
+while True:
+    _prompt = raw_input('ENTER PROMPT: ') if not use_py3 else input('ENTER PROMPT: ')
+
+    ppls = {}
+    # loop over all domains and compute perplexity
+    for domain, domain_prior in domains:
+       print(u'computing for domain: {}'.format(domain))
+       # tokenize data and add domain tag to it
+       prompt = domain + u' ' + _prompt
+       split_prompt = bpe.apply([prompt])[0].split()
+
+       # numericalize data and pad to the seq_len dimension
+       text = [word2idx[i] for i in split_prompt]
+       padding_text = text + [0] * (seq_length - len(text))
+       tokens_generated = np.tile(padding_text, (2,1)) 
+       output_scores = predict_fn({'input_1':tokens_generated})['tied_embedding_softmax'].squeeze()[0]
+       token_scores = output_scores[:-1]
+
+       # compute the perplexity for this sequence
+       xent = 0
+       for sequence_idx, token_idx in enumerate(text[1:]):
+           token = idx2word[token_idx]
+
+           # compute the probability of this token
+           Z = np.exp(token_scores[sequence_idx]).sum()
+           token_prob = np.exp(token_scores[sequence_idx, token_idx]) / Z 
+           xent -= np.log(token_prob) / len(text[1:])
+           
+       ppls[domain] = round(np.exp(xent), 6)
+       #print(u'{} ppl = {}'.format(domain, ppls[domain]))
+
+    # sort the domains based on perplexities and print
+    ppls = [(k, v) for k, v in ppls.items()]
+    ppls.sort(key=lambda x: x[1])
+    print('PROMPT: {}'.format(_prompt))
+    for t in ppls:
+        domain, ppl = t
+        print(u'{} ppl = {}'.format(domain, ppl))

generator/ctrl/training_utils/README.md

@@ -0,0 +1,85 @@
+# Fine-Tuning the Model on Custom Dataset
+
+
+This folder contains sample code to fine-tune the model on custom data. It is primarily targeted for GPU usage, but there are pointers throughout showing how to run on TPUs as well. 
+
+Fine-tuning can be used to augment existing control codes or add new control codes. There are 5 steps elaborated upon in the example below: 
+
+1. Patch `keras.py` as in the generation script
+2. Obtain raw versions of your text files
+3. Convert this text data into TFRecords; _if you wish to use TPUs, you must transfer these records to GCS._
+4. Fine-tuning the model on these TFRecords files
+5. Testing that the generation works. 
+
+## Example of adding a new control code
+
+Let's begin by adding a new control code `Moby` that is associated with the book [Moby Dick](https://www.gutenberg.org/ebooks/2701)
+
+If you run `generation.py` with the pretrained models available and try to use this control code, you will find that the model outputs gibberish. Great! It is indeed a fresh new control code. We will run this again after training as a sanity check. 
+
+### Step 1 - Patch your `keras.py`
+As is required for the generation script, you must patch your `keras.py`. If you patched it before, please roll-back and re-patch with the latest version. 
+
+There are two changes: (1) it defaults to `use_tpu=False` so training/inference takes place on GPUs, (2) the batch size defaults to 4 for GPU training. You might need to go lower depending on your machine. 
+
+You can leave `use_tpu=True` if you wish to train on TPUs and adjust the batch size accordingly. 
+
+### Step 2 - Obtain Your Data
+
+The book is available publicly; you can simply download it as 
+
+```
+wget -O moby_dick.txt https://www.gutenberg.org/files/2701/2701-0.txt
+```
+
+### Step 3 - Convert Data to TFRecords
+
+We include the file `make_tf_records.py` to facilitate this.
+
+Run:
+
+```
+python make_tf_records.py --text_file moby_dick.txt --control_code Moby --sequence_len 256
+```
+
+It has three arguments: `text_file` which specifies the name of the file to convert, `control_code` which specifies one token (must be in vocabulary) to append to each example, and `sequence_len` which specifies the sequence length to use to create the data. This must match the sequence length of the model being trained. 
+
+
+### Step 4 - Train!
+
+Simply run `python training.py --model_dir <path_to_model>.ckpt/ --iterations <number_of_iterations>`
+
+The script picks up all TFRecords in the current folder and fine-tunes the model provided in the `--model_dir` flag. 
+
+If you intend to use TPUs, you must transfer these TFRecords to GCS and edit the location of the data path used by `input_fn` to the GCS bucket. 
+
+To very important gotchas here:
+
+1. If you have very limited data, the model will very likely overfit and end up memorizing. At the moment, just keep the `--iterations` flag low, preferably equivalent to one epoch or so. 
+
+2. The model is updated and stored in the same directory, if you don't wish to overwrite your model files, please create a backup before you run the training code. 
+
+### Step 5 - Generate! 
+
+We ran `python training.py --model_dir seqlen256_v1.ckpt/ --iterations 250` and try generating with the `Moby` control code. 
+
+Running with the `Moby` control code and a prompt of `I` yields something reasonable in-domain:
+
+```
+Moby I <GENERATION_BEGINS> was a little fellow, and he was a
+ great man, what should that matter? And yet it seemed to me that
+ Queequegs words about his father were true. He had been very angry
+ with him, because the old man would not let him go a-whaling...
+```
+
+Providing a prompt also works:
+
+```
+Moby Then I realized, it wasn't one white whale but three! <GENERATION_BEGINS> And all three
+ were making straight for my boat, which was now some distance away.
+
+ But the three spouts seemed to be coming from different directions, and
+ as they drew nearer and nearer, their tongues began licking up the
+ brine like so many hungry wolves at a carcass...
+```
+

generator/ctrl/training_utils/make_tf_records.py

@@ -0,0 +1,81 @@
+import numpy as np
+import os
+import tensorflow as tf
+import tqdm
+import pdb
+import glob
+import time
+import sys
+import re
+import argparse
+import fastBPE
+import platform
+
+use_py3 = platform.python_version()[0] == '3'
+
+parser = argparse.ArgumentParser(description='TensorFlow code for creating TFRecords data')
+parser.add_argument('--text_file', type=str, required=True,
+                                        help='location of text file to convert to TFRecords')
+parser.add_argument('--control_code', type=str, required=True,
+                                        help='control code to use for this file. must be in the vocabulary, else it will error out.')
+parser.add_argument('--sequence_len', type=int, required=True,
+                                        help='sequence length of model being fine-tuned (256 or 512)')
+
+args = parser.parse_args()
+
+
+path_to_train_file = fname = args.text_file
+domain = [args.control_code]
+
+train_text = open(path_to_train_file, 'rb').read().decode(encoding='utf-8')
+bpe = fastBPE.fastBPE('../codes', '../vocab')
+tokenized_train_text = bpe.apply([train_text.encode('ascii', errors='ignore') if not use_py3 else train_text])[0] # will NOT work for non-English texts 
+# if you want to run non-english text, please tokenize separately using ./fast applybpe and then run this script on the .bpe file with utf8 encoding
+
+tokenized_train_text = re.findall(r'\S+|\n', tokenized_train_text)
+tokenized_train_text = list(filter(lambda x: x != u'@@', tokenized_train_text))
+
+# load the vocabulary from file
+vocab = open('../vocab').read().decode(encoding='utf-8').split('\n') if not use_py3 else open('../vocab', encoding='utf-8').read().split('\n')
+vocab = list(map(lambda x: x.split(' ')[0], vocab)) + ['<unk>'] + ['\n']
+print ('{} unique words'.format(len(vocab)))
+
+if args.control_code not in vocab:
+    print('Provided control code is not in the vocabulary')
+    print('Please provide a different one; refer to the vocab file for allowable tokens')
+    sys.exit(1)
+    
+# Creating a mapping from unique characters to indices
+word2idx = {u:i for i, u in enumerate(vocab)}
+idx2word = np.array(vocab)
+
+seq_length = args.sequence_len-1
+
+def numericalize(x):
+    count = 0
+    for i in x:
+        if i not in word2idx:
+            print(i)
+            count += 1
+    return count>1, [word2idx.get(i, word2idx['<unk>'])  for i in x]
+
+tfrecords_fname = fname.lower()+'.tfrecords'
+
+total = 0
+skipped = 0
+with tf.io.TFRecordWriter(tfrecords_fname) as writer:
+    for i in tqdm.tqdm(range(0, len(tokenized_train_text), seq_length)):
+        flag_input, inputs = numericalize(domain+tokenized_train_text[i:i+seq_length])
+        flag_output, outputs = numericalize(tokenized_train_text[i:i+seq_length+1])
+        total += 1
+        if flag_input or flag_output:
+            skipped += 1
+            continue
+
+        if len(inputs)!=seq_length+1 or len(outputs)!=seq_length+1:
+            break
+        example_proto = tf.train.Example(features=tf.train.Features(feature={'input': tf.train.Feature(int64_list=tf.train.Int64List(value=inputs)),
+                                                                             'output': tf.train.Feature(int64_list=tf.train.Int64List(value=outputs))}))
+        writer.write(example_proto.SerializeToString())
+print('Done')
+print('Skipped', skipped, 'of', total)

generator/ctrl/training_utils/training.py

@@ -0,0 +1,164 @@
+from __future__ import division
+from __future__ import print_function
+import sys
+sys.path.append('../')
+import tensorflow as tf
+import os
+import numpy as np
+tf.enable_eager_execution()
+import transformer
+import argparse
+import pdb
+import re
+from collections import Counter
+from tensorflow.python import debug as tf_debug
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import embedding_ops
+import fastBPE
+import platform
+
+use_py3 = platform.python_version()[0] == '3'
+
+parser = argparse.ArgumentParser(description='TensorFlow code for generating from CTRL')
+parser.add_argument('--model_dir', type=str, required=True,
+                                        help='location of model checkpoint')
+parser.add_argument('--seed', type=int, default=1337,
+                                        help='random seed for TensorFlow, numpy and PythonHash')
+parser.add_argument('--sequence_len', type=int, default=256,
+                                        help='sequence len of model being fine-tuned (must match also the TFRecords)')
+parser.add_argument('--iterations', type=int, default=1000,
+                                        help='random seed for TensorFlow, numpy and PythonHash')
+
+args = parser.parse_args()
+tf.random.set_random_seed(args.seed)
+os.environ['PYTHONHASHSEED'] = str(args.seed)
+np.random.seed(args.seed)
+
+# load the vocabulary from file
+vocab = open('../vocab').read().decode(encoding='utf-8').split('\n') if not use_py3 else open('../vocab', encoding='utf-8').read().split('\n')
+vocab = list(map(lambda x: x.split(' ')[0], vocab)) + ['<unk>'] + ['\n']
+print ('{} unique words'.format(len(vocab)))
+
+# length of the vocabulary
+vocab_size = len(vocab)
+
+# define the numericalization map
+# idx2word maps the numericalized ID to the word
+# word2idx maps the word to the numericalized ID
+word2idx = {u:i for i, u in enumerate(vocab)}
+idx2word = np.array(vocab)
+
+
+
+# sequence length to use for the transformer
+# must match the model being fine-tuned
+seq_length = args.sequence_len
+
+def input_fn(params=None):
+    print('READING!', params)
+    dataset = tf.data.Dataset.list_files(tf.io.gfile.glob('./*.tfrecords'), shuffle=True)
+
+    tf_data = tf.data.TFRecordDataset(dataset)
+    myfeatures = {
+        'input': tf.io.FixedLenFeature([256], tf.int64),
+        'output': tf.io.FixedLenFeature([256], tf.int64)
+        }
+
+    def _parse_text_function(example_proto):
+        blah = tf.io.parse_single_example(example_proto, myfeatures)
+        return blah['input'], blah['output']
+    
+    train_data = tf_data.map(_parse_text_function).batch(params['batch_size'], drop_remainder=True).repeat().shuffle(10000)#.prefetch(tf.contrib.data.AUTOTUNE)
+    
+    return train_data
+
+
+# the dimension of the transformer
+embedding_dim = 1280
+
+
+# Now, we begin defining the model
+# we defer the transformer definition to transformer.py
+# here, we only define the tied softmax layer
+# this layer ties the softmax weights to the input embeddings
+class TiedEmbeddingSoftmax(tf.keras.layers.Layer):
+
+  def __init__(self, vocab_size=vocab_size, embedding_size=embedding_dim, **kwargs):
+    super(TiedEmbeddingSoftmax, self).__init__()
+    self.w = self.add_weight(name='w', shape=(vocab_size, embedding_size),
+                             initializer='random_normal',
+                             trainable=True)
+    self.b = self.add_weight(name='b', shape=(vocab_size,),
+                             initializer='zeros',
+                             trainable=True)
+
+  def call(self, inputs, embed=True):
+    if embed:
+      dtype = tf.keras.backend.dtype(inputs)
+      if dtype != 'int32' and dtype != 'int64':
+        inputs = math_ops.cast(inputs, 'int32')
+      return embedding_ops.embedding_lookup(self.w, inputs)
+    else:
+      return tf.tensordot(inputs, tf.transpose(self.w), 1) + self.b
+
+# input for the keras model
+tokens = tf.keras.layers.Input(shape=(seq_length,), dtype='int32')
+
+# instantiates a tied softmax class
+tied_embedding_softmax = TiedEmbeddingSoftmax()
+
+# embedded tokens, before passing it to the transformer
+embedded = tied_embedding_softmax(tokens, embed=True)
+
+# the activations after passing it from the transformer
+# for some odd reason, TPUs don't play well with specifying the arguments of the Encoder() function
+# so you have to leave them at their defaults
+transformed = transformer.Encoder()(embedded, training=False)
+
+
+# pass the activations from our tiedsoftmax class
+# this time with embed=False denoting that we are doing the softmax operation
+# and not a lookup
+logits = tied_embedding_softmax(transformed, embed=False)
+
+
+# finally, define the Keras model with inputs as tokens and outputs as the logits we just computed
+model = tf.keras.Model(inputs=tokens, outputs=logits)
+
+
+# the loss function is a simple categorical crossentropy between the logits and the labels
+def loss(labels, logits):
+    return tf.keras.losses.sparse_categorical_crossentropy(labels, logits, from_logits=True)
+
+# the optimizer is not used since this code only supports inference
+# however, to compile the model, we still define it
+optimizer = tf.contrib.estimator.clip_gradients_by_norm(
+        tf.train.AdagradOptimizer(learning_rate=1e-2), 0.25)
+
+
+# compile the model with the optimizer and loss            
+model.compile(optimizer=optimizer, loss=loss)
+print(model.summary())
+
+
+# IMPORTANT
+# this is where the saved model is presented to the code
+# the model directory should have the model checkpoint and
+# a checkpoint file
+run_config = tf.contrib.tpu.RunConfig(
+        model_dir=args.model_dir)
+
+
+# this converts the Keras model to a TensorFlow estimator
+# this step is critical
+# remember to patch the TF 1.14 file before running the code, else you're going to see errors here
+
+run_config = tf.contrib.tpu.RunConfig(
+        model_dir=args.model_dir,
+        session_config=tf.ConfigProto(allow_soft_placement=True, log_device_placement=True),
+        tpu_config=tf.contrib.tpu.TPUConfig(iterations_per_loop=100, num_cores_per_replica=1, input_partition_dims=[[1, 1], [1, 1]], per_host_input_for_training=3))
+tf.logging.set_verbosity(tf.logging.INFO)
+
+estimator_model = tf.keras.estimator.model_to_estimator(keras_model=model, config=run_config)
+
+estimator_model.train(input_fn=input_fn, steps=args.iterations)

generator/ctrl/transformer.py

@@ -0,0 +1,139 @@
+import tensorflow as tf
+import numpy as np
+
+def angle_defn(pos, i, d_model_size):
+  angle_rates = 1 / np.power(10000, (2 * (i//2)) / np.float32(d_model_size))
+  return pos * angle_rates
+
+def positional_encoding(position, d_model_size):
+  # create the sinusoidal pattern for the positional encoding
+  angle_rads = angle_defn(np.arange(position)[:, np.newaxis], np.arange(d_model_size)[np.newaxis, :], d_model_size)
+  
+  sines = np.sin(angle_rads[:, 0::2])
+  cosines = np.cos(angle_rads[:, 1::2])
+  
+  pos_encoding = tf.cast(np.concatenate([sines, cosines], axis=-1)[np.newaxis, ...], dtype=tf.float32)
+  return pos_encoding 
+
+
+
+def scaled_dot_product_attention(q, k, v, mask):
+  # calculate attention
+  matmul_qk = tf.matmul(q, k, transpose_b=True)
+  
+  dk = tf.cast(tf.shape(k)[-1], tf.float32)
+  scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)
+
+  if mask is not None:
+    scaled_attention_logits += (mask * -1e9)
+    
+  attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1) 
+  output = tf.matmul(attention_weights, v) 
+  return output
+
+class MultiHeadAttention(tf.keras.layers.Layer):
+  def __init__(self, d_model_size, num_heads):
+    super(MultiHeadAttention, self).__init__()
+    self.num_heads = num_heads
+    self.d_model_size = d_model_size
+    
+    self.depth = int(d_model_size / self.num_heads)
+    
+    self.Wq = tf.keras.layers.Dense(d_model_size)
+    self.Wk = tf.keras.layers.Dense(d_model_size)
+    self.Wv = tf.keras.layers.Dense(d_model_size)
+    
+    self.dense = tf.keras.layers.Dense(d_model_size)
+        
+  def split_into_heads(self, x, batch_size):
+    x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))
+    return tf.transpose(x, perm=[0, 2, 1, 3])
+    
+  def call(self, v, k, q, mask):
+    batch_size = tf.shape(q)[0]
+    
+    q = self.Wq(q)
+    k = self.Wk(k)
+    v = self.Wv(v)
+    
+    q = self.split_into_heads(q, batch_size)
+    k = self.split_into_heads(k, batch_size)
+    v = self.split_into_heads(v, batch_size)
+    
+    scaled_attention = tf.transpose(scaled_dot_product_attention(q, k, v, mask), perm=[0, 2, 1, 3])
+    original_size_attention = tf.reshape(scaled_attention,  (batch_size, -1, self.d_model_size))
+    output = self.dense(original_size_attention) 
+        
+    return output
+
+
+
+def point_wise_feed_forward_network(d_model_size, dff):
+  return tf.keras.Sequential([tf.keras.layers.Dense(dff, activation='relu'), 
+      tf.keras.layers.Dense(d_model_size)])
+
+
+class EncoderLayer(tf.keras.layers.Layer):
+  def __init__(self, d_model_size, num_heads, dff, rate=0.1):
+    super(EncoderLayer, self).__init__()
+
+    self.multi_head_attention = MultiHeadAttention(d_model_size, num_heads)
+    self.ffn = point_wise_feed_forward_network(d_model_size, dff)
+
+    self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
+    self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
+    
+    self.dropout1 = tf.keras.layers.Dropout(rate)
+    self.dropout2 = tf.keras.layers.Dropout(rate)
+    
+  def call(self, x, training, mask):
+    normed = self.layernorm1(x)
+    attn_output  = self.multi_head_attention(normed, normed, normed, mask)
+    attn_output = self.dropout1(attn_output, training=training)
+    out1 = x + attn_output
+
+    out2 = self.layernorm2(out1)
+    ffn_output = self.ffn(out2)
+    ffn_output = self.dropout2(ffn_output, training=training)
+    out2 = out1 + ffn_output
+    
+    return out2
+
+
+
+
+class Encoder(tf.keras.layers.Layer):
+  def __init__(self, num_layers=48, d_model_size=1280, num_heads=16, dff=8192, input_vocab_size=50000,
+               rate=0.1, **kwargs):
+    super(Encoder, self).__init__()
+
+    self.d_model_size = d_model_size
+    self.num_layers = num_layers
+    
+    self.pos_encoding = positional_encoding(input_vocab_size, self.d_model_size)
+
+    for i in range(num_layers):
+      setattr(self, "layer%i" % i, EncoderLayer(d_model_size, num_heads, dff, rate))
+    
+    self.layernorm = tf.keras.layers.LayerNormalization(epsilon=1e-6)  
+    self.dropout = tf.keras.layers.Dropout(rate)
+
+  def get_config(self):
+    base_config = super(Encoder, self).get_config()
+    return base_config
+  
+  def call(self, x, training):
+
+    seq_len = tf.shape(x)[1]
+    
+    mask = 1 - tf.linalg.band_part(tf.ones((seq_len, seq_len)), -1, 0)
+    
+    x *= tf.math.sqrt(tf.cast(self.d_model_size, tf.float32))
+    x += self.pos_encoding[:, :seq_len, :]
+
+    x = self.dropout(x, training=training)
+    
+    for i in range(self.num_layers):
+      x = getattr(self, "layer%i" % i)(x, training, mask)
+    return self.layernorm(x)
+