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 sample_model(
  12      model_name='124M',
  13      seed=None,
  14      nsamples=0,
  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      Run the sample_model
  24      :model_name=124M : String, which model to use
  25      :seed=None : Integer seed for random number generators, fix seed to
  26       reproduce results
  27      :nsamples=0 : Number of samples to return, if 0, continues to
  28       generate samples indefinately.
  29      :batch_size=1 : Number of batches (only affects speed/memory).
  30      :length=None : Number of tokens in generated text, if None (default), is
  31       determined by model hyperparameters
  32      :temperature=1 : Float value controlling randomness in boltzmann
  33       distribution. Lower temperature results in less random completions. As the
  34       temperature approaches zero, the model will become deterministic and
  35       repetitive. Higher temperature results in more random completions.
  36      :top_k=0 : Integer value controlling diversity. 1 means only 1 word is
  37       considered for each step (token), resulting in deterministic completions,
  38       while 40 means 40 words are considered at each step. 0 (default) is a
  39       special setting meaning no restrictions. 40 generally is a good value.
  40       :models_dir : path to parent folder containing model subfolders
  41       (i.e. contains the <model_name> folder)
  42      """
  43      models_dir = os.path.expanduser(os.path.expandvars(models_dir))
  44      enc = encoder.get_encoder(model_name, models_dir)
  45      hparams = model.default_hparams()
  46      with open(os.path.join(models_dir, model_name, 'hparams.json')) as f:
  47          hparams.override_from_dict(json.load(f))
  48  
  49      if length is None:
  50          length = hparams.n_ctx
  51      elif length > hparams.n_ctx:
  52          raise ValueError("Can't get samples longer than window size: %s" % hparams.n_ctx)
  53  
  54      with tf.Session(graph=tf.Graph()) as sess:
  55          np.random.seed(seed)
  56          tf.set_random_seed(seed)
  57  
  58          output = sample.sample_sequence(
  59              hparams=hparams, length=length,
  60              start_token=enc.encoder['<|endoftext|>'],
  61              batch_size=batch_size,
  62              temperature=temperature, top_k=top_k, top_p=top_p
  63          )[:, 1:]
  64  
  65          saver = tf.train.Saver()
  66          ckpt = tf.train.latest_checkpoint(os.path.join(models_dir, model_name))
  67          saver.restore(sess, ckpt)
  68  
  69          generated = 0
  70          while nsamples == 0 or generated < nsamples:
  71              out = sess.run(output)
  72              for i in range(batch_size):
  73                  generated += batch_size
  74                  text = enc.decode(out[i])
  75                  print("=" * 40 + " SAMPLE " + str(generated) + " " + "=" * 40)
  76                  print(text)
  77  
  78  if __name__ == '__main__':
  79      fire.Fire(sample_model)
  80  
  81