File size: 4,003 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 | """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") # ROCm: faster SDPA (see RESULTS.md); set to 0 to disable
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} # request order
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()
|