#!/usr/bin/env python3 """ Rebuild RAG index with improved chunking: - Remove file headers and page markers - Strict chunk size enforcement - Better paragraph-aware splitting - Filter out empty/noise chunks """ import json import re import math from pathlib import Path from collections import Counter from config import cfg PROCESSED_DIR = cfg.processed_dir INDEX_FILE = cfg.index_file KEYWORD_FILE = cfg.keyword_file CHUNK_SIZE = 512 CHUNK_OVERLAP = 64 MIN_CHUNK_LEN = 30 # Discard chunks shorter than this def clean_text(text, filename): """Remove noise from extracted PDF text and merge broken lines""" # Remove file header line text = re.sub(r'^# 源文件:.*\n*', '', text) # Remove page markers text = re.sub(r'--- 第\d+页 ---\n?', '', text) # Remove leading/trailing whitespace per line lines = [line.strip() for line in text.split('\n')] # Merge broken lines from PDF extraction. # PDF text has line breaks from page layout, not real paragraph breaks. # Strategy: only treat as paragraph break if previous text ends with # sentence-ending punctuation AND there's a blank line. merged = [] buf = '' blank_count = 0 for line in lines: if not line: blank_count += 1 continue if buf: last_char = buf[-1] if buf else '' ends_sentence = last_char in '。！？；：）》」』"' # Real paragraph break = ended with punctuation + had blank line(s) if ends_sentence and blank_count > 0: merged.append(buf) merged.append('') # preserve paragraph separator buf = line else: # Merge: either no blank line, or previous didn't end sentence buf += line else: buf = line blank_count = 0 if buf: merged.append(buf) text = '\n'.join(merged) # Collapse multiple blank lines text = re.sub(r'\n{3,}', '\n\n', text) return text.strip() def _smart_split(text, max_len, overlap): """Split long text at natural break points (punctuation), not mid-word""" # Priority: sentence-end > clause-end > comma > any punct > force break_chars = '。！？；\n' soft_break_chars = '，、：）》」』"' pieces = [] start = 0 while start < len(text): if start + max_len >= len(text): pieces.append(text[start:]) break end = start + max_len # Look backwards for best break point best = -1 # First try sentence-ending punctuation for i in range(end, max(start + max_len // 2, start), -1): if i < len(text) and text[i] in break_chars: best = i + 1 break # If not found, try soft breaks (comma etc) if best == -1: for i in range(end, max(start + max_len // 2, start), -1): if i < len(text) and text[i] in soft_break_chars: best = i + 1 break # Last resort: hard cut if best == -1: best = end pieces.append(text[start:best]) # Overlap: go back a bit start = max(best - overlap, start + 1) if start >= len(text): break return pieces def split_chunks(text, source_file): """Split text into chunks with strict size enforcement. Tables marked with [TABLE_START]...[TABLE_END] are kept as whole chunks.""" chunks = [] chunk_id = 0 # First: extract tables as separate blocks # Split text into table blocks and text blocks parts = re.split(r'(\[TABLE_START\].*?\[TABLE_END\])', text, flags=re.DOTALL) all_paragraphs = [] for part in parts: part = part.strip() if not part: continue if part.startswith('[TABLE_START]'): # Keep table as a single "paragraph" (won't be split further) table_content = part.replace('[TABLE_START]', '').replace('[TABLE_END]', '').strip() if table_content: all_paragraphs.append(('table', table_content)) else: # Regular text: split into paragraphs for para in re.split(r'\n\s*\n', part): para = para.strip() if para and len(para) >= 5: all_paragraphs.append(('text', para)) # Now process paragraphs paragraphs_iter = all_paragraphs current_chunk = '' for ptype, para in paragraphs_iter: # Tables: always emit as their own chunk(s) if ptype == 'table': # Flush current text chunk first if current_chunk and len(current_chunk) >= MIN_CHUNK_LEN: chunks.append({ 'id': f'{source_file}#chunk{chunk_id}', 'source': source_file, 'chunk_id': chunk_id, 'text': current_chunk, 'length': len(current_chunk), }) chunk_id += 1 current_chunk = '' # If table fits in one chunk, keep whole; otherwise split by rows if len(para) <= CHUNK_SIZE * 2: # Allow tables up to 2x chunk size chunks.append({ 'id': f'{source_file}#chunk{chunk_id}', 'source': source_file, 'chunk_id': chunk_id, 'text': para, 'length': len(para), }) chunk_id += 1 else: # Large table: split keeping header + groups of rows lines = para.split('\n') header = '\n'.join(lines[:2]) if len(lines) > 2 else '' # header + separator table_chunk = header for line in lines[2:]: if len(table_chunk) + len(line) + 1 > CHUNK_SIZE * 2: if table_chunk: chunks.append({ 'id': f'{source_file}#chunk{chunk_id}', 'source': source_file, 'chunk_id': chunk_id, 'text': table_chunk, 'length': len(table_chunk), }) chunk_id += 1 table_chunk = header + '\n' + line # restart with header else: table_chunk += '\n' + line if table_chunk and len(table_chunk) >= MIN_CHUNK_LEN: chunks.append({ 'id': f'{source_file}#chunk{chunk_id}', 'source': source_file, 'chunk_id': chunk_id, 'text': table_chunk, 'length': len(table_chunk), }) chunk_id += 1 continue # Regular text paragraph # If adding this paragraph fits in chunk_size if len(current_chunk) + len(para) + 1 <= CHUNK_SIZE: current_chunk = (current_chunk + '\n' + para).strip() if current_chunk else para else: # Save current chunk if it has content if current_chunk and len(current_chunk) >= MIN_CHUNK_LEN: chunks.append({ 'id': f'{source_file}#chunk{chunk_id}', 'source': source_file, 'chunk_id': chunk_id, 'text': current_chunk, 'length': len(current_chunk), }) chunk_id += 1 # If paragraph itself is too long, split by sentences/characters if len(para) > CHUNK_SIZE: # Try sentence splitting first (at major punctuation) sentences = re.split(r'(?<=[。！？；\n])', para) # If no split happened (single long sentence), try comma-level if len(sentences) <= 1 and len(para) > CHUNK_SIZE: sentences = _smart_split(para, CHUNK_SIZE, CHUNK_OVERLAP) # Already split into right-sized pieces, emit directly for piece in sentences: if len(piece.strip()) >= MIN_CHUNK_LEN: chunks.append({ 'id': f'{source_file}#chunk{chunk_id}', 'source': source_file, 'chunk_id': chunk_id, 'text': piece.strip(), 'length': len(piece.strip()), }) chunk_id += 1 current_chunk = '' continue sub_chunk = '' for sent in sentences: sent = sent.strip() if not sent: continue if len(sub_chunk) + len(sent) + 1 <= CHUNK_SIZE: sub_chunk = (sub_chunk + sent).strip() if sub_chunk else sent else: if sub_chunk and len(sub_chunk) >= MIN_CHUNK_LEN: chunks.append({ 'id': f'{source_file}#chunk{chunk_id}', 'source': source_file, 'chunk_id': chunk_id, 'text': sub_chunk, 'length': len(sub_chunk), }) chunk_id += 1 # If single sentence > CHUNK_SIZE, smart split at punctuation if len(sent) > CHUNK_SIZE: pieces = _smart_split(sent, CHUNK_SIZE, CHUNK_OVERLAP) for piece in pieces: if len(piece) >= MIN_CHUNK_LEN: chunks.append({ 'id': f'{source_file}#chunk{chunk_id}', 'source': source_file, 'chunk_id': chunk_id, 'text': piece, 'length': len(piece), }) chunk_id += 1 sub_chunk = '' else: sub_chunk = sent # Don't forget remaining sub_chunk if sub_chunk and len(sub_chunk) >= MIN_CHUNK_LEN: current_chunk = sub_chunk else: current_chunk = '' else: # Start new chunk with overlap from previous overlap = current_chunk[-CHUNK_OVERLAP:] if len(current_chunk) > CHUNK_OVERLAP else '' current_chunk = (overlap + '\n' + para).strip() if overlap else para # Last chunk if current_chunk and len(current_chunk) >= MIN_CHUNK_LEN: chunks.append({ 'id': f'{source_file}#chunk{chunk_id}', 'source': source_file, 'chunk_id': chunk_id, 'text': current_chunk, 'length': len(current_chunk), }) return chunks # Build index all_chunks = [] txt_files = sorted(PROCESSED_DIR.glob('*.txt')) print(f'Processing {len(txt_files)} text files...') for i, txt_path in enumerate(txt_files): text = txt_path.read_text(encoding='utf-8') text = clean_text(text, txt_path.stem) if len(text) < 30: continue chunks = split_chunks(text, txt_path.stem) all_chunks.extend(chunks) if (i + 1) % 200 == 0: print(f' Progress: {i+1}/{len(txt_files)}, chunks: {len(all_chunks)}') # Build keyword inverted index print(f'Building keyword index... ({len(all_chunks)} chunks)') keyword_index = {} for chunk in all_chunks: words = set() text = chunk['text'].lower() # Chinese bigram + trigram chinese_chars = re.findall(r'[\u4e00-\u9fff]+', text) for cc in chinese_chars: for j in range(len(cc) - 1): words.add(cc[j:j+2]) if len(cc) >= 3: for j in range(len(cc) - 2): words.add(cc[j:j+3]) # English words eng_words = re.findall(r'[a-z]+', text) words.update(eng_words) for w in words: if w not in keyword_index: keyword_index[w] = [] keyword_index[w].append(chunk['id']) # Save index_data = { 'total_chunks': len(all_chunks), 'total_files': len(txt_files), 'chunk_size': CHUNK_SIZE, 'chunks': all_chunks, 'keyword_index_size': len(keyword_index), } with open(INDEX_FILE, 'w', encoding='utf-8') as f: json.dump(index_data, f, ensure_ascii=False) with open(KEYWORD_FILE, 'w', encoding='utf-8') as f: json.dump(keyword_index, f, ensure_ascii=False) # Stats lengths = [c['length'] for c in all_chunks] print(f'\nDone!') print(f' Total chunks: {len(all_chunks)}') print(f' Keywords: {len(keyword_index)}') print(f' Avg length: {sum(lengths)/len(lengths):.0f}') print(f' Min/Max: {min(lengths)}/{max(lengths)}') print(f' <100: {sum(1 for l in lengths if l < 100)} ({sum(1 for l in lengths if l < 100)*100/len(lengths):.1f}%)') print(f' >1000: {sum(1 for l in lengths if l > 1000)} ({sum(1 for l in lengths if l > 1000)*100/len(lengths):.1f}%)') print(f' Index size: {INDEX_FILE.stat().st_size/1024/1024:.1f} MB')