Views
No views yet
1import os
2from transformers import AutoTokenizer, AutoModelForCausalLM
3import torch
4
5os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
6
7
8# 设置模型和数据路径
9MODEL_PATH = "Fintech-Dreamer/FinSynth_model_fraud"
10
11
12def generate_response(model, tokenizer, instruction, input_text, max_length=2048):
13 """
14 使用模型生成回答
15
16 参数：
17 model: 加载的语言模型实例
18 tokenizer: 模型对应的分词器
19 instruction: 指令部分文本，一般是任务描述
20 input_text: 输入文本，一般是需要分析的内容
21 max_length: 生成文本的最大长度，默认为2048个token
22
23 返回：
24 prompt: 完整的输入提示词
25 response: 模型生成的回答
26 """
27 # 构造提示词格式 - 使用特殊标记组织对话形式
28 # <｜begin of sentence｜>标记句子开始，<｜User｜>和<｜Assistant｜>分别标记用户和助手角色
29 prompt = f"<｜begin of sentence｜><｜User｜>{instruction}\n{input_text}<｜Assistant｜>"
30
31 # 编码输入，将文本转换为模型可以理解的token序列
32 inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True)
33 # 将输入移动到模型所在的设备（CPU或GPU）
34 inputs = inputs.to(model.device)
35
36 # 使用torch.no_grad()避免计算梯度，节省内存并加速推理过程
37 with torch.no_grad():
38 # 调用模型的generate方法生成回答
39 outputs = model.generate(
40 **inputs,
41 max_length=max_length, # 设置生成文本的最大长度
42 num_return_sequences=1, # 只返回一个生成序列
43 do_sample=True, # 使用采样策略，增加多样性
44 temperature=0.6, # 温度参数，控制生成文本的随机性（较低的值使输出更确定）
45 top_p=0.95, # 使用nucleus sampling，只考虑概率和超过0.95的token
46 pad_token_id=tokenizer.eos_token_id, # 将填充标记设置为结束标记
47 use_cache=True, # 使用缓存加速生成过程
48 )
49
50 # 将生成的token序列解码为文本
51 response = tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
52
53 # 提取回答部分（去除提示词部分）
54 if "<｜Assistant｜>" in response:
55 response = response.split("<｜Assistant｜>")[1].strip()
56
57 return prompt, response
58
59
60def process_test_data():
61 """
62 处理测试数据集并生成预测结果
63
64 功能：
65 - 加载测试数据集
66 - 初始化模型和分词器
67 - 对每个测试样本进行预测
68 - 输出预测结果
69
70 返回：
71 None，结果直接打印
72 """
73 # 加载测试数据
74
75 # 加载模型和分词器
76 print(f"加载模型: {MODEL_PATH}")
77 # 加载预训练的分词器，用于将文本转换为token
78 tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
79 # 加载预训练的语言模型，设置为自动选择设备，使用bfloat16精度以提高性能
80 model = AutoModelForCausalLM.from_pretrained(
81 MODEL_PATH,
82 device_map="auto", # 自动选择可用的设备（CPU/GPU）
83 torch_dtype=torch.bfloat16, # 使用bfloat16精度，在保持准确性的同时减少内存占用
84 use_cache=True, # 启用缓存以提高生成速度
85 )
86
87 # 设置模型为评估模式，关闭dropout等训练特性，提高推理性能
88 model.eval()
89
90 # 处理每个测试样本
91 print("开始生成预测...")
92
93 try:
94 # 提取指令和输入文本
95 instruction = "Combine your financial knowledge to carefully analyze whether this is a fraudulent transaction，just answer yes or no" # 指令部分，告诉模型要执行的任务
96 input_text = "The transaction involves a 20-year-old customer with an income of $0.1, who has been at their current address for 46 months. The request was made 0.03 days ago, and the intended amount is $48.49. The payment type is 'AA', and the email used is from a free provider. The customer's bank account is 6 months old, and they do not have other credit cards. The proposed credit limit is $1500, and the request was made via the internet. The session lasted 3.36 minutes, and the device used had 1 distinct email in the last 8 weeks. The device has no prior fraud records." # 输入文本，需要分析的内容
97
98 print("\n正在生成预测...")
99
100 # 调用generate_response函数生成预测
101 full_prompt, response = generate_response(model, tokenizer, instruction, input_text)
102
103 # 打印完整的预测结果，不截断
104 print("\n===== 完整输入输出 =====")
105 print(f"提示词: {full_prompt}")
106 print(f"\n预测结果: {response}")
107
108 # 简要展示关键结果
109 print("\n===== 简要结果 =====")
110 # 简单的欺诈判断逻辑：检查输出中是否包含"fraudulent"且不包含"not"
111 print(f"预测标签: {'欺诈' if 'fraudulent' in response.lower() and 'not' not in response.lower() else '非欺诈'}")
112 print(f"输出长度: {len(response)} 字符")
113
114 # 以JSON格式输出结果，便于后续处理或保存
115 print("\n===== JSON格式 =====")
116 result_json = {"prompt": full_prompt, "predict": response}
117 print(result_json)
118
119 except Exception as e:
120 # 异常处理，确保一个样本的错误不会导致整个程序崩溃
121 print(f"\n处理样本时出错: {str(e)}")
122 import traceback
123
124 traceback.print_exc() # 打印详细错误信息，便于调试
125
126 print("\n预测完成！")
127 return None
128
129
130def main():
131 """
132 主函数，程序入口点
133
134 功能：
135 - 启动测试数据处理和预测流程
136 """
137 print("===== 模型调用 =====")
138 process_test_data()
139
140
141if __name__ == "__main__":
142 main()