This repository has been archived by the owner on Sep 11, 2024. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 1
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: create a docling reader for docs as JSON.
New class DoclingJSONReader from BaseDoclingReader to read docs as JSON. Unit test for this class Signed-off-by: Cesar Berrospi Ramis <[email protected]>
- Loading branch information
Showing
3 changed files
with
40 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,25 @@ | ||
# | ||
# Copyright IBM Corp. 2024 - 2024 | ||
# SPDX-License-Identifier: MIT | ||
# | ||
|
||
import json | ||
from typing import Iterable | ||
|
||
from docling_core.types import Document as DLDocument | ||
from llama_index.core.schema import Document as LIDocument | ||
|
||
from quackling.llama_index.readers.base import BaseDoclingReader | ||
|
||
|
||
class DoclingJSONReader(BaseDoclingReader): | ||
def lazy_load_data(self, file_path: str | list[str]) -> Iterable[LIDocument]: | ||
|
||
file_paths = file_path if isinstance(file_path, list) else [file_path] | ||
|
||
for source in file_paths: | ||
with open(source, encoding="utf-8") as file_obj: | ||
data = json.load(file_obj) | ||
dl_doc: DLDocument = DLDocument.model_validate(data) | ||
li_doc: LIDocument = self._create_li_doc_from_dl_doc(dl_doc=dl_doc) | ||
yield li_doc |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,14 @@ | ||
# | ||
# Copyright IBM Corp. 2024 - 2024 | ||
# SPDX-License-Identifier: MIT | ||
# | ||
|
||
from quackling.llama_index.readers import DoclingJSONReader | ||
|
||
|
||
def test_lazy_load_data(): | ||
reader = DoclingJSONReader(parse_type=DoclingJSONReader.ParseType.JSON) | ||
|
||
file_path = "tests/unit/data/0_inp_dl_doc.json" | ||
li_docs = list(reader.lazy_load_data(file_path)) | ||
assert len(li_docs) == 1 |