--- a +++ b/utils/prompter.py @@ -0,0 +1,51 @@ +""" +A dedicated helper to manage templates and prompt building. +""" + +import json +import os.path as osp +from typing import Union + + +class Prompter(object): + __slots__ = ("template", "_verbose") + + def __init__(self, template_name: str = "", verbose: bool = False): + self._verbose = verbose + if not template_name: + # Enforce the default here, so the constructor can be called with '' and will not break. + template_name = "alpaca" + file_name = osp.join("data", "templates", f"{template_name}.json") + if not osp.exists(file_name): + raise ValueError(f"Can't read {file_name}") + with open(file_name) as fp: + self.template = json.load(fp) + if self._verbose: + print( + f"Using prompt template {template_name}: {self.template['description']}" + ) + + def generate_prompt( + self, + instruction: str, + input: Union[None, str] = None, + label: Union[None, str] = None, + ) -> str: + # returns the full prompt from instruction and optional input + # if a label (=response, =output) is provided, it's also appended. + if input: + res = self.template["prompt_input"].format( + instruction=instruction, input=input + ) + else: + res = self.template["prompt_no_input"].format( + instruction=instruction + ) + if label: + res = f"{res}{label}" + if self._verbose: + print(res) + return res + + def get_response(self, output: str) -> str: + return output.split(self.template["response_split"])[-1].strip() #-1 to allow for multi-turn prompts \ No newline at end of file