# WebDataset支持实施计划

## 1. 核心差异

**ShareGPT4V**: JSON文件 + 图像文件夹（分离）
**VILA WebDataset**: .tar文件（图像+文本打包）+ .wdsmeta.json元数据

## 2. 实施步骤

### Phase 1: 依赖和基础结构（1小时）

1. **添加依赖**
   - `requirements.txt` 添加 `webdataset>=0.2.48`

2. **创建WebDataset数据集类**
   - 文件: `src/lmflow/datasets/webdataset_multi_modal_dataset.py`
   - 类: `CustomWebDatasetMultiModalDataset`
   - 功能: 读取.wdsmeta.json，加载tar文件，转换为messages格式

### Phase 2: 核心实现（2-3小时）

**关键接口**:
```python
class CustomWebDatasetMultiModalDataset(Dataset):
    def __init__(self, wdsmeta_path: str, data_args: DatasetArguments):
        # 1. 读取.wdsmeta.json
        # 2. 构建WebDataset对象
        # 3. 计算总样本数
    
    def __getitem__(self, idx):
        # 1. 从WebDataset读取样本
        # 2. 解析图像（可能是PIL Image或bytes）
        # 3. 解析conversations（可能是JSON字符串或dict）
        # 4. 转换为messages格式（与CustomMultiModalDataset一致）
        return {"messages": messages}
```

**需要确认的WebDataset格式**:
- 图像字段名: `image`, `jpg`, `png`?
- 文本字段名: `json`, `txt`, `conversations`?
- Conversations格式: 已经是标准格式还是需要转换?

### Phase 3: 集成到Dataset类（30分钟）

**修改**: `src/lmflow/datasets/dataset.py`
```python
elif backend == "webdataset_multi_modal":
    from .webdataset_multi_modal_dataset import CustomWebDatasetMultiModalDataset
    raw_dataset = CustomWebDatasetMultiModalDataset(self.dataset_path, data_args)
    self.backend_dataset = raw_dataset
```

### Phase 4: 训练脚本（30分钟）

**创建**: `chengyue_script/dvlm/train_vila_webds.sh`
- 基于 `train_sharegpt4v_vlm_autoresume.sh`
- 修改:
  - `DATASET_WDSMETA` 替代 `DATASET_JSON`
  - 移除 `IMAGE_FOLDER`（图像在tar中）
  - `backend="webdataset_multi_modal"`

## 3. 关键技术点

### 3.1 WebDataset读取
```python
import webdataset as wds
import json

# 从元数据构建数据集
with open(wdsmeta_path) as f:
    meta = json.load(f)

urls = [shard['url'] for shard in meta['shardlist']]
dataset = wds.WebDataset(urls)
```

### 3.2 数据格式转换
- WebDataset样本格式: `{"__key__": "sample_id", "jpg": image_bytes, "json": conversations_str}`
- 需要转换为: `{"messages": [...]}` 格式（与CustomMultiModalDataset一致）

### 3.3 图像处理
- WebDataset中图像可能是bytes，需要转换为PIL Image
- 或者直接传递给processor（如果processor支持bytes）

## 4. 需要调研的问题

1. **VILA WebDataset实际格式**:
   - 字段名是什么？
   - 图像是bytes还是PIL Image？
   - Conversations是JSON字符串还是dict？

2. **性能考虑**:
   - 是否需要缓存？
   - 是否需要多进程加载？

3. **错误处理**:
   - 损坏的tar文件如何处理？
   - 缺失字段如何处理？

## 5. 实施优先级

**P0（必须）**:
- WebDataset读取
- 格式转换
- 集成到Dataset类

**P1（重要）**:
- 错误处理
- 性能优化

**P2（可选）**:
- 缓存机制
- 多进程优化

## 6. 测试计划

1. **单元测试**: 读取单个shard，验证格式转换
2. **集成测试**: 使用小数据集训练1个epoch
3. **性能测试**: 对比JSON vs WebDataset加载速度

## 7. 预计工作量

- **Phase 1-3**: 4-5小时（核心实现）
- **Phase 4**: 30分钟（脚本）
- **测试调试**: 2-3小时
- **总计**: 6-8小时

## 8. 风险点

1. **格式不匹配**: VILA WebDataset格式可能与预期不同
2. **性能问题**: WebDataset加载可能比JSON慢
3. **内存问题**: 大规模数据集可能OOM

## 9. 下一步行动

1. **先调研**: 查看一个实际的WebDataset tar文件内容
2. **再实现**: 基于实际格式实现
3. **后优化**: 性能调优
