Issue with SynthID watermark implementation
Author: marluxiaboss
Pull request: no
Labels: name: Feature request

name: Watermark

### Feature request

Caveat: I wasn't sure if should have formulated it as a bug or a feature request.

The SynthID watermark implementation on transformers (v4.46.0) #34350 is inconsistent with the original implementation from https://github.com/google-deepmind/synthid-text. The issue is related to the order in which LogitsProcessor and LogitsWarper objects are applied from the processors list created in the `_get_logits_processor` function and used in the `_sample` function from `transformers/blob/main/src/transformers/generation/utils.py`.

I suggest that the overall logic for the order in which these processor objects are applied should be revisited when providing a `WatermarkingConfig` or a list of custom `LogitsProcessor` to the `generate` function.

### Motivation

The original implementation of SynthID from DeepMind's repo (https://github.com/google-deepmind/synthid-text/blob/main/src/synthid_text/synthid_mixin.py) modifies the `_sample` function from `transformers/blob/main/src/transformers/generation/utils.py` such that it applies the logits processors and logits warpers first before applying the watermarking.

```python
     # pre-process distribution
      next_token_scores = logits_processor(input_ids, next_token_logits)
      indices_mapping = None
      unwatermarked_scores = None
      if do_sample:
        *regular_warpers, watermarking_logits_warper = logits_warper
        if not isinstance(
            watermarking_logits_warper,
            logits_processing.SynthIDLogitsProcessor,
        ):
          raise ValueError(
              "SynthIDLogitsProcessor should be the final warper in the list"
              " while watermarking."
          )
        for logit_warper in regular_warpers:
          next_token_scores = logit_warper(input_ids, next_token_scores)
        # Watermark final scores with sparse top_k.
        next_token_scores, indices_mapping, unwatermarked_scores = (
            watermarking_logits_warper.watermarked_call(
                input_ids, next_token_scores
            )
        )
```

In the transformer's implementation of the `_sample` function (`transformers/blob/main/src/transformers/generation/utils.py`), a watermarking scheme is always applied after the logits processors but before the logits warpers – same for the custom logits processors provided as parameters to the `generate` function:

```python
        if generation_config.watermarking_config is not None:
            processors.append(
                generation_config.watermarking_config.construct_processor(self.config.vocab_size, device)
            )

        # TODO (joao): find a strategy to specify the order of the processors
        processors = self._merge_criteria_processor_list(processors, logits_processor)

        # Processors previously known as `LogitsWarpers`, only applied with sampling strategies
        if generation_config.do_sample:
            # In beam methods, we need to keep at least one non-eos token to explore continuations that might have a
            # better score (i.e. keep len(list(generation_config._eos_token_tensor)) + 1)
            if generation_config.num_beams > 1:
                if isinstance(generation_config._eos_token_tensor, list):
                    min_tokens_to_keep = len(generation_config._eos_token_tensor) + 1
                elif isinstance(generation_config._eos_token_tensor, torch.Tensor):
                    min_tokens_to_keep = generation_config._eos_token_tensor.shape[0] + 1
                else:
                    min_tokens_to_keep = 2
            else:
                min_tokens_to_keep = 1

            # the following idea is largely copied from this PR: https://github.com/huggingface/transformers/pull/5420/files
            # all samplers can be found in `generation_utils_samplers.py`
            if generation_config.temperature is not None and generation_config.temperature != 1.0:
                processors.append(TemperatureLogitsWarper(generation_config.temperature))
            if generation_config.top_k is not None and generation_config.top_k != 0:
                processors.append(
                    TopKLogitsWarper(top_k=generation_config.top_k, min_tokens_to_keep=min_tokens_to_keep)
                )
            if generation_config.top_p is not None and generation_config.top_p < 1.0:
                processors.append(
                    TopPLogitsWarper(top_p=generation_config.top_p, min_tokens_to_keep=min_tokens_to_keep)
                )

```

Whether the ordering is critical or not for SynthID is left as a question for the authors of the paper.



### Your contribution

I cannot provide an immediate solution, as it would require a major redesign of part of the generation logic, but I would like to have the opinion of the authors from Deepmind and the Hugging Face team.
@gante @zucchini-nlp @SunMarc (and @sumedhghaisas2 for the Deepmind team).

Thanks for your input and your help!

Comments:
@marluxiaboss That you for opening this issue. I can confirm this is indeed a problem, the watermarking logits processor should be applied as the last logits processor and I can send that PR in short time. The misalignment between might have happened as logits_processor - warper merge and watermarking was happening at the same time. Apologies for this. 

cc @gante as well

@sumedhghaisas2 hey, is there a PR to fix it or would you like me to work on this? :)

@zucchini-nlp If you have time I would be very happy if you take it up. Let me know if you need any help on the solution.