Spaces:
Running on Zero
Running on Zero
File size: 2,912 Bytes
72b6c80 25c4ebd 72b6c80 25c4ebd 72b6c80 | 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 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | # ruff: noqa: I001
import spaces
import json
import torch
from transformers import AutoModelForImageTextToText, AutoProcessor
processor = AutoProcessor.from_pretrained(
"numind/NuExtract3",
trust_remote_code=True,
)
model = (
AutoModelForImageTextToText.from_pretrained(
"numind/NuExtract3",
attn_implementation="sdpa",
dtype=torch.bfloat16,
trust_remote_code=True,
)
.to("cuda")
.eval()
)
@spaces.GPU
def extract(image, text, template, enable_thinking):
inputs = processor.apply_chat_template(
[
{
"role": "user",
"content": [
{"type": "image", "image": image},
{"type": "text", "text": text},
],
}
],
add_generation_prompt=True,
tokenize=True,
return_dict=True,
return_tensors="pt",
mode="structured",
template=template,
enable_thinking=enable_thinking,
).to(model.device)
with torch.inference_mode():
generated_ids = model.generate(
**inputs,
max_new_tokens=4096,
do_sample=False,
)
output = processor.batch_decode(
generated_ids[:, inputs.input_ids.shape[1] :],
skip_special_tokens=True,
clean_up_tokenization_spaces=False,
)[0].strip()
return json.loads(
output.split("</think>", 1)[1].strip() if enable_thinking else output
)
@spaces.GPU
def generate_template(image):
inputs = processor.apply_chat_template(
[
{
"role": "user",
"content": [
{"type": "image", "image": image},
{
"type": "text",
"text": (
"Create a reusable structured extraction template grounded only "
"in the visible document. Include fields supported by the document, "
"represent repeated records as arrays, use NuExtract template leaf "
"types, and return only the JSON template."
),
},
],
}
],
add_generation_prompt=True,
tokenize=True,
return_dict=True,
return_tensors="pt",
mode="template-generation",
).to(model.device)
with torch.inference_mode():
generated_ids = model.generate(
**inputs,
max_new_tokens=4096,
do_sample=False,
)
return json.dumps(
json.loads(
processor.batch_decode(
generated_ids[:, inputs.input_ids.shape[1] :],
skip_special_tokens=True,
clean_up_tokenization_spaces=False,
)[0].strip()
),
indent=2,
)
|