mirror of
https://github.com/DS4SD/docling.git
synced 2025-08-02 15:32:30 +00:00
ddb8ad9227
- Introduce the AcceleratorOptions, AcceleratorDevice and use them to set the device where the models run. - Introduce the accelerator_utils with function to decide the device and resolve the AUTO setting. - Refactor the way how the docling-ibm-models are called to match the new init signature of models. - Translate the accelerator options to the specific inputs for third-party models. - Extend the docling CLI with parameters to set the num_threads and device. - Add new unit tests. - Write new example how to use the accelerator options.
41 lines
1.3 KiB
Python
41 lines
1.3 KiB
Python
import logging
|
|
|
|
import torch
|
|
|
|
from docling.datamodel.pipeline_options import AcceleratorDevice
|
|
|
|
_log = logging.getLogger(__name__)
|
|
|
|
|
|
def decide_device(accelerator_device: AcceleratorDevice) -> str:
|
|
r"""
|
|
Resolve the device based on the acceleration options and the available devices in the system
|
|
Rules:
|
|
1. AUTO: Check for the best available device on the system.
|
|
2. User-defined: Check if the device actually exists, otherwise fall-back to CPU
|
|
"""
|
|
cuda_index = 0
|
|
device = "cpu"
|
|
|
|
has_cuda = torch.backends.cuda.is_built() and torch.cuda.is_available()
|
|
has_mps = torch.backends.mps.is_built() and torch.backends.mps.is_available()
|
|
|
|
if accelerator_device == AcceleratorDevice.AUTO:
|
|
# TODO: Enable MPS later
|
|
if has_cuda:
|
|
device = f"cuda:{cuda_index}"
|
|
else:
|
|
if accelerator_device == AcceleratorDevice.CUDA:
|
|
if has_cuda:
|
|
device = f"cuda:{cuda_index}"
|
|
else:
|
|
_log.warning("CUDA is not available in the system. Fall back to 'CPU'")
|
|
elif accelerator_device == AcceleratorDevice.MPS:
|
|
if has_mps:
|
|
device = "mps"
|
|
else:
|
|
_log.warning("MPS is not available in the system. Fall back to 'CPU'")
|
|
|
|
_log.info("Accelerator device: '%s'", device)
|
|
return device
|