Update split_files_to_excel.py

split_files_to_excel.py

@@ -455,7 +455,7 @@ def split_chunks_by_tokens_period(documents, max_length=170, overlap=10, min_chu
 
 # -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
 
-def split_doc_in_chunks(input_folder):
+def split_doc_in_chunks(input_folder, base_folders):
     docs = []
     for i, filename in enumerate(input_folder):
         path = filename#os.path.join(input_folder, filename)
@@ -465,8 +465,10 @@ def split_doc_in_chunks(input_folder):
         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)
+                raw_chunks = split_pdf(path, input_folder)
+                for raw_chunk in raw_chunks:
+                    raw_chunk.metadata["Base Folder"] = base_folders[i]
+                chunks = group_chunks_by_section(raw_chunks)
                 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")
@@ -475,9 +477,11 @@ def split_doc_in_chunks(input_folder):
         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)
+                raw_chunks = split_docx(path, input_folder)
+                for raw_chunk in raw_chunks:
+                    raw_chunk.metadata["Base Folder"] = base_folders[i]
+                #print(f"RAW :\n***\n{raw_chunks}")
+                chunks = group_chunks_by_section(raw_chunks)
                 print(f"Document splitted in {len(chunks)} chunks")
                 #if "cards-Jan 2022-SP.docx" in path:
                     #for chunk in chunks:
@@ -496,6 +500,7 @@ def split_doc_in_chunks(input_folder):
                     chunk.metadata["filename"] = filename.split("/")[-1]
                     chunk.metadata["file_directory"] = filename.split("/")[:-1]
                     chunk.metadata["filetype"] = filename.split(".")[-1]
+                    chunk.metadata["Base Folder"] = base_folders[i]
                     if "page" in chunk.metadata:
                         counter[chunk.metadata['page']] += 1
                         for i in range(len(chunks)):
@@ -566,15 +571,18 @@ def extract_zip(zip_path):
 
 def split_in_df(files):
     processed_files = []
+    base_folders = []
     print("Processing zip files...")
     for file_path in files:
         if file_path.endswith('.zip'):
             extracted_files = extract_zip(file_path)
             processed_files.extend(extracted_files)
+            base_folders.append(os.path.splitext(os.path.basename(file_path))[0])
         else:
             processed_files.append(file_path)
-    print("Finished processing zip files\Splitting files into chunks...")
-    documents = split_doc_in_chunks(processed_files)
+            base_folders.append("")
+    print("Finished processing zip files\nSplitting files into chunks...")
+    documents = split_doc_in_chunks(processed_files, base_folders)
     re_docs = resplit_by_end_of_sentence(documents, 1000, 100, 1500)
     print("Finished splitting")
     df = pd.DataFrame()
@@ -590,6 +598,7 @@ def split_in_df(files):
 
         doc_data["Token_Length"] = re_doc.metadata['token_length']
         doc_data["Titles"] = re_doc.metadata['titles'] if 'titles' in re_doc.metadata else ""
+        doc_data["Base Folder"] = re_doc.metadata["Base Folder"]
 
         # for key, value in zip(metadata_keys, metadata_values):
         #     doc_data[key] = value