from unstructured.partition.auto import partition from unstructured.chunking.title import chunk_by_title from unstructured.chunking.basic import chunk_elements from unstructured.documents.elements import Element, Title, CompositeElement from unstructured.staging.base import convert_to_dataframe from typing import Type, List, Literal, Tuple from unstructured.cleaners.core import replace_unicode_quotes, clean_non_ascii_chars, clean_ordered_bullets, group_broken_paragraphs, replace_unicode_quotes, clean, clean_trailing_punctuation, remove_punctuation, bytes_string_to_string import gradio as gr import time import pandas as pd import re import gzip import pickle from pydantic import BaseModel, Field from tools.helper_functions import get_file_path_end, get_file_path_end_with_ext # Creating an alias for pandas DataFrame using Type PandasDataFrame = Type[pd.DataFrame] # %% # pdf partitioning strategy vars pdf_partition_strat = "ocr_only" # ["fast", "ocr_only", "hi_res"] # %% # Element metadata modification vars meta_keys_to_filter = ["file_directory", "filetype"] element_types_to_filter = ['UncategorizedText', 'Header'] # %% # Clean function vars bytes_to_string=False replace_quotes=True clean_non_ascii=False clean_ordered_list=True group_paragraphs=True trailing_punctuation=False all_punctuation=False clean_text=True extra_whitespace=True dashes=True bullets=True lowercase=False # %% # Chunking vars minimum_chunk_length = 2000 start_new_chunk_after_end_of_this_element_length = 2000 hard_max_character_length_chunks = 3000 multipage_sections=True overlap_all=True include_orig_elements=True # %% class Document(BaseModel): """Class for storing a piece of text and associated metadata. Implementation adapted from Langchain code: https://github.com/langchain-ai/langchain/blob/master/libs/core/langchain_core/documents/base.py""" page_content: str """String text.""" metadata: dict = Field(default_factory=dict) """Arbitrary metadata about the page content (e.g., source, relationships to other documents, etc.). """ type: Literal["Document"] = "Document" # %% def create_title_id_dict(elements:List[Element]): # Assuming the object is stored in a variable named 'elements_list' titles = [item.text for item in elements if isinstance(item, Title)] #### Get all elements under these titles chapter_ids = {} for element in elements: for chapter in titles: if element.text == chapter and element.category == "Title": chapter_ids[element._element_id] = chapter break chapter_to_id = {v: k for k, v in chapter_ids.items()} return chapter_ids, chapter_to_id # %% def filter_elements(elements:List[Element], excluded_elements: List[str] = ['']): """ Filter out elements from a list based on their categories. Args: elements: The list of elements to filter. excluded_elements: A list of element categories to exclude. Returns: A new list containing the filtered elements. """ filtered_elements = [] for element in elements: if element.category not in excluded_elements: filtered_elements.append(element) return filtered_elements # %% def remove_keys_from_meta( elements: List[Element], meta_remove_keys: List[str], excluded_element_types: List[str] = [] ) -> List[Element]: ''' Remove specified metadata keys from an Unstructured Element object ''' for element in elements: if element.category not in excluded_element_types: for key in meta_remove_keys: try: del element.metadata.__dict__[key] # Directly modify metadata except KeyError: print(f"Key '{key}' not found in element metadata.") return elements def filter_elements_and_metadata( elements: List[Element], excluded_categories: List[str] = [], meta_remove_keys: List[str] = [], ) -> List[Element]: """ Filters elements based on categories and removes specified metadata keys. Args: elements: The list of elements to process. excluded_categories: A list of element categories to exclude. meta_remove_keys: A list of metadata keys to remove. Returns: A new list containing the processed elements. """ filtered_elements = [] for element in elements: if element.category not in excluded_categories: for key in meta_remove_keys: try: del element.metadata.__dict__[key] except KeyError: # Better logging/error handling instead of just printing # Use a proper logger or raise a warning/exception pass filtered_elements.append(element) return filtered_elements # %% def add_parent_title_to_meta(elements:List[Element], chapter_ids:List[str], excluded_element_types:List[str]=['']) -> List[Element]: ''' Add parent title to Unstructured metadata elements ''' for element in elements: if element.category in excluded_element_types: pass else: meta = element.metadata.to_dict() if "parent_id" in meta and meta["parent_id"] in chapter_ids and "title_name" not in meta: title_name = chapter_ids[meta["parent_id"]] # Directly modify the existing element metadata object element.metadata.title_name = title_name return elements def chunk_all_elements(elements:List[Element], file_name_base:str, chunk_type:str = "Basic_chunking", minimum_chunk_length:int=minimum_chunk_length, start_new_chunk_after_end_of_this_element_length:int=start_new_chunk_after_end_of_this_element_length, hard_max_character_length_chunks:int=hard_max_character_length_chunks, multipage_sections:bool=multipage_sections, overlap_all:bool=overlap_all, include_orig_elements:bool=include_orig_elements): ''' Use Unstructured.io functions to chunk an Element object by Title or across all elements. ''' output_files = [] output_summary = "" chapter_ids, chapter_to_id = create_title_id_dict(elements) ### Break text down into chunks try: if chunk_type == "Chunk within title": chunks = chunk_by_title( elements, include_orig_elements=include_orig_elements, combine_text_under_n_chars=minimum_chunk_length, new_after_n_chars=start_new_chunk_after_end_of_this_element_length, max_characters=hard_max_character_length_chunks, multipage_sections=multipage_sections, overlap_all=overlap_all ) else: chunks = chunk_elements( elements, include_orig_elements=include_orig_elements, new_after_n_chars=start_new_chunk_after_end_of_this_element_length, max_characters=hard_max_character_length_chunks, overlap_all=overlap_all ) except Exception as output_summary: print(output_summary) return output_summary, output_files, file_name_base chunk_sections, chunk_df, chunks_out = element_chunks_to_document(chunks, chapter_ids) file_name_suffix = "_chunk" # The new file name does not overwrite the old file name as the 'chunked' elements are only used as an output, and not an input to other functions output_summary, output_files, file_name_base_new = export_elements_as_table_to_file(chunks_out, file_name_base, file_name_suffix, chunk_sections) return output_summary, output_files, file_name_base # %% def element_chunks_to_document(chunks:CompositeElement, chapter_ids:List[str]) -> Tuple[List[Document], PandasDataFrame, List[str]]: ''' Take an Unstructured.io chunk_by_title output with the original parsed document elements and turn it into a Document format commonly used by vector databases, and a Pandas dataframe. ''' chunk_sections = [] current_title_id = '' current_title = '' last_page = '' chunk_df_list = [] for chunk in chunks: chunk_meta = chunk.metadata.to_dict() true_element_ids = [] element_categories = [] titles = [] titles_id = [] if "page_number" in chunk_meta: last_page = chunk_meta["page_number"] chunk_text = chunk.text #chunk_page_number = chunk.metadata.to_dict()["page_number"] # If the same element text is found, add the element_id to the chunk (NOT PERFECT. THIS WILL FAIL IF THE SAME TEXT IS SEEN MULTIPL TIMES) for element in chunk.metadata.orig_elements: #element_text = element.text element_id = element._element_id element_category = element.category element_meta = element.metadata.to_dict() if "page_number" in element_meta: element_page_number = element_meta["page_number"] last_page = element_page_number true_element_ids.append(element_id) element_categories.append(element_category) # Set new metadata for chunk if "page_number" in element_meta: chunk_meta["last_page_number"] = last_page chunk_meta["true_element_ids"] = true_element_ids for loop_id in chunk_meta['true_element_ids']: if loop_id in chapter_ids: current_title = chapter_ids[loop_id] current_title_id = loop_id titles.append(current_title) titles_id.append(current_title_id) chunk_meta['titles'] = titles chunk_meta['titles_id'] = titles_id # Remove original elements data for documents chunk_meta.pop('orig_elements') chunk_dict_for_df = chunk_meta.copy() chunk_dict_for_df['text'] = chunk.text chunk_df_list.append(chunk_dict_for_df) chunk_doc = [Document(page_content=chunk_text, metadata=chunk_meta)] chunk_sections.extend(chunk_doc) ## Write metadata back to elements chunk.metadata.__dict__ = chunk_meta chunk_df = pd.DataFrame(chunk_df_list) # print("Doc format: ", chunk_sections) return chunk_sections, chunk_df, chunks # %% def write_elements_to_documents(elements:List[Element]): ''' Take Unstructured.io parsed elements and write it into a 'Document' format commonly used by vector databases ''' doc_sections = [] for element in elements: meta = element.metadata.to_dict() meta["type"] = element.category meta["element_id"] = element._element_id element_doc = [Document(page_content=element.text, metadata= meta)] doc_sections.extend(element_doc) #print("Doc format: ", doc_sections) return doc_sections # %% def clean_elements(elements:List[Element], dropdown_options: List[str] = [''], output_name:str = "combined_elements", bytes_to_string:bool=False, replace_quotes:bool=True, clean_non_ascii:bool=False, clean_ordered_list:bool=True, group_paragraphs:bool=True, trailing_punctuation:bool=False, all_punctuation:bool=False, clean_text:bool=True, extra_whitespace:bool=True, dashes:bool=True, bullets:bool=True, lowercase:bool=False) -> List[Element]: ''' Apply Unstructured cleaning processes to a list of parse elements. ''' out_files = [] output_summary = "" # Set variables to True based on dropdown selections for option in dropdown_options: if option == "Convert bytes to string": bytes_to_string = True elif option == "Replace quotes": replace_quotes = True elif option == "Clean non ASCII": clean_non_ascii = True elif option == "Clean ordered list": clean_ordered_list = True elif option == "Group paragraphs": group_paragraphs = True elif option == "Remove trailing punctuation": trailing_punctuation = True elif option == "Remove all punctuation": all_punctuation = True elif option == "Clean text": clean_text = True elif option == "Remove extra whitespace": extra_whitespace = True elif option == "Remove dashes": dashes = True elif option == "Remove bullets": bullets = True elif option == "Make lowercase": lowercase = True cleaned_elements = elements.copy() for element in cleaned_elements: try: if element: # Check if element is not None or empty if bytes_to_string: element.apply(bytes_string_to_string) if replace_quotes: element.apply(replace_unicode_quotes) if clean_non_ascii: element.apply(clean_non_ascii_chars) if clean_ordered_list: element.apply(clean_ordered_bullets) if group_paragraphs: element.apply(group_broken_paragraphs) if trailing_punctuation: element.apply(clean_trailing_punctuation) if all_punctuation: element.apply(remove_punctuation) if group_paragraphs: element.apply(group_broken_paragraphs) if clean_text: element.apply(lambda x: clean(x, extra_whitespace=extra_whitespace, dashes=dashes, bullets=bullets, lowercase=lowercase)) except Exception as e: print(e) element = element alt_out_message, out_files, output_file_base = export_elements_as_table_to_file(cleaned_elements, output_name, file_name_suffix="_clean") output_summary = "Text elements successfully cleaned." print(output_summary) return cleaned_elements, output_summary, out_files, output_file_base # %% [markdown] def export_elements_as_table_to_file(elements:List[Element], file_name_base:str, file_name_suffix:str="", chunk_documents:List[Document]=[]): ''' Export elements as as a table. ''' output_summary = "" out_files = [] # Convert to dataframe format out_table = convert_to_dataframe(elements) # If the file suffix already exists in the output file name, don't add it again. if file_name_suffix not in file_name_base: out_file_name_base = file_name_base + file_name_suffix else: out_file_name_base = file_name_base out_file_name = "output/" + out_file_name_base + ".csv" out_table.to_csv(out_file_name) out_files.append(out_file_name) # Convert to document format if chunk_documents: out_documents = chunk_documents else: out_documents = write_elements_to_documents(elements) out_file_name_docs = "output/" + out_file_name_base + "_docs.pkl.gz" with gzip.open(out_file_name_docs, 'wb') as file: pickle.dump(out_documents, file) out_files.append(out_file_name_docs) output_summary = "File successfully exported." return output_summary, out_files, out_file_name_base # # Partition PDF def get_file_type(filename): pattern = r"\.(\w+)$" # Match a dot followed by one or more word characters at the end of the string match = re.search(pattern, filename) if match: file_type = match.group(1) # Extract the captured file type (without the dot) print(file_type) # Output: "png" else: print("No file type found.") return file_type # %% def partition_file(filenames:List[str], pdf_partition_strat:str = pdf_partition_strat, progress = gr.Progress()): ''' Partition document files into text elements using the Unstructured package. Currently supports PDF, docx, pptx, html, several image file types, text document types, email messages, code files. ''' out_message = "" combined_elements = [] out_files = [] for file in progress.tqdm(filenames, desc="Partitioning files", unit="files"): try: tic = time.perf_counter() print(file) file_name = get_file_path_end_with_ext(file) file_name_base = get_file_path_end(file) file_type = get_file_type(file_name) image_file_type_list = ["jpg", "jpeg", "png", "heic"] if file_type in image_file_type_list: print("File is an image. Using OCR method to partition.") file_elements = partition(file, strategy="ocr_only") else: file_elements = partition(file, strategy=pdf_partition_strat) toc = time.perf_counter() new_out_message = f"Successfully partitioned file: {file_name} in {toc - tic:0.1f} seconds\n" print(new_out_message) out_message = out_message + new_out_message combined_elements.extend(file_elements) except Exception as e: new_out_message = f"Failed to partition file: {file_name} due to {e}. Partitioning halted." print(new_out_message) out_message = out_message + new_out_message break out_table = convert_to_dataframe(combined_elements) # If multiple files, overwrite default file name for outputs if len(filenames) > 1: file_name_base = "combined_files" alt_out_message, out_files, output_file_base = export_elements_as_table_to_file(combined_elements, file_name_base, file_name_suffix="_elements") return out_message, combined_elements, out_files, output_file_base, out_table # %% def modify_metadata_elements(elements_out_cleaned:List[Element], meta_keys_to_filter:List[str]=meta_keys_to_filter, element_types_to_filter:List[str]=element_types_to_filter) -> List[Element]: ''' Take an element object, add parent title names to metadata. Remove specified metadata keys or element types from element list. ''' chapter_ids, chapter_to_id = create_title_id_dict(elements_out_cleaned.copy()) elements_out_meta_mod = add_parent_title_to_meta(elements_out_cleaned.copy(), chapter_ids) elements_out_meta_mod_meta_filt = remove_keys_from_meta(elements_out_meta_mod.copy(), meta_keys_to_filter) elements_out_filtered_meta_mod = filter_elements(elements_out_meta_mod_meta_filt, element_types_to_filter) return elements_out_filtered_meta_mod # %% # file_stub = "C:/Users/SPedrickCase/OneDrive - Lambeth Council/Apps/doc_rag_prep/examples/" # filenames = [] # pdf_filename = [file_stub + "Lambeth_2030-Our_Future_Our_Lambeth_foreword.pdf"] # filenames.extend(pdf_filename) # html_filename = [file_stub + "transport-strategy.html"] # filenames.extend(html_filename) # docx_filename = [file_stub + "FINAL Policy and Procedure for Writing Housing Policies.docx"] # filenames.extend(docx_filename) # out_message, elements_parse = partition_file(filenames=filenames, pdf_partition_strat="ocr_only") # for element in elements_parse[:10]: # print(f"{element.category.upper()}: {element.text} - Metadata: {element.metadata.to_dict()}") # elements_out = elements_parse.copy() # %% [markdown] # ### Process with document layout detection - fast strategy # # The "fast" strategy will extract the text using pdfminer and process the raw text with partition_text. If the PDF text is not extractable, partition_pdf will fall back to "ocr_only". We recommend using the "fast" strategy in most cases where the PDF has extractable text. # elements_out_parse = partition_pdf(filename=filename, strategy="fast") # for element in elements_out_parse[:10]: # print(f"{element.category.upper()}: {element.text} - Metadata: {element.metadata.to_dict()}") # elements_out = elements_out_parse.copy() # ### OCR only # # The "ocr_only" strategy runs the document through Tesseract for OCR and then runs the raw text through partition_text. Currently, "hi_res" has difficulty ordering elements for documents with multiple columns. If you have a document with multiple columns that does not have extractable text, we recommend using the "ocr_only" strategy. "ocr_only" falls back to "fast" if Tesseract is not available and the document has extractable text. # elements_out_parse = partition_pdf(filename=filename, strategy="ocr_only") # for element in elements_out_parse[:10]: # print(f"{element.category.upper()}: {element.text} - Metadata: {element.metadata.to_dict()}") # elements_out = elements_out_parse.copy() # ### Hi-res partitioning # # The "hi_res" strategy will identify the layout of the document using detectron2. The advantage of “hi_res” is that it uses the document layout to gain additional information about document elements. We recommend using this strategy if your use case is highly sensitive to correct classifications for document elements. If detectron2 is not available, the "hi_res" strategy will fall back to the "ocr_only" strategy. # elements_out = partition_pdf(filename=filename, strategy="hi_res") # for element in elements_out[:10]: # print(f"{element.category.upper()}: {element.text} - Metadata: {element.metadata.to_dict()}") # %% [markdown] # ## Clean data # %% # elements_out_cleaned = clean_elements(elements_out.copy(), bytes_to_string=False, # replace_quotes=True , # clean_non_ascii=False, # clean_ordered_list=True , # group_paragraphs=True, # trailing_punctuation=False, # all_punctuation=False, # clean_text=True , # extra_whitespace=True, # dashes=True , # bullets=True , # lowercase=False) # %% [markdown] # ## Add/remove elements to/from metadata # %% [markdown] # ### Write to table, dictionary, document format # %% ### Dataframe format # elements_out_filtered_df = convert_to_dataframe(elements_out_filtered_meta_mod) # elements_out_filtered_df.to_csv("table.csv") # elements_out_filtered_df.head(6) # # %% # ### Dictionary format # elements_out_filtered_dict = convert_to_dict(elements_out_filtered_meta_mod) # elements_out_filtered_dict[20] # # %% [markdown] # # ### Document format for embeddings # # %% # doc_sections = write_elements_to_documents(elements_out_filtered_meta_mod, element_types_to_filter) # doc_sections[0:10] # # %% [markdown] # # ### Break text down into chunks # # %% # chunks_by_title = chunk_by_title( # elements_out_filtered_meta_mod, # include_orig_elements=True, # combine_text_under_n_chars=minimum_chunk_length, # new_after_n_chars=start_new_chunk_after_end_of_this_element_length, # max_characters=hard_max_character_length_chunks, # multipage_sections=True, # overlap_all=True # ) # chunk_sections, chunk_df = element_chunks_to_document(chunks_by_title, chapter_ids) # chunk_df.to_csv("chunked_df.csv") # print(chunk_sections[2]) # # %% # chunks_basic = chunk_elements( # elements_out_filtered_meta_mod, # include_orig_elements=True, # new_after_n_chars=start_new_chunk_after_end_of_this_element_length, # max_characters=hard_max_character_length_chunks, # overlap_all=True # ) # chunk_basic_sections, chunk_basic_df = element_chunks_to_document(chunks_basic, chapter_ids) # chunk_basic_df.to_csv("chunked_basic_df.csv") # %% [markdown] # # Partition Word document # # You cannot get location metadata for bounding boxes from word documents # %% # word_filename = "../examples/FINAL Policy and Procedure for Writing Housing Policies.docx" # # %% # docx_elements = partition(filename=word_filename) # for element in docx_elements: # print(f"{element.category.upper()}: {element.text} - Metadata: {element.metadata.to_dict()}") # # %% # docx_elements[5].text # # %% # docx_elements[5].category # # %% # docx_elements[5].metadata.to_dict() # # %% [markdown] # # ## Find elements associated with chapters # # %% # chapter_ids, chapter_to_id = create_title_id_dict(docx_elements) # chapter_ids # # %% # doc_sections = write_elements_to_documents(docx_elements.copy(), chapter_ids) # # %% # doc_sections # # %% [markdown] # # ### Chunk documents # # %% # chunks = chunk_by_title( # docx_elements, # include_orig_elements=False, # combine_text_under_n_chars=0, # new_after_n_chars=500, # max_characters=1000, # multipage_sections=True, # overlap_all=True # ) # # %% # print(chunks) # # %% # chunk_sections = element_chunks_to_document(chunks.copy(), docx_elements.copy(), chapter_ids) # # %% # chunk_sections[5].page_content # # %% # chunk_sections[5].metadata["true_element_ids"] # # %% # for element in docx_elements: # if element._element_id in chunk_sections[5].metadata["true_element_ids"]: # print(element.text) # # %% [markdown] # # # Partition PPTX document # # %% # pptx_filename = "../examples/LOTI presentation Jan 2024.pptx" # # %% # pptx_elements = partition(filename=pptx_filename) # for element in pptx_elements[:10]: # print(f"{element.category.upper()}: {element.text} - Metadata: {element.metadata.to_dict()}") # # %% # chapter_ids, chapter_to_id = create_title_id_dict(pptx_elements) # chapter_ids # # %% # pptx_sections = write_elements_to_documents(pptx_elements.copy(), chapter_ids) # # %% # pptx_sections # # %% # pptx_chunks = chunk_by_title( # pptx_elements, # include_orig_elements=False, # combine_text_under_n_chars=0, # new_after_n_chars=500, # max_characters=1000, # multipage_sections=True, # overlap_all=True # ) # # %% # pptx_chunk_sections = element_chunks_to_document(pptx_chunks.copy(), pptx_elements.copy(), chapter_ids) # # %% [markdown] # # ### Load documents into a vectorDB (Not necessary) # # %% # import chromadb # # %% # client = chromadb.PersistentClient(path="chroma_tmp", settings=chromadb.Settings(allow_reset=True)) # client.reset() # # %% # collection = client.create_collection( # name="policy_statements", # metadata={"hnsw:space": "cosine"} # ) # # %% # chapter_ids # # %% # for element in docx_elements: # parent_id = element.metadata.parent_id # #print(element.text) # #print(parent_id) # #print(element.metadata.to_dict()) # if parent_id: # try: # print(parent_id) # chapter = chapter_ids[parent_id] # print(chapter) # except KeyError: # chapter = "None" # else: # chapter = "None" # collection.add( # documents=[element.text], # ids=[element._element_id], # metadatas=[{"chapter": chapter}] # ) # # %% [markdown] # # #### See the elements in the VectorDB and perform hybrid search # # %% # results = collection.peek() # print(results["documents"]) # # %% # print(collection.metadata) # # %% # import json # result = collection.query( # query_texts=["What should policies do?"], # n_results=2, # where={"chapter": '3.0 Policy Statements'}, # ) # print(json.dumps(result, indent=2)) # # %% # collection = client.create_collection( # name="policy_statements_chunk", # metadata={"hnsw:space": "cosine"} # ) # # %% # for element in chunks: # parent_id = element.metadata.parent_id # #print(element.text) # #print(parent_id) # #print(element.metadata.to_dict()) # if parent_id: # try: # print(parent_id) # chapter = chapter_ids[parent_id] # print(chapter) # except KeyError: # chapter = "None" # else: # chapter = "None" # print(element._element_id) # collection.add( # documents=[element.text], # ids=[element.orig_elements], # metadatas=[{"chapter": chapter}] # ) # # %% [markdown] # # # Partition HTML # # %% # html_filename = "../examples/transport-strategy.html" # # %% # html_elements = partition(filename=html_filename) # for element in html_elements[:10]: # print(f"{element.category.upper()}: {element.text} - Metadata: {element.metadata.to_dict()}") # # %% [markdown] # # # Partition image # # %% # img_filename = "../examples/example_complaint_letter.jpg" # # %% # img_elements = partition(filename=img_filename) # for element in img_elements[:10]: # print(f"{element.category.upper()}: {element.text} - Metadata: {element.metadata.to_dict()}") # # %% [markdown] # # # Partition XLSX # # %% # xlsx_filename = "../examples/fuel-poverty-sub-regional-tables-2020-2018-data.xlsx" # # %% # xlsx_elements = partition(filename=xlsx_filename) # for element in xlsx_elements[:10]: # print(f"{element.category.upper()}: {element.text} - Metadata: {element.metadata.to_dict()}") # # %% [markdown] # # # Partition .py # # %% # py_filename = "../examples/app.py" # # %% # py_elements = partition(filename=py_filename) # for element in py_elements[:10]: # print(f"{element.category.upper()}: {element.text} - Metadata: {element.metadata.to_dict()}")