table_prep.py

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 []