-
Notifications
You must be signed in to change notification settings - Fork 3
/
inference_ft.py
52 lines (47 loc) · 1.99 KB
/
inference_ft.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
from transformers import LlamaForCausalLM, LlamaTokenizer
from peft import PeftModel
import torch
overall_instruction = "你是复旦大学知识工场实验室训练出来的语言模型CuteGPT。给定任务描述,请给出对应请求的回答。\n"
def generate_prompt(query, history, input=None):
prompt = overall_instruction
for i, (old_query, response) in enumerate(history):
# 多轮对话需要跟训练时保持一致
prompt += "问:{}\n答:\n{}\n".format(old_query, response)
prompt += "问:{}\n答:\n".format(query)
return prompt
model_name = "XuYipei/kw-cutegpt-13b-ift"
tokenizer = LlamaTokenizer.from_pretrained(model_name)
print('additional_special_token:', tokenizer.additional_special_tokens)
model = LlamaForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16,
device_map="auto",
)
model.eval()
device = torch.device("cuda")
history = []
queries = ['请推荐五本名著,依次列出作品名、作者','再来三本呢?']
memory_limit = 3 # the number of (query, response) to remember
for query in queries:
prompt = generate_prompt(query, history)
print(prompt)
input_ids = tokenizer(prompt, return_tensors="pt", padding=False, truncation=False, add_special_tokens=False)
input_ids = input_ids["input_ids"].to(device)
with torch.no_grad():
outputs=model.generate(
input_ids=input_ids,
top_p=0.8,
top_k=50,
repetition_penalty=1.1,
max_new_tokens = 256,
early_stopping = True,
eos_token_id = tokenizer.convert_tokens_to_ids('<end>'),
pad_token_id = tokenizer.eos_token_id,
min_length = input_ids.shape[1] + 1
)
s = outputs[0][input_ids.shape[1]:]
response=tokenizer.decode(s)
response = response.replace('<s>', '').replace('<end>', '').replace('</s>', '')
print(response)
history.append((query, response))
history = history[-memory_limit:]