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

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

12
docling/pipeline/base_pipeline.py
View File

@ -93,8 +93,9 @@ 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
def get_default_options(cls) -> PipelineOptions: def get_default_options(cls) -> PipelineOptions:
@ -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: