src/transformers/generation/tf_logits_process.py/TFForceTokensLogitsProcessor/__call__
class TFForceTokensLogitsProcessor:    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
        def _force_token(generation_idx):
            batch_size = scores.shape[0]
            current_token = self.force_token_array[generation_idx]

            new_scores = tf.ones_like(scores, dtype=scores.dtype) * -float("inf")
            indices = tf.stack((tf.range(batch_size), tf.tile([current_token], [batch_size])), axis=1)
            updates = tf.zeros((batch_size,), dtype=scores.dtype)
            new_scores = tf.tensor_scatter_nd_update(new_scores, indices, updates)
            return new_scores

        scores = tf.cond(
            tf.greater_equal(cur_len, tf.shape(self.force_token_array)[0]),
            # If the current length is geq than the length of force_token_array, the processor does nothing.
            lambda: tf.identity(scores),
            # Otherwise, it may force a certain token.
            lambda: tf.cond(
                tf.greater_equal(self.force_token_array[cur_len], 0),
                # Only valid (positive) tokens are forced
                lambda: _force_token(cur_len),
                # Otherwise, the processor does nothing.
                lambda: scores,
            ),
        )
        return scores