import os import time import pandas as pd from typing import Type import gradio as gr import numpy as np from datetime import datetime import accelerate today_rev = datetime.now().strftime("%Y%m%d") from transformers import AutoModel from torch import cuda, backends, tensor, mm from search_funcs.helper_functions import read_file # Check for torch cuda print("Is CUDA enabled? ", cuda.is_available()) print("Is a CUDA device available on this computer?", backends.cudnn.enabled) if cuda.is_available(): torch_device = "cuda" os.system("nvidia-smi") else: torch_device = "cpu" print("Device used is: ", torch_device) #from search_funcs.helper_functions import get_file_path_end PandasDataFrame = Type[pd.DataFrame] # Load embeddings # 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") # 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 get_file_path_end(file_path): # First, get the basename of the file (e.g., "example.txt" from "/path/to/example.txt") basename = os.path.basename(file_path) # Then, split the basename and its extension and return only the basename without the extension filename_without_extension, _ = os.path.splitext(basename) #print(filename_without_extension) return filename_without_extension def load_embeddings(embeddings_name = embeddings_name): ''' 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, device_map="auto") global embeddings embeddings = embeddings_func return embeddings def docs_to_jina_embed_np_array(docs_out, in_file, return_intermediate_files = "No", embeddings_super_compress = "No", embeddings = embeddings_model, 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] ## Load in pre-embedded file if exists file_list = [string.name for string in in_file] #print(file_list) embeddings_file_names = [string.lower() for string in file_list if "embedding" in string.lower()] data_file_names = [string.lower() 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." if embeddings_file_names: print("Loading embeddings from file.") embeddings_out = 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]: embeddings_out /= 100 # print("embeddings loaded: ", embeddings_out) if not embeddings_file_names: 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": 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 print(out_message) return out_message, embeddings_out, None#, None 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): 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) 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() # 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) results_df_out = results_df_out.rename(columns={"documents":orig_df_col}) results_df_out = results_df_out.drop(["page_section", "row", "source", "id"], axis=1, errors="ignore") 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 in_join_file: join_filename = in_join_file.name # Import data join_df = read_file(join_filename) join_df[in_join_column] = join_df[in_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) results_df_out[search_df_join_column] = results_df_out[search_df_join_column].astype(str).str.replace("\.0$","", regex=True) results_df_out = results_df_out.merge(join_df,left_on=search_df_join_column, right_on=in_join_column, how="left").drop(in_join_column, axis=1) return results_df_out def jina_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 = None, in_join_column = None, search_df_join_column = None, device = torch_device, embeddings = embeddings_model, progress=gr.Progress()): # ,vectorstore, embeddings # print("vectorstore loaded: ", vectorstore) # 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) # 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" results_df_out.to_excel(results_df_name, index= None) results_first_text = results_df_out.iloc[0, 1] return results_first_text, results_df_name # Deprecated Chroma functions - kept just in case needed in future. 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 = "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