diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..02c054b --- /dev/null +++ b/__init__.py @@ -0,0 +1,4 @@ +from .qa_prompt import QAPrompt + +def question_answer(): + return QAPrompt() diff --git a/qa_prompt.py b/qa_prompt.py new file mode 100644 index 0000000..213dc25 --- /dev/null +++ b/qa_prompt.py @@ -0,0 +1,29 @@ +from typing import List, Tuple, Dict + + +class QAPrompt: + def __init__(self): + super().__init__() + self._template = """ +'Use the following pieces of context to answer the question at the end. +If you don''t know the answer, just say that you don''t know, don''t try to make +up an answer. + +{context} + +Question: {question} + +Helpful Answer:' +""" + + def __call__(self, question: str, docs: List[str], history=None) -> List[Dict[str, str]]: + context = '\n'.join(docs) + prompt_str = self._template.format(context=context, question=question) + ret = [{'question': prompt_str}] + if not isinstance(history, list): + return ret + else: + history_data = [] + for item in history: + history_data.append({'question': item[0], 'answer': item[1]}) + return history + ret