| # 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` 中调小 |
|
|
