renzhiyuan
/
l_ai_knowledge
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
l_ai_knowledge/dataset/README_zh.md

6.8 KiB
Raw Permalink Blame History

QA数据集采样工具

一个全面的QA数据集采样工具，使用OpenAI的GPT模型生成答案。该工具帮助您从大规模数据集（如MS MARCO）创建高质量的问答数据集。

功能特性

  • 智能采样：智能地从大型数据集中采样查询、文档和相关性判断
  • 答案生成：使用OpenAI的GPT模型自动生成高质量答案
  • 断点续传：支持中断后继续生成，从上次位置开始
  • 进度跟踪：实时进度更新和统计信息
  • 结果可视化：易于阅读的问答对展示，包含完整上下文

安装指南

系统要求

  • Python 3.7+
  • OpenAI API密钥

安装依赖

pip install pandas pyarrow openai

设置环境变量

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. 从大型数据集采样

首先，从完整数据集中采样查询、文档和相关性判断的子集：

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模型为采样的问答生成答案：

python dataset/qa_dataset.py generate \
  --input_dir ./dataset/samples \
  --output_dir ./dataset/samples

3. 查看结果

展示生成的问答对及其上下文：

python dataset/qa_dataset.py show \
  --input_dir ./dataset/samples \
  -n 5

详细使用说明

采样命令

从完整数据集中创建代表性样本。

python dataset/qa_dataset.py sample [选项]

必需参数：

  • --queries：查询parquet文件路径（列：id, text）
  • --corpus：语料库parquet文件路径（列：id, text）
  • --qrels：相关性判断parquet文件路径（列：qid, pid）

可选参数：

  • --nq：要采样的查询数量（默认：1000）
  • --output_dir：采样数据输出目录（默认：./save）

示例：

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为采样问题生成答案。

python dataset/qa_dataset.py generate [选项]

必需参数：

  • --input_dir：包含采样数据的目录（queries.parquet, corpus.parquet, qrels.parquet）

可选参数：

  • --output_dir：生成答案的输出目录（默认：./save）

特性：

  • 断点续传：中断后自动从上次位置继续
  • 错误处理：API调用失败自动重试3次
  • 进度保存：每成功生成一个答案就保存进度

示例：

python dataset/qa_dataset.py generate \
  --input_dir ./my_sample \
  --output_dir ./my_sample

展示命令

展示生成的问答对及完整上下文。

python dataset/qa_dataset.py show [选项]

必需参数：

  • --input_dir：包含QA数据的目录（queries.parquet, corpus.parquet, qrels.parquet, qas.parquet, answers.parquet）

可选参数：

  • -n：要展示的结果数量（默认：5）

示例：

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模型或端点：

# 使用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

大型数据集采样

对于非常大的数据集，建议分批采样：

# 第一批
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调试器启用详细输出：

python -m pdb dataset/qa_dataset.py sample --queries ...

示例工作流

# 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许可证 - 可自由用于研究和项目。