| import re |
| from pathlib import Path |
| from typing import Iterable, List, Literal, Optional, Tuple, Union |
|
|
| import numpy as np |
| from docling_core.types.doc import ( |
| CodeItem, |
| DocItemLabel, |
| DoclingDocument, |
| NodeItem, |
| TextItem, |
| ) |
| from docling_core.types.doc.labels import CodeLanguageLabel |
| from PIL import Image |
| from pydantic import BaseModel |
|
|
| from docling.datamodel.base_models import ItemAndImageEnrichmentElement |
| from docling.datamodel.pipeline_options import AcceleratorOptions |
| from docling.models.base_model import BaseItemAndImageEnrichmentModel |
| from docling.utils.accelerator_utils import decide_device |
|
|
|
|
| class CodeFormulaModelOptions(BaseModel): |
| """ |
| Configuration options for the CodeFormulaModel. |
| |
| Attributes |
| ---------- |
| kind : str |
| Type of the model. Fixed value "code_formula". |
| do_code_enrichment : bool |
| True if code enrichment is enabled, False otherwise. |
| do_formula_enrichment : bool |
| True if formula enrichment is enabled, False otherwise. |
| """ |
|
|
| kind: Literal["code_formula"] = "code_formula" |
| do_code_enrichment: bool = True |
| do_formula_enrichment: bool = True |
|
|
|
|
| class CodeFormulaModel(BaseItemAndImageEnrichmentModel): |
| """ |
| Model for processing and enriching documents with code and formula predictions. |
| |
| Attributes |
| ---------- |
| enabled : bool |
| True if the model is enabled, False otherwise. |
| options : CodeFormulaModelOptions |
| Configuration options for the CodeFormulaModel. |
| code_formula_model : CodeFormulaPredictor |
| The predictor model for code and formula processing. |
| |
| Methods |
| ------- |
| __init__(self, enabled, artifacts_path, accelerator_options, code_formula_options) |
| Initializes the CodeFormulaModel with the given configuration options. |
| is_processable(self, doc, element) |
| Determines if a given element in a document can be processed by the model. |
| __call__(self, doc, element_batch) |
| Processes the given batch of elements and enriches them with predictions. |
| """ |
|
|
| _model_repo_folder = "ds4sd--CodeFormula" |
| elements_batch_size = 5 |
| images_scale = 1.66 |
| expansion_factor = 0.03 |
|
|
| def __init__( |
| self, |
| enabled: bool, |
| artifacts_path: Optional[Path], |
| options: CodeFormulaModelOptions, |
| accelerator_options: AcceleratorOptions, |
| ): |
| """ |
| Initializes the CodeFormulaModel with the given configuration. |
| |
| Parameters |
| ---------- |
| enabled : bool |
| True if the model is enabled, False otherwise. |
| artifacts_path : Path |
| Path to the directory containing the model artifacts. |
| options : CodeFormulaModelOptions |
| Configuration options for the model. |
| accelerator_options : AcceleratorOptions |
| Options specifying the device and number of threads for acceleration. |
| """ |
| self.enabled = enabled |
| self.options = options |
|
|
| if self.enabled: |
| device = decide_device(accelerator_options.device) |
|
|
| from docling_ibm_models.code_formula_model.code_formula_predictor import ( |
| CodeFormulaPredictor, |
| ) |
|
|
| if artifacts_path is None: |
| artifacts_path = self.download_models() |
| else: |
| artifacts_path = artifacts_path / self._model_repo_folder |
|
|
| self.code_formula_model = CodeFormulaPredictor( |
| artifacts_path=str(artifacts_path), |
| device=device, |
| num_threads=accelerator_options.num_threads, |
| ) |
|
|
| @staticmethod |
| def download_models( |
| local_dir: Optional[Path] = None, |
| force: bool = False, |
| progress: bool = False, |
| ) -> Path: |
| from huggingface_hub import snapshot_download |
| from huggingface_hub.utils import disable_progress_bars |
|
|
| if not progress: |
| disable_progress_bars() |
| download_path = snapshot_download( |
| repo_id="ds4sd/CodeFormula", |
| force_download=force, |
| local_dir=local_dir, |
| revision="v1.0.1", |
| ) |
|
|
| return Path(download_path) |
|
|
| def is_processable(self, doc: DoclingDocument, element: NodeItem) -> bool: |
| """ |
| Determines if a given element in a document can be processed by the model. |
| |
| Parameters |
| ---------- |
| doc : DoclingDocument |
| The document being processed. |
| element : NodeItem |
| The element within the document to check. |
| |
| Returns |
| ------- |
| bool |
| True if the element can be processed, False otherwise. |
| """ |
| return self.enabled and ( |
| (isinstance(element, CodeItem) and self.options.do_code_enrichment) |
| or ( |
| isinstance(element, TextItem) |
| and element.label == DocItemLabel.FORMULA |
| and self.options.do_formula_enrichment |
| ) |
| ) |
|
|
| def _extract_code_language(self, input_string: str) -> Tuple[str, Optional[str]]: |
| """Extracts a programming language from the beginning of a string. |
| |
| This function checks if the input string starts with a pattern of the form |
| ``<_some_language_>``. If it does, it extracts the language string and returns |
| a tuple of (remainder, language). Otherwise, it returns the original string |
| and `None`. |
| |
| Args: |
| input_string (str): The input string, which may start with ``<_language_>``. |
| |
| Returns: |
| Tuple[str, Optional[str]]: |
| A tuple where: |
| - The first element is either: |
| - The remainder of the string (everything after ``<_language_>``), |
| if a match is found; or |
| - The original string, if no match is found. |
| - The second element is the extracted language if a match is found; |
| otherwise, `None`. |
| """ |
| pattern = r"^<_([^>]+)_>\s*(.*)" |
| match = re.match(pattern, input_string, flags=re.DOTALL) |
| if match: |
| language = str(match.group(1)) |
| remainder = str(match.group(2)) |
| return remainder, language |
| else: |
| return input_string, None |
|
|
| def _get_code_language_enum(self, value: Optional[str]) -> CodeLanguageLabel: |
| """ |
| Converts a string to a corresponding `CodeLanguageLabel` enum member. |
| |
| If the provided string does not match any value in `CodeLanguageLabel`, |
| it defaults to `CodeLanguageLabel.UNKNOWN`. |
| |
| Args: |
| value (Optional[str]): The string representation of the code language or None. |
| |
| Returns: |
| CodeLanguageLabel: The corresponding enum member if the value is valid, |
| otherwise `CodeLanguageLabel.UNKNOWN`. |
| """ |
| if not isinstance(value, str): |
| return CodeLanguageLabel.UNKNOWN |
|
|
| try: |
| return CodeLanguageLabel(value) |
| except ValueError: |
| return CodeLanguageLabel.UNKNOWN |
|
|
| def __call__( |
| self, |
| doc: DoclingDocument, |
| element_batch: Iterable[ItemAndImageEnrichmentElement], |
| ) -> Iterable[NodeItem]: |
| """ |
| Processes the given batch of elements and enriches them with predictions. |
| |
| Parameters |
| ---------- |
| doc : DoclingDocument |
| The document being processed. |
| element_batch : Iterable[ItemAndImageEnrichmentElement] |
| A batch of elements to be processed. |
| |
| Returns |
| ------- |
| Iterable[Any] |
| An iterable of enriched elements. |
| """ |
| if not self.enabled: |
| for element in element_batch: |
| yield element.item |
| return |
|
|
| labels: List[str] = [] |
| images: List[Union[Image.Image, np.ndarray]] = [] |
| elements: List[TextItem] = [] |
| for el in element_batch: |
| assert isinstance(el.item, TextItem) |
| elements.append(el.item) |
| labels.append(el.item.label) |
| images.append(el.image) |
|
|
| outputs = self.code_formula_model.predict(images, labels) |
|
|
| for item, output in zip(elements, outputs): |
| if isinstance(item, CodeItem): |
| output, code_language = self._extract_code_language(output) |
| item.code_language = self._get_code_language_enum(code_language) |
| item.text = output |
|
|
| yield item |
|
|
