feat: UI polish again, improve search in doc features & document list is now more explicit

app.py

@@ -1,20 +1,49 @@
 import time
 import uuid
+import openai
+import os
 import pandas as pd
+import numpy as np
 import gradio as gr
-from llama_index import GPTSimpleVectorIndex, MockLLMPredictor, ServiceContext
+from llama_index import GPTSimpleVectorIndex
+from gpt_index.indices.struct_store.pandas import GPTPandasIndex
+from openai.embeddings_utils import get_embedding, cosine_similarity
+
+openai.api_key = os.getenv("OPENAI_API_KEY")
 
 title = "Confidential forensics tool with ChatGPT"
 examples = ["Who is Phillip Allen?", "What the project in Austin is about?", "Give me more details about the real estate project"]
 
-llm_predictor = MockLLMPredictor()
-service_context_mock = ServiceContext.from_defaults(llm_predictor=llm_predictor)
-
 index = GPTSimpleVectorIndex.load_from_disk('email.json')
-docs_arr = []
-for doc in index.docstore.docs:
-    docs_arr.append(doc)
-dat_fr = pd.DataFrame({"Documents loaded": docs_arr})
+
+dat_fr = pd.DataFrame({"Documents loaded": ["email.json", "email_200.parquet"]})
+
+df = pd.read_parquet("email_200.parquet")
+# df["embedding"] = [get_embedding(x) forW x in df.body.values]
+# df.to_parquet("email_50.parquet")
+
+# df = pd.read_csv("email_ok.csv", nrows=50)
+# df["embedding"] = [get_embedding(x) for x in df.body.values]
+# df.to_parquet("email_50.parquet")
+
+def search_emails(opt, message, n=3):
+    "Outputs the top n emails that match the most the pattern"
+    if len(message.strip()) < 1:
+        message = "Oops, it looks like your query was not valid. Please make sure you typed something in your text box and then try again."
+    else:
+        try:
+            embedding = get_embedding(message)
+            message = ""
+            df['similarities'] = df.embedding.apply(lambda x: cosine_similarity(x, embedding))
+
+            message_tmp = df.sort_values('similarities', ascending=False).head(n)
+            message_tmp = [(row.file, row.body, row.similarities) for index, row in message_tmp.iterrows()]
+            for msg in message_tmp:
+                message += f"{msg[0]}\nContent: {msg[1].strip()}\n{msg[2]}\n\n"
+        except Exception as e:
+            message = "An error occured when handling your query, please try again."
+            print(e)
+    return message, ""
 
 def respond_upload(btn_upload, message, chat_history):
     time.sleep(2)
@@ -47,48 +76,32 @@ def respond(message, chat_history):
     chat_history.append((message, bot_message))
     return "", chat_history
 
-def find_doc(opt, msg2):
-    message = ""
-    if len(msg2.strip()) < 1:
-        message = "Oops, it looks like your query was not valid. Please make sure you typed something in your text box and then try again."
-    else:
-        try:
-            resp = index.query(msg2, service_context=service_context_mock)
-            for key, item in resp.extra_info.items():
-                message += f"Document: {key}\nExtra details:\n"
-                for sub_key, sub_item in item.items():
-                    message += f"---- {sub_key}: {sub_item}"
-
-        except Exception as e:
-            message = "An error occured when handling your query, please try again."
-            print(e)
-    return message, ""
-
 with gr.Blocks(title=title) as demo:
     gr.Markdown(
     """
 
     # """ + title + """
-    ...
     """)
     dat = gr.Dataframe(
-            value=dat_fr
+            value=dat_fr,
+            max_cols=1,
+            max_rows=4,
+            overflow_row_behaviour="paginate",
     )
+    btn_upload = gr.UploadButton("Upload a new document...", file_types=["text"])
     gr.Markdown(
     """
     ## Chatbot
     """)
     chatbot = gr.Chatbot().style(height=400)
     with gr.Row():
-        with gr.Column(scale=0.70):
+        with gr.Column(scale=0.85):
             msg = gr.Textbox(
                 show_label=False,
                 placeholder="Enter text and press enter, or click on Send.",
             ).style(container=False)
         with gr.Column(scale=0.15, min_width=0):
             btn_send = gr.Button("Send your query")
-        with gr.Column(scale=0.15, min_width=0):
-            btn_upload = gr.UploadButton("Upload a new document...", file_types=["text"])
     with gr.Row():
         gr.Markdown(
         """
@@ -121,7 +134,7 @@ with gr.Blocks(title=title) as demo:
         with gr.Column(scale=0.15, min_width=0):
             btn_send2 = gr.Button("Send your query")
 
-    btn_send2.click(find_doc, [opt, msg2], [opt, msg2])
+    btn_send2.click(search_emails, [opt, msg2], [opt, msg2])
     
 if __name__ == "__main__":
     demo.launch()

email_200.parquet

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:689f185923970ee7a8da61eb63ae2116d75af0a473dd2735c0848bc87b505135
+size 8670196

requirements.txt

@@ -1,3 +1,10 @@
 gradio
 llama_index
-pandas
+pandas
+gpt_index
+numpy
+openai
+plotly
+scipy
+scikit-learn
+fastparquet