run_training说明.md 6.7 KB

训练数据管理系统说明

概述

训练数据管理系统位于 data_pipeline/trainer/ 目录下,负责将生成的训练数据文件加载到向量数据库中。该系统支持多种文件格式的自动识别和处理。

主要组件

1. 核心文件

  • run_training.py - 主训练脚本,支持命令行调用
  • vanna_trainer.py - 训练器核心模块,封装训练逻辑

2. 配置来源

训练数据路径配置现已统一到 data_pipeline/config.py

SCHEMA_TOOLS_CONFIG = {
    "output_directory": "./data_pipeline/training_data/",
    # 其他配置...
}

文件格式与处理逻辑

文件处理优先级和判断逻辑

代码按以下顺序判断文件类型:

  1. .ddl → DDL文件
  2. .md.markdown → 文档文件
  3. _pair.json_pairs.json → JSON问答对文件
  4. _pair.sql_pairs.sql → 格式化问答对文件
  5. .sql (但不以 _pair.sql_pairs.sql 结尾) → SQL示例文件
  6. 其他 → 跳过处理

1. DDL文件 (.ddl)

  • 处理函数: train_ddl_statements()
  • 调用的训练函数: train_ddl()
  • 文件格式:

    • 使用分号 (;) 作为分隔符
    • 每个DDL语句之间用分号分隔
    • 示例格式:

      create table public.bss_company (
      id varchar(32) not null     -- 主键ID,主键,
      version integer not null    -- 版本号,
      company_name varchar(255)   -- 公司名称,
      primary key (id)
      );
      

      2. 文档文件 (.md, .markdown)

      • 处理函数: train_documentation_blocks()
      • 调用的训练函数: train_documentation()
      • 文件格式:
      • Markdown文件: 按标题级别自动分割 (#, ##, ###)
      • 非Markdown文件: 使用 --- 作为分隔符
      • 示例格式: ```markdown

        bss_company(存储高速公路管理公司信息)

        bss_company 表存储高速公路管理公司信息,用于服务区运营管理

      字段列表:

      • id (varchar(32)) - 主键ID [主键, 非空]
      • company_name (varchar(255)) - 公司名称 ```

3. SQL示例文件 (.sql, 但排除 _pair.sql_pairs.sql)

  • 处理函数: train_sql_examples()
  • 调用的训练函数: train_sql_example()
  • 文件格式:

    • 使用分号 (;) 作为分隔符
    • 每个SQL示例之间用分号分隔
    • 示例格式:

      SELECT * FROM bss_company WHERE delete_ts IS NULL;
      SELECT company_name, company_no FROM bss_company ORDER BY company_name;
      

      4. 格式化问答对文件 (_pair.sql, _pairs.sql)

      • 处理函数: train_formatted_question_sql_pairs()
      • 调用的训练函数: train_question_sql_pair()
      • 文件格式:
      • 使用 Question:SQL: 标记
      • 问答对之间用双空行分隔
      • 支持单行和多行SQL
      • 示例格式: ```

    Question: 查询所有公司信息 SQL: SELECT * FROM bss_company WHERE delete_ts IS NULL;

    Question: 统计每个公司的服务区数量 SQL: SELECT c.company_name, COUNT(sa.id) as area_count FROM bss_company c LEFT JOIN bss_service_area sa ON c.id = sa.company_id WHERE c.delete_ts IS NULL GROUP BY c.company_name;

    
    ### 5. JSON格式问答对文件 (`_pair.json`, `_pairs.json`)
    - **处理函数**: `train_json_question_sql_pairs()`
    - **调用的训练函数**: `train_question_sql_pair()`
    - **文件格式**:
    - 标准JSON数组格式
    - 每个对象包含 `question` 和 `sql` 字段
    - 示例格式:
    ```json
    [
        {
            "question": "查询所有公司信息",
            "sql": "SELECT * FROM bss_company WHERE delete_ts IS NULL"
        },
        {
            "question": "按公司统计服务区数量",
            "sql": "SELECT company_name, COUNT(*) FROM bss_service_area GROUP BY company_name"
        }
    ]
    

使用方式

1. 命令行使用

基本使用

# 使用默认配置路径
python -m data_pipeline.trainer.run_training

# 指定训练数据目录
python -m data_pipeline.trainer.run_training --data_path ./data_pipeline/training_data/

2. 在工作流中调用

训练数据加载已集成到工作流编排器中,作为第4步自动执行:

python -m data_pipeline.schema_workflow \
  --db-connection "postgresql://user:pass@localhost:5432/database" \
  --table-list tables.txt \
  --business-context "业务描述"

3. 编程方式调用

from data_pipeline.trainer.run_training import process_training_files

# 处理训练文件
success = process_training_files("./data_pipeline/training_data/")
if success:
    print("训练数据加载成功")

扫描策略

目录扫描范围

  • 只扫描指定目录的直接文件,不递归扫描子目录
  • 跳过所有子目录,只处理文件

不支持的文件类型

  • .txt 文件 - 不被处理
  • 其他扩展名文件 - 被跳过
  • 子目录 - 被忽略

配置架构

配置优先级

命令行参数 > data_pipeline/config.py > 默认值

统一配置

所有数据管道相关配置现统一在 data_pipeline/config.py

SCHEMA_TOOLS_CONFIG = {
    "output_directory": "./data_pipeline/training_data/",
    # 训练相关配置...
}

统计信息

训练完成后显示统计:

  • DDL文件数量
  • 文档文件数量
  • SQL示例文件数量
  • 格式化问答对文件数量
  • JSON问答对文件数量
  • 总训练记录数
  • 各数据类型分布

向量数据库集成

支持的向量数据库

  • ChromaDB - 本地文件存储
  • PgVector - PostgreSQL向量扩展

数据类型标识

训练数据加载时会自动标记数据类型:

  • ddl - DDL语句
  • documentation - 文档内容
  • sql - SQL示例和问答对

验证机制

训练完成后自动验证:

  • 从向量数据库检索训练数据
  • 统计各类型数据数量
  • 确保数据成功加载

最佳实践

文件组织

data_pipeline/training_data/
├── *.ddl              # DDL文件
├── *_detail.md        # MD文档文件
├── qs_*_pair.json     # 问答对文件
├── filename_mapping.txt  # 文件映射
└── logs/              # 日志目录(如果需要)

命名规范

  • DDL文件: table_name.ddlschema__table_name.ddl
  • MD文件: table_name_detail.mdschema__table_name_detail.md
  • JSON问答对: qs_dbname_timestamp_pair.json

性能优化

  • 批处理配置现已移至 app_config.py 中的训练配置区域
  • 单线程处理确保稳定性
  • 自动识别文件格式,无需手动分类

这个训练数据管理系统为完整的数据管道提供了最后一环,确保生成的训练数据能够有效地加载到向量数据库中供AI模型使用。