54 lines
1.4 KiB
Python
54 lines
1.4 KiB
Python
 
|
"""Table-aware chunk grouping helpers."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from rag_cut.models import Block, BlockType
|
||
|
|
|
||
|
|
|
||
|
|
def group_tables_with_context(blocks: list[Block]) -> list[list[Block]]:
|
||
|
|
"""
|
||
|
|
Group each table with its contextual title paragraph and footnote.
|
||
|
|
Keeps: heading/intro → table → footnote as one atomic group when adjacent.
|
||
|
|
"""
|
||
|
|
if not blocks:
|
||
|
|
return []
|
||
|
|
|
||
|
|
groups: list[list[Block]] = []
|
||
|
|
current: list[Block] = []
|
||
|
|
i = 0
|
||
|
|
|
||
|
|
while i < len(blocks):
|
||
|
|
block = blocks[i]
|
||
|
|
if block.type != BlockType.TABLE:
|
||
|
|
current.append(block)
|
||
|
|
i += 1
|
||
|
|
continue
|
||
|
|
|
||
|
|
group: list[Block] = []
|
||
|
|
if current:
|
||
|
|
tail = current[-1]
|
||
|
|
if tail.type in {BlockType.HEADING, BlockType.PARAGRAPH}:
|
||
|
|
group.append(tail)
|
||
|
|
current = current[:-1]
|
||
|
|
if current:
|
||
|
|
groups.append(current)
|
||
|
|
current = []
|
||
|
|
|
||
|
|
group.append(block)
|
||
|
|
j = i + 1
|
||
|
|
while j < len(blocks):
|
||
|
|
nxt = blocks[j]
|
||
|
|
if nxt.type == BlockType.PARAGRAPH and len((nxt.text or "")) <= 320:
|
||
|
|
if any(k in (nxt.text or "") for k in ("注", "备注", "说明", "Note", "※")):
|
||
|
|
group.append(nxt)
|
||
|
|
j += 1
|
||
|
|
break
|
||
|
|
break
|
||
|
|
|
||
|
|
groups.append(group)
|
||
|
|
i = j
|
||
|
|
|
||
|
|
if current:
|
||
|
|
groups.append(current)
|
||
|
|
return groups
|