Added basic semantic search functionality

app.py

@@ -1,5 +1,5 @@
 import nltk
-
+from typing import TypeVar
 nltk.download('names')
 nltk.download('stopwords')
 nltk.download('wordnet')
@@ -9,11 +9,26 @@ from search_funcs.fast_bm25 import BM25
 from search_funcs.clean_funcs import initial_clean, get_lemma_tokens#, stem_sentence
 from nltk import word_tokenize
 
+PandasDataFrame = TypeVar('pd.core.frame.DataFrame')
 
 import gradio as gr
 import pandas as pd
 import os
 
+from itertools import compress
+
+#from langchain.embeddings import HuggingFaceEmbeddings
+#from langchain.vectorstores import FAISS
+from transformers import AutoModel
+
+import search_funcs.ingest as ing
+import search_funcs.chatfuncs as chatf
+
+# Import Chroma and instantiate a client. The default Chroma client is ephemeral, meaning it will not save to disk.
+import chromadb
+
+#collection = client.create_collection(name="my_collection")
+
 def prepare_input_data(in_file, text_column, clean="No", progress=gr.Progress()):
 
     filename = in_file.name
@@ -63,7 +78,7 @@ def get_file_path_end(file_path):
     
     return filename_without_extension
 
-def save_prepared_data(in_file, prepared_text_list, in_df, in_column):
+def save_prepared_data(in_file, prepared_text_list, in_df, in_bm25_column):
 
     # Check if the list and the dataframe have the same length
     if len(prepared_text_list) != len(in_df):
@@ -73,10 +88,10 @@ def save_prepared_data(in_file, prepared_text_list, in_df, in_column):
 
     file_name = get_file_path_end(in_file.name) + "_cleaned" + file_end
 
-    prepared_text_df = pd.DataFrame(data={in_column + "_cleaned":prepared_text_list})
+    prepared_text_df = pd.DataFrame(data={in_bm25_column + "_cleaned":prepared_text_list})
 
     # Drop original column from input file to reduce file size
-    in_df = in_df.drop(in_column, axis = 1)
+    in_df = in_df.drop(in_bm25_column, axis = 1)
 
     prepared_df = pd.concat([in_df, prepared_text_df], axis = 1)
 
@@ -194,7 +209,7 @@ def read_file(filename):
     elif file_type == 'parquet':
         return pd.read_parquet(filename).reset_index().drop(["index", "Unnamed: 0"], axis=1, errors="ignore")
 
-def put_columns_in_df(in_file, in_column):
+def put_columns_in_df(in_file, in_bm25_column):
     '''
     When file is loaded, update the column dropdown choices and change 'clean data' dropdown option to 'no'.
     '''
@@ -213,12 +228,12 @@ def put_columns_in_df(in_file, in_column):
     return gr.Dropdown(choices=concat_choices), gr.Dropdown(value="No", choices = ["Yes", "No"]),\
         gr.Dropdown(choices=concat_choices)
 
-def put_columns_in_join_df(in_file, in_column):
+def put_columns_in_join_df(in_file, in_bm25_column):
     '''
     When file is loaded, update the column dropdown choices and change 'clean data' dropdown option to 'no'.
     '''
 
-    print("in_column")
+    print("in_bm25_column")
 
     new_choices = []
     concat_choices = []
@@ -241,11 +256,293 @@ def dummy_function(gradio_component):
 
 def display_info(info_component):
     gr.Info(info_component)
-# %%
+
+embeddings_name = "jinaai/jina-embeddings-v2-small-en"
+
+#embeddings_name = "BAAI/bge-base-en-v1.5"
+import chromadb
+from typing_extensions import Protocol
+from chromadb import Documents, EmbeddingFunction, Embeddings
+
+embeddings_model = AutoModel.from_pretrained(embeddings_name, trust_remote_code=True)
+
+class MyEmbeddingFunction(EmbeddingFunction):
+    def __call__(self, input) -> Embeddings:
+     
+
+        embeddings = []
+        for text in input:
+            embeddings.append(embeddings_model.encode(text))
+
+        return embeddings
+
+
+def load_embeddings(embeddings_name = "jinaai/jina-embeddings-v2-small-en"):
+    '''
+    Load embeddings model and create a global variable based on it.
+    '''
+
+    # Import Chroma and instantiate a client. The default Chroma client is ephemeral, meaning it will not save to disk.
+    
+    #else: 
+    embeddings_func = AutoModel.from_pretrained(embeddings_name, trust_remote_code=True)
+
+    global embeddings
+
+    embeddings = embeddings_func
+
+    return embeddings
+
+embeddings = load_embeddings(embeddings_name)
+
+def docs_to_chroma_save(docs_out, embeddings=embeddings, progress=gr.Progress()):
+    '''
+    Takes a Langchain document class and saves it into a Chroma sqlite file.
+    '''
+
+
+
+    print(f"> Total split documents: {len(docs_out)}")
+
+    #print(docs_out)
+
+    page_contents = [doc.page_content for doc in docs_out]
+    page_meta = [doc.metadata for doc in docs_out]
+    ids_range = range(0,len(page_contents)) 
+    ids = [str(element) for element in ids_range]
+
+    embeddings_list = []
+    for page in progress.tqdm(page_contents, desc = "Preparing search index", unit = "rows"):
+        embeddings_list.append(embeddings.encode(sentences=page, max_length=1024).tolist())
+
+
+    client = chromadb.PersistentClient(path=".")
+
+    # Create a new Chroma collection to store the supporting evidence. We don't need to specify an embedding fuction, and the default will be used.
+    try:
+        collection = client.get_collection(name="my_collection")
+        client.delete_collection(name="my_collection")
+    except: 
+        collection = client.create_collection(name="my_collection")
+                                          
+    collection.add(
+    documents = page_contents,
+    embeddings = embeddings_list,
+    metadatas = page_meta,
+    ids = ids)
+
+    #chatf.vectorstore = vectorstore_func
+
+    out_message = "Document processing complete"
+
+    return out_message, collection
+
+def jina_simple_retrieval(new_question_kworded, vectorstore, docs, k_val, out_passages,
+                           vec_score_cut_off, vec_weight): # ,vectorstore, embeddings
+
+            from numpy.linalg import norm
+
+            cos_sim = lambda a,b: (a @ b.T) / (norm(a)*norm(b))
+
+            query = embeddings.encode(new_question_kworded)
+
+            # Calculate cosine similarity with each string in the list
+            cosine_similarities = [cos_sim(query, string_vector) for string_vector in vectorstore]
+
+
+
+            print(cosine_similarities)
+
+
+
+            #vectorstore=globals()["vectorstore"]
+            #embeddings=globals()["embeddings"]
+            doc_df = pd.DataFrame()
+
+
+            docs = vectorstore.similarity_search_with_score(new_question_kworded, k=k_val)
+
+            print("Docs from similarity search:")
+            print(docs)
+
+            # Keep only documents with a certain score
+            docs_len = [len(x[0].page_content) for x in docs]
+            docs_scores = [x[1] for x in docs]
+
+            # Only keep sources that are sufficiently relevant (i.e. similarity search score below threshold below)
+            score_more_limit = pd.Series(docs_scores) < vec_score_cut_off
+            docs_keep = list(compress(docs, score_more_limit))
+
+            if not docs_keep:
+                return [], pd.DataFrame(), []
+
+            # Only keep sources that are at least 100 characters long
+            length_more_limit = pd.Series(docs_len) >= 100
+            docs_keep = list(compress(docs_keep, length_more_limit))
+
+            if not docs_keep:
+                return [], pd.DataFrame(), []
+
+            docs_keep_as_doc = [x[0] for x in docs_keep]
+            docs_keep_length = len(docs_keep_as_doc)
+
+
+                
+            if docs_keep_length == 1:
+
+                content=[]
+                meta_url=[]
+                score=[]
+                
+                for item in docs_keep:
+                    content.append(item[0].page_content)
+                    meta_url.append(item[0].metadata['source'])
+                    score.append(item[1])       
+
+                # Create df from 'winning' passages
+
+                doc_df = pd.DataFrame(list(zip(content, meta_url, score)),
+                columns =['page_content', 'meta_url', 'score'])
+
+                docs_content = doc_df['page_content'].astype(str)
+                docs_url = doc_df['meta_url']
+
+                return docs_keep_as_doc, docs_content, docs_url
+            
+            # Check for if more docs are removed than the desired output
+            if out_passages > docs_keep_length: 
+                out_passages = docs_keep_length
+                k_val = docs_keep_length
+                     
+            vec_rank = [*range(1, docs_keep_length+1)]
+            vec_score = [(docs_keep_length/x)*vec_weight for x in vec_rank]
+
+            ## Calculate final score based on three ranking methods
+            final_score = [a for a in zip(vec_score)]
+            final_rank = [sorted(final_score, reverse=True).index(x)+1 for x in final_score]
+            # Force final_rank to increment by 1 each time
+            final_rank = list(pd.Series(final_rank).rank(method='first'))
+
+            #print("final rank: " + str(final_rank))
+            #print("out_passages: " + str(out_passages))
+
+            best_rank_index_pos = []
+
+            for x in range(1,out_passages+1):
+                try:
+                    best_rank_index_pos.append(final_rank.index(x))
+                except IndexError: # catch the error
+                    pass
+
+            # Adjust best_rank_index_pos to 
+
+            best_rank_pos_series = pd.Series(best_rank_index_pos)
+
+
+            docs_keep_out = [docs_keep[i] for i in best_rank_index_pos]
+        
+            # Keep only 'best' options
+            docs_keep_as_doc = [x[0] for x in docs_keep_out]
+                               
+            # Make df of best options
+            doc_df = create_doc_df(docs_keep_out)
+
+            return docs_keep_as_doc, doc_df, docs_keep_out
+
+def chroma_retrieval(new_question_kworded, vectorstore, docs, k_val, out_passages,
+                           vec_score_cut_off, vec_weight): # ,vectorstore, embeddings
+
+            query = embeddings.encode(new_question_kworded).tolist()
+
+            docs = vectorstore.query(
+            query_embeddings=query,
+            n_results= 9999 # No practical limit on number of responses returned
+            #where={"metadata_field": "is_equal_to_this"},
+            #where_document={"$contains":"search_string"}
+            )
+
+            # Calculate cosine similarity with each string in the list
+            #cosine_similarities = [cos_sim(query, string_vector) for string_vector in vectorstore]
+
+            #print(docs)
+
+            #vectorstore=globals()["vectorstore"]
+            #embeddings=globals()["embeddings"]
+            df = pd.DataFrame(data={'ids': docs['ids'][0],
+                                    'documents': docs['documents'][0],
+                                    'metadatas':docs['metadatas'][0],
+                                    'distances':docs['distances'][0]#,                                    
+                                    #'embeddings': docs['embeddings']
+                                    })
+            
+            def create_docs_keep_from_df(df):
+                dict_out = {'ids' : [df['ids']],
+                            'documents': [df['documents']],
+                            'metadatas': [df['metadatas']],
+                            'distances': [df['distances']],
+                            'embeddings': None
+                            }
+                return dict_out
+                
+            # Prepare the DataFrame by transposing
+            df_docs = df#.apply(lambda x: x.explode()).reset_index(drop=True)
+
+            #print(df_docs)
+
+
+            # Keep only documents with a certain score
+            
+            docs_scores = df_docs["distances"] #.astype(float)
+
+            #print(docs_scores)
+
+            # Only keep sources that are sufficiently relevant (i.e. similarity search score below threshold below)
+            score_more_limit = df_docs.loc[docs_scores < vec_score_cut_off, :]
+            docs_keep = create_docs_keep_from_df(score_more_limit) #list(compress(docs, score_more_limit))
+
+            #print(docs_keep)
+
+            if not docs_keep:
+                return 'No result found!', ""
+
+            # Only keep sources that are at least 100 characters long
+            docs_len = score_more_limit["documents"].str.len() >= 100
+            length_more_limit = score_more_limit.loc[docs_len, :] #pd.Series(docs_len) >= 100
+            docs_keep = create_docs_keep_from_df(length_more_limit) #list(compress(docs_keep, length_more_limit))
+
+            #print(docs_keep)
+
+            print(length_more_limit)
+
+            if not docs_keep:
+                return 'No result found!', ""
+            
+            results_df_name = "semantic_search_result.csv"
+            length_more_limit.to_csv(results_df_name, index= None)
+            results_first_text = length_more_limit["documents"][0]
+
+            
+            return results_first_text, results_df_name
+
 # ## Gradio app - BM25 search
 block = gr.Blocks(theme = gr.themes.Base())
 
-with block:  
+with block:
+
+    ingest_text = gr.State()
+    ingest_metadata = gr.State()
+    ingest_docs = gr.State()
+    vectorstore_state = gr.State() # globals()["vectorstore"]
+    embeddings_state = gr.State() # globals()["embeddings"]
+
+    k_val = gr.State(100)
+    out_passages = gr.State(100)
+    vec_score_cut_off = gr.State(100)
+    vec_weight = gr.State(1)
+
+    docs_keep_as_doc_state = gr.State()
+    doc_df_state = gr.State()
+    docs_keep_out_state = gr.State()
 
     corpus_state = gr.State()
     data_state = gr.State(pd.DataFrame())
@@ -267,14 +564,18 @@ depends on factors such as the type of documents or queries. Information taken f
     # Fast text search
     Enter a text query below to search through a text data column and find relevant terms. It will only find terms containing the exact text you enter. Your data should contain at least 20 entries for the search to consistently return results.
     """)
+
     
     with gr.Tab(label="Search your data"):
+        with gr.Row():
+            current_source = gr.Textbox(label="Current data source(s)", value="None")
+
         with gr.Accordion(label = "Load in data", open=True):
-            in_corpus = gr.File(label="Upload your search data here")
+            in_bm25_file = gr.File(label="Upload your search data here")
             with gr.Row():
-                in_column = gr.Dropdown(label="Enter the name of the text column in the data file to search")
+                in_bm25_column = gr.Dropdown(label="Enter the name of the text column in the data file to search")
                 
-                load_data_button = gr.Button(value="Load data")
+                load_bm25_data_button = gr.Button(value="Load data")
                  
 
             with gr.Row():
@@ -291,6 +592,21 @@ depends on factors such as the type of documents or queries. Information taken f
             with gr.Row():
                 output_single_text = gr.Textbox(label="Top result")
                 output_file = gr.File(label="File output")
+
+    
+    with gr.Tab("Fuzzy/semantic search"):
+        with gr.Accordion("CSV/Excel file", open = True):
+            in_semantic_file = gr.File(label="Upload data file for semantic search")
+            in_semantic_column = gr.Dropdown(label="Enter the name of the text column in the data file to search")
+            load_semantic_data_button = gr.Button(value="Load in CSV/Excel file", variant="secondary", scale=0)
+        
+        ingest_embed_out = gr.Textbox(label="File/web page preparation progress")
+        semantic_query = gr.Textbox(label="Enter semantic search query here")
+        semantic_submit = gr.Button(value="Start semantic search", variant="secondary", scale = 1)
+
+        with gr.Row():
+            semantic_output_single_text = gr.Textbox(label="Top result")
+            semantic_output_file = gr.File(label="File output")
             
 
     with gr.Tab(label="Advanced options"):
@@ -327,28 +643,33 @@ depends on factors such as the type of documents or queries. Information taken f
     in_no_search_results_button.click(display_info, inputs=in_no_search_info)
     
 
-    in_corpus.upload(put_columns_in_df, inputs=[in_corpus, in_column], outputs=[in_column, in_clean_data, search_df_join_column])
+    # Update dropdowns upon initial file load
+    in_bm25_file.upload(put_columns_in_df, inputs=[in_bm25_file, in_bm25_column], outputs=[in_bm25_column, in_clean_data, search_df_join_column])
     in_join_file.upload(put_columns_in_join_df, inputs=[in_join_file, in_join_column], outputs=[in_join_column])
-
-    # Load in the data
-    load_data_button.click(fn=prepare_input_data, inputs=[in_corpus, in_column, in_clean_data], outputs=[corpus_state, load_finished_message, data_state, output_file]).\
+ 
+    # Load in BM25 data
+    load_bm25_data_button.click(fn=prepare_input_data, inputs=[in_bm25_file, in_bm25_column, in_clean_data], outputs=[corpus_state, load_finished_message, data_state, output_file]).\
     then(fn=prepare_bm25, inputs=[corpus_state, in_k1, in_b, in_alpha], outputs=[load_finished_message]).\
-    then(fn=put_columns_in_df, inputs=[in_corpus, in_column], outputs=[in_column, in_clean_data, search_df_join_column])
-
-    #save_clean_data_button.click(fn=save_prepared_data, inputs=[in_corpus, corpus_state, data_state, in_column], outputs=[output_file])   
+    then(fn=put_columns_in_df, inputs=[in_bm25_file, in_bm25_column], outputs=[in_bm25_column, in_clean_data, search_df_join_column])
+   
+    # BM25 search functions on click or enter
+    search_button.click(fn=bm25_search, inputs=[in_query, in_no_search_results, data_state, in_bm25_column, in_clean_data, in_join_file, in_join_column, search_df_join_column], outputs=[output_single_text, output_file, mod_query], api_name="search")
+    in_query.submit(fn=bm25_search, inputs=[in_query, in_no_search_results, data_state, in_bm25_column, in_clean_data, in_join_file, in_join_column, search_df_join_column], outputs=[output_single_text, output_file, mod_query])
     
-
-    # Search functions on click or enter
-    search_button.click(fn=bm25_search, inputs=[in_query, in_no_search_results, data_state, in_column, in_clean_data, in_join_file, in_join_column, search_df_join_column],
-                        outputs=[output_single_text, output_file, mod_query], api_name="search")
+    # Load in a csv/excel file for semantic search
+    in_semantic_file.upload(put_columns_in_df, inputs=[in_semantic_file, in_semantic_column], outputs=[in_semantic_column, in_clean_data, search_df_join_column])
+    load_semantic_data_button.click(ing.parse_csv_or_excel, inputs=[in_semantic_file, in_semantic_column], outputs=[ingest_text, current_source]).\
+             then(ing.csv_excel_text_to_docs, inputs=[ingest_text, in_semantic_column], outputs=[ingest_docs]).\
+             then(docs_to_chroma_save, inputs=[ingest_docs], outputs=[ingest_embed_out, vectorstore_state])
     
-    in_query.submit(fn=bm25_search, inputs=[in_query, in_no_search_results, data_state, in_column, in_clean_data, in_join_file, in_join_column, search_df_join_column],
-                        outputs=[output_single_text, output_file, mod_query])
+    # Semantic search query
+    semantic_submit.click(chroma_retrieval, inputs=[semantic_query, vectorstore_state, ingest_docs, k_val,out_passages, vec_score_cut_off, vec_weight], outputs=[semantic_output_single_text, semantic_output_file], api_name="semantic")
     
     # Dummy functions just to get dropdowns to work correctly with Gradio 3.50
-    in_column.change(dummy_function, in_column, None)
+    in_bm25_column.change(dummy_function, in_bm25_column, None)
     search_df_join_column.change(dummy_function, search_df_join_column, None)
     in_join_column.change(dummy_function, in_join_column, None)
+    in_semantic_column.change(dummy_function, in_join_column, None)
 
 block.queue().launch(debug=True)
 

search_funcs/chatfuncs.py

@@ -0,0 +1,393 @@
+import re
+import os
+from typing import TypeVar, List
+import pandas as pd
+
+
+# Model packages
+import torch.cuda
+
+# Alternative model sources
+#from dataclasses import asdict, dataclass
+
+# Langchain functions
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.docstore.document import Document
+
+# For keyword extraction (not currently used)
+#import nltk
+#nltk.download('wordnet')
+from nltk.corpus import stopwords
+from nltk.tokenize import RegexpTokenizer
+from nltk.stem import WordNetLemmatizer
+
+# For Name Entity Recognition model
+#from span_marker import SpanMarkerModel # Not currently used
+
+
+import gradio as gr
+
+torch.cuda.empty_cache()
+
+PandasDataFrame = TypeVar('pd.core.frame.DataFrame')
+
+embeddings = None  # global variable setup
+vectorstore = None # global variable setup
+model_type = None # global variable setup
+
+max_memory_length = 0 # How long should the memory of the conversation last?
+
+full_text = "" # Define dummy source text (full text) just to enable highlight function to load
+
+model = [] # Define empty list for model functions to run
+tokenizer = [] # Define empty list for model functions to run
+
+## Highlight text constants
+hlt_chunk_size = 12
+hlt_strat = [" ", ". ", "! ", "? ", ": ", "\n\n", "\n", ", "]
+hlt_overlap = 4
+
+## Initialise NER model ##
+ner_model = []#SpanMarkerModel.from_pretrained("tomaarsen/span-marker-mbert-base-multinerd") # Not currently used
+
+
+# Currently set gpu_layers to 0 even with cuda due to persistent bugs in implementation with cuda
+if torch.cuda.is_available():
+    torch_device = "cuda"
+    gpu_layers = 0
+else: 
+    torch_device =  "cpu"
+    gpu_layers = 0
+
+print("Running on device:", torch_device)
+threads = 6 #torch.get_num_threads()
+print("CPU threads:", threads)
+
+# Vectorstore funcs
+
+# Prompt functions
+
+def write_out_metadata_as_string(metadata_in):
+    metadata_string = [f"{'  '.join(f'{k}: {v}' for k, v in d.items() if k != 'page_section')}" for d in metadata_in] # ['metadata']
+    return metadata_string
+
+
+def determine_file_type(file_path):
+        """
+        Determine the file type based on its extension.
+    
+        Parameters:
+            file_path (str): Path to the file.
+    
+        Returns:
+            str: File extension (e.g., '.pdf', '.docx', '.txt', '.html').
+        """
+        return os.path.splitext(file_path)[1].lower()
+
+
+def create_doc_df(docs_keep_out):
+    # Extract content and metadata from 'winning' passages.
+            content=[]
+            meta=[]
+            meta_url=[]
+            page_section=[]
+            score=[]
+
+            doc_df = pd.DataFrame()
+
+            
+
+            for item in docs_keep_out:
+                content.append(item[0].page_content)
+                meta.append(item[0].metadata)
+                meta_url.append(item[0].metadata['source'])
+
+                file_extension = determine_file_type(item[0].metadata['source'])
+                if (file_extension != ".csv") & (file_extension != ".xlsx"):
+                    page_section.append(item[0].metadata['page_section'])
+                else: page_section.append("")
+                score.append(item[1])       
+
+            # Create df from 'winning' passages
+
+            doc_df = pd.DataFrame(list(zip(content, meta, page_section, meta_url, score)),
+               columns =['page_content', 'metadata', 'page_section', 'meta_url', 'score'])
+
+            docs_content = doc_df['page_content'].astype(str)
+            doc_df['full_url'] = "https://" + doc_df['meta_url'] 
+
+            return doc_df
+
+
+def get_expanded_passages(vectorstore, docs, width):
+
+    """
+    Extracts expanded passages based on given documents and a width for context.
+    
+    Parameters:
+    - vectorstore: The primary data source.
+    - docs: List of documents to be expanded.
+    - width: Number of documents to expand around a given document for context.
+    
+    Returns:
+    - expanded_docs: List of expanded Document objects.
+    - doc_df: DataFrame representation of expanded_docs.
+    """
+
+    from collections import defaultdict
+    
+    def get_docs_from_vstore(vectorstore):
+        vector = vectorstore.docstore._dict
+        return list(vector.items())
+
+    def extract_details(docs_list):
+        docs_list_out = [tup[1] for tup in docs_list]
+        content = [doc.page_content for doc in docs_list_out]
+        meta = [doc.metadata for doc in docs_list_out]
+        return ''.join(content), meta[0], meta[-1]
+    
+    def get_parent_content_and_meta(vstore_docs, width, target):
+        #target_range = range(max(0, target - width), min(len(vstore_docs), target + width + 1))
+        target_range = range(max(0, target), min(len(vstore_docs), target + width + 1)) # Now only selects extra passages AFTER the found passage
+        parent_vstore_out = [vstore_docs[i] for i in target_range]
+        
+        content_str_out, meta_first_out, meta_last_out = [], [], []
+        for _ in parent_vstore_out:
+            content_str, meta_first, meta_last = extract_details(parent_vstore_out)
+            content_str_out.append(content_str)
+            meta_first_out.append(meta_first)
+            meta_last_out.append(meta_last)
+        return content_str_out, meta_first_out, meta_last_out
+
+    def merge_dicts_except_source(d1, d2):
+            merged = {}
+            for key in d1:
+                if key != "source":
+                    merged[key] = str(d1[key]) + " to " + str(d2[key])
+                else:
+                    merged[key] = d1[key]  # or d2[key], based on preference
+            return merged
+
+    def merge_two_lists_of_dicts(list1, list2):
+        return [merge_dicts_except_source(d1, d2) for d1, d2 in zip(list1, list2)]
+
+    # Step 1: Filter vstore_docs
+    vstore_docs = get_docs_from_vstore(vectorstore)
+    doc_sources = {doc.metadata['source'] for doc, _ in docs}
+    vstore_docs = [(k, v) for k, v in vstore_docs if v.metadata.get('source') in doc_sources]
+
+    # Step 2: Group by source and proceed
+    vstore_by_source = defaultdict(list)
+    for k, v in vstore_docs:
+        vstore_by_source[v.metadata['source']].append((k, v))
+        
+    expanded_docs = []
+    for doc, score in docs:
+        search_source = doc.metadata['source']
+        
+
+        #if file_type == ".csv" | file_type == ".xlsx":
+        #     content_str, meta_first, meta_last = get_parent_content_and_meta(vstore_by_source[search_source], 0, search_index)
+
+        #else:
+        search_section = doc.metadata['page_section']
+        parent_vstore_meta_section = [doc.metadata['page_section'] for _, doc in vstore_by_source[search_source]]
+        search_index = parent_vstore_meta_section.index(search_section) if search_section in parent_vstore_meta_section else -1
+
+        content_str, meta_first, meta_last = get_parent_content_and_meta(vstore_by_source[search_source], width, search_index)
+        meta_full = merge_two_lists_of_dicts(meta_first, meta_last)
+
+        expanded_doc = (Document(page_content=content_str[0], metadata=meta_full[0]), score)
+        expanded_docs.append(expanded_doc)
+
+    doc_df = pd.DataFrame()
+
+    doc_df = create_doc_df(expanded_docs)  # Assuming you've defined the 'create_doc_df' function elsewhere
+
+    return expanded_docs, doc_df
+
+def highlight_found_text(search_text: str, full_text: str, hlt_chunk_size:int=hlt_chunk_size, hlt_strat:List=hlt_strat, hlt_overlap:int=hlt_overlap) -> str:
+    """
+    Highlights occurrences of search_text within full_text.
+    
+    Parameters:
+    - search_text (str): The text to be searched for within full_text.
+    - full_text (str): The text within which search_text occurrences will be highlighted.
+    
+    Returns:
+    - str: A string with occurrences of search_text highlighted.
+    
+    Example:
+    >>> highlight_found_text("world", "Hello, world! This is a test. Another world awaits.")
+    'Hello, <mark style="color:black;">world</mark>! This is a test. Another <mark style="color:black;">world</mark> awaits.'
+    """
+
+    def extract_text_from_input(text, i=0):
+        if isinstance(text, str):
+            return text.replace("  ", " ").strip()
+        elif isinstance(text, list):
+            return text[i][0].replace("  ", " ").strip()
+        else:
+            return ""
+
+    def extract_search_text_from_input(text):
+        if isinstance(text, str):
+            return text.replace("  ", " ").strip()
+        elif isinstance(text, list):
+            return text[-1][1].replace("  ", " ").strip()
+        else:
+            return ""
+
+    full_text = extract_text_from_input(full_text)
+    search_text = extract_search_text_from_input(search_text)
+
+
+
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=hlt_chunk_size,
+        separators=hlt_strat,
+        chunk_overlap=hlt_overlap,
+    )
+    sections = text_splitter.split_text(search_text)
+
+    found_positions = {}
+    for x in sections:
+        text_start_pos = 0
+        while text_start_pos != -1:
+            text_start_pos = full_text.find(x, text_start_pos)
+            if text_start_pos != -1:
+                found_positions[text_start_pos] = text_start_pos + len(x)
+                text_start_pos += 1
+
+    # Combine overlapping or adjacent positions
+    sorted_starts = sorted(found_positions.keys())
+    combined_positions = []
+    if sorted_starts:
+        current_start, current_end = sorted_starts[0], found_positions[sorted_starts[0]]
+        for start in sorted_starts[1:]:
+            if start <= (current_end + 10):
+                current_end = max(current_end, found_positions[start])
+            else:
+                combined_positions.append((current_start, current_end))
+                current_start, current_end = start, found_positions[start]
+        combined_positions.append((current_start, current_end))
+
+    # Construct pos_tokens
+    pos_tokens = []
+    prev_end = 0
+    for start, end in combined_positions:
+        if end-start > 15: # Only combine if there is a significant amount of matched text. Avoids picking up single words like 'and' etc.
+            pos_tokens.append(full_text[prev_end:start])
+            pos_tokens.append('<mark style="color:black;">' + full_text[start:end] + '</mark>')
+            prev_end = end
+    pos_tokens.append(full_text[prev_end:])
+
+    return "".join(pos_tokens)
+
+
+# # Chat history functions
+
+def clear_chat(chat_history_state, sources, chat_message, current_topic):
+    chat_history_state = []
+    sources = ''
+    chat_message = ''
+    current_topic = ''
+
+    return chat_history_state, sources, chat_message, current_topic
+
+
+# Keyword functions
+
+def remove_q_stopwords(question): # Remove stopwords from question. Not used at the moment 
+    # Prepare keywords from question by removing stopwords
+    text = question.lower()
+
+    # Remove numbers
+    text = re.sub('[0-9]', '', text)
+
+    tokenizer = RegexpTokenizer(r'\w+')
+    text_tokens = tokenizer.tokenize(text)
+    #text_tokens = word_tokenize(text)
+    tokens_without_sw = [word for word in text_tokens if not word in stopwords]
+
+    # Remove duplicate words while preserving order
+    ordered_tokens = set()
+    result = []
+    for word in tokens_without_sw:
+        if word not in ordered_tokens:
+            ordered_tokens.add(word)
+            result.append(word)
+     
+
+
+    new_question_keywords = ' '.join(result)
+    return new_question_keywords
+
+def remove_q_ner_extractor(question):
+    
+    predict_out = ner_model.predict(question)
+
+
+
+    predict_tokens = [' '.join(v for k, v in d.items() if k == 'span') for d in predict_out]
+
+    # Remove duplicate words while preserving order
+    ordered_tokens = set()
+    result = []
+    for word in predict_tokens:
+        if word not in ordered_tokens:
+            ordered_tokens.add(word)
+            result.append(word)
+     
+
+
+    new_question_keywords = ' '.join(result).lower()
+    return new_question_keywords
+
+def apply_lemmatize(text, wnl=WordNetLemmatizer()):
+
+    def prep_for_lemma(text):
+
+        # Remove numbers
+        text = re.sub('[0-9]', '', text)
+        print(text)
+
+        tokenizer = RegexpTokenizer(r'\w+')
+        text_tokens = tokenizer.tokenize(text)
+        #text_tokens = word_tokenize(text)
+
+        return text_tokens
+
+    tokens = prep_for_lemma(text)
+
+    def lem_word(word):
+    
+        if len(word) > 3: out_word = wnl.lemmatize(word)
+        else: out_word = word
+
+        return out_word
+
+    return [lem_word(token) for token in tokens]
+
+def keybert_keywords(text, n, kw_model):
+    tokens_lemma = apply_lemmatize(text)
+    lemmatised_text = ' '.join(tokens_lemma)
+
+    keywords_text = KeyBERT(model=kw_model).extract_keywords(lemmatised_text, stop_words='english', top_n=n, 
+                                                   keyphrase_ngram_range=(1, 1))
+    keywords_list = [item[0] for item in keywords_text]
+
+    return keywords_list
+    
+# Gradio functions
+def turn_off_interactivity(user_message, history):
+        return gr.update(value="", interactive=False), history + [[user_message, None]]
+
+def restore_interactivity():
+        return gr.update(interactive=True)
+
+def update_message(dropdown_value):
+        return gr.Textbox.update(value=dropdown_value)
+
+def hide_block():
+        return gr.Radio.update(visible=False)

search_funcs/ingest.py

@@ -0,0 +1,417 @@
+# ---
+# jupyter:
+#   jupytext:
+#     formats: ipynb,py:light
+#     text_representation:
+#       extension: .py
+#       format_name: light
+#       format_version: '1.5'
+#       jupytext_version: 1.14.6
+#   kernelspec:
+#     display_name: Python 3 (ipykernel)
+#     language: python
+#     name: python3
+# ---
+
+# # Ingest website to FAISS
+
+# ## Install/ import stuff we need
+
+import os
+from pathlib import Path
+import re
+import pandas as pd
+from typing import TypeVar, List
+
+#from langchain.embeddings import HuggingFaceEmbeddings # HuggingFaceInstructEmbeddings, 
+from langchain.vectorstores.faiss import FAISS
+from langchain.vectorstores import Chroma
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.docstore.document import Document
+
+#from bs4 import BeautifulSoup
+#from docx import Document as Doc
+#from pypdf import PdfReader
+
+PandasDataFrame = TypeVar('pd.core.frame.DataFrame')
+# -
+
+split_strat = ["\n\n", "\n", ". ", "! ", "? "]
+chunk_size = 500
+chunk_overlap = 0
+start_index = True
+
+## Parse files
+def determine_file_type(file_path):
+        """
+        Determine the file type based on its extension.
+    
+        Parameters:
+            file_path (str): Path to the file.
+    
+        Returns:
+            str: File extension (e.g., '.pdf', '.docx', '.txt', '.html').
+        """
+        return os.path.splitext(file_path)[1].lower()
+
+def parse_file(file_paths, text_column='text'):
+    """
+    Accepts a list of file paths, determines each file's type based on its extension,
+    and passes it to the relevant parsing function.
+    
+    Parameters:
+        file_paths (list): List of file paths.
+        text_column (str): Name of the column in CSV/Excel files that contains the text content.
+    
+    Returns:
+        dict: A dictionary with file paths as keys and their parsed content (or error message) as values.
+    """
+    
+    
+
+    if not isinstance(file_paths, list):
+        raise ValueError("Expected a list of file paths.")
+    
+    extension_to_parser = {
+        # '.pdf': parse_pdf,
+        # '.docx': parse_docx,
+        # '.txt': parse_txt,
+        # '.html': parse_html,
+        # '.htm': parse_html,  # Considering both .html and .htm for HTML files
+        '.csv': lambda file_path: parse_csv_or_excel(file_path, text_column),
+        '.xlsx': lambda file_path: parse_csv_or_excel(file_path, text_column)
+    }
+    
+    parsed_contents = {}
+    file_names = []
+
+    for file_path in file_paths:
+        print(file_path.name)
+        #file = open(file_path.name, 'r')
+        #print(file)
+        file_extension = determine_file_type(file_path.name)
+        if file_extension in extension_to_parser:
+            parsed_contents[file_path.name] = extension_to_parser[file_extension](file_path.name)
+        else:
+            parsed_contents[file_path.name] = f"Unsupported file type: {file_extension}"
+
+        filename_end = get_file_path_end(file_path.name)
+
+        file_names.append(filename_end)
+    
+    return parsed_contents, file_names
+
+def text_regex_clean(text):
+    # Merge hyphenated words
+        text = re.sub(r"(\w+)-\n(\w+)", r"\1\2", text)
+        # If a double newline ends in a letter, add a full stop.
+        text = re.sub(r'(?<=[a-zA-Z])\n\n', '.\n\n', text)
+        # Fix newlines in the middle of sentences
+        text = re.sub(r"(?<!\n\s)\n(?!\s\n)", " ", text.strip())
+        # Remove multiple newlines
+        text = re.sub(r"\n\s*\n", "\n\n", text)
+        text = re.sub(r"  ", " ", text)
+        # Add full stops and new lines between words with no space between where the second one has a capital letter
+        text = re.sub(r'(?<=[a-z])(?=[A-Z])', '. \n\n', text)
+
+        return text
+
+def parse_csv_or_excel(file_path, text_column = "text"):
+        """
+        Read in a CSV or Excel file.
+        
+        Parameters:
+            file_path (str): Path to the CSV file.
+            text_column (str): Name of the column in the CSV file that contains the text content.
+        
+        Returns:
+            Pandas DataFrame: Dataframe output from file read
+        """
+
+        #out_df = pd.DataFrame()
+
+        #for file_path in file_paths:
+        file_extension = determine_file_type(file_path.name)
+        file_name = get_file_path_end(file_path.name)
+        file_names = [file_name]
+
+        if file_extension == ".csv":
+                df = pd.read_csv(file_path.name, low_memory=False)
+                if text_column not in df.columns: return pd.DataFrame(), ['Please choose a valid column name']
+                df['source'] = file_name
+                df['page_section'] = ""
+        elif file_extension == ".xlsx":
+                df = pd.read_excel(file_path.name, engine='openpyxl')
+                if text_column not in df.columns: return pd.DataFrame(), ['Please choose a valid column name']
+                df['source'] = file_name
+                df['page_section'] = ""
+        else:
+                print(f"Unsupported file type: {file_extension}")
+                return pd.DataFrame(), ['Please choose a valid file type']
+            
+        #    file_names.append(file_name)
+        #    out_df = pd.concat([out_df, df])
+        
+        #if text_column not in df.columns:
+        #    return f"Column '{text_column}' not found in {file_path}"
+        #text_out = " ".join(df[text_column].dropna().astype(str))
+        return df, file_names
+
+def parse_excel(file_path, text_column):
+        """
+        Read text from an Excel file.
+        
+        Parameters:
+            file_path (str): Path to the Excel file.
+            text_column (str): Name of the column in the Excel file that contains the text content.
+        
+        Returns:
+            Pandas DataFrame: Dataframe output from file read
+        """
+        df = pd.read_excel(file_path, engine='openpyxl')
+        #if text_column not in df.columns:
+        #    return f"Column '{text_column}' not found in {file_path}"
+        #text_out = " ".join(df[text_column].dropna().astype(str))
+        return df
+
+def get_file_path_end(file_path):
+    match = re.search(r'(.*[\/\\])?(.+)$', file_path)
+        
+    filename_end = match.group(2) if match else ''
+
+    return filename_end
+
+# +
+# Convert parsed text to docs
+# -
+
+def text_to_docs(text_dict: dict, chunk_size: int = chunk_size) -> List[Document]:
+    """
+    Converts the output of parse_file (a dictionary of file paths to content)
+    to a list of Documents with metadata.
+    """
+    
+    doc_sections = []
+    parent_doc_sections = []
+
+    for file_path, content in text_dict.items():
+        ext = os.path.splitext(file_path)[1].lower()
+
+        # Depending on the file extension, handle the content
+        # if ext == '.pdf':
+        #     docs, page_docs = pdf_text_to_docs(content, chunk_size)
+        # elif ext in ['.html', '.htm', '.txt', '.docx']:
+        #     docs = html_text_to_docs(content, chunk_size)
+        if ext in ['.csv', '.xlsx']:
+            docs, page_docs = csv_excel_text_to_docs(content, chunk_size)
+        else:
+            print(f"Unsupported file type {ext} for {file_path}. Skipping.")
+            continue
+
+        
+        filename_end = get_file_path_end(file_path)
+
+        #match = re.search(r'(.*[\/\\])?(.+)$', file_path)
+        #filename_end = match.group(2) if match else ''
+
+        # Add filename as metadata
+        for doc in docs: doc.metadata["source"] = filename_end
+        #for parent_doc in parent_docs: parent_doc.metadata["source"] = filename_end
+        
+        doc_sections.extend(docs)
+        #parent_doc_sections.extend(parent_docs)
+
+    return doc_sections#, page_docs
+
+
+def write_out_metadata_as_string(metadata_in):
+    # If metadata_in is a single dictionary, wrap it in a list
+    if isinstance(metadata_in, dict):
+        metadata_in = [metadata_in]
+
+    metadata_string = [f"{'  '.join(f'{k}: {v}' for k, v in d.items() if k != 'page_section')}" for d in metadata_in] # ['metadata']
+    return metadata_string
+
+def csv_excel_text_to_docs(df, text_column='text', chunk_size=None) -> List[Document]:
+    """Converts a DataFrame's content to a list of Documents with metadata."""
+    
+    doc_sections = []
+    df[text_column] = df[text_column].astype(str) # Ensure column is a string column
+
+    # For each row in the dataframe
+    for idx, row in df.iterrows():
+        # Extract the text content for the document
+        doc_content = row[text_column]
+        
+        # Generate metadata containing other columns' data
+        metadata = {"row": idx + 1}
+        for col, value in row.items():
+            if col != text_column:
+                metadata[col] = value
+
+        metadata_string = write_out_metadata_as_string(metadata)[0]
+
+        
+
+        # If chunk_size is provided, split the text into chunks
+        if chunk_size:
+            # Assuming you have a text splitter function similar to the PDF handling
+            text_splitter = RecursiveCharacterTextSplitter(
+                chunk_size=chunk_size,
+                # Other arguments as required by the splitter
+            )
+            sections = text_splitter.split_text(doc_content)
+
+            
+            # For each section, create a Document object
+            for i, section in enumerate(sections):
+                #section = '. '.join([metadata_string, section])
+                doc = Document(page_content=section, 
+                               metadata={**metadata, "section": i, "row_section": f"{metadata['row']}-{i}"})
+                doc_sections.append(doc)
+        else:
+            # If no chunk_size is provided, create a single Document object for the row
+            #doc_content = '. '.join([metadata_string, doc_content])
+            doc = Document(page_content=doc_content, metadata=metadata)
+            doc_sections.append(doc)
+    
+    return doc_sections
+
+# # Functions for working with documents after loading them back in
+
+def pull_out_data(series):
+
+    # define a lambda function to convert each string into a tuple
+    to_tuple = lambda x: eval(x)
+
+    # apply the lambda function to each element of the series
+    series_tup = series.apply(to_tuple)
+
+    series_tup_content = list(zip(*series_tup))[1]
+
+    series = pd.Series(list(series_tup_content))#.str.replace("^Main post content", "", regex=True).str.strip()
+
+    return series
+
+def docs_from_csv(df):
+
+    import ast
+    
+    documents = []
+    
+    page_content = pull_out_data(df["0"])
+    metadatas = pull_out_data(df["1"])
+
+    for x in range(0,len(df)):       
+        new_doc = Document(page_content=page_content[x], metadata=metadatas[x])
+        documents.append(new_doc)
+        
+    return documents
+
+def docs_from_lists(docs, metadatas):
+
+    documents = []
+
+    for x, doc in enumerate(docs):
+        new_doc = Document(page_content=doc, metadata=metadatas[x])
+        documents.append(new_doc)
+        
+    return documents
+
+def docs_elements_from_csv_save(docs_path="documents.csv"):
+
+    documents = pd.read_csv(docs_path)
+
+    docs_out = docs_from_csv(documents)
+
+    out_df = pd.DataFrame(docs_out)
+
+    docs_content = pull_out_data(out_df[0].astype(str))
+
+    docs_meta = pull_out_data(out_df[1].astype(str))
+
+    doc_sources = [d['source'] for d in docs_meta]
+
+    return out_df, docs_content, docs_meta, doc_sources
+
+# ## Create embeddings and save faiss vector store to the path specified in `save_to`
+
+def load_embeddings(model_name = "BAAI/bge-base-en-v1.5"):
+
+    #if model_name == "hkunlp/instructor-large":
+    #    embeddings_func = HuggingFaceInstructEmbeddings(model_name=model_name,
+    #    embed_instruction="Represent the paragraph for retrieval: ",
+    #    query_instruction="Represent the question for retrieving supporting documents: "
+    #    )
+
+    #else: 
+    embeddings_func = HuggingFaceEmbeddings(model_name=model_name)
+
+    global embeddings
+
+    embeddings = embeddings_func
+
+    return embeddings_func
+
+def embed_faiss_save_to_zip(docs_out, save_to="faiss_lambeth_census_embedding", model_name = "BAAI/bge-base-en-v1.5"):
+
+    load_embeddings(model_name=model_name)
+
+    #embeddings_fast = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+
+    print(f"> Total split documents: {len(docs_out)}")
+
+    vectorstore = FAISS.from_documents(documents=docs_out, embedding=embeddings)
+        
+
+    if Path(save_to).exists():
+        vectorstore.save_local(folder_path=save_to)
+
+    print("> DONE")
+    print(f"> Saved to: {save_to}")
+
+    ### Save as zip, then remove faiss/pkl files to allow for upload to huggingface
+
+    import shutil
+
+    shutil.make_archive(save_to, 'zip', save_to)
+
+    os.remove(save_to + "/index.faiss")
+    os.remove(save_to + "/index.pkl")
+
+    shutil.move(save_to + '.zip', save_to + "/" + save_to + '.zip')
+
+    return vectorstore
+
+def docs_to_chroma_save(embeddings, docs_out:PandasDataFrame, save_to:str):
+    print(f"> Total split documents: {len(docs_out)}")
+    
+    vectordb = Chroma.from_documents(documents=docs_out, 
+                                 embedding=embeddings,
+                                 persist_directory=save_to)
+    
+    # persiste the db to disk
+    vectordb.persist()
+    
+    print("> DONE")
+    print(f"> Saved to: {save_to}")
+    
+    return vectordb
+
+def sim_search_local_saved_vec(query, k_val, save_to="faiss_lambeth_census_embedding"):
+
+    load_embeddings()
+
+    docsearch = FAISS.load_local(folder_path=save_to, embeddings=embeddings)
+
+
+    display(Markdown(question))
+
+    search = docsearch.similarity_search_with_score(query, k=k_val)
+
+    for item in search:
+        print(item[0].page_content)
+        print(f"Page: {item[0].metadata['source']}")
+        print(f"Date: {item[0].metadata['date']}")
+        print(f"Score: {item[1]}")
+        print("---")