File size: 6,659 Bytes
e379257 | 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 | """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() |