Create split_files_to_excel.py

split_files_to_excel.py

@@ -0,0 +1,474 @@
+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))