zbtrs/dayunhe

政务/应急文档问答系统

基于 RAG (Retrieval-Augmented Generation) 的文档问答系统，支持政务和应急任务的文档检索与问答。

功能特性

• 🔍 智能文档检索：基于 Embedding 的语义检索，支持多文档并行处理
• 📄 多格式文档支持：支持 docx、pdf、txt 等多种文档格式
• 🚀 高性能推理：使用 vLLM 进行模型推理，支持并发处理
• ⚡ 异步处理：支持异步并发处理多个问题，提高处理效率
• 🎯 任务切换：支持政务（zw）和应急（yj）两种任务模式

环境要求

• Python 3.10
• CUDA 支持的 GPU（推荐多卡）
• 足够的 GPU 显存（建议 24GB+）

安装步骤

1. 创建虚拟环境

conda create -n mars python=3.10
conda activate mars

2. 安装依赖

pip install -r requirements.txt

注意：

• textract 可能需要安装系统级依赖，如 libmagic
• vllm 和 torch 建议根据 CUDA 版本安装对应版本

3. 准备模型和数据

确保以下目录结构存在：

.
├── models/
│   ├── Qwen3-zw/          # 政务任务模型
│   └── Qwen3-Embedding-0.6B/  # Embedding 模型
├── data/
│   ├── zw/                # 政务任务数据
│   │   ├── testa.xlsx 或 testa.csv
│   │   └── rag/          # RAG 文档目录
│   └── yj/                # 应急任务数据（可选）
│       ├── testa.xlsx 或 testa.csv
│       └── rag/
└── code/
    ├── main.py
    ├── llm.py
    ├── embedding.py
    ├── utils.py
    └── serve.sh

使用方法

1. 启动 vLLM 服务

cd code
bash serve.sh

服务配置说明：

• 默认启动政务任务模型（../models/Qwen3-zw）
• 如需切换为应急任务，修改 serve.sh 中的模型路径为 ../models/Qwen3-yj（或其他应急模型路径）
• 服务默认运行在 http://localhost:8000/v1

serve.sh 参数说明：

• CUDA_VISIBLE_DEVICES: 指定使用的 GPU 设备
• --tensor-parallel-size: 张量并行大小，根据 GPU 数量调整
• --gpu-memory-utilization: GPU 内存利用率
• --max-model-len: 最大模型长度

2. 运行推理

python main.py

任务切换：

• 默认任务：政务任务（TASK = 'zw'）
• 切换为应急任务：修改 main.py 中的全局变量 TASK = 'yj'

配置说明

main.py 主要配置参数

TASK = 'zw'                    # 任务类型：'zw'（政务）或 'yj'（应急）
MAX_CHUNK_SIZE = 2048          # 文本分块大小
CHUNK_OVERLAP = 256            # 分块重叠大小
K_NUM = 3                      # 检索 top-k 数量
CONCURRENCY_LEVEL = 16         # 并发处理级别
MODEL_NAME = "qwen"            # 模型名称（需与 serve.sh 中的 --served-model-name 一致）

数据路径配置

• 测试数据：../data/{TASK}/testa.xlsx 或 ../data/{TASK}/testa.csv
• RAG 文档目录：../data/{TASK}/rag/
• Embedding 模型：../models/Qwen3-Embedding-0.6B
• 输出结果：../result.csv

项目结构

code/
├── main.py           # 主程序入口
├── llm.py            # LLM 调用接口（基于 OpenAI API）
├── embedding.py      # Embedding 检索器
├── utils.py          # 工具函数（文件读取、数据处理等）
└── serve.sh          # vLLM 服务启动脚本

工作流程

1. 文档加载：从 RAG 目录加载所有文档（支持 docx、pdf、txt 等格式）
2. 文本分割：使用 RecursiveCharacterTextSplitter 将文档分割成块
3. 构建索引：使用 Embedding 模型对文档块进行向量化并缓存
4. 问题检索：对每个问题检索最相关的 top-k 文档块
5. 生成答案：将检索到的文档块和问题一起发送给 LLM 生成答案
6. 结果输出：将结果保存为 CSV 文件

注意事项

1. GPU 显存：确保有足够的 GPU 显存运行 vLLM 服务和 Embedding 模型
2. 服务端口：确保 8000 端口未被占用，或修改 llm.py 中的 base_url
3. 文档格式：不支持 .wps 格式，其他常见格式均可处理
4. 并发控制：根据 GPU 显存和性能调整 CONCURRENCY_LEVEL
5. 模型路径：确保模型路径正确，且模型文件完整

输出格式

结果文件 result.csv 包含以下字段：

• ID: 问题 ID
• answer: 生成的答案

常见问题

Q: 服务启动失败？
A: 检查 GPU 显存是否足够，CUDA 环境是否正确配置，模型路径是否存在。

Q: 推理速度慢？
A: 可以调整 CONCURRENCY_LEVEL 和 K_NUM，或增加 GPU 数量。

Q: 内存不足？
A: 减小 MAX_CHUNK_SIZE、CONCURRENCY_LEVEL 或 --gpu-memory-utilization。

依赖说明

主要依赖包：

• vllm: 高性能 LLM 推理引擎
• torch: PyTorch 深度学习框架
• openai: OpenAI API 客户端（用于调用本地 vLLM 服务）
• langchain-text-splitters: 文本分割工具
• pandas: 数据处理
• python-docx: Word 文档处理
• textract: 文档文本提取
• tqdm: 进度条显示