mirrors/docling
1
0
Fork 0
mirror of https://github.com/DS4SD/docling.git synced 2025-08-01 15:02:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

use default device until we enable global management

Browse Source 
Signed-off-by: Michele Dolfi <dol@zurich.ibm.com>
This commit is contained in:
Michele Dolfi 2024-11-27 11:39:23 +01:00
parent 74e005df63
commit c1b6442670
2 changed files with 40 additions and 31 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
docling/datamodel/pipeline_options.py
View File

@ -46,13 +46,13 @@ class RapidOcrOptions(OcrOptions):
use_cls: Optional[bool] = None # same default as rapidocr use_cls: Optional[bool] = None # same default as rapidocr
use_rec: Optional[bool] = None # same default as rapidocr use_rec: Optional[bool] = None # same default as rapidocr
class Device(Enum): # class Device(Enum):
CPU = "CPU" # CPU = "CPU"
CUDA = "CUDA" # CUDA = "CUDA"
DIRECTML = "DIRECTML" # DIRECTML = "DIRECTML"
AUTO = "AUTO" # AUTO = "AUTO"
device: Device = Device.AUTO # Default value is AUTO # device: Device = Device.AUTO # Default value is AUTO
print_verbose: bool = False # same default as rapidocr print_verbose: bool = False # same default as rapidocr

59
docling/models/rapid_ocr_model.py
View File

@ -30,34 +30,43 @@ class RapidOcrModel(BaseOcrModel):
"Alternatively, Docling has support for other OCR engines. See the documentation." "Alternatively, Docling has support for other OCR engines. See the documentation."
) )
# Same as Defaults in RapidOCR # This configuration option will be revamped while introducing device settings for all models.
cls_use_cuda = False # For the moment we will default to auto and let onnx-runtime pick the best.
rec_use_cuda = False cls_use_cuda = True
det_use_cuda = False rec_use_cuda = True
det_use_dml = False det_use_cuda = True
cls_use_dml = False det_use_dml = True
rec_use_dml = False cls_use_dml = True
rec_use_dml = True
# If we set everything to true onnx-runtime would automatically choose the fastest accelerator # # Same as Defaults in RapidOCR
if self.options.device == self.options.Device.AUTO: # cls_use_cuda = False
cls_use_cuda = True # rec_use_cuda = False
rec_use_cuda = True # det_use_cuda = False
det_use_cuda = True # det_use_dml = False
det_use_dml = True # cls_use_dml = False
cls_use_dml = True # rec_use_dml = False
rec_use_dml = True
# If we set use_cuda to true onnx would use the cuda device available in runtime if no cuda device is available it would run on CPU. # # If we set everything to true onnx-runtime would automatically choose the fastest accelerator
elif self.options.device == self.options.Device.CUDA: # if self.options.device == self.options.Device.AUTO:
cls_use_cuda = True # cls_use_cuda = True
rec_use_cuda = True # rec_use_cuda = True
det_use_cuda = True # det_use_cuda = True
# det_use_dml = True
# cls_use_dml = True
# rec_use_dml = True
# If we set use_dml to true onnx would use the dml device available in runtime if no dml device is available it would work on CPU. # # If we set use_cuda to true onnx would use the cuda device available in runtime if no cuda device is available it would run on CPU.
elif self.options.device == self.options.Device.DIRECTML: # elif self.options.device == self.options.Device.CUDA:
det_use_dml = True # cls_use_cuda = True
cls_use_dml = True # rec_use_cuda = True
rec_use_dml = True # det_use_cuda = True
# # If we set use_dml to true onnx would use the dml device available in runtime if no dml device is available it would work on CPU.
# elif self.options.device == self.options.Device.DIRECTML:
# det_use_dml = True
# cls_use_dml = True
# rec_use_dml = True
self.reader = RapidOCR( self.reader = RapidOCR(
text_score=self.options.text_score, text_score=self.options.text_score,