Views
No views yet
transformers:1import os
2
3DEVICE_ID = "0"
4os.environ['CUDA_VISIBLE_DEVICES'] = DEVICE_ID
5from transformers import AutoModelForCausalLM, AutoTokenizer
6
7model_name = "mjoys/jinpan-13B"
8device = "cuda" # the device to load the model onto
9
10model = AutoModelForCausalLM.from_pretrained(
11 model_name,
12 torch_dtype="auto",
13 device_map="auto"
14)
15tokenizer = AutoTokenizer.from_pretrained(model_name)
16
17prompt = "怎么办理信用卡"
18messages = [
19 {"role": "system", "content": "You are a helpful assistant. 你是一个乐于助人的助手。"},
20 {"role": "user", "content": prompt}
21]
22text = tokenizer.apply_chat_template(
23 messages,
24 tokenize=False,
25 add_generation_prompt=True
26)
27model_inputs = tokenizer([text], return_tensors="pt").to(device)
28
29generated_ids = model.generate(
30 **model_inputs,
31 max_new_tokens=512
32)
33generated_ids = [
34 output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
35]
36
37response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
38print(response)