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 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 from llama_index.core.text_splitter import SentenceSplitter from config import CHUNK_SIZE, CHUNK_OVERLAP def chunk_table_document(doc, max_rows_per_chunk=5, max_chunk_size=2000): """Simple table chunking: max 5 rows or 2000 chars per chunk""" table_num = doc.metadata.get('table_number', 'unknown') # Parse table lines = doc.text.strip().split('\n') table_header_lines = [] data_rows = [] in_data = False for line in lines: if line.startswith('Данные таблицы:'): in_data = True table_header_lines.append(line) elif in_data and line.startswith('Строка'): data_rows.append(line) elif not in_data: table_header_lines.append(line) table_header = '\n'.join(table_header_lines) + '\n' if not data_rows: # No rows, return as is return [doc] log_message(f"Таблица {table_num}: {len(data_rows)} строк") # Simple chunking chunks = [] current_chunk_rows = [] current_size = len(table_header) for row in data_rows: row_size = len(row) + 1 # Check if adding this row exceeds limits if (len(current_chunk_rows) >= max_rows_per_chunk or current_size + row_size > max_chunk_size) and current_chunk_rows: # Save current chunk chunk_text = table_header + '\n'.join(current_chunk_rows) chunks.append(chunk_text) log_message(f" Чанк: {len(current_chunk_rows)} строк, {len(chunk_text)} символов") # Start new chunk with overlap of 1 row if len(current_chunk_rows) > 0: current_chunk_rows = [current_chunk_rows[-1]] current_size = len(table_header) + len(current_chunk_rows[0]) + 1 else: current_chunk_rows = [] current_size = len(table_header) current_chunk_rows.append(row) current_size += row_size # Final chunk if current_chunk_rows: chunk_text = table_header + '\n'.join(current_chunk_rows) chunks.append(chunk_text) log_message(f" Последний чанк: {len(current_chunk_rows)} строк") log_message(f"Таблица {table_num} разделена на {len(chunks)} чанков") # Create documents chunked_docs = [] for i, chunk_text in enumerate(chunks): chunk_metadata = doc.metadata.copy() chunk_metadata.update({ "chunk_id": i, "total_chunks": len(chunks), "chunk_size": len(chunk_text), "is_chunked": True }) chunked_doc = Document(text=chunk_text, metadata=chunk_metadata) chunked_docs.append(chunked_doc) return chunked_docs def table_to_document(table_data, document_id=None): if not isinstance(table_data, dict): log_message(f"⚠️ ПРОПУЩЕНА: table_data не является словарем") 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 or len(table_rows) == 0: log_message(f"⚠️ ПРОПУЩЕНА: Таблица {table_num} из '{doc_id}' - нет данных в 'data'") return [] content = create_table_content(table_data) content_size = len(content) row_count = len(table_rows) base_doc = Document( text=content, metadata={ "type": "table", "table_number": table_num, "table_title": table_title, "document_id": doc_id, "section": section, "section_id": section, "total_rows": row_count, "content_size": content_size } ) if content_size > CHUNK_SIZE: chunked_docs = chunk_table_document(base_doc) log_message(f" ✂️ Разделена на {len(chunked_docs)} чанков") for i, chunk_doc in enumerate(chunked_docs): log_message(f" Чанк {i+1}: {chunk_doc.metadata['chunk_size']} символов") return chunked_docs else: log_message(f"✓ ДОБАВЛЕНА: Таблица {table_num} из документа '{doc_id}' | " f"Размер: {content_size} символов | Строк: {row_count}") 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 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 []