mirror of
https://github.com/DS4SD/docling.git
synced 2025-12-16 16:48:21 +00:00
fix: enrichment models batch size and expose picture classifier (#878)
* expose picture classifier in CLI Signed-off-by: Michele Dolfi <dol@zurich.ibm.com> * use different batch size in each model Signed-off-by: Michele Dolfi <dol@zurich.ibm.com> * remove batch size from CLI Signed-off-by: Michele Dolfi <dol@zurich.ibm.com> * cleanup imports Signed-off-by: Michele Dolfi <dol@zurich.ibm.com> --------- Signed-off-by: Michele Dolfi <dol@zurich.ibm.com>
This commit is contained in:
@@ -6,6 +6,7 @@ from typing_extensions import TypeVar
|
||||
|
||||
from docling.datamodel.base_models import ItemAndImageEnrichmentElement, Page
|
||||
from docling.datamodel.document import ConversionResult
|
||||
from docling.datamodel.settings import settings
|
||||
|
||||
|
||||
class BasePageModel(ABC):
|
||||
@@ -21,6 +22,8 @@ EnrichElementT = TypeVar("EnrichElementT", default=NodeItem)
|
||||
|
||||
class GenericEnrichmentModel(ABC, Generic[EnrichElementT]):
|
||||
|
||||
elements_batch_size: int = settings.perf.elements_batch_size
|
||||
|
||||
@abstractmethod
|
||||
def is_processable(self, doc: DoclingDocument, element: NodeItem) -> bool:
|
||||
pass
|
||||
|
||||
@@ -61,6 +61,7 @@ class CodeFormulaModel(BaseItemAndImageEnrichmentModel):
|
||||
Processes the given batch of elements and enriches them with predictions.
|
||||
"""
|
||||
|
||||
elements_batch_size = 5
|
||||
images_scale = 1.66 # = 120 dpi, aligned with training data resolution
|
||||
expansion_factor = 0.03
|
||||
|
||||
|
||||
Reference in New Issue
Block a user