Update app.py

app.py

@@ -2,6 +2,7 @@ import requests
 import os
 import gradio as gr
 import time
+import heapq
 from utils import package_installer
 
 package_installer('sentence_transformers')
@@ -23,13 +24,29 @@ def query(payload):
     
 
 def chat(message,history):
-    chat_history = ''
-    for chats in history:
-        # formatted_assistant_msg = chats[1].replace(chats[0],'').strip().removesuffix('<|end|>')
-        chat_history += chats[1]
-
+    
     prompt = f"<|user|>\n{message}<|end|>\n<|assistant|>"
-    user_input = chat_history + prompt
+    chat_history = []
+
+    if history:
+
+        for chats in history:
+            chat_history.append(chats[1])
+
+        emb_prompt = embedding_model.encode(prompt,convert_to_tensor=True)
+        emb_chat_history = embedding_model.encode(chat_history,convert_to_tensor=True)
+        cosine_similarity_scores = util.cos_sim(emb_prompt,emb_chat_history)
+        top_2_scores = heapq.nlargest(2,cosine_similarity_scores[0])
+        top_2_chats = [chat_history[i] for i in sorted(list(cosine_similarity_scores[0]).index(val) for val in top_2_scores)]
+        
+    similar_chat_history = ''
+    if chat_history:
+        for chats in chat_history:
+            # formatted_assistant_msg = chats[1].replace(chats[0],'').strip().removesuffix('<|end|>')
+            similar_chat_history += chats
+
+    #prompt = f"<|user|>\n{message}<|end|>\n<|assistant|>"
+    user_input = similar_chat_history + prompt
     inp_dict = {"inputs":user_input,
                 "parameters": {"max_new_tokens":750,"temperature":0.01}}
     output = query(inp_dict)