code/multi_vlm2vec_train/cot_gen/README.md

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. 安装依赖

pip install vllm oss2 pillow tqdm numpy

2. 小批量测试

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

3. 全量生成

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

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

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

4. 验证质量

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

5. 重试失败记录

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

输出格式

每行一条 JSONL：

{
    "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 中调小