| """ClefModel + optional BoxHead. Kept free of transformers/peft imports so it can be unit-tested with a mock backbone.""" | |
| from __future__ import annotations | |
| from typing import Any | |
| import torch | |
| from .joint_schema_model import ClefModel, JointSchemaHead | |
| from .box import BoxHead, BoxPrediction | |
| class GlintModel(ClefModel): | |
| """`forward(batch)` returns the joint-schema logits exactly like ClefModel (list per record). | |
| `decide(batch)` returns `(logits, boxes)`; `boxes[i]` is a BoxPrediction or None.""" | |
| def __init__(self, language_model: Any, head: JointSchemaHead, box_head: BoxHead | None = None) -> None: | |
| super().__init__(language_model, head) | |
| self.box_head = box_head | |
| def _base_model(self): | |
| lm = self.language_model | |
| return lm.get_base_model() if hasattr(lm, "get_base_model") else lm | |
| def _hidden(self, batch: dict[str, Any]): | |
| base_model = self._base_model() | |
| media = batch.get("media") or {} | |
| text_model = base_model.model | |
| if not media and hasattr(text_model, "language_model"): | |
| text_model = text_model.language_model | |
| outputs = text_model(input_ids=batch["input_ids"], attention_mask=batch["attention_mask"], | |
| use_cache=False, return_dict=True, **media) | |
| return outputs.last_hidden_state, base_model | |
| def decide(self, batch: dict[str, Any]) -> tuple[list[list[torch.Tensor]], list[BoxPrediction | None]]: | |
| hidden, base_model = self._hidden(batch) | |
| records = batch["records"] | |
| logits: list[list[torch.Tensor]] = [[] for _ in records] | |
| keep = [i for i, r in enumerate(records) if r.questions] | |
| if keep: | |
| sel = keep if len(keep) < len(records) else slice(None) | |
| out = self.head(hidden[sel], batch["input_ids"][sel], batch["attention_mask"][sel], [records[i] for i in keep], | |
| base_model.get_output_embeddings().weight) | |
| for i, z in zip(keep, out): | |
| logits[i] = z | |
| boxes: list[BoxPrediction | None] = [None] * len(records) | |
| if any(getattr(r, "box_questions", ()) for r in records): | |
| if self.box_head is None: | |
| raise ValueError("record has box questions but this checkpoint has no box head (train with --box-head on)") | |
| boxes = self.box_head(hidden, batch["attention_mask"], records) | |
| return logits, boxes | |
| def forward(self, batch: dict[str, Any]) -> list[list[torch.Tensor]]: | |
| return self.decide(batch)[0] | |