From 75be0914f6bd2cdeda1539f83b38fcbc854d5cfa Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Tue, 23 Apr 2024 22:08:03 +0200 Subject: Add reranking as a separate step --- rag/cli.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) (limited to 'rag/cli.py') diff --git a/rag/cli.py b/rag/cli.py index 690563e..932e2a9 100644 --- a/rag/cli.py +++ b/rag/cli.py @@ -37,13 +37,13 @@ def rag(generator: str, query: str, limit): retriever = Retriever() generator = get_generator(generator) documents = retriever.retrieve(query, limit=limit) - prompt = Prompt(query, documents) + prompt = generator.rerank(Prompt(query, documents)) print("Answer: ") for chunk in generator.generate(prompt): print(chunk, end="", flush=True) print("\n\n") - for i, doc in enumerate(documents): + for i, doc in enumerate(prompt.documents): print(f"### Document {i}") print(f"**Title: {doc.title}**") print(doc.text) -- cgit v1.2.3-70-g09d2