from collections import defaultdict import json from huggingface_hub import hf_hub_download, list_repo_files from llama_index.core import Document from my_logging import log_message CUSTOM_TABLE_CONFIGS = { "ГОСТ Р 50.05.01-2018": { "№3": {"method": "group_by_column", "group_column": "Класс герметичности и чувствительности"}, "№Б.1": {"method": "group_by_column", "group_column": "Класс чувствительности системы контроля"} }, "ГОСТ Р 50.06.01-2017": {"№ Б.2": {"method": "split_by_rows"}}, "НП-104-18": {"*": {"method": "group_entire_table"}}, "НП-068-05": { "Таблица 1": {"method": "group_by_column", "group_column": "Рабочее давление среды, МПа"}, "Таблица 2": {"method": "group_by_column", "group_column": "Рабочее давление среды, МПа"}, "Таблица Приложения 1": {"method": "group_by_column", "group_column": "Тип"} }, "ГОСТ Р 59023.1-2020": { "№ 1": {"method": "split_by_rows"}, "№ 2": {"method": "split_by_rows"}, "№ 3": {"method": "split_by_rows"} }, "НП-089-15": {"-": {"method": "split_by_rows"}}, "НП-105-18": {"№ 4.8": {"method": "group_entire_table"}}, "ГОСТ Р 50.05.23-2020": {"№8": {"method": "group_entire_table"}}, "ГОСТ Р 50.03.01-2017": {"А.8": {"method": "group_entire_table"}} } def create_meta_info(document_name, section, table_number, table_title, extra_info=""): base_info = f'Документ "{document_name}", Раздел: {section}, Таблица: {table_number}' if table_title and table_title.strip(): base_info += f', Название: {table_title}' if extra_info: base_info += f', {extra_info}' return base_info def create_chunk_text(meta_info, headers, rows, add_row_numbers=False): chunk_lines = [meta_info.rstrip()] chunk_lines.append("Заголовки: " + " | ".join(headers)) for i, row in enumerate(rows, start=1): row_parts = [f"{h}: {row.get(h, '')}" for h in headers if row.get(h, '')] if add_row_numbers: chunk_lines.append(f"Строка {i}: {' | '.join(row_parts)}") else: chunk_lines.append(' | '.join(row_parts)) return "\n".join(chunk_lines) def get_custom_config(document_id, table_number): for doc_pattern, tables_config in CUSTOM_TABLE_CONFIGS.items(): if document_id.startswith(doc_pattern): return tables_config.get(table_number, tables_config.get("*")) return None def group_by_column_method(table_data, document_name, group_column): documents = [] headers = table_data.get("headers", []) rows = table_data.get("data", []) section = table_data.get("section", "") table_number = table_data.get("table_number", "") table_title = table_data.get("table_title", "") grouped = defaultdict(list) for row in rows: grouped[row.get(group_column, "UNKNOWN")].append(row) for group_value, group_rows in grouped.items(): meta_info = create_meta_info(document_name, section, table_number, table_title, f'Группа по "{group_column}": {group_value}') chunk_text = create_chunk_text(meta_info, headers, group_rows, add_row_numbers=True) documents.append(Document( text=chunk_text, metadata={ "type": "table", "table_number": table_number, "table_title": table_title, "document_id": document_name, "section": section, "section_id": section, "group_column": group_column, "group_value": group_value, "total_rows": len(group_rows), "processing_method": "group_by_column" } )) return documents def split_by_rows_method(table_data, document_name): documents = [] headers = table_data.get("headers", []) rows = table_data.get("data", []) section = table_data.get("section", "") table_number = table_data.get("table_number", "") table_title = table_data.get("table_title", "") for i, row in enumerate(rows, start=1): meta_info = create_meta_info(document_name, section, table_number, table_title, f'Строка: {i}') chunk_text = create_chunk_text(meta_info, headers, [row]) documents.append(Document( text=chunk_text, metadata={ "type": "table", "table_number": table_number, "table_title": table_title, "document_id": document_name, "section": section, "section_id": section, "row_number": i, "total_rows": len(rows), "processing_method": "split_by_rows" } )) return documents def group_entire_table_method(table_data, document_name): headers = table_data.get("headers", []) rows = table_data.get("data", []) section = table_data.get("section", "") table_number = table_data.get("table_number", "") table_title = table_data.get("table_title", "") meta_info = create_meta_info(document_name, section, table_number, table_title) chunk_text = create_chunk_text(meta_info, headers, rows) return [Document( text=chunk_text, metadata={ "type": "table", "table_number": table_number, "table_title": table_title, "document_id": document_name, "section": section, "section_id": section, "total_rows": len(rows), "processing_method": "group_entire_table" } )] def process_table(table_data, document_name, method_config): method = method_config.get("method") if method == "group_by_column": return group_by_column_method(table_data, document_name, method_config.get("group_column")) elif method == "split_by_rows": return split_by_rows_method(table_data, document_name) elif method == "group_entire_table": return group_entire_table_method(table_data, document_name) return None def table_to_document(table_data, document_id=None): if not isinstance(table_data, dict): return [] doc_id = document_id or table_data.get('document_id', table_data.get('document', 'Неизвестно')) table_num = table_data.get('table_number', 'Неизвестно') table_title = table_data.get('table_title', 'Неизвестно') section = table_data.get('section', 'Неизвестно') method_config = get_custom_config(doc_id, table_num) if method_config: log_message(f"✓ Таблица {table_num} '{table_title}' в документе {doc_id}: метод {method_config['method']}") custom_docs = process_table(table_data, doc_id, method_config) if custom_docs: return custom_docs header_content = f"Таблица: {table_num}\nНазвание: {table_title}\nДокумент: {doc_id}\nРаздел: {section}\n" if 'data' in table_data and isinstance(table_data['data'], list): table_content = header_content + "\nДанные таблицы:\n" for row_idx, row in enumerate(table_data['data']): if isinstance(row, dict): row_text = " | ".join([f"{k}: {v}" for k, v in row.items()]) table_content += f"Строка {row_idx + 1}: {row_text}\n" return [Document( text=table_content, metadata={ "type": "table", "table_number": table_num, "table_title": table_title, "document_id": doc_id, "section": section, "section_id": section, "total_rows": len(table_data['data']), "processing_method": "default" } )] return [Document( text=header_content, metadata={ "type": "table", "table_number": table_num, "table_title": table_title, "document_id": doc_id, "section": section, "section_id": section, "processing_method": "default" } )] def load_table_data(repo_id, hf_token, table_data_dir): log_message("Загрузка табличных данных") try: files = list_repo_files(repo_id=repo_id, repo_type="dataset", token=hf_token) table_files = [f for f in files if f.startswith(table_data_dir) and f.endswith('.json')] log_message(f"Найдено {len(table_files)} JSON файлов с таблицами") table_documents = [] for file_path in table_files: try: local_path = hf_hub_download( repo_id=repo_id, filename=file_path, local_dir='', repo_type="dataset", token=hf_token ) with open(local_path, 'r', encoding='utf-8') as f: table_data = json.load(f) if isinstance(table_data, dict): document_id = table_data.get('document', 'unknown') if 'sheets' in table_data: for sheet in table_data['sheets']: sheet['document'] = document_id docs_list = table_to_document(sheet, document_id) table_documents.extend(docs_list) else: docs_list = table_to_document(table_data, document_id) table_documents.extend(docs_list) elif isinstance(table_data, list): for table_json in table_data: docs_list = table_to_document(table_json) table_documents.extend(docs_list) except Exception as e: log_message(f"Ошибка файла {file_path}: {str(e)}") continue log_message(f"✓✓✓ ИТОГО создано {len(table_documents)} табличных документов (до chunking)") return table_documents except Exception as e: log_message(f"Ошибка загрузки табличных данных: {str(e)}") return []