Update app.py

app.py

@@ -1,5 +1,3 @@
-# inspiration from Ekimetrics climate qa
-
 import streamlit as st
 import os
 import re
@@ -13,13 +11,15 @@ import pandas as pd
 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
 
+# Get HF token
+hf_token = os.environ["HF_TOKEN"]
+login(token=hf_token, add_to_git_credential=True)
 
-# Enter openai API key
+# Get openai API key
 openai_key = os.environ["OPENAI_API_KEY"]
 
-# Select model
-model_name = "gpt-3.5-turbo"
 
 # Define the template
 template = PromptTemplate(
@@ -30,10 +30,12 @@ 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_name']+' ref. '+str(d.meta['ref_id'])+' &&&: '+d.content for d in documents])}; Question: {query}; Answer:""",
 )
 
+# Create a list of options for the dropdown
+model_options = ['chatGPT','Llama2']
+
 # Create a list of options for the dropdown
 country_options = ['Angola','Botswana','Lesotho','Kenya','Malawi','Mozambique','Namibia','South Africa','Zambia','Zimbabwe']
 
@@ -98,14 +100,33 @@ def get_refs(docs, res):
 
   return result_str
 
-def run_query(input_text, country):
+# 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."
+  # Add the meta data for references
+  context = ' - '.join(['&&& [ref. '+str(d.meta['ref_id'])+'] '+d.meta['document_name']+' &&&: '+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)
     # st.write('Selected country: ', country) # Debugging country
-    res = pipe.run(query=input_text, documents=docs)
-    output = res["results"][0]
+    if model_sel == "chatGPT":
+        res = pipe.run(query=input_text, documents=docs)
+        output = res["results"][0]
+        references = get_refs(docs, res["results"][0])
+    else:
+        res = client.text_generation(get_prompt_llama2(docs, query=input_text), max_new_tokens=4000, temperature=0.01, model=model)
+        output = res
+        references = get_refs(docs, res)
     st.write('Response')
     st.success(output)
-    references = get_refs(docs, res["results"][0])
     st.write('References')
     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)
@@ -121,19 +142,29 @@ retriever = EmbeddingRetriever(
     progress_bar=False,
 )
 
-# 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"])
-
 
 # Guiding text             
 st.title('Climate Policy Documents: Vulnerabilities Analysis Q&A (test)')
 st.markdown('This tool seeks to provide an interface for quering national climate policy documents (NDCs, LTS etc.). The current version is powered by chatGPT (3.5) and limited to 10 Southern African countries (Angola, Botswana, Eswatini, Lesotho, Malawi, Mozambique, Namibia, South Africa, Zambia, Zimbabwe) and Kenya. The intended use case is to allow users to interact with the documents and obtain valuable insights on various vulnerable groups affected by climate change.')
 st.markdown('**DISCLAIMER:** *This prototype tool based on LLMs (Language Models) is provided "as is" for experimental and exploratory purposes only, and should not be used for critical or production applications. Users are advised that the tool may contain errors, bugs, or limitations and should be used with caution and awareness of potential risks, and the developers make no warranties or guarantees regarding its performance, reliability, or suitability for any specific purpose.*')
     
+# Dropdown selectbox: model
+model_sel = st.selectbox('Select a model:', model_options)
+
+#----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 pipeline
+    pipe = Pipeline()
+    pipe.add_node(component=pn, name="prompt_node", inputs=["Query"])
+else:
+    model = "meta-llama/Llama-2-70b-chat-hf"
+    # Instantiate the inference client
+    client = InferenceClient()
+
 
 # Dropdown selectbox
 country = st.selectbox('Select a country:', country_options) 
@@ -148,10 +179,5 @@ else:
 
 
 if st.button('Submit'):
-    run_query(text, country=country)
-
-import sys
-import streamlit as st
+    run_query(text, country=country, model_sel=model_sel)
 
-python_version = sys.version
-st.text(f"Python Version: {python_version}")