Skip to content

Commit

Permalink
Update generate example to llama 3 (#448)
Browse files Browse the repository at this point in the history
  • Loading branch information
casper-hansen authored Apr 18, 2024
1 parent 0e52a5c commit edcf780
Showing 1 changed file with 18 additions and 10 deletions.
28 changes: 18 additions & 10 deletions examples/generate.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,28 +2,36 @@
from transformers import AutoTokenizer, TextStreamer


quant_path = "TheBloke/Mistral-7B-Instruct-v0.2-AWQ"
quant_path = "casperhansen/llama-3-8b-instruct-awq"

# Load model
model = AutoAWQForCausalLM.from_quantized(quant_path, fuse_layers=True)
tokenizer = AutoTokenizer.from_pretrained(quant_path, trust_remote_code=True)
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

# Convert prompt to tokens
prompt_template = "[INST] {prompt} [/INST]"

prompt = "You're standing on the surface of the Earth. "\
"You walk one mile south, one mile west and one mile north. "\
"You end up exactly where you started. Where are you?"

tokens = tokenizer(
prompt_template.format(prompt=prompt),
return_tensors='pt'
).input_ids.cuda()
chat = [
{"role": "system", "content": "You are a concise assistant that helps answer questions."},
{"role": "user", "content": prompt},
]

terminators = [
tokenizer.eos_token_id,
tokenizer.convert_tokens_to_ids("<|eot_id|>")
]

tokens = tokenizer.apply_chat_template(
chat,
return_tensors="pt"
).cuda()

# Generate output
generation_output = model.generate(
tokens,
streamer=streamer,
max_new_tokens=512
)
max_new_tokens=64,
eos_token_id=terminators
)

0 comments on commit edcf780

Please sign in to comment.