Views
No views yet
1from transformers import LlamaForCausalLM, LlamaTokenizer
2import torch
3
4def generate_prompt(query, history, input=None):
5 prompt = ""
6 for i, (old_query, response) in enumerate(history):
7 prompt += "{}{}\n<end>".format(old_query, response)
8 prompt += "{}".format(query)
9 return prompt
10
11
12
13# Load model
14device = torch.device("cuda:0")
15model_name = "/data/dell/xuyipei/my_llama/my_llama_13b/llama_13b_112_sft_v1"
16tokenizer = LlamaTokenizer.from_pretrained(model_name)
17model = LlamaForCausalLM.from_pretrained(
18 model_name,
19 torch_dtype=torch.float16
20)
21model.eval()
22model = model.to(device)
23
24
25# Inference
26history = []
27queries = ['请推荐五本名著，依次列出作品名、作者\n', '请再来三本\n']
28memory_limit = 3 # the number of (query, response) to remember
29for query in queries:
30 prompt = generate_prompt(prompt, history)
31 input_ids = tokenizer(query, return_tensors="pt", padding=False, truncation=False, add_special_tokens=False)
32 input_ids = input_ids["input_ids"].to(device)
33
34 with torch.no_grad():
35 outputs=model.generate(
36 input_ids=input_ids,
37 top_p=0.8,
38 top_k=50,
39 repetition_penalty=1.1,
40 max_new_tokens = 256,
41 early_stopping = True,
42 eos_token_id = tokenizer.convert_tokens_to_ids('<end>'),
43 pad_token_id = tokenizer.eos_token_id,
44 min_length = input_ids.shape[1] + 1
45 )
46 s = outputs[0]
47 response=tokenizer.decode(s)
48 response = response.replace('<s>', '').replace('<end>', '').replace('</s>', '')
49 print(response)
50 history.append((query, response))
51 history = history[-memory_limit:]