src/data/pipeline.py

"""
OpenMind Data Pipeline - Preprocessing, Deduplication, and Tokenization.

Complete pipeline for:
1. Downloading datasets from Hugging Face
2. Deduplication using MinHash LSH
3. Language and quality filtering
4. Tokenization with our BPE tokenizer
5. Packing into memory-mapped binary files
"""

import os
import json
import hashlib
import struct
from pathlib import Path
from typing import Iterable, Iterator, Optional
from collections import defaultdict

import numpy as np
from tqdm import tqdm


class MinHashDeduplicator:
    """
    MinHash LSH-based approximate deduplication.

    Uses Jaccard similarity with configurable n-gram shingles and
    Locality-Sensitive Hashing for efficient near-duplicate detection.
    """

    def __init__(
        self,
        num_perm: int = 128,
        threshold: float = 0.85,
        ngram_size: int = 5,
        seed: int = 42,
    ):
        self.num_perm = num_perm
        self.threshold = threshold
        self.ngram_size = ngram_size
        self.seed = seed

        # Generate random hash parameters
        rng = np.random.RandomState(seed)
        self.a = rng.randint(1, 2**31 - 1, size=num_perm, dtype=np.int64)
        self.b = rng.randint(0, 2**31 - 1, size=num_perm, dtype=np.int64)
        self.prime = np.int64(2**31 - 1)  # Mersenne prime

        # LSH bands/rows for the given threshold
        self.bands, self.rows = self._optimal_bands_rows()
        self.buckets: dict[int, dict[str, list]] = defaultdict(lambda: defaultdict(list))

    def _optimal_bands_rows(self) -> tuple[int, int]:
        """Find optimal bands and rows for the given threshold."""
        best = (1, self.num_perm)
        for b in range(1, self.num_perm + 1):
            if self.num_perm % b == 0:
                r = self.num_perm // b
                t = (1.0 / b) ** (1.0 / r)
                if abs(t - self.threshold) < abs((1.0 / best[0]) ** (1.0 / best[1]) - self.threshold):
                    best = (b, r)
        return best

    def _shingle(self, text: str) -> set[int]:
        """Create character n-gram shingles from text."""
        text = text.lower().strip()
        shingles = set()
        for i in range(len(text) - self.ngram_size + 1):
            shingle = text[i: i + self.ngram_size]
            shingles.add(hash(shingle) & 0xFFFFFFFF)
        return shingles

    def _minhash(self, shingles: set[int]) -> np.ndarray:
        """Compute MinHash signature from shingles."""
        if not shingles:
            return np.full(self.num_perm, np.iinfo(np.int64).max)

        shingle_array = np.array(list(shingles), dtype=np.int64)
        # Hash each shingle with each hash function and take minimum
        hashvals = np.zeros(self.num_perm, dtype=np.int64)
        for i in range(self.num_perm):
            hashed = (self.a[i] * shingle_array + self.b[i]) % self.prime
            hashvals[i] = hashed.min()
        return hashvals

    def is_duplicate(self, text: str, doc_id: str = "") -> bool:
        """Check if text is a near-duplicate of any seen document."""
        shingles = self._shingle(text)
        if len(shingles) < 2:
            return False

        signature = self._minhash(shingles)

        # Check against existing buckets
        for band_idx in range(self.bands):
            start = band_idx * self.rows
            end = start + self.rows
            band_hash = hashlib.md5(signature[start:end].tobytes()).hexdigest()

            if band_hash in self.buckets[band_idx]:
                return True

        # Add to buckets
        for band_idx in range(self.bands):
            start = band_idx * self.rows
            end = start + self.rows
            band_hash = hashlib.md5(signature[start:end].tobytes()).hexdigest()
            self.buckets[band_idx][band_hash].append(doc_id)

        return False


class QualityFilter:
    """
    Document quality filter with configurable rules.

    Filters based on:
    - Document length (min/max characters)
    - Repeated line ratio
    - Special character ratio
    - Word count thresholds
    """

    def __init__(
        self,
        min_chars: int = 100,
        max_chars: int = 100_000,
        max_repeated_line_ratio: float = 0.3,
        max_special_char_ratio: float = 0.3,
        min_words: int = 20,
        min_avg_word_len: float = 3.0,
    ):
        self.min_chars = min_chars
        self.max_chars = max_chars
        self.max_repeated_line_ratio = max_repeated_line_ratio
        self.max_special_char_ratio = max_special_char_ratio
        self.min_words = min_words
        self.min_avg_word_len = min_avg_word_len

    def filter(self, text: str) -> bool:
        """Return True if document passes quality checks, False to reject."""
        # Length check
        if len(text) < self.min_chars or len(text) > self.max_chars:
            return False

        # Word count
        words = text.split()
        if len(words) < self.min_words:
            return False

        # Average word length (filters gibberish)
        avg_word_len = sum(len(w) for w in words) / max(len(words), 1)
        if avg_word_len < self.min_avg_word_len:
            return False

        # Repeated line ratio
        lines = text.strip().split("\n")
        if len(lines) > 1:
            unique_lines = set(line.strip() for line in lines if line.strip())
            if len(unique_lines) / max(len(lines), 1) < (1 - self.max_repeated_line_ratio):
                return False

        # Special character ratio
        special_chars = sum(1 for c in text if not c.isalnum() and not c.isspace())
        if special_chars / max(len(text), 1) > self.max_special_char_ratio:
            return False

        return True


class DataPipeline:
    """
    Complete data preprocessing pipeline.

    Handles downloading, filtering, deduplication, tokenization,
    and packing into memory-mapped binary files for training.
    """

    def __init__(
        self,
        tokenizer_path: Optional[str] = None,
        output_dir: str = "data",
        max_seq_len: int = 2048,
        seed: int = 42,
    ):
        self.output_dir = output_dir
        self.max_seq_len = max_seq_len
        self.seed = seed
        self.tokenizer = None

        if tokenizer_path and os.path.exists(tokenizer_path):
            try:
                from .tokenizer import BPETokenizer
            except ImportError:
                import sys as _sys
                _sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
                from tokenizer import BPETokenizer
            self.tokenizer = BPETokenizer.load(tokenizer_path)

        self.quality_filter = QualityFilter()
        self.deduplicator = MinHashDeduplicator(seed=seed)

        os.makedirs(output_dir, exist_ok=True)

    def load_tokenizer(self, tokenizer_path: str):
        """Load tokenizer from disk."""
        try:
            from .tokenizer import BPETokenizer
        except ImportError:
            import sys as _sys
            _sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
            from tokenizer import BPETokenizer
        self.tokenizer = BPETokenizer.load(tokenizer_path)

    def process_dataset(
        self,
        dataset_name: str,
        split: str = "train",
        text_column: str = "text",
        streaming: bool = True,
        max_documents: Optional[int] = None,
        output_name: Optional[str] = None,
    ) -> str:
        """
        Process a Hugging Face dataset through the full pipeline.

        Args:
            dataset_name: HF dataset identifier (e.g., "roneneldan/TinyStories")
            split: Dataset split to process
            text_column: Column containing text data
            streaming: Use streaming mode for large datasets
            max_documents: Optional limit on documents to process
            output_name: Output filename (without extension)

        Returns:
            Path to the output .bin file
        """
        from datasets import load_dataset

        print(f"Loading dataset: {dataset_name} (split={split}, streaming={streaming})")
        dataset = load_dataset(dataset_name, split=split, streaming=streaming)

        if output_name is None:
            output_name = dataset_name.replace("/", "_") + f"_{split}"

        # Process documents through pipeline
        processed_docs = self._pipeline_iter(dataset, text_column, max_documents)

        # Tokenize and pack
        output_path = self.tokenize_and_pack(processed_docs, output_name)

        return output_path

    def _pipeline_iter(
        self,
        dataset,
        text_column: str,
        max_documents: Optional[int] = None,
    ) -> Iterator[str]:
        """Run documents through filtering and deduplication."""
        stats = {"total": 0, "passed_quality": 0, "passed_dedup": 0}

        for i, example in enumerate(tqdm(dataset, desc="Processing documents")):
            if max_documents and i >= max_documents:
                break

            stats["total"] += 1
            text = example.get(text_column, "")

            if not text or not isinstance(text, str):
                continue

            # Quality filter
            if not self.quality_filter.filter(text):
                continue
            stats["passed_quality"] += 1

            # Deduplication
            if self.deduplicator.is_duplicate(text, doc_id=str(i)):
                continue
            stats["passed_dedup"] += 1

            yield text

        print(f"Pipeline stats: {json.dumps(stats, indent=2)}")

    def tokenize_and_pack(
        self,
        documents: Iterable[str],
        output_name: str,
    ) -> str:
        """
        Tokenize documents and pack into a memory-mapped binary file.

        Documents are concatenated with <|endoftext|> separators and
        split into fixed-length sequences of max_seq_len.

        Args:
            documents: Iterator of text documents
            output_name: Output filename (without extension)

        Returns:
            Path to the output .bin file
        """
        assert self.tokenizer is not None, "Tokenizer must be loaded first!"

        eos_id = self.tokenizer.eos_token_id
        all_tokens = []
        doc_count = 0

        for doc in tqdm(documents, desc="Tokenizing"):
            tokens = self.tokenizer.encode(doc)
            tokens.append(eos_id)  # Add document separator
            all_tokens.extend(tokens)
            doc_count += 1

        if not all_tokens:
            print("Warning: No tokens produced!")
            return ""

        print(f"Tokenized {doc_count} documents -> {len(all_tokens):,} tokens")

        # Pack into sequences of max_seq_len
        num_sequences = len(all_tokens) // self.max_seq_len
        trimmed = all_tokens[: num_sequences * self.max_seq_len]

        print(f"Packed into {num_sequences} sequences of length {self.max_seq_len}")

        # Save as memory-mapped binary
        output_path = self.save_mmap(np.array(trimmed, dtype=np.uint16), output_name)

        return output_path

    def save_mmap(self, tokens: np.ndarray, filename: str) -> str:
        """
        Save token array as a memory-mapped binary file.

        Format: uint16 (supports vocab up to 65k)

        Args:
            tokens: 1D numpy array of token IDs
            filename: Output filename (without extension)

        Returns:
            Path to the saved file
        """
        output_path = os.path.join(self.output_dir, f"{filename}.bin")

        # Write header with metadata
        header_path = os.path.join(self.output_dir, f"{filename}_meta.json")
        meta = {
            "num_tokens": int(len(tokens)),
            "dtype": str(tokens.dtype),
            "max_seq_len": self.max_seq_len,
            "num_sequences": int(len(tokens)) // self.max_seq_len,
        }
        with open(header_path, "w") as f:
            json.dump(meta, f, indent=2)

        # Write binary data
        tokens.tofile(output_path)

        file_size_mb = os.path.getsize(output_path) / (1024 * 1024)
        print(f"Saved {len(tokens):,} tokens to {output_path} ({file_size_mb:.1f} MB)")

        return output_path

    def process_jsonl(
        self,
        input_path: str,
        text_field: str = "text",
        output_name: Optional[str] = None,
        max_documents: Optional[int] = None,
    ) -> str:
        """
        Process a local JSONL file through the pipeline.

        Args:
            input_path: Path to .jsonl file
            text_field: JSON field containing text
            output_name: Output filename
            max_documents: Optional document limit

        Returns:
            Path to output .bin file
        """
        if output_name is None:
            output_name = Path(input_path).stem

        def doc_iter():
            count = 0
            with open(input_path, "r", encoding="utf-8") as f:
                for line in f:
                    if max_documents and count >= max_documents:
                        break
                    try:
                        data = json.loads(line.strip())
                        text = data.get(text_field, "")
                        if text and self.quality_filter.filter(text):
                            if not self.deduplicator.is_duplicate(text, str(count)):
                                yield text
                                count += 1
                    except json.JSONDecodeError:
                        continue

        return self.tokenize_and_pack(doc_iter(), output_name)


class TokenDataset:
    """
    Memory-mapped token dataset for efficient training data loading.

    Reads pre-tokenized binary files and yields fixed-length sequences.
    Supports distributed training with proper sharding.
    """

    def __init__(
        self,
        data_path: str,
        max_seq_len: int = 2048,
        dtype: np.dtype = np.uint16,
    ):
        self.max_seq_len = max_seq_len

        # Memory-map the file for efficient random access
        self.data = np.memmap(data_path, dtype=dtype, mode="r")
        self.num_tokens = len(self.data)
        self.num_sequences = self.num_tokens // max_seq_len

        print(f"Loaded {self.num_tokens:,} tokens from {data_path}")
        print(f"  {self.num_sequences} sequences of length {max_seq_len}")

    def __len__(self) -> int:
        return self.num_sequences

    def __getitem__(self, idx: int) -> np.ndarray:
        start = idx * self.max_seq_len
        end = start + self.max_seq_len
        return self.data[start:end].astype(np.int64)

    def get_batch(self, batch_size: int, device: str = "cpu"):
        """Get a random batch of sequences."""
        import torch
        indices = np.random.randint(0, self.num_sequences, size=batch_size)
        batch = np.stack([self[i] for i in indices])
        x = torch.from_numpy(batch).long().to(device)
        # Labels are the same as inputs (shifted internally by the model)
        return x, x.clone()