Skip to content

Commit 45a4834

Browse files
committed
added PromptTemplateAdapter
1 parent 4f48d7c commit 45a4834

File tree

1 file changed

+31
-0
lines changed

1 file changed

+31
-0
lines changed
Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
class PromptTemplateAdapter:
2+
""" Makes DSPy base prompts better with custom model templates, for an enhanced text output. """
3+
4+
templates = {
5+
"mistral": "[INST] {prompt} [/INST]",
6+
"llama": "<s>[INST] {prompt} [/INST]",
7+
"chatml": "<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant",
8+
"default": "{prompt}",
9+
}
10+
# Templates currently made for only 3 models as samples. This can be extended.
11+
# (function written below
12+
13+
14+
def __init__(self, model_family: str = "default"):
15+
if model_family not in self.templates:
16+
self.template = self.templates["default"]
17+
else:
18+
self.template = self.templates[model_family]
19+
20+
self.model_family = model_family
21+
22+
23+
def wrap(self, prompt: str) -> str:
24+
return self.template.format(prompt=prompt)
25+
26+
27+
# Makes addition of models and the templates possible.
28+
@classmethod
29+
def register_template(cls, name: str, template: str) -> None:
30+
"""Allow users to add custom model families."""
31+
cls.templates[name] = template

0 commit comments

Comments
 (0)