Update app.py

app.py

@@ -8,7 +8,6 @@ from langchain.chat_models import ChatOpenAI
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import Pinecone
 from langchain.chains import LLMChain
-from langchain.chains.retrieval_qa.base import RetrievalQA
 from langchain.chains.question_answering import load_qa_chain
 import pinecone
 
@@ -33,6 +32,7 @@ EMBEDDING_MODEL = os.environ.get("PINECONE_INDEX", "sentence-transformers/all-mp
 # return top-k text chunks from vector store
 TOP_K_DEFAULT = 15
 TOP_K_MAX = 30
+SCORE_DEFAULT = 0.3
 
 
 BUTTON_MIN_WIDTH = 215
@@ -106,6 +106,45 @@ init_message = f"""This demonstration website is based on \
     2. Insert your **Question** and click  `{KEY_SUBMIT}`
 """
 
+PROMPT_DOC = PromptTemplate(
+    input_variables=["context", "chat_history", "question"],
+    template="""Context:
+##
+{context}
+##
+
+Chat History:
+##
+{chat_history}
+##
+
+Question:
+{question}
+
+Optinal:
+Don't use standalone clause/figure name in the answer, expand it with corresponding metadata TS name
+
+Desired format:
+Clause/figure name: <dot_separated_numbers>
+TS name: [\w\.]
+
+Answer:"""
+)
+
+PROMPT_BASE = PromptTemplate(
+    input_variables=['question', "chat_history"],
+    template="""Chat History:
+##
+{chat_history}
+##
+
+Question:
+##
+{question}
+##
+
+Answer:"""
+)
 
 #----------------------------------------------------------------------------------------------------------
 #----------------------------------------------------------------------------------------------------------
@@ -122,7 +161,8 @@ def init_model(api_key, emb_name, db_api_key, db_env, db_index):
             if llm_name == "gpt-3.5-turbo":
                 llm_dict[llm_name] = ChatOpenAI(model_name=llm_name,
                                                 temperature = OPENAI_TEMP,
-                                                openai_api_key = api_key)
+                                                openai_api_key = api_key
+                                                )
             else:
                 llm_dict[llm_name] = OpenAI(model_name=llm_name,
                                             temperature = OPENAI_TEMP,
@@ -148,22 +188,25 @@ def init_model(api_key, emb_name, db_api_key, db_env, db_index):
 def get_chat_history(inputs) -> str:
     res = []
     for human, ai in inputs:
-        res.append(f"Human: {human}\nAI: {ai}")
+        res.append(f"Q: {human}\nA: {ai}")
     return "\n".join(res)
 
-def remove_duplicates(documents):
+def remove_duplicates(documents, score_min):
     seen_content = set()
     unique_documents = []
-    for doc in documents:
-        if doc.page_content not in seen_content:
+    for (doc, score) in documents:
+        if (doc.page_content not in seen_content) and (score >= score_min):
             seen_content.add(doc.page_content)
             unique_documents.append(doc)
     return unique_documents
 
-def doc_similarity(query, db, top_k):
-    docsearch = db.as_retriever(search_kwargs={'k':top_k})
-    docs = docsearch.get_relevant_documents(query)
-    udocs = remove_duplicates(docs)
+def doc_similarity(query, db, top_k, score):
+    docs = db.similarity_search_with_score(query = query,
+                                           k=top_k)
+    #docsearch = db.as_retriever(search_kwargs={'k':top_k})
+    #docs = docsearch.get_relevant_documents(query)
+    # print(docs)
+    udocs = remove_duplicates(docs, score)
     return udocs
 
 def user(user_message, history):
@@ -171,7 +214,7 @@ def user(user_message, history):
 
 def bot(box_message, ref_message,
         llm_dropdown, llm_dict, doc_list,
-        db, top_k):
+        db, top_k, score):
 
     # bot_message = random.choice(["Yes", "No"])
     # 0 is user question, 1 is bot response
@@ -184,9 +227,9 @@ def bot(box_message, ref_message,
 
     if not ref_message:
         ref_message = question
-        details = f"Q:  {question}"
+        details = f"Q: {question}"
     else:
-        details = f"Q:  {question}\nR: {ref_message}"
+        details = f"Q: {question}\nR: {ref_message}"
         
         
     llm = llm_dict[llm_dropdown]
@@ -196,27 +239,31 @@ def bot(box_message, ref_message,
             box_message[-1][1] = DOCS_WARNING
             return box_message, "", ""
         
-        chain = load_qa_chain(llm, chain_type="stuff")
-        docs = doc_similarity(ref_message, db, top_k)
+        docs = doc_similarity(ref_message, db, top_k, score)
         delta_top_k = top_k - len(docs)
 
         if delta_top_k > 0:
-            docs = doc_similarity(ref_message, db, top_k+delta_top_k)
-            
+            docs = doc_similarity(ref_message, db, top_k+delta_top_k, score)
+
+        prompt = PROMPT_DOC
+        #chain = load_qa_chain(llm, chain_type="stuff")
+   
     else:
-        chain = LLMChain(llm = llm,
-                         prompt = PromptTemplate(template='{question}',
-                                                input_variables=['question']),
-                         output_key = 'output_text')
+        prompt = PROMPT_BASE
         docs = []
-
-    all_output = chain({"input_documents": docs,
-                        "question": question,
-                        "chat_history": get_chat_history(history)})
+    
+    chain = LLMChain(llm = llm,
+                     prompt = prompt,
+                     output_key = 'output_text')
+
+    all_output = chain({"question": question,
+                        "context": docs,
+                        "chat_history": get_chat_history(history)
+                       })
+        
     
     bot_message = all_output['output_text']
 
-
     source = "".join([f"""<details> <summary>{doc.metadata["source"]}</summary>
 {doc.page_content}
 
@@ -288,12 +335,21 @@ with gr.Blocks(
                 
 
     with gr.Tab(TAB_2):
-        top_k = gr.Slider(1,
-                          TOP_K_MAX,
-                          value=TOP_K_DEFAULT,
-                          step=1,
-                          label="Vector similarity top_k",
-                          interactive=True)
+        with gr.Row():
+            with gr.Column():
+                top_k = gr.Slider(1,
+                                  TOP_K_MAX,
+                                  value=TOP_K_DEFAULT,
+                                  step=1,
+                                  label="Vector similarity top_k",
+                                  interactive=True)
+            with gr.Column():
+                score = gr.Slider(0.01,
+                                  0.99,
+                                  value=SCORE_DEFAULT,
+                                  step=0.01,
+                                  label="Vector similarity score",
+                                  interactive=True)
         detail_panel = gr.Chatbot(label="Related Docs")
     
     with gr.Tab(TAB_3):
@@ -349,7 +405,7 @@ with gr.Blocks(
         bot,
         [chatbot, ref,
          llm_dropdown, llm, doc_check,
-         vector_db, top_k],
+         vector_db, top_k, score],
         [chatbot, ref, detail_panel]
     )