rewrite llm_map to deal with iterative prompting better

Signed-off-by: Henry Lindeman <[email protected]>
aryn-ai · HenryL27 · Jan 30, 2025 · Jan 17, 2025 · Jan 17, 2025 · Jan 17, 2025
commit 8b6f085eb14e31e7a4bef675d46cf3b15148dab1
diff --git a/lib/sycamore/sycamore/llms/prompts/prompts.py b/lib/sycamore/sycamore/llms/prompts/prompts.py
@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Any, Union, Optional, Callable, Sequence
+from typing import Any, Union, Optional, Callable
 import copy
 
 import pydantic
@@ -42,7 +42,7 @@ class SycamorePrompt:
     convert sycamore objects (``Document``s, ``Element``s) into ``RenderedPrompts``
     """
 
-    def render_document(self, doc: Document) -> Union[RenderedPrompt, Sequence[RenderedPrompt]]:
+    def render_document(self, doc: Document) -> RenderedPrompt:
         """Render this prompt, given this document as context.
         Used in llm_map
 
@@ -54,7 +54,7 @@ def render_document(self, doc: Document) -> Union[RenderedPrompt, Sequence[Rende
         """
         raise NotImplementedError(f"render_document is not implemented for {self.__class__.__name__}")
 
-    def render_element(self, elt: Element, doc: Document) -> Union[RenderedPrompt, Sequence[RenderedPrompt]]:
+    def render_element(self, elt: Element, doc: Document) -> RenderedPrompt:
         """Render this prompt, given this element and its parent document as context.
         Used in llm_map_elements
 
@@ -66,7 +66,7 @@ def render_element(self, elt: Element, doc: Document) -> Union[RenderedPrompt, S
         """
         raise NotImplementedError(f"render_element is not implemented for {self.__class__.__name__}")
 
-    def render_multiple_documents(self, docs: list[Document]) -> Union[RenderedPrompt, Sequence[RenderedPrompt]]:
+    def render_multiple_documents(self, docs: list[Document]) -> RenderedPrompt:
         """Render this prompt, given a list of documents as context.
         Used in llm_reduce
 
@@ -113,19 +113,6 @@ def set(self, **kwargs) -> "SycamorePrompt":
                 new.__dict__[k] = v
         return new
 
-    def is_done(self, s: str) -> bool:
-        """Decide whether a given response is sufficient. Used when rendering
-        the prompt generates a sequence of prompts rather than a single prompt.
-        The default implementation always returns True
-
-        Args:
-            s: a string response from the LLM
-
-        Returns:
-            Whether to continue making LLM calls
-        """
-        return True
-
 
 def _build_format_str(
     system: Optional[str], user: Union[None, str, list[str]], format_args: dict[str, Any]
@@ -201,7 +188,7 @@ def _render_element_list_to_string(self, doc: Document):
         elts = self.element_select(doc.elements)
         return self.element_list_constructor(elts)
 
-    def render_document(self, doc: Document) -> Union[RenderedPrompt, Sequence[RenderedPrompt]]:
+    def render_document(self, doc: Document) -> RenderedPrompt:
         """Render this prompt, given this document as context, using python's
         ``str.format()`` method. The keys passed into ``format()`` are as follows:
 
@@ -280,11 +267,18 @@ class ElementListIterPrompt(ElementListPrompt):
             # ]
     """
 
-    def __init__(self, *, element_batcher: Optional[Callable[[list[Element]], list[list[Element]]]] = None, **kwargs):
+    def __init__(
+        self,
+        *,
+        element_batcher: Optional[Callable[[list[Element]], list[list[Element]]]] = None,
+        iteration_var_name: str = "i",
+        **kwargs,
+    ):
         self.element_batcher = element_batcher or (lambda e: [e])
+        self.iteration_var_name = iteration_var_name
         super().__init__(**kwargs)
 
-    def render_document(self, doc: Document) -> Sequence[RenderedPrompt]:
+    def render_document(self, doc: Document) -> RenderedPrompt:
         """Render this prompt, given this document as context, using python's
         ``str.format()`` method. The keys passed into ``format()`` are as follows:
 
@@ -304,19 +298,22 @@ def render_document(self, doc: Document) -> Sequence[RenderedPrompt]:
             ``self.user.format()`` using the format keys as specified above. Each instance
             is rendered from a batch of elements generated by ``self.element_batcher``
         """
+        i = doc.properties.get(self.iteration_var_name, 0)
 
         format_args = self.kwargs
         format_args["doc_text"] = doc.text_representation
         flat_props = flatten_data(doc.properties, prefix="doc_property", separator="_")
         format_args.update(flat_props)
 
-        prompts = []
-        for elt_batch in self.element_batcher(doc.elements):
-            elements = self.element_select(elt_batch)
-            elementstr = self.element_list_constructor(elements)
-            messages = _build_format_str(self.system, self.user, {"elements": elementstr, **format_args})
-            prompts.append(RenderedPrompt(messages=messages))
-        return prompts
+        for j, elt_batch in enumerate(self.element_batcher(doc.elements)):
+            if j < i:
+                continue
+            else:
+                elements = self.element_select(elt_batch)
+                elementstr = self.element_list_constructor(elements)
+                messages = _build_format_str(self.system, self.user, {"elements": elementstr, **format_args})
+                return RenderedPrompt(messages=messages)
+        return RenderedPrompt(messages=[])
 
 
 class ElementPrompt(SycamorePrompt):

diff --git a/lib/sycamore/sycamore/tests/unit/test_docset.py b/lib/sycamore/sycamore/tests/unit/test_docset.py
@@ -44,7 +44,6 @@ def __init__(self):
         super().__init__(model_name="mock_model")
 
     def generate(self, *, prompt: RenderedPrompt, llm_kwargs: Optional[dict] = None) -> str:
-        print(prompt)
         if llm_kwargs is None:
             llm_kwargs = {}
         if prompt.messages[-1].content.endswith("Element_index: 1\nText: third element\n"):
@@ -98,6 +97,8 @@ def generate(self, *, prompt: RenderedPrompt, llm_kwargs: Optional[dict] = None)
                 return "group2"
             elif value == "3" or value == "three":
                 return "group3"
+            else:
+                return ""
         else:
             return prompt.messages[-1].content
 

diff --git a/lib/sycamore/sycamore/tests/unit/transforms/test_base_llm.py b/lib/sycamore/sycamore/tests/unit/transforms/test_base_llm.py
@@ -54,19 +54,19 @@ def test_happy_path(self):
         assert outdocs[1].text_representation == "booga"
         assert outdocs[1].properties["out"] == "booga"
 
-    def test_postprocess(self):
+    def test_validate(self):
         prompt = FakeDocPrompt()
         llm = FakeLLM()
         doc1 = Document({"text_representation": "ooga"})
         doc2 = Document({"text_representation": "booga"})
         count = 0
 
-        def ppfn(d: Document, i: int) -> Document:
+        def valfn(d: Document) -> bool:
             nonlocal count
             count += 1
-            return d
+            return count > 1
 
-        map = LLMMap(None, prompt, "out", llm, postprocess_fn=ppfn)
+        map = LLMMap(None, prompt, "out", llm, validate=valfn)
         _ = map.llm_map([doc1, doc2])
 
         assert count == 2
@@ -112,12 +112,12 @@ def test_postprocess(self):
         doc2 = Document({"doc_id": "2", "elements": [{"text_representation": "booga"}, {}]})
         count = 0
 
-        def ppfn(e: Element, i: int) -> Element:
+        def valfn(e: Element) -> bool:
             nonlocal count
             count += 1
-            return e
+            return count > 1
 
-        map = LLMMapElements(None, prompt, "out", llm, postprocess_fn=ppfn)
+        map = LLMMapElements(None, prompt, "out", llm, validate=valfn)
         _ = map.llm_map_elements([doc1, doc2])
 
         assert count == 4
diff --git a/lib/sycamore/sycamore/tests/unit/transforms/test_extract_entity.py b/lib/sycamore/sycamore/tests/unit/transforms/test_extract_entity.py
@@ -202,6 +202,7 @@ def test_extract_entity_with_tokenizer(self, mocker):
             entity_extractor=entity_extractor,
         )
         taken = entity_docset.take()
+
         assert taken[0].properties[f"{new_field}_source_element_index"] == {0, 1, 2}
         assert taken[1].properties[f"{new_field}_source_element_index"] == {2}
         assert taken[0].properties[new_field] == "4"

diff --git a/lib/sycamore/sycamore/transforms/base_llm.py b/lib/sycamore/sycamore/transforms/base_llm.py
@@ -8,22 +8,15 @@
 
 
 def _infer_prompts(
-    prompts: list[Sequence[RenderedPrompt]],
+    prompts: list[RenderedPrompt],
     llm: LLM,
     llm_mode: LLMMode,
-    is_done: Callable[[str], bool] = lambda s: True,
 ) -> list[tuple[str, int]]:
     if llm_mode == LLMMode.SYNC:
         res = []
-        for piter in prompts:
-            s = ""
-            i = -1
-            for p in piter:
-                i += 1
-                s = llm.generate(prompt=p)
-                if is_done(s):
-                    break
-            res.append((s, i))
+        for p in prompts:
+            s = llm.generate(prompt=p)
+            res.append(s)
         return res
     elif llm_mode == LLMMode.ASYNC:
         raise NotImplementedError("Haven't done async yet")
@@ -73,27 +66,47 @@ def __init__(
         output_field: str,
         llm: LLM,
         llm_mode: LLMMode = LLMMode.SYNC,
-        postprocess_fn: Callable[[Document, int], Document] = lambda d, i: d,
+        iteration_var: Optional[str] = None,
+        validate: Callable[[Document], bool] = lambda d: True,
+        max_tries: int = 5,
         **kwargs,
     ):
         self._prompt = prompt
         self._validate_prompt()
         self._output_field = output_field
         self._llm = llm
         self._llm_mode = llm_mode
-        self._postprocess_fn = postprocess_fn
+        self._iteration_var = iteration_var
+        self._validate = validate
+        self._max_tries = max_tries
         super().__init__(child, f=self.llm_map, **kwargs)
 
     def llm_map(self, documents: list[Document]) -> list[Document]:
-        rendered_inc = [self._prompt.render_document(d) for d in documents]
-        rendered = _as_sequences(rendered_inc)
-        results = _infer_prompts(rendered, self._llm, self._llm_mode, self._prompt.is_done)
-        postprocessed = []
-        for d, (r, i) in zip(documents, results):
-            d.properties[self._output_field] = r
-            new_d = self._postprocess_fn(d, i)
-            postprocessed.append(new_d)
-        return postprocessed
+        if self._iteration_var is not None:
+            for d in documents:
+                d.properties[self._iteration_var] = 0
+
+        valid = [False] * len(documents)
+        tries = 0
+        while not all(valid) and tries < self._max_tries:
+            tries += 1
+            rendered = [self._prompt.render_document(d) for v, d in zip(valid, documents) if not v]
+            if sum([0, *(len(r.messages) for r in rendered)]) == 0:
+                break
+            results = _infer_prompts(rendered, self._llm, self._llm_mode)
+            ri = 0
+            for i in range(len(documents)):
+                if valid[i]:
+                    continue
+                documents[i].properties[self._output_field] = results[ri]
+                valid[i] = self._validate(documents[i])
+                ri += 1
+                if self._iteration_var is not None and not valid[i]:
+                    documents[i].properties[self._iteration_var] += 1
+            if self._iteration_var is None:
+                break
+
+        return documents
 
     def _validate_prompt(self):
         doc = Document()
@@ -143,30 +156,56 @@ def __init__(
         output_field: str,
         llm: LLM,
         llm_mode: LLMMode = LLMMode.SYNC,
-        postprocess_fn: Callable[[Element, int], Element] = lambda e, i: e,
+        iteration_var: Optional[str] = None,
+        validate: Callable[[Element], bool] = lambda d: True,
+        max_tries: int = 5,
         **kwargs,
     ):
         self._prompt = prompt
         self._validate_prompt()
         self._output_field = output_field
         self._llm = llm
         self._llm_mode = llm_mode
-        self._postprocess_fn = postprocess_fn
+        self._iteration_var = iteration_var
+        self._validate = validate
+        self._max_tries = max_tries
         super().__init__(child, f=self.llm_map_elements, **kwargs)
 
     def llm_map_elements(self, documents: list[Document]) -> list[Document]:
-        rendered = [(d, e, self._prompt.render_element(e, d)) for d in documents for e in d.elements]
-        results = _infer_prompts(
-            _as_sequences([p for _, _, p in rendered]), self._llm, self._llm_mode, self._prompt.is_done
-        )
-        new_elts = []
+        elt_doc_pairs = [(e, d) for d in documents for e in d.elements]
+        if self._iteration_var is not None:
+            for e, _ in elt_doc_pairs:
+                e.properties[self._iteration_var] = 0
+
+        valid = [False] * len(elt_doc_pairs)
+        tries = 0
+        while not all(valid) and tries < self._max_tries:
+            tries += 1
+            rendered = [self._prompt.render_element(e, d) for v, (e, d) in zip(valid, elt_doc_pairs) if not v]
+            if sum([0, *(len(r.messages) for r in rendered)]) == 0:
+                break
+            results = _infer_prompts(rendered, self._llm, self._llm_mode)
+            ri = 0
+            for i in range(len(elt_doc_pairs)):
+                if valid[i]:
+                    continue
+                print(ri)
+                elt, doc = elt_doc_pairs[i]
+                elt.properties[self._output_field] = results[ri]
+                valid[i] = self._validate(elt)
+                ri += 1
+                if self._iteration_var is not None:
+                    elt.properties[self._iteration_var] += 1
+            if self._iteration_var is None:
+                break
+
         last_doc = None
-        for (r, i), (d, e, _) in zip(results, rendered):
+        new_elts = []
+        for e, d in elt_doc_pairs:
             if last_doc is not None and last_doc.doc_id != d.doc_id:
                 last_doc.elements = new_elts
                 new_elts = []
-            e.properties[self._output_field] = r
-            new_elts.append(self._postprocess_fn(e, i))
+            new_elts.append(e)
             last_doc = d
         if last_doc is not None:
             last_doc.elements = new_elts