eval.py

import argparse, json, asyncio, base64, fitz
from dotenv import load_dotenv
from litellm import acompletion
from src.processor import process_document

async def run():
    load_dotenv()
    p = argparse.ArgumentParser()
    p.add_argument("pdf")
    p.add_argument("-p", "--page", type=int, required=True)
    a = p.parse_args()

    m1, m2 = "cometapi/gpt-5.4", "cometapi/gemini-3-flash"
    j_m = "cometapi/gemini-3.1-pro-preview"

    r1 = await process_document(a.pdf, a.page, m_cls=m1, m_ext=m1, multi_step=True)
    r2 = await process_document(a.pdf, a.page, m_cls=m1, m_ext=m2, multi_step=True)

    doc = fitz.open(a.pdf)
    pg = doc[a.page - 1]
    
    t_dpi = 0
    for i in pg.get_image_info(xrefs=True):
        x, y, x1, y1 = i["bbox"]
        if x1 - x > 0 and y1 - y > 0:
            t_dpi = max(t_dpi, i["width"] * 72 / (x1 - x), i["height"] * 72 / (y1 - y))
    
    m_dpi = int(115200 / max(pg.rect.width, pg.rect.height))
    if t_dpi == 0 and len(pg.get_drawings()) > 0: 
        t_dpi = m_dpi
    
    pix = pg.get_pixmap(dpi=min(max(int(t_dpi), 72), m_dpi))
    b64 = base64.b64encode(pix.tobytes("png")).decode("utf-8")
    doc.close()

    with open("prompt_eval.md", "r", encoding="utf-8") as f:
        sys = f.read()

    u_txt = f"JSON 1:\n{json.dumps(r1['pages'][0], indent=2, ensure_ascii=False)}\n\nJSON 2:\n{json.dumps(r2['pages'][0], indent=2, ensure_ascii=False)}"

    resp = await acompletion(
        model=j_m,
        messages=[
            {"role": "system", "content": sys},
            {"role": "user", "content": [
                {"type": "text", "text": u_txt}, 
                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64}"}}
            ]}
        ]
    )

    with open("eval.md", "a", encoding="utf-8") as f:
        f.write(f"\n```page {a.page}\n{resp.choices[0].message.content.strip()}\n```\n")

if __name__ == "__main__":
    asyncio.run(run())