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

MAX_ROWS_PER_CHUNK = 10
MAX_CHUNK_SIZE = 4000

def create_table_content(table_data):
    """Create formatted content from table data"""
    doc_id = 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', 'Неизвестно')
    
    content = f"Таблица: {table_num}\n"
    content += f"Название: {table_title}\n"
    content += f"Документ: {doc_id}\n"
    content += f"Раздел: {section}\n"
    
    headers = table_data.get('headers', [])
    if headers:
        content += f"\nЗаголовки: {' | '.join(headers)}\n"
    
    if 'data' in table_data and isinstance(table_data['data'], list):
        content += "\nДанные таблицы:\n"
        for row_idx, row in enumerate(table_data['data'], start=1):
            if isinstance(row, dict):
                row_text = " | ".join([f"{k}: {v}" for k, v in row.items() if v])
                content += f"Строка {row_idx}: {row_text}\n"
    
    return content

def chunk_table_document(doc, max_chunk_size=MAX_CHUNK_SIZE, max_rows_per_chunk=MAX_ROWS_PER_CHUNK):
    lines = doc.text.strip().split('\n')

    # Separate header and data rows
    header_lines = []
    data_rows = []
    in_data = False

    for line in lines:
        if line.startswith('Данные таблицы:'):
            in_data = True
            header_lines.append(line)
        elif in_data and line.startswith('Строка'):
            data_rows.append(line)
        elif not in_data:
            header_lines.append(line)

    header = '\n'.join(header_lines) + '\n'

    if not data_rows:
        return [doc]

    chunks = []
    current_rows = []
    current_size = len(header)

    for row in data_rows:
        row_size = len(row) + 1
        # Check both limits: chunk size and row count
        if ((current_size + row_size > max_chunk_size or len(current_rows) >= max_rows_per_chunk) and current_rows):
            chunk_text = header + '\n'.join(current_rows)
            chunks.append(chunk_text)
            log_message(f"Создана часть таблицы размером {len(chunk_text)} символов с {len(current_rows)} строками")
            current_rows = []
            current_size = len(header)

        current_rows.append(row)
        current_size += row_size
        log_message(f"Добавлена строка к текущему чанку, текущий размер {current_size} символов")

    # Add final chunk
    if current_rows:
        chunk_text = header + '\n'.join(current_rows)
        chunks.append(chunk_text)
        log_message(f"Создана финальная часть таблицы размером {len(chunk_text)} символов с {len(current_rows)} строками")

    # Create Document objects
    chunked_docs = []
    for i, chunk_text in enumerate(chunks):
        chunk_doc = Document(
            text=chunk_text,
            metadata={
                "type": "table",
                "table_number": doc.metadata.get('table_number'),
                "document_id": doc.metadata.get('document_id'),
                "section": doc.metadata.get('section'),
                "chunk_id": i,
                "total_chunks": len(chunks),
                "is_chunked": True
            }
        )
        chunked_docs.append(chunk_doc)

    return chunked_docs


def table_to_document(table_data, document_id=None):
    """Convert table data to Document, chunk if needed"""
    
    if not isinstance(table_data, dict):
        return []
    
    doc_id = document_id or table_data.get('document_id') or table_data.get('document', 'Неизвестно')
    table_num = table_data.get('table_number', 'Неизвестно')
    table_title = table_data.get('table_title', 'Неизвестно')
    section = table_data.get('section', 'Неизвестно')
    table_rows = table_data.get('data', [])
    
    if not table_rows:
        return []
    
    # Build table content
    content = f"Таблица: {table_num}\n"
    content += f"Название: {table_title}\n"
    content += f"Документ: {doc_id}\n"
    content += f"Раздел: {section}\n"
    
    headers = table_data.get('headers', [])
    if headers:
        content += f"\nЗаголовки: {' | '.join(headers)}\n"
    
    content += "\nДанные таблицы:\n"
    for row_idx, row in enumerate(table_rows, start=1):
        if isinstance(row, dict):
            row_text = " | ".join([f"{k}: {v}" for k, v in row.items() if v])
            content += f"Строка {row_idx}: {row_text}\n"
    
    # Create base document
    base_doc = Document(
        text=content,
        metadata={
            "type": "table",
            "table_number": table_num,
            "document_id": doc_id,
            "section": section
        }
    )
    if len(content) > 4000:
        chunks = chunk_table_document(base_doc)
        log_message(f"Таблица {table_num} разбита на {len(chunks)} частей")
        return chunk_table_document(base_doc)
    return [base_doc]
    
    
def load_table_data(repo_id, hf_token, table_data_dir):
    log_message("=" * 60)
    log_message("НАЧАЛО ЗАГРУЗКИ ТАБЛИЧНЫХ ДАННЫХ")
    log_message("=" * 60)
    
    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 = []
        stats = {
            'total_tables': 0,
            'total_size': 0,
            'by_document': defaultdict(lambda: {'count': 0, 'size': 0})
        }
        
        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
                )
                
                log_message(f"\nОбработка файла: {file_path}")
                
                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:
                            sorted_sheets = sorted(
                                table_data['sheets'],
                                key=lambda sheet: sheet.get('table_number', '')  # or use 'table_number'
                            )
                            
                            for sheet in sorted_sheets:
                                sheet['document'] = document_id
                                docs_list = table_to_document(sheet, document_id)
                                table_documents.extend(docs_list)
                                
                                for doc in docs_list:
                                    stats['total_tables'] += 1
                                    size = doc.metadata.get('content_size', 0)
                                    stats['total_size'] += size
                                    stats['by_document'][document_id]['count'] += 1
                                    stats['by_document'][document_id]['size'] += size
                                    log_message(f"Добавлена таблица {sheet.get('table_number', 'Неизвестно')} из документа {document_id}, размер {size} символов")
                        else:
                            docs_list = table_to_document(table_data, document_id)
                            table_documents.extend(docs_list)
                            
                            for doc in docs_list:
                                stats['total_tables'] += 1
                                size = doc.metadata.get('content_size', 0)
                                stats['total_size'] += size
                                stats['by_document'][document_id]['count'] += 1
                                stats['by_document'][document_id]['size'] += size
                                
                        
            except Exception as e:
                log_message(f"❌ ОШИБКА файла {file_path}: {str(e)}")
                continue
        
        # Log summary statistics
        log_message("\n" + "=" * 60)
        log_message("СТАТИСТИКА ПО ТАБЛИЦАМ")
        log_message("=" * 60)
        log_message(f"Всего таблиц добавлено: {stats['total_tables']}")
        log_message(f"Общий размер: {stats['total_size']:,} символов")
        log_message(f"Средний размер таблицы: {stats['total_size'] // stats['total_tables'] if stats['total_tables'] > 0 else 0:,} символов")
        
        log_message("\nПо документам:")
        for doc_id, doc_stats in sorted(stats['by_document'].items()):
            log_message(f"  • {doc_id}: {doc_stats['count']} таблиц, "
                       f"{doc_stats['size']:,} символов")
        
        log_message("=" * 60)
        
        return table_documents
        
    except Exception as e:
        log_message(f"❌ КРИТИЧЕСКАЯ ОШИБКА загрузки табличных данных: {str(e)}")
        return []