大模型应用：批量文档摘要与分类实践：本地合同、报告数据处理与导出.70

# 导入所需库
import os
import re
import pandas as pd
from docx import Document
import PyPDF2

# ======================================
# 配置参数（只需修改这部分路径和参数）
# ======================================
MODEL_PATH = r"D:\AIWorld\models\Qwen1.5-7B-Chat-4bit"  # Qwen 1.5 7B模型路径
DOCUMENT_FOLDER = r"D:\Documents\Contracts\ToProcess"  # 文档文件夹路径
OUTPUT_CSV_PATH = r"D:\Documents\Contracts\Batch_Result.csv"  # 输出CSV文件路径
MAX_SUMMARY_LENGTH = 200  # 摘要最大字数
MAX_CHUNK_TOKENS = 3000  # 每个文本块的最大token数
CHUNK_OVERLAP_TOKENS = 50  # 文本块之间的重叠token数
# 预设多标签列表（可根据你的业务需求修改/扩展）
PREDEFINED_TAGS = ["合同", "租赁类", "买卖类", "服务类", "报告", "财务报告", "审计报告", "项目报告", "办公相关", "个人相关"]

# ======================================
# 步骤1：工具函数 - 读取单个文档（支持.docx和.pdf）
# ======================================
def read_single_document(file_path):
    """
    读取单个文档的文本内容，支持.docx和.pdf格式
    :param file_path: 单个文档的完整路径
    :return: 清理后的文档文本内容
    """
    text_content = ""
    
    # 区分文档格式，分别处理
    if file_path.endswith(".docx"):
        try:
            doc = Document(file_path)
            for paragraph in doc.paragraphs:
                if paragraph.text.strip():  # 跳过空段落
                    text_content += paragraph.text + "\n\n"
        except Exception as e:
            print(f"读取.docx文档失败：{file_path}，错误信息：{str(e)}")
            return ""
    
    elif file_path.endswith(".pdf"):
        try:
            with open(file_path, "rb") as pdf_file:
                pdf_reader = PyPDF2.PdfReader(pdf_file)
                for page in pdf_reader.pages:
                    page_text = page.extract_text()
                    if page_text.strip():  # 跳过空页面
                        text_content += page_text + "\n\n"
        except Exception as e:
            print(f"读取.pdf文档失败：{file_path}，错误信息：{str(e)}")
            return ""
    
    else:
        print(f"不支持的文档格式：{file_path}")
        return ""
    
    # 文本预处理：清理无效字符、冗余空格
    text_content = clean_text(text_content)
    return text_content

# ======================================
# 步骤2：工具函数 - 文本预处理（清理无效内容）
# ======================================
def clean_text(text):
    """
    清理文本中的无效字符、冗余空格、分页符等
    :param text: 原始文本内容
    :return: 清理后的文本内容
    """
    # 去除多余的空格、换行符
    text = re.sub(r"\n+", "\n\n", text)
    text = re.sub(r"\s+", " ", text)
    # 去除常见的无效字符（如分页符、特殊符号）
    text = re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]", "", text)
    # 去除页眉页脚常见的重复内容（可根据你的文档情况扩展）
    text = re.sub(r"第\d+页/\d+页", "", text)
    text = re.sub(r"版权所有\s*©.*", "", text)
    return text.strip()

from langchain.text_splitter import RecursiveCharacterTextSplitter

# ======================================
# 步骤3：工具函数 - 文本分块（使用Langchain的RecursiveCharacterTextSplitter）
# ======================================
def split_text_into_chunks(text, tokenizer):
    """
    将长文本分割为符合模型上下文窗口限制的小文本块
    :param text: 清理后的完整文档文本
    :param tokenizer: 模型对应的tokenizer
    :return: 文本块列表
    """
    # 初始化文本分割器
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=MAX_CHUNK_TOKENS,
        chunk_overlap=CHUNK_OVERLAP_TOKENS,
        length_function=lambda x: len(tokenizer.encode(x)),  # 以token数作为长度计算标准
        separators=["\n\n", "\n", " ", ""]  # 分割符优先级
    )
    
    # 执行分割
    chunks = text_splitter.split_text(text)
    return chunks

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import torch

# ======================================
# 步骤4：工具函数 - 构建Schema提示词
# ======================================
def build_prompt(chunk_text):
    """
    构建符合Schema要求的提示词，引导模型输出标准化结果
    :param chunk_text: 单个文本块的内容
    :return: 完整的提示词
    """
    prompt = f"""
你是一名专业的文档处理助手，负责总结文档核心信息并标注多标签分类。
请严格按照以下要求完成任务：
1.  文档摘要：总结该文本块的核心信息，控制在{MAX_SUMMARY_LENGTH}字以内，语言简洁、准确，不遗漏关键信息。
2.  多标签分类：从预设标签列表中选择1个或多个符合文本内容的标签，标签之间用英文逗号","分隔，不得使用预设列表外的标签。
3.  输出格式：严格按照以下固定格式输出，不得添加任何额外内容、注释或标点符号：
### 文档摘要
[你的摘要内容]
### 多标签分类
[你的标签内容]

预设标签列表：{",".join(PREDEFINED_TAGS)}
文本块内容：
{chunk_text}
"""
    return prompt.strip()

# ======================================
# 步骤5：工具函数 - 调用Qwen 1.5 7B模型处理单个文本块
# ======================================
def process_single_chunk(chunk_text, model, tokenizer):
    """
    调用Qwen 1.5 7B模型，处理单个文本块，生成摘要与分类
    :param chunk_text: 单个文本块的内容
    :param model: 加载后的Qwen 1.5 7B模型
    :param tokenizer: 加载后的模型tokenizer
    :return: 该文本块的摘要（str）、标签（str）
    """
    # 构建提示词
    prompt = build_prompt(chunk_text)
    
    # 构建模型输入（符合Qwen模型的输入格式要求）
    messages = [
        {"role": "user", "content": prompt}
    ]
    input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device)
    
    # 模型生成配置（控制输出长度、避免重复）
    generation_config = {
        "max_new_tokens": 500,  # 足够容纳摘要和标签
        "temperature": 0.3,  # 温度越低，输出越稳定、越符合格式要求
        "top_p": 0.9,
        "do_sample": False,  # 关闭采样，保证输出一致性
        "eos_token_id": tokenizer.eos_token_id,
        "pad_token_id": tokenizer.pad_token_id
    }
    
    # 执行生成
    with torch.no_grad():  # 关闭梯度计算，减少内存占用
        outputs = model.generate(input_ids, **generation_config)
    
    # 解析输出结果
    response = tokenizer.decode(outputs[0][input_ids.shape[1]:], skip_special_tokens=True).strip()
    
    # 提取摘要和标签（按照Schema格式解析）
    summary = ""
    tags = ""
    try:
        # 匹配"### 文档摘要"后的内容
        summary_match = re.search(r"### 文档摘要\n(.*?)\n### 多标签分类", response, re.DOTALL)
        # 匹配"### 多标签分类"后的内容
        tags_match = re.search(r"### 多标签分类\n(.*)", response, re.DOTALL)
        
        if summary_match:
            summary = summary_match.group(1).strip()
        if tags_match:
            tags = tags_match.group(1).strip()
    except Exception as e:
        print(f"解析模型输出失败，错误信息：{str(e)}，模型原始输出：{response}")
    
    return summary, tags

# ======================================
# 步骤6：工具函数 - 整合多个文本块的结果
# ======================================
def integrate_chunk_results(chunk_summaries, chunk_tags):
    """
    整合所有文本块的摘要和标签，生成完整文档的结果
    :param chunk_summaries: 所有文本块的摘要列表
    :param chunk_tags: 所有文本块的标签列表
    :return: 完整摘要（str）、完整标签（str）
    """
    # 整合摘要：去重后拼接，保持简洁
    integrated_summary = "\n".join([s for s in chunk_summaries if s])
    # 若整合后摘要过长，再次精简（取前MAX_SUMMARY_LENGTH*2字，避免过长）
    if len(integrated_summary) > MAX_SUMMARY_LENGTH * 2:
        integrated_summary = integrated_summary[:MAX_SUMMARY_LENGTH * 2] + "..."
    
    # 整合标签：去重，保持格式统一
    all_tags = []
    for tag_str in chunk_tags:
        if tag_str:
            tags = [t.strip() for t in tag_str.split(",")]
            all_tags.extend(tags)
    # 去重并排序，保持结果一致性
    unique_tags = sorted(list(set(all_tags)))
    integrated_tags = ",".join(unique_tags)
    
    return integrated_summary, integrated_tags

# ======================================
# 步骤7：主函数 - 批量处理所有文档
# ======================================
def batch_process_documents():
    """
    批量处理指定文件夹中的所有文档，生成摘要与分类，最终导出CSV
    """
    # 步骤7.1：加载Qwen 1.5 7B模型和tokenizer
    print("开始加载Qwen 1.5 7B模型，请稍候（首次加载可能需要1-2分钟）...")
    try:
        # 加载tokenizer
        tokenizer = AutoTokenizer.from_pretrained(
            MODEL_PATH,
            trust_remote_code=True
        )
        # 加载模型
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_PATH,
            trust_remote_code=True,
            torch_dtype=torch.float16,  # 使用float16精度，减少内存占用
            device_map="auto"  # 自动分配设备（优先使用GPU，无GPU则使用CPU）
        )
        # 模型设置为评估模式
        model.eval()
        print("模型加载成功！")
    except Exception as e:
        print(f"模型加载失败，错误信息：{str(e)}")
        return
    
    # 步骤7.2：初始化结果列表
    result_list = []
    
    # 步骤7.3：遍历文档文件夹，批量处理每个文档
    print(f"开始遍历文档文件夹：{DOCUMENT_FOLDER}")
    for file_name in os.listdir(DOCUMENT_FOLDER):
        file_path = os.path.join(DOCUMENT_FOLDER, file_name)
        
        # 仅处理.docx和.pdf文件
        if not (file_path.endswith(".docx") or file_path.endswith(".pdf")):
            continue
        
        print(f"\n开始处理文档：{file_name}")
        
        # 步骤7.3.1：读取文档文本
        doc_text = read_single_document(file_path)
        if not doc_text:
            print(f"文档{file_name}无有效文本，跳过处理")
            continue
        
        # 步骤7.3.2：文本分块
        text_chunks = split_text_into_chunks(doc_text, tokenizer)
        print(f"文档{file_name}分割为{len(text_chunks)}个文本块")
        
        # 步骤7.3.3：处理每个文本块
        chunk_summaries = []
        chunk_tags = []
        for idx, chunk in enumerate(text_chunks):
            print(f"正在处理第{idx+1}/{len(text_chunks)}个文本块")
            summary, tags = process_single_chunk(chunk, model, tokenizer)
            chunk_summaries.append(summary)
            chunk_tags.append(tags)
        
        # 步骤7.3.4：整合文本块结果
        integrated_summary, integrated_tags = integrate_chunk_results(chunk_summaries, chunk_tags)
        
        # 步骤7.3.5：添加到结果列表
        result_list.append({
            "文档文件名": file_name,
            "文档完整路径": file_path,
            "文档摘要": integrated_summary,
            "多标签分类": integrated_tags
        })
        
        print(f"文档{file_name}处理完成，摘要：{integrated_summary[:50]}...，标签：{integrated_tags}")
    
    # 步骤7.4：将结果导出为CSV文件
    if result_list:
        try:
            df = pd.DataFrame(result_list)
            df.to_csv(OUTPUT_CSV_PATH, index=False, encoding="utf_8_sig")  # utf_8_sig支持Excel中文显示
            print(f"\n所有文档处理完成！结果已导出至CSV文件：{OUTPUT_CSV_PATH}")
        except Exception as e:
            print(f"CSV文件导出失败，错误信息：{str(e)}")
    else:
        print("\n无有效文档处理结果，未生成CSV文件")

# ======================================
# 运行主函数
# ======================================
if __name__ == "__main__":
    batch_process_documents()