Introduce dspy.Template

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 @@ def __call__(
         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:

dspy/primitives/__init__.py

@@ -3,3 +3,4 @@
 from .prediction import *
 from .assertions import *
 from .python_interpreter import *
+from .template import *

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])

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