fix: fix legacy doc ref (#162)

Signed-off-by: Panos Vagenas <35837085+vagenas@users.noreply.github.com>
DS4SD · Oct 18, 2024 · eb154a1 · eb154a1
1 parent 77fa1db
commit eb154a1
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 15 deletions.
diff --git a/docling/datamodel/document.py b/docling/datamodel/document.py
@@ -6,12 +6,6 @@
 from typing import TYPE_CHECKING, Dict, Iterable, List, Optional, Tuple, Type, Union
 
 import filetype
-from docling_core.types import BaseText
-from docling_core.types import Document as DsDocument
-from docling_core.types import DocumentDescription as DsDocumentDescription
-from docling_core.types import FileInfoObject as DsFileInfoObject
-from docling_core.types import PageDimensions, PageReference, Prov, Ref
-from docling_core.types import Table as DsSchemaTable
 from docling_core.types.doc import (
     DocItem,
     DocItemLabel,
@@ -22,7 +16,22 @@
     TextItem,
 )
 from docling_core.types.doc.document import ListItem
-from docling_core.types.legacy_doc.base import Figure, GlmTableCell, TableCell
+from docling_core.types.legacy_doc.base import (
+    BaseText,
+    Figure,
+    GlmTableCell,
+    PageDimensions,
+    PageReference,
+    Prov,
+    Ref,
+)
+from docling_core.types.legacy_doc.base import Table as DsSchemaTable
+from docling_core.types.legacy_doc.base import TableCell
+from docling_core.types.legacy_doc.document import (
+    CCSDocumentDescription as DsDocumentDescription,
+)
+from docling_core.types.legacy_doc.document import CCSFileInfoObject as DsFileInfoObject
+from docling_core.types.legacy_doc.document import ExportedCCSDocument as DsDocument
 from docling_core.utils.file import resolve_file_source
 from pydantic import BaseModel
 from typing_extensions import deprecated

diff --git a/docling/models/ds_glm_model.py b/docling/models/ds_glm_model.py
@@ -5,15 +5,23 @@
 from deepsearch_glm.nlp_utils import init_nlp_model
 from deepsearch_glm.utils.doc_utils import to_docling_document
 from deepsearch_glm.utils.load_pretrained_models import load_pretrained_nlp_models
-from docling_core.types import BaseText
-from docling_core.types import Document as DsDocument
-from docling_core.types import DocumentDescription as DsDocumentDescription
-from docling_core.types import FileInfoObject as DsFileInfoObject
-from docling_core.types import PageDimensions, PageReference, Prov, Ref
-from docling_core.types import Table as DsSchemaTable
 from docling_core.types.doc import BoundingBox, CoordOrigin, DoclingDocument
 from docling_core.types.legacy_doc.base import BoundingBox as DsBoundingBox
-from docling_core.types.legacy_doc.base import Figure, TableCell
+from docling_core.types.legacy_doc.base import (
+    Figure,
+    PageDimensions,
+    PageReference,
+    Prov,
+    Ref,
+)
+from docling_core.types.legacy_doc.base import Table as DsSchemaTable
+from docling_core.types.legacy_doc.base import TableCell
+from docling_core.types.legacy_doc.document import BaseText
+from docling_core.types.legacy_doc.document import (
+    CCSDocumentDescription as DsDocumentDescription,
+)
+from docling_core.types.legacy_doc.document import CCSFileInfoObject as DsFileInfoObject
+from docling_core.types.legacy_doc.document import ExportedCCSDocument as DsDocument
 from PIL import ImageDraw
 from pydantic import BaseModel, ConfigDict
 

diff --git a/tests/verify_utils.py b/tests/verify_utils.py
@@ -3,8 +3,8 @@
 from pathlib import Path
 from typing import List
 
-from docling_core.types import Document as DsDocument
 from docling_core.types.doc import DoclingDocument
+from docling_core.types.legacy_doc.document import ExportedCCSDocument as DsDocument
 from pydantic import TypeAdapter
 from pydantic.json import pydantic_encoder