Views
No views yet
1from transformers import LlamaForCausalLM, AutoTokenizer
2import torch
3
4ckpt = "tiansz/ChatYuan-7B-merge"
5device = torch.device('cuda')
6model = LlamaForCausalLM.from_pretrained(ckpt)
7tokenizer = AutoTokenizer.from_pretrained(ckpt)
8
9def answer(prompt):
10 prompt = f"用户：{prompt}\n小元："
11 input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
12 generate_ids = model.generate(input_ids, max_new_tokens=1024, do_sample = True, temperature = 0.7)
13 output = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
14 response = output[len(prompt):]
15 return response
16
17result = answer("你好")
18print(result)1from transformers import LlamaForCausalLM, AutoTokenizer
2import torch
3
4ckpt = "tiansz/ChatYuan-7B-merge"
5device = torch.device('cuda')
6max_memory = f'{int(torch.cuda.mem_get_info()[0]/1024**3)-1}GB'
7n_gpus = torch.cuda.device_count()
8max_memory = {i: max_memory for i in range(n_gpus)}
9model = LlamaForCausalLM.from_pretrained(ckpt, device_map='auto', load_in_8bit=True, max_memory=max_memory)
10tokenizer = AutoTokenizer.from_pretrained(ckpt)
11
12def answer(prompt):
13 prompt = f"用户：{prompt}\n小元："
14 input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
15 generate_ids = model.generate(input_ids, max_new_tokens=1024, do_sample = True, temperature = 0.7)
16 output = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
17 response = output[len(prompt):]
18 return response
19
20result = answer("你好")
21print(result)