Update split_files_to_excel.py

split_files_to_excel.py

@@ -68,7 +68,7 @@ text_splitter = CharacterTextSplitter(
 
 def function_split_call(fi_input, dropdown, choice, chunk_size):
     if choice == "Intelligent split":
-        return split_in_df(fi_input)
+        return split_in_df(fi_input, nb_pages)
     elif choice == "Non intelligent split":
         return non_intelligent_split(fi_input, chunk_size)
     else:
@@ -78,7 +78,7 @@ def change_textbox(dropdown,radio):
     if len(dropdown) == 0 :
         dropdown = ["introduction", "objective", "summary", "conclusion"]
     if radio == "Intelligent split by keywords":
-        return gr.Dropdown(dropdown, multiselect=True, visible=True, allow_custom_value=True), gr.Number(visible=False)
+        return gr.Dropdown(dropdown, multiselect=True, visible=True, allow_custom_value=True), gr.Number(label="First pages to keep (0 for all)", value=2, interactive=True, visible=True)
     elif radio == "Non intelligent split":
         return gr.Dropdown(dropdown, visible=False),gr.Number(label="Chunk size", value=1000, interactive=True, visible=True)
     else:
@@ -464,7 +464,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, base_folders):
+def split_doc_in_chunks(input_folder, base_folders, nb_pages):
     docs = []
     for i, filename in enumerate(input_folder):
         path = filename#os.path.join(input_folder, filename)
@@ -478,7 +478,15 @@ def split_doc_in_chunks(input_folder, base_folders):
                 for raw_chunk in raw_chunks:
                     print(f"BASE zzzzz LIST : {base_folders} = i = {i}")
                     raw_chunk.metadata["Base Folder"] = base_folders[i]
-                chunks = group_chunks_by_section(raw_chunks)
+                sb_chunks = group_chunks_by_section(raw_chunks)
+                if nb_pages > 0:
+                    for sb_chunk in sb_chunks:
+                        if int(sb_chunk.metadata["page_number"])<nb_pages:
+                            chunks.append(sb_chunk)
+                    else:
+                        break
+                else:
+                    chunks = sb_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")
@@ -579,7 +587,7 @@ def extract_zip(zip_path):
             zip_ref.extract(file_info.filename)
     return extracted_files
 
-def split_in_df(files):
+def split_in_df(files, nb_pages):
     processed_files = []
     base_folders = []
     print("Processing zip files...")
@@ -593,7 +601,7 @@ def split_in_df(files):
             base_folders.append("")
     print(f"BASE FOLDERS LIST : {base_folders}")
     print("Finished processing zip files\nSplitting files into chunks...")
-    documents = split_doc_in_chunks(processed_files, base_folders)
+    documents = split_doc_in_chunks(processed_files, base_folders, nb_pages)
     re_docs = resplit_by_end_of_sentence(documents, 1000, 100, 1500)
     print("Finished splitting")
     df = pd.DataFrame()