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}")