| |
| """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()) |
|
|
