Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2device = "cuda" # the device to load the model onto
3model = AutoModelForCausalLM.from_pretrained(
4 model_path, torch_dtype="auto", device_map="auto"
5)
6tokenizer = AutoTokenizer.from_pretrained(model_path)
7
8def build_user_query(question, pred_answer, answer, base_prompt):
9 input_text = base_prompt.replace("{{question}}", question)
10 input_text = input_text.replace("{{pred_step}}", pred_answer)
11 input_text = input_text.replace("{{answer}}", answer)
12 input_text = input_text.replace("{{analysis}}", "") # default set analysis to blank, if exist, you can pass in the corresponding parameter.
13 return input_text
14
15chat_prompt = """<|im_start|>system
16You are a helpful assistant.<|im_end|>
17<|im_start|>human
18{}<|im_end|>
19<|im_start|>gpt
20"""
21
22basic_prompt = """## 任务描述\n\n你是一个数学老师，学生提交了题目的解题步骤，你需要参考`题干`，`解析`和`答案`，判断`学生解题步骤`的结果是否正确。忽略`学生解题步骤`中的错误，只关注最后的答案。答案可能出现在`解析`中，也可能出现在`答案`中。\n\n## 输入内容\n\n题干:\n\n```\n{{question}}\n```\n\n解析:\n\n```\n{{analysis}}\n\n```\n\n答案:\n\n```\n{{answer}}\n```\n\n学生解题步骤:\n\n```\n{{pred_step}}\n```\n\n输出:"""
23base_prompt = chat_prompt.format(basic_prompt)
24
25def build_user_query(question, pred_answer, answer, base_prompt):
26 input_text = base_prompt.replace("{{question}}", question)
27 input_text = input_text.replace("{{pred_step}}", pred_answer)
28 input_text = input_text.replace("{{answer}}", answer)
29 input_text = input_text.replace("{{analysis}}", "") # default set analysis to blank, if exist, you can pass in the corresponding parameter.
30 return input_text
31
32prompt = build_user_query("1+1=", "3", "2", base_prompt)
33
34model_inputs = tokenizer([prompt], return_tensors="pt").to(device)
35
36generated_ids = model.generate(model_inputs.input_ids, temperature=0, max_new_tokens=16, eos_token_id=100005)
37generated_ids = [
38 output_ids[len(input_ids) :]
39 for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
40]
41
42response = tokenizer.batch_decode(generated_ids, skip_special_tokens=False)[0]