| """Load a Glint checkpoint (LoRA + head) and answer decision questions about screenshots.
|
| Output: JSON with per-option probabilities (SystemOne-style answers) + latency.
|
| `box` questions are answered Jev-1.14 style: {"type": "box", "box": {"top","bottom","left","right"}, "confidence"}."""
|
| import os; os.environ.setdefault("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "1")
|
| import argparse, json, time
|
| from pathlib import Path
|
| import torch
|
| from PIL import Image
|
| from glint.model import load_checkpoint, default_device
|
| from glint.data import load_jsonl, to_record
|
| from glint.joint_schema_model import collate_records, systemone_answer
|
| from glint.schema import encode_glint_record
|
| from glint.box import box_answer, decode_box
|
|
|
| @torch.inference_mode()
|
| def decide(model, proc, record, device, max_length=4096):
|
| sync = torch.cuda.synchronize if device == "cuda" else (lambda: None)
|
| sync(); t0 = time.perf_counter()
|
| enc = encode_glint_record(proc.tokenizer, record, max_length=max_length, processor=proc)
|
| batch = collate_records([enc], proc.tokenizer.pad_token_id, torch.device(device))
|
| t1 = time.perf_counter()
|
| all_logits, boxes = model.decide(batch)
|
| logits, box_pred = all_logits[0], boxes[0]
|
| sync(); t2 = time.perf_counter()
|
| answers, probs = {}, {}
|
| for q, z in zip(enc.questions, logits):
|
| p = dict(zip(q.option_ids, z.float().softmax(-1).tolist()))
|
| probs[q.question_id] = {k: round(v, 4) for k, v in p.items()}
|
| answers[q.question_id] = systemone_answer(record["questions"][q.question_id], p)
|
| if box_pred is not None:
|
| edges = decode_box(box_pred.logits).cpu(); conf = box_pred.conf_logit.float().sigmoid().cpu()
|
| for i, qid in enumerate(box_pred.question_ids):
|
| answers[qid] = box_answer(edges[i], float(conf[i]))
|
| answers = {qid: answers[qid] for qid in record["questions"] if qid in answers}
|
| return {"answers": answers, "probabilities": probs,
|
| "usage": {"input_tokens": len(enc.input_ids)},
|
| "latency_ms": {"preprocess": round((t1 - t0) * 1000, 1), "forward": round((t2 - t1) * 1000, 1),
|
| "total": round((t2 - t0) * 1000, 1)}}
|
|
|
| def main():
|
| ap = argparse.ArgumentParser()
|
| ap.add_argument("--ckpt", required=True)
|
| ap.add_argument("--image", help="screenshot path (with --questions)")
|
| ap.add_argument("--questions", help="questions JSON file or inline JSON")
|
| ap.add_argument("--state", default="", help="state string or JSON")
|
| ap.add_argument("--jsonl", help="alternatively: run on records of a JSONL file")
|
| ap.add_argument("--limit", type=int, default=5)
|
| ap.add_argument("--max-pixels", type=int, default=None, help="override the checkpoint's max_pixels")
|
| ap.add_argument("--no-merge", action="store_true")
|
| ap.add_argument("--warmup", type=int, default=1)
|
| args = ap.parse_args()
|
| dev = default_device()
|
| model, proc, meta = load_checkpoint(args.ckpt, dev, args.max_pixels, merge=not args.no_merge)
|
| if args.jsonl:
|
| rows = load_jsonl(args.jsonl)[: args.limit]
|
| records = [(r.get("id"), to_record(r), r.get("gold")) for r in rows]
|
| else:
|
| qs = args.questions
|
| qs = json.loads(Path(qs).read_text(encoding="utf-8")) if Path(qs).exists() else json.loads(qs)
|
| try:
|
| state = json.loads(args.state)
|
| except Exception:
|
| state = args.state
|
| records = [(args.image, {"state": state, "images": [Image.open(args.image).convert("RGB")], "questions": qs}, None)]
|
| for _ in range(args.warmup):
|
| decide(model, proc, records[0][1], dev)
|
| for rid, rec, gold in records:
|
| out = {"id": rid, **decide(model, proc, rec, dev)}
|
| if gold:
|
| out["gold"] = gold
|
| print(json.dumps(out, ensure_ascii=False, indent=1), flush=True)
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
