A Python package for writing Language Models prompts in a nested way.
from procoder.functional import format_prompt, replaced_submodule
from procoder.prompt import *

requirements = NamedBlock(
    "Requirements",
    Collection(
        NamedVariable(
            refname="input_req",
            name="Input Requirement",
            content="The input should be two numbers.",
        ),
        NamedVariable(
            refname="output_req",
            name="Output Requirement",
            content=Single(
                "The output should be the sum of the two numbers."
            ).set_refname("output_req_content"),
        ),
    ),
)
instruction = NamedBlock(
    "Instruction",
    "Write a function in {language} that satisfies the {input_req} and {output_req}.",
)
prompt = (
    Collection(requirements, instruction)
    .set_sep("\n\n")
    .set_indexing_method(sharp2_indexing)
)
another_prompt = replaced_submodule(
    prompt,
    "output_req_content",
    Single("The output should be the multiplication of the two numbers."),
)

inputs = {"language": "python"}
print("First prompt:")
print(format_prompt(prompt, inputs))
print("")
print("Second prompt:")
print(format_prompt(another_prompt, inputs))

examples/example.py
dhh1995/PromptCoder
################################################################################
# Define basic utilities
################################################################################

from .functional import *

__version__ = "0.0.1"

procoder/__init__.py
dhh1995/PromptCoder
import copy

from procoder.utils.my_typing import *

from .prompt.base import TS, Module, Single, T, as_module
from .prompt.proxy import AddIndentProxy, SilenceProxy


def as_prompt(x: TS):
    return as_module(x)


# def flatten(x: T):

indent2 = " " * 2
indent4 = " " * 4
indent_tab = "\t"


def silence(x: T):
    return SilenceProxy(x)


def add_indent(x: T, indent: str):
    return AddIndentProxy(x, indent)


def add_indent2(x: T):
    return add_indent(x, indent=indent2)


def add_indent4(x: T):
    return add_indent(x, indent=indent4)


def add_indent_tab(x: T):
    return add_indent(x, indent=indent_tab)


def process_module_refname(v: T, target: Dict[str, Any], key: str = None):
    refname = v.refname or key
    name = v._get_name()

    if refname and name:
        if refname in target:
            raise ValueError(f"Duplicate key: {refname}")
        else:
            target[refname] = name


def collect_refnames(vars: Union[Dict[str, Any], T]) -> Dict[str, Any]:
    """Collect a mapping from refnames to names from all submodules of the vars"""
    target = {}
    modules = set()

    if isinstance(vars, Module):
        vars = {"prompt": vars}
    else:
        vars = {**vars}  # copy

    # find all submodules
    for k, v in vars.items():
        if isinstance(v, Module):
            modules |= v.get_all_submodules()

    for k, v in vars.items():
        if isinstance(v, Module):
            process_module_refname(v, target, key=k)
            # avoid duplicate
            modules.discard(v)

    for v in modules:
        process_module_refname(v, target)

    return target


def format_refnames(inputs: Dict[str, Any], format_func: Callable[[Any], str]):
    """Format the names with format_func"""
    return {k: format_func(v) for k, v in inputs.items()}


def add_brackets(refnames: Dict[str, Any]) -> Dict[str, Any]:
    return format_refnames(refnames, lambda x: f"[{x}]")


def check_duplicate_keys(x: Dict[str, Any], y: Dict[str, Any]):
    """Check if there are duplicate keys in two dicts"""
    for k in x.keys():
        if k in y:
            raise ValueError(f"Duplicate key: {k}")


def add_refnames(
    vars: Union[Dict[str, Any], T], target: Dict[str, Any], include_brackets=True
):
    """Add refnames to the target dict"""
    refnames = collect_refnames(vars)
    if include_brackets:
        refnames = add_brackets(refnames)
    check_duplicate_keys(target, refnames)
    target.update(refnames)
    return target


def format_prompt(
    prompt: Optional[TS],
    inputs: Dict[str, Any],
    refnames: Dict[str, Any] = None,
    include_brackets: bool = True,
) -> Optional[str]:
    """Format the prompt with inputs and (optional) refnames"""
    if prompt is None:
        return None
    if isinstance(prompt, str):
        return prompt
    if refnames is None:
        refnames = collect_refnames(dict(prompt=prompt))
    if include_brackets:
        refnames = add_brackets(refnames)
    check_duplicate_keys(inputs, refnames)
    return prompt(x={**inputs, **refnames})


def format_multiple_prompts(
    prompts: List[Optional[TS]],
    inputs: Dict[str, Any],
    refnames: Dict[str, Any] = None,
    include_brackets: Union[bool, List[bool]] = None,
) -> List[str]:
    """Format multiple prompts with inputs and (optional) refnames"""
    if refnames is None:
        refnames = collect_refnames(
            {
                f"prompt_{i}": prompt
                for i, prompt in enumerate(prompts)
                if prompt is not None
            }
        )
    if include_brackets is None:
        include_brackets = True
    if isinstance(include_brackets, bool):
        include_brackets = [include_brackets] * len(prompts)
    if len(prompts) != len(include_brackets):
        raise ValueError(
            f"Length of prompts ({len(prompts)}) and include_brackets ({len(include_brackets)}) should be the same"
        )

    results = []
    for prompt, include in zip(prompts, include_brackets):
        results.append(format_prompt(prompt, inputs, refnames, include))
    return results


def remove_direct_submodule(
    x: T, names: List[str] = None, refnames: List[str] = None
) -> T:
    """Remove a list of direct submodules by names or refnames"""
    if names is None and refnames is None:
        raise ValueError("Either names or refnames should be specified")
    y = copy.deepcopy(x)
    remains = OrderedDict()
    for name, module in x.named_children():
        if names is not None and name in names:
            continue
        if refnames is not None and module.refname in refnames:
            continue
        remains[name] = module
    y.set_submodules(remains)
    return y


def remove_submodule(x: T, refname: str, recursive=True) -> Tuple[T, bool]:
    """Remove a submodule [recursively] by refname"""
    # TODO: support remove multiple occurrences of the module
    y, success = x, False

    for name, module in x.named_children():
        if module.refname == refname:
            y = remove_direct_submodule(x, refnames=[refname])
            success = True
        elif recursive:
            removed, success_ = remove_submodule(module, refname, recursive)
            if success_:
                y = copy.deepcopy(x)
                y._replace_submodule(name, removed)
            success = success or success_
        if success:
            break
    return y, success


def removed_submodule(x: T, refname: str, recursive=True) -> T:
    """Remove a submodule [recursively] by refname"""
    assert isinstance(refname, str), f"refname should be a string, got {refname}"
    y, _ = remove_submodule(x, refname, recursive)
    return y


def removed_submodules(x: T, refnames: Union[List[str], str], recursive=True) -> T:
    """Remove a list of submodules [recursively] by refnames"""
    if isinstance(refnames, str):
        refnames = [refnames]
    for refname in refnames:
        x = removed_submodule(x, refname, recursive)
    return x


def replace_submodule(
    x: T, refname: str, new_module: T, recursive=True
) -> Tuple[T, bool]:
    """Replace a module [recursively] with a new module, by refname."""
    # TODO: support replacement for multiple occurrences of the module
    y = copy.deepcopy(x)
    success = False

    for name, module in x.named_children():
        if module.refname == refname:
            y._replace_submodule(name, new_module)
            # set new_module's refname to refname
            assert new_module.refname is None or new_module.refname == refname
            new_module.set_refname(refname)
            success = True
        elif recursive:
            replaced, success_ = replace_submodule(
                module, refname, new_module, recursive
            )
            if success_:
                y._replace_submodule(name, replaced)
            success = success or success_
        if success:
            break
    return y, success


def replaced_submodule(x: T, refname: str, new_module: T, recursive=True) -> T:
    """Replace a module [recursively] with a new module, by refname."""
    y, _ = replace_submodule(x, refname, new_module, recursive)
    return y


def find_submodule(x: T, refname: str) -> T:
    """Find a submodule by refname"""
    for module in x.get_all_submodules():
        if module.refname == refname:
            return module
    return None


def replace_prompt(x: T, replace_func: Callable[[str], str], inplace=False) -> T:
    """Replace the prompt of a module or its direct submodules by a replace function"""
    # TODO: maybe support recursive
    y = x if inplace else copy.deepcopy(x)
    if isinstance(y, Single):
        y.prompt = replace_func(y.prompt)
    else:
        for name, module in y.named_children():
            y._replace_submodule(name, replace_prompt(module, replace_func))
    return y

procoder/functional.py
dhh1995/PromptCoder
from .base import Module, Single, T
from .modules import *
from .sequential import *
from .utils import *

procoder/prompt/__init__.py
dhh1995/PromptCoder
from procoder.utils.my_typing import *

T = TypeVar("T", bound="Module")
TS = Union[T, str]

# Mimicing the pytorch design of nn.Module


class Module:
    r"""Base class for all prompt modules.

    Your prompts should also subclass this class.

    Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes::

        import procoder.prompt as pp
        import procoder.prompt.functional as PF

        class Prompt(pp.Module):
            def __init__(self):
                super().__init__()
                self.prompt1 = pp.Single("What is your name?")
                self.prompt2 = pp.Single("What is your quest?")

            def repr(self, x):
                return PF.concat([self.prompt1(x), self.prompt2(x)])

    Submodules assigned in this way will be registered.

    .. note::
        As per the example above, an ``__init__()`` call to the parent class
        must be made before assignment on the child.
    """

    _modules: Dict[str, Optional["Module"]]
    _sep: str = ""
    _delta_indent: str = ""
    _name: str = None
    _refname: str = None
    _debug: bool = False

    def __init__(self, *args, **kwargs) -> None:
        super().__setattr__("_modules", OrderedDict())

    def set_sep(self, sep: str) -> "Module":
        self._sep = sep
        return self

    def set_delta_indent(self, delta_indent: str) -> "Module":
        self._delta_indent = delta_indent
        return self

    def set_refname(self, refname: str) -> "Module":
        self._refname = refname
        return self

    def forward(
        self, newline: bool = True, indent: str = "", x: Dict[str, Any] = None
    ) -> TS:
        """Forward pass of the module.

        Args:
            newline: Whether this is a beginning of a new line.
            indent: Indentation string.
            x: Input to the module.

        Returns:
            Output of the module.
        """
        raise NotImplementedError

    def _call_impl(self, *args, **kwargs):
        if self._debug:
            print(f"entered {self} with refname={self.refname}")
        return self.forward(*args, **kwargs)

    __call__: Callable[..., Any] = _call_impl

    def add_module(self, name: str, module: Optional["Module"]) -> None:
        r"""Adds a child module to the current module.

        The module can be accessed as an attribute using the given name.

        Args:
            name (str): name of the child module. The child module can be
                accessed from this module using the given name
            module (Module): child module to be added to the module.
        """
        if not isinstance(module, Module) and module is not None:
            raise TypeError("{} is not a Module subclass".format(typename(module)))
        elif not isinstance(name, str):
            raise TypeError(
                "module name should be a string. Got {}".format(typename(name))
            )
        elif hasattr(self, name) and name not in self._modules:
            raise KeyError("attribute '{}' already exists".format(name))
        elif "." in name:
            raise KeyError('module name can\'t contain ".", got: {}'.format(name))
        elif name == "":
            raise KeyError('module name can\'t be empty string ""')
        self._modules[name] = module

    def _replace_submodule(self, name: str, module: Optional["Module"]) -> None:
        """Replace one of the submodule of this module by name."""
        if not isinstance(module, Module) and module is not None:
            raise TypeError("{} is not a Module subclass".format(typename(module)))
        if name not in self._modules:
            raise KeyError("attribute '{}' not exists".format(name))
        self._modules[name] = module

    def set_submodules(self, modules: Dict[str, Optional["Module"]]) -> None:
        """Sets the submodules of this module."""
        for name, module in modules.items():
            if not isinstance(module, Module) and module is not None:
                raise TypeError("{} is not a Module subclass".format(typename(module)))
        self._modules = modules

    def get_submodule(self, target: str) -> "Module":
        """
        Returns the submodule given by ``target`` if it exists,
        otherwise throws an error.

        Args:
            target: The fully-qualified string name of the submodule
                to look for.

        Returns:
            Module: The submodule referenced by ``target``

        Raises:
            AttributeError: If the target string references an invalid
                path or resolves to something that is not an
                ``Module``
        """
        if target == "":
            return self

        atoms: List[str] = target.split(".")
        mod: "Module" = self

        for item in atoms:
            if not hasattr(mod, item):
                raise AttributeError(
                    mod._get_name() + " has no " "attribute `" + item + "`"
                )

            mod = getattr(mod, item)

            if not isinstance(mod, Module):
                raise AttributeError("`" + item + "` is not " "an Module")

        return mod

    @property
    def refname(self):
        return self._refname

    def _get_name(self):
        if self._name:
            return self._name
        return self.__class__.__name__

    def __getattr__(self, name: str) -> "Module":
        if "_modules" in self.__dict__:
            modules = self.__dict__["_modules"]
            if name in modules:
                return modules[name]
            raise AttributeError(
                "'{}' object has no attribute '{}'".format(type(self).__name__, name)
            )
        else:  # to support copy
            return object.__getattribute__(self, name)

    def children(self) -> Iterator["Module"]:
        r"""Returns an iterator over immediate children modules.

        Yields:
            Module: a child module
        """
        for name, module in self.named_children():
            yield module

    def named_children(self) -> Iterator[Tuple[str, "Module"]]:
        r"""Returns an iterator over immediate children modules, yielding both
        the name of the module as well as the module itself.

        Yields:
            (str, Module): Tuple containing a name and child module
        """
        memo = set()
        for name, module in self._modules.items():
            if module is not None and module not in memo:
                memo.add(module)
                yield name, module

    def get_all_submodules(self) -> Set["Module"]:
        r"""Returns an iterator over all submodules.

        Yields:
            Module: a submodule
        """
        modules = set()
        for module in self.children():
            modules.add(module)
            modules |= module.get_all_submodules()
        return modules


class Single(Module):
    def __init__(self, prompt: str, need_format: bool = True):
        super().__init__()
        self.prompt = prompt
        self._need_format = need_format

    def forward(self, newline: bool = True, indent: str = "", x: Dict[str, Any] = None):
        content = self.prompt
        if self._need_format and x is not None:
            content = content.format(**x)
        return indent + content if newline else content


def as_module(x: TS):
    if isinstance(x, Module):
        return x
    elif isinstance(x, str):
        return Single(x)
    else:
        raise TypeError(f"Expected Module or str, but got {type(x)}")

procoder/prompt/base.py
dhh1995/PromptCoder
from functools import partial

from procoder.utils.my_typing import *
from procoder.utils.my_typing import Any, Dict

from .base import TS, Module, T, as_module
from .sequential import *
from .utils import make_ordered_dict


class Paired(Sequential):
    """A pair of modules"""

    def __init__(self, first: TS, second: TS):
        super().__init__(make_ordered_dict(["first", "second"], [first, second]))

    def _post_init(self):
        self._sep = " "


class Collection(Sequential):
    """A collection of multiple modules, seperated by newlines."""

    def _post_init(self):
        self.set_indexing_method(number_indexing)


class Block(Sequential):
    """A prompt block."""

    def __init__(self, head: TS, body: TS):
        super().__init__(make_ordered_dict(["head", "body"], [head, body]))


class NamedBlock(Sequential):
    """A named prompt block.

    Args:
        name: Name of the block.
        content: Content of the block.
        refname: the name to reference the block.
        shortname: Short name of the block, used when referenced.
    """

    def __init__(self, name: TS, content: TS, refname: str = None):
        super().__init__(make_ordered_dict(["name", "content"], [name, content]))
        if isinstance(name, str):
            self._name = name
        self._refname = refname


class NamedVariable(NamedBlock):
    def _post_init(self):
        self._sep = ": "

procoder/prompt/modules.py
dhh1995/PromptCoder
from typing import overload

from procoder.utils.my_typing import *

from .base import TS, Module, T, as_module


class AddIndentProxy(Module):
    """Proxy to add indent."""

    _modules: Dict[str, T]

    def __init__(self, module: T, indent: str):
        super().__init__()
        self.add_module("prompt", as_module(module))
        self._delta = indent

    def forward(self, newline=True, indent="", x=None):
        return self.prompt(newline, indent + self._delta, x)


class SilenceProxy(Module):
    """Proxy to silence the output."""

    _modules: Dict[str, T]

    def __init__(self, module: T):
        super().__init__()
        self.add_module("prompt", as_module(module))

    def forward(self, newline=True, indent="", x=None):
        return ""

procoder/prompt/proxy.py
dhh1995/PromptCoder
from typing import overload

from procoder.utils.my_typing import *

from .base import TS, Module, T, as_module
from .utils import *


class Sequential(Module):
    """A sequential combination of multiple modules."""

    _modules: Dict[str, Module]
    _indexing_method: Callable[[int], str] = None
    _name_enabled: bool = False

    @overload
    def __init__(self, *args: Module) -> None:
        ...

    @overload
    def __init__(self, arg: OrderedDict[str, Module]) -> None:
        ...

    def __init__(self, *args):
        super().__init__()
        if len(args) == 1 and isinstance(args[0], OrderedDict):
            for key, val in args[0].items():
                self.add_module(key, as_module(val))
        else:
            # if len(args) == 1 and isinstance(args[0], str):
            #     print(
            #         "Warning[PromptCoder]: Only one string is passed to "
            #         "a Sequential Module, did you forget to add commas? "
            #         f'The first 20 chars of the string argument are: "{args[0][:20]}".'
            #     )
            for idx, val in enumerate(args):
                if val is None:
                    continue
                v = as_module(val)
                key = v.refname or f"_{idx}"
                self.add_module(key, v)
        self._sep = "\n"
        self._post_init()

    def _post_init(self):
        pass

    def enable_name(self):
        self._name_enabled = True
        return self

    @property
    def indexing_method(self):
        return self._indexing_method

    @property
    def name_enabled(self):
        return self._name_enabled

    def set_indexing_method(self, func: Callable[[int], str]):
        self._indexing_method = func
        return self

    def forward(self, newline=True, indent="", x=None):
        indent += self._delta_indent
        sep = self._sep
        res = []
        for i, (k, p) in enumerate(self.named_children()):
            if i > 0:
                res.append(sep)

            # compute prefix based on indexing method and naming method
            prefix = ""
            if self.indexing_method:
                prefix += self.indexing_method(i)
            if self.name_enabled:
                prefix += f"{k}: "
            if prefix != "":
                if newline:
                    prefix = indent + prefix
                newline = False
            # combine prefix and content
            res.append(prefix + p(newline, indent, x))
            newline = sep.endswith("\n")
        return "".join(res)

procoder/prompt/sequential.py
dhh1995/PromptCoder
import collections

import roman
from procoder.utils.my_typing import *

# from .base import Module, T, TS, as_module


def make_ordered_dict(names, values):
    return collections.OrderedDict(zip(names, values))


def number_indexing(i: int) -> str:
    return f"{i + 1}. "


def letter_indexing(i: int) -> str:
    return f"{chr(ord('a') + i)}. "


def uppercase_letter_indexing(i: int) -> str:
    return f"{chr(ord('A') + i)}. "


def roman_numeral_indexing(i: int) -> str:
    return f"{roman.toRoman(i + 1)}. "


def star_indexing(i: int) -> str:
    return "* "


def dash_indexing(i: int) -> str:
    return "- "


def get_sharp_indexing(n: int) -> Callable[[int], str]:
    return lambda i: "#" * n + " "


def get_equal_indexing(n: int) -> Callable[[int], str]:
    return lambda i: "=" * n + " "


sharp1_indexing = get_sharp_indexing(1)
sharp2_indexing = get_sharp_indexing(2)
sharp3_indexing = get_sharp_indexing(3)
sharp4_indexing = get_sharp_indexing(4)

equal2_indexing = get_equal_indexing(2)
equal3_indexing = get_equal_indexing(3)
equal4_indexing = get_equal_indexing(4)

procoder/prompt/utils.py
dhh1995/PromptCoder
from .my_typing import *

procoder/utils/__init__.py
dhh1995/PromptCoder
from typing import (
    Any,
    Callable,
    Dict,
    Iterator,
    List,
    Mapping,
    NamedTuple,
    Optional,
    OrderedDict,
    Sequence,
    Set,
    Tuple,
    Type,
    TypeVar,
    Union,
)


def typename(o):
    module = ""
    class_name = ""
    if (
        hasattr(o, "__module__")
        and o.__module__ != "builtins"
        and o.__module__ != "__builtin__"
        and o.__module__ is not None
    ):
        module = o.__module__ + "."

    if hasattr(o, "__qualname__"):
        class_name = o.__qualname__
    elif hasattr(o, "__name__"):
        class_name = o.__name__
    else:
        class_name = o.__class__.__name__

    return module + class_name

procoder/utils/my_typing.py
dhh1995/PromptCoder
import setuptools

with open("README.md", "r") as fh:
    long_description = fh.read()

with open("requirements.txt", "r") as f:
    requirements = f.read().splitlines()

setuptools.setup(
    name="procoder",
    version="0.0.1",
    author="Honghua Dong, Yangjun Ruan",
    author_email="honghuad@cs.toronto.edu, yjruan@cs.toronto.edu",
    description="A package for writing prompts in a coding style.",
    long_description=long_description,
    long_description_content_type="text/markdown",
    url="https://github.com/dhh1995/PromptCoder",
    packages=setuptools.find_packages(),
    install_requires=requirements,
    classifiers=[
        "Programming Language :: Python :: 3",
        "License :: OSI Approved :: MIT License",
        "Operating System :: OS Independent",
    ],
    python_requires=">=3.6",
)

setup.py
dhh1995/PromptCoder