| """Student = Qwen3.5 backbone + LoRA + Clef-style JointSchemaHead (option logits -> softmax per question)
|
| (+ optional BoxHead for Jev-style `box` questions, see glint/box.py)."""
|
| from __future__ import annotations
|
| import json, os
|
| from pathlib import Path
|
| import torch
|
| from safetensors.torch import load_file, save_file
|
| from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration
|
| from peft import LoraConfig, PeftModel, get_peft_model
|
| from .joint_schema_model import JointSchemaHead
|
| from .box import BoxHead
|
| from .glint_model import GlintModel
|
|
|
|
|
| LORA_TARGET = (r".*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj"
|
| r"|in_proj_qkv|in_proj_z|out_proj)")
|
|
|
| def default_device():
|
| return "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
| def parse_dtype(name: str):
|
| return {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[name]
|
|
|
| def load_processor(base: str, max_pixels: int, min_pixels: int = 65536):
|
| proc = AutoProcessor.from_pretrained(base)
|
| ip = proc.image_processor
|
| ip.size = {"shortest_edge": min_pixels, "longest_edge": max_pixels}
|
| if hasattr(ip, "max_pixels"):
|
| ip.max_pixels = max_pixels
|
| if hasattr(ip, "min_pixels"):
|
| ip.min_pixels = min_pixels
|
| return proc
|
|
|
| def build_student(base="Qwen/Qwen3.5-2B", max_pixels=430080, device=None, dtype="bf16",
|
| lora_r=32, lora_alpha=64, lora_dropout=0.05, head_cfg=None, grad_ckpt=False, box_head_cfg=None):
|
| """box_head_cfg: None = no box head (previous behaviour); a dict (may be empty) = BoxHead kwargs."""
|
| device = device or default_device()
|
| td = parse_dtype(dtype)
|
| proc = load_processor(base, max_pixels)
|
| backbone = Qwen3_5ForConditionalGeneration.from_pretrained(base, dtype=td).to(device)
|
| backbone.config.use_cache = False
|
| if grad_ckpt:
|
| backbone.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
|
| backbone.enable_input_require_grads()
|
| backbone = get_peft_model(backbone, LoraConfig(r=lora_r, lora_alpha=lora_alpha, lora_dropout=lora_dropout,
|
| target_modules=LORA_TARGET))
|
| hidden = backbone.config.text_config.hidden_size
|
| head_cfg = dict(head_cfg or {"width": 512, "routing_layers": 2, "layers": 2, "heads": 8, "feedforward": 2048})
|
| head_cfg["hidden_size"] = hidden
|
| head = JointSchemaHead(**head_cfg, dropout=0.1).to(device, td)
|
| box_head = None
|
| if box_head_cfg is not None:
|
| box_head_cfg = {**box_head_cfg, "hidden_size": hidden}
|
| box_head = BoxHead(**box_head_cfg, dropout=0.1).to(device, td)
|
| model = GlintModel(backbone, head, box_head)
|
| meta = {"base_model": base, "max_pixels": max_pixels, "dtype": dtype, "head_config": head_cfg,
|
| "lora": {"r": lora_r, "alpha": lora_alpha, "dropout": lora_dropout, "target_modules": LORA_TARGET}}
|
| if box_head_cfg is not None:
|
| meta["box_head_config"] = box_head_cfg
|
| return model, proc, meta
|
|
|
| def save_checkpoint(model: GlintModel, meta: dict, out_dir: str | os.PathLike, extra: dict | None = None):
|
| out = Path(out_dir); out.mkdir(parents=True, exist_ok=True)
|
| model.language_model.save_pretrained(out / "adapter")
|
| save_file({k: v.detach().cpu().contiguous() for k, v in model.head.state_dict().items()}, out / "head.safetensors")
|
| if getattr(model, "box_head", None) is not None:
|
| save_file({k: v.detach().cpu().contiguous() for k, v in model.box_head.state_dict().items()}, out / "box_head.safetensors")
|
| (out / "glint_config.json").write_text(json.dumps({**meta, **(extra or {})}, indent=2, ensure_ascii=False), encoding="utf-8")
|
|
|
| def load_checkpoint(ckpt_dir: str | os.PathLike, device=None, max_pixels: int | None = None, merge=True):
|
| ckpt = Path(ckpt_dir); device = device or default_device()
|
| meta = json.loads((ckpt / "glint_config.json").read_text(encoding="utf-8"))
|
| td = parse_dtype(meta["dtype"])
|
| proc = load_processor(meta["base_model"], max_pixels or meta["max_pixels"])
|
| backbone = Qwen3_5ForConditionalGeneration.from_pretrained(meta["base_model"], dtype=td).to(device)
|
| backbone.config.use_cache = False
|
| backbone = PeftModel.from_pretrained(backbone, ckpt / "adapter")
|
| if merge:
|
| backbone = backbone.merge_and_unload()
|
| head = JointSchemaHead(**meta["head_config"]).to(device, td)
|
| head.load_state_dict(load_file(ckpt / "head.safetensors"), strict=True)
|
| box_head = None
|
| if meta.get("box_head_config") is not None:
|
| box_head = BoxHead(**meta["box_head_config"]).to(device, td)
|
| box_head.load_state_dict(load_file(ckpt / "box_head.safetensors"), strict=True)
|
| return GlintModel(backbone, head, box_head).eval(), proc, meta
|
|
|
