in lm_human_preferences/language/model.py [0:0]
def __call__(self, *, X, Y=None, past=None, past_tokens=None, mask=None,
padding_token: Optional[int]=None, do_dropout=False):
X = tf.convert_to_tensor(X, dtype=tf.int32)
if mask is not None:
mask = tf.convert_to_tensor(mask, dtype=tf.bool)
assert mask.dtype == tf.bool
if padding_token is not None:
assert mask is None, 'At most one of mask and padding_token should be set'
mask = tf.not_equal(X, padding_token)
X = tf.where(mask, X, tf.zeros_like(X))
if past is not None:
assert past_tokens is not None, 'padding_token requires past_tokens'
mask = tf.concat([tf.not_equal(past_tokens, padding_token), mask], axis=1)
with tf.variable_scope(self.scope, reuse=self.built, auxiliary_name_scope=not self.built):
self.built = True
results = {}
batch, sequence = utils.shape_list(X)
seed = tf.random.uniform(dtype=tf.int64, shape=[2], minval=-2**63, maxval=2**63-1)
wpe_seed, wte_seed, blocks_seed, heads_seed = split_seed(seed, 4)
wpe = tf.get_variable('wpe', [self.hparams.n_ctx, self.hparams.n_embd],
initializer=tf.random_normal_initializer(stddev=0.01))
wte = tf.get_variable('wte', [self.hparams.n_vocab, self.hparams.n_embd],
initializer=tf.random_normal_initializer(stddev=0.02))
wpe = dropout(wpe, self.hparams.embd_pdrop,
do_dropout=do_dropout, stateless=True, seed=wpe_seed, name='wpe_drop')
wte = dropout(wte, self.hparams.embd_pdrop,
do_dropout=do_dropout, stateless=True, seed=wte_seed, name='wte_drop')
past_length = 0 if past is None else tf.shape(past)[-2]
positions = positions_for(batch=batch, sequence=sequence, past_length=past_length, mask=mask)
h = embed(X, wte) + embed(positions, wpe)
# Transformer
presents = []
pasts = tf.unstack(past, axis=1) if past is not None else [None] * self.hparams.n_layer
assert len(pasts) == self.hparams.n_layer
block_seeds = split_seed(blocks_seed, self.hparams.n_layer)
for layer, (past, block_seed) in enumerate(zip(pasts, block_seeds)):
h, present = block(
h, 'h%d' % layer, past=past, mask=mask, do_dropout=do_dropout, scale=True,
hparams=self.hparams, seed=block_seed)
presents.append(present)
results['present'] = tf.stack(presents, axis=1)
h = norm(h, 'ln_f')
if mask is not None:
# For non-present tokens, use the output from the last present token instead.
present_indices = utils.where(mask[:,past_length:], tf.tile(tf.range(sequence)[None,:], [batch, 1]), -1)
use_indices = utils.cumulative_max(present_indices)
# assert since GPUs don't
with tf.control_dependencies([tf.assert_none_equal(use_indices, -1)]):
h = utils.index_each(h, use_indices)
results['h'] = h
# Language model loss. Do tokens <n predict token n?
h_flat = tf.reshape(h, [batch*sequence, self.hparams.n_embd])
flat_lm_logits = tf.matmul(h_flat, wte, transpose_b=True)
labels = tf.concat([X[:, 1:], X[:, :1]], axis=1)
flat_labels = tf.reshape(labels, [batch*sequence])
flat_losses = tf.nn.sparse_softmax_cross_entropy_with_logits(
labels=flat_labels,
logits=flat_lm_logits)
lm_losses = tf.reshape(flat_losses, [batch, sequence])
lm_logits = tf.reshape(flat_lm_logits, [batch, sequence, -1])
relevant_losses = lm_losses[:, :-1]
results['lm_all_losses'] = relevant_losses
results['lm_logits'] = lm_logits
results['lm_losses'] = tf.reduce_mean(relevant_losses, axis=-1)
head_seeds = split_seed(heads_seed, len(self.scalar_heads))
for head_name, head_seed in zip(self.scalar_heads, head_seeds):
with tf.variable_scope(f"heads/{head_name}"):
dropped_h = \
dropout(h, self.hparams.head_pdrop, do_dropout=do_dropout, seed=head_seed, name='drop')
# TODO: refactor this, perhaps move to Policy
res, reg_loss = fc_layer(dropped_h, (), scale=0 if head_name == 'value' else None)
results[head_name] = tf.cast(res, dtype=tf.float32, name='res_cast')
results[f"{head_name}_regularizer"] = tf.cast(reg_loss, dtype=tf.float32, name='reg_loss_cast')
# All done!
return results