Views
No views yet



eagenerate function for accelerated generation, similar to the generate method in Hugging Face Transformers.1import torch
2from model.ea_model_griffin import EaModel
3from fastchat.model import get_conversation_template
4
5# Replace with your actual model paths
6base_model_path = "meta-llama/Llama-3-8B-Instruct" # Example base model
7EAGLE_model_path = "husj576/GRIFFIN-llama3-instruct-8B" # Example GRIFFIN draft model
8
9model = EaModel.from_pretrained(
10 base_model_path=base_model_path,
11 ea_model_path=EAGLE_model_path,
12 torch_dtype=torch.float16,
13 low_cpu_mem_usage=True,
14 device_map="auto",
15 total_token=-1
16)
17model.eval()
18
19your_message="Hello"
20conv = get_conversation_template("llama3") # Use the correct conversation template for your base model
21conv.append_message(conv.roles[0], your_message)
22conv.append_message(conv.roles[1], None)
23prompt = conv.get_prompt()
24
25input_ids=model.tokenizer([prompt]).input_ids
26input_ids = torch.as_tensor(input_ids).cuda()
27output_ids=model.eagenerate(input_ids,temperature=0.5,max_new_tokens=512)
28output=model.tokenizer.decode(output_ids[0])
29
30print(output)