diff options
author | Gustaf Rydholm <gustaf.rydholm@gmail.com> | 2024-06-18 01:37:32 +0200 |
---|---|---|
committer | Gustaf Rydholm <gustaf.rydholm@gmail.com> | 2024-06-18 01:37:32 +0200 |
commit | b1ff0c55422d7b0af2c379679b8721014ef36926 (patch) | |
tree | 52aa88b2a8a0bba07f968c6ae24c002ce2d44226 /rag/generator/prompt.py | |
parent | b8c6a78f70d84f3360461aa91864e8538569d450 (diff) |
Wip rewrite
Diffstat (limited to 'rag/generator/prompt.py')
-rw-r--r-- | rag/generator/prompt.py | 24 |
1 files changed, 24 insertions, 0 deletions
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:" + ) |