diff --git a/lib/sycamore/sycamore/functions/document.py b/lib/sycamore/sycamore/functions/document.py index 3c4016c76..265e9dcd1 100644 --- 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 index 2002376f6..0b7192e33 100644 --- 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 index 534872d68..abc4a1fb5 100644 --- 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 index 63a2bcb86..d34800fd5 100644 --- 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 ):