File size: 8,100 Bytes
df529cc | 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 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 | #!/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())
|