Skip to content

Instantly share code, notes, and snippets.

@gante
Created July 10, 2024 14:57
Show Gist options
  • Save gante/e0cdcdf5438a0ae24390b2aee6e5852e to your computer and use it in GitHub Desktop.
Save gante/e0cdcdf5438a0ae24390b2aee6e5852e to your computer and use it in GitHub Desktop.
DoLa demo
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Meta-Llama-3-8B-Instruct", torch_dtype=torch.bfloat16, device_map="auto"
)
model.generation_config.eos_token_id = model.generation_config.pad_token_id
question = 'What does Darth Vader say to Luke in "The Empire Strikes Back"?'
text = f"Answer with a short answer.\n\nQuestion: {question}\n\nAnswer: "
inputs = tokenizer(text, return_tensors="pt").to(model.device)
generate_kwargs={
"do_sample": False, "max_new_tokens": 40, "top_p": None, "temperature": None
}
# Vanilla: gets the correct quote, but it's wrong about the misquote :(
vanilla_output = model.generate(**inputs, **generate_kwargs)
print("\n" + tokenizer.decode(vanilla_output[0, inputs.input_ids.shape[-1]:], skip_special_tokens=True))
# "No, I am your father." (Note: This is a famous misquote. The actual quote is "No, I am your father"
# is not in the movie. The correct quote is
# Vanilla: gets the quote and the misquote right!
dola_output = model.generate(**inputs, **generate_kwargs, dola_layers='high', repetition_penalty=1.2)
print("\n" + tokenizer.decode(dola_output[0, inputs.input_ids.shape[-1]:], skip_special_tokens=True))
# "No, I am your father." (Note: This is one of the most famous lines in movie history, and it's often misquoted as
# "Luke, I am your father.")
@Santiagowwi
Copy link

im just copy paste your code and don't work:
image

@andysingal
Copy link

im just copy paste your code and don't work: image

got the same error: https://colab.research.google.com/drive/16szW7upSw4ITmG0kN5LGGA_aGYcMalkZ?usp=sharing

@gante
Copy link
Author

gante commented Jul 11, 2024

@Santiagowwi @andysingal you need to install from main, i.e. start your script with

!pip install --upgrade git+https://github.com/huggingface/transformers.git

This feature will be part of v4.43, after it is released you don't need this special installation instruction 🤗

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment