File size: 4,055 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 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | """
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):
# Load failed records
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
# Init
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)
|