Update app.py

app.py

@@ -12,6 +12,7 @@ from haystack.document_stores import FAISSDocumentStore
 from haystack.nodes import EmbeddingRetriever
 from haystack.schema import Document
 from huggingface_hub import login, HfApi, hf_hub_download, InferenceClient
+import openai
 
 # Get HF token
 hf_token = os.environ["HF_TOKEN"]
@@ -19,18 +20,30 @@ login(token=hf_token, add_to_git_credential=True)
 
 # Get openai API key
 openai_key = os.environ["OPENAI_API_KEY"]
+openai.api_key = openai_api_key
+
+# Get openai API key
+pinecone_key = os.environ["PINECONE_API_KEY"]
 
 
-template = PromptTemplate(
-    prompt="""
-Answer the given question using the following documents. \
+# template = PromptTemplate(
+#     prompt="""
+# Answer the given question using the following documents. \
+# Formulate your answer in the style of an academic report. \
+# Provide example quotes and citations using extracted text from the documents. \
+# Use facts and numbers from the documents in your answer. \
+# Reference information used from documents at the end of each applicable sentence (ex: [source: document_name]), where 'document_name' is the text provided at the start of each document (demarcated by '- &&&' and '&&&:')'. \
+# If no relevant information to answer the question is present in the documents, just say you don't have enough information to answer. \
+# Context: {' - '.join(['&&& '+d.meta['document']+' ref. '+str(d.meta['ref_id'])+' &&&: '+d.content for d in documents])}; Question: {query}; Answer:""",
+# )
+
+prompt_template="Answer the given question using the following documents. \
 Formulate your answer in the style of an academic report. \
 Provide example quotes and citations using extracted text from the documents. \
 Use facts and numbers from the documents in your answer. \
-Reference information used from documents at the end of each applicable sentence (ex: [source: document_name]), where 'document_name' is the text provided at the start of each document (demarcated by '- &&&' and '&&&:')'. \
-If no relevant information to answer the question is present in the documents, just say you don't have enough information to answer. \
-Context: {' - '.join(['&&& '+d.meta['document']+' ref. '+str(d.meta['ref_id'])+' &&&: '+d.content for d in documents])}; Question: {query}; Answer:""",
-)
+ALWAYS include references for information used from documents at the end of each applicable sentence using the format: '[ref. #]', where '[ref. #]' is included in the text provided at the start of each document (demarcated by the pattern '- &&& [ref. #] document_name &&&:')'. \
+Do not include page numbers in the references. \
+If no relevant information to answer the question is present in the documents, just say you don't have enough information to answer."
 
 # Create a list of options for the dropdown
 model_options = ['chatGPT','Llama2']
@@ -39,7 +52,7 @@ model_options = ['chatGPT','Llama2']
 country_options = ['All Countries','Angola','Botswana','Lesotho','Kenya','Malawi','Mozambique','Namibia','Rwanda','South Africa','Zambia','Zimbabwe']
 
 # Create a list of options for the dropdown
-class_options = ['All Categories','Agricultural communities', 'Children', 'Coastal communities', 'Ethnic, racial or other minorities', 'Fishery communities', 'Informal sector workers', 'Members of indigenous and local communities', 'Migrants and displaced persons', 'Older persons', 'Persons living in poverty', 'Persons with disabilities', 'Persons with pre-existing health conditions', 'Residents of drought-prone regions', 'Rural populations', 'Sexual minorities (LGBTQI+)', 'Urban populations', 'Women and other genders','Other']
+vulnerability_options = ['All Categories','Agricultural communities', 'Children', 'Coastal communities', 'Ethnic, racial or other minorities', 'Fishery communities', 'Informal sector workers', 'Members of indigenous and local communities', 'Migrants and displaced persons', 'Older persons', 'Persons living in poverty', 'Persons with disabilities', 'Persons with pre-existing health conditions', 'Residents of drought-prone regions', 'Rural populations', 'Sexual minorities (LGBTQI+)', 'Urban populations', 'Women and other genders','Other']
 
 # List of examples
 examples = [
@@ -48,40 +61,74 @@ examples = [
     "In addition to gender, children, and youth, is there any mention of other groups facing disproportional impacts from climate change due to their geographic location, socio-economic status, age, gender, health, and occupation?"
 ]
 
-def get_docs(input_query, country = None):
-    '''
-    Construct a hacky query to focus the retriever on the target country (see notes below)
-    We take the top 150 k because we want to make sure we have 10 pertaining to the selected country
-    '''
-    if country == 'All Countries':
-        query = input_query
+# def get_docs(input_query, country = None):
+#     '''
+#     Construct a hacky query to focus the retriever on the target country (see notes below)
+#     We take the top 150 k because we want to make sure we have 10 pertaining to the selected country
+#     '''
+#     if country == 'All Countries':
+#         query = input_query
+#     else:
+#         query = "For the country of "+country+", "+input_query
+#     # Retrieve top k documents
+#     docs = retriever.retrieve(query=query,top_k = 150)
+#     # Break out the key fields and convert to pandas for filtering
+#     docs = [{**x.meta,"score":x.score,"content":x.content} for x in docs]
+#     df_docs = pd.DataFrame(docs)
+#     if country != 'All Countries':
+#       df_docs = df_docs.query('country in @country')
+#     # Take the top 10
+#     df_docs = df_docs.head(10)
+#     # Get ourselves an index setup from which to base the source reference number from (in the prompt and matching afterwards)
+#     df_docs = df_docs.reset_index()
+#     df_docs['ref_id'] = df_docs.index + 1 # start the index at 1
+#     # Convert back to Document format
+#     ls_dict = []
+#     # Iterate over df and add relevant fields to the dict object
+#     for index, row in df_docs.iterrows():
+#         # Create a Document object for each row
+#         doc = Document(
+#             row['content'],
+#             meta={'country': row['country'],'document': row['document'], 'page': row['page'], 'file_name': row['file_name'], 'ref_id': row['ref_id'], 'score': row['score']}
+#         )
+
+#         # Append the Document object to the documents list
+#         ls_dict.append(doc)
+#     return(ls_dict)
+
+def get_docs(input_query, country = [], vulnerability_cat = []):
+  if not country:
+    country = "All Countries"
+  if not vulnerability_cat:
+    if country == "All Countries":
+      filters = None
     else:
-        query = "For the country of "+country+", "+input_query
-    # Retrieve top k documents
-    docs = retriever.retrieve(query=query,top_k = 150)
-    # Break out the key fields and convert to pandas for filtering
-    docs = [{**x.meta,"score":x.score,"content":x.content} for x in docs]
-    df_docs = pd.DataFrame(docs)
-    if country != 'All Countries':
-      df_docs = df_docs.query('country in @country')
-    # Take the top 10
-    df_docs = df_docs.head(10)
-    # Get ourselves an index setup from which to base the source reference number from (in the prompt and matching afterwards)
-    df_docs = df_docs.reset_index()
-    df_docs['ref_id'] = df_docs.index + 1 # start the index at 1
-    # Convert back to Document format
-    ls_dict = []
-    # Iterate over df and add relevant fields to the dict object
-    for index, row in df_docs.iterrows():
-        # Create a Document object for each row
-        doc = Document(
-            row['content'],
-            meta={'country': row['country'],'document': row['document'], 'page': row['page'], 'file_name': row['file_name'], 'ref_id': row['ref_id'], 'score': row['score']}
-        )
-
-        # Append the Document object to the documents list
-        ls_dict.append(doc)
-    return(ls_dict)
+      filters = {'country': {'$in': country}}
+  else:
+    if country == "All Countries":
+          filters = {'vulnerability_cat': {'$in': vulnerability_cat}}
+    else:
+      filters = {'country': {'$in': country},'vulnerability_cat': {'$in': vulnerability_cat}}
+  docs = retriever.retrieve(query=query, filters = filters, top_k = 10)
+  # Break out the key fields and convert to pandas for filtering
+  docs = [{**x.meta,"score":x.score,"content":x.content} for x in docs]
+  df_docs = pd.DataFrame(docs)
+  # Get ourselves an index setup from which to base the source reference number from (in the prompt and matching afterwards)
+  df_docs = df_docs.reset_index()
+  df_docs['ref_id'] = df_docs.index + 1 # start the index at 1
+  # Convert back to Document format
+  ls_dict = []
+  # Iterate over df and add relevant fields to the dict object
+  for index, row in df_docs.iterrows():
+      # Create a Document object for each row
+      doc = Document(
+          row['content'],
+          meta={'country': row['country'],'document': row['document'], 'page': row['page'], 'file_name': row['file_name'], 'ref_id': row['ref_id'], 'vulnerability_cat': row['vulnerability_cat'], 'score': row['score']}
+      )
+
+      # Append the Document object to the documents list
+      ls_dict.append(doc)
+  return ls_dict
 
 def get_refs(docs, res):
   '''
@@ -107,26 +154,22 @@ def get_refs(docs, res):
   return result_str
 
 # define a special function for putting the prompt together (as we can't use haystack)
-def get_prompt_llama2(docs, query):
-  base_prompt="Answer the given question using the following documents. \
-  Formulate your answer in the style of an academic report. \
-  Provide example quotes and citations using extracted text from the documents. \
-  Use facts and numbers from the documents in your answer. \
-  ALWAYS include references for information used from documents at the end of each applicable sentence using the format: '[ref. #]', where '[ref. #]' is included in the text provided at the start of each document (demarcated by the pattern '- &&& [ref. #] document_name &&&:')'. \
-  Do not include page numbers in the references. \
-  If no relevant information to answer the question is present in the documents, just say you don't have enough information to answer."
+def get_prompt(docs, query):
+  base_prompt=prompt_template
   # Add the meta data for references
-  context = ' - '.join(['&&& [ref. '+str(d.meta['ref_id'])+'] '+d.meta['document_name']+' &&&: '+d.content for d in docs])
+  context = ' - '.join(['&&& [ref. '+str(d.meta['ref_id'])+'] '+d.meta['document']+' &&&: '+d.content for d in docs])
   prompt = base_prompt+"; Context: "+context+"; Question: "+query+"; Answer:"
   return(prompt)
 
 def run_query(input_text, country, model_sel):
-    docs = get_docs(input_text, country)
+    # docs = get_docs(input_text, country)
+    docs = get_docs(query, country=country,vulnerability_cat=vulnerability_options)
     # st.write('Selected country: ', country) # Debugging country
     if model_sel == "chatGPT":
-        res = pipe.run(query=input_text, documents=docs)
+        # res = pipe.run(query=input_text, documents=docs)
+        res = openai.ChatCompletion.create(model="gpt-3.5-turbo", messages=[{"role": "user", "content": get_prompt(docs, query=query)}])
         output = res["results"][0]
-        references = get_refs(docs, res["results"][0])
+        references = get_refs(docs, output)
     else:
         res = client.text_generation(get_prompt_llama2(docs, query=input_text), max_new_tokens=4000, temperature=0.01, model=model)
         output = res
@@ -137,17 +180,30 @@ def run_query(input_text, country, model_sel):
     st.markdown('References are based on text automatically extracted from climate policy documents. These extracts may contain non-legible characters or disjointed text as an artifact of the extraction procedure')
     st.markdown(references, unsafe_allow_html=True)
 
-# Setup retriever, pulling from local faiss datastore
+# # Setup retriever, pulling from local faiss datastore
+# retriever = EmbeddingRetriever(
+#     document_store=FAISSDocumentStore.load(
+#         index_path="./cpv_full_southern_africa_kenya.faiss",
+#         config_path="./cpv_full_southern_africa_kenya.json",
+#     ),
+#     embedding_model="sentence-transformers/multi-qa-mpnet-base-dot-v1",
+#     model_format="sentence_transformers",
+#     progress_bar=False,
+# )
+
+# Setup retriever, pulling from pinecone
+doc_file_name="cpv_full_southern_africa"
+document_store = PineconeDocumentStore(api_key=pinecone_key,
+                                       environment="asia-southeast1-gcp-free",
+                                       index=doc_file_name)
+
 retriever = EmbeddingRetriever(
-    document_store=FAISSDocumentStore.load(
-        index_path="./cpv_full_southern_africa_kenya.faiss",
-        config_path="./cpv_full_southern_africa_kenya.json",
-    ),
-    embedding_model="sentence-transformers/multi-qa-mpnet-base-dot-v1",
-    model_format="sentence_transformers",
-    progress_bar=False,
+  document_store=document_store,
+  embedding_model="sentence-transformers/multi-qa-mpnet-base-dot-v1",
+  model_format="sentence_transformers"
 )
 
+
 with st.sidebar:
     # Dropdown selectbox
     country = st.sidebar.selectbox('Select a country:', country_options)
@@ -157,7 +213,7 @@ with st.sidebar:
         * *Then be sure to mention the country names of interest in your query*
         """
     )
-    vulnerabilities_cat = st.sidebar.selectbox('Select a vulnerabilities category:', class_options)
+    vulnerabilities_cat = st.sidebar.selectbox('Select a vulnerabilities category:', vulnerability_options)
     # choice = st.sidebar.radio(label = 'Select the Document',
     #                         help = 'You can upload the document \
     #                         or else you can try a example document', 
@@ -189,12 +245,12 @@ model_sel = "chatGPT"
 #----Model Select logic-------
 if model_sel == "chatGPT":
     model_name = "gpt-3.5-turbo"
-    # Initialize the PromptNode
-    pn = PromptNode(model_name_or_path=model_name, default_prompt_template=template, api_key=openai_key, max_length=2000, model_kwargs={"generation_kwargs": {"do_sample": False, "temperature": 0}})
+    # # Initialize the PromptNode
+    # pn = PromptNode(model_name_or_path=model_name, default_prompt_template=template, api_key=openai_key, max_length=2000, model_kwargs={"generation_kwargs": {"do_sample": False, "temperature": 0}})
     
-    # Initialize the pipeline
-    pipe = Pipeline()
-    pipe.add_node(component=pn, name="prompt_node", inputs=["Query"])
+    # # Initialize the pipeline
+    # pipe = Pipeline()
+    # pipe.add_node(component=pn, name="prompt_node", inputs=["Query"])
 else:
     # Currently disabled
     model = "meta-llama/Llama-2-70b-chat-hf"