157 lines
5.7 KiB
Python
157 lines
5.7 KiB
Python
from openai import OpenAI
|
||
import os
|
||
import json
|
||
import random
|
||
import pandas as pd
|
||
import re
|
||
import sys
|
||
import sqlite3
|
||
|
||
import time
|
||
|
||
from sql_prompt import create_prompt_json
|
||
from sql_prompt import read_json_file
|
||
|
||
# 自动下载模型时，指定使用modelscope; 否则，会从HuggingFace下载
|
||
os.environ['VLLM_USE_MODELSCOPE']='True'
|
||
if __name__ == "__main__":
|
||
|
||
sys.stdout.reconfigure(encoding='utf-8')
|
||
|
||
# 连接到数据库
|
||
sort_json_path = "sort.json" # 分类信息文件
|
||
category_data = read_json_file(sort_json_path)
|
||
|
||
# 文件路径
|
||
table_name = 'data'
|
||
conn = sqlite3.connect(f'{table_name}.db') # 替换为您的数据库文件名
|
||
cursor = conn.cursor()
|
||
sort_name = '产品类型'
|
||
sort_value = ('住宿', '门票', '抢购')
|
||
num_limit = "500"
|
||
# SQL查询 - 随机选择500个指定产品类型的记录
|
||
query = f"""
|
||
SELECT * FROM data
|
||
WHERE {sort_name} IN {sort_value}
|
||
ORDER BY RANDOM()
|
||
LIMIT {num_limit}
|
||
"""
|
||
cursor.execute(query)
|
||
result = cursor.fetchall()
|
||
|
||
# 修改 OpenAI 的 API 密钥和 API 基础 URL 以使用 vLLM 的 API 服务器。
|
||
openai_api_key = "EMPTY"
|
||
openai_api_base = "http://localhost:8000/v1"
|
||
client = OpenAI(
|
||
api_key=openai_api_key,
|
||
base_url=openai_api_base,
|
||
)
|
||
|
||
#work_dir = './prompt'
|
||
|
||
results = []
|
||
|
||
for i, row in enumerate(result):
|
||
# 记录开始时间
|
||
loop_start_time = time.time()
|
||
|
||
prompt = create_prompt_json(row ,category_data) # 您的prompt文件路径
|
||
prompt_str = json.dumps(prompt, ensure_ascii=False) # 将字典转换为JSON字符串，保留中文字符
|
||
|
||
# 使用聊天接口
|
||
messages = [
|
||
{"role": "user", "content": prompt_str}
|
||
]
|
||
|
||
completion = client.chat.completions.create(
|
||
model="Qwen3-8B",
|
||
messages=messages
|
||
)
|
||
|
||
# 获取生成的文本
|
||
for choice in completion.choices:
|
||
generated_text = choice.message.content
|
||
|
||
# 确保文本使用UTF-8编码
|
||
if isinstance(generated_text, bytes):
|
||
generated_text = generated_text.decode('utf-8')
|
||
elif isinstance(generated_text, str):
|
||
# 如果是字符串，但编码可能不是UTF-8，先转为bytes再解码
|
||
try:
|
||
generated_text = generated_text.encode('latin-1').decode('utf-8')
|
||
except UnicodeError:
|
||
# 如果上面的转换失败，保持原样
|
||
pass
|
||
|
||
result_entry = []
|
||
result_entry.extend(row) # 将row的所有元素分别添加到result_entry
|
||
|
||
# 提取</think>后的内容作为最终回答
|
||
think_parts = generated_text.split("</think>")
|
||
if len(think_parts) > 1:
|
||
final_answer = think_parts[1].strip()
|
||
# 继续处理的文本改为最终回答部分
|
||
processing_text = final_answer
|
||
else:
|
||
# 如果没有</think>标记，使用原始文本
|
||
processing_text = generated_text
|
||
|
||
# 处理提取出的文本内容
|
||
try:
|
||
# JSON解析失败，使用正则表达式
|
||
patterns = {
|
||
"primary_category": r'primary_category["\s]*[:：]\s*["]*([^",\n}]+)',
|
||
"secondary_category": r'secondary_category["\s]*[:：]\s*["]*([^",\n}]+)',
|
||
"tertiary_category": r'tertiary_category["\s]*[:：]\s*["]*([^",\n}]+)',
|
||
"confidence": r'confidence["\s]*[:：]\s*([0-9.]+)',
|
||
"reasoning": r'reasoning["\s]*[:：]\s*["]*([^"}\n]+(?:\n[^"}\n]+)*)'
|
||
}
|
||
|
||
# 按顺序提取各个字段的值并追加到result_entry列表
|
||
for field, pattern in patterns.items():
|
||
match = re.search(pattern, processing_text, re.IGNORECASE)
|
||
if match:
|
||
# 确保匹配结果使用UTF-8编码
|
||
match_text = match.group(1).strip()
|
||
if isinstance(match_text, bytes):
|
||
match_text = match_text.decode('utf-8')
|
||
elif isinstance(match_text, str):
|
||
try:
|
||
match_text = match_text.encode('latin-1').decode('utf-8', errors='replace')
|
||
except UnicodeError:
|
||
pass
|
||
result_entry.append(match_text)
|
||
else:
|
||
result_entry.append("")
|
||
except Exception as e:
|
||
print(f"处理样本 {i} 时出错: {e}")
|
||
|
||
# 添加处理后的结果
|
||
results.append(result_entry)
|
||
|
||
# 计算并打印本次循环的执行时间
|
||
loop_end_time = time.time()
|
||
loop_duration = loop_end_time - loop_start_time
|
||
print(f"第 {i+1} 次模型推理时间: {loop_duration:.2f} 秒")
|
||
|
||
df = pd.DataFrame(results)
|
||
excel_file = '500_Qwen3_8B_sort.xlsx'
|
||
# 获取数据库表的列名作为DataFrame的表头
|
||
columns = [description[0] for description in cursor.description]
|
||
columns.append('primary_category')
|
||
columns.append('secondary_category')
|
||
columns.append('tertiary_category')
|
||
columns.append('confidence')
|
||
columns.append('reasoning')
|
||
# 添加index列
|
||
df.columns = columns
|
||
|
||
# 将第三列移到第一列
|
||
cols = df.columns.tolist()
|
||
third_col = cols.pop(2) # 第三列索引为2
|
||
cols.insert(0, third_col)
|
||
df = df[cols]
|
||
|
||
|
||
df.to_excel(excel_file, index=False)
|
||
print(f"\n所有生成文本已保存到 {excel_file}") |