src/transformers/generation/tf_logits_process.py/TFNoBadWordsLogitsProcessor/__call__
class TFNoBadWordsLogitsProcessor:    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
        # We want to mask some banned tokens, at a score level. Since the banned tokens depend on the previous
        # `input_ids`, they may have a different length for each row, and they may even be empty for some rows.
        # To remain simple and XLA-compatible, we work on a per-row fashion.
        # TODO (Joao): this function might trigger XLA retracing as `cur_len` increases. Fix it if it becomes
        # a frequent choke point. (make `cur_len` a tensor?)
        def _get_row_updated_score(row_inputs: Tuple[tf.Tensor]) -> tf.Tensor:
            row_input_ids, row_score = row_inputs
            banned_tokens = self._calc_row_banned_bad_tokens(row_input_ids[:cur_len])
            banned_tokens_mask = tf.scatter_nd(
                indices=tf.expand_dims(banned_tokens, axis=-1),
                updates=tf.ones_like(banned_tokens, dtype=tf.bool),
                shape=row_score.shape,
            )
            row_score = tf.where(banned_tokens_mask, -float("inf"), row_score)
            return row_score

        scores = tf.map_fn(_get_row_updated_score, (input_ids, scores), fn_output_signature=tf.float32)
        return scores