INFO:__main__:Merged pair ('tokenization', '.</w>') (Total merges: 66).
INFO:__main__:No more pairs to merge.
INFO:__main__:BPE vocabulary saved to bpe_vocab.json.gz.
INFO:__main__:BPE vocabulary loaded from bpe_vocab.json.gz.
INFO:__main__:BPETokenizer initialized.
Encoded tokens: []
Decoded text:
import gzip
import html
import os
import logging
from collections import defaultdict, Counter
from functools import lru_cache
from typing import Dict, List, Tuple, Set, Optional, Iterator

import ftfy
import regex as re

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def default_bpe_vocab_path() -> str:
    """
    Returns the default path for the BPE vocabulary file.

    Returns:
        str: Path to the default BPE vocabulary file.
    """
    return os.path.join(os.path.dirname(os.path.abspath(__file__)), "bpe_vocab.json.gz")


def bytes_to_unicode_mapping() -> Dict[int, str]:
    """
    Creates a reversible mapping between byte values and unicode strings.

    Returns:
        Dict[int, str]: Mapping from byte values to unicode characters.
    """
    bs = list(range(ord("!"), ord("~") + 1)) + \
         list(range(ord("¡"), ord("¬") + 1)) + \
         list(range(ord("®"), ord("ÿ") + 1))
    cs = bs[:]
    n = 0
    for b in range(2**8):
        if b not in bs:
            bs.append(b)
            cs.append(2**8 + n)
            n += 1
    cs = [chr(c) for c in cs]
    mapping = dict(zip(bs, cs))
    logger.debug("Created byte to unicode mapping.")
    return mapping


def get_symbol_pairs(word: Tuple[str, ...]) -> Set[Tuple[str, str]]:
    """
    Retrieves all adjacent symbol pairs in a given word.

    Args:
        word (Tuple[str, ...]): The word represented as a tuple of symbols.

    Returns:
        Set[Tuple[str, str]]: Set of adjacent symbol pairs.
    """
    return { (word[i], word[i + 1]) for i in range(len(word) - 1) }


def basic_clean(text: str) -> str:
    """
    Cleans the input text using ftfy and HTML unescaping.

    Args:
        text (str): The input text to clean.

    Returns:
        str: Cleaned text.
    """
    text = ftfy.fix_text(text)
    text = html.unescape(text)
    return text.strip()


def whitespace_clean(text: str) -> str:
    """
    Cleans up whitespace in the input text.

    Args:
        text (str): The input text to clean.

    Returns:
        str: Whitespace-cleaned text.
    """
    text = re.sub(r'\s+', ' ', text)
    return text.strip()


class BPETrainer:
    """
    A class to train Byte Pair Encoding (BPE) vocabulary from a corpus.
    """

    def __init__(self, num_merges: int = 10000) -> None:
        """
        Initializes the BPETrainer.

        Args:
            num_merges (int, optional): Number of BPE merge operations. Defaults to 10000.
        """
        self.num_merges = num_merges
        self.vocab: Dict[Tuple[str, ...], int] = {}
        logger.info(f"BPETrainer initialized with {self.num_merges} merges.")

    def train(self, corpus: Iterator[str]) -> None:
        """
        Trains the BPE vocabulary on the provided corpus.

        Args:
            corpus (Iterator[str]): An iterator over lines of text.
        """
        # Count frequency of each word in the corpus
        word_freq = Counter()
        for line in corpus:
            line = whitespace_clean(basic_clean(line)).lower()
            tokens = line.split()
            for token in tokens:
                # Represent token as tuple of characters with end-of-word marker
                symbol_sequence = tuple(token) + ('</w>',)
                word_freq[symbol_sequence] += 1
        self.vocab = dict(word_freq)
        logger.info(f"Counted {len(self.vocab)} unique words.")

        for i in range(self.num_merges):
            pairs = self._get_stats()
            if not pairs:
                logger.info("No more pairs to merge.")
                break
            best = max(pairs, key=pairs.get)
            frequency = pairs[best]
            logger.debug(f"Step {i + 1}: Merging pair {best} with frequency {frequency}.")
            self._merge_vocab(best)
            logger.info(f"Merged pair {best} (Total merges: {i + 1}).")

    def _get_stats(self) -> Dict[Tuple[str, str], int]:
        """
        Computes frequencies of adjacent symbol pairs in the current vocabulary.

        Returns:
            Dict[Tuple[str, str], int]: Frequencies of symbol pairs.
        """
        pairs = defaultdict(int)
        for word, freq in self.vocab.items():
            for pair in get_symbol_pairs(word):
                pairs[pair] += freq
        logger.debug(f"Computed {len(pairs)} symbol pairs.")
        return pairs

    def _merge_vocab(self, pair: Tuple[str, str]) -> None:
        """
        Merges all occurrences of the specified pair in the vocabulary.

        Args:
            pair (Tuple[str, str]): The pair of symbols to merge.
        """
        bigram = ''.join(pair)
        new_vocab = {}
        for word, freq in self.vocab.items():
            new_word = []
            i = 0
            while i < len(word):
                if i < len(word) - 1 and word[i] == pair[0] and word[i + 1] == pair[1]:
                    new_word.append(bigram)
                    i += 2
                else:
                    new_word.append(word[i])
                    i += 1
            new_vocab[tuple(new_word)] = freq
        self.vocab = new_vocab
        logger.debug(f"Vocabulary size after merge: {len(self.vocab)}.")

    def save_vocab(self, file_path: str) -> None:
        """
        Saves the trained BPE vocabulary to a gzip-compressed JSON file.

        Args:
            file_path (str): Path to save the vocabulary file.
        """
        import json
        # Convert tuple keys to space-separated strings for JSON serialization
        vocab_to_save = {' '.join(k): v for k, v in self.vocab.items()}
        with gzip.open(file_path, 'wt', encoding='utf-8') as f:
            json.dump(vocab_to_save, f)
        logger.info(f"BPE vocabulary saved to {file_path}.")

    def load_vocab(self, file_path: str) -> None:
        """
        Loads the BPE vocabulary from a gzip-compressed JSON file.

        Args:
            file_path (str): Path to the vocabulary file.
        """
        import json
        with gzip.open(file_path, 'rt', encoding='utf-8') as f:
            self.vocab = {tuple(k.split()): v for k, v in json.load(f).items()}
        logger.info(f"BPE vocabulary loaded from {file_path}.")

class BPETokenizer:
    """
    A Byte Pair Encoding (BPE) tokenizer for encoding and decoding text.
    """

    def __init__(self, vocab_path: Optional[str] = None) -> None:
        """
        Initializes the BPETokenizer.

        Args:
            vocab_path (Optional[str], optional): Path to the BPE vocabulary file. Defaults to None.
        """
        self.vocab: Dict[Tuple[str, ...], int] = {}
        self.bpe_ranks: Dict[Tuple[str, str], int] = {}
        self.byte_encoder = bytes_to_unicode_mapping()
        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
        self.cache: Dict[str, str] = {}
        self.pat = re.compile(
            r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]+|[^\s\p{L}\p{N}]+""",
            re.IGNORECASE
        )
        if vocab_path:
            self.load_vocab(vocab_path)
        logger.info("BPETokenizer initialized.")

    def load_vocab(self, file_path: str) -> None:
        """
        Loads the BPE vocabulary from a gzip-compressed JSON file.

        Args:
            file_path (str): Path to the vocabulary file.
        """
        import json
        try:
            with gzip.open(file_path, 'rt', encoding='utf-8') as f:
                self.vocab = {tuple(k.split()): v for k, v in json.load(f).items()}
            # Sort pairs by their rank
            self.bpe_ranks = {pair: rank for rank, pair in enumerate(sorted(
                self._get_all_pairs(),
                key=lambda p: self.vocab.get(p, 0),
                reverse=True
            ))}
            logger.info(f"BPE vocabulary loaded from {file_path}.")
        except (IOError, json.JSONDecodeError) as e:
            logger.error(f"Failed to load vocabulary from {file_path}: {e}")
            raise

    def _get_all_pairs(self) -> Set[Tuple[str, str]]:
        """
        Retrieves all symbol pairs from the vocabulary.

        Returns:
            Set[Tuple[str, str]]: All symbol pairs.
        """
        pairs = set()
        for word in self.vocab.keys():
            pairs.update(get_symbol_pairs(word))
        return pairs

    @lru_cache(maxsize=None)
    def bpe(self, token: str) -> str:
        """
        Applies BPE to a single token.

        Args:
            token (str): The token to apply BPE on.

        Returns:
            str: The BPE-encoded token.
        """
        if token in self.cache:
            return self.cache[token]

        # Convert to unicode bytes
        token_bytes = token.encode('utf-8')
        token_chars = ''.join([self.byte_encoder[b] for b in token_bytes])
        word = tuple(token_chars) + ('</w>',)
        pairs = get_symbol_pairs(word)

        if not pairs:
            return token + '</w>'

        while True:
            # Find the highest ranked pair
            min_pair = min(
                pairs,
                key=lambda pair: self.bpe_ranks.get(pair, float('inf'))
            )
            if min_pair not in self.bpe_ranks:
                break
            first, second = min_pair
            new_word = []
            i = 0
            while i < len(word):
                try:
                    j = word.index(first, i)
                except ValueError:
                    new_word.extend(word[i:])
                    break
                new_word.extend(word[i:j])
                if j < len(word) - 1 and word[j + 1] == second:
                    new_word.append(first + second)
                    i = j + 2
                else:
                    new_word.append(word[j])
                    i = j + 1
            new_word = tuple(new_word)
            word = new_word
            if len(word) == 1:
                break
            else:
                pairs = get_symbol_pairs(word)

        word_str = ' '.join(word)
        self.cache[token] = word_str
        return word_str

    def encode(self, text: str) -> List[int]:
        """
        Encodes input text into a list of BPE token IDs.

        Args:
            text (str): The input text to encode.

        Returns:
            List[int]: List of token IDs.
        """
        bpe_tokens = []
        try:
            cleaned_text = whitespace_clean(basic_clean(text)).lower()
            for token in re.findall(self.pat, cleaned_text):
                bpe_token = self.bpe(token)
                token_ids = [self.vocab[pair] for pair in bpe_token.split(' ') if pair in self.vocab]
                bpe_tokens.extend(token_ids)
            logger.debug(f"Encoded text into {len(bpe_tokens)} tokens.")
            return bpe_tokens
        except Exception as e:
            logger.error(f"Error during encoding: {e}")
            raise

    def decode(self, tokens: List[int]) -> str:
        """
        Decodes a list of BPE token IDs back into text.

        Args:
            tokens (List[int]): List of token IDs to decode.

        Returns:
            str: The decoded text.
        """
        try:
            reversed_vocab = {v: k for k, v in self.vocab.items()}
            word = ''.join([reversed_vocab.get(token, '') for token in tokens])
            # Convert unicode bytes back to normal bytes
            byte_sequence = bytearray([self.byte_decoder[char] for char in word if char in self.byte_decoder])
            text = byte_sequence.decode('utf-8', errors='replace').replace('</w>', ' ').strip()
            logger.debug("Decoding completed.")
            return text
        except Exception as e:
            logger.error(f"Error during decoding: {e}")
            raise


def main() -> None:
    """
    Example usage of BPETrainer and BPETokenizer.
    """
    # Sample corpus
    sample_corpus = [
        "This is a sample sentence.",
        "This is another sentence for testing BPE tokenizer.",
        "Byte pair encoding is useful for tokenization."
    ]

    # Initialize and train BPETrainer
    trainer = BPETrainer(num_merges=100)
    trainer.train(iter(sample_corpus))
    vocab_path = "bpe_vocab.json.gz"
    trainer.save_vocab(vocab_path)

    # Initialize BPETokenizer with the trained vocabulary
    tokenizer = BPETokenizer(vocab_path=vocab_path)

    # Encode and decode a sample text
    sample_text = "This is a test sentence for the BPE tokenizer."
    encoded = tokenizer.encode(sample_text)
    print(f"Encoded tokens: {encoded}")

    decoded = tokenizer.decode(encoded)
    print(f"Decoded text: {decoded}")


if __name__ == "__main__":
    main()