File size: 2,290 Bytes
0ab2566
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
# 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` 中调小