69 lines
2.4 KiB
Python
69 lines
2.4 KiB
Python
import string
|
|
from typing import Any, Union
|
|
|
|
|
|
class PromptTemplate:
|
|
"""A prompt template for a language model.
|
|
|
|
A prompt template consists of an array of strings that will be concatenated together. It accepts
|
|
a set of parameters from the user that can be used to generate a prompt for a language model.
|
|
|
|
The template can be formatted using f-strings.
|
|
|
|
Example:
|
|
|
|
.. code-block:: python
|
|
|
|
from mlflow.metrics.genai.prompt_template import PromptTemplate
|
|
|
|
# Instantiation using initializer
|
|
prompt = PromptTemplate(template_str="Say {foo} {baz}")
|
|
|
|
# Instantiation using partial_fill
|
|
prompt = PromptTemplate(template_str="Say {foo} {baz}").partial_fill(foo="bar")
|
|
|
|
# Format the prompt
|
|
prompt.format(baz="qux")
|
|
"""
|
|
|
|
def __init__(self, template_str: Union[str, list[str]]):
|
|
self.template_strs = [template_str] if isinstance(template_str, str) else template_str
|
|
|
|
@property
|
|
def variables(self):
|
|
return {
|
|
fname
|
|
for template_str in self.template_strs
|
|
for _, fname, _, _ in string.Formatter().parse(template_str)
|
|
if fname
|
|
}
|
|
|
|
def format(self, **kwargs: Any) -> str:
|
|
safe_kwargs = {k: v for k, v in kwargs.items() if v is not None}
|
|
formatted_strs = []
|
|
for template_str in self.template_strs:
|
|
extracted_variables = [
|
|
fname for _, fname, _, _ in string.Formatter().parse(template_str) if fname
|
|
]
|
|
if all(item in safe_kwargs.keys() for item in extracted_variables):
|
|
formatted_strs.append(template_str.format(**safe_kwargs))
|
|
|
|
return "".join(formatted_strs)
|
|
|
|
def partial_fill(self, **kwargs: Any) -> "PromptTemplate":
|
|
safe_kwargs = {k: v for k, v in kwargs.items() if v is not None}
|
|
new_template_strs = []
|
|
for template_str in self.template_strs:
|
|
extracted_variables = [
|
|
fname for _, fname, _, _ in string.Formatter().parse(template_str) if fname
|
|
]
|
|
safe_available_kwargs = {
|
|
k: safe_kwargs.get(k, "{" + k + "}") for k in extracted_variables
|
|
}
|
|
new_template_strs.append(template_str.format_map(safe_available_kwargs))
|
|
|
|
return PromptTemplate(template_str=new_template_strs)
|
|
|
|
def __str__(self) -> str:
|
|
return "".join(self.template_strs)
|