electblake/image-data-extractor

1# ruff: noqa: I0012 3import spaces4 5import json6 7import torch8from transformers import AutoModelForImageTextToText, AutoProcessor9 10 11processor = AutoProcessor.from_pretrained(12    "numind/NuExtract3",13    trust_remote_code=True,14)15 16model = (17    AutoModelForImageTextToText.from_pretrained(18        "numind/NuExtract3",19        attn_implementation="sdpa",20        dtype=torch.bfloat16,21        trust_remote_code=True,22    )23    .to("cuda")24    .eval()25)26 27 28@spaces.GPU29def extract(image, text, template, enable_thinking):30    inputs = processor.apply_chat_template(31        [32            {33                "role": "user",34                "content": [35                    {"type": "image", "image": image},36                    {"type": "text", "text": text},37                ],38            }39        ],40        add_generation_prompt=True,41        tokenize=True,42        return_dict=True,43        return_tensors="pt",44        mode="structured",45        template=template,46        enable_thinking=enable_thinking,47    ).to(model.device)48 49    with torch.inference_mode():50        generated_ids = model.generate(51            **inputs,52            max_new_tokens=4096,53            do_sample=False,54        )55 56    output = processor.batch_decode(57        generated_ids[:, inputs.input_ids.shape[1] :],58        skip_special_tokens=True,59        clean_up_tokenization_spaces=False,60    )[0].strip()61 62    return json.loads(63        output.split("</think>", 1)[1].strip() if enable_thinking else output64    )65 66 67@spaces.GPU68def generate_template(image):69    inputs = processor.apply_chat_template(70        [71            {72                "role": "user",73                "content": [74                    {"type": "image", "image": image},75                    {76                        "type": "text",77                        "text": (78                            "Create a reusable structured extraction template grounded only "79                            "in the visible document. Include fields supported by the document, "80                            "represent repeated records as arrays, use NuExtract template leaf "81                            "types, and return only the JSON template."82                        ),83                    },84                ],85            }86        ],87        add_generation_prompt=True,88        tokenize=True,89        return_dict=True,90        return_tensors="pt",91        mode="template-generation",92    ).to(model.device)93 94    with torch.inference_mode():95        generated_ids = model.generate(96            **inputs,97            max_new_tokens=4096,98            do_sample=False,99        )100 101    return json.dumps(102        json.loads(103            processor.batch_decode(104                generated_ids[:, inputs.input_ids.shape[1] :],105                skip_special_tokens=True,106                clean_up_tokenization_spaces=False,107            )[0].strip()108        ),109        indent=2,110    )111