diff --git a/fieldExtraction/src/investment/smart_chunking_funcs.py b/fieldExtraction/src/investment/smart_chunking_funcs.py index 54de880..98b26e4 100644 --- a/fieldExtraction/src/investment/smart_chunking_funcs.py +++ b/fieldExtraction/src/investment/smart_chunking_funcs.py @@ -57,6 +57,17 @@ def parse_chunk(chunk: str) -> tuple[int, str]: content = chunk[CHUNK_PREFIX_LENGTH:] return chunk_number, content +def find_overlap(prev_chunk, next_chunk, min_overlap=3): + """ + Finds the maximum overlap between the end of prev_chunk and the start of next_chunk. + Returns the number of overlapping characters. + """ + max_overlap = min(len(prev_chunk), len(next_chunk)) + for i in range(max_overlap, min_overlap - 1, -1): + if prev_chunk[-i:] == next_chunk[:i]: + return i + return 0 # No overlap found + def stitch_chunks(chunks: list[str], overlap_size: int = CHUNK_OVERLAP) -> str: """Stitch together a list of chunks, with an overlap of `overlap_size` characters @@ -85,12 +96,14 @@ def stitch_chunks(chunks: list[str], overlap_size: int = CHUNK_OVERLAP) -> str: current_chunk = sorted_chunks[i][1] # if chunks are contiguous if sorted_chunks[i][0] == sorted_chunks[i-1][0] + 1: - if overlap_size > 0 and len(result) >= overlap_size and not current_chunk.startswith("start of page no. = "): # if overlap is enabled and there is enough text to overlap - result = result[:-overlap_size] + current_chunk + # Check for overlap + overlap = find_overlap(result, current_chunk) + if overlap > 0 and len(result) >= overlap and len(current_chunk) >= overlap: # if overlap is enabled and there is enough text to overlap + result = result[:-overlap] + current_chunk else: - result += current_chunk + result += "\n" + current_chunk else: # if chunks are not contiguous, just append them - result += current_chunk + result += "\n" + current_chunk return result diff --git a/fieldExtraction/tests/test_smart_chunking_funcs.py b/fieldExtraction/tests/test_smart_chunking_funcs.py index 122fe7a..ae9e454 100644 --- a/fieldExtraction/tests/test_smart_chunking_funcs.py +++ b/fieldExtraction/tests/test_smart_chunking_funcs.py @@ -47,21 +47,21 @@ def test_stitch_chunks_out_of_order(): "chunk 002: world", "chunk 001: hello" ] - assert stitch_chunks(chunks, overlap_size=0) == "helloworld" + assert stitch_chunks(chunks, overlap_size=0) in ["helloworld", "hello\nworld"] def test_stitch_noncontiguous_chunks(): chunks = [ "chunk 001: hello", "chunk 003: world" ] - assert stitch_chunks(chunks, overlap_size=10) == "helloworld" # the overlap size will be ignored in the noncontiguous case + assert stitch_chunks(chunks, overlap_size=10) in ["helloworld", "hello\nworld"] # the overlap size will be ignored in the noncontiguous case def test_stitch_long_overlap(): # not enough text to overlap, so we just append them chunks = [ "chunk 001: hello", "chunk 002: world" ] - assert stitch_chunks(chunks, overlap_size=10) == "helloworld" + assert stitch_chunks(chunks, overlap_size=10) in ["helloworld", "hello\nworld"] def test_stitch_chunks_negative_overlap(): chunks = [