import numpy as np import io import os import logging import collections import tempfile from langchain.document_loaders import UnstructuredFileLoader from langchain.text_splitter import CharacterTextSplitter from langchain.vectorstores import FAISS from langchain.embeddings import HuggingFaceEmbeddings from langchain.document_loaders import PDFMinerPDFasHTMLLoader from bs4 import BeautifulSoup import re from langchain.docstore.document import Document import unstructured from unstructured.partition.docx import partition_docx from unstructured.partition.auto import partition from transformers import AutoTokenizer MODEL = "thenlper/gte-base" CHUNK_SIZE = 1000 CHUNK_OVERLAP = 200 embeddings = HuggingFaceEmbeddings( model_name=MODEL, cache_folder=os.getenv("SENTENCE_TRANSFORMERS_HOME") ) model_id = "mistralai/Mistral-7B-Instruct-v0.1" tokenizer = AutoTokenizer.from_pretrained( model_id, padding_side="left" ) text_splitter = CharacterTextSplitter( separator = "\n", chunk_size = CHUNK_SIZE, chunk_overlap = CHUNK_OVERLAP, length_function = len, ) ## PDF Functions def group_text_by_font_size(content): cur_fs = [] cur_text = '' cur_page = -1 cur_c = content[0] multi_fs = False snippets = [] # first collect all snippets that have the same font size for c in content: # print(f"c={c}\n\n") if c.find('a') != None and c.find('a').get('name'): cur_page = int(c.find('a').get('name')) sp_list = c.find_all('span') if not sp_list: continue for sp in sp_list: # print(f"sp={sp}\n\n") if not sp: continue st = sp.get('style') if not st: continue fs = re.findall('font-size:(\d+)px',st) # print(f"fs={fs}\n\n") if not fs: continue fs = [int(fs[0])] if len(cur_fs)==0: cur_fs = fs if fs == cur_fs: cur_text += sp.text elif not sp.find('br') and cur_c==c: cur_text += sp.text cur_fs.extend(fs) multi_fs = True elif sp.find('br') and multi_fs == True: # if a br tag is found and the text is in a different fs, it is the last part of the multifontsize line cur_fs.extend(fs) snippets.append((cur_text+sp.text,max(cur_fs), cur_page)) cur_fs = [] cur_text = '' cur_c = c multi_fs = False else: snippets.append((cur_text,max(cur_fs), cur_page)) cur_fs = fs cur_text = sp.text cur_c = c multi_fs = False snippets.append((cur_text,max(cur_fs), cur_page)) return snippets def get_titles_fs(fs_list): filtered_fs_list = [item[0] for item in fs_list if item[0] > fs_list[0][0]] return sorted(filtered_fs_list, reverse=True) def calculate_total_characters(snippets): font_sizes = {} #dictionary to store font-size and total characters for text, font_size, _ in snippets: #remove newline# and digits cleaned_text = text.replace('\n', '') #cleaned_text = re.sub(r'\d+', '', cleaned_text) total_characters = len(cleaned_text) #update the dictionary if font_size in font_sizes: font_sizes[font_size] += total_characters else: font_sizes[font_size] = total_characters #convert the dictionary into a sorted list of tuples size_charac_list = sorted(font_sizes.items(), key=lambda x: x[1], reverse=True) return size_charac_list def create_documents(source, snippets, font_sizes): docs = [] titles_fs = get_titles_fs(font_sizes) for snippet in snippets: cur_fs = snippet[1] if cur_fs>font_sizes[0][0] and len(snippet[0])>2: content = min((titles_fs.index(cur_fs)+1), 3)*"#" + " " + snippet[0].replace(" ", " ") category = "Title" else: content = snippet[0].replace(" ", " ") category = "Paragraph" metadata={"source":source, "filename":source.split("/")[-1], "file_directory": "/".join(source.split("/")[:-1]), "file_category":"", "file_sub-cat":"", "file_sub2-cat":"", "category":category, "filetype":source.split(".")[-1], "page_number":snippet[2]} categories = source.split("/") cat_update="" if len(categories)>4: cat_update = {"file_category":categories[1], "file_sub-cat":categories[2], "file_sub2-cat":categories[3]} elif len(categories)>3: cat_update = {"file_category":categories[1], "file_sub-cat":categories[2]} elif len(categories)>2: cat_update = {"file_category":categories[1]} metadata.update(cat_update) docs.append(Document(page_content=content, metadata=metadata)) return docs ## Group Chunks docx or pdf # -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE def group_chunks_by_section(chunks, min_chunk_size=512): filtered_chunks = [chunk for chunk in chunks if chunk.metadata['category'] != 'PageBreak']# Add more filters if needed #print(f"filtered = {len(filtered_chunks)} - before = {len(chunks)}") new_chunks = [] seen_paragraph = False new_title = True #switches when there is a new paragraph to create a new chunk for i, chunk in enumerate(filtered_chunks): # print(f"\n\n\n#{i}:METADATA: {chunk.metadata['category']}") if new_title: #print(f"<-- NEW title DETECTED -->") new_chunk = chunk new_title = False add_content = False new_chunk.metadata['titles'] = "" #print(f"CONTENT: {new_chunk.page_content}\nMETADATA: {new_chunk.metadata['category']} \n title: {new_chunk.metadata['title']}") if chunk.metadata['category'].lower() =='title': new_chunk.metadata['titles'] += f"{chunk.page_content} ~~ " else: #Activates when a paragraph is seen after one or more titles seen_paragraph = True #Avoid adding the title 2 times to the page content if add_content:#and chunk.page_content not in new_chunk.page_content new_chunk.page_content += f"\n{chunk.page_content}" #edit the end_page number, the last one keeps its place try: new_chunk.metadata['end_page'] = chunk.metadata['page_number'] except: print("", end="") #print("Exception: No page number in metadata") add_content = True #If filtered_chunks[i+1] raises an error, this is probably because this is the last chunk try: #If the next chunk is a title and we have already seen a paragraph and the current chunk content is long enough, we create a new document if filtered_chunks[i+1].metadata['category'].lower() =="title" and seen_paragraph and len(new_chunk.page_content)>min_chunk_size: if 'category' in new_chunk.metadata: new_chunk.metadata.pop('category') new_chunks.append(new_chunk) new_title = True seen_paragraph = False #index out of range except: new_chunks.append(new_chunk) #print('🆘 Gone through all chunks 🆘') break return new_chunks # -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE ## Split documents by font def split_pdf(file_path, folder): loader = PDFMinerPDFasHTMLLoader(file_path) data = loader.load()[0] # entire pdf is loaded as a single Document soup = BeautifulSoup(data.page_content,'html.parser') content = soup.find_all('div')#List of all elements in div tags try: snippets = group_text_by_font_size(content) except Exception as e: print("ERROR WHILE GROUPING BY FONT SIZE", e) snippets = [("ERROR WHILE GROUPING BY FONT SIZE", 0, -1)] font_sizes = calculate_total_characters(snippets)#get the amount of characters for each font_size chunks = create_documents(file_path, snippets, font_sizes) return chunks # -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE def split_docx(file_path, folder): chunks_elms = partition_docx(filename=file_path) chunks = [] file_categories = file_path.split("/") for chunk_elm in chunks_elms: category = chunk_elm.category if category == "Title": chunk = Document(page_content= min(chunk_elm.metadata.to_dict()['category_depth']+1, 3)*"#" + ' ' + chunk_elm.text, metadata=chunk_elm.metadata.to_dict()) else: chunk = Document(page_content=chunk_elm.text, metadata=chunk_elm.metadata.to_dict()) metadata={"source":file_path, "filename":file_path.split("/")[-1], "file_category":"", "file_sub-cat":"", "file_sub2-cat":"", "category":category, "filetype":file_path.split(".")[-1]} cat_update="" if len(file_categories)>4: cat_update = {"file_category":file_categories[1], "file_sub-cat":file_categories[2], "file_sub2-cat":file_categories[3]} elif len(file_categories)>3: cat_update = {"file_category":file_categories[1], "file_sub-cat":file_categories[2]} elif len(file_categories)>2: cat_update = {"file_category":file_categories[1]} metadata.update(cat_update) chunk.metadata.update(metadata) chunks.append(chunk) return chunks # Load the index of documents (if it has already been built) def rebuild_index(input_folder, output_folder): paths_time = [] to_keep = set() print(f'number of files {len(paths_time)}') if len(output_folder.list_paths_in_partition()) > 0: with tempfile.TemporaryDirectory() as temp_dir: for f in output_folder.list_paths_in_partition(): with output_folder.get_download_stream(f) as stream: with open(os.path.join(temp_dir, os.path.basename(f)), "wb") as f2: f2.write(stream.read()) index = FAISS.load_local(temp_dir, embeddings) to_remove = [] logging.info(f"{len(index.docstore._dict)} vectors loaded") for idx, doc in index.docstore._dict.items(): source = (doc.metadata["source"], doc.metadata["last_modified"]) if source in paths_time: # Identify documents already indexed and still present in the source folder to_keep.add(source) else: # Identify documents removed from the source folder to_remove.append(idx) docstore_id_to_index = {v: k for k, v in index.index_to_docstore_id.items()} # Remove documents that have been deleted from the source folder vectors_to_remove = [] for idx in to_remove: del index.docstore._dict[idx] ind = docstore_id_to_index[idx] del index.index_to_docstore_id[ind] vectors_to_remove.append(ind) index.index.remove_ids(np.array(vectors_to_remove, dtype=np.int64)) index.index_to_docstore_id = { i: ind for i, ind in enumerate(index.index_to_docstore_id.values()) } logging.info(f"{len(to_remove)} vectors removed") else: index = None to_add = [path[0] for path in paths_time if path not in to_keep] print(f'to_keep: {to_keep}') print(f'to_add: {to_add}') return index, to_add # -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE def split_chunks_by_tokens(documents, max_length=170, overlap=10): # Create an empty list to store the resized documents resized = [] # Iterate through the original documents list for doc in documents: encoded = tokenizer.encode(doc.page_content) if len(encoded) > max_length: remaining_encoded = tokenizer.encode(doc.page_content) while len(remaining_encoded) > 0: split_doc = Document(page_content=tokenizer.decode(remaining_encoded[:max(10, max_length)]), metadata=doc.metadata.copy()) resized.append(split_doc) remaining_encoded = remaining_encoded[max(10, max_length - overlap):] else: resized.append(doc) print(f"Number of chunks before resplitting: {len(documents)} \nAfter splitting: {len(resized)}") return resized # -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE def split_chunks_by_tokens_period(documents, max_length=170, overlap=10, min_chunk_size=20): # Create an empty list to store the resized documents resized = [] previous_file="" # Iterate through the original documents list for doc in documents: current_file = doc.metadata['source'] if current_file != previous_file: #chunk counting previous_file = current_file chunk_counter = 0 is_first_chunk = True # Keep track of the first chunk in the document encoded = tokenizer.encode(doc.page_content)#encode the current document if len(encoded) > max_length: remaining_encoded = encoded is_last_chunk = False while len(remaining_encoded) > 1 and not is_last_chunk: # Check for a period in the first 'overlap' tokens overlap_text = tokenizer.decode(remaining_encoded[:overlap])# Index by token period_index_b = overlap_text.find('.')# Index by character if len(remaining_encoded)>max_length + min_chunk_size: current_encoded = remaining_encoded[:max(10, max_length)] else: current_encoded = remaining_encoded[:max(10, max_length + min_chunk_size)] #if the last chunk is to small, concatenate it with the previous one is_last_chunk = True period_index_e = len(doc.page_content) # an amount of character that I am sure will be greater or equal to the max lengh of a chunk, could have done len(tokenizer.decode(current_encoded)) if len(remaining_encoded)>max_length+min_chunk_size:# If it is not the last sub chunk overlap_text_last = tokenizer.decode(current_encoded[-overlap:]) period_index_last = overlap_text_last.find('.') if period_index_last != -1 and period_index_last < len(overlap_text_last) - 1: #print(f"period index last found at {period_index_last}") period_index_e = period_index_last - len(overlap_text_last) + 1 #print(f"period_index_e :{period_index_e}") #print(f"last :{overlap_text_last}") if not is_first_chunk:#starting after the period in overlap if period_index_b == -1:# Period not found in overlap #print(". not found in overlap") split_doc = Document(page_content=tokenizer.decode(current_encoded)[:period_index_e], metadata=doc.metadata.copy()) # Keep regular splitting else: if is_last_chunk : #not the first but the last split_doc = Document(page_content=tokenizer.decode(current_encoded)[period_index_b+1:], metadata=doc.metadata.copy()) #print("Should start after \".\"") else: split_doc = Document(page_content=tokenizer.decode(current_encoded)[period_index_b+1:period_index_e], metadata=doc.metadata.copy()) # Split at the begining and the end else:#first chunk split_doc = Document(page_content=tokenizer.decode(current_encoded)[:period_index_e], metadata=doc.metadata.copy()) # split only at the end if its first chunk if 'titles' in split_doc.metadata: chunk_counter += 1 split_doc.metadata['chunk_id'] = chunk_counter #A1 We could round chunk length in token if we ignore the '.' position in the overlap and save time of computation split_doc.metadata['token_length'] = len(tokenizer.encode(split_doc.page_content)) resized.append(split_doc) remaining_encoded = remaining_encoded[max(10, max_length - overlap):] is_first_chunk = False #print(len(tokenizer.encode(split_doc.page_content)), split_doc.page_content, "\n-----------------") elif len(encoded)>min_chunk_size:#ignore the chunks that are too small #print(f"◀Document:{{ {doc.page_content} }} was not added because to short▶") if 'titles' in doc.metadata:#check if it was splitted by or split_docx chunk_counter += 1 doc.metadata['chunk_id'] = chunk_counter doc.metadata['token_length'] = len(encoded) resized.append(doc) print(f"Number of chunks before resplitting: {len(documents)} \nAfter splitting: {len(resized)}") return resized # -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE def split_doc_in_chunks(input_folder): docs = [] for i, filename in enumerate(input_folder): path = filename#os.path.join(input_folder, filename) print(f"Treating file {i}/{len(input_folder)}") # Select the appropriate document loader chunks=[] if path.endswith(".pdf"): try: print("Treatment of pdf file", path) raw_chuncks = split_pdf(path, input_folder) chunks = group_chunks_by_section(raw_chuncks) print(f"Document splitted in {len(chunks)} chunks") # for chunk in chunks: # print(f"\n\n____\n\n\nPDF CONTENT: \n{chunk.page_content}\ntitle: {chunk.metadata['title']}\nFile Name: {chunk.metadata['filename']}\n\n") except Exception as e: print("Error while splitting the pdf file: ", e) elif path.endswith(".docx"): try: print ("Treatment of docx file", path) raw_chuncks = split_docx(path, input_folder) #print(f"RAW :\n***\n{raw_chuncks}") chunks = group_chunks_by_section(raw_chuncks) print(f"Document splitted in {len(chunks)} chunks") #if "cards-Jan 2022-SP.docx" in path: #for chunk in chunks: #print(f"\n\n____\n\n\nDOCX CONTENT: \n{chunk.page_content}\ntitle: {chunk.metadata['title']}\nFile Name: {chunk.metadata['filename']}\n\n") except Exception as e: print("Error while splitting the docx file: ", e) elif path.endswith(".doc"): try: loader = UnstructuredFileLoader(path) # Load the documents and split them in chunks chunks = loader.load_and_split(text_splitter=text_splitter) counter, counter2 = collections.Counter(), collections.Counter() filename = os.path.basename(path) # Define a unique id for each chunk for chunk in chunks: chunk.metadata["filename"] = filename.split("/")[-1] chunk.metadata["file_directory"] = filename.split("/")[:-1] chunk.metadata["filetype"] = filename.split(".")[-1] if "page" in chunk.metadata: counter[chunk.metadata['page']] += 1 for i in range(len(chunks)): counter2[chunks[i].metadata['page']] += 1 chunks[i].metadata['source'] = filename else: if len(chunks) == 1: chunks[0].metadata['source'] = filename #The file type is not supported (e.g. .xlsx) except Exception as e: print(f"An error occurred: {e}") try: if len(chunks)>0: docs += chunks except NameError as e: print(f"An error has occured: {e}") return docs # -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE def resplit_by_end_of_sentence(docs): print("❌❌\nResplitting docs by end of sentence\n❌❌") resized_docs = split_chunks_by_tokens_period(docs, max_length=200, overlap=40, min_chunk_size=20) try: # add chunk title to all resplitted chunks #todo move this to split_chunks_by_tokens_period(inject_title = True) with a boolean parameter cur_source = "" cpt_chunk = 1 for resized_doc in resized_docs: try: title = resized_doc.metadata['titles'].split(' ~~ ')[-2] #Getting the last title of the chunk and adding it to the content if it is not the case if title not in resized_doc.page_content: resized_doc.page_content = title + "\n" + resized_doc.page_content if cur_source == resized_doc.metadata["source"]: resized_doc.metadata['chunk_number'] = cpt_chunk else: cpt_chunk = 1 cur_source = resized_doc.metadata["source"] resized_doc.metadata['chunk_number'] = cpt_chunk except Exception as e:#either the title was notfound or title absent in metadata print("An error occured: ", e) #print(f"METADATA:\n{resized_doc.metadata}") cpt_chunk += 1 except Exception as e: print('AN ERROR OCCURRED: ', e) return resized_docs # -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE def build_index(docs, index, output_folder): if len(docs) > 0: if index is not None: # Compute the embedding of each chunk and index these chunks new_index = FAISS.from_documents(docs, embeddings) index.merge_from(new_index) else: index = FAISS.from_documents(docs, embeddings) with tempfile.TemporaryDirectory() as temp_dir: index.save_local(temp_dir) for f in os.listdir(temp_dir): output_folder.upload_file(f, os.path.join(temp_dir, f))