Create copies instead

docprompt · Oct 7, 2024 · 0a20b3a · 0a20b3a
1 parent 9dedc24
commit 0a20b3a
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/docprompt/tasks/ocr/result.py b/docprompt/tasks/ocr/result.py
@@ -70,7 +70,9 @@ def contribute_to_document_node(
 
         page_node = document_node.page_nodes[page_number - 1]
         if hasattr(page_node.metadata, "ocr_results"):
-            page_node.metadata.ocr_results = self
+            page_node.metadata.ocr_results = self.model_copy(
+                update={"raster_image": None}
+            )
         else:
             super().contribute_to_document_node(document_node, page_number=page_number)
 
@@ -80,5 +82,3 @@ def contribute_to_document_node(
                 page_number=page_number,
                 image_bytes=self.raster_image,
             )
-
-        self.raster_image = None  # We need to clear this for memory reasons
diff --git a/docprompt/tasks/result.py b/docprompt/tasks/result.py
@@ -54,7 +54,7 @@ def contribute_to_document_node(
         ), "Page number must be less than or equal to the number of pages in the document"
 
         page_node = document_node.page_nodes[page_number - 1]
-        page_node.metadata.task_results[self.task_key] = self
+        page_node.metadata.task_results[self.task_key] = self.model_copy()
 
 
 TTaskInput = TypeVar("TTaskInput")  # What invoke requires