Update app.py

app.py

@@ -1,62 +1,30 @@
 import gradio as gr
-import numpy as np
-import json
-import pickle as pkl
-from transformers import AutoTokenizer
-import re
-# Vector Loader
-vectors = pkl.load(open("vectors.pkl", "rb"))
-vocab = [word.lower() for word in vectors.keys()]
+from vectordb import Memory
 
-# Tokenizer
-tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
-def make_alphanumeric(input_string):
-    return re.sub(r'[^a-zA-Z0-9 ]', '', input_string)
+# Initialize Memory
+memory = Memory()
 
-def tokenize(text):
-    # Check data
-    if len(text) == 0:
-        gr.Error("No text provided.")
-    elif len(text) > 4096:
-        gr.Error("Text too long.")
-    
-    # Filter
-    text = make_alphanumeric(text.lower())
-    pre_tokenize_result = tokenizer._tokenizer.pre_tokenizer.pre_tokenize_str(text)
-    pre_tokenized_text = [word for word, offset in pre_tokenize_result]
-    
-    tokens = []
-    for word in pre_tokenized_text:
-        if word in vocab:
-            tokens.append(word)
-    return tokens
+# Save some example data
+memory.save(
+    ["apples are green", "oranges are orange"],  # save your text content
+    [{"url": "https://apples.com"}, {"url": "https://oranges.com"}],  # associate metadata
+)
 
+# Define a function for querying
+def search_query(query):
+    results = memory.search(query, top_n=1)  # Search for top result
+    return results
 
-# Interface
-def onInput(paragraph, progress = gr.Progress()):
-    tokens = tokenize(paragraph)
-    
-    if not tokens:  # Handle case with no tokens found
-        return np.zeros(300).tolist()  # Return a zero vector of appropriate dimension
-    
-    merged_vector = np.zeros(300)  # Assuming vectors are 300-dimensional
+# Create Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("### VectorDB Search")
     
-    # Merge vectors using NumPy
-    totalTokens = len(tokens)
-    for ind, token in enumerate(tokens):
-        completion = 0.2*((ind+1)/totalTokens)
-        progress(0.6 + completion, f"Merging {token}, Token #{tokens.index(token)+1}/{len(tokens)}")
+    with gr.Row():
+        input_query = gr.Textbox(label="Enter your query")
+        output_result = gr.Textbox(label="Search Results", interactive=False)
 
-        if token not in vectors:
-            continue
-
-        vector = vectors[token]
-        merged_vector += vector
-    
-    # Normalize
-    merged_vector /= len(tokens)
-    
-    return merged_vector.tolist(), json.dumps(tokens)
+    search_button = gr.Button("Search")
+    search_button.click(search_query, inputs=input_query, outputs=output_result)
 
-demo = gr.Interface(fn=onInput, inputs="text", outputs=["text", "json"])
-demo.launch()
+# Run the Gradio app
+demo.launch()