284 lines
6.8 KiB
Markdown
284 lines
6.8 KiB
Markdown
# QA数据集采样工具
|
||
|
||
一个全面的QA数据集采样工具，使用OpenAI的GPT模型生成答案。该工具帮助您从大规模数据集（如MS MARCO）创建高质量的问答数据集。
|
||
|
||
## 功能特性
|
||
|
||
- **智能采样**：智能地从大型数据集中采样查询、文档和相关性判断
|
||
- **答案生成**：使用OpenAI的GPT模型自动生成高质量答案
|
||
- **断点续传**：支持中断后继续生成，从上次位置开始
|
||
- **进度跟踪**：实时进度更新和统计信息
|
||
- **结果可视化**：易于阅读的问答对展示，包含完整上下文
|
||
|
||
## 安装指南
|
||
|
||
### 系统要求
|
||
|
||
- Python 3.7+
|
||
- OpenAI API密钥
|
||
|
||
### 安装依赖
|
||
|
||
```bash
|
||
pip install pandas pyarrow openai
|
||
```
|
||
|
||
### 设置环境变量
|
||
|
||
```bash
|
||
export OPENAI_API_KEY="你的openai-api-key"
|
||
# 可选：使用自定义OpenAI端点
|
||
export OPENAI_BASE_URL="https://api.openai.com/v1"
|
||
```
|
||
|
||
### 准备数据集
|
||
|
||
您可以使用任何符合格式要求的QA数据集，或下载预处理好的样本：
|
||
|
||
**使用HuggingFace/ModelScope样本**
|
||
我们提供了来自流行QA数据集的预处理样本：
|
||
- MarkrAI/eli5_sample_autorag
|
||
- MarkrAI/msmarco_sample_autorag
|
||
- MarkrAI/triviaqa_sample_autorag
|
||
- gnekt/hotpotqa_small_sample_autorag
|
||
|
||
**使用您自己的数据集**
|
||
确保您的数据集包含以下文件：
|
||
- `queries.parquet`（列：id, text）
|
||
- `corpus.parquet`（列：id, text）
|
||
- `qrels.parquet`（列：qid, pid）
|
||
|
||
## 快速开始
|
||
|
||
### 1. 从大型数据集采样
|
||
|
||
首先，从完整数据集中采样查询、文档和相关性判断的子集：
|
||
|
||
```bash
|
||
python dataset/qa_dataset.py sample \
|
||
--queries ~/dataset/mmarco-queries.parquet \
|
||
--corpus ~/dataset/mmarco-corpus.parquet \
|
||
--qrels ~/dataset/mmarco-qrels.parquet \
|
||
--nq 100 \
|
||
--output_dir ./dataset/samples
|
||
```
|
||
|
||
### 2. 生成答案
|
||
|
||
使用OpenAI的GPT模型为采样的问答生成答案：
|
||
|
||
```bash
|
||
python dataset/qa_dataset.py generate \
|
||
--input_dir ./dataset/samples \
|
||
--output_dir ./dataset/samples
|
||
```
|
||
|
||
### 3. 查看结果
|
||
|
||
展示生成的问答对及其上下文：
|
||
|
||
```bash
|
||
python dataset/qa_dataset.py show \
|
||
--input_dir ./dataset/samples \
|
||
-n 5
|
||
```
|
||
|
||
## 详细使用说明
|
||
|
||
### 采样命令
|
||
|
||
从完整数据集中创建代表性样本。
|
||
|
||
```bash
|
||
python dataset/qa_dataset.py sample [选项]
|
||
```
|
||
|
||
**必需参数：**
|
||
- `--queries`：查询parquet文件路径（列：`id`, `text`）
|
||
- `--corpus`：语料库parquet文件路径（列：`id`, `text`）
|
||
- `--qrels`：相关性判断parquet文件路径（列：`qid`, `pid`）
|
||
|
||
**可选参数：**
|
||
- `--nq`：要采样的查询数量（默认：1000）
|
||
- `--output_dir`：采样数据输出目录（默认：./save）
|
||
|
||
**示例：**
|
||
```bash
|
||
python dataset/qa_dataset.py sample \
|
||
--queries data/queries.parquet \
|
||
--corpus data/corpus.parquet \
|
||
--qrels data/qrels.parquet \
|
||
--nq 500 \
|
||
--output_dir ./my_sample
|
||
```
|
||
|
||
### 生成命令
|
||
|
||
使用OpenAI API为采样问题生成答案。
|
||
|
||
```bash
|
||
python dataset/qa_dataset.py generate [选项]
|
||
```
|
||
|
||
**必需参数：**
|
||
- `--input_dir`：包含采样数据的目录（queries.parquet, corpus.parquet, qrels.parquet）
|
||
|
||
**可选参数：**
|
||
- `--output_dir`：生成答案的输出目录（默认：./save）
|
||
|
||
**特性：**
|
||
- **断点续传**：中断后自动从上次位置继续
|
||
- **错误处理**：API调用失败自动重试3次
|
||
- **进度保存**：每成功生成一个答案就保存进度
|
||
|
||
**示例：**
|
||
```bash
|
||
python dataset/qa_dataset.py generate \
|
||
--input_dir ./my_sample \
|
||
--output_dir ./my_sample
|
||
```
|
||
|
||
### 展示命令
|
||
|
||
展示生成的问答对及完整上下文。
|
||
|
||
```bash
|
||
python dataset/qa_dataset.py show [选项]
|
||
```
|
||
|
||
**必需参数：**
|
||
- `--input_dir`：包含QA数据的目录（queries.parquet, corpus.parquet, qrels.parquet, qas.parquet, answers.parquet）
|
||
|
||
**可选参数：**
|
||
- `-n`：要展示的结果数量（默认：5）
|
||
|
||
**示例：**
|
||
```bash
|
||
python dataset/qa_dataset.py show \
|
||
--input_dir ./my_sample \
|
||
-n 3
|
||
```
|
||
|
||
## 输入数据格式
|
||
|
||
### 查询文件 (queries.parquet)
|
||
| 列名 | 类型 | 描述 |
|
||
|------|------|------|
|
||
| id | string | 唯一查询标识符 |
|
||
| text | string | 实际的问题文本 |
|
||
|
||
### 语料库文件 (corpus.parquet)
|
||
| 列名 | 类型 | 描述 |
|
||
|------|------|------|
|
||
| id | string | 唯一段落/文档标识符 |
|
||
| text | string | 段落/文档内容 |
|
||
|
||
### 相关性判断文件 (qrels.parquet)
|
||
| 列名 | 类型 | 描述 |
|
||
|------|------|------|
|
||
| qid | string | 查询ID（匹配queries.id） |
|
||
| pid | string | 段落ID（匹配corpus.id） |
|
||
|
||
## 输出文件
|
||
|
||
运行所有命令后，输出目录将包含：
|
||
|
||
### 采样数据
|
||
- `queries.parquet`：采样的查询子集
|
||
- `corpus.parquet`：采样的文档子集
|
||
- `qrels.parquet`：采样的相关性判断
|
||
|
||
### 生成的答案
|
||
- `answers.parquet`：生成的答案（含唯一ID）
|
||
- `qas.parquet`：问答映射（qid → aid）
|
||
|
||
## 高级用法
|
||
|
||
### 自定义OpenAI配置
|
||
|
||
您可以使用不同的OpenAI模型或端点：
|
||
|
||
```bash
|
||
# 使用GPT-4 Turbo
|
||
export OPENAI_API_KEY="你的密钥"
|
||
python dataset/qa_dataset.py generate --input_dir ./samples
|
||
|
||
# 使用Azure OpenAI
|
||
export OPENAI_API_KEY="azure密钥"
|
||
export OPENAI_BASE_URL="https://你的资源.openai.azure.com/openai/deployments/gpt-4"
|
||
python dataset/qa_dataset.py generate --input_dir ./samples
|
||
```
|
||
|
||
### 大型数据集采样
|
||
|
||
对于非常大的数据集，建议分批采样：
|
||
|
||
```bash
|
||
# 第一批
|
||
python dataset/qa_dataset.py sample --nq 1000 --output_dir ./batch1
|
||
python dataset/qa_dataset.py generate --input_dir ./batch1
|
||
|
||
# 第二批
|
||
python dataset/qa_dataset.py sample --nq 1000 --output_dir ./batch2
|
||
python dataset/qa_dataset.py generate --input_dir ./batch2
|
||
```
|
||
|
||
## 故障排除
|
||
|
||
### 常见问题
|
||
|
||
**1. OpenAI API错误**
|
||
- 确保API密钥设置正确：`echo $OPENAI_API_KEY`
|
||
- 检查API配额和账单状态
|
||
- 验证与OpenAI的网络连接
|
||
|
||
**2. 大数据集内存问题**
|
||
- 减小`--nq`参数以获得更小的样本
|
||
- 确保pandas操作有足够的RAM
|
||
- 考虑使用更小的parquet文件
|
||
|
||
**3. 文件未找到错误**
|
||
- 验证所有输入文件路径是否正确
|
||
- 确保parquet文件有正确的列名
|
||
- 检查文件权限
|
||
|
||
### 调试模式
|
||
|
||
通过添加打印语句或使用Python调试器启用详细输出：
|
||
|
||
```bash
|
||
python -m pdb dataset/qa_dataset.py sample --queries ...
|
||
```
|
||
|
||
## 示例工作流
|
||
|
||
```bash
|
||
# 1. 设置环境
|
||
export OPENAI_API_KEY="sk-..."
|
||
|
||
# 2. 从MS MARCO采样200个查询
|
||
python dataset/qa_dataset.py sample \
|
||
--queries ~/mmarco/queries.parquet \
|
||
--corpus ~/mmarco/corpus.parquet \
|
||
--qrels ~/mmarco/qrels.parquet \
|
||
--nq 200 \
|
||
--output_dir ./marco_sample
|
||
|
||
# 3. 生成答案（根据API速率限制可能需要一些时间）
|
||
python dataset/qa_dataset.py generate \
|
||
--input_dir ./marco_sample \
|
||
--output_dir ./marco_sample
|
||
|
||
# 4. 查看结果
|
||
python dataset/qa_dataset.py show \
|
||
--input_dir ./marco_sample \
|
||
-n 10
|
||
```
|
||
|
||
## 贡献
|
||
|
||
欢迎提交问题和功能增强请求！
|
||
|
||
## 许可证
|
||
|
||
MIT许可证 - 可自由用于研究和项目。 |