mirrors/docling
1
0
Fork 0
mirror of https://github.com/DS4SD/docling.git synced 2025-07-26 20:14:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

Use device_map for transformer models

Browse Source 
Signed-off-by: Christoph Auer <cau@zurich.ibm.com>
This commit is contained in:
Christoph Auer 2025-07-09 16:49:21 +02:00
parent 931eb55b88
commit 05123c9342
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
docling/models/picture_description_vlm_model.py
View File

@ -65,6 +65,7 @@ class PictureDescriptionVlmModel(
self.processor = AutoProcessor.from_pretrained(artifacts_path) self.processor = AutoProcessor.from_pretrained(artifacts_path)
self.model = AutoModelForVision2Seq.from_pretrained( self.model = AutoModelForVision2Seq.from_pretrained(
artifacts_path, artifacts_path,
device_map=self.device,
torch_dtype=torch.bfloat16, torch_dtype=torch.bfloat16,
_attn_implementation=( _attn_implementation=(
"flash_attention_2" "flash_attention_2"
@ -72,7 +73,7 @@ class PictureDescriptionVlmModel(
and accelerator_options.cuda_use_flash_attention2 and accelerator_options.cuda_use_flash_attention2
else "eager" else "eager"
), ),
).to(self.device) )
self.provenance = f"{self.options.repo_id}" self.provenance = f"{self.options.repo_id}"