File size: 1,933 Bytes
c35c7ce | 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 | #!/usr/bin/env python3
from __future__ import annotations
import argparse
import re
from pathlib import Path
from PIL import Image
from puker_judge_utils import (
BASE_MODEL_ID,
BINARY_PROMPT,
MODEL_REPO_ID,
encode_image_prompt,
generate_answer,
load_adapter,
print_json,
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Judge one assembled playing-card candidate as VALID/INVALID."
)
parser.add_argument("image", type=Path)
parser.add_argument("--repo-id", default=MODEL_REPO_ID)
parser.add_argument("--base-model-id", default=BASE_MODEL_ID)
parser.add_argument(
"--int4",
action="store_true",
help="Use bitsandbytes NF4 weights with BF16 compute.",
)
return parser.parse_args()
def main() -> None:
args = parse_args()
model, processor, device, timings = load_adapter(
"binary_adapter",
repo_id=args.repo_id,
base_model_id=args.base_model_id,
int4=args.int4,
)
with Image.open(args.image) as source:
image = source.convert("RGB").copy()
inputs = encode_image_prompt(
processor,
image,
BINARY_PROMPT,
device,
)
raw_output, generation_seconds = generate_answer(
model,
processor,
inputs,
)
match = re.search(r"\b(INVALID|VALID)\b", raw_output.upper())
if match is None:
raise SystemExit(f"Model returned an invalid answer: {raw_output!r}")
print_json(
{
"prediction": match.group(1),
"raw_output": raw_output,
"image": str(args.image.resolve()),
"quantization": "int4-nf4" if args.int4 else "bf16",
"generation_seconds": round(generation_seconds, 4),
**{key: round(value, 4) for key, value in timings.items()},
}
)
if __name__ == "__main__":
main()
|