configurable chunking and embedding

gradio_app/app.py

@@ -10,12 +10,13 @@ from time import perf_counter
 
 import gradio as gr
 import markdown
+import lancedb
 from jinja2 import Environment, FileSystemLoader
 
 from gradio_app.backend.ChatGptInteractor import num_tokens_from_messages
 from gradio_app.backend.cross_encoder import rerank_with_cross_encoder
 from gradio_app.backend.query_llm import *
-from gradio_app.backend.semantic_search import table, embedder
+from gradio_app.backend.embedders import EmbedderFactory
 
 from settings import *
 
@@ -30,6 +31,8 @@ env = Environment(loader=FileSystemLoader('gradio_app/templates'))
 context_template = env.get_template('context_template.j2')
 context_html_template = env.get_template('context_html_template.j2')
 
+db = lancedb.connect(LANCEDB_DIRECTORY)
+
 # Examples
 examples = [
     'What is BERT?',
@@ -46,7 +49,7 @@ def add_text(history, text):
     return history, gr.Textbox(value="", interactive=False)
 
 
-def bot(history, llm, cross_enc):
+def bot(history, llm, cross_enc, chunk, embed):
     history[-1][1] = ""
     query = history[-1][0]
 
@@ -55,27 +58,33 @@ def bot(history, llm, cross_enc):
 
     logger.info('Retrieving documents...')
     gr.Info('Start documents retrieval ...')
-    time = perf_counter()
+    t = perf_counter()
+
+    table_name = f'{LANCEDB_TABLE_NAME}_{chunk}_{embed}'
+    table = db.open_table(table_name)
+
+    embedder = EmbedderFactory.get_embedder(embed)
 
-    query_vec = embedder.embed(query)[0]
+    query_vec = embedder.embed([query])[0]
     documents = table.search(query_vec, vector_column_name=VECTOR_COLUMN_NAME)
-    documents = documents.limit(TOP_K_RANK).to_list()
-    thresh_dist = thresh_distances[EMBED_NAME]
+    top_k_rank = TOP_K_RANK if cross_enc is not None else TOP_K_RERANK
+    documents = documents.limit(top_k_rank).to_list()
+    thresh_dist = thresh_distances[embed]
     thresh_dist = max(thresh_dist, min(d['_distance'] for d in documents))
     documents = [d for d in documents if d['_distance'] <= thresh_dist]
     documents = [doc[TEXT_COLUMN_NAME] for doc in documents]
 
-    time = perf_counter() - time
-    logger.info(f'Finished Retrieving documents in {round(time, 2)} seconds...')
+    t = perf_counter() - t
+    logger.info(f'Finished Retrieving documents in {round(t, 2)} seconds...')
 
     logger.info('Reranking documents...')
     gr.Info('Start documents reranking ...')
-    time = perf_counter()
+    t = perf_counter()
 
     documents = rerank_with_cross_encoder(cross_enc, documents, query)
 
-    time = perf_counter() - time
-    logger.info(f'Finished Reranking documents in {round(time, 2)} seconds...')
+    t = perf_counter() - t
+    logger.info(f'Finished Reranking documents in {round(t, 2)} seconds...')
 
     msg_constructor = get_message_constructor(llm)
     while len(documents) != 0:
@@ -91,11 +100,14 @@ def bot(history, llm, cross_enc):
         raise gr.Error('Model context length exceeded, reload the page')
 
     llm_gen = get_llm_generator(llm)
+    logger.info('Generating answer...')
+    t = perf_counter()
     for part in llm_gen(messages):
         history[-1][1] += part
         yield history, context_html
     else:
-        print('Finished generation stream.')
+        t = perf_counter() - t
+        logger.info(f'Finished Generating answer in {round(t, 2)} seconds...')
 
 
 with gr.Blocks() as demo:
@@ -109,7 +121,7 @@ with gr.Blocks() as demo:
                 bubble_full_width=False,
                 show_copy_button=True,
                 show_share_button=True,
-                height=600,
+                height=500,
             )
 
             with gr.Row():
@@ -121,14 +133,22 @@ with gr.Blocks() as demo:
                 )
                 txt_btn = gr.Button(value="Submit text", scale=1)
 
-            llm_name = gr.Radio(
+            chunk_name = gr.Radio(
                 choices=[
-                    "gpt-3.5-turbo",
-                    "mistralai/Mistral-7B-Instruct-v0.1",
-                    "GeneZC/MiniChat-3B",
+                    "md",
+                    "txt",
                 ],
-                value="gpt-3.5-turbo",
-                label='LLM'
+                value="md",
+                label='Chunking policy'
+            )
+
+            embed_name = gr.Radio(
+                choices=[
+                    "text-embedding-ada-002",
+                    "sentence-transformers/all-MiniLM-L6-v2",
+                ],
+                value="text-embedding-ada-002",
+                label='Embedder'
             )
 
             cross_enc_name = gr.Radio(
@@ -141,6 +161,16 @@ with gr.Blocks() as demo:
                 label='Cross-Encoder'
             )
 
+            llm_name = gr.Radio(
+                choices=[
+                    "gpt-3.5-turbo",
+                    "mistralai/Mistral-7B-Instruct-v0.1",
+                    "GeneZC/MiniChat-3B",
+                ],
+                value="gpt-3.5-turbo",
+                label='LLM'
+            )
+
             # Examples
             gr.Examples(examples, input_textbox)
 
@@ -151,7 +181,7 @@ with gr.Blocks() as demo:
     txt_msg = txt_btn.click(
         add_text, [chatbot, input_textbox], [chatbot, input_textbox], queue=False
     ).then(
-        bot, [chatbot, llm_name, cross_enc_name], [chatbot, context_html]
+        bot, [chatbot, llm_name, cross_enc_name, chunk_name, embed_name], [chatbot, context_html]
     )
 
     # Turn it back on
@@ -159,7 +189,7 @@ with gr.Blocks() as demo:
 
     # Turn off interactivity while generating if you hit enter
     txt_msg = input_textbox.submit(add_text, [chatbot, input_textbox], [chatbot, input_textbox], queue=False).then(
-        bot, [chatbot, llm_name, cross_enc_name], [chatbot, context_html])
+        bot, [chatbot, llm_name, cross_enc_name, chunk_name, embed_name], [chatbot, context_html])
 
     # Turn it back on
     txt_msg.then(lambda: gr.Textbox(interactive=True), None, [input_textbox], queue=False)

gradio_app/backend/semantic_search.py

@@ -1,14 +0,0 @@
-import logging
-import lancedb
-
-from gradio_app.backend.embedders import EmbedderFactory
-from settings import *
-
-
-# Setting up the logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-embedder = EmbedderFactory.get_embedder(EMBED_NAME)
-
-db = lancedb.connect(LANCEDB_DIRECTORY)
-table = db.open_table(LANCEDB_TABLE_NAME)

prep_scripts/lancedb_setup.py

@@ -1,4 +1,5 @@
 import shutil
+import time
 
 import lancedb
 import openai
@@ -18,7 +19,7 @@ with open('data/openaikey.txt') as f:
 openai.api_key = OPENAI_KEY
 
 
-shutil.rmtree(LANCEDB_DIRECTORY, ignore_errors=True)
+# shutil.rmtree(LANCEDB_DIRECTORY, ignore_errors=True)
 db = lancedb.connect(LANCEDB_DIRECTORY)
 batch_size = 32
 
@@ -27,7 +28,8 @@ schema = pa.schema([
     pa.field(TEXT_COLUMN_NAME, pa.string()),
     pa.field(DOCUMENT_PATH_COLUMN_NAME, pa.string()),
 ])
-tbl = db.create_table(LANCEDB_TABLE_NAME, schema=schema, mode="overwrite")
+table_name = f'{LANCEDB_TABLE_NAME}_{CHUNK_POLICY}_{EMBED_NAME}'
+tbl = db.create_table(table_name, schema=schema, mode="overwrite")
 
 input_dir = Path(MARKDOWN_SOURCE_DIR)
 files = list(input_dir.rglob("*"))
@@ -45,15 +47,21 @@ for file in files:
     with open(file, encoding='utf-8') as f:
         f = f.read()
         f = remove_comments(f)
-        f = split_markdown(f)
+        if CHUNK_POLICY == "txt":
+            f = md2txt_then_split(f)
+        else:
+            assert CHUNK_POLICY == "md"
+            f = split_markdown(f)
         chunks.extend((chunk, os.path.abspath(file)) for chunk in f)
 
 from matplotlib import pyplot as plt
 plt.hist([len(c) for c, d in chunks], bins=100)
+plt.title(table_name)
 plt.show()
 
 embedder = EmbedderFactory.get_embedder(EMBED_NAME)
 
+time_embed, time_ingest = [], []
 for i in tqdm.tqdm(range(0, int(np.ceil(len(chunks) / batch_size)))):
     texts, doc_paths = [], []
     for text, doc_path in chunks[i * batch_size:(i + 1) * batch_size]:
@@ -61,14 +69,24 @@ for i in tqdm.tqdm(range(0, int(np.ceil(len(chunks) / batch_size)))):
             texts.append(text)
             doc_paths.append(doc_path)
 
+    t = time.perf_counter()
     encoded = embedder.embed(texts)
+    time_embed.append(time.perf_counter() - t)
+
     df = pd.DataFrame({
         VECTOR_COLUMN_NAME: encoded,
         TEXT_COLUMN_NAME: texts,
         DOCUMENT_PATH_COLUMN_NAME: doc_paths,
     })
 
+    t = time.perf_counter()
     tbl.add(df)
+    time_ingest.append(time.perf_counter() - t)
+
+
+time_embed = sum(time_embed)
+time_ingest = sum(time_ingest)
+print(f'Embedding: {time_embed}, Ingesting: {time_ingest}')
 
 
 

prep_scripts/markdown_to_text.py

@@ -1,6 +1,9 @@
 import os
 import re
 
+from bs4 import BeautifulSoup
+from markdown import markdown
+
 from settings import *
 
 
@@ -95,3 +98,32 @@ def split_markdown(md):
 
     return res
 
+
+def markdown_to_text(markdown_string):
+    """ Converts a markdown string to plaintext """
+
+    # md -> html -> text since BeautifulSoup can extract text cleanly
+    html = markdown(markdown_string)
+
+    html = re.sub(r'<!--((.|\n)*)-->', '', html)
+    html = re.sub('<code>bash', '<code>', html)
+
+    # extract text
+    soup = BeautifulSoup(html, "html.parser")
+    text = ''.join(soup.findAll(string=True))
+
+    text = re.sub('```(py|diff|python)', '', text)
+    text = re.sub('```\n', '\n', text)
+    text = re.sub('-         .*', '', text)
+    text = text.replace('...', '')
+    text = re.sub('\n(\n)+', '\n\n', text)
+
+    return text
+
+
+def md2txt_then_split(md):
+    txt = markdown_to_text(md)
+    return split_content(txt)
+
+
+

settings.py

@@ -5,8 +5,11 @@ VECTOR_COLUMN_NAME = "embedding"
 TEXT_COLUMN_NAME = "text"
 DOCUMENT_PATH_COLUMN_NAME = "document_path"
 
-# EMBED_NAME = "sentence-transformers/all-MiniLM-L6-v2"
-EMBED_NAME = "text-embedding-ada-002"
+CHUNK_POLICY = "md"
+# CHUNK_POLICY = "txt"
+
+EMBED_NAME = "sentence-transformers/all-MiniLM-L6-v2"
+# EMBED_NAME = "text-embedding-ada-002"
 
 TOP_K_RANK = 50
 TOP_K_RERANK = 5
@@ -28,5 +31,5 @@ context_lengths = {
     "gpt-3.5-turbo": 4096,
     "sentence-transformers/all-MiniLM-L6-v2": 128,
     "thenlper/gte-large": 512,
-    "text-embedding-ada-002": 8191,
+    "text-embedding-ada-002": 1000,  # actual context length is 8191, but it's too much
 }