Refine the code of Universal Assisted Generation

[xinpengzz]

What does this PR do?

co-Authors: @WenjingKangIntel

Universal Assisted Generation(UAG) allowed the assisted generation with any assistant model regardless of its tokenizer.
This work is awesome and makes the different models work together possible !

This PR is to refine its code.
The code of the Universal Assisted Generation (UAG) contains some useless attributes and some processes that need improvement.

I have contacted with the @danielkorat, the one of authors of UAG, about it, and he suggested me to create a pull request to discuss it conveniently and detailedly.

Code changes:

• removed the useless attributes: self.prev_tokens, self.prev_target_ids.
• added a judge statement when the start_assistant_look_index < 0.

Todo:

• the size of window (the number of past tokens) is an experience value. Perhaps this value could be set to be adjustable, or dynamic

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@danielkorat
@gante
and anyone interested in the program 😄

[danielkorat]

cc @mosheber

[xinpengzz]

The usage:

from transformers import AutoModelForCausalLM, AutoTokenizer

prompt = "Alice and Bob"
checkpoint = "google/gemma-2-9b"
assistant_checkpoint = "double7/vicuna-68m"

assistant_tokenizer = AutoTokenizer.from_pretrained(assistant_checkpoint)
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
inputs = tokenizer(prompt, return_tensors="pt")

model = AutoModelForCausalLM.from_pretrained(checkpoint)
assistant_model = AutoModelForCausalLM.from_pretrained(assistant_checkpoint)

# assistant_model.generation_config.assistant_lookbehind = 10 # the default value
# assistant_model.generation_config.target_lookbehind = 10 # the default value

outputs = model.generate(**inputs, assistant_model=assistant_model, tokenizer=tokenizer, assistant_tokenizer=assistant_tokenizer)
tokenizer.batch_decode(outputs, skip_special_tokens=True)

[danielkorat]

Hi @ArthurZucker @zucchini-nlp 👋

[zucchini-nlp]

Thanks!

[xinpengzz]

I have tested these changes on 2000 samples in the alpaca-cleaned dataset, and the output is identical to those without assistant decoding. main model is Qwen1.5-0.5B and draft model is vicuna-68m.

[xinpengzz]

Hi @ArthurZucker 👋 , Would it be convenient to review the code?

[ArthurZucker]

Hey yeah sorry reviewing now!

[ArthurZucker]

Nice and simple, thanks for adding the arg descriptions! 🤗