| """ |
| Retry failed records using local vLLM inference. |
| |
| Usage: |
| python retry_failed.py --failed_path /path/to/failed_records.jsonl |
| """ |
|
|
| import os |
| import json |
| import argparse |
|
|
| from PIL import Image |
| from tqdm import tqdm |
| from vllm import LLM, SamplingParams |
|
|
| from config import ( |
| MODEL_PATH, TENSOR_PARALLEL_SIZE, GPU_MEMORY_UTILIZATION, MAX_MODEL_LEN, |
| BATCH_SIZE, MAX_TOKENS, TEMPERATURE, TOP_P, |
| SYSTEM_PROMPT, COT_PROMPT_TEMPLATE, |
| ) |
| from generate_cot import ( |
| load_image_pair, build_messages, save_results, |
| ) |
|
|
|
|
| def main(args): |
| |
| failed_samples = [] |
| with open(args.failed_path, "r", encoding="utf-8") as f: |
| for line in f: |
| if line.strip(): |
| try: |
| item = json.loads(line) |
| item.pop("error", None) |
| failed_samples.append(item) |
| except json.JSONDecodeError: |
| continue |
|
|
| print(f"Loaded {len(failed_samples)} failed records to retry") |
| if not failed_samples: |
| return |
|
|
| base, ext = os.path.splitext(args.failed_path) |
| output_path = base.replace("_failed", "") + "_retried" + ext |
| still_failed_path = base + "_still_failed" + ext |
|
|
| |
| llm = LLM( |
| model=MODEL_PATH, |
| tensor_parallel_size=args.tp or TENSOR_PARALLEL_SIZE, |
| gpu_memory_utilization=GPU_MEMORY_UTILIZATION, |
| max_model_len=MAX_MODEL_LEN, |
| trust_remote_code=True, |
| dtype="bfloat16", |
| limit_mm_per_prompt={"image": 2}, |
| ) |
| sampling_params = SamplingParams( |
| max_tokens=MAX_TOKENS, |
| temperature=TEMPERATURE, |
| top_p=TOP_P, |
| ) |
|
|
| total_success = 0 |
| total_failed = 0 |
| results = [] |
| still_failed = [] |
|
|
| for batch_start in tqdm(range(0, len(failed_samples), BATCH_SIZE), desc="Retrying"): |
| batch = failed_samples[batch_start:batch_start + BATCH_SIZE] |
|
|
| valid_samples = [] |
| batch_inputs = [] |
| for sample in batch: |
| try: |
| original, cropped = load_image_pair(sample["query_tfsid"], sample["query_bbox"]) |
| messages = build_messages(sample, original, cropped) |
| valid_samples.append(sample) |
| batch_inputs.append(messages) |
| except Exception: |
| still_failed.append(sample) |
| total_failed += 1 |
|
|
| if not batch_inputs: |
| continue |
|
|
| try: |
| outputs = llm.chat(messages=batch_inputs, sampling_params=sampling_params) |
| for sample, output in zip(valid_samples, outputs): |
| text = output.outputs[0].text.strip() |
| if text and len(text) >= 10: |
| results.append({ |
| "rn": sample["rn"], |
| "query_tfsid": sample["query_tfsid"], |
| "query_bbox": sample["query_bbox"], |
| "instruction": sample["instruction"], |
| "cot": text, |
| "prompt_index": sample["prompt_index"], |
| }) |
| total_success += 1 |
| else: |
| still_failed.append(sample) |
| total_failed += 1 |
| except Exception: |
| still_failed.extend(valid_samples) |
| total_failed += len(valid_samples) |
|
|
| if results: |
| save_results(results, output_path) |
| if still_failed: |
| with open(still_failed_path, "w", encoding="utf-8") as f: |
| for item in still_failed: |
| f.write(json.dumps(item, ensure_ascii=False) + "\n") |
|
|
| print(f"\nRetry done! Success: {total_success}, Still failed: {total_failed}") |
| if total_success > 0: |
| print(f"Results: {output_path}") |
| if total_failed > 0: |
| print(f"Still failed: {still_failed_path}") |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--failed_path", type=str, required=True) |
| parser.add_argument("--tp", type=int, default=None) |
| args = parser.parse_args() |
| main(args) |
|
|
