Created
July 10, 2024 14:57
-
-
Save gante/e0cdcdf5438a0ae24390b2aee6e5852e to your computer and use it in GitHub Desktop.
DoLa demo
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 @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
im just copy paste your code and don't work: