功能总结

1. 导出缓存数据集（cached_dataset.py）

作用：将预处理后的数据集保存为缓存格式，避免训练时重复 tokenize。

工作流程：

展开代码
原始数据集 → 预处理/Tokenize → 保存为缓存格式

详细逻辑：

1. 路径识别

   • 识别 S3/OSS 路径（s3://, oss://, s3:/, oss:/）
   • 规范化路径格式（统一为 s3:// 或 oss://）

2. 保存策略

   • 本地路径：直接保存到 {output_dir}/train 和 {output_dir}/val
   • S3/OSS 路径：
     • 在本地临时目录保存
     • 使用 megfile 递归复制到 S3/OSS
     • 复制完成后清理临时目录

3. 目录结构

   展开代码
   {output_dir}/
   ├── train/          # 训练集缓存
   └── val/            # 验证集缓存（可选）
   

4. 特性

   • 支持 S3/OSS 远程存储
   • 自动处理 train/val 分离
   • 显示复制进度（每 100 个文件）

2. 加载缓存数据集（infer/utils.py）

作用：从缓存格式加载数据集，用于训练/推理。

工作流程：

展开代码
缓存路径 → 识别路径类型 → 构建加载路径 → 加载数据集

详细逻辑：

1. 路径处理

   • 支持 dataset_sample 语法：path#1000（采样 1000 条）
   • 识别并规范化 S3/OSS 路径

2. 路径构建策略

   对于 S3/OSS 路径：

   展开代码
   输入: s3://bucket/dataset
   训练集: s3://bucket/dataset/train
   验证集: s3://bucket/dataset/val
   

   对于本地路径：

   • 优先检查是否存在 train/ 或 val/ 子目录
   • 存在则加载子目录，否则尝试直接加载根目录
   • 都不存在则回退到 {path}/train 或 {path}/val

3. 数据集过滤

   • 根据 max_length 过滤超长样本
   • 验证集不存在时跳过（不报错）

4. 额外功能

   • 支持数据集采样（#sample_size 语法）
   • 支持多个缓存数据集路径
   • 支持分别指定训练集和验证集路径

关键设计点

1. 统一的目录结构

   • 导出和加载都使用 train/ 和 val/ 子目录结构
   • 保证导出和加载的一致性

2. S3/OSS 支持

   • 导出：本地临时目录 → S3/OSS
   • 加载：直接从 S3/OSS 读取（HuggingFace datasets 支持）

3. 兼容性

   • 兼容旧格式（直接加载根目录）
   • 自动检测并适配不同目录结构

4. 错误处理

   • 验证集不存在时跳过
   • 远程路径加载失败时给出警告

使用示例

导出：

bash
展开代码
swift export --to_cached_dataset true \
    --output_dir s3://bucket/my_dataset \
    --dataset my_dataset

加载：

yaml
展开代码
cached_dataset:
  - s3://bucket/my_dataset  # 自动加载 train/ 子目录
cached_val_dataset:
  - s3://bucket/my_dataset  # 自动加载 val/ 子目录

或使用采样：

yaml
展开代码
cached_dataset:
  - s3://bucket/my_dataset#10000  # 只加载10000条样本

这两个实现配合使用，实现了数据集的预处理缓存和高效加载。