Update app.py

app.py

@@ -6,22 +6,18 @@ from langchain_community.embeddings import HuggingFaceBgeEmbeddings
 import os
 HUGGINGFACEHUB_API_TOKEN = os.environ["token"]
 
-def clean_(l):
-    s = list(l)[0][1]
-    s = s.replace("\n", "=")
+def clean_(s):
+    s = s.replace("\n0: ", "=")
     return re.split('=', s, maxsplit=1)[-1].strip()
 
-def similarity_search2(vectordb, query, k, unique="True"):
+def similarity_search2(vectordb, query, k=1, unique="True"):
     print(f"\nQuery Key: {query}, \nrows requested:{k}\nUnique values:{unique}")
-    D = vectordb.similarity_search(query,k)
-    temp = []
-    for d in D:
-        temp.append(clean_(d))
-    del D
-    if unique == "True":
-        return str(np.unique(np.array(temp)))[1:-1]
+    if unique == "False":
+        vals = vectordb.similarity_search(query,k=k)
     else:
-        return str(np.array(temp))[1:-1]
+        vals = vectordb.similarity_search(query,k=1)
+    for val in vals:
+        return clean_(val.page_content)
     
 with gr.Blocks() as demo:
     gr.Markdown(
@@ -31,7 +27,7 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             query = gr.Textbox(placeholder="your query", label="Query")
-            k = gr.Slider(10,100000,5, label="number of samples to check")
+            k = gr.Slider(1,306,1, label="number of samples to check")
             unique = gr.Radio(["True", "False"], label="Return Unique values")
             with gr.Row():
                 btn = gr.Button("Submit")
@@ -41,14 +37,16 @@ with gr.Blocks() as demo:
                 embedding = HuggingFaceBgeEmbeddings(
                     model_name = model_id,
                     model_kwargs = model_kwargs,
-                    encode_kwargs = {'normalize_embeddings':True}
+                    cache_folder=r"models",
+                    encode_kwargs = {'normalize_embeddings':True},
                 )
-                persist_directory = "db_book_mmt"
+                persist_directory = "MMT_unique"
                 vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding)
                 return similarity_search2(vectordb, query, k, unique)
         with gr.Column():
             output = gr.Textbox(scale=10, label="Output")
             btn.click(mmt_query, [query, k, unique], output)
-            
+    
+    # interface = gr.Interface(fn=auto_eda, inputs="dataframe", outputs="json")
 # demo.queue()
-demo.launch()
+demo.launch(share=True)