314 lines
10 KiB
Python
314 lines
10 KiB
Python
"""Table helpers: normalization, header detection, Markdown export, cross-page merge."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import re
|
||
from typing import Any
|
||
|
||
FOOTNOTE_ROW_RE = re.compile(r"^[\s*※①②③④⑤]*(?:注[:：]?|备注[:：]?|说明[:：]?|Note[:：]?)", re.I)
|
||
DATA_FIRST_CELL_RE = re.compile(r"^[a-z_][a-z0-9_.-]*$", re.I)
|
||
TABLE_TITLE_RE = re.compile(
|
||
r"^(?:表\s*\d+[:：.]?|Table\s*\d+[:：.]?|图\s*\d+[:：.]?)?\s*.{2,80}$",
|
||
re.I,
|
||
)
|
||
KEYWORD_TERMS = (
|
||
"字段", "参数", "必填", "选填", "状态", "类型", "说明", "含义", "取值",
|
||
"field", "parameter", "required", "optional", "status", "description",
|
||
)
|
||
|
||
|
||
def normalize_cell(value: str | None) -> str:
|
||
"""Merge in-cell line breaks; escape pipe chars for Markdown tables."""
|
||
if not value:
|
||
return ""
|
||
text = str(value).replace("\r\n", "\n").replace("\r", "\n")
|
||
parts = [p.strip() for p in text.split("\n") if p.strip()]
|
||
merged = " ".join(parts) if parts else ""
|
||
return merged.replace("|", "\\|")
|
||
|
||
|
||
def normalize_rows(rows: list[list[str | None]]) -> list[list[str]]:
|
||
"""Pad columns and normalize every cell without losing row alignment."""
|
||
if not rows:
|
||
return []
|
||
cleaned = [[normalize_cell(c) for c in row] for row in rows]
|
||
col_count = max((len(r) for r in cleaned), default=0)
|
||
return [row + [""] * (col_count - len(row)) for row in cleaned]
|
||
|
||
|
||
def looks_like_table(rows: list[list[str | None]]) -> bool:
|
||
"""Return True only for rows that have a real table-like grid."""
|
||
if len(rows) < 2:
|
||
return False
|
||
|
||
cleaned = normalize_rows(rows)
|
||
col_count = max((len(r) for r in cleaned), default=0)
|
||
if col_count < 2:
|
||
return False
|
||
|
||
non_empty_cells = sum(1 for row in cleaned for cell in row if cell)
|
||
rows_with_two_cells = sum(1 for row in cleaned if sum(1 for cell in row if cell) >= 2)
|
||
return non_empty_cells >= 4 and rows_with_two_cells >= 2
|
||
|
||
|
||
DESCRIPTION_HINTS = ("必填", "选填", "格式要求", "required", "optional", "格式", "用于标识")
|
||
QA_HEADER_TERMS = ("query", "question", "用户输入", "reference_output", "answer", "标准答案", "session")
|
||
|
||
|
||
def _row_fill(row: list[str]) -> int:
|
||
return sum(1 for c in row if c)
|
||
|
||
|
||
def looks_like_column_header_row(row: list[str]) -> bool:
|
||
"""True when a row looks like short spreadsheet column names."""
|
||
filled = [c.strip() for c in row if c and c.strip()]
|
||
if len(filled) < 2:
|
||
return False
|
||
if any(len(c) > 40 for c in filled):
|
||
return False
|
||
identifier_like = sum(
|
||
1
|
||
for c in filled
|
||
if DATA_FIRST_CELL_RE.match(c) or re.match(r"^[a-z][a-z0-9_]*$", c, re.I)
|
||
)
|
||
return identifier_like >= max(2, (len(filled) + 1) // 2)
|
||
|
||
|
||
def looks_like_description_row(row: list[str]) -> bool:
|
||
"""True when a row is a template field-description line (not data/header)."""
|
||
filled = [c.strip() for c in row if c and c.strip()]
|
||
if not filled:
|
||
return False
|
||
if max(len(c) for c in filled) >= 48:
|
||
return True
|
||
return sum(1 for c in filled if any(h in c for h in DESCRIPTION_HINTS)) >= 2
|
||
|
||
|
||
def detect_header_row_count(rows: list[list[str]]) -> int:
|
||
"""Detect 1-2 header rows from content patterns."""
|
||
if len(rows) < 2:
|
||
return 1 if rows else 0
|
||
|
||
first_fill = _row_fill(rows[0])
|
||
second_fill = _row_fill(rows[1]) if len(rows) > 1 else 0
|
||
if first_fill < 2:
|
||
return 0
|
||
|
||
if len(rows) > 2 and second_fill >= 2:
|
||
first_short = all(len(c) <= 24 for c in rows[0] if c)
|
||
second_short = all(len(c) <= 24 for c in rows[1] if c)
|
||
second_is_data = bool(rows[1][0]) and DATA_FIRST_CELL_RE.match(rows[1][0])
|
||
third_data_like = _row_fill(rows[2]) >= max(1, first_fill - 1)
|
||
if first_short and second_short and third_data_like and not second_is_data:
|
||
return 2
|
||
return 1
|
||
|
||
|
||
def detect_spreadsheet_layout(rows: list[list[str]]) -> dict[str, int]:
|
||
"""
|
||
Detect spreadsheet preamble/header/data boundaries (1-based row numbers).
|
||
|
||
Common template: row 1 = field descriptions, row 2 = column names, row 3+ = data.
|
||
"""
|
||
normalized = normalize_rows(rows)
|
||
if not normalized:
|
||
return {
|
||
"preamble_rows": 0,
|
||
"header_rows": 1,
|
||
"header_row_start": 1,
|
||
"header_row_end": 1,
|
||
"data_start_row": 2,
|
||
}
|
||
|
||
preamble = 0
|
||
header_index = 0
|
||
|
||
if (
|
||
len(normalized) >= 3
|
||
and looks_like_description_row(normalized[0])
|
||
and looks_like_column_header_row(normalized[1])
|
||
):
|
||
preamble = 1
|
||
header_index = 1
|
||
header_rows = 1
|
||
else:
|
||
header_rows = detect_header_row_count(normalized)
|
||
header_index = preamble
|
||
|
||
header_end_index = header_index + header_rows - 1
|
||
data_start_index = header_end_index + 1
|
||
|
||
return {
|
||
"preamble_rows": preamble,
|
||
"header_rows": header_rows,
|
||
"header_row_start": header_index + 1,
|
||
"header_row_end": header_end_index + 1,
|
||
"data_start_row": data_start_index + 1,
|
||
}
|
||
|
||
|
||
def is_qa_style_table(rows: list[list[str]], layout: dict[str, int] | None = None) -> bool:
|
||
"""True for evaluation/Q&A sheets where each row should become one chunk."""
|
||
if not rows:
|
||
return False
|
||
layout = layout or detect_spreadsheet_layout(rows)
|
||
h_start = layout["header_row_start"] - 1
|
||
h_end = layout["header_row_end"]
|
||
header_text = " ".join(
|
||
(cell or "").lower() for row in rows[h_start:h_end] for cell in row if cell
|
||
)
|
||
return sum(1 for term in QA_HEADER_TERMS if term in header_text) >= 2
|
||
|
||
|
||
def split_body_and_footnotes(rows: list[list[str]], header_rows: int) -> tuple[list[list[str]], list[list[str]], str]:
|
||
"""Separate data rows from trailing footnote rows."""
|
||
if header_rows >= len(rows):
|
||
return [], [], ""
|
||
|
||
body = rows[header_rows:]
|
||
footnote_rows: list[list[str]] = []
|
||
while body:
|
||
first_cell = (body[-1][0] if body[-1] else "") or ""
|
||
joined = " ".join(c for c in body[-1] if c)
|
||
if FOOTNOTE_ROW_RE.match(first_cell) or FOOTNOTE_ROW_RE.match(joined):
|
||
footnote_rows.insert(0, body.pop())
|
||
elif len(joined) <= 80 and any(k in joined for k in ("注", "备注", "说明", "Note")):
|
||
footnote_rows.insert(0, body.pop())
|
||
else:
|
||
break
|
||
|
||
footnotes = " ".join(" ".join(c for c in row if c) for row in footnote_rows).strip()
|
||
return body, footnote_rows, footnotes
|
||
|
||
|
||
def header_signature(rows: list[list[str]], header_rows: int) -> tuple[tuple[str, ...], ...]:
|
||
if header_rows <= 0:
|
||
return ()
|
||
return tuple(tuple(row) for row in rows[:header_rows])
|
||
|
||
|
||
def rows_to_markdown(
|
||
rows: list[list[str | None]],
|
||
header_rows: int | None = None,
|
||
include_footnotes: str = "",
|
||
) -> str:
|
||
"""Render rows as a standard Markdown table."""
|
||
if not looks_like_table(rows):
|
||
return ""
|
||
|
||
normalized = normalize_rows(rows)
|
||
if header_rows is None:
|
||
header_rows = detect_header_row_count(normalized)
|
||
|
||
body_rows, _, inline_footnotes = split_body_and_footnotes(normalized, header_rows)
|
||
data_rows = normalized[:header_rows] + body_rows
|
||
if not data_rows:
|
||
return ""
|
||
|
||
col_count = max(len(r) for r in data_rows)
|
||
lines: list[str] = []
|
||
for i, row in enumerate(data_rows):
|
||
padded = row + [""] * (col_count - len(row))
|
||
lines.append("| " + " | ".join(padded) + " |")
|
||
if i == header_rows - 1:
|
||
lines.append("| " + " | ".join(["---"] * col_count) + " |")
|
||
|
||
md = "\n".join(lines)
|
||
footnotes = include_footnotes or inline_footnotes
|
||
if footnotes:
|
||
md += f"\n\n*{footnotes}*"
|
||
return md
|
||
|
||
|
||
def guess_table_title(text: str) -> str | None:
|
||
"""Guess table title from a short preceding line."""
|
||
cleaned = normalize_cell(text)
|
||
if not cleaned or len(cleaned) > 120:
|
||
return None
|
||
if TABLE_TITLE_RE.match(cleaned):
|
||
return cleaned
|
||
if cleaned.endswith("表") or cleaned.endswith("列表") or cleaned.endswith("说明"):
|
||
return cleaned
|
||
if re.match(r"^表\s*\d+", cleaned):
|
||
return cleaned
|
||
return None
|
||
|
||
|
||
def extract_table_keywords(*texts: str, limit: int = 20) -> list[str]:
|
||
"""Extract retrieval keywords from table title, headers and body."""
|
||
source = " ".join(t for t in texts if t)
|
||
words = re.findall(r"[A-Za-z][A-Za-z0-9_-]{2,}|[\u4e00-\u9fff]{2,}", source)
|
||
seen: set[str] = set()
|
||
result: list[str] = []
|
||
for word in words:
|
||
if word in seen:
|
||
continue
|
||
seen.add(word)
|
||
result.append(word)
|
||
if len(result) >= limit:
|
||
break
|
||
for term in KEYWORD_TERMS:
|
||
if term.lower() in source.lower() and term not in seen:
|
||
result.append(term)
|
||
seen.add(term)
|
||
if len(result) >= limit:
|
||
break
|
||
return result[:limit]
|
||
|
||
|
||
def build_table_embedding_text(
|
||
*,
|
||
chapter: str = "",
|
||
table_title: str = "",
|
||
markdown: str = "",
|
||
description: str = "",
|
||
footnotes: str = "",
|
||
keywords: list[str] | None = None,
|
||
ocr_text: str = "",
|
||
) -> str:
|
||
"""Compose embedding text: chapter + title + markdown + description + keywords."""
|
||
parts: list[str] = []
|
||
if chapter:
|
||
parts.append(f"章节标题：{chapter}")
|
||
if table_title:
|
||
parts.append(f"表格标题：{table_title}")
|
||
if markdown:
|
||
parts.append(markdown)
|
||
if description:
|
||
parts.append(f"表格说明：{description}")
|
||
if footnotes:
|
||
parts.append(f"脚注说明：{footnotes}")
|
||
if ocr_text:
|
||
parts.append(f"表格 OCR：{ocr_text}")
|
||
if keywords:
|
||
parts.append(f"关键词：{'，'.join(keywords)}")
|
||
return "\n".join(parts)
|
||
|
||
|
||
def table_meta_summary(meta: dict[str, Any]) -> dict[str, Any]:
|
||
"""Pick table-specific fields for chunk metadata."""
|
||
keys = (
|
||
"table_title",
|
||
"table_description",
|
||
"header_rows",
|
||
"header_signature",
|
||
"footnotes",
|
||
"keywords",
|
||
"chapter",
|
||
"pages",
|
||
"page",
|
||
"bbox",
|
||
"bboxes",
|
||
"crop_path",
|
||
"image_path",
|
||
"row_count",
|
||
"col_count",
|
||
"cross_page",
|
||
"table_source",
|
||
"preceding_text",
|
||
"following_text",
|
||
"nearest_heading",
|
||
"parent_heading",
|
||
)
|
||
return {k: meta[k] for k in keys if k in meta and meta[k] not in (None, "", [], {})}
|