124 lines
5.5 KiB
Python
124 lines
5.5 KiB
Python
from vllm import LLM, SamplingParams
|
||
from transformers import AutoTokenizer
|
||
import os
|
||
import json
|
||
import random
|
||
import pandas as pd
|
||
import re
|
||
|
||
# 读取您的prompt.json文件
|
||
def load_from_json(json_path):
|
||
with open(json_path, 'r', encoding='utf-8') as f:
|
||
data = json.load(f)
|
||
return data
|
||
|
||
# 从JSON文件构建提示词
|
||
def build_prompt_from_json(prompt_data):
|
||
# 根据您的JSON结构提取和组装内容
|
||
instruction = prompt_data.get("instruction", "")
|
||
task_description = prompt_data.get("task_description", "")
|
||
input_data = prompt_data.get("input", {})
|
||
output_requirements = prompt_data.get("output_requirements", {})
|
||
|
||
# 组装完整提示词
|
||
prompt = f"{instruction}\n\n{task_description}\n\n"
|
||
prompt += f"基本信息:\n{json.dumps(input_data.get('item_info', {}), ensure_ascii=False, indent=2)}\n\n"
|
||
prompt += f"分类体系:\n{json.dumps(input_data.get('category_system', {}), ensure_ascii=False, indent=2)}\n\n"
|
||
prompt += f"输出要求:\n{json.dumps(output_requirements, ensure_ascii=False, indent=2)}"
|
||
|
||
return prompt
|
||
|
||
def get_completion(prompts, model, tokenizer=None, temperature=0.6, top_p=0.95, top_k=20, min_p=0, max_tokens=4096, max_model_len=8192):
|
||
stop_token_ids = [151645, 151643]
|
||
# 创建采样参数。temperature 控制生成文本的多样性，top_p 控制核心采样的概率，top_k 通过限制候选词的数量来控制生成文本的质量和多样性, min_p 通过设置概率阈值来筛选候选词，从而在保证文本质量的同时增加多样性
|
||
sampling_params = SamplingParams(temperature=temperature, top_p=top_p, top_k=top_k, min_p=min_p, max_tokens=max_tokens, stop_token_ids=stop_token_ids) # max_tokens 用于限制模型在推理过程中生成的最大输出长度
|
||
# 初始化 vLLM 推理引擎
|
||
llm = LLM(model=model, tokenizer=tokenizer, max_model_len=max_model_len,trust_remote_code=True) # max_model_len 用于限制模型在推理过程中可以处理的最大输入和输出长度之和。
|
||
outputs = llm.generate(prompts, sampling_params)
|
||
return outputs
|
||
|
||
# 自动下载模型时，指定使用modelscope; 否则，会从HuggingFace下载
|
||
os.environ['VLLM_USE_MODELSCOPE']='True'
|
||
if __name__ == "__main__":
|
||
random.seed(114514)
|
||
random_numbers = [random.randint(0, 61929) for _ in range(20)]
|
||
|
||
# 初始化 vLLM 推理引擎
|
||
model='/root/autodl-tmp/Qwen/Qwen3-8B' # 指定模型路径
|
||
tokenizer = AutoTokenizer.from_pretrained(model, use_fast=False) # 加载分词器
|
||
|
||
work_dir = './prompt'
|
||
data_dir = './data'
|
||
|
||
results = []
|
||
|
||
for i in random_numbers:
|
||
|
||
prompt_json_path = f"{work_dir}/{i}_prompt.json" # 您的prompt文件路径
|
||
prompt_data = load_from_json(prompt_json_path)
|
||
prompt = build_prompt_from_json(prompt_data)
|
||
|
||
data_json_path = f"{data_dir}/{i}.json"
|
||
data_json = load_from_json(data_json_path)
|
||
|
||
messages = [
|
||
{"role": "user", "content": prompt}
|
||
]
|
||
text = tokenizer.apply_chat_template(
|
||
messages,
|
||
tokenize=False,
|
||
add_generation_prompt=True,
|
||
enable_thinking=True # 是否开启思考模式，默认为 True
|
||
)
|
||
|
||
outputs = get_completion(text, model, tokenizer=None, temperature=0.6, top_p = 0.95, top_k=20, min_p=0) # 对于思考模式，官方建议使用以下参数：temperature = 0.6，TopP = 0.95，TopK = 20，MinP = 0。
|
||
|
||
for output in outputs:
|
||
generated_text = output.outputs[0].text
|
||
|
||
result_entry = {}
|
||
|
||
# 首先添加样本ID字段
|
||
result_entry["样本ID"] = i
|
||
|
||
# 添加原始数据字段
|
||
if isinstance(data_json, dict):
|
||
for key, value in data_json.items():
|
||
result_entry[key] = value
|
||
|
||
# 提取</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]+)*)'
|
||
}
|
||
|
||
for field, pattern in patterns.items():
|
||
match = re.search(pattern, processing_text, re.IGNORECASE)
|
||
if match:
|
||
result_entry[field] = match.group(1).strip()
|
||
except Exception as e:
|
||
print(f"处理样本 {i} 时出错: {e}")
|
||
|
||
# 添加处理后的结果
|
||
results.append(result_entry)
|
||
|
||
# 创建DataFrame并保存到Excel
|
||
df = pd.DataFrame(results)
|
||
excel_file = '20_random_Qwen3_8B_responses.xlsx'
|
||
df.to_excel(excel_file, index=False)
|
||
print(f"\n所有生成文本已保存到 {excel_file}") |