diff --git a/fieldExtraction/src/investment/preprocess.py b/fieldExtraction/src/investment/preprocess.py index 364ba6a..f88acb1 100644 --- a/fieldExtraction/src/investment/preprocess.py +++ b/fieldExtraction/src/investment/preprocess.py @@ -83,7 +83,7 @@ def one_to_n_exhibit_chunking(text_dict, filename) -> tuple[dict, dict]: """ exhibit_pages = preprocessing_funcs.get_exhibit_pages(text_dict, filename) exhibit_chunk_mapping = preprocessing_funcs.chunk_by_exhibit(text_dict, exhibit_pages) - exhibit_dict = {page : string_utils.get_exhibit_chunk(text_dict, exhibit_chunk_mapping, page) for page in exhibit_pages} + exhibit_dict = preprocessing_funcs.get_exhibit_dict(text_dict, exhibit_chunk_mapping) all_exhibit_headers = {page: get_exhibit_header(exhibit_dict[page], filename) for page in exhibit_pages} return exhibit_dict, all_exhibit_headers diff --git a/fieldExtraction/src/utils/string_utils.py b/fieldExtraction/src/utils/string_utils.py index 54024a9..070ac56 100644 --- a/fieldExtraction/src/utils/string_utils.py +++ b/fieldExtraction/src/utils/string_utils.py @@ -370,32 +370,6 @@ def is_empty(value, pd_mask=True): return value in empty_values -def get_exhibit_chunk(text_dict: dict, - exhibit_chunk_mapping: dict, - exhibit_page: str) -> str: - """ - Compiles text from a specific exhibit based on its starting page number. - - This function aggregates the text of all pages that are mapped to a specific exhibit - start page in a document. It uses the exhibit chunk mapping to determine which pages - belong to the exhibit defined by the given start page number, and concatenates their - texts into a single string. - - Parameters: - text_dict (dict): A dictionary where keys are page numbers and values are the text on those pages. - exhibit_chunk_mapping (dict): A dictionary mapping each page number to the starting page number of the exhibit it belongs to. - exhibit_page (str): The page number that marks the beginning of the exhibit to compile text for. - - Returns: - str: A string that combines all the text from the pages belonging to the specified exhibit, separated by newlines. - - Notes: - - The function assumes that `exhibit_chunk_mapping` correctly maps all page numbers in `text_dict` to their respective exhibit starting pages. - - It is important that `exhibit_page` exists as a key in the `exhibit_chunk_mapping` dictionary and corresponds to the starting page of an exhibit. - """ - return '\n'.join([page_text for page_num, page_text in text_dict.items() if exhibit_chunk_mapping[page_num] == exhibit_page]) - - def datetime_str(): return datetime.now().strftime("[%Y-%m-%d %H:%M:%S]") diff --git a/fieldExtraction/tests/get_exhibit_chunk_test.py b/fieldExtraction/tests/get_exhibit_chunk_test.py deleted file mode 100644 index 86a7c44..0000000 --- a/fieldExtraction/tests/get_exhibit_chunk_test.py +++ /dev/null @@ -1,34 +0,0 @@ - -from src.utils.string_utils import get_exhibit_chunk - -def test1(): - """Test non-empty text_dict with empty exhibit_pages.""" - text_dict = {'1': 'This', '2': 'is', '3': 'an', '4': 'exhibit', '5': 'test'} - exhibit_chunk_mapping = {'1': '1', '2': '2', '3': '3', '4': '4', '5': '5'} - exhibit_page_num = '4' - assert get_exhibit_chunk(text_dict, exhibit_chunk_mapping, exhibit_page_num) == 'exhibit' - -def test2(): - """Test when exhibit_pages contains some keys that are also in text_dict.""" - text_dict = {'1': 'Welcome', '2': 'to', '3': 'the', '4': 'exhibit'} - exhibit_chunk_mapping = {'1': '1', '2': '2', '3': '2', '4': '4'} - exhibit_page_num = '2' - expected_output = "to\nthe" - assert get_exhibit_chunk(text_dict, exhibit_chunk_mapping, exhibit_page_num) == expected_output - - -def test3(): - """Test multiple exhibit pages with transitions between them.""" - text_dict = {'1': 'Page 1', '2': 'Page 2', '3': 'Page 3', '4': 'Page 4', '5': 'Page 5'} - exhibit_chunk_mapping = {'1': '0', '2': '2', '3': '2', '4': '4', '5': '4'} - exhibit_page_num = '4' - expected_output = 'Page 4\nPage 5' - assert get_exhibit_chunk(text_dict, exhibit_chunk_mapping, exhibit_page_num) == expected_output - - -test1() - -test2() - -test3() -