mirrors/docling
1
0
Fork 0
mirror of https://github.com/DS4SD/docling.git synced 2025-07-27 04:24:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

Update base_pipeline.py

Browse Source 
Signed-off-by: jane-temcious <jane.temcious@iterate.ai>
This commit is contained in:
jane-temcious 2025-06-26 06:49:07 +05:30 committed by GitHub
parent 5ab8792697
commit 1fdfb2bc50
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 9 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
docling/pipeline/base_pipeline.py
View File

@ -93,7 +93,8 @@ class BasePipeline(ABC):
pass pass
def _unload(self, conv_res: ConversionResult): def _unload(self, conv_res: ConversionResult):
pass # Clear GPU memory after processingAdd commentMore actions
clear_gpu_memory()
@classmethod @classmethod
@abstractmethod @abstractmethod
@ -196,13 +197,18 @@ class PaginatedPipeline(BasePipeline): # TODO this is a bad name.
return conv_res return conv_res
def _unload(self, conv_res: ConversionResult) -> ConversionResult: def _unload(self, conv_res: ConversionResult) -> ConversionResult:
# Unload page backends
for page in conv_res.pages: for page in conv_res.pages:
if page._backend is not None: if page._backend is not None:
page._backend.unload() page._backend.unload()
# Unload input backend
if conv_res.input._backend: if conv_res.input._backend:
conv_res.input._backend.unload() conv_res.input._backend.unload()
# Call parent _unload to clean up GPU memory
super()._unload(conv_res)
return conv_res return conv_res
def _determine_status(self, conv_res: ConversionResult) -> ConversionStatus: def _determine_status(self, conv_res: ConversionResult) -> ConversionStatus: