Repository: stanfordnlp/dspy
Pull request #552 (2024-03-04 21:42:19)

Introduce dspy.Template
The original implementation of `TemplateBackend` relied upon the `signature_to_template` function. This would take in a `dspy.Signature` and return a `dsp.Template`. However, we would pass `dspy.Example` into the `dsp.Template` and the extract function would return `dsp.Example` a different type. Leaving us with two different types of `Example` within a single prediction. To combat this, a new `Template` primitive was introduced into `dspy` to mirror and replace the original `dsp.Template`.

I do not expect this current PR to achieve feature parity with the original class, but initial tests have been introduced to mock and mirror the current prompt structure used in `dsp`.

Let me know what you think @CyrusOfEden, the last piece I believe to start using the `TemplateBackend` is to fully explore the `do_generate` function dealing with incomplete or inaccurate completions.

CyrusOfEden: YES man this is 🔥

Patch:
diff --git a/dspy/backends/template.py b/dspy/backends/template.py
index c940495d..98fd6969 100644
--- a/dspy/backends/template.py
+++ b/dspy/backends/template.py
@@ -1,5 +1,6 @@
 from dspy.signatures.signature import Signature, signature_to_template
 from dspy.primitives.example import Example
+from dspy.primitives.template import Template
 
 import typing as t
 from .base import BaseBackend, GeneratedOutput
@@ -30,7 +31,7 @@ class TemplateBackend(BaseBackend):
         example = Example(demos=demos, **kwargs)
 
         # Generate Template
-        template = signature_to_template(signature)
+        template = Template(signature)
 
         # Clean Up Kwargs Before Sending Through Language Model
         for input in signature.input_fields:
diff --git a/dspy/primitives/__init__.py b/dspy/primitives/__init__.py
index b7baf229..374e40ec 100644
--- a/dspy/primitives/__init__.py
+++ b/dspy/primitives/__init__.py
@@ -3,3 +3,4 @@ from .program import *
 from .prediction import *
 from .assertions import *
 from .python_interpreter import *
+from .template import *
diff --git a/dspy/primitives/template.py b/dspy/primitives/template.py
new file mode 100644
index 00000000..1283c469
--- /dev/null
+++ b/dspy/primitives/template.py
@@ -0,0 +1,169 @@
+import regex
+import typing as t
+from dspy.signatures.signature import Signature
+from dspy.primitives.example import Example
+
+
+def passages_to_text(passages: t.Iterable[str]) -> str:
+    assert len(passages) > 0
+    if len(passages) > 1:
+        return "\n".join(
+            [f"[{idx + 1}] <<{text}>>" for idx, text in enumerate(passages)]
+        )
+    else:
+        return passages[0]
+
+
+def format_answers(answers: t.Iterable[str]) -> str:
+    assert len(answers) > 0
+    return (answers[0]).strip()
+
+
+def default_format_handler(x: str) -> str:
+    assert type(x) == str
+    return " ".join(x.split())
+
+
+class Template:
+    def __init__(self, signature: Signature, **kwargs):
+        self.signature = signature
+        self.kwargs = kwargs
+
+        self.format_handlers: dict[str, t.Callable] = {
+            "context": passages_to_text,
+            "passages": passages_to_text,
+            "answers": format_answers,
+        }
+
+        for key, field in signature.fields.items():
+            format = field.json_schema_extra.get("format")
+            if format:
+                self.format_handlers[key] = format
+
+    def _get_format_handler(self, name: str) -> t.Callable[[str], str]:
+        if name in self.format_handlers:
+            return self.format_handlers[name]
+
+        return default_format_handler
+
+    def _example_has_input_fields(self, example: Example):
+        for name in self.signature.input_fields:
+            if name not in example:
+                raise Exception(f"Example missing necessary input field: {name}")
+
+    def _example_has_output_fields(self, example: Example):
+        for name in self.signature.output_fields:
+            if name not in example:
+                raise Exception(f"Example missing necessary output field: {name}")
+
+    def query(self, example: Example, is_demo: bool) -> str:
+        if is_demo:
+            self._example_has_input_fields(example)
+            self._example_has_output_fields(example)
+
+        result = []
+
+        # Append all Input Values, Regardless of Demo or not
+        for name, field in self.signature.input_fields.items():
+            format_handler = self._get_format_handler(name)
+
+            result.append(
+                f"{field.json_schema_extra['prefix']} {format_handler(example[name])}"
+            )
+
+        for name, field in self.signature.output_fields.items():
+            format_handler = self._get_format_handler(name)
+
+            if name not in example:
+                result.append(f"{field.json_schema_extra['prefix']} ")
+                break
+            else:
+                result.append(
+                    f"{field.json_schema_extra['prefix']} {format_handler(example[name])}"
+                )
+
+        return "\n\n".join(result)
+
+    def guidelines(self) -> str:
+        """Returns the task guidelines as described in the lm prompt"""
+        result = "Follow the following format.\n\n"
+
+        field_strings = []
+        for field in self.signature.fields.values():
+            field_strings.append(
+                f"{field.json_schema_extra['prefix']} {field.json_schema_extra['desc']}"
+            )
+
+        return result + "\n\n".join(field_strings)
+
+    def extract(self, example: Example, raw_pred: str) -> Example:
+        """Extracts the answer from the LM raw prediction using the template structure
+
+        Args:
+            example (Example): Contains the input variables that raw_pred was completed on.
+            raw_pred (str): LM generated field_strings
+
+        Returns:
+            Example: The example with the output variables filled in
+
+        """
+
+        # We have to deepcopy, so that the values dont continously overwrite each other.
+        example = example.copy()
+
+        full_text = self.__call__(example) + raw_pred
+
+        if not full_text.endswith("\n\n---"):
+            full_text = full_text + "\n\n---"
+
+        # Generate Search Strings
+        search_strings = []
+        output_fields = list(self.signature.output_fields.keys())
+        for idx, (key, field) in enumerate(self.signature.output_fields.items()):
+            if len(search_strings) > 0:
+                search_strings[-1] += f"{field.json_schema_extra['prefix']}"
+
+            target_str = f"(?s){field.json_schema_extra['prefix']}\\s?(.+?)"
+            if idx != len(self.signature.output_fields) - 1:
+                target_str += "\\n\\n"
+            else:
+                target_str += "\\n\\n\\-\\-\\-"
+
+            search_strings.append(target_str)
+
+        # Generate Results
+        if len(self.signature.output_fields) == 1:
+            matches = regex.findall(search_strings[0], full_text)
+
+            # If no matches are found, and there are is only one prediction, return entire prediction
+            if matches is None:
+                example[output_fields[0]] = full_text
+            else:
+                example[output_fields[0]] = matches[-1]
+
+        else:
+            for idx, field in enumerate(output_fields):
+                matches = regex.findall(search_strings[idx], full_text)
+
+                if len(matches) > 0:
+                    example[field] = matches[-1]
+
+        return example
+
+    def __call__(self, example: Example, show_guidelines: bool = True) -> str:
+        prompt_spans = []
+
+        # Start by getting the instructions
+        prompt_spans.append(self.signature.instructions)
+
+        # Generate the Guidelines
+        prompt_spans.append(self.guidelines())
+
+        # Generate Spans for Each Demo
+        for demo in example.get("demos", []):
+            prompt_spans.append(self.query(demo, is_demo=True))
+
+        # Generate Empty Demo for Generation
+        prompt_spans.append(self.query(example, is_demo=False))
+
+        return "\n\n---\n\n".join([span.strip() for span in prompt_spans])


Test patch:
diff --git a/tests/primitives/test_template.py b/tests/primitives/test_template.py
new file mode 100644
index 00000000..8bded3ef
--- /dev/null
+++ b/tests/primitives/test_template.py
@@ -0,0 +1,95 @@
+import pytest
+from dspy import Example, Signature, InputField, OutputField
+from dspy.primitives import Template
+
+
+class Emotion(Signature):
+    """Classify emotion among sadness, joy, love, anger, fear, surprise."""
+
+    sentence = InputField()
+    sentiment = OutputField()
+
+
+class CheckCitationFaithfulness(Signature):
+    """Verify that the text is based on the provided context."""
+
+    context = InputField(desc="facts here are assumed to be true")
+    text = InputField()
+    faithfulness = OutputField(
+        desc="True/False indicating if text is faithful to context"
+    )
+
+
+class COTCheckCitationFaithfulness(Signature):
+    """Verify that the text is based on the provided context."""
+
+    context = InputField(desc="facts here are assumed to be true")
+    text = InputField()
+    rationale = OutputField(
+        desc="Think step by step in order to generate the faithfulness.",
+    )
+    faithfulness = OutputField(
+        desc="True/False indicating if text is faithful to context"
+    )
+
+
+TEMPLATE_SCENARIOS = [
+    {
+        "signature": Emotion,
+        "output": "Joy",
+        "input_kwargs": {
+            "sentence": "This is a positive test sentence.",
+        },
+        "output_kwargs": {"sentiment": "Joy"},
+        "prompt": "Classify emotion among sadness, joy, love, anger, fear, surprise.\n\n---\n\nFollow the following format.\n\nSentence: ${sentence}\n\nSentiment: ${sentiment}\n\n---\n\nSentence: This is a positive test sentence.\n\nSentiment:",
+    },
+    {
+        "signature": CheckCitationFaithfulness,
+        "output": "False",
+        "input_kwargs": {
+            "context": [
+                "The 21-year-old made seven appearances for the Hammers and netted his only goal for them in a Europa League qualification round match against Andorran side FC Lustrains last season. Lee had two loan spells in League One last term, with Blackpool and then Colchester United. He scored twice for the U's but was unable to save them from relegation. The length of Lee's contract with the promoted Tykes has not been revealed. Find all the latest football transfers on our dedicated page."
+            ],
+            "text": "Lee scored 3 goals for Colchester United.",
+        },
+        "output_kwargs": {
+            "faithfulness": "False",
+        },
+        "prompt": "Verify that the text is based on the provided context.\n\n---\n\nFollow the following format.\n\nContext: facts here are assumed to be true\n\nText: ${text}\n\nFaithfulness: True/False indicating if text is faithful to context\n\n---\n\nContext: The 21-year-old made seven appearances for the Hammers and netted his only goal for them in a Europa League qualification round match against Andorran side FC Lustrains last season. Lee had two loan spells in League One last term, with Blackpool and then Colchester United. He scored twice for the U's but was unable to save them from relegation. The length of Lee's contract with the promoted Tykes has not been revealed. Find all the latest football transfers on our dedicated page.\n\nText: Lee scored 3 goals for Colchester United.\n\nFaithfulness:",
+    },
+    {
+        "signature": COTCheckCitationFaithfulness,
+        "output": "Faithfulness: False",
+        "output": "produce the faithfulness. We know that Lee had two loan spells in League One last term, with Blackpool and then Colchester United. He scored twice for the U's but was unable to save them from relegation. However, there is no mention of him scoring three goals for Colchester United.\n\nFaithfulness: False",
+        "input_kwargs": {
+            "context": [
+                "The 21-year-old made seven appearances for the Hammers and netted his only goal for them in a Europa League qualification round match against Andorran side FC Lustrains last season. Lee had two loan spells in League One last term, with Blackpool and then Colchester United. He scored twice for the U's but was unable to save them from relegation. The length of Lee's contract with the promoted Tykes has not been revealed. Find all the latest football transfers on our dedicated page."
+            ],
+            "text": "Lee scored 3 goals for Colchester United.",
+        },
+        "output_kwargs": {
+            "rationale": "produce the faithfulness. We know that Lee had two loan spells in League One last term, with Blackpool and then Colchester United. He scored twice for the U's but was unable to save them from relegation. However, there is no mention of him scoring three goals for Colchester United.",
+            "faithfulness": "False",
+        },
+        "prompt": "Verify that the text is based on the provided context.\n\n---\n\nFollow the following format.\n\nContext: facts here are assumed to be true\n\nText: ${text}\n\nRationale: Think step by step in order to generate the faithfulness.\n\nFaithfulness: True/False indicating if text is faithful to context\n\n---\n\nContext: The 21-year-old made seven appearances for the Hammers and netted his only goal for them in a Europa League qualification round match against Andorran side FC Lustrains last season. Lee had two loan spells in League One last term, with Blackpool and then Colchester United. He scored twice for the U's but was unable to save them from relegation. The length of Lee's contract with the promoted Tykes has not been revealed. Find all the latest football transfers on our dedicated page.\n\nText: Lee scored 3 goals for Colchester United.\n\nRationale:",
+    },
+]
+
+
+def test_example_initialization():
+    for scenario in TEMPLATE_SCENARIOS:
+        template = Template(scenario["signature"])
+        example = Example(**scenario["input_kwargs"])
+        assert template(example) == scenario["prompt"], print(template(example))
+
+
+def test_template_extraction():
+    for scenario in TEMPLATE_SCENARIOS:
+        template = Template(scenario["signature"])
+        example = Example(**scenario["input_kwargs"])
+        extracted = template.extract(example, scenario["output"])
+        correct_example = Example(
+            **scenario["input_kwargs"], **scenario["output_kwargs"]
+        )
+
+        assert extracted == correct_example