Improvements with embeddings load and file save

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 load_embedding_model, docs_to_bge_embed_np_array, bge_semantic_search
+from search_funcs.semantic_functions import load_embedding_model, docs_to_embed_np_array, bge_semantic_search
 from search_funcs.helper_functions import display_info, initial_data_load, put_columns_in_join_df, get_connection_params, output_folder, get_or_create_env_var # Not currently used: get_temp_folder_path, empty_folder, 
 from search_funcs.spacy_search_funcs import spacy_fuzzy_search
 from search_funcs.aws_functions import load_data_from_aws
@@ -99,7 +99,7 @@ depends on factors such as the type of documents or queries. Information taken f
     """
     **Thematic/semantic search**
 
-    This search type enables you to search for general terms (e.g. happiness, nature) and the search will pick out text passages that are most semantically similar to them. 1. Load in data file (ideally a file with '_cleaned' at the end of the name, a pkl.gz file), with (optionally) the 'embeddings... .npz' file in the same folder to save loading time. 2. Select the field in your data to search. 3. Wait for the data file to be prepared for search. 4. Enter the search term in the 'Enter semantic search query here' box below and press Enter/click on 'Start semantic search'. 4. Your search results will be saved in a csv file and will be presented in the 'File output' area below.
+    This search type enables you to search for general terms (e.g. happiness, nature) and the search will pick out text passages that are most semantically similar to them. 1. Load in data file (ideally a file with '_cleaned' at the end of the name, a pkl.gz file), with (optionally) the 'embed... .npz' file in the same folder to save loading time. 2. Select the field in your data to search. 3. Wait for the data file to be prepared for search. 4. Enter the search term in the 'Enter semantic search query here' box below and press Enter/click on 'Start semantic search'. 4. Your search results will be saved in a csv file and will be presented in the 'File output' area below.
     """)       
 
         with gr.Row():
@@ -202,7 +202,7 @@ depends on factors such as the type of documents or queries. Information taken f
     
     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=[semantic_input_document_format, semantic_load_progress, output_file_state], api_name="convert_texts_to_documents").\
-        then(docs_to_bge_embed_np_array, inputs=[semantic_input_document_format, in_semantic_file, output_file_state, in_clean_data, embeddings_state, embeddings_model_name_state, embeddings_model_loc_state, return_intermediate_files, embeddings_compress], outputs=[semantic_load_progress, embeddings_state, semantic_output_file, output_file_state, embeddings_model_state], api_name="embed_documents")
+        then(docs_to_embed_np_array, inputs=[semantic_input_document_format, in_semantic_file, output_file_state, in_clean_data, embeddings_state, embeddings_model_name_state, embeddings_model_loc_state, return_intermediate_files, embeddings_compress], outputs=[semantic_load_progress, embeddings_state, semantic_output_file, output_file_state, embeddings_model_state], api_name="embed_documents")
 
     # Semantic search query
     semantic_submit.click(bge_semantic_search, inputs=[semantic_query, embeddings_state, semantic_input_document_format, semantic_k_val, semantic_min_distance, embeddings_model_state, embeddings_model_name_state, embeddings_compress, join_data_state, in_join_column, search_df_join_column], outputs=[semantic_output_single_text, semantic_output_file], api_name="semantic_search")

search_funcs/helper_functions.py

@@ -261,13 +261,13 @@ def initial_data_load(in_file:List[str], progress = gr.Progress(track_tqdm=True)
 
     progress(0.3, desc="Loading in data files")
 
-    data_file_names = [string for string in file_list if "tokenised" not in string.lower() and "npz" not in string.lower() and "search_index" not in string.lower()]
+    data_file_names = [string for string in file_list if "tokenised" not in string.lower() and "npz" not in string.lower() or "prep_docs" in string.lower()]
     print("Data file names:", data_file_names)
 
     if not data_file_names:
         out_message = "Please load in at least one csv/Excel/parquet data file."
         print(out_message)
-        return gr.Dropdown(choices=concat_choices), gr.Dropdown(choices=concat_choices), pd.DataFrame(), pd.DataFrame(), index_load, embed_load, tokenised_load, out_message, None
+        return gr.Dropdown(choices=concat_choices), gr.Dropdown(choices=concat_choices), pd.DataFrame(), pd.DataFrame(), index_load, embed_load, tokenised_load, out_message, None, file_list
 
     # This if you have loaded in a documents object for the semantic search
     if "pkl" in data_file_names[0]: 
@@ -288,11 +288,9 @@ def initial_data_load(in_file:List[str], progress = gr.Progress(track_tqdm=True)
             if file_size > file_size_bytes_500mb:
                 out_message = "Data file greater than 500mb in size. Please use smaller sizes."
                 print(out_message)
-                return gr.Dropdown(choices=concat_choices), gr.Dropdown(choices=concat_choices), pd.DataFrame(), pd.DataFrame(), index_load, embed_load, tokenised_load, out_message, None
-
+                return gr.Dropdown(choices=concat_choices), gr.Dropdown(choices=concat_choices), pd.DataFrame(), pd.DataFrame(), index_load, embed_load, tokenised_load, out_message, None, file_list
 
             df_new = read_file(file)
-
             df = pd.concat([df, df_new], ignore_index = True)
 
         new_choices = list(df.columns)
@@ -302,22 +300,22 @@ def initial_data_load(in_file:List[str], progress = gr.Progress(track_tqdm=True)
     progress(0.6, desc="Loading in embedding/search index files")
 
     # Check if there is a search index file already
-    index_file_names = [string for string in file_list if ".gz" in string.lower()]
+    index_file_names = [string for string in file_list if "pkl.gz" in string.lower()]
 
     if index_file_names:
         index_file_name = index_file_names[0]
         print("Search index file name found:", index_file_name)
         index_load = read_file(index_file_name)
 
-    embeddings_file_names = [string for string in file_list if "embedding" in string.lower()]
+    embeddings_file_names = [string for string in file_list if ".npz" in string.lower()]
 
     if embeddings_file_names:
         print("Loading embeddings from file.")
         embed_load = np.load(embeddings_file_names[0])['arr_0']
 
         # If embedding files have 'super_compress' in the title, they have been multiplied by 100 before save
-        if "compress" in embeddings_file_names[0]:
-            embed_load /= 100
+        #if "compress" in embeddings_file_names[0]:
+        #    embed_load /= 100
     else:
         embed_load = np.array([])
 

search_funcs/semantic_functions.py

@@ -50,7 +50,7 @@ def load_embedding_model(embeddings_name = "BAAI/bge-small-en-v1.5", embedding_l
 
     return embeddings_model, torch_device
 
-def docs_to_bge_embed_np_array(
+def docs_to_embed_np_array(
     docs_out: list, 
     in_file: list, 
     output_file_state: str, 
@@ -136,9 +136,9 @@ def docs_to_bge_embed_np_array(
 
             progress(0.9, desc = "Saving embeddings to file")
             if embeddings_compress == "No":
-                semantic_search_file_name = output_folder + data_file_name_no_ext + '_bge_embeddings.npz'
+                semantic_search_file_name = output_folder + data_file_name_no_ext + '_embeddings.npz'
             else:
-                semantic_search_file_name = output_folder + data_file_name_no_ext + '_bge_embedding_compress.npz'
+                semantic_search_file_name = output_folder + data_file_name_no_ext + '_embedding_compress.npz'
                 
             np.savez_compressed(semantic_search_file_name, embeddings_out)
 

search_funcs/semantic_ingest_functions.py

@@ -97,6 +97,11 @@ def csv_excel_text_to_docs(df:PandasDataFrame, in_file:List[str], text_column:st
     ensure_output_folder_exists(output_folder)
     output_list = []
 
+    if not isinstance(text_column, str):
+        text_column = str(text_column)
+
+    print("text_column:", text_column)
+
     if not in_file:
         return None, "Please load in at least one file.", output_list
 
@@ -115,12 +120,12 @@ def csv_excel_text_to_docs(df:PandasDataFrame, in_file:List[str], text_column:st
     data_file_name = data_file_names[0]
 
     # Check if file is a document format, and explode out as needed
-    if "prepared_docs" in data_file_name:
+    if "prep_docs" in data_file_name:
         print("Loading in documents from file.")
 
         doc_sections = df
 
-        # Convert each element in the Series to a Document instance
+        print("doc_sections:", doc_sections[0])
 
         return doc_sections, "Finished preparing documents", output_list
 
@@ -147,17 +152,29 @@ def csv_excel_text_to_docs(df:PandasDataFrame, in_file:List[str], text_column:st
         clean_toc = time.perf_counter()
         clean_time_out = f"Cleaning the text took {clean_toc - clean_tic:0.1f} seconds."
         print(clean_time_out)
+    
+    else:
+        df_list = list(df[text_column])
+        prepared_text_df = pd.DataFrame(data={text_column:df_list})
 
-    cols = [col for col in df.columns if col != original_text_column]
+        # Drop original column from input file to reduce file size
+        in_df = df.drop(text_column, axis = 1)
+        df = pd.concat([in_df, prepared_text_df], axis = 1)
 
+    cols = [col for col in df.columns if col != original_text_column]
     df["metadata"] = combine_metadata_columns(df, cols)
 
     progress(0.3, desc = "Converting data to document format")
 
+    #print("text_column name:", text_column)
+    #print("text_column:", df[text_column])
+    #print("metadata", df["metadata"])
+
     # Create a list of Document objects
-    doc_sections = [Document(page_content=row[text_column], 
-                        metadata= parse_metadata(row["metadata"]))
+    doc_sections = [Document(page_content=row[text_column], metadata= parse_metadata(row["metadata"]))
                 for index, row in progress.tqdm(df.iterrows(), desc = "Splitting up text", unit = "rows")]
+    
+    print("doc_sections:", doc_sections[0])
 
     ingest_toc = time.perf_counter()
 
@@ -169,15 +186,11 @@ def csv_excel_text_to_docs(df:PandasDataFrame, in_file:List[str], text_column:st
         data_file_out_name_no_ext = get_file_path_end(data_file_name)
         file_name = data_file_out_name_no_ext
 
-        if clean == "No":
-            out_doc_file_name = output_folder + file_name + "_prepared_docs.pkl.gz"
-            with gzip.open(out_doc_file_name, 'wb') as file:
-                pickle.dump(doc_sections, file)
+        if clean == "No": out_doc_file_name = output_folder + file_name + "_prep_docs.pkl.gz"
+        elif clean == "Yes": out_doc_file_name = output_folder + file_name + "_cleaned_prep_docs.pkl.gz"
 
-        elif clean == "Yes":
-            out_doc_file_name = output_folder + file_name + "_cleaned_prepared_docs.pkl.gz"
-            with gzip.open(out_doc_file_name, 'wb') as file:
-                pickle.dump(doc_sections, file)
+        with gzip.open(out_doc_file_name, 'wb') as file:
+            pickle.dump(doc_sections, file)
 
         output_list.append(out_doc_file_name)
         print("Documents saved to file.")