import json import zipfile import pandas as pd from huggingface_hub import hf_hub_download, list_repo_files from llama_index.core import Document from llama_index.core.text_splitter import SentenceSplitter from my_logging import log_message # Configuration CHUNK_SIZE = 1500 CHUNK_OVERLAP = 256 def chunk_text_documents(documents): """Chunk with deduplication""" text_splitter = SentenceSplitter( chunk_size=CHUNK_SIZE, chunk_overlap=300 # Increased overlap ) seen_texts = set() chunked = [] for doc in documents: # Skip duplicates or too-short content text_normalized = doc.text.strip() if len(text_normalized) < 50 or text_normalized in seen_texts: continue seen_texts.add(text_normalized) chunks = text_splitter.get_nodes_from_documents([doc]) for i, chunk in enumerate(chunks): chunk.metadata.update({ 'chunk_id': i, 'total_chunks': len(chunks), 'chunk_size': len(chunk.text), 'document_group': normalize_doc_id(doc.metadata.get('document_id', 'unknown')) }) chunked.append(chunk) if chunked: avg_size = sum(len(c.text) for c in chunked) / len(chunked) log_message(f"✓ Text: {len(documents)} docs → {len(chunked)} chunks (avg: {avg_size:.0f} chars)") return chunked def chunk_table_by_rows(table_data, doc_id, max_chars=2000): """Chunk tables by content size, not fixed rows""" headers = table_data.get('headers', []) rows = table_data.get('data', []) table_num = table_data.get('table_number', 'unknown') table_title = table_data.get('table_title', '') section = table_data.get('section', '') table_num_clean = str(table_num).strip() # Create unique identifier import re if 'приложени' in section.lower(): appendix_match = re.search(r'приложени[еия]\s*(\d+|[а-яА-Я])', section.lower()) if appendix_match: table_identifier = f"{table_num_clean} (Приложение {appendix_match.group(1).upper()})" else: table_identifier = table_num_clean else: table_identifier = table_num_clean if not rows: return [] # Estimate base metadata size base_content = f"Документ: {doc_id}\nТаблица: {table_identifier}\n" if table_title: base_content += f"Название: {table_title}\n" if section: base_content += f"Раздел: {section}\n" header_content = "" if headers: header_content = "Столбцы: " + " | ".join(str(h) for h in headers) + "\n\n" base_size = len(base_content) + len(header_content) # Group rows by size chunks = [] current_rows = [] current_size = base_size for row in rows: # Estimate row size if isinstance(row, dict): row_str = " | ".join(f"{k}: {v}" for k, v in row.items() if v and str(v).strip() and str(v).lower() not in ['nan', 'none', '']) elif isinstance(row, list): row_str = " | ".join(str(v) for v in row if v and str(v).strip() and str(v).lower() not in ['nan', 'none', '']) else: row_str = str(row) row_size = len(row_str) + 2 # +2 for newline # If adding this row exceeds limit and we have rows, create chunk if current_size + row_size > max_chars and current_rows: chunks.append(current_rows[:]) current_rows = [] current_size = base_size current_rows.append(row) current_size += row_size # Add remaining rows if current_rows: chunks.append(current_rows) # Create documents documents = [] for chunk_idx, chunk_rows in enumerate(chunks): content = base_content content += f"Таблица {table_identifier} документа {doc_id}\n" if len(chunks) > 1: content += f"Часть {chunk_idx+1} из {len(chunks)}\n" content += "\n" + header_content for idx, row in enumerate(chunk_rows, 1): if isinstance(row, dict): parts = [f"{k}: {v}" for k, v in row.items() if v and str(v).strip() and str(v).lower() not in ['nan', 'none', '']] if parts: content += f"{idx}. {' | '.join(parts)}\n" elif isinstance(row, list): parts = [str(v) for v in row if v and str(v).strip() and str(v).lower() not in ['nan', 'none', '']] if parts: content += f"{idx}. {' | '.join(parts)}\n" metadata = { 'type': 'table', 'document_id': doc_id, 'document_group': normalize_doc_id(doc_id), 'table_number': table_num_clean, 'table_identifier': table_identifier, 'table_title': table_title, 'section': section, 'chunk_id': chunk_idx, 'total_chunks': len(chunks), 'chunk_size': len(content), 'is_complete_table': len(chunks) == 1 } documents.append(Document(text=content, metadata=metadata)) log_message(f" Chunk {chunk_idx+1}: {len(chunk_rows)} rows, {len(content)} chars") log_message(f" Meta: doc={doc_id}, table={table_identifier}, group={metadata['document_group']}") log_message(f" Table {table_identifier} ({doc_id}): {len(rows)} rows → {len(chunks)} chunks") return documents def normalize_doc_id(doc_id): import re normalized = re.sub(r'\s+', ' ', str(doc_id).strip().upper()) normalized = normalized.replace('ГОСТ Р', 'ГОСТР').replace('ГОСТР', 'ГОСТ Р') return normalized def format_table_content(table_data, headers, rows, doc_id, table_identifier, chunk_info=""): table_num = table_data.get('table_number', 'unknown') table_title = table_data.get('table_title', '') section = table_data.get('section', '') # Build content with multiple search variations content = f"ДОКУМЕНТ: {doc_id}\n" content += f"ТАБЛИЦА: {table_identifier}\n" # Add search variations for document ID doc_variations = [doc_id] if 'Р' in doc_id: doc_variations.append(doc_id.replace(' Р ', ' Р')) doc_variations.append(doc_id.replace(' Р ', 'Р')) for var in set(doc_variations): content += f"ДОКУМЕНТ_ВАРИАНТ: {var}\n" if table_title: content += f"НАЗВАНИЕ: {table_title}\n" if section: content += f"РАЗДЕЛ: {section}\n" content += f"{'='*70}\n\n" # Enhanced search text content += f"Документ {doc_id}. " content += f"Таблица {table_identifier}. " content += f"Номер таблицы {table_num}. " if table_title: content += f"Название: {table_title}. " if section: content += f"Раздел: {section}. " # Add more search patterns content += f"Таблицы документа {doc_id}. " content += f"Содержание {doc_id}. " if chunk_info: content += f"{chunk_info}. " content += f"\n\nДАННЫЕ ТАБЛИЦЫ {table_identifier}:\n{'='*70}\n\n" if headers: content += f"СТОЛБЦЫ: {' | '.join(str(h) for h in headers)}\n\n" for idx, row in enumerate(rows, 1): if isinstance(row, dict): parts = [f"{k}: {v}" for k, v in row.items() if v and str(v).strip().lower() not in ['nan', 'none', '', 'null']] if parts: content += f"{idx}. {' | '.join(parts)}\n" elif isinstance(row, list): parts = [str(v) for v in row if v and str(v).strip().lower() not in ['nan', 'none', '', 'null']] if parts: content += f"{idx}. {' | '.join(parts)}\n" return content def load_json_documents(repo_id, hf_token, json_dir): import zipfile import tempfile import os log_message("Loading JSON documents...") files = list_repo_files(repo_id=repo_id, repo_type="dataset", token=hf_token) json_files = [f for f in files if f.startswith(json_dir) and f.endswith('.json')] zip_files = [f for f in files if f.startswith(json_dir) and f.endswith('.zip')] log_message(f"Found {len(json_files)} JSON files and {len(zip_files)} ZIP files") documents = [] stats = {'success': 0, 'failed': 0, 'empty': 0} for file_path in json_files: try: log_message(f" Loading: {file_path}") local_path = hf_hub_download( repo_id=repo_id, filename=file_path, repo_type="dataset", token=hf_token ) docs = extract_sections_from_json(local_path) if docs: documents.extend(docs) stats['success'] += 1 log_message(f" ✓ Extracted {len(docs)} sections") else: stats['empty'] += 1 log_message(f" ⚠ No sections found") except Exception as e: stats['failed'] += 1 log_message(f" ✗ Error: {e}") for zip_path in zip_files: try: log_message(f" Processing ZIP: {zip_path}") local_zip = hf_hub_download( repo_id=repo_id, filename=zip_path, repo_type="dataset", token=hf_token ) with zipfile.ZipFile(local_zip, 'r') as zf: json_files_in_zip = [f for f in zf.namelist() if f.endswith('.json') and not f.startswith('__MACOSX') and not f.startswith('.') and not '._' in f] log_message(f" Found {len(json_files_in_zip)} JSON files in ZIP") for json_file in json_files_in_zip: try: file_content = zf.read(json_file) # Skip if file is too small if len(file_content) < 10: log_message(f" ✗ Skipping: {json_file} (file too small)") stats['failed'] += 1 continue # Try UTF-8 first (most common) try: text_content = file_content.decode('utf-8') except UnicodeDecodeError: try: text_content = file_content.decode('utf-8-sig') except UnicodeDecodeError: try: # Try UTF-16 (the issue you're seeing) text_content = file_content.decode('utf-16') except UnicodeDecodeError: try: text_content = file_content.decode('windows-1251') except UnicodeDecodeError: log_message(f" ✗ Skipping: {json_file} (encoding failed)") stats['failed'] += 1 continue # Validate JSON structure if not text_content.strip().startswith('{') and not text_content.strip().startswith('['): log_message(f" ✗ Skipping: {json_file} (not valid JSON)") stats['failed'] += 1 continue with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json', encoding='utf-8') as tmp: tmp.write(text_content) tmp_path = tmp.name docs = extract_sections_from_json(tmp_path) if docs: documents.extend(docs) stats['success'] += 1 log_message(f" ✓ {json_file}: {len(docs)} sections") else: stats['empty'] += 1 log_message(f" ⚠ {json_file}: No sections") os.unlink(tmp_path) except json.JSONDecodeError as e: stats['failed'] += 1 log_message(f" ✗ {json_file}: Invalid JSON") except Exception as e: stats['failed'] += 1 log_message(f" ✗ {json_file}: {str(e)[:100]}") except Exception as e: log_message(f" ✗ Error with ZIP: {e}") log_message(f"="*60) log_message(f"JSON Loading Stats:") log_message(f" Success: {stats['success']}") log_message(f" Empty: {stats['empty']}") log_message(f" Failed: {stats['failed']}") log_message(f" Total sections: {len(documents)}") log_message(f"="*60) return documents def extract_sections_from_json(json_path): documents = [] try: with open(json_path, 'r', encoding='utf-8') as f: data = json.load(f) doc_id = data.get('document_metadata', {}).get('document_id', 'unknown') doc_id = normalize_doc_id(doc_id) # NORMALIZE for section in data.get('sections', []): if section.get('section_text', '').strip(): documents.append(Document( text=section['section_text'], metadata={ 'type': 'text', 'document_id': doc_id, 'section_id': section.get('section_id', ''), 'chunk_size': len(section['section_text']) } )) for subsection in section.get('subsections', []): if subsection.get('subsection_text', '').strip(): documents.append(Document( text=subsection['subsection_text'], metadata={ 'type': 'text', 'document_id': doc_id, 'section_id': subsection.get('subsection_id', ''), 'chunk_size': len(subsection['subsection_text']) } )) for sub_sub in subsection.get('sub_subsections', []): if sub_sub.get('sub_subsection_text', '').strip(): documents.append(Document( text=sub_sub['sub_subsection_text'], metadata={ 'type': 'text', 'document_id': doc_id, 'section_id': sub_sub.get('sub_subsection_id', ''), 'chunk_size': len(sub_sub['sub_subsection_text']) } )) except Exception as e: log_message(f"Error extracting from {json_path}: {e}") return documents def load_table_documents(repo_id, hf_token, table_dir): """Load ALL tables including from multi-document files""" log_message("Loading tables...") 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_dir) and f.endswith('.json')] log_message(f"Found {len(table_files)} table files") all_chunks = [] doc_id_stats = {} for file_path in table_files: try: local_path = hf_hub_download( repo_id=repo_id, filename=file_path, repo_type="dataset", token=hf_token ) with open(local_path, 'r', encoding='utf-8') as f: data = json.load(f) file_doc_id = data.get('document_id', data.get('document', 'unknown')) for sheet in data.get('sheets', []): sheet_doc_id = sheet.get('document_id', sheet.get('document', file_doc_id)) # Track which documents we're loading if sheet_doc_id not in doc_id_stats: doc_id_stats[sheet_doc_id] = 0 chunks = chunk_table_by_rows(sheet, sheet_doc_id) all_chunks.extend(chunks) doc_id_stats[sheet_doc_id] += len(chunks) except Exception as e: log_message(f"Error loading {file_path}: {e}") # Log what we loaded log_message(f"\nTable loading summary:") for doc_id, count in sorted(doc_id_stats.items()): log_message(f" {doc_id}: {count} chunks") log_message(f"\n✓ Total table chunks: {len(all_chunks)}") return all_chunks def load_image_documents(repo_id, hf_token, image_dir): """Load with proper linking""" log_message("Loading images...") files = list_repo_files(repo_id=repo_id, repo_type="dataset", token=hf_token) csv_files = [f for f in files if f.startswith(image_dir) and f.endswith('.csv')] documents = [] seen = set() for file_path in csv_files: try: local_path = hf_hub_download( repo_id=repo_id, filename=file_path, repo_type="dataset", token=hf_token ) df = pd.read_csv(local_path) for _, row in df.iterrows(): doc_id = str(row.get('Обозначение документа', 'unknown')) img_num = str(row.get('№ Изображения', 'unknown')) key = f"{doc_id}_{img_num}" if key in seen: continue seen.add(key) content = f"Документ: {doc_id}\n" content += f"Рисунок: {img_num}\n" content += f"Название: {row.get('Название изображения', '')}\n" content += f"Описание: {row.get('Описание изображение', '')}\n" documents.append(Document( text=content, metadata={ 'type': 'image', 'document_id': doc_id, 'document_group': normalize_doc_id(doc_id), 'image_number': img_num, 'section': str(row.get('Раздел документа', '')), 'chunk_size': len(content) } )) except Exception as e: log_message(f"Error loading {file_path}: {e}") if documents: avg_size = sum(d.metadata['chunk_size'] for d in documents) / len(documents) log_message(f"✓ Images: {len(documents)} loaded (avg: {avg_size:.0f} chars)") return documents def load_all_documents(repo_id, hf_token, json_dir, table_dir, image_dir): """Main loader - combines all document types""" log_message("="*60) log_message("STARTING DOCUMENT LOADING") log_message("="*60) # Load text sections text_docs = load_json_documents(repo_id, hf_token, json_dir) text_chunks = chunk_text_documents(text_docs) # Load tables (already chunked) table_chunks = load_table_documents(repo_id, hf_token, table_dir) # Load images (no chunking needed) image_docs = load_image_documents(repo_id, hf_token, image_dir) all_docs = text_chunks + table_chunks + image_docs log_message("="*60) log_message(f"TOTAL DOCUMENTS: {len(all_docs)}") log_message(f" Text chunks: {len(text_chunks)}") log_message(f" Table chunks: {len(table_chunks)}") log_message(f" Images: {len(image_docs)}") log_message("="*60) return all_docs