| |
| """Generate a few ChartQA responses and audit the DyME CoT format.""" |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| from pathlib import Path |
|
|
| import torch |
| from PIL import Image |
| from transformers import AutoProcessor, LlavaOnevisionForConditionalGeneration |
|
|
| _PROJECT_ROOT = Path(__file__).resolve().parent.parent |
| if str(_PROJECT_ROOT) not in sys.path: |
| sys.path.insert(0, str(_PROJECT_ROOT)) |
|
|
| from data_utils.chart.data_collector import prepare_chart_rl_data |
|
|
|
|
| SECTIONS = ("Goal:", "Observation:", "Reasoning:", "Conclusion:", "Answer:") |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--model_path", required=True) |
| parser.add_argument( |
| "--processor_path", |
| default=None, |
| help="Tokenizer/processor source; defaults to --model_path.", |
| ) |
| parser.add_argument("--dataset", required=True) |
| parser.add_argument("--indices", default="0,1,100,1000") |
| parser.add_argument("--max_new_tokens", type=int, default=300) |
| parser.add_argument("--output", default=None) |
| return parser.parse_args() |
|
|
|
|
| def ordered_sections(text: str) -> bool: |
| positions = [text.find(section) for section in SECTIONS] |
| return all(position >= 0 for position in positions) and positions == sorted(positions) |
|
|
|
|
| def main() -> int: |
| args = parse_args() |
| indices = [int(value.strip()) for value in args.indices.split(",") if value.strip()] |
| rows = prepare_chart_rl_data(args.dataset) |
| if not indices or min(indices) < 0 or max(indices) >= len(rows): |
| raise ValueError(f"indices must be within [0, {len(rows) - 1}]") |
|
|
| model_path = str(Path(args.model_path).resolve()) |
| processor_path = str(Path(args.processor_path or args.model_path).resolve()) |
| processor = AutoProcessor.from_pretrained(processor_path, local_files_only=True) |
| processor.tokenizer.padding_side = "left" |
| model = LlavaOnevisionForConditionalGeneration.from_pretrained( |
| model_path, |
| torch_dtype=torch.bfloat16, |
| low_cpu_mem_usage=True, |
| attn_implementation="sdpa", |
| local_files_only=True, |
| ).to("cuda:0") |
| model.eval() |
|
|
| results = [] |
| for index in indices: |
| row = rows[index] |
| image = Image.open(row["image"]).convert("RGB") |
| messages = [ |
| { |
| "role": "user", |
| "content": [ |
| {"type": "image"}, |
| {"type": "text", "text": row["prompt"]}, |
| ], |
| } |
| ] |
| text = processor.apply_chat_template( |
| messages, add_generation_prompt=True, tokenize=False |
| ) |
| inputs = processor(text=[text], images=[image], return_tensors="pt") |
| inputs = {key: value.to("cuda:0") for key, value in inputs.items()} |
| prompt_length = inputs["input_ids"].shape[1] |
| with torch.inference_mode(): |
| output_ids = model.generate( |
| **inputs, |
| max_new_tokens=args.max_new_tokens, |
| do_sample=False, |
| repetition_penalty=1.0, |
| use_cache=True, |
| ) |
| generated_ids = output_ids[:, prompt_length:] |
| response = processor.batch_decode( |
| generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False |
| )[0].strip() |
| result = { |
| "index": index, |
| "question": row["question_wo_prompt"], |
| "gold": row["reference_answer"], |
| "response": response, |
| "ordered_cot_format": ordered_sections(response), |
| "response_words": len(response.split()), |
| } |
| results.append(result) |
| print( |
| f"\n[COT sample {index}] format={result['ordered_cot_format']} " |
| f"words={result['response_words']} gold={result['gold']}\n{response}", |
| flush=True, |
| ) |
|
|
| format_count = sum(result["ordered_cot_format"] for result in results) |
| summary = { |
| "model_path": model_path, |
| "processor_path": processor_path, |
| "samples": len(results), |
| "ordered_cot_count": format_count, |
| "ordered_cot_rate": format_count / len(results), |
| "mean_response_words": sum(item["response_words"] for item in results) / len(results), |
| "results": results, |
| } |
| print("\n[COT summary] " + json.dumps({ |
| key: value for key, value in summary.items() if key != "results" |
| }, ensure_ascii=False), flush=True) |
| if args.output: |
| output_path = Path(args.output) |
| output_path.parent.mkdir(parents=True, exist_ok=True) |
| output_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
