| """Reference-free confidence of recognizer outputs: mean negative log-likelihood of the (loop-cut) predicted |
| transliteration under the recognizer that produced it, given the image and the recognition prompt. |
| |
| Input: a prediction file with {"id", "image", "pred", ...}. Output: one line per row, |
| {"id", "nll", "n_tokens"}; nll is null for empty predictions. Reuses the training example builder |
| (train_lora.Examples) with the prediction as the target and without the end-of-turn token. Resumable. |
| """ |
| import argparse |
| import json |
| import math |
| import os |
| import sys |
|
|
| import torch |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from qwen_common import load_model |
| from train_lora import Examples |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--model", required=True) |
| ap.add_argument("--adapter", required=True) |
| ap.add_argument("--input", required=True) |
| ap.add_argument("--image_root", required=True) |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--device", default="cuda") |
| ap.add_argument("--limit", type=int, default=0) |
| args = ap.parse_args() |
| rows = [json.loads(l) for l in open(args.input) if l.strip()] |
| done = set() |
| if os.path.exists(args.out): |
| done = {json.loads(l)["id"] for l in open(args.out) if l.strip()} |
| todo = [dict(r, target=r["pred"] or "") for r in rows if r["id"] not in done] |
| if args.limit: |
| todo = todo[:args.limit] |
| print(f"rows {len(rows)} done {len(done)} todo {len(todo)}", flush=True) |
| if not todo: |
| return |
| processor, model = load_model(args.model, adapter=args.adapter) |
| model.to(args.device).eval() |
| ds = Examples(todo, processor, "ocr", args.image_root) |
| ds.end = "" |
| fout = open(args.out, "a") |
| for i, r in enumerate(todo): |
| rec = {"id": r["id"], "nll": None, "n_tokens": 0} |
| if r["target"].strip(): |
| try: |
| item = ds[i] |
| n = int((item["labels"] != -100).sum()) |
| batch = {k: v.to(args.device) for k, v in item.items()} |
| with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): |
| loss = model(**batch).loss.item() |
| rec.update(nll=loss if math.isfinite(loss) else None, n_tokens=n) |
| except Exception as e: |
| rec["error"] = repr(e)[:200] |
| fout.write(json.dumps(rec) + "\n") |
| fout.flush() |
| if (i + 1) % 50 == 0: |
| print(f"{i + 1}/{len(todo)}", flush=True) |
| print("DONE", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
