General code improvements and refinements.

Dockerfile

@@ -58,7 +58,5 @@ WORKDIR $HOME/app
 
 # Copy the current directory contents into the container at $HOME/app setting the owner to the user
 COPY --chown=user . $HOME/app
-#COPY . $HOME/app
-
 
 CMD ["python", "app.py"]

app.py

@@ -7,7 +7,7 @@ PandasDataFrame = Type[pd.DataFrame]
 
 from search_funcs.bm25_functions import prepare_bm25_input_data, prepare_bm25, bm25_search
 from search_funcs.semantic_ingest_functions import csv_excel_text_to_docs
-from search_funcs.semantic_functions import docs_to_bge_embed_np_array, bge_simple_retrieval
+from search_funcs.semantic_functions import docs_to_bge_embed_np_array, bge_semantic_search
 from search_funcs.helper_functions import display_info, initial_data_load, put_columns_in_join_df, get_temp_folder_path, empty_folder, get_connection_params, output_folder
 from search_funcs.spacy_search_funcs import spacy_fuzzy_search
 from search_funcs.aws_functions import load_data_from_aws
@@ -17,39 +17,33 @@ temp_folder_path = get_temp_folder_path()
 empty_folder(temp_folder_path)
 
 ## Gradio app - BM25 search
-block = gr.Blocks(theme = gr.themes.Base()) # , css="theme.css"
+app = gr.Blocks(theme = gr.themes.Base()) # , css="theme.css"
 
-
-with block:
+with app:
     print("Please don't close this window! Open the below link in the web browser of your choice.")
 
-    ingest_text = gr.State()
-    ingest_metadata = gr.State()
-    ingest_docs = gr.State()
-    vectorstore_state = gr.State() # globals()["vectorstore"]
-    embeddings_state = gr.State(np.array([])) # globals()["embeddings"]
-    search_index_state = gr.State()
-    tokenised_state = gr.State()
-    bm25_search_object_state = gr.State()
-
-    k_val = gr.State(9999)
-    out_passages = gr.State(9999)
-    vec_weight = gr.State(1)
-
-    corpus_state = gr.State()
-    keyword_data_list_state = gr.State([])
-    join_data_state = gr.State(pd.DataFrame())
-    output_file_state = gr.State([])
-
-    orig_keyword_data_state = gr.State(pd.DataFrame())
-    keyword_data_state = gr.State(pd.DataFrame())
-   
-    orig_semantic_data_state = gr.State(pd.DataFrame())
-    semantic_data_state = gr.State(pd.DataFrame())
+    # BM25 state objects
+    orig_keyword_data_state = gr.State(pd.DataFrame()) # Original data that is not changed #gr.Dataframe(pd.DataFrame(),visible=False) #gr.State(pd.DataFrame())
+    prepared_keyword_data_state = gr.State(pd.DataFrame()) # Data frame the contains modified data #gr.Dataframe(pd.DataFrame(),visible=False) #gr.State(pd.DataFrame())
+    #tokenised_prepared_keyword_data_state = gr.State([]) # This is data that has been loaded in as tokens #gr.Dataframe(pd.DataFrame(),visible=False) #gr.State()
+    tokenised_prepared_keyword_data_state = gr.State([]) # Data that has been prepared for search (tokenised) #gr.Dataframe(np.array([]), type="array", visible=False) #gr.State([])
+    bm25_search_index_state = gr.State()   
+
 
+    # Semantic search state objects
+    orig_semantic_data_state = gr.State(pd.DataFrame()) #gr.Dataframe(pd.DataFrame(),visible=False) # gr.State(pd.DataFrame())
+    semantic_data_state = gr.State(pd.DataFrame()) #gr.Dataframe(pd.DataFrame(),visible=False) # gr.State(pd.DataFrame())
+    semantic_input_document_format = gr.State([])
+    embeddings_state = gr.State(np.array([])) #gr.Dataframe(np.array([]), type="numpy", visible=False) #gr.State(np.array([])) # globals()["embeddings"]
+    semantic_k_val = gr.Number(9999, visible=False)
+
+    # State objects for app in general
     session_hash_state = gr.State("")
     s3_output_folder_state = gr.State("")
+    join_data_state = gr.State(pd.DataFrame()) #gr.Dataframe(pd.DataFrame(), visible=False) #gr.State(pd.DataFrame())
+    output_file_state = gr.Dropdown([], visible=False, allow_custom_value=True) #gr.Dataframe(type="array", visible=False) #gr.State([])
 
+    # Informational state objects
     in_k1_info = gr.State("""k1: Constant used for influencing the term frequency saturation. After saturation is reached, additional
 presence for the term adds a significantly less additional score. According to [1]_, experiments suggest
 that 1.2 < k1 < 2 yields reasonably good results, although the optimal value depends on factors such as
@@ -167,7 +161,7 @@ depends on factors such as the type of documents or queries. Information taken f
             out_aws_data_message = gr.Textbox(label="AWS data load progress")
 
     # Changing search parameters button
-    in_search_param_button.click(fn=prepare_bm25, inputs=[corpus_state, in_bm25_file, in_bm25_column, search_index_state, return_intermediate_files, in_k1, in_b, in_alpha], outputs=[load_finished_message])
+    in_search_param_button.click(fn=prepare_bm25, inputs=[tokenised_prepared_keyword_data_state, in_bm25_file, in_bm25_column, bm25_search_index_state, return_intermediate_files, in_k1, in_b, in_alpha], outputs=[load_finished_message])
                
     # ---
     in_k1_button.click(display_info, inputs=in_k1_info)
@@ -178,43 +172,41 @@ depends on factors such as the type of documents or queries. Information taken f
     ### Loading AWS data ###
     load_aws_keyword_data_button.click(fn=load_data_from_aws, inputs=[in_aws_keyword_file, aws_password_box], outputs=[in_bm25_file, out_aws_data_message])
     load_aws_semantic_data_button.click(fn=load_data_from_aws, inputs=[in_aws_semantic_file, aws_password_box], outputs=[in_semantic_file, out_aws_data_message])
-          
     
     ### BM25 SEARCH ###
     # Update dropdowns upon initial file load
-    in_bm25_file.change(initial_data_load, inputs=[in_bm25_file], outputs=[in_bm25_column, search_df_join_column, keyword_data_state, orig_keyword_data_state, search_index_state, embeddings_state, tokenised_state, load_finished_message, current_source], api_name="initial_load")
+    in_bm25_file.change(initial_data_load, inputs=[in_bm25_file], outputs=[in_bm25_column, search_df_join_column, prepared_keyword_data_state, orig_keyword_data_state, bm25_search_index_state, embeddings_state, tokenised_prepared_keyword_data_state, load_finished_message, current_source], api_name="initial_load")
     in_join_file.change(put_columns_in_join_df, inputs=[in_join_file], outputs=[in_join_column, join_data_state, in_join_message])
  
     # Load in BM25 data
-    load_bm25_data_button.click(fn=prepare_bm25_input_data, inputs=[in_bm25_file, in_bm25_column, keyword_data_state, tokenised_state, in_clean_data, return_intermediate_files], outputs=[corpus_state, load_finished_message, keyword_data_state, output_file, output_file, keyword_data_list_state, in_bm25_column], api_name="load_keyword").\
-    then(fn=prepare_bm25, inputs=[corpus_state, in_bm25_file, in_bm25_column, search_index_state, in_clean_data, return_intermediate_files, in_k1, in_b, in_alpha], outputs=[load_finished_message, output_file, bm25_search_object_state], api_name="prepare_keyword")
-    
+    load_bm25_data_button.click(fn=prepare_bm25_input_data, inputs=[in_bm25_file, in_bm25_column, prepared_keyword_data_state, tokenised_prepared_keyword_data_state, in_clean_data, return_intermediate_files], outputs=[tokenised_prepared_keyword_data_state, load_finished_message, prepared_keyword_data_state, output_file, output_file, in_bm25_column], api_name="load_keyword").\
+    then(fn=prepare_bm25, inputs=[tokenised_prepared_keyword_data_state, in_bm25_file, in_bm25_column, bm25_search_index_state, in_clean_data, return_intermediate_files, in_k1, in_b, in_alpha], outputs=[load_finished_message, output_file, bm25_search_index_state, tokenised_prepared_keyword_data_state], api_name="prepare_keyword") # keyword_data_list_state
     
     # BM25 search functions on click or enter
-    keyword_search_button.click(fn=bm25_search, inputs=[keyword_query, in_no_search_results, orig_keyword_data_state, keyword_data_state, in_bm25_column, join_data_state, in_clean_data, bm25_search_object_state, in_join_column, search_df_join_column], outputs=[output_single_text, output_file], api_name="keyword_search")
-    keyword_query.submit(fn=bm25_search, inputs=[keyword_query, in_no_search_results, orig_keyword_data_state, keyword_data_state, in_bm25_column, join_data_state, in_clean_data, bm25_search_object_state, in_join_column, search_df_join_column], outputs=[output_single_text, output_file])
+    keyword_search_button.click(fn=bm25_search, inputs=[keyword_query, in_no_search_results, orig_keyword_data_state, prepared_keyword_data_state, in_bm25_column, join_data_state, in_clean_data, bm25_search_index_state, tokenised_prepared_keyword_data_state, in_join_column, search_df_join_column, in_k1, in_b, in_alpha], outputs=[output_single_text, output_file], api_name="keyword_search")
+    keyword_query.submit(fn=bm25_search, inputs=[keyword_query, in_no_search_results, orig_keyword_data_state, prepared_keyword_data_state, in_bm25_column, join_data_state, in_clean_data, bm25_search_index_state, tokenised_prepared_keyword_data_state, in_join_column, search_df_join_column, in_k1, in_b, in_alpha], outputs=[output_single_text, output_file])
 
     # Fuzzy search functions on click
-    fuzzy_search_button.click(fn=spacy_fuzzy_search, inputs=[keyword_query, keyword_data_list_state, keyword_data_state, in_bm25_column, join_data_state, search_df_join_column, in_join_column, no_spelling_mistakes], outputs=[output_single_text, output_file], api_name="fuzzy_search")
+    fuzzy_search_button.click(fn=spacy_fuzzy_search, inputs=[keyword_query, tokenised_prepared_keyword_data_state, prepared_keyword_data_state, in_bm25_column, join_data_state, search_df_join_column, in_join_column, no_spelling_mistakes], outputs=[output_single_text, output_file], api_name="fuzzy_search")
     
     ### SEMANTIC SEARCH ###
 
     # Load in a csv/excel file for semantic search
-    in_semantic_file.change(initial_data_load, inputs=[in_semantic_file], outputs=[in_semantic_column,  search_df_join_column,  semantic_data_state, orig_semantic_data_state, search_index_state, embeddings_state, tokenised_state, semantic_load_progress, current_source_semantic])
+    in_semantic_file.change(initial_data_load, inputs=[in_semantic_file], outputs=[in_semantic_column,  search_df_join_column,  semantic_data_state, orig_semantic_data_state, bm25_search_index_state, embeddings_state, tokenised_prepared_keyword_data_state, semantic_load_progress, current_source_semantic])
     load_semantic_data_button.click(
-        csv_excel_text_to_docs, inputs=[semantic_data_state, in_semantic_file, in_semantic_column, in_clean_data, return_intermediate_files], outputs=[ingest_docs, semantic_load_progress, output_file_state]).\
-        then(docs_to_bge_embed_np_array, inputs=[ingest_docs, in_semantic_file, embeddings_state, output_file_state, in_clean_data, return_intermediate_files, embedding_super_compress], outputs=[semantic_load_progress, vectorstore_state, semantic_output_file, output_file_state])
+        csv_excel_text_to_docs, inputs=[semantic_data_state, in_semantic_file, in_semantic_column, in_clean_data, return_intermediate_files], outputs=[semantic_input_document_format, semantic_load_progress, output_file_state]).\
+        then(docs_to_bge_embed_np_array, inputs=[semantic_input_document_format, in_semantic_file, embeddings_state, output_file_state, in_clean_data, return_intermediate_files, embedding_super_compress], outputs=[semantic_load_progress, embeddings_state, semantic_output_file, output_file_state]) #  vectorstore_state
     
     # Semantic search query
-    semantic_submit.click(bge_simple_retrieval, inputs=[semantic_query, vectorstore_state, ingest_docs, in_semantic_column, k_val, out_passages, semantic_min_distance, vec_weight, join_data_state, in_join_column, search_df_join_column], outputs=[semantic_output_single_text, semantic_output_file], api_name="semantic_search")
-    semantic_query.submit(bge_simple_retrieval, inputs=[semantic_query, vectorstore_state, ingest_docs, in_semantic_column, k_val, out_passages, semantic_min_distance, vec_weight, join_data_state, in_join_column, search_df_join_column], outputs=[semantic_output_single_text, semantic_output_file])
+    semantic_submit.click(bge_semantic_search, inputs=[semantic_query, embeddings_state, semantic_input_document_format, semantic_k_val, semantic_min_distance, join_data_state, in_join_column, search_df_join_column], outputs=[semantic_output_single_text, semantic_output_file], api_name="semantic_search")
+    semantic_query.submit(bge_semantic_search, inputs=[semantic_query, embeddings_state, semantic_input_document_format, semantic_k_val, semantic_min_distance, join_data_state, in_join_column, search_df_join_column], outputs=[semantic_output_single_text, semantic_output_file])
 
-    block.load(get_connection_params, inputs=None, outputs=[session_hash_state, s3_output_folder_state])
+    app.load(get_connection_params, inputs=None, outputs=[session_hash_state, s3_output_folder_state])
 
 # Launch the Gradio app
 if __name__ == "__main__":
-    block.queue().launch(show_error=True) # root_path="/data-text-search" # server_name="0.0.0.0", 
+    app.queue().launch(show_error=True) # root_path="/data-text-search" # server_name="0.0.0.0", 
     
 # Running on local server with https: https://discuss.huggingface.co/t/how-to-run-gradio-with-0-0-0-0-and-https/38003 or https://dev.to/rajshirolkar/fastapi-over-https-for-development-on-windows-2p7d # Need to download OpenSSL and create own keys 
-# block.queue().launch(ssl_verify=False, share=False, debug=False, server_name="0.0.0.0",server_port=443,
+# app.queue().launch(ssl_verify=False, share=False, debug=False, server_name="0.0.0.0",server_port=443,
 #                      ssl_certfile="cert.pem", ssl_keyfile="key.pem") # port 443 for https. Certificates currently not valid

requirements.txt

@@ -1,12 +1,11 @@
 pandas==2.2.2
 polars==0.20.3
 pyarrow==14.0.2
-openpyxl==3.1.2
+openpyxl==3.1.3
 torch==2.3.1
-transformers==4.41.2
 spacy
 en_core_web_sm @ https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.7.1/en_core_web_sm-3.7.1.tar.gz
 gradio
 sentence_transformers==3.0.1
-lxml==5.1.0
+lxml==5.2.2
 boto3==1.34.103

requirements_gpu.txt

@@ -1,11 +1,11 @@
 pandas==2.2.2
 polars==0.20.3
 pyarrow==14.0.2
-openpyxl==3.1.2
+openpyxl==3.1.3
 torch==2.3.1 --index-url https://download.pytorch.org/whl/cu121
 spacy
 en_core_web_sm @ https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.7.1/en_core_web_sm-3.7.1.tar.gz
 gradio
-sentence_transformers==2.3.1
-lxml==5.1.0
+sentence_transformers==3.0.1
+lxml==5.2.2
 boto3==1.34.103

search_funcs/bm25_functions.py

@@ -8,6 +8,7 @@ import time
 import pandas as pd
 from numpy import inf
 import gradio as gr
+from typing import List
 
 from datetime import datetime
 
@@ -165,7 +166,7 @@ class BM25:
 		return [documents[i] for i in heapq.nlargest(n, scores.keys(), key=scores.__getitem__)]
 	
 
-	def get_top_n_with_score(self, query, documents, n=5):
+	def get_top_n_with_score(self, query:str, documents:List[str], n=5):
 		"""
 		Retrieve the top n documents for the query along with their scores.
 
@@ -229,15 +230,47 @@ class BM25:
 		with open(f"{output_folder}{filename}.pkl", "rb") as fsave:
 			return pickle.load(fsave)
 
-# These following functions are my own work
+def prepare_bm25_input_data(
+    in_file: list, 
+    text_column: str, 
+    data_state: pd.DataFrame, 
+    tokenised_state: list, 
+    clean: str = "No",  
+    return_intermediate_files: str = "No", 
+    progress: gr.Progress = gr.Progress(track_tqdm=True)
+) -> tuple:
+	"""
+	Prepare BM25 input data by loading, cleaning, and tokenizing the text data.
+
+	Parameters
+	----------
+	in_file: list
+		List of input files to be processed.
+	text_column: str
+		The name of the text column in the data file to search.
+	data_state: pd.DataFrame
+		The current state of the data.
+	tokenised_state: list
+		The current state of the tokenized data.
+	clean: str, optional
+		Whether to clean the text data (default is "No").
+	return_intermediate_files: str, optional
+		Whether to return intermediate processing files (default is "No").
+	progress: gr.Progress, optional
+		Progress tracker for the function (default is gr.Progress(track_tqdm=True)).
+
+	Returns
+	-------
+	tuple
+		A tuple containing the prepared search text list, a message, the updated data state, 
+		the tokenized data, the search index, and a dropdown component for the text column.
+	"""
 
-def prepare_bm25_input_data(in_file, text_column, data_state, tokenised_state, clean="No",  return_intermediate_files = "No", progress=gr.Progress(track_tqdm=True)):
-	#print(in_file)
 	ensure_output_folder_exists(output_folder)
 
 	if not in_file:
 		print("No input file found. Please load in at least one file.")
-		return None, "No input file found. Please load in at least one file.", data_state, None, None, [], gr.Dropdown(allow_custom_value=True, value=text_column, choices=data_state.columns.to_list())
+		return None, "No input file found. Please load in at least one file.", data_state, None, None, gr.Dropdown(allow_custom_value=True, value=text_column, choices=data_state.columns.to_list())
 
 	progress(0, desc = "Loading in data")
 	file_list = [string.name for string in in_file]
@@ -247,25 +280,24 @@ def prepare_bm25_input_data(in_file, text_column, data_state, tokenised_state, c
 	data_file_names = [string for string in file_list if "tokenised" not in string.lower() and "npz" not in string.lower() and "gz" not in string.lower()]
 
 	if not data_file_names:
-		return None, "Please load in at least one csv/Excel/parquet data file.", data_state, None, None, [], gr.Dropdown(allow_custom_value=True, value=text_column, choices=data_state.columns.to_list())
+		return None, "Please load in at least one csv/Excel/parquet data file.", data_state, None, None, gr.Dropdown(allow_custom_value=True, value=text_column, choices=data_state.columns.to_list())
 
 	if not text_column:
-		return None, "Please enter a column name to search.", data_state, None, None,[], gr.Dropdown(allow_custom_value=True, value=text_column, choices=data_state.columns.to_list())
+		return None, "Please enter a column name to search.", data_state, None, None, gr.Dropdown(allow_custom_value=True, value=text_column, choices=data_state.columns.to_list())
 
 	data_file_name = data_file_names[0]
 
 	df = data_state #read_file(data_file_name)
-	data_file_out_name = get_file_path_end_with_ext(data_file_name)
+	#data_file_out_name = get_file_path_end_with_ext(data_file_name)
 	data_file_out_name_no_ext = get_file_path_end(data_file_name)
 
-	## Load in pre-tokenised corpus if exists
-	tokenised_df = pd.DataFrame()
+	## Load in pre-tokenised prepared_search_text_list if exists
+	#tokenised_df = pd.DataFrame()
 
-	tokenised_file_names = [string for string in file_list if "tokenised" in string.lower()]
+	#tokenised_file_names = [string for string in file_list if "tokenised" in string.lower()]
 	search_index_file_names = [string for string in file_list if "gz" in string.lower()]
 
-	print("Dataframe columns:", df.columns)
-
+	# Set all search text to lower case
 	df[text_column] = df[text_column].astype(str).str.lower()
 
 	if "copy_of_case_note_id" in df.columns:
@@ -273,10 +305,10 @@ def prepare_bm25_input_data(in_file, text_column, data_state, tokenised_state, c
 		df.loc[~df["copy_of_case_note_id"].isna(), text_column] = ""
 
 	if search_index_file_names:
-		corpus = list(df[text_column])
+		prepared_search_text_list = list(df[text_column])
 		message = "Tokenisation skipped - loading search index from file."
 		print(message)
-		return corpus, message, df, None, None, [], gr.Dropdown(allow_custom_value=True, value=text_column, choices=data_state.columns.to_list())
+		return prepared_search_text_list, message, df, None, None, gr.Dropdown(allow_custom_value=True, value=text_column, choices=data_state.columns.to_list())
 
 	
 	if clean == "Yes":
@@ -285,11 +317,11 @@ def prepare_bm25_input_data(in_file, text_column, data_state, tokenised_state, c
 		print("Starting data clean.")
 
 		#df = df.drop_duplicates(text_column)
-		df_list = list(df[text_column])
-		df_list = initial_clean(df_list)
+		prepared_text_as_list = list(df[text_column])
+		prepared_text_as_list = initial_clean(prepared_text_as_list)
 
 		# Save to file if you have cleaned the data
-		out_file_name, text_column, df  = save_prepared_bm25_data(data_file_name, df_list, df, text_column)
+		out_file_name, text_column, df  = save_prepared_bm25_data(data_file_name, prepared_text_as_list, df, text_column)
 	
 		clean_toc = time.perf_counter()
 		clean_time_out = f"Cleaning the text took {clean_toc - clean_tic:0.1f} seconds."
@@ -297,7 +329,7 @@ def prepare_bm25_input_data(in_file, text_column, data_state, tokenised_state, c
 
 	else:
 		# Don't clean or save file to disk
-		df_list = list(df[text_column])
+		prepared_text_as_list = list(df[text_column])
 		print("No data cleaning performed")
 		out_file_name = None
 		
@@ -305,24 +337,27 @@ def prepare_bm25_input_data(in_file, text_column, data_state, tokenised_state, c
 	
 	progress(0.4, desc = "Tokenising text")
 
+	print("Tokenised state:", tokenised_state)
+
 	if tokenised_state:
-		tokenised_df = tokenised_state
-		corpus = tokenised_df.iloc[:,0].tolist()
+		prepared_search_text_list = tokenised_state.iloc[:,0].tolist()
 		print("Tokenised data loaded from file")
-		#print("Corpus is: ", corpus[0:5])
+		
+		#print("prepared_search_text_list is: ", prepared_search_text_list[0:5])
 
 	else:
 		tokeniser_tic = time.perf_counter()
-		corpus = []
+		prepared_search_text_list = []
 		batch_size = 256
-		for doc in tokenizer.pipe(progress.tqdm(df_list, desc = "Tokenising text", unit = "rows"), batch_size=batch_size):
-			corpus.append([token.text for token in doc])
+		for doc in tokenizer.pipe(progress.tqdm(prepared_text_as_list, desc = "Tokenising text", unit = "rows"), batch_size=batch_size):
+			prepared_search_text_list.append([token.text for token in doc])
 
 		tokeniser_toc = time.perf_counter()
 		tokenizer_time_out = f"Tokenising the text took {tokeniser_toc - tokeniser_tic:0.1f} seconds."
 		print(tokenizer_time_out)
+		#print("prepared_search_text_list is: ", prepared_search_text_list[0:5])
 
-	if len(df_list) >= 20:
+	if len(prepared_text_as_list) >= 20:
 		message = "Data loaded"
 	else:
 		message = "Data loaded. Warning: dataset may be too short to get consistent search results."
@@ -334,13 +369,29 @@ def prepare_bm25_input_data(in_file, text_column, data_state, tokenised_state, c
 		else:
 			tokenised_data_file_name = output_folder + data_file_out_name_no_ext + "_tokenised.parquet"
 
-		pd.DataFrame(data={"Corpus":corpus}).to_parquet(tokenised_data_file_name)
+		pd.DataFrame(data={"prepared_search_text_list":prepared_search_text_list}).to_parquet(tokenised_data_file_name)
+
+		return prepared_search_text_list, message, df, out_file_name, tokenised_data_file_name, gr.Dropdown(allow_custom_value=True, value=text_column, choices=data_state.columns.to_list()) # prepared_text_as_list, 
+
+	return prepared_search_text_list, message, df, out_file_name, None, gr.Dropdown(allow_custom_value=True, value=text_column, choices=data_state.columns.to_list()) # prepared_text_as_list, 
 
-		return corpus, message, df, out_file_name, tokenised_data_file_name, df_list, gr.Dropdown(allow_custom_value=True, value=text_column, choices=data_state.columns.to_list())
+def save_prepared_bm25_data(in_file_name: str, prepared_text_list: list, in_df: pd.DataFrame, in_bm25_column: str, progress: gr.Progress = gr.Progress(track_tqdm=True)) -> tuple:
+	"""
+	Save the prepared BM25 data to a file.
+
+	This function ensures the output folder exists, checks if the length of the prepared text list matches the input dataframe,
+	and saves the prepared data to a file in the specified format. The original column in the input dataframe is dropped to reduce file size.
 
-	return corpus, message, df, out_file_name, None, df_list, gr.Dropdown(allow_custom_value=True, value=text_column, choices=data_state.columns.to_list())
+	Parameters:
+	- in_file_name (str): The name of the input file.
+	- prepared_text_list (list): The list of prepared text.
+	- in_df (pd.DataFrame): The input dataframe.
+	- in_bm25_column (str): The name of the column to be processed.
+	- progress (gr.Progress, optional): The progress tracker for the operation.
 
-def save_prepared_bm25_data(in_file_name, prepared_text_list, in_df, in_bm25_column, progress=gr.Progress(track_tqdm=True)):
+	Returns:
+	- tuple: A tuple containing the file name, new text column name, and the prepared dataframe.
+	"""
 
 	ensure_output_folder_exists(output_folder)
 
@@ -368,26 +419,54 @@ def save_prepared_bm25_data(in_file_name, prepared_text_list, in_df, in_bm25_col
 
 	return file_name, new_text_column, prepared_df
 
-def prepare_bm25(corpus, in_file, text_column, search_index, clean, return_intermediate_files, k1=1.5, b = 0.75, alpha=-5, progress=gr.Progress(track_tqdm=True)):
-	#bm25.save("saved_df_bm25")
-	#bm25 = BM25.load(re.sub(r'\.pkl$', '', file_in.name))
-
-	
+def prepare_bm25(
+    prepared_search_text_list: List[str], 
+    in_file: List[gr.File], 
+    text_column: str, 
+    search_index: BM25, 
+    clean: str, 
+    return_intermediate_files: str, 
+    k1: float = 1.5, 
+    b: float = 0.75, 
+    alpha: float = -5, 
+    progress: gr.Progress = gr.Progress(track_tqdm=True)
+) -> tuple:
+	"""
+	Prepare the BM25 search index.
+
+	This function prepares the BM25 search index from the provided text list and input file. It ensures the necessary
+	files and columns are present, processes the data, and optionally saves intermediate files.
+
+	Parameters:
+	- prepared_search_text_list (List[str]): The list of prepared search text.
+	- in_file (List[gr.File]): The list of input files.
+	- text_column (str): The name of the column to search.
+	- search_index (BM25): The BM25 search index.
+	- clean (str): Indicates whether to clean the data.
+	- return_intermediate_files (str): Indicates whether to return intermediate files.
+	- k1 (float, optional): The k1 parameter for BM25. Default is 1.5.
+	- b (float, optional): The b parameter for BM25. Default is 0.75.
+	- alpha (float, optional): The alpha parameter for BM25. Default is -5.
+	- progress (gr.Progress, optional): The progress tracker for the operation.
+
+	Returns:
+	- tuple: A tuple containing the output message, BM25 search index, and other relevant information.
+	"""
 
 	if not in_file:
 		out_message ="No input file found. Please load in at least one file."
 		print(out_message)
-		return out_message, None
+		return out_message, None, None
 
-	if not corpus:
+	if not prepared_search_text_list:
 		out_message = "No data file found. Please load in at least one csv/Excel/Parquet file."
 		print(out_message)
-		return out_message, None
+		return out_message, None, None, None
 
 	if not text_column:
 		out_message = "Please enter a column name to search."
 		print(out_message)
-		return  out_message, None
+		return  out_message, None, None, None
 
 	file_list = [string.name for string in in_file]
 
@@ -397,36 +476,23 @@ def prepare_bm25(corpus, in_file, text_column, search_index, clean, return_inter
 	data_file_names = [string for string in file_list if "tokenised" not in string.lower() and "npz" not in string.lower() and "gz" not in string.lower()]
 
 	if not data_file_names:
-		return "Please load in at least one csv/Excel/parquet data file.", None
+		return "Please load in at least one csv/Excel/parquet data file.", None, None, None
 
 	data_file_name = data_file_names[0]
 	data_file_out_name = get_file_path_end_with_ext(data_file_name)
 	data_file_name_no_ext = get_file_path_end(data_file_name)
 
-	# Check if there is a search index file already
-	#index_file_names = [string for string in file_list if "gz" in string.lower()]
-
 	progress(0.6, desc = "Preparing search index")
 
-	#if index_file_names:
 	if search_index:
-		#index_file_name = index_file_names[0]
-
-		#print(index_file_name)
-
-		bm25_load = search_index
-		
-
-		#index_file_out_name = get_file_path_end_with_ext(index_file_name)
-		#index_file_name_no_ext = get_file_path_end(index_file_name)
-
+		bm25 = search_index
 	else:
-		print("Preparing BM25 corpus")
+		print("Preparing BM25 search corpus")
 
-		bm25_load = BM25(corpus, k1=k1, b=b, alpha=alpha)
+		bm25 = BM25(prepared_search_text_list, k1=k1, b=b, alpha=alpha)
 
-	global bm25
-	bm25 = bm25_load
+	#global bm25
+	#bm25 = bm25_load
 
 	if return_intermediate_files == "Yes":
 		print("Saving search index file")
@@ -451,7 +517,7 @@ def prepare_bm25(corpus, in_file, text_column, search_index, clean, return_inter
 
 	print(message)
 
-	return message, None, bm25
+	return message, None, bm25, prepared_search_text_list
 
 def convert_bm25_query_to_tokens(free_text_query, clean="No"):
     '''
@@ -474,9 +540,75 @@ def convert_bm25_query_to_tokens(free_text_query, clean="No"):
 
     return out_query
 
-def bm25_search(free_text_query, in_no_search_results, original_data, searched_data, text_column, in_join_file, clean, bm25, in_join_column = "", search_df_join_column = "", progress=gr.Progress(track_tqdm=True)):   
+def bm25_search(
+    free_text_query: str, 
+    in_no_search_results: int, 
+    original_data: pd.DataFrame, 
+    searched_data: pd.DataFrame, 
+    text_column: str, 
+    in_join_file: str, 
+    clean: str, 
+    bm25: BM25, 
+    prepared_search_text_list_state: list, 
+    in_join_column: str = "", 
+    search_df_join_column: str = "", 
+    k1: float = 1.5, 
+    b: float = 0.75, 
+    alpha: float = -5, 
+    progress: gr.Progress = gr.Progress(track_tqdm=True)
+) -> tuple:
+	"""
+	Perform a BM25 search on the provided text data.
+
+	Parameters
+	----------
+	free_text_query : str
+		The query text to search for.
+	in_no_search_results : int
+		The number of search results to return.
+	original_data : pd.DataFrame
+		The original data containing the text to be searched.
+	searched_data : pd.DataFrame
+		The data that has been prepared for searching.
+	text_column : str
+		The name of the column in the data to search.
+	in_join_file : str
+		The file to join the search results with.
+	clean : str
+		Whether to clean the text data.
+	bm25 : BM25
+		The BM25 object used for searching.
+	prepared_search_text_list_state : list
+		The state of the prepared search text list.
+	in_join_column : str, optional
+		The column to join on in the input file (default is "").
+	search_df_join_column : str, optional
+		The column to join on in the search dataframe (default is "").
+	k1 : float, optional
+		The k1 parameter for BM25 (default is 1.5).
+	b : float, optional
+		The b parameter for BM25 (default is 0.75).
+	alpha : float, optional
+		The alpha parameter for BM25 (default is -5).
+	progress : gr.Progress, optional
+		Progress tracker for the function (default is gr.Progress(track_tqdm=True)).
+
+	Returns
+	-------
+	tuple
+		A tuple containing a message, the search results file name (if any), the BM25 object, and the prepared search text list.
+	"""
 
 	progress(0, desc = "Conducting keyword search")
+
+	print("in_join_file at start of bm25_search:", in_join_file)
+
+	if not bm25:
+		print("Preparing BM25 search corpus")
+
+		bm25 = BM25(prepared_search_text_list_state, k1=k1, b=b, alpha=alpha)
+
+	# print("bm25:", bm25)
 	
 	# Prepare query
 	if (clean == "Yes") | (text_column.endswith("_cleaned")):
@@ -484,8 +616,6 @@ def bm25_search(free_text_query, in_no_search_results, original_data, searched_d
 	else:
 		token_query = convert_bm25_query_to_tokens(free_text_query, clean="No")
 
-	#print(token_query)
-
 	# Perform search
 	print("Searching")
 
@@ -504,7 +634,6 @@ def bm25_search(free_text_query, in_no_search_results, original_data, searched_d
 
 	# Join scores onto searched data
 	results_df_out = results_df[['index', 'search_text', 'search_score_abs']].merge(searched_data,left_on="index", right_index=True, how="left", suffixes = ("", "_y")).drop("index_y", axis=1, errors="ignore")
-
 	
 
 	# Join on data from duplicate case notes
@@ -516,33 +645,27 @@ def bm25_search(free_text_query, in_no_search_results, original_data, searched_d
 			print("Clean is yes")
 			orig_text_column = text_column.replace("_cleaned", "")
 
-		#print(orig_text_column)
-		#print(original_data.columns)
-
 		original_data["original_note_id"] = original_data["copy_of_case_note_id"]
 		original_data["original_note_id"] = original_data["original_note_id"].combine_first(original_data["note_id"])
 
 		results_df_out = results_df_out.merge(original_data[["original_note_id", "note_id", "copy_of_case_note_id", "person_id"]],left_on="note_id", right_on="original_note_id", how="left", suffixes=("_primary", "")) # .drop(orig_text_column, axis = 1)
 		results_df_out.loc[~results_df_out["copy_of_case_note_id"].isnull(), "search_text"] = ""
 		results_df_out.loc[~results_df_out["copy_of_case_note_id"].isnull(), text_column] = ""
-
-		#results_df_out = pd.concat([results_df_out, original_data[~original_data["copy_of_case_note_id"].isna()][["copy_of_case_note_id", "person_id"]]])
-		# Replace NaN with an empty string
-		# results_df_out.fillna('', inplace=True)
-		
-		
 	
+	print("in_join_file:", in_join_file)
+
 	# Join on additional files
 	if not in_join_file.empty:
 		progress(0.5, desc = "Joining on additional data file")
-		join_df = in_join_file
-		join_df[in_join_column] = join_df[in_join_column].astype(str).str.replace("\.0$","", regex=True)
+		#join_df = in_join_file
+		# Prepare join columns as string and remove .0 at end of stringified numbers
+		in_join_file[in_join_column] = in_join_file[in_join_column].astype(str).str.replace("\.0$","", regex=True)
 		results_df_out[search_df_join_column] = results_df_out[search_df_join_column].astype(str).str.replace("\.0$","", regex=True)
 
 		# Duplicates dropped so as not to expand out dataframe
-		join_df = join_df.drop_duplicates(in_join_column)
+		in_join_file = in_join_file.drop_duplicates(in_join_column)
 
-		results_df_out = results_df_out.merge(join_df,left_on=search_df_join_column, right_on=in_join_column, how="left", suffixes=('','_y'))#.drop(in_join_column, axis=1)
+		results_df_out = results_df_out.merge(in_join_file,left_on=search_df_join_column, right_on=in_join_column, how="left", suffixes=('','_y'))#.drop(in_join_column, axis=1)
 
 	# Reorder results by score, and whether there is text
 	results_df_out = results_df_out.sort_values(['search_score_abs', "search_text"], ascending=False)	
@@ -559,7 +682,7 @@ def bm25_search(free_text_query, in_no_search_results, original_data, searched_d
 	# Highlight found text and save to file
 	results_df_out_wb = create_highlighted_excel_wb(results_df_out, free_text_query, "search_text")
 	results_df_out_wb.save(results_df_name)
-	#results_df_out.to_excel(results_df_name, index= None)
+
 	results_first_text = results_df_out[text_column].iloc[0]
 
 	print("Returning results")

search_funcs/helper_functions.py

@@ -9,6 +9,8 @@ import gzip
 import pickle
 import numpy as np
 
+from typing import List
+
 # Openpyxl functions for output
 from openpyxl import Workbook
 from openpyxl.cell.text import InlineFont 
@@ -175,15 +177,15 @@ def read_file(filename):
 
     return file
 
-def initial_data_load(in_file):
+def initial_data_load(in_file:List[str]):
     '''
-    When file is loaded, update the column dropdown choices
+    When file is loaded, update the column dropdown choices and relevant state variables
     '''
     new_choices = []
     concat_choices = []
     index_load = None
     embed_load = np.array([])
-    tokenised_load =[]
+    tokenised_load = []
     out_message = ""
     current_source = ""
     df = pd.DataFrame()
@@ -257,7 +259,7 @@ def initial_data_load(in_file):
         
     return gr.Dropdown(choices=concat_choices), gr.Dropdown(choices=concat_choices), df, df, index_load, embed_load, tokenised_load, out_message, current_source
 
-def put_columns_in_join_df(in_file):
+def put_columns_in_join_df(in_file:str):
     '''
     When file is loaded, update the column dropdown choices
     '''
@@ -354,7 +356,20 @@ def highlight_found_text(search_text: str, full_text: str) -> str:
 
     return "".join(pos_tokens), combined_positions
 
-def create_rich_text_cell_from_positions(full_text, combined_positions):
+def create_rich_text_cell_from_positions(full_text: str, combined_positions: list[tuple[int, int]]) -> CellRichText:
+    """
+    Create a rich text cell with highlighted positions.
+
+    This function takes the full text and a list of combined positions, and creates a rich text cell
+    with the specified positions highlighted in red.
+
+    Parameters:
+    full_text (str): The full text to be processed.
+    combined_positions (list[tuple[int, int]]): A list of tuples representing the start and end positions to be highlighted.
+
+    Returns:
+    CellRichText: The created rich text cell with highlighted positions.
+    """
     # Construct pos_tokens
     red = InlineFont(color='00FF0000')
     rich_text_cell = CellRichText()
@@ -369,7 +384,21 @@ def create_rich_text_cell_from_positions(full_text, combined_positions):
 
     return rich_text_cell
 
-def create_highlighted_excel_wb(df, search_text, column_to_highlight):
+def create_highlighted_excel_wb(df: pd.DataFrame, search_text: str, column_to_highlight: str) -> Workbook:
+    """
+    Create a new Excel workbook with highlighted search text.
+
+    This function takes a DataFrame, a search text, and a column name to highlight. It creates a new Excel workbook,
+    highlights the occurrences of the search text in the specified column, and returns the workbook.
+
+    Parameters:
+    df (pd.DataFrame): The DataFrame containing the data to be written to the Excel workbook.
+    search_text (str): The text to search for and highlight in the specified column.
+    column_to_highlight (str): The name of the column in which to highlight the search text.
+
+    Returns:
+    Workbook: The created Excel workbook with highlighted search text.
+    """
 
     # Create a new Excel workbook
     wb = Workbook()

search_funcs/semantic_functions.py

@@ -5,11 +5,10 @@ from typing import Type
 import gradio as gr
 import numpy as np
 from datetime import datetime
-#from transformers import AutoModel, AutoTokenizer
-from search_funcs.helper_functions import get_file_path_end
-#import torch
-from torch import cuda, backends#, tensor, mm, utils
+from search_funcs.helper_functions import get_file_path_end, create_highlighted_excel_wb, ensure_output_folder_exists, output_folder
+from torch import cuda, backends
 from sentence_transformers import SentenceTransformer
+PandasDataFrame = Type[pd.DataFrame]
 
 today_rev = datetime.now().strftime("%Y%m%d")
 
@@ -25,22 +24,6 @@ else:
 
 print("Device used is: ", torch_device)
 
-from search_funcs.helper_functions import create_highlighted_excel_wb, ensure_output_folder_exists, output_folder
-
-PandasDataFrame = Type[pd.DataFrame]
-
-# Load embeddings - Jina - deprecated
-# Pinning a Jina revision for security purposes: https://www.baseten.co/blog/pinning-ml-model-revisions-for-compatibility-and-security/
-# Save Jina model locally as described here: https://huggingface.co/jinaai/jina-embeddings-v2-base-en/discussions/29
-# embeddings_name = "jinaai/jina-embeddings-v2-small-en"
-# local_embeddings_location = "model/jina/"
-# revision_choice = "b811f03af3d4d7ea72a7c25c802b21fc675a5d99"
-
-# try:
-#     embeddings_model = AutoModel.from_pretrained(local_embeddings_location, revision = revision_choice, trust_remote_code=True,local_files_only=True, device_map="auto")
-# except:
-#     embeddings_model = AutoModel.from_pretrained(embeddings_name, revision = revision_choice, trust_remote_code=True, device_map="auto")
-
 # Load embeddings
 embeddings_name = "BAAI/bge-small-en-v1.5"
 
@@ -65,32 +48,53 @@ else:
     embeddings_model = SentenceTransformer(embeddings_name)
     print("Could not find local model installation. Downloading from Huggingface")
     
-def docs_to_bge_embed_np_array(docs_out, in_file, embeddings_state, output_file_state, clean, return_intermediate_files = "No", embeddings_super_compress = "No", embeddings_model = embeddings_model, progress=gr.Progress(track_tqdm=True)):
-    '''
-    Takes a Langchain document class and saves it into a Numpy array.
-    '''
+
+def docs_to_bge_embed_np_array(
+    docs_out: list, 
+    in_file: list, 
+    embeddings_state: np.ndarray, 
+    output_file_state: str, 
+    clean: str, 
+    return_intermediate_files: str = "No", 
+    embeddings_super_compress: str = "No", 
+    embeddings_model: SentenceTransformer = embeddings_model, 
+    progress: gr.Progress = gr.Progress(track_tqdm=True)
+) -> tuple:
+    """
+    Process documents to create BGE embeddings and save them as a numpy array.
+
+    Parameters:
+    - docs_out (list): List of documents to be embedded.
+    - in_file (list): List of input files.
+    - embeddings_state (np.ndarray): Current state of embeddings.
+    - output_file_state (str): State of the output file.
+    - clean (str): Indicates if the data should be cleaned.
+    - return_intermediate_files (str, optional): Whether to return intermediate files. Default is "No".
+    - embeddings_super_compress (str, optional): Whether to super compress the embeddings. Default is "No".
+    - embeddings_model (SentenceTransformer, optional): The embeddings model to use. Default is embeddings_model.
+    - progress (gr.Progress, optional): Progress tracker for the function. Default is gr.Progress(track_tqdm=True).
+
+    Returns:
+    - tuple: A tuple containing the output message, embeddings, and output file state.
+    """
+
 
     ensure_output_folder_exists(output_folder)
 
     if not in_file:
         out_message = "No input file found. Please load in at least one file."
         print(out_message)
-        return out_message, None, None, output_file_state
-        
+        return out_message, None, None, output_file_state        
 
     progress(0.6, desc = "Loading/creating embeddings")
 
     print(f"> Total split documents: {len(docs_out)}")
 
-    #print(docs_out)
-
     page_contents = [doc.page_content for doc in docs_out]
 
     ## Load in pre-embedded file if exists
     file_list = [string.name for string in in_file]
 
-    #print(file_list)
-
     embeddings_file_names = [string for string in file_list if "embedding" in string.lower()]
     data_file_names = [string for string in file_list if "tokenised" not in string.lower() and "npz" not in string.lower()]# and "gz" not in string.lower()]
     data_file_name = data_file_names[0]
@@ -98,22 +102,12 @@ def docs_to_bge_embed_np_array(docs_out, in_file, embeddings_state, output_file_
 
     out_message = "Document processing complete. Ready to search."
 
-     # print("embeddings loaded: ", embeddings_out)
 
     if embeddings_state.size == 0:
         tic = time.perf_counter()
         print("Starting to embed documents.")
-        #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())
-
-        
-        
-        #embeddings_out = calc_bge_norm_embeddings(page_contents, embeddings_model, tokenizer)
 
         embeddings_out = embeddings_model.encode(sentences=page_contents, show_progress_bar = True, batch_size = 32, normalize_embeddings=True) # For BGE
-        #embeddings_list = embeddings.encode(sentences=page_contents, normalize_embeddings=True).tolist() # For BGE embeddings
-        #embeddings_list = embeddings.encode(sentences=page_contents).tolist() # For minilm
 
         toc = time.perf_counter()
         time_out = f"The embedding took {toc - tic:0.1f} seconds"
@@ -147,31 +141,43 @@ def docs_to_bge_embed_np_array(docs_out, in_file, embeddings_state, output_file_
 
     return out_message, embeddings_out, output_file_state, output_file_state
 
-def process_data_from_scores_df(df_docs, in_join_file, out_passages, vec_score_cut_off, vec_weight, orig_df_col, in_join_column, search_df_join_column, progress = gr.Progress(track_tqdm=True)):
-
-    def create_docs_keep_from_df(df):
-        dict_out = {'ids' : [df['ids']],
-                    'documents': [df['documents']],
-                    'metadatas': [df['metadatas']],
-                    'distances': [round(df['distances'].astype(float), 4)],
-                    'embeddings': None
-                    }
-        return dict_out
-        
-    # Prepare the DataFrame by transposing
-    #df_docs = df#.apply(lambda x: x.explode()).reset_index(drop=True)
-
-    # Keep only documents with a certain score
-
-    #print(df_docs)
-    
+def process_data_from_scores_df(
+    df_docs: pd.DataFrame, 
+    in_join_file: pd.DataFrame, 
+    vec_score_cut_off: float, 
+    in_join_column: str, 
+    search_df_join_column: str, 
+    progress: gr.Progress = gr.Progress(track_tqdm=True)
+) -> pd.DataFrame:
+    """
+    Process the data from the scores DataFrame by filtering based on score cutoff and document length,
+    and optionally joining with an additional file.
+
+    Parameters
+    ----------
+    df_docs : pd.DataFrame
+        DataFrame containing document scores and metadata.
+    in_join_file : pd.DataFrame
+        DataFrame to join with the results based on specified columns.
+    vec_score_cut_off : float
+        Cutoff value for the vector similarity score.
+    in_join_column : str
+        Column name in the join file to join on.
+    search_df_join_column : str
+        Column name in the search DataFrame to join on.
+    progress : gr.Progress, optional
+        Progress tracker for the function (default is gr.Progress(track_tqdm=True)).
+
+    Returns
+    -------
+    pd.DataFrame
+        Processed DataFrame with filtered and joined data.
+    """
+           
     docs_scores = df_docs["distances"] #.astype(float)
 
     # 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 score_more_limit.empty:
         return pd.DataFrame()
@@ -179,26 +185,17 @@ def process_data_from_scores_df(df_docs, in_join_file, out_passages, vec_score_c
     # Only keep sources that are at least 100 characters long
     docs_len = score_more_limit["documents"].str.len() >= 100
 
-    #print(docs_len)
-
     length_more_limit = score_more_limit.loc[docs_len == True, :] #pd.Series(docs_len) >= 100
-    #docs_keep = create_docs_keep_from_df(length_more_limit) #list(compress(docs_keep, length_more_limit))
-
-    #print(length_more_limit)
 
     if length_more_limit.empty:
         return pd.DataFrame()
             
     length_more_limit['ids'] = length_more_limit['ids'].astype(int)
 
-    #length_more_limit.to_csv("length_more_limit.csv", index = None)
 
     # Explode the 'metadatas' dictionary into separate columns
     df_metadata_expanded = length_more_limit['metadatas'].apply(pd.Series)
 
-    #print(length_more_limit)
-    #print(df_metadata_expanded)
-
     # Concatenate the original DataFrame with the expanded metadata DataFrame
     results_df_out = pd.concat([length_more_limit.drop('metadatas', axis=1), df_metadata_expanded], axis=1)
 
@@ -208,9 +205,6 @@ def process_data_from_scores_df(df_docs, in_join_file, out_passages, vec_score_c
     results_df_out['distances'] = round(results_df_out['distances'].astype(float), 3)
     
 
-    # Join back to original df
-    # results_df_out = orig_df.merge(length_more_limit[['ids', 'distances']], left_index = True, right_on = "ids", how="inner").sort_values("distances")
-
     # Join on additional files
     if not in_join_file.empty:
         progress(0.5, desc = "Joining on additional data file")
@@ -227,68 +221,73 @@ def process_data_from_scores_df(df_docs, in_join_file, out_passages, vec_score_c
 
     return results_df_out
 
-def bge_simple_retrieval(query_str:str, vectorstore, docs, orig_df_col:str, k_val:int, out_passages:int,
-                           vec_score_cut_off:float, vec_weight:float, in_join_file, in_join_column = None, search_df_join_column = None, device = torch_device, embeddings = embeddings_model, progress=gr.Progress(track_tqdm=True)): # ,vectorstore, embeddings
+def bge_semantic_search(
+    query_str: str, 
+    embeddings: np.ndarray, 
+    documents: list, 
+    k_val: int, 
+    vec_score_cut_off: float, 
+    in_join_file: pd.DataFrame, 
+    in_join_column: str = None, 
+    search_df_join_column: str = None, 
+    device: str = torch_device, 
+    embeddings_model: SentenceTransformer = embeddings_model, 
+    progress: gr.Progress = gr.Progress(track_tqdm=True)
+) -> pd.DataFrame:
+    """
+    Perform a semantic search using the BGE model.
+
+    Parameters:
+    - query_str (str): The query string to search for.
+    - embeddings (np.ndarray): The embeddings to search within.
+    - documents (list): The list of documents to search.
+    - k_val (int): The number of top results to return.
+    - vec_score_cut_off (float): The score cutoff for filtering results.
+    - in_join_file (pd.DataFrame): The DataFrame to join with the search results.
+    - in_join_column (str, optional): The column name in the join DataFrame to join on. Default is None.
+    - search_df_join_column (str, optional): The column name in the search DataFrame to join on. Default is None.
+    - device (str, optional): The device to run the model on. Default is torch_device.
+    - embeddings_model (SentenceTransformer, optional): The embeddings model to use. Default is embeddings_model.
+    - progress (gr.Progress, optional): Progress tracker for the function. Default is gr.Progress(track_tqdm=True).
+
+    Returns:
+    - pd.DataFrame: The DataFrame containing the search results.
+    """
 
-    # print("vectorstore loaded: ", vectorstore)
     progress(0, desc = "Conducting semantic search")
 
     ensure_output_folder_exists(output_folder)
 
     print("Searching")
 
-    # Convert it to a PyTorch tensor and transfer to GPU
-    #vectorstore_tensor = tensor(vectorstore).to(device)
-
     # Load the sentence transformer model and move it to GPU
-    embeddings = embeddings.to(device)
+    embeddings_model = embeddings_model.to(device)
 
     # Encode the query using the sentence transformer and convert to a PyTorch tensor
-    query = embeddings.encode(query_str, normalize_embeddings=True)
-
-    # query = calc_bge_norm_embeddings(query_str, embeddings_model=embeddings_model, tokenizer=tokenizer)
-
-    #query_tensor = tensor(query).to(device)
-
-    # if query_tensor.dim() == 1:
-    #     query_tensor = query_tensor.unsqueeze(0)  # Reshape to 2D with one row
+    query = embeddings_model.encode(query_str, normalize_embeddings=True)
 
     # Sentence transformers method, not used:
-    cosine_similarities = query @ vectorstore.T
-    #cosine_similarities = util.cos_sim(query_tensor, vectorstore_tensor)[0]
-    #top_results = torch.topk(cos_scores, k=top_k)
-   
-
-    # Normalize the query tensor and vectorstore tensor
-    #query_norm = query_tensor / query_tensor.norm(dim=1, keepdim=True)
-    #vectorstore_norm = vectorstore_tensor / vectorstore_tensor.norm(dim=1, keepdim=True)
-
-    # Calculate cosine similarities (batch processing)
-    #cosine_similarities = mm(query_norm, vectorstore_norm.T)
-    #cosine_similarities = mm(query_tensor, vectorstore_tensor.T)
+    cosine_similarities = query @ embeddings.T
 
     # Flatten the tensor to a 1D array
     cosine_similarities = cosine_similarities.flatten()
 
-    # Convert to a NumPy array if it's still a PyTorch tensor
-    #cosine_similarities = cosine_similarities.cpu().numpy()
-
     # Create a Pandas Series
     cosine_similarities_series = pd.Series(cosine_similarities)
 
-    # Pull out relevent info from docs
-    page_contents = [doc.page_content for doc in docs]
-    page_meta = [doc.metadata for doc in docs]
+    # Pull out relevent info from documents
+    page_contents = [doc.page_content for doc in documents]
+    page_meta = [doc.metadata for doc in documents]
     ids_range = range(0,len(page_contents)) 
     ids = [str(element) for element in ids_range]
 
-    df_docs = pd.DataFrame(data={"ids": ids,
+    df_documents = pd.DataFrame(data={"ids": ids,
                                 "documents": page_contents,
                                     "metadatas":page_meta,
                                     "distances":cosine_similarities_series}).sort_values("distances", ascending=False).iloc[0:k_val,:]
 
     
-    results_df_out = process_data_from_scores_df(df_docs, in_join_file, out_passages, vec_score_cut_off, vec_weight, orig_df_col, in_join_column, search_df_join_column)
+    results_df_out = process_data_from_scores_df(df_documents, in_join_file, vec_score_cut_off, in_join_column, search_df_join_column)
 
     print("Search complete")
 
@@ -312,291 +311,4 @@ def bge_simple_retrieval(query_str:str, vectorstore, docs, orig_df_col:str, k_va
 
     print("Returning results")
 
-    return results_first_text, results_df_name
-
-
-def docs_to_jina_embed_np_array_deprecated(docs_out, in_file, embeddings_state, return_intermediate_files = "No", embeddings_super_compress = "No", embeddings = embeddings_model, progress=gr.Progress(track_tqdm=True)):
-    '''
-    Takes a Langchain document class and saves it into a Chroma sqlite file.
-    '''
-    if not in_file:
-        out_message = "No input file found. Please load in at least one file."
-        print(out_message)
-        return out_message, None, None
-        
-
-    progress(0.6, desc = "Loading/creating embeddings")
-
-    print(f"> Total split documents: {len(docs_out)}")
-
-    #print(docs_out)
-
-    page_contents = [doc.page_content for doc in docs_out]
-
-    ## Load in pre-embedded file if exists
-    file_list = [string.name for string in in_file]
-
-    #print(file_list)
-
-    embeddings_file_names = [string for string in file_list if "embedding" in string.lower()]
-    data_file_names = [string for string in file_list if "tokenised" not in string.lower() and "npz" not in string.lower()]# and "gz" not in string.lower()]
-    data_file_name = data_file_names[0]
-    data_file_name_no_ext = get_file_path_end(data_file_name)
-
-    out_message = "Document processing complete. Ready to search."
-
-     # print("embeddings loaded: ", embeddings_out)
-
-    if embeddings_state.size == 0:
-        tic = time.perf_counter()
-        print("Starting to embed documents.")
-        #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())
-
-        embeddings_out = embeddings.encode(sentences=page_contents, max_length=1024, show_progress_bar = True, batch_size = 32) # For Jina embeddings
-        #embeddings_list = embeddings.encode(sentences=page_contents, normalize_embeddings=True).tolist() # For BGE embeddings
-        #embeddings_list = embeddings.encode(sentences=page_contents).tolist() # For minilm
-
-        toc = time.perf_counter()
-        time_out = f"The embedding took {toc - tic:0.1f} seconds"
-        print(time_out)
-
-        # If you want to save your files for next time
-        if return_intermediate_files == "Yes":
-            progress(0.9, desc = "Saving embeddings to file")
-            if embeddings_super_compress == "No":
-                semantic_search_file_name = data_file_name_no_ext + '_' + 'embeddings.npz'
-                np.savez_compressed(semantic_search_file_name, embeddings_out)
-            else:
-                semantic_search_file_name = data_file_name_no_ext + '_' + 'embedding_compress.npz'
-                embeddings_out_round = np.round(embeddings_out, 3) 
-                embeddings_out_round *= 100 # Rounding not currently used
-                np.savez_compressed(semantic_search_file_name, embeddings_out_round)
-
-            return out_message, embeddings_out, semantic_search_file_name
-
-        return out_message, embeddings_out, None
-    else:
-        # Just return existing embeddings if already exist
-        embeddings_out = embeddings_state
-    
-    print(out_message)
-
-    return out_message, embeddings_out, None#, None
-
-def jina_simple_retrieval_deprecated(query_str:str, vectorstore, docs, orig_df_col:str, k_val:int, out_passages:int,
-                           vec_score_cut_off:float, vec_weight:float, in_join_file, in_join_column = None, search_df_join_column = None, device = torch_device, embeddings = embeddings_model, progress=gr.Progress(track_tqdm=True)): # ,vectorstore, embeddings
-
-    # print("vectorstore loaded: ", vectorstore)
-    progress(0, desc = "Conducting semantic search")
-
-    print("Searching")
-
-    # Convert it to a PyTorch tensor and transfer to GPU
-    vectorstore_tensor = tensor(vectorstore).to(device)
-
-    # Load the sentence transformer model and move it to GPU
-    embeddings = embeddings.to(device)
-
-    # Encode the query using the sentence transformer and convert to a PyTorch tensor
-    query = embeddings.encode(query_str)
-    query_tensor = tensor(query).to(device)
-
-    if query_tensor.dim() == 1:
-        query_tensor = query_tensor.unsqueeze(0)  # Reshape to 2D with one row
-
-    # Normalize the query tensor and vectorstore tensor
-    query_norm = query_tensor / query_tensor.norm(dim=1, keepdim=True)
-    vectorstore_norm = vectorstore_tensor / vectorstore_tensor.norm(dim=1, keepdim=True)
-
-    # Calculate cosine similarities (batch processing)
-    cosine_similarities = mm(query_norm, vectorstore_norm.T)
-
-    # Flatten the tensor to a 1D array
-    cosine_similarities = cosine_similarities.flatten()
-
-    # Convert to a NumPy array if it's still a PyTorch tensor
-    cosine_similarities = cosine_similarities.cpu().numpy()
-
-    # Create a Pandas Series
-    cosine_similarities_series = pd.Series(cosine_similarities)
-
-    # Pull out relevent info from docs
-    page_contents = [doc.page_content for doc in docs]
-    page_meta = [doc.metadata for doc in docs]
-    ids_range = range(0,len(page_contents)) 
-    ids = [str(element) for element in ids_range]
-
-    df_docs = pd.DataFrame(data={"ids": ids,
-                                "documents": page_contents,
-                                    "metadatas":page_meta,
-                                    "distances":cosine_similarities_series}).sort_values("distances", ascending=False).iloc[0:k_val,:]
-
-    
-    results_df_out = process_data_from_scores_df(df_docs, in_join_file, out_passages, vec_score_cut_off, vec_weight, orig_df_col, in_join_column, search_df_join_column)
-
-    print("Search complete")
-
-    # If nothing found, return error message
-    if results_df_out.empty:
-        return 'No result found!', None
-    
-    query_str_file = query_str.replace(" ", "_")
-
-    results_df_name = "semantic_search_result_" + today_rev + "_" +  query_str_file + ".xlsx"
-
-    print("Saving search output to file")
-    progress(0.7, desc = "Saving search output to file")
-
-    results_df_out.to_excel(results_df_name, index= None)
-    results_first_text = results_df_out.iloc[0, 1]
-
-    print("Returning results")
-
-    return results_first_text, results_df_name
-
-# Deprecated Chroma functions - kept just in case needed in future.
-# Chroma support is currently deprecated
-# Import Chroma and instantiate a client. The default Chroma client is ephemeral, meaning it will not save to disk.
-#import chromadb
-#from chromadb.config import Settings
-#from typing_extensions import Protocol
-#from chromadb import Documents, EmbeddingFunction, Embeddings
-
-# Remove Chroma database file. If it exists as it can cause issues
-#chromadb_file = "chroma.sqlite3"
-
-#if os.path.isfile(chromadb_file):
-#    os.remove(chromadb_file)
-
-
-def docs_to_chroma_save_deprecated(docs_out, embeddings = embeddings_model, progress=gr.Progress()):
-    '''
-    Takes a Langchain document class and saves it into a Chroma sqlite file. Not currently used.
-    '''
-
-    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]
-
-    tic = time.perf_counter()
-    #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())
-
-    embeddings_list = embeddings.encode(sentences=page_contents, max_length=256, show_progress_bar = True, batch_size = 32).tolist() # For Jina embeddings
-    #embeddings_list = embeddings.encode(sentences=page_contents, normalize_embeddings=True).tolist() # For BGE embeddings
-    #embeddings_list = embeddings.encode(sentences=page_contents).tolist() # For minilm
-
-    toc = time.perf_counter()
-    time_out = f"The embedding took {toc - tic:0.1f} seconds"
-
-    #pd.Series(embeddings_list).to_csv("embeddings_out.csv")
-
-    # Jina tiny
-    # This takes about 300 seconds for 240,000 records = 800 / second, 1024 max length
-    # For 50k records:
-    # 61 seconds at 1024 max length
-    # 55 seconds at 512 max length
-    # 43 seconds at 256 max length
-    # 31 seconds at 128 max length
-
-    # The embedding took 1372.5 seconds at 256 max length for 655,020 case notes
-
-    # BGE small
-    # 96 seconds for 50k records at 512 length
-
-    # all-MiniLM-L6-v2
-    # 42.5 seconds at (256?) max length
-
-    # paraphrase-MiniLM-L3-v2
-    # 22 seconds for 128 max length
-
-
-    print(time_out)
-
-    chroma_tic = time.perf_counter()
-
-    # Create a new Chroma collection to store the documents and metadata. We don't need to specify an embedding fuction, and the default will be used.
-    client = chromadb.PersistentClient(path="./last_year", settings=Settings(
-    anonymized_telemetry=False))
-
-    try:
-        print("Deleting existing collection.")
-        #collection = client.get_collection(name="my_collection")
-        client.delete_collection(name="my_collection")
-        print("Creating new collection.")
-        collection = client.create_collection(name="my_collection")
-    except: 
-        print("Creating new collection.")
-        collection = client.create_collection(name="my_collection")
-
-    # Match batch size is about 40,000, so add that amount in a loop
-    def create_batch_ranges(in_list, batch_size=40000):
-        total_rows = len(in_list)
-        ranges = []
-        
-        for start in range(0, total_rows, batch_size):
-            end = min(start + batch_size, total_rows)
-            ranges.append(range(start, end))
-        
-        return ranges
-
-    batch_ranges = create_batch_ranges(embeddings_list)
-    print(batch_ranges)
-
-    for row_range in progress.tqdm(batch_ranges, desc = "Creating vector database", unit = "batches of 40,000 rows"):
-    
-        collection.add(
-        documents = page_contents[row_range[0]:row_range[-1]],
-        embeddings = embeddings_list[row_range[0]:row_range[-1]],
-        metadatas = page_meta[row_range[0]:row_range[-1]],
-        ids = ids[row_range[0]:row_range[-1]])
-        #print("Here")
-        
-    # print(collection.count())
-    
-
-    #chatf.vectorstore = vectorstore_func
-
-    chroma_toc = time.perf_counter()
-
-    chroma_time_out = f"Loading to Chroma db took {chroma_toc - chroma_tic:0.1f} seconds"
-    print(chroma_time_out)
-
-    out_message = "Document processing complete"
-
-    return out_message, collection
-
-def chroma_retrieval_deprecated(query_str:str, vectorstore, docs, orig_df_col:str, k_val:int, out_passages:int,
-                           vec_score_cut_off:float, vec_weight:float, in_join_file = None, in_join_column = None, search_df_join_column = None, embeddings = embeddings_model): # ,vectorstore, embeddings
-
-            query = embeddings.encode(query_str).tolist()
-
-            docs = vectorstore.query(
-            query_embeddings=query,
-            n_results= k_val # No practical limit on number of responses returned
-            #where={"metadata_field": "is_equal_to_this"},
-            #where_document={"$contains":"search_string"}
-            )
-
-            df_docs = pd.DataFrame(data={'ids': docs['ids'][0],
-                                    'documents': docs['documents'][0],
-                                    'metadatas':docs['metadatas'][0],
-                                    'distances':docs['distances'][0]#,                                    
-                                    #'embeddings': docs['embeddings']
-                                    })
-            
-            results_df_out = process_data_from_scores_df(df_docs, in_join_file, out_passages, vec_score_cut_off, vec_weight, orig_df_col, in_join_column, search_df_join_column)
-
-            results_df_name = output_folder + "semantic_search_result.csv"
-            results_df_out.to_csv(results_df_name, index= None)
-            results_first_text = results_df_out[orig_df_col].iloc[0]
-
-            return results_first_text, results_df_name
+    return results_first_text, results_df_name

search_funcs/spacy_search_funcs.py

@@ -27,9 +27,14 @@ except:
 	nlp = spacy.load("en_core_web_sm")
 	print("Successfully imported spaCy model")
 
-def spacy_fuzzy_search(string_query:str, df_list: List[str], original_data: PandasDataFrame, text_column:str, in_join_file: PandasDataFrame, search_df_join_column:str, in_join_column:str, no_spelling_mistakes:int = 1, progress=gr.Progress(track_tqdm=True)):
+def spacy_fuzzy_search(string_query:str, tokenised_data: List[List[str]], original_data: PandasDataFrame, text_column:str, in_join_file: PandasDataFrame, search_df_join_column:str, in_join_column:str, no_spelling_mistakes:int = 1, progress=gr.Progress(track_tqdm=True)):
     ''' Conduct fuzzy match on a list of data.'''
 
+    #print("df_list:", df_list)
+
+    # Convert tokenised data back into a list of strings
+    df_list = list(map(" ".join, tokenised_data))
+
     if len(df_list) > 10000:
          out_message = "Your data has more than 10,000 rows and will take more than three minutes to do a fuzzy search. Please try keyword or semantic search for data of this size." 
          return out_message, None