File tree Expand file tree Collapse file tree 1 file changed +31
-0
lines changed
integrations/dspy/src/haystack_integrations/components/generators/dspy Expand file tree Collapse file tree 1 file changed +31
-0
lines changed Original file line number Diff line number Diff line change 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
You can’t perform that action at this time.
0 commit comments