in _force_token_id_to_be_generated we have much simpler torch code:
scores[:, [x for if x != token_id]] = -float("inf")Is it possible to make the TF Code simpler? TF doesn't support assignment, but maybe to and from numpy could be faster. Would definitely be simpler.
@staticmethod
def _force_token_id_to_be_generated(scores, token_id) -> None:
"""force one of token_ids to be generated by setting prob of all other tokens to 0 (logprob=-float("inf"))"""
output_list = []
# Is there a better way to do in TF?
bs, vocab_size = scores.shape
inf_tensor = tf.convert_to_tensor([-float("inf")] * bs, dtype=scores.dtype)
for x in range(vocab_size):
if x != token_id:
output_list.append(inf_tensor)
else:
output_list.append(scores[:, x])
scores = tf.stack(output_list, axis=1, name="scores")
assert scores.shape == (bs, vocab_size)
return scores
in
_force_token_id_to_be_generatedwe have much simpler torch code:Is it possible to make the TF Code simpler? TF doesn't support assignment, but maybe to and from
numpycould be faster. Would definitely be simpler.