Added filters

app.py

@@ -1,4 +1,5 @@
-import pinecone
+import json
+import pinecone
 import streamlit as st
 from transformers import pipeline, AutoTokenizer
 from sentence_transformers import SentenceTransformer
@@ -6,6 +7,8 @@ from sentence_transformers import SentenceTransformer
 PINECONE_KEY = st.secrets["PINECONE_API_KEY"]  # app.pinecone.io
 PINE_CONE_ENVIRONMENT = st.secrets["PINE_CONE_ENVIRONMENT"]  # app.pinecone.io
 
+st.set_page_config(layout="wide")
+
 @st.cache_resource
 def init_pinecone():
     pinecone.init(api_key=PINECONE_KEY, environment=PINE_CONE_ENVIRONMENT)  # get a free api key from app.pinecone.io
@@ -24,7 +27,7 @@ st.session_state.index = init_pinecone()
 retriever, reader, tokenizer = init_models()
 
 
-def card(name, description, score):
+def card(name, description, score, data_type, region, country):
     return st.markdown(f"""
     <div class="container-fluid">
         <div class="row align-items-start">
@@ -36,31 +39,57 @@ def card(name, description, score):
                      [<b>Score: </b>{score}]
                  </span>
              </div>
-             <div  class="col-md-4 col-sm-4">
-                    <small>{metadata}</metadata>
+             <div  class="col-md-1 col-sm-1">
+                    <small>{data_type}</metadata>
+             </div>        
+             <div  class="col-md-1 col-sm-1">
+                    <small>{region}</metadata>
+             </div>        
+             <div  class="col-md-1 col-sm-1">
+                    <small>{country}</metadata>
              </div>        
          </div>
      </div>
         """, unsafe_allow_html=True)
 
-st.title("")
 
-def run_query(query):
+def index_query(xq, top_k, regions=[], countries=[]):
+    #st.write(f"Regions: {regions}")
+    filters = []
+    if len(regions)>0:
+        filters.append({'region': {"$in": regions}})
+    if len(countries)>0:
+        filters.append({'country': {"$in": countries}})
+    if len(filters)==1:
+        filter = filters[0]
+    elif len(filters)>1:
+        filter = {"$and": filters} 
+    else:
+        filter = {}
+    #st.write(filter)
+    xc = st.session_state.index.query(xq, top_k=20, filter = filter, include_metadata=True, include_vectors = True)
+    #xc = st.session_state.index.query(xq, top_k=top_k, include_metadata=True, include_vectors = True)
+    return xc
+
+def run_query(query, scrape_boost, top_k , regions, countries):
     xq = retriever.encode([query]).tolist()
     try:
-        xc = st.session_state.index.query(xq, top_k=3, include_metadata=True, include_vectors = True)
+        xc = index_query(xq, top_k, regions, countries)
     except:
         # force reload
         pinecone.init(api_key=PINECONE_KEY, environment=PINE_CONE_ENVIRONMENT)
         st.session_state.index = pinecone.Index("company-description")
-        xc = st.session_state.index.query(xq, top_k=10, include_metadata=True, include_vectors = True)
+        xc = index_query(xq, top_k, regions, countries)
 
     results = []
     for match in xc['matches']:
         #answer = reader(question=query, context=match["metadata"]['context'])
-        answer = {'score': match['score']}
+        score = match['score']
+        if 'type' in match['metadata'] and match['metadata']['type']=='description-webcontent':
+            score = score * scrape_boost
+        answer = {'score': score}
         answer["name"] = match["metadata"]['company_name'].strip('_description')
-        answer["description"] = match["metadata"]['description']
+        answer["description"] = match["metadata"]['description'] if "description" in match['metadata'] else ""
         answer["metadata"] = match["metadata"]
         results.append(answer)
 
@@ -70,7 +99,10 @@ def run_query(query):
         company_name = r["name"]
         description = r["description"].replace(company_name, f"<mark>{company_name}</mark>")
         score = round(r["score"], 4)
-        card(company_name, description, score)
+        data_type = r["metadata"]["type"] if "type" in r["metadata"] else ""
+        region = r["metadata"]["region"]
+        country = r["metadata"]["country"]
+        card(company_name, description, score, data_type, region, country)
 
 def check_password():
     """Returns `True` if the user had the correct password."""
@@ -100,7 +132,9 @@ def check_password():
         # Password correct.
         return True
 
-if check_password():
+if True or check_password():
+    st.title("")
+
     st.write("""
     Search for a company in free text
     """)
@@ -108,9 +142,27 @@ if check_password():
     st.markdown("""
     <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/[email protected]/dist/css/bootstrap.min.css" integrity="sha384-Gn5384xqQ1aoWXA+058RXPxPg6fy4IWvTNh0E263XmFcJlSAwiGgFAW/dAiS6JXm" crossorigin="anonymous">
     """, unsafe_allow_html=True)
-    
+    with open("data/countries.json", "r") as f:
+        countries = json.load(f)['countries']
+    countries_selectbox = st.sidebar.multiselect("Country", countries, default=[])
+    all_regions = ('Africa', 'Europe', 'Asia & Pacific', 'North America', 'South/Latin America')
+    region_selectbox = st.sidebar.multiselect("Region", all_regions, default=all_regions)
+    scrape_boost = st.sidebar.number_input('webcontent_boost', value=2.)
+    top_k = st.sidebar.number_input('Top K Results', value=20)
+
+    # with st.container():
+        # col1, col2, col3, col4 = st.columns(4)
+        # with col1:
+            # scrape_boost = st.number_input('webcontent_boost', value=2.)
+        # with col2:
+            # top_k = st.number_input('Top K Results', value=20)
+        # with col3:
+            # regions = st.number_input('Region', value=20)
+        # with col4:
+            # countries = st.number_input('Country', value=20)
+
     query = st.text_input("Search!", "")
 
     if query != "":
-        run_query(query)
+        run_query(query, scrape_boost, top_k, region_selectbox, countries_selectbox)