Created Duckdb Vector Store

app.py

@@ -1,7 +1,10 @@
 import re
+import duckdb
+import pandas as pd
 import gradio as gr
 from io import StringIO
-import pandas as pd
+from langchain_community.vectorstores.duckdb import DuckDB
+from langchain_community.embeddings import HuggingFaceBgeEmbeddings
 from langchain_community.document_loaders import RecursiveUrlLoader
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.document_transformers import Html2TextTransformer
@@ -9,6 +12,24 @@ from langchain_community.document_transformers import Html2TextTransformer
 
 TAB_LINES = 22
 
+# Embedding Model args
+model_name = "BAAI/bge-small-en-v1.5"
+model_kwargs = {'device': 'cpu'}
+encode_kwargs = {'normalize_embeddings': True}
+
+# HuggingFace Embeddings
+hf = HuggingFaceBgeEmbeddings(
+    model_name=model_name,
+    model_kwargs=model_kwargs,
+    encode_kwargs=encode_kwargs
+)
+
+# DuckDB Connection
+con = duckdb.connect('Collections.duckdb')
+
+# DuckDB Vector Store
+vector_store = DuckDB(connection = con, embedding=hf)
+
 def html_only_metadata_extractor(raw_html, url, response):
     content_type = response.headers.get("Content-Type", "")
     if "text/html" in content_type:
@@ -90,6 +111,10 @@ def concat_dfs(df_list):
     concatenated_df = pd.concat(df_list, ignore_index=True)
     return concatenated_df
 
+def create_embeddings(docs):
+    ids = vector_store.add_documents(docs)
+    result = con.execute(f"SELECT * FROM embeddings").fetchdf()
+    return result[result['id'].isin(ids)]
 
 def get_docs(url, max_depth):
     raw_html = scrape_text(url, max_depth)
@@ -108,8 +133,10 @@ def get_docs(url, max_depth):
     text_splitter = RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=200)
     documents_splits = text_splitter.split_documents(clean_docs)
     formatted_chunks = format_chunks_with_spaces(documents_splits)
+    embeddings = create_embeddings(documents_splits)
+    
     
-    return format_page_content(raw_html), format_page_content(clean_docs), concat_tables, format_metdata(raw_html), formatted_chunks
+    return format_page_content(raw_html), format_page_content(clean_docs), concat_tables, format_metdata(raw_html), formatted_chunks, embeddings
 
 
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple", secondary_hue="indigo")) as demo:
@@ -147,9 +174,11 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple", secondary_hue="indigo"
                 with gr.Tab("Metadata"):
                     metadata = gr.Textbox(lines=TAB_LINES, label="Metadata", value="", interactive=False,
                                           autoscroll=False)
+                with gr.Tab("Embeddings"):
+                    embeddings = gr.Dataframe(label="Vector Store", interactive=False)
         
         scarpe_url_button.click(get_docs, inputs=[url_input, max_depth], outputs=[raw_page_content, page_content, tables,
-                                                                     metadata, parsed_chunks])
+                                                                     metadata, parsed_chunks, embeddings])
 
         
 if __name__ == "__main__":

requirements.txt

@@ -5,4 +5,6 @@ langchain-text-splitters
 html2text
 lxml
 beautifulsoup4
-html5lib
+html5lib
+duckdb
+sentence_transformers