Update greg_funcs.py

change data acquisition to use vector db
show chain of thought in the chat output

greg_funcs.py

@@ -1,6 +1,6 @@
 from sentence_transformers import SentenceTransformer, CrossEncoder, util
-from torch import tensor as torch_tensor
-from datasets import load_dataset
+#from torch import tensor as torch_tensor
+#from datasets import load_dataset
 
 from langchain.llms import OpenAI
 from langchain.docstore.document import Document
@@ -13,6 +13,7 @@ from langchain.agents import initialize_agent, Tool
 import sqlite3
 #import pandas as pd
 import json
+import chromadb
 
 # database
 cxn = sqlite3.connect('./data/mbr.db')
@@ -29,81 +30,82 @@ cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
 
 """# import datasets"""
 
-dataset = load_dataset("gfhayworth/hack_policy", split='train')
-mypassages = list(dataset.to_pandas()['psg'])
-
-dataset_embed = load_dataset("gfhayworth/hack_policy_embed", split='train')
-dataset_embed_pd = dataset_embed.to_pandas()
-mycorpus_embeddings = torch_tensor(dataset_embed_pd.values)
-
-def search(query, passages = mypassages, doc_embedding = mycorpus_embeddings, top_k=20, top_n = 1):
-    question_embedding = bi_encoder.encode(query, convert_to_tensor=True)
-    question_embedding = question_embedding #.cuda()
-    hits = util.semantic_search(question_embedding, doc_embedding, top_k=top_k)
-    hits = hits[0]  # Get the hits for the first query
-
-    ##### Re-Ranking #####
-    cross_inp = [[query, passages[hit['corpus_id']]] for hit in hits]
-    cross_scores = cross_encoder.predict(cross_inp)
-
-    # Sort results by the cross-encoder scores
-    for idx in range(len(cross_scores)):
-        hits[idx]['cross-score'] = cross_scores[idx]
-
-    hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True)
-    predictions = hits[:top_n]
-    return predictions
-    # for hit in hits[0:3]:
-    #         print("\t{:.3f}\t{}".format(hit['cross-score'], mypassages[hit['corpus_id']].replace("\n", " ")))
-
-
-
-def get_text_fmt(qry, passages = mypassages, doc_embedding=mycorpus_embeddings):
-    predictions = search(qry, passages = passages, doc_embedding = doc_embedding, top_n=5, )
-    prediction_text = []
-    for hit in predictions:
-        page_content = passages[hit['corpus_id']]
-        metadata = {"source": hit['corpus_id']}
-        result = Document(page_content=page_content, metadata=metadata)
-        prediction_text.append(result)
-    return prediction_text
+# dataset = load_dataset("gfhayworth/hack_policy", split='train')
+# mypassages = list(dataset.to_pandas()['psg'])
+
+# dataset_embed = load_dataset("gfhayworth/hack_policy_embed", split='train')
+# dataset_embed_pd = dataset_embed.to_pandas()
+# mycorpus_embeddings = torch_tensor(dataset_embed_pd.values)
+###########################################################################################################################
+"""# set up vector db"""
+from chromadb.config import Settings
+
+chroma_client = chromadb.Client(settings=Settings(
+    chroma_db_impl="duckdb+parquet",
+    persist_directory="./data/mychromadb/" # Optional, defaults to .chromadb/ in the current directory
+))
+collection = chroma_client.get_collection(name="benefit_collection")
+
+def vdb_rslt(qry,src,top_k=20):
+  results = collection.query(
+    query_embeddings=[ bi_encoder.encode(qry) ],
+    n_results=top_k,
+    where={"source": src},
+    )
+  return results
+##################################################################################################################################
+# Semantic Search Functions
+def rtrv(qry, src = 'H1036236000SB23.pdf', top_k=20):
+  rslts = vdb_rslt(qry,src) 
+  return rslts
+
+def rernk(query, collection=collection, top_k=20, top_n = 5):
+  rtrv_rslts = rtrv(query, top_k=top_k)
+  rtrv_ids = rtrv_rslts.get('ids')[0]
+  rtrv_docs = rtrv_rslts.get('documents')[0]
+
+  ##### Re-Ranking #####
+  cross_inp = [[query, doc] for doc in rtrv_docs]
+  cross_scores = cross_encoder.predict(cross_inp)
+
+  # Sort results by the cross-encoder scores
+  combined = list(zip(rtrv_ids, list(cross_scores)))
+  sorted_tuples = sorted(combined, key=lambda x: x[1], reverse=True)
+  sorted_ids = [t[0] for t in sorted_tuples[:top_n]]
+  predictions = collection.get(ids=sorted_ids, include=["documents","metadatas"])
+  return predictions
+
+def get_text_fmt(qry):
+  prediction_text = []
+  predictions = rernk(qry, collection=collection, top_k=20, top_n = 5)
+  docs = predictions['documents']
+  meta = predictions['metadatas']
+  for i in range(len(docs)):
+    result = Document(page_content=docs[i], metadata=meta[i]) 
+    prediction_text.append(result)
+  return prediction_text
 
+##################################################################################################################################
 """# LLM based qa functions"""
 
-template = """You are a friendly AI assistant for the insurance company Humana. Given the following extracted parts of a long document and a question, create a succinct final answer. 
+template = """You are a friendly AI assistant for the insurance company Humana. 
+Given the following extracted parts of a long document and a question, create a succinct final answer. 
 If you don't know the answer, just say that you don't know. Don't try to make up an answer.
-If the question is not about Humana, politely inform the user that you are tuned to only answer questions about Humana benefits.
+If the question is not about Humana, politely inform the user that you are tuned to only answer questions about Humana.
 QUESTION: {question}
 =========
-{context}
+{summaries}
 =========
 FINAL ANSWER:"""
-PROMPT = PromptTemplate(template=template, input_variables=["context", "question"])
+PROMPT = PromptTemplate(template=template, input_variables=["summaries", "question"])
 
-chain_qa = load_qa_chain(OpenAI(temperature=0), chain_type="stuff", prompt=PROMPT)
-
-def get_text_fmt(qry, passages = mypassages, doc_embedding=mycorpus_embeddings):
-  predictions = search(qry, passages = passages, doc_embedding = doc_embedding, top_n=5, )
-  prediction_text = []
-  for hit in predictions:
-    page_content = passages[hit['corpus_id']]
-    metadata = {"source": hit['corpus_id']}
-    result = Document(page_content=page_content, metadata=metadata)
-    prediction_text.append(result)
-  return prediction_text
+chain_qa = load_qa_with_sources_chain(OpenAI(temperature=0), chain_type="stuff", prompt=PROMPT, verbose=True)
 
 def get_llm_response(message):
   mydocs = get_text_fmt(message)
-  responses = chain_qa.run(input_documents=mydocs, question=message)
+  responses = chain_qa({"input_documents":mydocs, "question":message})
   return responses
 
-# for x in xmpl_list:
-#   print(32*'=')
-#   print(x)
-#   print(32*'=')
-#   r = get_llm_response(x)
-#   print(r)
-
 """# Database query"""
 
 db = SQLDatabase.from_uri("sqlite:///./data/mbr.db")
@@ -113,10 +115,10 @@ llm = OpenAI(temperature=0)
 # model_name: str = "text-davinci-003"
 # instruction fine-tuned, sometimes referred to as GPT-3.5
 
-db_chain = SQLDatabaseChain(llm=llm, database=db, verbose=True)
+db_chain = SQLDatabaseChain(llm=llm, database=db, verbose=True, return_intermediate_steps=True)
 
 def db_qry(qry):
-  responses = db_chain.run(query='my mbr_id is 456 ;'+str(qry) ) ############### hardcode mbr id 456 for demo
+  responses = db_chain('my mbr_id is 456 ;'+str(qry) ) ############### hardcode mbr id 456 for demo
   return responses
 
 #db_qry('how many footcare visits have I had?')
@@ -178,13 +180,38 @@ def mrkl_rspnd(qry):
   response = mrkl({"input":str(qry) })
   return response
 
+def get_cot(r):
+  cot = '<p>'
+  try:
+    intermedObj = r['intermediate_steps']
+    cot +='<b>Input:</b> '+r['input']+'<br>'
+    for agnt_action, obs in intermedObj:
+      al = '<br>  '.join(agnt_action.log.split('\n') )
+      cot += '<b>AI chain of thought:</b> '+ al +'<br>'
+      if type(obs) is dict:
+        if obs.get('input_documents') is not None: 
+          for d in obs['input_documents']:
+            cot += '&nbsp;&nbsp;&nbsp;&nbsp;'+'<i>- '+str(d.page_content)+'</i>'+' <a href="'+ str(d.metadata['url']) +'">'+str(d.metadata['page'])+'</a> '+'<br>'
+          cot += '<b>Observation:</b> '+str(obs['output_text']) +'<br><br>'
+        elif obs.get('intermediate_steps') is not None:
+          cot += '<b>Query:</b> '+str(obs.get('intermediate_steps')) +'<br><br>'
+        else:  
+          pass 
+      else:
+        cot += '<b>Observation:</b> '+str(obs) +'<br><br>'
+  except:
+    pass
+  cot += '</p>'  
+  return cot
+
 def chat(message, history):
     history = history or []
     message = message.lower()
     
     response = mrkl_rspnd(message)
+    cot = get_cot(response)
     history.append((message, response['output']))
-    return history, history
+    return history, history, cot