Skip to content

Commit da9cbd4

Browse files
committed
Re-Ranking of Vector Search Results! SentenceTransformer('all-MiniLM-L6-v2') used to re-rank and filter similarity-search results. LARS is now officially a RRAG app: Retrieval & Re-Ranking AR!
1 parent 6227c76 commit da9cbd4

File tree

1 file changed

+32
-3
lines changed

1 file changed

+32
-3
lines changed

web_app/app.py

Lines changed: 32 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22
from flask import send_from_directory
33
from flask import jsonify
44

5+
from sentence_transformers import SentenceTransformer, util
6+
57
from pdfminer.high_level import extract_text
68
from werkzeug.utils import secure_filename
79

@@ -2928,6 +2930,30 @@ def filter_relevant_documents(query, search_results, threshold=1):
29282930
return page_contents, do_rag
29292931

29302932

2933+
def rerank_results_ml(query, documents, top_n=5):
2934+
# Load pre-trained SBERT model
2935+
model = SentenceTransformer('all-MiniLM-L6-v2')
2936+
2937+
# Encode the query
2938+
query_embedding = model.encode(query, convert_to_tensor=True)
2939+
2940+
# Encode the documents
2941+
doc_embeddings = model.encode([doc.page_content for doc in documents], convert_to_tensor=True)
2942+
2943+
# Compute cosine similarities
2944+
cosine_scores = util.pytorch_cos_sim(query_embedding, doc_embeddings)[0]
2945+
2946+
# Create a list of (index, score) tuples
2947+
indexed_scores = list(enumerate(cosine_scores))
2948+
2949+
# Sort by score in descending order
2950+
sorted_indexes = sorted(indexed_scores, key=lambda x: x[1], reverse=True)
2951+
2952+
# Reorder the original documents based on the sorted indexes
2953+
ranked_documents = [documents[idx] for idx, _ in sorted_indexes[:top_n]]
2954+
2955+
return ranked_documents
2956+
29312957

29322958
@app.route('/setup_for_llama_cpp_response', methods=['POST'])
29332959
def setup_for_llama_cpp_response():
@@ -2993,7 +3019,7 @@ def setup_for_llama_cpp_response():
29933019
try:
29943020
# docs = VECTOR_STORE.similarity_search(user_query, embedding_fn=embedding_function)
29953021
# docs_with_relevance_score = VECTOR_STORE.similarity_search_with_relevance_scores(user_query, 10, embedding_fn=embedding_function)
2996-
docs_list_with_cosine_distance = VECTOR_STORE.similarity_search_with_score(user_query, 7, embedding_fn=embedding_function)
3022+
docs_list_with_cosine_distance = VECTOR_STORE.similarity_search_with_score(user_query, 11, embedding_fn=embedding_function)
29973023
# print(f'\n\nsimple similarity search results: \n {docs}\n\n')
29983024
# print(f'\n\nRelevance Score similarity search results (range 0 to 1): \n {docs_with_relevance_score}\n\n')
29993025
# print(f'\n\nDocs list most similar to query based on cosine distance: \n {docs_list_with_cosine_distance}\n\n')
@@ -3005,7 +3031,10 @@ def setup_for_llama_cpp_response():
30053031
# if score >= similarity_threshold
30063032
# ]
30073033

3008-
docs = [doc for doc, score in docs_list_with_cosine_distance]
3034+
filtered_docs = [doc for doc, score in docs_list_with_cosine_distance]
3035+
3036+
docs = rerank_results_ml(user_query, filtered_docs, top_n=5)
3037+
30093038

30103039
print("\n\nDetermining do_rag \n\n")
30113040
# We do not modify the force_enable_rag or force_disable_rag flags in this method, we simply respond to them here. UI updates should handle those flags.
@@ -3320,7 +3349,7 @@ def get_references():
33203349

33213350
if docs_have_relevant_info:
33223351

3323-
refer_pages_string = "<br><br><h6>Refer to the following pages in the mentioned docs:</h6>"
3352+
refer_pages_string = "<br><br><h6>Additional data may be found in the following documents & pages:</h6>"
33243353

33253354
for index, doc in enumerate(user_should_refer_pages_in_doc, start=1):
33263355
pdf_iframe_id = f"stream{stream_session_id}PdfViewer{str(index)}"

0 commit comments

Comments
 (0)