How to implement beam search on logits output from BART model?
Author: ZiyueWangUoB
Pull request: no

Hi,

I've converted my BART-LARGE-CNN model to ONNX and now am trying to perform beam search then decoding to generate a summary. The output from logits layer is in the shape of [1,804,50265] which I'm assuming is [batch size, time step, vocab length]. I then perform softmax on each time step and log the result.


If my understanding is correct, I can perform greedy search easily with this (at each time step, select the highest probability, then decode using vocab to find the summary). But how can I implement beam search? For example, given the choice of a word (i.e. "The) at time step 1, the probability of each word following that at time step 2 should be different depending on the word selected at time step 1. But in this case the logits are fixed for every different word at time step 2?

Sorry I'm new to this, pretty confused. Any help would be appreciated. 


Comments:
Gently pinging @lewtun, and I'm unknowledgeable about the intersection between ONNX and text generation :)

@ZiyueWangUoB yeah, you can implement the basic version of greedy search as you described. Beam search is more complex, a good reference is the following [blog post](https://huggingface.co/blog/how-to-generate). But I suspect we already have a solution for ONNX models!

@gante Yes I've read that article and understand the theory behind beam search. However I feel like I'm missing something with the onnx output, as the logits alone shouldn't be able to cover the beam search algorithm. 

@ZiyueWangUoB There are several ways to kickstart beam search, but in all of them you have to do shenanigans at the start to obtain `N` (number of beams) sets of logits from a single input row.
Option 1 -  The first iteration is a normal greedy search where you keep the top_k (k=`N`) tokens
Option 2 - You replicate your input `N` times, but set a large score penalty in all but the first row. You can use beam search from the first iteration.

From there, run the usual beam search: obtain `[N, vocab_size]` logits, pick the top `N` based on the score. Following our code, especially our TF and JAX implementations, is great to understand everything that goes into making it right!

Note: As per our [issues guidelines](https://github.com/huggingface/transformers/blob/main/ISSUES.md), we reserve GitHub issues for bugs in the repository and/or feature requests. For any other matters (like these questions), we'd like to invite you to use our [forum](https://discuss.huggingface.co/) 🤗

@gante how would I go about getting N sets of logits? The model only outputs a [1,804,50265] set, which I’m assuming is 1 set of logits.

Thanks for the ping @gante !

@ZiyueWangUoB we actually have a beam search example with BART + TorchScript from the ONNX team that you can inspect here: https://github.com/huggingface/transformers/blob/main/examples/research_projects/onnx/summarization/bart_onnx/generation_onnx.py

tl;dr implementing beam search from scratch is quite involved and essentially requires reimplementing large chunks of the `generate()` method we have in `transformers`

If ONNX Runtime is an option, an alternative would be to run the generation with our `optimum` lib:

```python
from transformers import AutoTokenizer, pipeline
from optimum.onnxruntime import ORTModelForSeq2SeqLM

model_id = "facebook/bart-large-cnn"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = ORTModelForSeq2SeqLM.from_pretrained(model_id, from_transformers=True)
summarizer = pipeline("summarization", model=model, tokenizer=tokenizer)

text = "The tower is 324 metres (1,063 ft) tall, about the same height as an 81-storey building, and the tallest structure in Paris. Its base is square, measuring 125 metres (410 ft) on each side. During its construction, the Eiffel Tower surpassed the Washington Monument to become the tallest man-made structure in the world, a title it held for 41 years until the Chrysler Building in New York City was finished in 1930. It was the first structure to reach a height of 300 metres. Due to the addition of a broadcasting aerial at the top of the tower in 1957, it is now taller than the Chrysler Building by 5.2 metres (17 ft). Excluding transmitters, the Eiffel Tower is the second tallest free-standing structure in France after the Millau Viaduct."
summarizer(text)
```


@lewtun Thanks for the input! I've been looking at that example, but the problem with directly using that script is the result means the length of the encoded text is fixed. I'll look into the torchscript beam search further, but as of right now that can't be directly used.

As for optimum, I will try that. I'm currently trying to convert the model into Tensorrt at the end, but if that's not reasonable I will use optimum. 

> As for optimum, I will try that. I'm currently trying to convert the model into Tensorrt at the end, but if that's not reasonable I will use optimum.

Cool! We're hoping to integrate TensorRT as a backend in `optimum` when we get some bandwidth - in the meantime, you're more than welcome to open an issue/PR if you feel so inclined :)