file_chunk.py

"""
Improved Document Chunking for LeanRAG

This module provides advanced document chunking capabilities with multiple strategies:

1. SEMANTIC: Respects sentence boundaries for better context preservation
2. HYBRID: Combines semantic and token-based approaches
3. FIXED_TOKEN: Original token-based chunking with sliding window

Key improvements:
- Semantic awareness (doesn't break sentences mid-way)
- Quality scoring for chunks
- Configurable chunk sizes and overlap
- Better handling of document structure
- Support for multiple documents
- Support for PDF and JSONL input formats

Usage:
    from file_chunk import chunk_documents, ChunkingStrategy, load_documents

    # Load documents from file (PDF or JSONL)
    documents = load_documents("document.pdf")
    # or
    documents = load_documents("data.jsonl")

    # Semantic chunking (recommended)
    results = chunk_documents(
        documents,
        strategy=ChunkingStrategy.SEMANTIC,
        max_token_size=1024,
        overlap_token_size=128
    )

    # Access chunks
    for doc_result in results:
        for chunk in doc_result['chunks']:
            print(f"Text: {chunk['text']}")
            print(f"Quality: {chunk['quality_score']}")
"""

import json
import tiktoken
import re
from hashlib import md5
from typing import List, Dict, Any, Optional
from enum import Enum
from pathlib import Path
import PyPDF2

class ChunkingStrategy(Enum):
    FIXED_TOKEN = "fixed_token"
    SEMANTIC = "semantic"
    HYBRID = "hybrid"

def read_pdf_file(file_path: str) -> List[str]:
    """
    Read and parse a PDF file, returning a list of page texts.

    Args:
        file_path: Path to the PDF file

    Returns:
        List of strings, one per page
    """
    documents = []

    try:
        with open(file_path, 'rb') as file:
            pdf_reader = PyPDF2.PdfReader(file)

            for page_num in range(len(pdf_reader.pages)):
                page = pdf_reader.pages[page_num]
                text = page.extract_text()

                # Clean up the text
                text = re.sub(r'\n+', ' ', text)  # Replace multiple newlines with spaces
                text = re.sub(r'\s+', ' ', text)  # Replace multiple spaces with single space
                text = text.strip()

                if text:  # Only add non-empty pages
                    documents.append(text)

    except Exception as e:
        raise ValueError(f"Error reading PDF file {file_path}: {str(e)}")

    return documents

def read_jsonl_file(file_path: str) -> List[str]:
    """
    Read and parse a JSONL file, returning a list of document texts.

    Args:
        file_path: Path to the JSONL file

    Returns:
        List of document texts
    """
    documents = []

    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            for line_num, line in enumerate(f, 1):
                if line.strip():
                    try:
                        data = json.loads(line)
                        # Extract text content (try different field names)
                        text = (data.get('context') or
                               data.get('text') or
                               data.get('content') or
                               data.get('input', ''))
                        if text:
                            documents.append(text)
                    except json.JSONDecodeError as e:
                        raise ValueError(f"Error parsing JSON at line {line_num} in {file_path}: {str(e)}")

    except FileNotFoundError:
        raise ValueError(f"File not found: {file_path}")
    except Exception as e:
        raise ValueError(f"Error reading file {file_path}: {str(e)}")

    return documents

def load_documents(file_path: str) -> List[str]:
    """
    Load documents from a file, automatically detecting the format (PDF or JSONL).

    Args:
        file_path: Path to the input file

    Returns:
        List of document texts
    """
    path = Path(file_path)

    if not path.exists():
        raise ValueError(f"File not found: {file_path}")

    if path.suffix.lower() == '.pdf':
        return read_pdf_file(file_path)
    elif path.suffix.lower() == '.jsonl':
        return read_jsonl_file(file_path)
    else:
        raise ValueError(f"Unsupported file format: {path.suffix}. Supported formats: .pdf, .jsonl")

def find_supported_files(directory_path: str) -> List[Path]:
    """
    Recursively find all supported files (.pdf, .jsonl) in a directory.

    Args:
        directory_path: Path to the directory to search

    Returns:
        List of Path objects for supported files
    """
    path = Path(directory_path)

    if not path.exists():
        raise ValueError(f"Directory not found: {directory_path}")

    if not path.is_dir():
        raise ValueError(f"Path is not a directory: {directory_path}")

    supported_extensions = {'.pdf', '.jsonl'}
    supported_files = []

    # Recursively find all files with supported extensions
    for file_path in path.rglob('*'):
        if file_path.is_file() and file_path.suffix.lower() in supported_extensions:
            supported_files.append(file_path)

    return supported_files

def load_documents_from_directory(directory_path: str) -> List[str]:
    """
    Load documents from all supported files in a directory recursively.

    Args:
        directory_path: Path to the directory containing documents

    Returns:
        List of document texts from all files
    """
    supported_files = find_supported_files(directory_path)
    all_documents = []

    for file_path in supported_files:
        try:
            documents = load_documents(str(file_path))
            all_documents.extend(documents)
            print(f"Loaded {len(documents)} documents from {file_path}")
        except Exception as e:
            print(f"Warning: Failed to load {file_path}: {e}")
            continue

    return all_documents

def compute_mdhash_id(content, prefix: str = ""):
    return prefix + md5(content.encode()).hexdigest()

def split_into_sentences(text: str) -> List[str]:
    """Split text into sentences using regex patterns."""
    # Handle common sentence endings with lookbehind to avoid splitting on abbreviations
    sentence_pattern = r'(?<=[.!?])\s+(?=[A-Z])'
    sentences = re.split(sentence_pattern, text.strip())

    # Filter out empty sentences and very short fragments
    sentences = [s.strip() for s in sentences if s.strip() and len(s.strip()) > 10]

    return sentences

def estimate_tokens(text: str, encoder) -> int:
    """Estimate token count for text."""
    return len(encoder.encode(text))

def semantic_chunking(
    text: str,
    encoder,
    max_tokens: int = 512,
    min_tokens: int = 100,
    overlap_tokens: int = 64
) -> List[str]:
    """
    Semantic chunking that respects sentence boundaries.
    """
    sentences = split_into_sentences(text)
    if not sentences:
        return [text]

    chunks = []
    current_chunk = ""
    current_tokens = 0

    for sentence in sentences:
        sentence_tokens = estimate_tokens(sentence, encoder)

        # If adding this sentence would exceed max_tokens and we have content
        if current_tokens + sentence_tokens > max_tokens and current_chunk:
            chunks.append(current_chunk.strip())

            # Start new chunk with overlap from previous chunk
            if overlap_tokens > 0 and len(chunks) > 0:
                overlap_text = create_overlap_text(chunks[-1], overlap_tokens, encoder)
                current_chunk = overlap_text + " " + sentence
                current_tokens = estimate_tokens(current_chunk, encoder)
            else:
                current_chunk = sentence
                current_tokens = sentence_tokens
        else:
            if current_chunk:
                current_chunk += " " + sentence
            else:
                current_chunk = sentence
            current_tokens += sentence_tokens

    # Add the last chunk if it has content
    if current_chunk.strip():
        chunks.append(current_chunk.strip())

    # Filter out chunks that are too small
    filtered_chunks = [chunk for chunk in chunks if estimate_tokens(chunk, encoder) >= min_tokens]

    return filtered_chunks if filtered_chunks else [text]

def create_overlap_text(text: str, overlap_tokens: int, encoder) -> str:
    """Create overlap text from the end of a chunk."""
    tokens = encoder.encode(text)
    if len(tokens) <= overlap_tokens:
        return text

    overlap_tokens_list = tokens[-overlap_tokens:]
    return encoder.decode(overlap_tokens_list)

def hybrid_chunking(
    text: str,
    encoder,
    max_tokens: int = 512,
    overlap_tokens: int = 64,
    semantic_weight: float = 0.7
) -> List[str]:
    """
    Hybrid approach combining semantic and token-based chunking.
    """
    # First try semantic chunking
    semantic_chunks = semantic_chunking(text, encoder, max_tokens, overlap_tokens=overlap_tokens)

    # If semantic chunking produces too few or too many chunks, fall back to token-based
    if len(semantic_chunks) == 1 or len(semantic_chunks) > len(text) // 200:  # heuristic
        return chunk_documents_fixed([text], encoder, max_tokens, overlap_tokens)[0]['chunks']

    return semantic_chunks

def chunk_documents_fixed(
    docs: List[str],
    encoder,
    max_token_size: int = 512,
    overlap_token_size: int = 64
) -> List[Dict[str, Any]]:
    """Original fixed token chunking method."""
    tokens_list = encoder.encode_batch(docs, num_threads=16)

    results = []
    for index, tokens in enumerate(tokens_list):
        chunk_token_ids = []
        lengths = []

        for start in range(0, len(tokens), max_token_size - overlap_token_size):
            chunk = tokens[start : start + max_token_size]
            chunk_token_ids.append(chunk)
            lengths.append(len(chunk))

        chunk_texts = encoder.decode_batch(chunk_token_ids)

        chunks = []
        for i, text in enumerate(chunk_texts):
            chunks.append({
                "text": text.strip().replace("\n", " "),
                "tokens": lengths[i],
                "chunk_index": i,
                "start_token": i * (max_token_size - overlap_token_size),
                "end_token": min((i + 1) * max_token_size - i * overlap_token_size, len(tokens))
            })

        results.append({
            "doc_index": index,
            "total_tokens": len(tokens),
            "num_chunks": len(chunks),
            "chunks": chunks
        })

    return results

def chunk_documents(
    docs: List[str],
    model_name: str = "cl100k_base",
    max_token_size: int = 512,
    overlap_token_size: int = 64,
    strategy: ChunkingStrategy = ChunkingStrategy.SEMANTIC,
    min_chunk_tokens: int = 50
) -> List[Dict[str, Any]]:
    """
    Improved document chunking with multiple strategies.

    Args:
        docs: List of documents to chunk
        model_name: Tiktoken model name
        max_token_size: Maximum tokens per chunk
        overlap_token_size: Token overlap between chunks
        strategy: Chunking strategy to use
        min_chunk_tokens: Minimum tokens for a chunk to be valid

    Returns:
        List of chunk dictionaries with metadata
    """
    encoder = tiktoken.get_encoding(model_name)

    results = []
    for doc_index, doc in enumerate(docs):
        if strategy == ChunkingStrategy.FIXED_TOKEN:
            doc_results = chunk_documents_fixed([doc], encoder, max_token_size, overlap_token_size)
            results.extend(doc_results)
            continue

        elif strategy == ChunkingStrategy.SEMANTIC:
            chunks = semantic_chunking(doc, encoder, max_token_size, min_chunk_tokens, overlap_token_size)

        elif strategy == ChunkingStrategy.HYBRID:
            chunks = hybrid_chunking(doc, encoder, max_token_size, overlap_token_size)

        # Convert chunks to standardized format
        chunk_dicts = []
        for i, chunk_text in enumerate(chunks):
            tokens = estimate_tokens(chunk_text, encoder)
            chunk_dicts.append({
                "text": chunk_text.strip().replace("\n", " "),
                "tokens": tokens,
                "chunk_index": i,
                "hash_code": compute_mdhash_id(chunk_text),
                "strategy": strategy.value,
                "quality_score": calculate_chunk_quality(chunk_text, tokens, max_token_size)
            })

        results.append({
            "doc_index": doc_index,
            "total_tokens": estimate_tokens(doc, encoder),
            "num_chunks": len(chunk_dicts),
            "chunks": chunk_dicts,
            "strategy": strategy.value
        })

    return results

def calculate_chunk_quality(text: str, tokens: int, max_tokens: int) -> float:
    """
    Calculate a quality score for a chunk based on various metrics.
    """
    if tokens == 0:
        return 0.0

    # Length appropriateness (prefer chunks that use 70-90% of max capacity)
    length_score = 1.0 - abs(tokens / max_tokens - 0.8) / 0.8

    # Sentence completeness (prefer chunks ending with sentence terminators)
    ends_with_sentence = text.strip().endswith(('.', '!', '?', '"', "'"))
    completeness_score = 1.0 if ends_with_sentence else 0.5

    # Word diversity (avoid repetitive chunks)
    words = re.findall(r'\b\w+\b', text.lower())
    unique_words = len(set(words))
    diversity_score = min(unique_words / len(words) if words else 0, 1.0)

    # Combine scores (weighted average)
    quality_score = (0.4 * length_score + 0.3 * completeness_score + 0.3 * diversity_score)

    return round(quality_score, 3)
if __name__ == "__main__":
    # Configuration - can be modified for different use cases
    config = {
        "max_token_size": 1024,      # Maximum tokens per chunk
        "overlap_token_size": 128,   # Token overlap between chunks
        "strategy": ChunkingStrategy.SEMANTIC,  # SEMANTIC, HYBRID, or FIXED_TOKEN
        "min_chunk_tokens": 50,      # Minimum tokens for a chunk to be valid
    }

    # File paths
    dataset = 'mix'
    original_text_file = f"datasets/{dataset}/{dataset}.jsonl"
    output_file = f"datasets/{dataset}/{dataset}_chunk.json"

    # Load data
    data = []
    try:
        with open(original_text_file, 'r', encoding='utf-8') as f:
            for line in f:
                if line.strip():
                    data.append(json.loads(line))
    except FileNotFoundError:
        print(f"Warning: {original_text_file} not found. Using sample data for testing.")
        # Sample data for testing
        data = [
            {"context": "This is a sample document. It contains multiple sentences for testing the chunking functionality."},
            {"context": "Another document with different content. This helps verify that the chunking works across multiple documents."}
        ]

    # Extract contexts
    contexts = [item['context'] for item in data if 'context' in item]

    if not contexts:
        contexts = [item.get('input', '') for item in data if item.get('input')]

    print(f"Loaded {len(contexts)} documents for chunking")

    # Perform chunking
    results = chunk_documents(
        contexts,
        max_token_size=config["max_token_size"],
        overlap_token_size=config["overlap_token_size"],
        strategy=config["strategy"],
        min_chunk_tokens=config["min_chunk_tokens"]
    )

    # Flatten results for output compatibility
    flattened_results = []
    total_quality = 0
    total_tokens = 0

    for doc_result in results:
        for chunk in doc_result['chunks']:
            flattened_results.append({
                "hash_code": chunk["hash_code"],
                "text": chunk["text"],
                "tokens": chunk["tokens"],
                "quality_score": chunk["quality_score"],
                "strategy": chunk["strategy"]
            })
            total_quality += chunk["quality_score"]
            total_tokens += chunk["tokens"]

    # Save results
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(flattened_results, f, ensure_ascii=False, indent=2)

    # Print statistics
    print(f"\nChunking completed using {config['strategy'].value} strategy:")
    print(f"- Processed {len(contexts)} documents")
    print(f"- Generated {len(flattened_results)} chunks")
    if flattened_results:
        print(f"- Average chunk size: {total_tokens / len(flattened_results):.1f} tokens")
        print(f"- Average quality score: {total_quality / len(flattened_results):.3f}")
    print(f"- Results saved to: {output_file}")