TF: XLA bad words logits processor and list of processors

[gante]

What does this PR do?

This PR converts to XLA-compatible the bad_words logits processor. As per the discussion below, I was unable to convert the ngrams one -- added an exception and a TODO.

Also makes a change to the list of processors -- XLA raised issues when the processors had different arguments, so I had to add cur_len to all processors. After the change, the list wrapper is also compatible with XLA.

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[gante]

@Rocketknight1 @patrickvonplaten I'm stuck on the ngram logits processor, so I'd like to request your suggestions regarding what to try out next :D The bad words logits processor is ready and XLA-compatible.

Context:

1. Without XLA, it works well;
2. With XLA, yields incorrect outputs (it masks the wrong tokens in some cases). It is not a CPU/GPU thing -- it has the same output regardless of the hardware;
3. The XLA/non-XLA mismatch is at the output of _calc_row_banned_ngram_tokens, which gets the tokens that should be banned for each row;
4. All intermediary variables I was able to pull out had the same contents. However, if I try to pull out all ngrams, I get a core dumped on XLA 🤔

Things I've tried (without any symptom change):

1. The current implementation is a tf.while_loop with tf.TensorArray. On ddc8911, we can see my original implementation with a tf.map_fn (which is closer to the original code). Both versions have the exact same symptoms described above, and return the same errors for the same inputs when XLA is on (!);
2. Pulling the initialization of the tf.TensorArray to the start of __call__, pass ngram_size as an argument, and use tf.function as a decorator to __call__. The two first changes are to attempt a retrace trigger, the last one to rule out problems associated with attempting to compile a class instance (as opposed to a function);
3. Using tf.shape instead of tensor.shape, as the former is more suited for symbolic tensors;
4. Using batches with a single row as input;
5. Looking for other ways to implement the sliding window on the inputs (i.e. getting the ngrams), with no success.

[patrickvonplaten]

I'd be very much in favor of just not converting the ngram Processor. I don't think it's a necessary requirement to publish the new TF generate method. Let's maybe leave this as a hard second issue in case the community is very interested in this feature.

I think it's now more important to think about how to advertise, document XLA TF generate well and not loose too much time on this.

[patrickvonplaten]

Also not that many models use this processor (only know of BART and T5 for some summarization tasks)

[Rocketknight1]

Agree that it's not necessary to convert this one, but examining it, I suspect that there are some sneaky changes in output size depending on inputs, and XLA is struggling to deal with it. It seems very tough to convert to XLA, but if we decide we need it later let me know and I'll do my best to dig into it.

[gante]

Great 👍 I'm going to revert that one, add a TODO pointing at this PR, add a few final tests for the list of logits processors with XLA, and will ping you back.

[gante]

@Rocketknight1 @patrickvonplaten ready for review

[patrickvonplaten]

100% with leaving as is for now!

[patrickvonplaten]

Nice! Great job on getting no_bad_word_tokens to work

[Rocketknight1]

If performance is slow, allowing more parallel_iterations here might improve things, since this is a lightweight comparison run over a potentially large number of bad_words.

[Rocketknight1]

Overall this looks good, especially since the tests are present and passing! The need to filter an arbitrary number of words, each of which can span multiple tokens is very challenging to implement in XLA, so the fact that it's working at all seems almost miraculous, lol.