Add results

pages/6_🔎_Find_Demo.py

@@ -1,40 +1,28 @@
 import os
 
+import pandas as pd
+
 import streamlit as st
 import streamlit_analytics
 from utils import add_logo_to_sidebar, add_footer, add_email_signup_form
 
+from huggingface_hub import snapshot_download
+
 from haystack.document_stores import InMemoryDocumentStore
-from haystack.nodes import BM25Retriever, EmbeddingRetreiver
+from haystack.nodes import BM25Retriever, EmbeddingRetriever
 
 HF_TOKEN = os.environ.get("HF_TOKEN")
 DATA_REPO_ID = "simplexico/cuad-qa-answers"
-DATA_FILENAME = "cuad_question_answers.json"
+DATA_FILENAME = "cuad_questions_answers.json"
 EMBEDDING_MODEL = "prajjwal1/bert-tiny"
 if EMBEDDING_MODEL == "prajjwal1/bert-tiny":
     EMBEDDING_DIM = 128
 else:
     EMBEDDING_DIM = 768
 
-streamlit_analytics.start_tracking()
-
-st.set_page_config(
-    page_title="Find Demo",
-    page_icon="🔎",
-    layout="wide",
-    initial_sidebar_state="expanded",
-    menu_items={
-        'Get Help': 'mailto:[email protected]',
-        'Report a bug': None,
-        'About': "## This a demo showcasing different Legal AI Actions"
-    }
-)
-
-add_logo_to_sidebar()
-st.sidebar.success("👆 Select a demo above.")
+EXAMPLE_TEXT = "the governing law is the State of Texas"
 
-st.title('🔎 Find Demo')
-st.markdown("🏗 This demo is currently under construction. Please visit back soon.")
+streamlit_analytics.start_tracking()
 
 @st.cache(allow_output_mutation=True)
 def load_dataset():
@@ -43,7 +31,7 @@ def load_dataset():
     return df
 
 @st.cache(allow_output_mutation=True)
-def generate_document_store(df):
+def generate_document_store(df, dummy=None):
     """Create haystack document store using contract clause data 
     """
     document_dicts = []
@@ -68,10 +56,53 @@ def generate_bm25_retriever(document_store):
 
 @st.cache(allow_output_mutation=True)
 def generate_embeddings(embedding_model, document_store):
-    embedding_retriever = EmbeddingRetreiver(embedding_model=embedding_model, document_store=document_store)
+    embedding_retriever = EmbeddingRetriever(embedding_model=embedding_model, document_store=document_store)
     document_store.update_embeddings(embedding_retriever)
     return embedding_retriever
 
+def process_query(query, retriever):
+    """Generates dataframe with top ten results"""
+    texts = []
+    contract_titles = []
+    candidate_documents = retriever.retrieve(
+        query=query,
+        top_k=10,
+    )
+        
+    for document in candidate_documents:
+        texts.append(document.content)
+        contract_titles.append(document.meta["contract_title"])
+    
+    return pd.DataFrame({"Text": texts, "Source Contract": contract_titles})
+
+st.set_page_config(
+    page_title="Find Demo",
+    page_icon="🔎",
+    layout="wide",
+    initial_sidebar_state="expanded",
+    menu_items={
+        'Get Help': 'mailto:[email protected]',
+        'Report a bug': None,
+        'About': "## This a demo showcasing different Legal AI Actions"
+    }
+)
+
+add_logo_to_sidebar()
+st.sidebar.success("👆 Select a demo above.")
+
+st.title('🔎 Find Demo')
+
+st.write("""
+This demo shows how a set of documents can be searched.
+We've set up a database of clauses from a set of open source legal documents.
+These clauses can be searched using **keywords** or using **semantic search**.
+Semantic search leverages an AI model which matches on clauses with a similar meaning to the input text. 
+""")
+st.write("**👈 Enter search query on the left** and hit the button **Find Clauses** to see the demo in action")
+
+query = st.sidebar.text_area(label='Enter Searcb Query', value=EXAMPLE_TEXT, height=250)
+button = st.sidebar.button('**Find Clauses**', type='primary', use_container_width=True)
+
 df = load_dataset()
 
 document_store = generate_document_store(df)
@@ -80,8 +111,32 @@ bm25_retriever = generate_bm25_retriever(document_store)
 
 embedding_retriever = generate_embeddings(EMBEDDING_MODEL, document_store)
 
+if button:
+
+    hide_dataframe_row_index = """
+        <style>
+        .row_heading.level0 {display:none}
+        .blank {display:none}
+        </style>
+        """
+
+    col1, col2 = st.columns(2)
+
+    with col1:
 
+        st.subheader('Keyword Search Results:')
+        # Inject CSS with Markdown
+        st.markdown(hide_dataframe_row_index, unsafe_allow_html=True)
+        df_bm25 = process_query(query, bm25_retriever)
+        st.table(df_bm25)
+        
+    with col2:
 
+        st.subheader('Semantic Search Results:')
+        # Inject CSS with Markdown
+        st.markdown(hide_dataframe_row_index, unsafe_allow_html=True)
+        df_embed = process_query(query, embedding_retriever)
+        st.table(df_embed)
 
 add_email_signup_form()