1 import numpy as np 2 import tensorflow as tf 3 from tensorflow.contrib.training import HParams 4 5 def default_hparams(): 6 return HParams( 7 n_vocab=0, 8 n_ctx=1024, 9 n_embd=768, 10 n_head=12, 11 n_layer=12, 12 ) 13 14 def shape_list(x): 15 """Deal with dynamic shape in tensorflow cleanly.""" 16 static = x.shape.as_list() 17 dynamic = tf.shape(x) 18 return [dynamic[i] if s is None else s for i, s in enumerate(static)] 19 20 def softmax(x, axis=-1): 21 x = x - tf.reduce_max(x, axis=axis, keepdims=True) 22 ex = tf.exp(x) 23 return ex / tf.reduce_sum(ex, axis=axis, keepdims=True) 24 25 def gelu(x): 26 return 0.5*x*(1+tf.tanh(np.sqrt(2/np.pi)*(x+0.044715*tf.pow(x, 3)))) 27 28 def norm(x, scope, *, axis=-1, epsilon=1e-5): 29 """Normalize to mean = 0, std = 1, then do a diagonal affine transform.""" 30 with tf.variable_scope(scope): 31 n_state = x.shape[-1].value 32 g = tf.get_variable('g', [n_state], initializer=tf.constant_initializer(1)) 33 b = tf.get_variable('b', [n_state], initializer=tf.constant_initializer(0)) 34 u = tf.reduce_mean(x, axis=axis, keepdims=True) 35 s = tf.reduce_mean(tf.square(x-u), axis=axis, keepdims=True) 36 x = (x - u) * tf.rsqrt(s + epsilon) 37 x = x*g + b 38 return x 39 40 def split_states(x, n): 41 """Reshape the last dimension of x into [n, x.shape[-1]/n].""" 42 *start, m = shape_list(x) 43 return tf.reshape(x, start + [n, m//n]) 44 45 def merge_states(x): 46 """Smash the last two dimensions of x into a single dimension.""" 47 *start, a, b = shape_list(x) 48 return tf.reshape(x, start + [a*b]) 49 50 def conv1d(x, scope, nf, *, w_init_stdev=0.02): 51 with tf.variable_scope(scope): 52 *start, nx = shape_list(x) 53 w = tf.get_variable('w', [1, nx, nf], initializer=tf.random_normal_initializer(stddev=w_init_stdev)) 54 b = tf.get_variable('b', [nf], initializer=tf.constant_initializer(0)) 55 c = tf.reshape(tf.matmul(tf.reshape(x, [-1, nx]), tf.reshape(w, [-1, nf]))+b, start+[nf]) 56 return c 57 58 def attention_mask(nd, ns, *, dtype): 59 """1's in the lower triangle, counting from the lower right corner. 60 61 Same as tf.matrix_band_part(tf.ones([nd, ns]), -1, ns-nd), but doesn't produce garbage on TPUs. 62 """ 63 i = tf.range(nd)[:,None] 64 j = tf.range(ns) 65 m = i >= j - ns + nd 66 return tf.cast(m, dtype) 67 68 69 def attn(x, scope, n_state, *, past, hparams): 70 assert x.shape.ndims == 3 # Should be [batch, sequence, features] 71 assert n_state % hparams.n_head == 0 72 if past is not None: 73 assert past.shape.ndims == 5 # Should be [batch, 2, heads, sequence, features], where 2 is [k, v] 74 75 def split_heads(x): 76 # From [batch, sequence, features] to [batch, heads, sequence, features] 77 return tf.transpose(split_states(x, hparams.n_head), [0, 2, 1, 3]) 78 79 def merge_heads(x): 80 # Reverse of split_heads 81 return merge_states(tf.transpose(x, [0, 2, 1, 3])) 82 83 def mask_attn_weights(w): 84 # w has shape [batch, heads, dst_sequence, src_sequence], where information flows from src to dst. 85 _, _, nd, ns = shape_list(w) 86 b = attention_mask(nd, ns, dtype=w.dtype) 87 b = tf.reshape(b, [1, 1, nd, ns]) 88 w = w*b - tf.cast(1e10, w.dtype)*(1-b) 89 return w 90 91 def multihead_attn(q, k, v): 92 # q, k, v have shape [batch, heads, sequence, features] 93 w = tf.matmul(q, k, transpose_b=True) 94 w = w * tf.rsqrt(tf.cast(v.shape[-1].value, w.dtype)) 95 96 w = mask_attn_weights(w) 97 w = softmax(w) 98 a = tf.matmul(w, v) 99 return a 100 101 with tf.variable_scope(scope): 102 c = conv1d(x, 'c_attn', n_state*3) 103 q, k, v = map(split_heads, tf.split(c, 3, axis=2)) 104 present = tf.stack([k, v], axis=1) 105 if past is not None: 106 pk, pv = tf.unstack(past, axis=1) 107 k = tf.concat([pk, k], axis=-2) 108 v = tf.concat([pv, v], axis=-2) 109 a = multihead_attn(q, k, v) 110 a = merge_heads(a) 111 a = conv1d(a, 'c_proj', n_state) 112 return a, present 113 114 115 def mlp(x, scope, n_state, *, hparams): 116 with tf.variable_scope(scope): 117 nx = x.shape[-1].value 118 h = gelu(conv1d(x, 'c_fc', n_state)) 119 h2 = conv1d(h, 'c_proj', nx) 120 return h2 121 122 123 def block(x, scope, *, past, hparams): 124 with tf.variable_scope(scope): 125 nx = x.shape[-1].value 126 a, present = attn(norm(x, 'ln_1'), 'attn', nx, past=past, hparams=hparams) 127 x = x + a 128 m = mlp(norm(x, 'ln_2'), 'mlp', nx*4, hparams=hparams) 129 x = x + m 130 return x, present 131 132 def past_shape(*, hparams, batch_size=None, sequence=None): 133 return [batch_size, hparams.n_layer, 2, hparams.n_head, sequence, hparams.n_embd // hparams.n_head] 134 135 def expand_tile(value, size): 136 """Add a new axis of given size.""" 137 value = tf.convert_to_tensor(value, name='value') 138 ndims = value.shape.ndims 139 return tf.tile(tf.expand_dims(value, axis=0), [size] + [1]*ndims) 140 141 def positions_for(tokens, past_length): 142 batch_size = tf.shape(tokens)[0] 143 nsteps = tf.shape(tokens)[1] 144 return expand_tile(past_length + tf.range(nsteps), batch_size) 145 146 147 def model(hparams, X, past=None, scope='model', reuse=False): 148 with tf.variable_scope(scope, reuse=reuse): 149 results = {} 150 batch, sequence = shape_list(X) 151 152 wpe = tf.get_variable('wpe', [hparams.n_ctx, hparams.n_embd], 153 initializer=tf.random_normal_initializer(stddev=0.01)) 154 wte = tf.get_variable('wte', [hparams.n_vocab, hparams.n_embd], 155 initializer=tf.random_normal_initializer(stddev=0.02)) 156 past_length = 0 if past is None else tf.shape(past)[-2] 157 h = tf.gather(wte, X) + tf.gather(wpe, positions_for(X, past_length)) 158 159 # Transformer 160 presents = [] 161 pasts = tf.unstack(past, axis=1) if past is not None else [None] * hparams.n_layer 162 assert len(pasts) == hparams.n_layer 163 for layer, past in enumerate(pasts): 164 h, present = block(h, 'h%d' % layer, past=past, hparams=hparams) 165 presents.append(present) 166 results['present'] = tf.stack(presents, axis=1) 167 h = norm(h, 'ln_f') 168 169 # Language model loss. Do tokens <n predict token n? 170 h_flat = tf.reshape(h, [batch*sequence, hparams.n_embd]) 171 logits = tf.matmul(h_flat, wte, transpose_b=True) 172 logits = tf.reshape(logits, [batch, sequence, hparams.n_vocab]) 173 results['logits'] = logits 174 return results 175