"""Replace teacher distributions with a locally-run Clef / Clef-Flash (open weights, Apache-2.0). No API cost. Clef-Flash is 9B (Qwen3.5-9B backbone, bf16 ~19 GB weights); needs a large GPU / unified memory. Stores probs (``teacher``) and raw logits (``teacher_logits``) in the format train.py / normalize_teacher expect. Usage: python scripts/label_local_clef.py in.jsonl out.jsonl --teacher Cloudflare/clef-flash --max-pixels 430080 [--limit 3] [--offload] (--offload = accelerate device_map="auto" with CPU offload fallback)""" import os; os.environ.setdefault("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "1") # ROCm: faster SDPA import argparse, json, sys, time from pathlib import Path import torch from huggingface_hub import snapshot_download sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from glint.data import load_jsonl, to_record, teacher_questions from glint.model import load_processor from glint.joint_schema_model import encode_record, collate_records, load_release_model def main(): ap = argparse.ArgumentParser(); ap.add_argument("inp"); ap.add_argument("out") ap.add_argument("--teacher", default="Cloudflare/clef-flash") ap.add_argument("--max-pixels", type=int, default=430080, help="430080 ~= 390 image tokens (same as student); 1048576 ~= 987") ap.add_argument("--min-pixels", type=int, default=65536) ap.add_argument("--max-length", type=int, default=16384) ap.add_argument("--limit", type=int, default=None) ap.add_argument("--offload", action="store_true", help="device_map=auto + CPU offload (fallback if bf16 OOMs)") ap.add_argument("--stats", default=None, help="write latency / memory summary JSON here") ap.add_argument("--resume", action="store_true", help="append to an existing out file and skip ids already labeled") args = ap.parse_args() path = snapshot_download(args.teacher) dev = "cuda" if torch.cuda.is_available() else "cpu" t0 = time.time() if args.offload: from transformers import Qwen3_5ForConditionalGeneration from safetensors.torch import load_file from glint.joint_schema_model import ClefModel, JointSchemaHead bb = Qwen3_5ForConditionalGeneration.from_pretrained(path, dtype=torch.bfloat16, device_map="auto", max_memory={0: "12GiB", "cpu": "48GiB"}) bb.config.use_cache = False head = JointSchemaHead(**json.loads((Path(path) / "joint_head_config.json").read_text())) head.load_state_dict(load_file(Path(path) / "joint_head.safetensors"), strict=True) model = ClefModel(bb, head.to(dev, torch.bfloat16)).eval() else: model, _ = load_release_model(path, device=dev) default_ip = json.loads((Path(path) / "processor_config.json").read_text()).get("image_processor", {}) proc = load_processor(path, args.max_pixels, args.min_pixels) # sets size + max/min_pixels like the student load_s = time.time() - t0 print(f"loaded {args.teacher} in {load_s:.1f}s; release processor image cfg={default_ip}; " f"mem_alloc={torch.cuda.memory_allocated()/2**30:.2f}GB" if dev == "cuda" else "", flush=True) rows = load_jsonl(args.inp)[: args.limit] done = set() if args.resume and Path(args.out).exists(): for line in Path(args.out).read_text(encoding="utf-8").splitlines(): try: done.add(json.loads(line)["id"]) except (ValueError, KeyError): pass # truncated last line from a crash rows = [ex for ex in rows if ex.get("id") not in done] print(f"resume: {len(done)} already labeled, {len(rows)} to go", flush=True) lat, img_tok, seq_tok = [], [], [] if dev == "cuda": torch.cuda.reset_peak_memory_stats() img_pad = proc.tokenizer.convert_tokens_to_ids("<|image_pad|>") with open(args.out, "a" if args.resume else "w", encoding="utf-8") as f: for i, ex in enumerate(rows): ts = time.time() rec = to_record(ex); rec["questions"] = teacher_questions(rec["questions"]) # Clef has no `box` type if not rec["questions"]: ex.pop("_base", None); f.write(json.dumps(ex, ensure_ascii=False) + "\n"); f.flush(); continue enc = encode_record(proc.tokenizer, rec, max_length=args.max_length, processor=proc) with torch.inference_mode(): logits = model(collate_records([enc], proc.tokenizer.pad_token_id, torch.device(dev)))[0] if dev == "cuda": torch.cuda.synchronize() dt = time.time() - ts; lat.append(dt) img_tok.append(sum(1 for t in enc.input_ids if t == img_pad)); seq_tok.append(len(enc.input_ids)) ex["teacher"] = {q.question_id: dict(zip(q.option_ids, z.float().softmax(-1).tolist())) for q, z in zip(enc.questions, logits)} ex["teacher_logits"] = {q.question_id: dict(zip(q.option_ids, z.float().tolist())) for q, z in zip(enc.questions, logits)} ex.setdefault("meta", {}).update({"teacher_model": args.teacher, "teacher_is_dummy": False, "teacher_max_pixels": args.max_pixels, "teacher_image_tokens": img_tok[-1]}) ex.pop("_base", None); f.write(json.dumps(ex, ensure_ascii=False) + "\n"); f.flush() print(json.dumps({"i": i, "id": ex.get("id"), "sec": round(dt, 3), "img_tokens": img_tok[-1], "seq_tokens": seq_tok[-1], "argmax": {q: max(d, key=d.get) for q, d in ex["teacher"].items()}, "gold": ex.get("gold")}, ensure_ascii=False), flush=True) srt = sorted(lat[1:] or lat) stats = {"teacher": args.teacher, "n": len(lat), "max_pixels": args.max_pixels, "offload": args.offload, "load_seconds": load_s, "first_sec": lat[0] if lat else None, "median_sec_excl_first": srt[len(srt) // 2] if srt else None, "mean_sec_excl_first": sum(srt) / len(srt) if srt else None, "image_tokens": sorted(set(img_tok)), "seq_tokens_median": sorted(seq_tok)[len(seq_tok) // 2] if seq_tok else None, "peak_mem_gb": torch.cuda.max_memory_allocated() / 2**30 if dev == "cuda" else None, "peak_reserved_gb": torch.cuda.max_memory_reserved() / 2**30 if dev == "cuda" else None} print(json.dumps(stats), flush=True) if args.stats: Path(args.stats).parent.mkdir(parents=True, exist_ok=True); Path(args.stats).write_text(json.dumps(stats, indent=1)) if __name__ == "__main__": main()