Spaces:
Running
Running
| from unstructured.partition.pdf import partition_pdf | |
| output_path = "./content/" | |
| file_path = output_path + 'attention.pdf' | |
| def pdf_extraction(file_path): | |
| chunks = partition_pdf( | |
| filename = file_path, | |
| infer_table_structure = True, | |
| strategy = "fast", | |
| extract_image_block_types = ["Image"], | |
| extract_image_block_to_payload = True, | |
| chunking_strategy = "by_title", | |
| max_characters = 10000, | |
| combine_text_under_n_chars = 2000, | |
| new_after_n_chars = 6000, | |
| ) | |
| return chunks | |
| def extract_tables_and_text(chunks): | |
| tables = [] | |
| text = [] | |
| for chunk in chunks: | |
| # Check if the chunk is actually an object with metadata | |
| if not hasattr(chunk, 'metadata'): | |
| continue | |
| if 'Table' in str(type(chunk)): | |
| tables.append(chunk) | |
| if 'CompositeElement' in str(type(chunk)): | |
| text.append(chunk) | |
| return tables, text | |
| def get_images_base64(chunks): | |
| images_b64 = [] | |
| for chunk in chunks: | |
| if 'CompositeElement' in str(type(chunk)): | |
| chunk_els = chunk.metadata.orig_elements | |
| for el in chunk_els: | |
| if 'Image' in str(type(el)): | |
| images_b64.append(el.metadata.image_base64) | |
| return images_b64 | |
| def run_extraction_pipeline(file_path): | |
| raw_chunks = pdf_extraction(file_path) | |
| # NEW: Safety Filter - only keep items that are objects (have metadata) | |
| # This removes any stray strings or non-element items created by the partition | |
| valid_chunks = [chunk for chunk in raw_chunks if hasattr(chunk, 'metadata')] | |
| tables, text_chunks = extract_tables_and_text(valid_chunks) | |
| images = get_images_base64(valid_chunks) | |
| return tables, text_chunks, images |