1  #!/usr/bin/env python3
   2  
   3  import fire
   4  import json
   5  import os
   6  import numpy as np
   7  import tensorflow as tf
   8  
   9  import model, sample, encoder
  10  
  11  def interact_model(
  12      model_name='124M',
  13      seed=None,
  14      nsamples=1,
  15      batch_size=1,
  16      length=None,
  17      temperature=1,
  18      top_k=0,
  19      top_p=1,
  20      models_dir='models',
  21  ):
  22      """
  23      Interactively run the model
  24      :model_name=124M : String, which model to use
  25      :seed=None : Integer seed for random number generators, fix seed to reproduce
  26       results
  27      :nsamples=1 : Number of samples to return total
  28      :batch_size=1 : Number of batches (only affects speed/memory).  Must divide nsamples.
  29      :length=None : Number of tokens in generated text, if None (default), is
  30       determined by model hyperparameters
  31      :temperature=1 : Float value controlling randomness in boltzmann
  32       distribution. Lower temperature results in less random completions. As the
  33       temperature approaches zero, the model will become deterministic and
  34       repetitive. Higher temperature results in more random completions.
  35      :top_k=0 : Integer value controlling diversity. 1 means only 1 word is
  36       considered for each step (token), resulting in deterministic completions,
  37       while 40 means 40 words are considered at each step. 0 (default) is a
  38       special setting meaning no restrictions. 40 generally is a good value.
  39       :models_dir : path to parent folder containing model subfolders
  40       (i.e. contains the <model_name> folder)
  41      """
  42      models_dir = os.path.expanduser(os.path.expandvars(models_dir))
  43      if batch_size is None:
  44          batch_size = 1
  45      assert nsamples % batch_size == 0
  46  
  47      enc = encoder.get_encoder(model_name, models_dir)
  48      hparams = model.default_hparams()
  49      with open(os.path.join(models_dir, model_name, 'hparams.json')) as f:
  50          hparams.override_from_dict(json.load(f))
  51  
  52      if length is None:
  53          length = hparams.n_ctx // 2
  54      elif length > hparams.n_ctx:
  55          raise ValueError("Can't get samples longer than window size: %s" % hparams.n_ctx)
  56  
  57      with tf.Session(graph=tf.Graph()) as sess:
  58          context = tf.placeholder(tf.int32, [batch_size, None])
  59          np.random.seed(seed)
  60          tf.set_random_seed(seed)
  61          output = sample.sample_sequence(
  62              hparams=hparams, length=length,
  63              context=context,
  64              batch_size=batch_size,
  65              temperature=temperature, top_k=top_k, top_p=top_p
  66          )
  67  
  68          saver = tf.train.Saver()
  69          ckpt = tf.train.latest_checkpoint(os.path.join(models_dir, model_name))
  70          saver.restore(sess, ckpt)
  71  
  72          while True:
  73              raw_text = input("Model prompt >>> ")
  74              while not raw_text:
  75                  print('Prompt should not be empty!')
  76                  raw_text = input("Model prompt >>> ")
  77              context_tokens = enc.encode(raw_text)
  78              generated = 0
  79              for _ in range(nsamples // batch_size):
  80                  out = sess.run(output, feed_dict={
  81                      context: [context_tokens for _ in range(batch_size)]
  82                  })[:, len(context_tokens):]
  83                  for i in range(batch_size):
  84                      generated += 1
  85                      text = enc.decode(out[i])
  86                      print("=" * 40 + " SAMPLE " + str(generated) + " " + "=" * 40)
  87                      print(text)
  88              print("=" * 80)
  89  
  90  if __name__ == '__main__':
  91      fire.Fire(interact_model)
  92  
  93