Skip to content

Commit 74d1e89

Browse files
committed
doc2rag: if not using cuda, explicitly set the accelerator to cpu
docling doesn't support intel or rocm accelerators, and is incorrectly selecting cuda in some cases. Only enable cuda if ramalama has already detected it. Fixes: #2089 Signed-off-by: Mike Bonnet <mikeb@redhat.com>
1 parent 33f5e9f commit 74d1e89

File tree

1 file changed

+9
-2
lines changed

1 file changed

+9
-2
lines changed

container-images/scripts/doc2rag

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ from pathlib import Path
1818
import docling
1919
import qdrant_client
2020
from docling.chunking import HybridChunker
21+
from docling.datamodel.accelerator_options import AcceleratorDevice, AcceleratorOptions
2122
from docling.datamodel.base_models import InputFormat
2223
from docling.datamodel.pipeline_options import PdfPipelineOptions
2324
from docling.document_converter import DocumentConverter, PdfFormatOption
@@ -40,8 +41,14 @@ class Converter:
4041

4142
def __init__(self, args):
4243
# Docling Setup (Turn off OCR (image processing) for drastically reduced RAM usage and big speed increase)
43-
pipeline_options = PdfPipelineOptions()
44-
pipeline_options.do_ocr = args.ocr
44+
if os.environ.get("CUDA_VISIBLE_DEVICES", "").lower() in ["", "none", "-1"]:
45+
dev = AcceleratorDevice.CPU
46+
else:
47+
dev = AcceleratorDevice.CUDA
48+
pipeline_options = PdfPipelineOptions(
49+
accelerator_options=AcceleratorOptions(device=dev),
50+
do_ocr=args.ocr,
51+
)
4552
self.sources = []
4653
for source in args.sources:
4754
self.add(source)

0 commit comments

Comments
 (0)