Generate: assisted decoding with sample
Author: gante
Pull request: yes

# What does this PR do?

This PR expands the previous [assisted generation PR](https://github.com/huggingface/transformers/pull/22211) so as to work with sampling.

Two important notes to review the PR:
1. I'd suggest starting the review by the docs, so you understand what's going on at a high level. Sampling adds an additional (controllable) heuristic, so the user can control between speed and pure sampling behavior. 
2. In terms of implementation, I've decided to overload the assisted generation function with a few extra lines to handle the sample case. This is to avoid adding a close copy of a 500-line function.

_____________________________________________________________________________

Bellow are some results, so you can understand the balancing act. Execution time obtained on a 3090.

<details>
  <summary>Script</summary>

  ```py
  from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
  import torch
  import time
  
  model_id = "EleutherAI/pythia-6.9b-deduped"
  assistant_id = "EleutherAI/pythia-160m-deduped"
  
  tokenizer = AutoTokenizer.from_pretrained(model_id)
  
  assistant_model = AutoModelForCausalLM.from_pretrained(assistant_id)
  assistant_model = assistant_model.to("cuda")
  
  model_kwargs = {
      "pretrained_model_name_or_path": model_id,
      "device_map": "auto",
      "max_memory": {0: "20GiB", "cpu": "50GiB"},
      "torch_dtype": torch.float16,
  }
  model = AutoModelForCausalLM.from_pretrained(**model_kwargs)
  
  inputs = tokenizer("Here's how to cook a good ramen:", return_tensors="pt").to("cuda")
  
  streamer = TextStreamer(tokenizer=tokenizer)
  
  print("Greedy with assistance:")
  start = time.time()
  model.generate(**inputs, assistant_model=assistant_model, streamer=streamer, max_new_tokens=64)
  print(f"Elapsed time: {time.time() - start:.2f} seconds")
  
  for p in (0.0, 0.2, 0.4, 0.6, 0.8, 1.0):
      print(f"Sample with assistance (assisted_keep_proba = {p})")
      torch.manual_seed(0)
      start = time.time()
      model.generate(
          **inputs,
          do_sample=True,
          assistant_model=assistant_model,
          assisted_keep_proba=p,
          streamer=streamer,
          max_new_tokens=64
      )
      print(f"Elapsed time: {time.time() - start:.2f} seconds")
  
  print("Original sample")
  torch.manual_seed(0)
  start = time.time()
  model.generate(**inputs, do_sample=True, streamer=streamer, max_new_tokens=64)
  print(f"Elapsed time: {time.time() - start:.2f} seconds")
  ```
</details>
<details>
  <summary>Sample results</summary>

  

  Decoding strategy   |  Result  |  Execution time
  :-------------------:|:-------:|:------:|
  Greedy (w/assistance) | Here's how to cook a good ramen:<br><br>1. Make sure you have a good stock.<br><br>2. Make sure you have a good broth.<br><br>3. Make sure you have a good ramen.<br><br>4. Make sure you have a good ramen.<br><br>5. Make sure you have a good ramen. | 1.44 seconds
  Sample (w/assistance<br>`assisted_keep_proba=0.0`) | Here's how to cook a good ramen:<br><br>1. Get a noodle.<br><br>2. Get a stock.<br><br>3. Get a packet of dried ingredients.<br><br>4. Cook the noodles.<br><br>5. Cook the stock.<br><br>6. Cook the packet of dried ingredients.<br><br>7. Enjoy!<br><br>And | 1.44 seconds
  Sample (w/assistance<br>`assisted_keep_proba=0.2`) | Here's how to cook a good ramen:<br><br>1. Get a noodle vendor.<br><br>The noodle vendor makes the noodles. Japanese restaurants often have the noodle vendor on-site.<br><br>2. Get a pot.<br><br>The pot is used to cook ramen.<br><br>3. Get a pot of boiling water. | 1.59 seconds
  Sample (w/assistance<br>`assisted_keep_proba=0.4`) | Here's how to cook a good ramen:<br><br>Step 1: Collect your ingredients.<br><br>For this recipe you need a big stock pot. That's good.<br><br>And some water.<br><br>Step 2: Peel the eggs.<br><br>Yes, that's it. Four eggs.<br><br>Step 3: Separate the yolks. | 1.71 seconds
  Sample (w/assistance<br>`assisted_keep_proba=0.6`) | Here's how to cook a good ramen:<br><br>Nothing much to take out of the packet. Just a big block of pork fat, some Chinese chilli paste and seasonings.<br><br>Preheat the oven to 210ºC (410ºF/Gas 6).<br><br>Place the pork fat, chilli paste and seasoning into a mixing bowl and | 2.08 seconds
  Sample (w/assistance<br>`assisted_keep_proba=0.8`) | Here's how to cook a good ramen:<br><br>**You'll need:** A large pot for boiling noodles<br>A small saucepan for cooking the noodles<br>BBQ chicken or roasted fish, or any grilled healthy protein<br>A box of ramen noodles, noodles that come in<br>shapes and sizes<br>Soups or broth, | 2.32 seconds
  Sample (w/assistance<br>`assisted_keep_proba=1.0`) | Here's how to cook a good ramen:<br><br>You take your pre-scalloped noodles, pour boiling water (or your preferred water-to-noodle ratio) over them, and leave them alone for four to five minutes. Once that's done, drain them, season with salt, and heat them up on the stove (microwave won | 2.56 seconds
  Original Sample) | Here's how to cook a good ramen:<br><br>You take your pre-scalloped noodles, pour boiling water (or your preferred cooking liquid) over it, and after that you go get your ramen broth, add-ins, and other condiments. You make your seasoning sauce, and heat that up. Mix it all together, and put | 2.05 seconds

  As it can be seen above, there is a trade off between time and quality. This will certainly be application specific: factual applications will be able to take the most of assisted decoding. In my brief experiments, `assisted_keep_proba=0.3` seems like a sensible default.

</details>



Comments:
_The documentation is not available anymore as the PR was closed or merged._

I'm closing this PR because I found a much much better way to handle the sample case 🧠 

Stay tuned 🚀