| """Read page images through `serve.sh`, one request per page. |
| |
| python examples/client.py page1.png page2.jpg ... # Markdown (transcription mode) |
| python examples/client.py --mode grounding page1.png # Markdown with block boxes |
| python examples/client.py --url http://127.0.0.1:8000/v1 --out pages.jsonl --concurrency 64 *.png |
| |
| Pages are downscaled to a 2,048 px longest edge (the training and evaluation resolution; 1,540 px is |
| faster and slightly less accurate). Greedy decoding, at most 4,096 output tokens. A page that starts to |
| loop is trimmed after the first copy of the repeated span (loop_guard.scan), finish = "loop". |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import asyncio |
| import base64 |
| import io |
| import json |
| import os |
| import sys |
|
|
| from openai import AsyncOpenAI |
| from PIL import Image |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from loop_guard import scan |
|
|
|
|
| def encode(path: str, edge: int) -> str: |
| im = Image.open(path).convert("RGB") |
| if max(im.size) > edge: |
| im.thumbnail((edge, edge), Image.LANCZOS) |
| buf = io.BytesIO() |
| im.save(buf, format="PNG") |
| return base64.b64encode(buf.getvalue()).decode() |
|
|
|
|
| async def read(cli: AsyncOpenAI, path: str, mode: str, edge: int, sem: asyncio.Semaphore) -> dict: |
| async with sem: |
| b64 = await asyncio.to_thread(encode, path, edge) |
| content = [{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64}"}}] |
| if mode == "grounding": |
| content.append({"type": "text", "text": "grounding"}) |
| resp = await cli.chat.completions.create(model="peppercorn", temperature=0, max_tokens=4096, |
| messages=[{"role": "user", "content": content}]) |
| text, finish = resp.choices[0].message.content or "", resp.choices[0].finish_reason |
| cut = scan(text) |
| if cut is not None: |
| text, finish = text[:cut], "loop" |
| return {"image": path, "text": text, "finish": finish, "tokens": resp.usage.completion_tokens} |
|
|
|
|
| async def main() -> None: |
| ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) |
| ap.add_argument("images", nargs="+") |
| ap.add_argument("--url", default="http://127.0.0.1:8000/v1") |
| ap.add_argument("--mode", choices=["text", "grounding"], default="text") |
| ap.add_argument("--edge", type=int, default=2048) |
| ap.add_argument("--concurrency", type=int, default=32) |
| ap.add_argument("--out", help="write one JSON line per page here instead of printing Markdown") |
| a = ap.parse_args() |
| cli = AsyncOpenAI(base_url=a.url, api_key="x", timeout=3600) |
| sem = asyncio.Semaphore(a.concurrency) |
| rows = await asyncio.gather(*(read(cli, p, a.mode, a.edge, sem) for p in a.images)) |
| if a.out: |
| with open(a.out, "w", encoding="utf-8") as fh: |
| for r in rows: |
| fh.write(json.dumps(r, ensure_ascii=False) + "\n") |
| else: |
| for r in rows: |
| print(f"<!-- {r['image']} finish={r['finish']} tokens={r['tokens']} -->\n{r['text']}\n") |
|
|
|
|
| if __name__ == "__main__": |
| asyncio.run(main()) |
|
|
