# CoT Data Generation for Embedding Model

使用本地 Qwen3.5-397B-A17B (vLLM) 为电商图搜 embedding 模型生成 Chain-of-Thought 数据。

## 文件说明

| 文件 | 功能 |
|------|------|
| `config.py` | 配置（模型路径、OSS、prompt模板等） |
| `generate_cot.py` | 主生成脚本（vLLM batch推理 + 断点续传） |
| `retry_failed.py` | 重试失败记录 |
| `merge_shards.py` | 合并分片输出 |
| `validate_cot.py` | 质量验证与统计 |
| `run.sh` | 启动脚本 |

## 快速开始

### 1. 安装依赖

```bash
pip install vllm oss2 pillow tqdm numpy
```

### 2. 小批量测试

```bash
cd /home/lyq494390/multi_vlm2vec_train/cot_gen
bash run.sh --test
```

### 3. 全量生成

```bash
# 全量（支持断点续传，中断后重新运行自动跳过已处理的）
bash run.sh

# 指定范围
bash run.sh --start 0 --end 100000

# 指定TP size
bash run.sh --tp 4
```

### 4. 验证质量

```bash
python validate_cot.py --cot_path /data/oss_bucket_0/liqi/dataset/embedding_pf/cot_output/zhaoxiangsi_train_cot.jsonl --num_samples 10
```

### 5. 重试失败记录

```bash
python retry_failed.py --failed_path /data/oss_bucket_0/liqi/dataset/embedding_pf/cot_output/failed_records.jsonl
```

## 输出格式

每行一条 JSONL：

```json
{
    "rn": "037a6b4e192533b2e0dc9ab3fb692806",
    "query_tfsid": "O1CN01velvra1klBlaW3SfK_!!4611686018427387507-0-imgsearch.jpg",
    "query_bbox": "22 114 414 790",
    "instruction": "更多类似碎花半身裙",
    "cot": "图片中展示了一条碎花印花的A字半身裙，面料轻薄透气，裙身布满小碎花图案...",
    "prompt_index": 0
}
```

## 资源说明

- 模型：`/data/oss_bucket_0/wuxiangyang/model/Qwen3.5-397B-A17B`
- 4x A800 80GB，tensor_parallel=4
- 图片通过 OSS bucket `pailitao-ai` 读取
- 原始数据约 80 万条记录，每条约 10 个 instruction，展开后约 800 万条 CoT 样本

## 注意事项

- 首次运行建议用 `bash run.sh --test` 验证图片加载和模型推理是否正常
- 支持断点续传：中断后重新运行会自动跳过已生成的记录
- 如果 OSS 图片路径与实际不符，修改 `generate_cot.py` 中 `get_image_from_oss` 的 `oss_path`
- `BATCH_SIZE` 默认64，如果显存不足可在 `config.py` 中调小