#!/usr/bin/env python3 """Audit independent 7B generations under the exact core-OPD teacher prompt.""" from __future__ import annotations import argparse import json import re import sys from pathlib import Path import torch from transformers import AutoProcessor, LlavaOnevisionForConditionalGeneration PROJECT_ROOT = Path(__file__).resolve().parents[1] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) from config.loader import load_config from data_utils.chart.data_collector import prepare_chart_rl_data from data_utils.chart.teacher_gate import teacher_visible_chartqa_sample from data_utils.rl_prompt import build_rl_prompt_template, resolve_response_contract from opsd_utils.privileged import build_privileged_context from opsd_utils.prompt_builder import tokenize_teacher_prompt HEADINGS = ("Goal", "Observation", "Reasoning", "Conclusion", "Answer") SECTION_RE = re.compile( r"(?ims)^\s*(Goal|Observation|Reasoning|Conclusion|Answer)\s*:\s*(.*?)" r"(?=^\s*(?:Goal|Observation|Reasoning|Conclusion|Answer)\s*:|\Z)" ) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--config", required=True) parser.add_argument("--indices", default="0,1,2,3") parser.add_argument("--max-new-tokens", type=int, default=384) parser.add_argument("--output", default=None) parser.add_argument("--device", default="cuda:0") return parser.parse_args() def classify_structure(text: str) -> dict[str, object]: matches = list(SECTION_RE.finditer(text or "")) names = [match.group(1) for match in matches] nonempty = [bool(match.group(2).strip()) for match in matches] stripped = str(text or "").strip() return { "heading_order": names, "five_part_valid": names == list(HEADINGS) and all(nonempty), "starts_with_answer": stripped.lower().startswith("answer:"), "single_line_answer": len(stripped.splitlines()) == 1 and stripped.lower().startswith("answer:"), } def main() -> int: args = parse_args() config = load_config(args.config) opsd = config["opsd"] contract_cfg = opsd.get("response_contract") or {} if not bool(contract_cfg.get("enabled", False)): raise ValueError("The audit config must enable opsd.response_contract") profile = str(contract_cfg.get("profile") or "").strip() resolve_response_contract(profile) providers = list(opsd.get("privileged_providers") or []) if providers != ["visual_facts_deplot"]: raise ValueError( "Core teacher audit expects exactly visual_facts_deplot; " f"got {providers!r}" ) prompt_template = build_rl_prompt_template( config["rl"].get("answer_flag"), response_profile=profile ) rows = prepare_chart_rl_data( config["dataset"]["train_dataset"], prompt_template=prompt_template ) indices = [int(value.strip()) for value in args.indices.split(",") if value.strip()] if not indices or min(indices) < 0 or max(indices) >= len(rows): raise ValueError(f"indices must be within [0, {len(rows) - 1}]") student_path = str(config["model"]["pretrained_model_path"]) teacher_path = str(config["model"]["teacher_model_path"]) processor = AutoProcessor.from_pretrained(student_path, local_files_only=True) processor.tokenizer.padding_side = "left" model = LlavaOnevisionForConditionalGeneration.from_pretrained( teacher_path, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, attn_implementation="sdpa", local_files_only=True, ).to(args.device) model.eval() results: list[dict[str, object]] = [] for index in indices: row = rows[index] visible = teacher_visible_chartqa_sample(row) suffix, teacher_images = build_privileged_context( visible, providers, privileged_profile=opsd.get("privileged_profile", "hybrid"), crop_cfg=opsd.get("privileged_image") or {}, opsd_config=opsd, ) qwen_hint = str(row.get("hint") or "").strip() if qwen_hint and qwen_hint in suffix: raise RuntimeError(f"row {index}: Qwen rewrite leaked into teacher suffix") if "[Reference Answer]" in suffix or "[Reference Reasoning]" in suffix: raise RuntimeError(f"row {index}: reference annotation leaked into teacher suffix") batch = tokenize_teacher_prompt( processor, row["prompt"], suffix, teacher_images, ) model_inputs = {} for key, value in batch.items(): if not isinstance(value, torch.Tensor): continue value = value.to(args.device) if value.is_floating_point(): value = value.to(torch.bfloat16) model_inputs[key] = value prompt_length = int(model_inputs["input_ids"].shape[1]) with torch.inference_mode(): generated = model.generate( **model_inputs, max_new_tokens=int(args.max_new_tokens), do_sample=False, repetition_penalty=1.0, use_cache=True, pad_token_id=processor.tokenizer.pad_token_id, eos_token_id=processor.tokenizer.eos_token_id, ) completion_ids = generated[:, prompt_length:] response = processor.batch_decode( completion_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False, )[0].strip() structure = classify_structure(response) first_token_id = ( int(completion_ids[0, 0].item()) if completion_ids.shape[1] else None ) first_token = ( processor.tokenizer.decode( [first_token_id], skip_special_tokens=False, clean_up_tokenization_spaces=False, ) if first_token_id is not None else "" ) record = { "index": index, "question": row["question_wo_prompt"], "reference_answer_for_audit_only": row["reference_answer"], "response": response, "completion_tokens": int(completion_ids.shape[1]), "response_words": len(response.split()), "first_token_id": first_token_id, "first_token": first_token, **structure, } results.append(record) print( f"\n[CORE-TEACHER sample={index}] five_part={record['five_part_valid']} " f"single_line_answer={record['single_line_answer']} " f"tokens={record['completion_tokens']} first_token={first_token!r}\n{response}", flush=True, ) summary = { "config": str(Path(args.config).resolve()), "student_processor": student_path, "teacher_model": teacher_path, "response_contract_profile": profile, "providers": providers, "samples": len(results), "five_part_rate": sum(bool(item["five_part_valid"]) for item in results) / len(results), "single_line_answer_rate": sum( bool(item["single_line_answer"]) for item in results ) / len(results), "mean_completion_tokens": sum( int(item["completion_tokens"]) for item in results ) / len(results), "results": results, } print( "\n[CORE-TEACHER summary] " + json.dumps( {key: value for key, value in summary.items() if key != "results"}, ensure_ascii=False, sort_keys=True, ), flush=True, ) if args.output: output = Path(args.output) output.parent.mkdir(parents=True, exist_ok=True) output.write_text( json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8" ) return 0 if summary["five_part_rate"] == 1.0 else 2 if __name__ == "__main__": raise SystemExit(main())