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