From b1ff0c55422d7b0af2c379679b8721014ef36926 Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Tue, 18 Jun 2024 01:37:32 +0200 Subject: Wip rewrite --- rag/generator/prompt.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) (limited to 'rag/generator/prompt.py') diff --git a/rag/generator/prompt.py b/rag/generator/prompt.py index 6523842..4840fdc 100644 --- a/rag/generator/prompt.py +++ b/rag/generator/prompt.py @@ -15,3 +15,27 @@ ANSWER_INSTRUCTION = ( class Prompt: query: str documents: List[Document] + generator_model: str + + def __context(self, documents: List[Document]) -> str: + results = [ + f"Document: {i}\ntitle: {doc.title}\ntext: {doc.text}" + for i, doc in enumerate(documents) + ] + return "\n".join(results) + + def to_str(self) -> str: + if self.generator_model == "cohere": + return f"{self.query}\n\n{ANSWER_INSTRUCTION}" + else: + return ( + "Context information is below.\n" + "---------------------\n" + f"{self.__context(self.documents)}\n\n" + "---------------------\n" + f"{ANSWER_INSTRUCTION}" + "Do not attempt to answer the query without relevant context and do not use" + " prior knowledge or training data!\n" + f"Query: {self.query.strip()}\n\n" + "Answer:" + ) -- cgit v1.2.3-70-g09d2