-
Notifications
You must be signed in to change notification settings - Fork 18
/
mixtral_example.py
64 lines (53 loc) · 2.47 KB
/
mixtral_example.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
# -*- coding: utf-8 -*-
"""
Copyright (c) Ant Financial Service Group and its affiliates.
"""
import time
import torch
from transformers import AutoTokenizer, BitsAndBytesConfig, AutoConfig
from pia.lookahead.models.mixtral.modeling_mixtral import MixtralForCausalLM
from pia.lookahead.examples import local_path_dict
# only worker with transformers>=4.36.0
model_dir = local_path_dict.get('mixtral', 'your/model/path')
tokenizer = AutoTokenizer.from_pretrained(model_dir)
# note: this model cannot be fully loaded into a A100, refer to mixtral_quant_example.py for solution
model = MixtralForCausalLM.from_pretrained(model_dir
, cache_dir='./'
, torch_dtype=torch.float16
, low_cpu_mem_usage=True
, device_map={"":"cuda:0"})
prompt = "Hello, I'm am conscious and"
inputs = tokenizer(prompt, return_tensors="pt")
input_ids = inputs.input_ids.cuda()
attention_mask = inputs.attention_mask.cuda()
position_ids = None
for use_lookahead in [False, False, True, True]:
debug_lookahead = False
decoding_length = 64
branch_length = 12
ts = time.time()
max_new_tokens = 256
decoding_kwargs = {"use_lookahead": use_lookahead,
"debug_lookahead": debug_lookahead,
"decoding_mode": 'hier',
"decoding_length": decoding_length,
"branch_length": branch_length}
outputs = model.generate(input_ids=input_ids,
attention_mask=attention_mask,
position_ids=position_ids,
pad_token_id=tokenizer.eos_token_id,
eos_token_id=tokenizer.eos_token_id,
use_cache=True,
max_new_tokens=max_new_tokens,
repetition_penalty=1.0,
do_sample=False,
decoding_kwargs=decoding_kwargs
)
output_ids = outputs
input_length = input_ids.size(-1)
output_ids = output_ids[0, input_length:].tolist()
response = tokenizer.decode(output_ids)
input_text = tokenizer.decode(input_ids[0])
te = time.time()
token_count = len(output_ids)
print(f'lookahead:{use_lookahead} time:{te - ts:.3f}s speed:{token_count/(te-ts):.1f}token/s response:\n{response}\n')