Update app.py

app.py

@@ -1,5 +1,7 @@
 from sentence_transformers import SentenceTransformer
+from sentence_transformers.cross_encoder import CrossEncoder
 import os
+import numpy as np
 from pymilvus import connections, utility, Collection, DataType, FieldSchema, CollectionSchema
 import streamlit as st
 
@@ -7,8 +9,9 @@ import streamlit as st
 @st.cache_resource
 def load_sentence_transformer():
     sent_model = SentenceTransformer('all-mpnet-base-v2')
+    ce_model = CrossEncoder('cross-encoder/stsb-distilroberta-base')
     print('loaded sentence transformer')
-    return sent_model
+    return sent_model, ce_model
 
 
 class TextVectorizer:
@@ -39,20 +42,23 @@ def find_similar_news(text: str, top_n: int=5):
                                 param=search_params,
                                 limit=top_n,
                                 guarantee_timestamp=1, 
-                                output_fields=['article_desc', 'article_category']) # which fields to return in output
+                                output_fields=['article_desc']) # which fields to return in output
 
     
-    output_dict = {"input_text": text, "similar_texts": [hit.entity.get('article_desc') for hits in result for hit in hits], 
-                  "text_category": [hit.entity.get('article_category') for hits in result for hit in hits]} 
-    txt_category = [f'<li><b>{txt}</b> (<i>{cat}</i>)</li>' for txt, cat in zip(output_dict.get('similar_texts'), output_dict.get('text_category'))]
-    similar_txt = ''.join(txt_category)
+    output_dict = {"input_text": text, "similar_texts": [hit.entity.get('article_desc') for hits in result for hit in hits]} 
+    texts = np.array(output_dict.get('similar_texts'))
+    ce_similarity_scores = np.array(ce_model.predict([[text, output_text] for output_text in texts]))
+    similarity_idxs = [*np.argsort(ce_similarity_scores)[::-1]]
+    texts = texts[similarity_idxs]
+    ce_similarity_scores = ce_similarity_scores[similarity_idxs]
+    txt_similarity = [f'<li><b>{txt}</b> (<i>similarity: {np.round(sim, 5)})</i></li>' for txt, sim in zip(texts, ce_similarity_scores)]
+    similar_txt = ''.join(txt_similarity)
     return f"<h4>Similar News Articles</h4><ol>{similar_txt}</ol>"
-    # return output_dict
     
 
 vectorizer = TextVectorizer()
 collection = get_milvus_collection()
-sent_model = load_sentence_transformer()
+sent_model, ce_model = load_sentence_transformer()
 
 
 def main():
@@ -62,18 +68,17 @@ def main():
     desc = '''<p style="font-size: 13px;">
     Embeddings of 300,000 news headlines are stored in Milvus vector database, used as a feature store. 
     Embeddings of the input headline are computed using sentence transformers (all-mpnet-base-v2). 
-    Similar news headlines are retrieved from the vector database using Euclidean distance as similarity metric. 
+    Similar news headlines are retrieved from the vector database using Euclidean distance as similarity metric and are reranked using cross encoder. 
     <span style="color: red;">This method (all-mpnet-base-v2) has the best performance compared to multi-qa-distilbert-cos-v1 fine-tuned using TSDAE 
     and extracting embeddings from fine-tuned DistilBERT classifier.</span>
     </p> 
     '''
     st.markdown(desc, unsafe_allow_html=True)
-    news_txt = st.text_area("Paste the headline of a news article:", "", height=50)
-    top_n = st.slider('Select number of similar articles to display', 1, 100, 10)
+    news_txt = st.text_area("Paste the headline of a news article and hit Ctrl+Enter:", "", height=30)
+    top_n = st.slider('Select the number of similar articles to display', 1, 100, 15)
     
-    if st.button("Submit"):
+    if news_txt:
         result = find_similar_news(news_txt, top_n)
-        # st.write(result)
         st.markdown(result, unsafe_allow_html=True)