|
|
|
|
from typing import Optional
|
|
|
|
|
from dataclasses import dataclass
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@dataclass
|
|
|
|
|
class Template:
|
|
|
|
|
|
|
|
|
|
name: str
|
|
|
|
|
|
|
|
|
|
def get_prompt(self, query: str, history: Optional[list] = None, prefix: Optional[str] = "") -> str:
|
|
|
|
|
return getattr(self, "_format_{}".format(self.name))(query, history, prefix)
|
|
|
|
|
|
|
|
|
|
def _format_alpaca(self, query: str, history: Optional[list], prefix: Optional[str] = "") -> str:
|
|
|
|
|
if prefix:
|
|
|
|
|
prompt = prefix
|
|
|
|
|
else:
|
|
|
|
|
prompt = "Below is an instruction that describes a task. "
|
|
|
|
|
prompt += "Write a response that appropriately completes the request.\n"
|
|
|
|
|
prompt += "Instruction:\n"
|
|
|
|
|
if history:
|
|
|
|
|
for old_query, response in history:
|
|
|
|
|
prompt += "Human:{}\nAssistant:{}\n".format(old_query, response)
|
|
|
|
|
prompt += "Human:{}\nAssistant:".format(query)
|
|
|
|
|
return prompt
|
|
|
|
|
|
|
|
|
|
def _format_vicuna(self, query: str, history: Optional[list], prefix: Optional[str] = "") -> str:
|
|
|
|
|
if prefix:
|
|
|
|
|
prompt = prefix
|
|
|
|
|
else:
|
|
|
|
|
prompt = "A chat between a curious user and an artificial intelligence assistant. "
|
|
|
|
|
prompt += "The assistant gives helpful, detailed, and polite answers to the user's questions. "
|
|
|
|
|
if history:
|
|
|
|
|
for old_query, response in history:
|
|
|
|
|
prompt += "USER: {} ASSISTANT: {}</s>".format(old_query, response)
|
|
|
|
|
prompt += "USER: {} ASSISTANT:".format(query)
|
|
|
|
|
return prompt
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _format_ziya(self, query: str, history: Optional[list], prefix: Optional[str] = "") -> str:
|
|
|
|
|
prompt = prefix
|
|
|
|
|
if history:
|
|
|
|
|
for old_query, response in history:
|
|
|
|
|
prompt += "<human>:{}\n<bot>:{}\n".format(old_query, response)
|
|
|
|
|
prompt += "<human>:{}\n<bot>:".format(query)
|
|
|
|
|
return prompt
|
