From 75be0914f6bd2cdeda1539f83b38fcbc854d5cfa Mon Sep 17 00:00:00 2001
From: Gustaf Rydholm <gustaf.rydholm@gmail.com>
Date: Tue, 23 Apr 2024 22:08:03 +0200
Subject: Add reranking as a separate step

---
 rag/ui.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

(limited to 'rag/ui.py')

diff --git a/rag/ui.py b/rag/ui.py
index fb02e5c..2fbf8de 100644
--- a/rag/ui.py
+++ b/rag/ui.py
@@ -99,11 +99,13 @@ def generate_chat(query: str):
     documents = retriever.retrieve(query, limit=15)
     prompt = Prompt(query, documents)
 
+    prompt = generator.rerank(prompt)
+
     with st.chat_message(ss.bot):
         response = st.write_stream(generator.generate(prompt))
 
-    display_context(documents)
-    store_chat(query, response, documents)
+    display_context(prompt.documents)
+    store_chat(query, response, prompt.documents)
 
 
 def store_chat(query: str, response: str, documents: List[Document]):
-- 
cgit v1.2.3-70-g09d2