Spaces:
Sleeping
Sleeping
| # core/table_extractor.py | |
| """ | |
| Извлечение таблиц из текстов документов. | |
| Единый источник для всех операций с таблицами. | |
| """ | |
| from __future__ import annotations | |
| import logging | |
| import re | |
| from dataclasses import dataclass, field | |
| from typing import Any | |
| logger = logging.getLogger(__name__) | |
| class ExtractedTable: | |
| """Структура извлечённой таблицы.""" | |
| title: str = "" | |
| headers: list[str] = field(default_factory=list) | |
| rows: list[list[str]] = field(default_factory=list) | |
| raw_text: str = "" | |
| source: str = "" | |
| confidence: float = 0.0 | |
| metadata: dict[str, Any] = field(default_factory=dict) | |
| def to_dict(self) -> dict[str, Any]: | |
| """Преобразует в словарь для JSON-сериализации.""" | |
| return { | |
| "title": self.title, | |
| "headers": self.headers, | |
| "rows": self.rows, | |
| "content": self.raw_text, | |
| "raw_text": self.raw_text, | |
| "doc_name": self.source, | |
| "source": self.source, | |
| "confidence": self.confidence, | |
| "metadata": self.metadata, | |
| "row_count": len(self.rows), | |
| } | |
| def is_valid(self, min_rows: int = 2) -> bool: | |
| """Проверяет, что таблица содержит достаточно данных.""" | |
| return len(self.rows) >= min_rows and bool(self.raw_text.strip()) | |
| class TableExtractor: | |
| """ | |
| Извлекает таблицы из текста различными методами. | |
| """ | |
| def __init__(self): | |
| self._seen: set[tuple[str, str]] = set() | |
| def extract( | |
| self, | |
| text: str, | |
| source: str = "", | |
| min_rows: int = 2, | |
| max_tables: int = 10, | |
| ) -> list[ExtractedTable]: | |
| """Извлекает таблицы из текста всеми доступными методами.""" | |
| if not text or not text.strip(): | |
| return [] | |
| all_tables: list[ExtractedTable] = [] | |
| # 1. Таблицы с маркером [ТАБЛИЦА ...] | |
| marker_tables = self._extract_marker_tables(text, source) | |
| all_tables.extend(marker_tables) | |
| # 2. Таблицы в Markdown-формате (с |) | |
| pipe_tables = self._extract_pipe_tables(text, source) | |
| all_tables.extend(pipe_tables) | |
| # 3. Таблицы по выравниванию (пробелы) | |
| aligned_tables = self._extract_aligned_tables(text, source) | |
| all_tables.extend(aligned_tables) | |
| # 4. Таблицы по ключевым словам | |
| keyword_tables = self._extract_keyword_tables(text, source) | |
| all_tables.extend(keyword_tables) | |
| # Дедупликация и фильтрация | |
| unique_tables = self._deduplicate_tables(all_tables) | |
| valid_tables = [t for t in unique_tables if t.is_valid(min_rows)] | |
| valid_tables.sort(key=lambda t: t.confidence, reverse=True) | |
| return valid_tables[:max_tables] | |
| def _extract_marker_tables(self, text: str, source: str) -> list[ExtractedTable]: | |
| """Извлекает таблицы по маркеру [ТАБЛИЦА ...].""" | |
| tables = [] | |
| lines = text.split("\n") | |
| in_table = False | |
| table_lines = [] | |
| table_title = "" | |
| table_metadata: dict[str, Any] = {} | |
| def flush() -> None: | |
| nonlocal table_lines, table_title, in_table, table_metadata | |
| if table_lines: | |
| if all(table_line.count("|") >= 2 for table_line in table_lines): | |
| table = self._build_table_from_pipe_lines( | |
| table_lines, source, confidence=0.95 | |
| ) | |
| table.title = table_title or table.title | |
| else: | |
| table = self._build_table_from_lines( | |
| table_lines, table_title or "Таблица", source, confidence=0.95 | |
| ) | |
| if table.is_valid(min_rows=1): | |
| table.metadata.update(table_metadata) | |
| tables.append(table) | |
| table_lines = [] | |
| table_title = "" | |
| table_metadata = {} | |
| in_table = False | |
| for i, line in enumerate(lines): | |
| stripped = line.strip() | |
| if stripped.upper().startswith("[ТАБЛИЦА"): | |
| if in_table: | |
| flush() | |
| in_table = True | |
| table_lines = [] | |
| page_match = re.search(r"\bстраница\s+(\d+)", stripped, flags=re.I) | |
| if page_match: | |
| table_metadata["page"] = int(page_match.group(1)) | |
| table_metadata["page_numbers"] = [int(page_match.group(1))] | |
| marker_title = re.sub(r"^\[ТАБЛИЦА\s*:?\s*", "", stripped, flags=re.I) | |
| marker_title = marker_title.rstrip("]").split("|", 1)[0].strip() | |
| if marker_title: | |
| table_title = marker_title | |
| elif i > 0 and len(lines[i - 1].strip()) < 100: | |
| table_title = lines[i - 1].strip() | |
| else: | |
| table_title = f"Таблица {len(tables) + 1}" | |
| continue | |
| if in_table: | |
| # Служебный маркер следующей страницы/источника не является | |
| # строкой таблицы, даже если между блоками нет пустой строки. | |
| if not stripped or stripped.startswith("[СТРАНИЦА") or stripped.startswith("[ИСТОЧНИК"): | |
| flush() | |
| else: | |
| table_lines.append(stripped) | |
| if in_table and table_lines: | |
| flush() | |
| return tables | |
| def _extract_pipe_tables(self, text: str, source: str) -> list[ExtractedTable]: | |
| """Извлекает Markdown-таблицы с разделителями |.""" | |
| tables = [] | |
| lines = text.split("\n") | |
| table_lines = [] | |
| in_table = False | |
| for line in lines: | |
| if "|" in line: | |
| stripped = line.strip() | |
| if stripped.count("|") >= 2: | |
| if not in_table: | |
| in_table = True | |
| table_lines = [] | |
| table_lines.append(stripped) | |
| else: | |
| if in_table and table_lines: | |
| table = self._build_table_from_pipe_lines( | |
| table_lines, source, confidence=0.85 | |
| ) | |
| if table.is_valid(min_rows=1): | |
| tables.append(table) | |
| table_lines = [] | |
| in_table = False | |
| else: | |
| if in_table and table_lines: | |
| table = self._build_table_from_pipe_lines( | |
| table_lines, source, confidence=0.85 | |
| ) | |
| if table.is_valid(min_rows=1): | |
| tables.append(table) | |
| table_lines = [] | |
| in_table = False | |
| if in_table and table_lines: | |
| table = self._build_table_from_pipe_lines( | |
| table_lines, source, confidence=0.85 | |
| ) | |
| if table.is_valid(min_rows=1): | |
| tables.append(table) | |
| return tables | |
| def _extract_aligned_tables(self, text: str, source: str) -> list[ExtractedTable]: | |
| """Извлекает таблицы по выравниванию текста.""" | |
| tables = [] | |
| lines = text.split("\n") | |
| i = 0 | |
| while i < len(lines): | |
| line = lines[i].strip() | |
| if not line: | |
| i += 1 | |
| continue | |
| if re.search(r"\d+\s*[мМ]?[м²³]?[°С]?", line): | |
| parts = re.split(r"\s{2,}", line) | |
| if len(parts) >= 3: | |
| table_lines = [line] | |
| j = i + 1 | |
| while j < len(lines): | |
| next_line = lines[j].strip() | |
| if not next_line: | |
| break | |
| next_parts = re.split(r"\s{2,}", next_line) | |
| if len(next_parts) >= 3 and any( | |
| re.search(r"\d", p) for p in next_parts | |
| ): | |
| table_lines.append(next_line) | |
| j += 1 | |
| else: | |
| break | |
| if len(table_lines) >= 3: | |
| table = self._build_table_from_aligned_lines( | |
| table_lines, source, confidence=0.7 | |
| ) | |
| if table.is_valid(min_rows=1): | |
| tables.append(table) | |
| i = j | |
| continue | |
| i += 1 | |
| return tables | |
| def _extract_keyword_tables(self, text: str, source: str) -> list[ExtractedTable]: | |
| """Извлекает таблицы по ключевым словам.""" | |
| tables = [] | |
| keywords = [ | |
| "таблица", "табл", "таблиц", "показатели", "значения", | |
| "параметры", "характеристики", "нормативы", "требования", | |
| ] | |
| lines = text.split("\n") | |
| i = 0 | |
| while i < len(lines): | |
| line_lower = lines[i].strip().lower() | |
| if any(kw in line_lower for kw in keywords) and ":" in lines[i]: | |
| title = lines[i].strip() | |
| table_lines = [] | |
| j = i + 1 | |
| while j < len(lines): | |
| next_line = lines[j].strip() | |
| if not next_line: | |
| break | |
| if any(kw in next_line.lower() for kw in keywords): | |
| break | |
| if re.search(r"\d+\s*[мМ]?[м²³]?[°С]?", next_line): | |
| table_lines.append(next_line) | |
| elif len(next_line) > 50 and j < i + 5: | |
| table_lines.append(next_line) | |
| else: | |
| break | |
| j += 1 | |
| if len(table_lines) >= 2: | |
| table = self._build_table_from_keyword_lines( | |
| table_lines, title, source, confidence=0.6 | |
| ) | |
| if table.is_valid(min_rows=1): | |
| tables.append(table) | |
| i = j | |
| continue | |
| i += 1 | |
| return tables | |
| def _build_table_from_lines( | |
| lines: list[str], | |
| title: str, | |
| source: str, | |
| confidence: float = 0.8, | |
| ) -> ExtractedTable: | |
| """Строит таблицу из строк.""" | |
| rows = [] | |
| for line in lines: | |
| line = line.strip() | |
| if line: | |
| if "\t" in line: | |
| parts = line.split("\t") | |
| else: | |
| parts = re.split(r"\s{2,}", line) | |
| if len(parts) > 1: | |
| rows.append([p.strip() for p in parts]) | |
| else: | |
| rows.append([line]) | |
| headers = rows[0] if rows and len(rows) > 0 else [] | |
| data_rows = rows[1:] if len(rows) > 1 else [] | |
| return ExtractedTable( | |
| title=title, | |
| headers=headers, | |
| rows=data_rows, | |
| raw_text="\n".join(lines), | |
| source=source, | |
| confidence=confidence, | |
| ) | |
| def _build_table_from_pipe_lines( | |
| lines: list[str], | |
| source: str, | |
| confidence: float = 0.85, | |
| ) -> ExtractedTable: | |
| """Строит таблицу из строк с разделителями |.""" | |
| rows = [] | |
| title = f"Таблица из {source}" if source else "Таблица" | |
| for line in lines: | |
| line = line.strip() | |
| if line: | |
| if line.startswith("|"): | |
| line = line[1:] | |
| if line.endswith("|"): | |
| line = line[:-1] | |
| parts = [p.strip().replace("\\|", "|") for p in re.split(r"(?<!\\)\|", line)] | |
| if any(parts): | |
| rows.append(parts) | |
| filtered_rows = [] | |
| for row in rows: | |
| if not all(re.fullmatch(r":?-+:?", p.strip()) for p in row): | |
| filtered_rows.append(row) | |
| headers = filtered_rows[0] if filtered_rows else [] | |
| data_rows = filtered_rows[1:] if len(filtered_rows) > 1 else [] | |
| return ExtractedTable( | |
| title=title, | |
| headers=headers, | |
| rows=data_rows, | |
| raw_text="\n".join(lines), | |
| source=source, | |
| confidence=confidence, | |
| ) | |
| def _build_table_from_aligned_lines( | |
| lines: list[str], | |
| source: str, | |
| confidence: float = 0.7, | |
| ) -> ExtractedTable: | |
| """Строит таблицу из выровненных строк.""" | |
| rows = [] | |
| title = f"Таблица из {source}" if source else "Таблица" | |
| for line in lines: | |
| line = line.strip() | |
| if line: | |
| parts = re.split(r"\s{2,}", line) | |
| if len(parts) > 1: | |
| rows.append(parts) | |
| else: | |
| rows.append([line]) | |
| headers = rows[0] if rows and len(rows) > 0 else [] | |
| data_rows = rows[1:] if len(rows) > 1 else [] | |
| return ExtractedTable( | |
| title=title, | |
| headers=headers, | |
| rows=data_rows, | |
| raw_text="\n".join(lines), | |
| source=source, | |
| confidence=confidence, | |
| ) | |
| def _build_table_from_keyword_lines( | |
| lines: list[str], | |
| title: str, | |
| source: str, | |
| confidence: float = 0.6, | |
| ) -> ExtractedTable: | |
| """Строит таблицу из строк по ключевым словам.""" | |
| rows = [] | |
| for line in lines: | |
| line = line.strip() | |
| if line: | |
| if "\t" in line: | |
| parts = line.split("\t") | |
| else: | |
| parts = re.split(r"\s{2,}", line) | |
| if len(parts) > 1: | |
| rows.append([p.strip() for p in parts]) | |
| else: | |
| rows.append([line]) | |
| headers = rows[0] if rows and len(rows) > 0 else [] | |
| data_rows = rows[1:] if len(rows) > 1 else [] | |
| return ExtractedTable( | |
| title=title, | |
| headers=headers, | |
| rows=data_rows, | |
| raw_text="\n".join(lines), | |
| source=source, | |
| confidence=confidence, | |
| ) | |
| def _table_fingerprint(table: ExtractedTable) -> str: | |
| """Создаёт уникальный отпечаток таблицы для дедупликации.""" | |
| return TableExtractor.table_fingerprint(table) | |
| def table_fingerprint(table: ExtractedTable) -> str: | |
| """Публичный ключ содержимого таблицы без обрезания последних строк.""" | |
| def clean(value: Any) -> str: | |
| return re.sub(r"\s+", " ", str(value).strip().lower()) | |
| headers = "|".join(clean(cell) for cell in table.headers) | |
| rows = "\n".join( | |
| "|".join(clean(cell) for cell in row) | |
| for row in table.rows | |
| ) | |
| return f"{clean(table.source)}::{headers}::{rows}" | |
| def _deduplicate_tables( | |
| cls, | |
| tables: list[ExtractedTable], | |
| ) -> list[ExtractedTable]: | |
| """Удаляет дубликаты таблиц по fingerprint.""" | |
| unique: list[ExtractedTable] = [] | |
| seen: set[str] = set() | |
| for table in tables: | |
| fingerprint = cls._table_fingerprint(table) | |
| if fingerprint in seen: | |
| continue | |
| seen.add(fingerprint) | |
| unique.append(table) | |
| return unique | |
| # ========= УДОБНЫЕ ФУНКЦИИ ДЛЯ ИМПОРТА ========= | |
| def extract_tables(text: str, doc_name: str = "", min_rows: int = 2) -> list[ExtractedTable]: | |
| """Быстрый импорт: извлекает таблицы из текста.""" | |
| extractor = TableExtractor() | |
| return extractor.extract(text, source=doc_name, min_rows=min_rows) | |
| def extract_tables_from_results(results: list[Any], min_rows: int = 2, limit: int = 3) -> list[ExtractedTable]: | |
| """Извлекает таблицы из результатов поиска.""" | |
| extractor = TableExtractor() | |
| all_tables: list[ExtractedTable] = [] | |
| seen: set[str] = set() | |
| for result in results[:limit]: | |
| if hasattr(result, "text"): | |
| text = getattr(result, "text", "") | |
| doc_name = getattr(result, "doc_name", "Документ") | |
| result_metadata = getattr(result, "metadata", {}) or {} | |
| elif isinstance(result, dict): | |
| text = result.get("text", "") | |
| doc_name = result.get("doc_name", "Документ") | |
| result_metadata = result.get("metadata", {}) or {} | |
| else: | |
| continue | |
| if not isinstance(text, str) or not text.strip(): | |
| continue | |
| tables = extractor.extract(text, source=doc_name, min_rows=min_rows) | |
| for table in tables: | |
| if isinstance(result_metadata, dict): | |
| table.metadata = {**result_metadata, **table.metadata} | |
| fingerprint = TableExtractor.table_fingerprint(table) | |
| if fingerprint not in seen: | |
| seen.add(fingerprint) | |
| all_tables.append(table) | |
| return all_tables | |
| def find_city_in_tables(tables: list[ExtractedTable], city: str, city_aliases: dict[str, set[str]] | None = None) -> dict[str, Any] | None: | |
| """ | |
| Ищет город в таблицах с учётом алиасов и словоформ. | |
| Args: | |
| tables: Список извлечённых таблиц | |
| city: Название города (каноническое) | |
| city_aliases: Словарь алиасов {каноническое_имя: {словоформы}} | |
| Returns: | |
| dict с информацией о найденной строке или None | |
| """ | |
| if not tables or not city: | |
| return None | |
| city_lower = city.lower().replace("ё", "е").strip() | |
| # Собираем все формы для поиска | |
| city_forms = {city_lower} | |
| if city_aliases and city_lower in city_aliases: | |
| city_forms.update(city_aliases[city_lower]) | |
| # Дополнительные варианты для составных названий | |
| if "-" in city_lower: | |
| city_forms.add(city_lower.replace("-", " ")) | |
| city_forms.add(city_lower.replace("-", "")) | |
| # Поиск по таблицам | |
| for table in tables: | |
| for idx, row in enumerate(table.rows): | |
| if isinstance(row, list): | |
| row_text = " ".join(str(c) for c in row) | |
| row_lower = row_text.lower().replace("ё", "е") | |
| # Проверяем каждую форму | |
| for form in city_forms: | |
| if re.search(rf"(?<![а-яa-z]){re.escape(form)}(?![а-яa-z])", row_lower): | |
| return { | |
| "table": table, | |
| "row": row, | |
| "index": idx, | |
| "row_text": row_text, | |
| "matched_form": form, | |
| "confidence": 0.9, | |
| } | |
| # Проверка по отдельным ячейкам (выше точность) | |
| for cell in row: | |
| cell_lower = str(cell).lower().replace("ё", "е") | |
| for form in city_forms: | |
| if re.search(rf"(?<![а-яa-z]){re.escape(form)}(?![а-яa-z])", cell_lower): | |
| return { | |
| "table": table, | |
| "row": row, | |
| "index": idx, | |
| "row_text": row_text, | |
| "matched_form": form, | |
| "confidence": 1.0, | |
| } | |
| elif isinstance(row, str): | |
| row_lower = row.lower() | |
| for form in city_forms: | |
| if re.search(rf"(?<![а-яa-z]){re.escape(form)}(?![а-яa-z])", row_lower): | |
| return { | |
| "table": table, | |
| "row": [row], | |
| "index": idx, | |
| "row_text": row, | |
| "matched_form": form, | |
| "confidence": 0.8, | |
| } | |
| # Fallback: поиск по токенам (для "ростов-на-дону" и т.п.) | |
| city_tokens = city_lower.replace("-", " ").split() | |
| if len(city_tokens) > 1: | |
| for table in tables: | |
| for idx, row in enumerate(table.rows): | |
| if isinstance(row, list): | |
| row_text = " ".join(str(c) for c in row).lower() | |
| if all(re.search(rf"(?<![а-яa-z]){re.escape(token)}(?![а-яa-z])", row_text) for token in city_tokens): | |
| return { | |
| "table": table, | |
| "row": row, | |
| "index": idx, | |
| "row_text": " ".join(str(c) for c in row), | |
| "matched_form": " ".join(city_tokens), | |
| "confidence": 0.7, | |
| } | |
| elif isinstance(row, str): | |
| row_lower = row.lower() | |
| if all(re.search(rf"(?<![а-яa-z]){re.escape(token)}(?![а-яa-z])", row_lower) for token in city_tokens): | |
| return { | |
| "table": table, | |
| "row": [row], | |
| "index": idx, | |
| "row_text": row, | |
| "matched_form": " ".join(city_tokens), | |
| "confidence": 0.7, | |
| } | |
| return None | |
| def tables_to_dicts(tables: list[ExtractedTable]) -> list[dict[str, Any]]: | |
| """Преобразует список ExtractedTable в список словарей.""" | |
| return [t.to_dict() for t in tables] | |
| def _run_self_tests() -> None: | |
| """Проверяет границы таблиц, одну строку, метаданные и города.""" | |
| tables = extract_tables( | |
| "[ТАБЛИЦА: Климат | страница 4]\n| Город | t_от | z_от |\n| --- | --- | --- |\n| Томск | -8,4 | 225 |", | |
| "СП 131.13330", | |
| min_rows=1, | |
| ) | |
| assert tables and tables[0].title == "Климат" | |
| assert find_city_in_tables(tables, "омск") is None | |
| assert find_city_in_tables(tables, "томск") is not None | |
| assert tables[0].metadata["page"] == 4 | |
| assert extract_tables("Обычный абзац без таблицы", min_rows=1) == [] | |
| aligned = extract_tables("| Название | Значение |\n| :--- | ---: |\n| A\\|B | 12 |", min_rows=1) | |
| assert aligned[0].rows == [["A|B", "12"]] | |
| if __name__ == "__main__": | |
| _run_self_tests() | |
| # ИСПРАВЛЕНО: инициализация stripped; названия/страницы таблиц; однострочные | |
| # таблицы; Markdown-разделители и экранированные ячейки; точное сопоставление | |
| # города; метаданные результатов поиска; дедупликация всех строк; самотесты. | |