Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add Nemoretriever-parse (aka Eclair) NIM #230

Open
wants to merge 79 commits into
base: main
Choose a base branch
from
Open
Changes from 1 commit
Commits
Show all changes
79 commits
Select commit Hold shift + click to select a range
3fbcf68
Add doughnut http endpoint
edknv Nov 14, 2024
e5a281d
fix table and chart extraction
edknv Nov 14, 2024
d933a33
handle 202 reponses by repolling status
edknv Nov 18, 2024
af9ca0e
add table format in unit tests
edknv Nov 18, 2024
b5c992d
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Nov 18, 2024
6ad99c0
fix table and image max dimensions
edknv Nov 19, 2024
e511dbf
Merge branch 'feat/doughnut-http-endpoint-1' of github.com:edknv/nv-i…
edknv Nov 19, 2024
b3c632b
add unit tests for the helper
edknv Nov 19, 2024
0deb5c9
add placeholder for url in docker compose
edknv Nov 19, 2024
2e7db6d
add check for empty dataframe in table/chart extraction
edknv Nov 19, 2024
c5b5d3e
clean up doughnut specific confiditions in inference func
edknv Nov 20, 2024
bcb0037
clean up doughnut specific confiditions in inference func
edknv Nov 20, 2024
187d13d
fix unit tests
edknv Nov 20, 2024
51b5bca
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Nov 22, 2024
321052f
Add support for text bounding boxes
edknv Nov 22, 2024
a9a74b7
also add table and image bounding boxes to metadata
edknv Nov 23, 2024
1d532a4
Use seprate environment variable for private endpoint
edknv Dec 2, 2024
e141cad
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Dec 3, 2024
ea057e6
Use seprate environment variable for private endpoint
edknv Dec 4, 2024
be8cc68
Merge branch 'feat/doughnut-http-endpoint-1' of github.com:edknv/nv-i…
edknv Dec 4, 2024
67ff768
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Dec 4, 2024
001a3b5
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Dec 6, 2024
ccb9f8f
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Dec 9, 2024
8088eed
migrate to using NimClient and ModelInterface
edknv Dec 9, 2024
0163880
Merge branch 'feat/doughnut-http-endpoint-1' of github.com:edknv/nv-i…
edknv Dec 9, 2024
792b724
update unit tests
edknv Dec 9, 2024
8f77199
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Dec 9, 2024
f96d406
Merge branch 'main' into feat/doughnut-http-endpoint-1
jdye64 Dec 10, 2024
9a17068
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Dec 10, 2024
71d7212
lint
edknv Dec 10, 2024
e9f021f
fix unit test
edknv Dec 10, 2024
e456d9e
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Dec 10, 2024
a7e0f36
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Dec 11, 2024
7e2e518
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Dec 13, 2024
b833850
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Dec 13, 2024
247c431
Merge branch 'feat/doughnut-http-endpoint-1' of github.com:edknv/nv-i…
edknv Dec 13, 2024
1fc1b67
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Dec 13, 2024
a705846
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Dec 16, 2024
e4e3406
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Dec 17, 2024
0ef5b86
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Dec 20, 2024
a0c584c
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Jan 6, 2025
a0befbf
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Jan 13, 2025
c6b0a35
migrate doughnut config to pydantic 2
edknv Jan 13, 2025
410610e
Merge branch 'feat/doughnut-http-endpoint-1' of github.com:edknv/nv-i…
edknv Jan 13, 2025
23eae8f
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Jan 13, 2025
26b5e50
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Jan 14, 2025
e340c34
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Jan 15, 2025
4f7bfeb
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Jan 22, 2025
1c119ff
Replace nvcf with NIM
edknv Jan 24, 2025
6b7a29e
remove postprocessing
edknv Jan 24, 2025
237dfa4
drop codename
edknv Jan 24, 2025
71d81d2
update tests
edknv Jan 24, 2025
c92ae91
lint
edknv Jan 24, 2025
848d18d
refactor fot batch size 1
edknv Jan 26, 2025
6a962fb
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Feb 6, 2025
db68afb
update model name and batch input
edknv Feb 6, 2025
e9fa05e
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Feb 7, 2025
e97867b
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Feb 10, 2025
4dc708f
Merge branch 'feat/doughnut-http-endpoint-1' of github.com:edknv/nv-i…
edknv Feb 10, 2025
23e901c
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Feb 10, 2025
e99d775
rc5 uses the regular ngc key
edknv Feb 11, 2025
5dea3b7
find and replace all eclair -> nemoretriever_parse
edknv Feb 11, 2025
1c8f4fd
clean up
edknv Feb 11, 2025
1aa002c
more clean up
edknv Feb 11, 2025
d92da02
enable table/chart extraction with yolox
edknv Feb 12, 2025
398b179
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Feb 12, 2025
ee32864
fix cached and paddle batching error
edknv Feb 12, 2025
6511ca6
update batching in nimclient
edknv Feb 12, 2025
41296ed
add option to select table extraction mode
edknv Feb 12, 2025
dfc9fd7
update image tag
edknv Feb 12, 2025
755ceea
fix page arrangement
edknv Feb 12, 2025
1de974d
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Feb 12, 2025
632880f
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Feb 13, 2025
10dea4c
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Feb 13, 2025
631222a
move the new nim to device 1 and new profile nemoretriever-parse
edknv Feb 13, 2025
541a36c
merge fix
edknv Feb 14, 2025
933e511
Merge branch 'main' of https://github.com/NVIDIA/nv-ingest into feat/…
edknv Feb 18, 2025
d3652c3
Merge branch 'main' into feat/doughnut-http-endpoint-1
edknv Feb 19, 2025
0bb6dee
increase number of workers and increase timeout
edknv Feb 19, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
also add table and image bounding boxes to metadata
edknv committed Nov 23, 2024
commit a9a74b739183176fba1fab34363a68d3c9a6f1a3
111 changes: 58 additions & 53 deletions src/nv_ingest/extraction_workflows/pdf/doughnut_helper.py
Original file line number Diff line number Diff line change
@@ -161,69 +161,74 @@ def doughnut(pdf_stream, extract_text: bool, extract_images: bool, extract_table

page_nearby_blocks = {
"text": {"content": [], "bbox": [], "type": []},
"images": {"content": [], "bbox": []},
"structured": {"content": [], "bbox": []},
"images": {"content": [], "bbox": [], "type": []},
"structured": {"content": [], "bbox": [], "type": []},
}

for cls, bbox, txt in zip(classes, bboxes, texts):

if extract_text and (cls in doughnut_utils.ACCEPTED_TEXT_CLASSES):
txt = doughnut_utils.postprocess_text(txt, cls)
transformed_bbox = doughnut_utils.reverse_transform_bbox(
bbox=bbox,
bbox_offset=bbox_offset,
original_width=DEFAULT_MAX_WIDTH,
original_height=DEFAULT_MAX_HEIGHT,
)

if cls in doughnut_utils.ACCEPTED_TEXT_CLASSES:
if identify_nearby_objects:
bbox = doughnut_utils.reverse_transform_bbox(
bbox=bbox,
bbox_offset=bbox_offset,
original_width=DEFAULT_MAX_WIDTH,
original_height=DEFAULT_MAX_HEIGHT,
)
page_nearby_blocks["text"]["content"].append(txt)
page_nearby_blocks["text"]["bbox"].append(bbox)
page_nearby_blocks["text"]["bbox"].append(transformed_bbox)
page_nearby_blocks["text"]["type"].append(cls)

accumulated_text.append(txt)

if extract_tables and (cls == "Table"):
try:
txt = txt.encode().decode("unicode_escape") # remove double backlashes
except UnicodeDecodeError:
pass
bbox = doughnut_utils.reverse_transform_bbox(
bbox=bbox,
bbox_offset=bbox_offset,
original_width=DEFAULT_MAX_WIDTH,
original_height=DEFAULT_MAX_HEIGHT,
)
table = LatexTable(latex=txt, bbox=bbox, max_width=DEFAULT_MAX_WIDTH, max_height=DEFAULT_MAX_HEIGHT)
accumulated_tables.append(table)

if extract_images and (cls == "Picture"):
if page_image is None:
scale_tuple = (DEFAULT_MAX_WIDTH, DEFAULT_MAX_HEIGHT)
padding_tuple = (DEFAULT_MAX_WIDTH, DEFAULT_MAX_HEIGHT)
page_image, *_ = pdfium_pages_to_numpy(
[pages[page_idx]], scale_tuple=scale_tuple, padding_tuple=padding_tuple
)
page_image = page_image[0]

img_numpy = crop_image(page_image, bbox)
if img_numpy is not None:
base64_img = numpy_to_base64(img_numpy)
bbox = doughnut_utils.reverse_transform_bbox(
bbox=bbox,
bbox_offset=bbox_offset,
original_width=DEFAULT_MAX_WIDTH,
original_height=DEFAULT_MAX_HEIGHT,
)
image = Base64Image(
image=base64_img,
bbox=bbox,
width=img_numpy.shape[1],
height=img_numpy.shape[0],
max_width=DEFAULT_MAX_WIDTH,
max_height=DEFAULT_MAX_HEIGHT,
if extract_text:
txt = doughnut_utils.postprocess_text(txt, cls)
accumulated_text.append(txt)

if cls == "Table":
if identify_nearby_objects:
page_nearby_blocks["structured"]["content"].append(txt)
page_nearby_blocks["structured"]["bbox"].append(transformed_bbox)
page_nearby_blocks["structured"]["type"].append(cls)

if extract_tables:
try:
txt = txt.encode().decode("unicode_escape") # remove double backlashes
except UnicodeDecodeError:
pass

table = LatexTable(
latex=txt, bbox=transformed_bbox, max_width=DEFAULT_MAX_WIDTH, max_height=DEFAULT_MAX_HEIGHT
)
accumulated_images.append(image)
accumulated_tables.append(table)

if cls == "Picture":
if identify_nearby_objects:
page_nearby_blocks["images"]["content"].append(txt)
page_nearby_blocks["images"]["bbox"].append(transformed_bbox)
page_nearby_blocks["images"]["type"].append(cls)

if extract_images:
if page_image is None:
scale_tuple = (DEFAULT_MAX_WIDTH, DEFAULT_MAX_HEIGHT)
padding_tuple = (DEFAULT_MAX_WIDTH, DEFAULT_MAX_HEIGHT)
page_image, *_ = pdfium_pages_to_numpy(
[pages[page_idx]], scale_tuple=scale_tuple, padding_tuple=padding_tuple
)
page_image = page_image[0]

img_numpy = crop_image(page_image, bbox)

if img_numpy is not None:
base64_img = numpy_to_base64(img_numpy)
image = Base64Image(
image=base64_img,
bbox=transformed_bbox,
width=img_numpy.shape[1],
height=img_numpy.shape[0],
max_width=DEFAULT_MAX_WIDTH,
max_height=DEFAULT_MAX_HEIGHT,
)
accumulated_images.append(image)

# Construct tables
if extract_tables: