reranker added

app.py

@@ -11,7 +11,7 @@ from jinja2 import Environment, FileSystemLoader
 
 from backend.query_llm import generate_hf, generate_openai
 from backend.semantic_search import retrieve
-
+from study.backend.reranker import rerank
 
 TOP_K = int(os.getenv("TOP_K", 4))
 
@@ -44,11 +44,16 @@ def bot(history, api_kind):
     # Retrieve documents relevant to query
     document_start = perf_counter()
 
-    documents = retrieve(query, TOP_K)
+    documents = retrieve(query, 10)
 
     document_time = perf_counter() - document_start
     logger.info(f'Finished Retrieving documents in {round(document_time, 2)} seconds...')
 
+    rerank_start = perf_counter()
+    documents = rerank(query, documents, TOP_K)
+    rerank_time = perf_counter() - rerank_start
+    logger.info(f'Finished Reranking documents in {round(rerank_time, 2)} seconds...')
+
     # Create Prompt
     prompt = template.render(documents=documents, query=query)
     prompt_html = template_html.render(documents=documents, query=query)

backend/reranker.py

@@ -0,0 +1,10 @@
+from FlagEmbedding import FlagReranker
+
+reranker = FlagReranker('BAAI/bge-reranker-large',
+                        use_fp16=True)
+
+
+def rerank(query: str, documents: [str], k: int) -> [str]:
+    scores = reranker.compute_score([(query, document) for document in documents])
+    sorted_docs = sorted(enumerate(scores), key=lambda x: x[1], reverse=True)
+    return [sorted_docs[i] for i, _ in sorted_docs[:k]]