Instrumented more code with TimeTrace decorators. (#474)

aryn-ai · Jun 18, 2024 · b0e3eca · b0e3eca
1 parent 36162a6
commit b0e3eca
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 0 deletions.
diff --git a/lib/sycamore/sycamore/functions/document.py b/lib/sycamore/sycamore/functions/document.py
@@ -6,9 +6,11 @@
 
 from sycamore.data import Document, Element
 from sycamore.utils.image_utils import try_draw_boxes
+from sycamore.utils.time_trace import timetrace
 from PIL import Image as PImage, ImageDraw
 
 
+@timetrace("Pdf2Imgs")
 def split_and_convert_to_image(doc: Document) -> list[Document]:
     """Split a document into individual pages as images and convert them into Document objects.
 

diff --git a/lib/sycamore/sycamore/transforms/extract_schema.py b/lib/sycamore/sycamore/transforms/extract_schema.py
@@ -13,6 +13,7 @@
 from sycamore.plan_nodes import Node
 from sycamore.transforms.map import Map
 from sycamore.utils.extract_json import extract_json
+from sycamore.utils.time_trace import timetrace
 
 
 def element_list_formatter(elements: list[Element]) -> str:
@@ -80,6 +81,7 @@ def __init__(
         self._prompt_formatter = prompt_formatter
         self._max_num_properties = max_num_properties
 
+    @timetrace("ExtrSchema")
     def extract_schema(self, document: Document) -> Document:
         entities = self._handle_zero_shot_prompting(document)
 
@@ -142,6 +144,7 @@ def __init__(
         self._num_of_elements = num_of_elements
         self._prompt_formatter = prompt_formatter
 
+    @timetrace("ExtrProps")
     def extract_properties(self, document: Document) -> Document:
         entities = self._handle_zero_shot_prompting(document)
 

diff --git a/lib/sycamore/sycamore/transforms/summarize.py b/lib/sycamore/sycamore/transforms/summarize.py
@@ -7,6 +7,7 @@
 from sycamore.llms import LLM
 from sycamore.llms.prompts import TextSummarizerGuidancePrompt
 from sycamore.transforms.map import Map
+from sycamore.utils.time_trace import timetrace
 
 
 class Summarizer(ABC):
@@ -55,6 +56,7 @@ def summarize(self, document: Document) -> Document:
         document.elements = elements
         return document
 
+    @timetrace("SummText")
     def _summarize_text_element(self, element: Element) -> Element:
         prompt = TextSummarizerGuidancePrompt()
 

diff --git a/lib/sycamore/sycamore/transforms/summarize_images.py b/lib/sycamore/sycamore/transforms/summarize_images.py
@@ -8,6 +8,7 @@
 from sycamore.transforms.map import Map
 from sycamore.utils.image_utils import base64_data_url
 from sycamore.utils.extract_json import extract_json
+from sycamore.utils.time_trace import timetrace
 
 
 class OpenAIImageSummarizer:
@@ -81,6 +82,7 @@ def __init__(
         self.prompt = prompt
         self.include_context = include_context
 
+    @timetrace("SummImg")
     def summarize_image(
         self, image: Image.Image, preceding_context: Optional[str] = None, following_context: Optional[str] = None
     ):