| """用本地大 Qwen-VL 模型（vLLM 离线批推）为 listwise reranker 训练集构造 CoT 数据。 |
| |
| 输入: 训练集 jsonl（格式同 train/train_reranker_qwen3vl.py 所用数据） |
| 输出: jsonl，每行 = 原样本 + 实际使用的 docs 顺序 + cot（思考过程）+ cot_answer（标签数字串） |
| |
| 本地用法见 cot_gen/run_gen_cot.sh |
| Nebula 用法见 scripts/submit/submit_cot_gen.sh： |
| nebula 按每卡一进程拉起本脚本，WORLD_SIZE > 1 时自动按 rank 分片， |
| 每个进程用 tp=1 独占一张卡，各自写 <output>.shard<rank>-<world>.jsonl，跑完再合并。 |
| """ |
| import os |
|
|
| os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") |
|
|
| import argparse |
| import glob |
| import json |
| import random |
| import re |
| import sys |
|
|
| current_file_path = os.path.dirname(os.path.abspath(__file__)) |
| sys.path.append(os.path.join(current_file_path, "../")) |
|
|
|
|
| def get_dist_env(): |
| rank = int(os.environ.get("RANK", 0)) |
| world_size = int(os.environ.get("WORLD_SIZE", 1)) |
| local_rank = int(os.environ.get("LOCAL_RANK", rank)) |
| return rank, world_size, local_rank |
|
|
|
|
| def bind_single_gpu(local_rank: int): |
| """分片模式下每个进程只用一张卡；调用点必须早于 torch/vllm 的 import。""" |
| visible = os.environ.get("CUDA_VISIBLE_DEVICES") |
| if visible: |
| gpu_ids = [x for x in visible.split(",") if x] |
| else: |
| gpu_ids = [str(i) for i in sorted(int(re.findall(r"\d+$", p)[0]) for p in glob.glob("/dev/nvidia[0-9]*"))] |
| if len(gpu_ids) > 1: |
| os.environ["CUDA_VISIBLE_DEVICES"] = gpu_ids[local_rank % len(gpu_ids)] |
| print(f"[rank {local_rank}] CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES')}") |
|
|
|
|
| def extract_prompt(full_prompt: str, pattern: str): |
| match = re.search(pattern, full_prompt, re.DOTALL) |
| return match.group(1).strip() if match else None |
|
|
|
|
| def load_data(data_path: str): |
| if data_path.endswith(".json"): |
| return json.load(open(data_path)) |
| return [json.loads(line) for line in open(data_path)] |
|
|
|
|
| def select_docs(item, num_docs: int, shuffle: bool, rng: random.Random): |
| docs = list(item["docs"]) |
| if shuffle: |
| rng.shuffle(docs) |
| if num_docs > 0: |
| docs = docs[:num_docs] |
| return docs |
|
|
|
|
| def build_cot_messages(sys_prompt, user_prompt, item, docs, max_pixels): |
| """将 user_prompt 中的 <image>/<instruction>/<doc_list>/<gt_labels> 占位符展开成 qwen-vl messages。""" |
| labels = [doc["label"] for doc in docs] |
| gt_str = "".join(str(l) for l in labels) |
|
|
| content = [] |
| for segment in re.split(r"(<image>|<instruction>|<doc_list>|<gt_labels>)", user_prompt): |
| if segment == "<image>": |
| content.append({"type": "image", "image": item["query_image"], "max_pixels": max_pixels}) |
| elif segment == "<instruction>": |
| content.append({"type": "text", "text": item.get("instruction", "")}) |
| elif segment == "<doc_list>": |
| for i, doc in enumerate(docs): |
| content.append({"type": "text", "text": f"\n{i}: "}) |
| content.append({"type": "image", "image": doc["crop_image"], "max_pixels": max_pixels}) |
| content.append({"type": "text", "text": f" {doc['doc_title']}"}) |
| content.append({"type": "text", "text": "\n"}) |
| elif segment == "<gt_labels>": |
| lines = [f"{i}: {l}" for i, l in enumerate(labels)] |
| content.append({"type": "text", "text": "\n".join(lines) + f"\n（拼接后应为：{gt_str}）"}) |
| else: |
| content.append({"type": "text", "text": segment}) |
|
|
| return [ |
| {"role": "system", "content": sys_prompt}, |
| {"role": "user", "content": content}, |
| ] |
|
|
|
|
| ANSWER_RE = re.compile(r"最终答案[：:]\s*([0-3]+)") |
|
|
|
|
| def parse_output(text: str): |
| """返回 (cot, answer)，解析失败返回 (None, None)。""" |
| matches = list(ANSWER_RE.finditer(text)) |
| if not matches: |
| return None, None |
| last = matches[-1] |
| answer = last.group(1) |
|
|
| |
| cot = text[: last.start()].replace("<think>", "").replace("</think>", "").strip() |
| if "【逐项分析】" not in cot: |
| return None, answer |
| return cot, answer |
|
|
|
|
| LEAK_RE = re.compile( |
| r"参考标签|参考答案|标注员|给定的?标签|给定的?答案|已知标签|正确答案|所给的?标签|提供的标签|ground\s*truth", |
| re.IGNORECASE, |
| ) |
|
|
|
|
| def is_valid(cot, answer, gt_str): |
| if cot is None or answer is None: |
| return False |
| if answer != gt_str: |
| return False |
| if LEAK_RE.search(cot): |
| return False |
| return True |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--model-path", type=str, required=True, help="本地大 qwen-vl 模型权重路径") |
| parser.add_argument("--train-data-path", type=str, required=True) |
| parser.add_argument("--output-path", type=str, required=True) |
| parser.add_argument("--failed-path", type=str, default=None, help="多次重试仍失败的样本记录，默认 output_path + .failed") |
| parser.add_argument("--prompt-path", type=str, default=os.path.join(current_file_path, "cot_prompt.jinja")) |
| parser.add_argument("--num-docs", type=int, default=10, help="每个样本使用的候选数，<=0 表示全部") |
| parser.add_argument("--shuffle-docs", action="store_true", help="是否先 shuffle 再截取（与训练随机采样对齐）") |
| parser.add_argument("--seed", type=int, default=42) |
| parser.add_argument("--max-crop-tokens", type=int, default=80, help="单图最大视觉 token 数，与训练配置对齐") |
| parser.add_argument("--patch-size", type=int, default=16) |
| parser.add_argument("--tensor-parallel-size", type=int, default=8) |
| parser.add_argument("--max-model-len", type=int, default=16384) |
| parser.add_argument("--max-new-tokens", type=int, default=4096) |
| parser.add_argument("--gpu-memory-utilization", type=float, default=0.9) |
| parser.add_argument("--max-retries", type=int, default=3, help="答案与 GT 不一致时的重试轮数") |
| parser.add_argument("--retry-temperature", type=float, default=0.7) |
| parser.add_argument("--batch-size", type=int, default=64, help="每轮送入 vLLM 的样本数") |
| parser.add_argument("--shard-id", type=int, default=None, help="不传则取环境变量 RANK") |
| parser.add_argument("--num-shards", type=int, default=None, help="不传则取环境变量 WORLD_SIZE") |
| parser.add_argument("--test", type=int, nargs="?", const=20, default=None, |
| help="试跑模式：只处理数据集前 N 条（默认 20），结果写到 <output>.test.jsonl") |
| args = parser.parse_args() |
|
|
| rank, world_size, local_rank = get_dist_env() |
| if args.num_shards is None: |
| args.num_shards = world_size |
| if args.shard_id is None: |
| args.shard_id = rank |
| if args.num_shards > 1: |
| |
| bind_single_gpu(local_rank) |
| args.tensor_parallel_size = 1 |
|
|
| from vllm import LLM, SamplingParams |
| from transformers import AutoProcessor |
| from collators.vision_process import process_vision_info |
|
|
| if args.test: |
| root, ext = os.path.splitext(args.output_path) |
| args.output_path = f"{root}.test{ext}" |
| if args.num_shards > 1: |
| root, ext = os.path.splitext(args.output_path) |
| args.output_path = f"{root}.shard{args.shard_id}-{args.num_shards}{ext}" |
| failed_path = args.failed_path or args.output_path + ".failed" |
| os.makedirs(os.path.dirname(os.path.abspath(args.output_path)), exist_ok=True) |
|
|
| with open(args.prompt_path, encoding="utf-8") as f: |
| full_prompt = f.read() |
| sys_prompt = extract_prompt(full_prompt, r"<sys_start>(.*?)<sys_end>") |
| user_prompt = extract_prompt(full_prompt, r"<user_start>(.*?)<user_end>") |
| assert sys_prompt and user_prompt, "prompt 模板缺少 <sys_start>/<user_start> 块" |
|
|
| max_pixels = args.max_crop_tokens * (args.patch_size * 2) ** 2 |
|
|
| data = load_data(args.train_data_path) |
| print(f"Load {len(data)} instances from {args.train_data_path}") |
| if args.test: |
| data = data[: args.test] |
| print(f"[TEST MODE] only generate for first {len(data)} samples -> {args.output_path}") |
|
|
| done_ids = set() |
| if os.path.exists(args.output_path): |
| for line in open(args.output_path): |
| try: |
| done_ids.add(json.loads(line)["sample_idx"]) |
| except Exception: |
| pass |
| print(f"Resume: {len(done_ids)} samples already done") |
|
|
| rng = random.Random(args.seed) |
| pending = [] |
| for idx, item in enumerate(data): |
| docs = select_docs(item, args.num_docs, args.shuffle_docs, rng) |
| if idx % args.num_shards != args.shard_id or idx in done_ids: |
| continue |
| gt_str = "".join(str(doc["label"]) for doc in docs) |
| messages = build_cot_messages(sys_prompt, user_prompt, item, docs, max_pixels) |
| pending.append((idx, item, docs, gt_str, messages)) |
| print(f"Shard {args.shard_id}/{args.num_shards}: {len(pending)} samples to generate") |
| if not pending: |
| return |
|
|
| max_images = 1 + max(len(d) for _, _, d, _, _ in pending) |
| llm = LLM( |
| model=args.model_path, |
| tensor_parallel_size=args.tensor_parallel_size, |
| max_model_len=args.max_model_len, |
| gpu_memory_utilization=args.gpu_memory_utilization, |
| limit_mm_per_prompt={"image": max_images}, |
| trust_remote_code=True, |
| ) |
| processor = AutoProcessor.from_pretrained(args.model_path, trust_remote_code=True) |
|
|
| def to_vllm_input(messages): |
| text = processor.apply_chat_template( |
| messages, tokenize=False, add_generation_prompt=True, enable_thinking=False |
| ) |
| image_inputs, _ = process_vision_info([messages]) |
| return {"prompt": text, "multi_modal_data": {"image": image_inputs}} |
|
|
| out_f = open(args.output_path, "a", encoding="utf-8") |
| failed_f = open(failed_path, "a", encoding="utf-8") |
|
|
| n_ok, n_fail = 0, 0 |
| for batch_start in range(0, len(pending), args.batch_size): |
| batch = pending[batch_start: batch_start + args.batch_size] |
| results = {} |
| last_text = {} |
|
|
| remaining = list(batch) |
| for attempt in range(args.max_retries + 1): |
| if not remaining: |
| break |
| temperature = 0.0 if attempt == 0 else args.retry_temperature |
| sampling_params = SamplingParams( |
| temperature=temperature, |
| top_p=0.95 if temperature > 0 else 1.0, |
| max_tokens=args.max_new_tokens, |
| seed=args.seed + attempt, |
| ) |
| vllm_inputs = [to_vllm_input(m) for _, _, _, _, m in remaining] |
| outputs = llm.generate(vllm_inputs, sampling_params) |
|
|
| still_failed = [] |
| for (idx, item, docs, gt_str, messages), out in zip(remaining, outputs): |
| text = out.outputs[0].text |
| last_text[idx] = text |
| cot, answer = parse_output(text) |
| if is_valid(cot, answer, gt_str): |
| results[idx] = (cot, answer) |
| else: |
| still_failed.append((idx, item, docs, gt_str, messages)) |
| remaining = still_failed |
| if remaining: |
| print(f"[batch {batch_start}] attempt {attempt}: {len(remaining)} samples need retry") |
|
|
| for idx, item, docs, gt_str, messages in batch: |
| record = dict(item) |
| record["docs"] = docs |
| record["sample_idx"] = idx |
| if idx in results: |
| cot, answer = results[idx] |
| record["cot"] = cot |
| record["cot_answer"] = answer |
| out_f.write(json.dumps(record, ensure_ascii=False) + "\n") |
| n_ok += 1 |
| else: |
| record["gt_answer"] = gt_str |
| record["raw_output"] = last_text.get(idx) |
| failed_f.write(json.dumps(record, ensure_ascii=False) + "\n") |
| n_fail += 1 |
| out_f.flush() |
| failed_f.flush() |
| print(f"Progress: {batch_start + len(batch)}/{len(pending)}, ok={n_ok}, fail={n_fail}") |
|
|
| out_f.close() |
| failed_f.close() |
| print(f"Done. success={n_ok}, failed={n_fail}, output={args.output_path}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
