fix artifacts_path type

Signed-off-by: Michele Dolfi <dol@zurich.ibm.com>
This commit is contained in:
Michele Dolfi 2025-01-24 16:00:22 +01:00
parent 2c35b57e34
commit 6c3d31d68b

View File

@ -1,6 +1,6 @@
import re
from pathlib import Path
from typing import Iterable, List, Literal, Optional, Tuple
from typing import Iterable, List, Literal, Optional, Tuple, Union
from docling_core.types.doc import CodeItem, DoclingDocument, NodeItem, TextItem
from docling_core.types.doc.base import BoundingBox
@ -62,7 +62,7 @@ class CodeFormulaModel(BaseItemAndImageEnrichmentModel):
def __init__(
self,
enabled: bool,
artifacts_path: Optional[Path],
artifacts_path: Optional[Union[Path, str]],
options: CodeFormulaModelOptions,
accelerator_options: AcceleratorOptions,
):
@ -92,6 +92,8 @@ class CodeFormulaModel(BaseItemAndImageEnrichmentModel):
if artifacts_path is None:
artifacts_path = self.download_models_hf()
else:
artifacts_path = Path(artifacts_path)
self.code_formula_model = CodeFormulaPredictor(
artifacts_path=artifacts_path,