dhb-xr 0.2.1__py3-none-any.whl

dhb_xr/tokenization/compression.py

@@ -0,0 +1,749 @@
+"""
+Compression layers for DHB-Token sequences.
+
+Implements lossless and lossy compression on discrete token streams from VQ-VAE/RVQ:
+- BPE (Byte-Pair Encoding): Merge frequent token pairs into super-tokens
+- Entropy Coding: Huffman/arithmetic coding based on token frequencies
+- RLE (Run-Length Encoding): Compress repeated tokens in static segments
+- Hierarchical: Secondary quantization for higher compression ratios
+
+References:
+- FAST (Physical Intelligence, 2025): DCT + BPE for action tokenization
+- VQ-VLA (ICCV 2025): Scaling vector-quantized action tokenizers
+- FlashVLA (2025): Token-aware compression and action reuse
+"""
+
+from collections import Counter, defaultdict
+from typing import Dict, List, Optional, Tuple, Union
+import heapq
+import numpy as np
+
+
+class BPECompressor:
+    """
+    Byte-Pair Encoding for token sequences.
+    
+    Merges frequent token pairs into super-tokens, reducing sequence length
+    while preserving exact recoverability (lossless).
+    
+    Inspired by FAST (Physical Intelligence, 2025) which achieves ~10x compression
+    on action sequences via DCT + BPE.
+    
+    Example:
+        >>> compressor = BPECompressor(vocab_size=512, num_merges=100)
+        >>> compressor.fit(token_corpus)  # List of token sequences
+        >>> compressed = compressor.encode([1, 2, 1, 2, 3])  # [256, 256, 3] if (1,2)->256
+        >>> original = compressor.decode(compressed)
+    """
+    
+    def __init__(self, vocab_size: int = 256, num_merges: int = 100):
+        """
+        Args:
+            vocab_size: Original VQ codebook size (tokens 0 to vocab_size-1)
+            num_merges: Number of BPE merges to learn
+        """
+        self.vocab_size = vocab_size
+        self.num_merges = num_merges
+        self.merges: Dict[Tuple[int, int], int] = {}  # (a, b) -> merged_token
+        self.reverse_merges: Dict[int, Tuple[int, int]] = {}  # merged_token -> (a, b)
+        self._fitted = False
+    
+    def fit(self, token_sequences: List[List[int]]) -> "BPECompressor":
+        """
+        Learn BPE merges from a corpus of token sequences.
+        
+        Args:
+            token_sequences: List of token sequences (each a list of ints)
+            
+        Returns:
+            self
+        """
+        # Flatten and count pair frequencies
+        all_tokens = []
+        for seq in token_sequences:
+            all_tokens.extend(list(seq))
+        
+        # Iteratively merge most frequent pairs
+        current_tokens = list(all_tokens)
+        next_id = self.vocab_size
+        
+        for _ in range(self.num_merges):
+            # Count pairs
+            pair_counts = Counter()
+            for i in range(len(current_tokens) - 1):
+                pair = (current_tokens[i], current_tokens[i + 1])
+                pair_counts[pair] += 1
+            
+            if not pair_counts:
+                break
+            
+            # Get most frequent pair
+            best_pair = pair_counts.most_common(1)[0][0]
+            if pair_counts[best_pair] < 2:
+                break  # No benefit from merging singletons
+            
+            # Merge
+            self.merges[best_pair] = next_id
+            self.reverse_merges[next_id] = best_pair
+            
+            # Replace in sequence
+            new_tokens = []
+            i = 0
+            while i < len(current_tokens):
+                if i < len(current_tokens) - 1 and (current_tokens[i], current_tokens[i + 1]) == best_pair:
+                    new_tokens.append(next_id)
+                    i += 2
+                else:
+                    new_tokens.append(current_tokens[i])
+                    i += 1
+            
+            current_tokens = new_tokens
+            next_id += 1
+        
+        self._fitted = True
+        return self
+    
+    def encode(self, tokens: Union[List[int], np.ndarray]) -> List[int]:
+        """
+        Encode a token sequence using learned BPE merges.
+        
+        Args:
+            tokens: Original token sequence
+            
+        Returns:
+            Compressed token sequence
+        """
+        if not self._fitted:
+            raise RuntimeError("BPECompressor must be fitted before encoding")
+        
+        tokens = list(tokens)
+        
+        # Apply merges in order learned
+        for (a, b), merged in self.merges.items():
+            new_tokens = []
+            i = 0
+            while i < len(tokens):
+                if i < len(tokens) - 1 and tokens[i] == a and tokens[i + 1] == b:
+                    new_tokens.append(merged)
+                    i += 2
+                else:
+                    new_tokens.append(tokens[i])
+                    i += 1
+            tokens = new_tokens
+        
+        return tokens
+    
+    def decode(self, tokens: Union[List[int], np.ndarray]) -> List[int]:
+        """
+        Decode a compressed sequence back to original tokens.
+        
+        Args:
+            tokens: Compressed token sequence
+            
+        Returns:
+            Original token sequence
+        """
+        tokens = list(tokens)
+        
+        # Recursively expand merged tokens
+        changed = True
+        while changed:
+            changed = False
+            new_tokens = []
+            for t in tokens:
+                if t in self.reverse_merges:
+                    a, b = self.reverse_merges[t]
+                    new_tokens.extend([a, b])
+                    changed = True
+                else:
+                    new_tokens.append(t)
+            tokens = new_tokens
+        
+        return tokens
+    
+    def compression_ratio(self, tokens: Union[List[int], np.ndarray]) -> float:
+        """Compute compression ratio (original_len / compressed_len)."""
+        compressed = self.encode(tokens)
+        return len(tokens) / len(compressed) if compressed else 1.0
+    
+    @property
+    def extended_vocab_size(self) -> int:
+        """Total vocabulary size including merged tokens."""
+        return self.vocab_size + len(self.merges)
+    
+    def get_stats(self) -> Dict:
+        """Get compression statistics."""
+        return {
+            "original_vocab": self.vocab_size,
+            "num_merges": len(self.merges),
+            "extended_vocab": self.extended_vocab_size,
+            "fitted": self._fitted,
+        }
+
+
+class HuffmanNode:
+    """Node for Huffman tree."""
+    def __init__(self, token: Optional[int], freq: int, left=None, right=None):
+        self.token = token
+        self.freq = freq
+        self.left = left
+        self.right = right
+    
+    def __lt__(self, other):
+        return self.freq < other.freq
+
+
+class EntropyCompressor:
+    """
+    Entropy coding (Huffman) for token sequences.
+    
+    Assigns variable-length codes based on token frequencies,
+    achieving near-optimal bits-per-token based on entropy.
+    
+    For RVQ indices with K=256, naive encoding = 8 bits/token.
+    With entropy coding: typically 4-6 bits/token (1.5-2x compression).
+    """
+    
+    def __init__(self):
+        self.codes: Dict[int, str] = {}  # token -> binary string
+        self.reverse_codes: Dict[str, int] = {}  # binary string -> token
+        self.frequencies: Dict[int, int] = {}
+        self._fitted = False
+    
+    def fit(self, token_sequences: List[List[int]]) -> "EntropyCompressor":
+        """
+        Build Huffman tree from token frequencies.
+        
+        Args:
+            token_sequences: List of token sequences
+            
+        Returns:
+            self
+        """
+        # Count frequencies
+        self.frequencies = Counter()
+        for seq in token_sequences:
+            self.frequencies.update(seq)
+        
+        if not self.frequencies:
+            self._fitted = True
+            return self
+        
+        # Build Huffman tree
+        heap = [HuffmanNode(token, freq) for token, freq in self.frequencies.items()]
+        heapq.heapify(heap)
+        
+        while len(heap) > 1:
+            left = heapq.heappop(heap)
+            right = heapq.heappop(heap)
+            merged = HuffmanNode(None, left.freq + right.freq, left, right)
+            heapq.heappush(heap, merged)
+        
+        # Generate codes
+        self.codes = {}
+        if heap:
+            self._generate_codes(heap[0], "")
+        
+        # Handle single-token case
+        if len(self.codes) == 1:
+            token = list(self.codes.keys())[0]
+            self.codes[token] = "0"
+        
+        self.reverse_codes = {v: k for k, v in self.codes.items()}
+        self._fitted = True
+        return self
+    
+    def _generate_codes(self, node: HuffmanNode, code: str):
+        """Recursively generate Huffman codes."""
+        if node.token is not None:
+            self.codes[node.token] = code if code else "0"
+            return
+        if node.left:
+            self._generate_codes(node.left, code + "0")
+        if node.right:
+            self._generate_codes(node.right, code + "1")
+    
+    def encode(self, tokens: Union[List[int], np.ndarray]) -> str:
+        """
+        Encode tokens to binary string.
+        
+        Args:
+            tokens: Token sequence
+            
+        Returns:
+            Binary string (e.g., "0110101...")
+        """
+        if not self._fitted:
+            raise RuntimeError("EntropyCompressor must be fitted before encoding")
+        return "".join(self.codes.get(t, "") for t in tokens)
+    
+    def decode(self, binary_string: str) -> List[int]:
+        """
+        Decode binary string back to tokens.
+        
+        Args:
+            binary_string: Encoded binary string
+            
+        Returns:
+            Original token sequence
+        """
+        tokens = []
+        current = ""
+        for bit in binary_string:
+            current += bit
+            if current in self.reverse_codes:
+                tokens.append(self.reverse_codes[current])
+                current = ""
+        return tokens
+    
+    def bits_per_token(self, tokens: Union[List[int], np.ndarray]) -> float:
+        """Compute average bits per token."""
+        encoded = self.encode(tokens)
+        return len(encoded) / len(tokens) if tokens else 0.0
+    
+    def theoretical_entropy(self) -> float:
+        """Compute theoretical entropy H = -sum(p * log2(p))."""
+        total = sum(self.frequencies.values())
+        if total == 0:
+            return 0.0
+        entropy = 0.0
+        for freq in self.frequencies.values():
+            p = freq / total
+            if p > 0:
+                entropy -= p * np.log2(p)
+        return entropy
+    
+    def get_stats(self) -> Dict:
+        """Get compression statistics."""
+        return {
+            "unique_tokens": len(self.codes),
+            "theoretical_entropy": self.theoretical_entropy(),
+            "avg_code_length": np.mean([len(c) for c in self.codes.values()]) if self.codes else 0,
+            "fitted": self._fitted,
+        }
+
+
+class RLECompressor:
+    """
+    Run-Length Encoding for token sequences.
+    
+    Compresses repeated tokens (common in static/low-motion segments).
+    Uses a special "repeat" token followed by (token, count) pairs.
+    
+    Example:
+        [5, 5, 5, 5, 3, 3] -> [(REPEAT, 5, 4), (REPEAT, 3, 2)]
+        or simplified: [RLE_MARKER, 5, 4, RLE_MARKER, 3, 2]
+    """
+    
+    def __init__(self, min_run: int = 3, max_count: int = 255):
+        """
+        Args:
+            min_run: Minimum run length to compress (shorter runs kept as-is)
+            max_count: Maximum count per run (limits encoding overhead)
+        """
+        self.min_run = min_run
+        self.max_count = max_count
+        self.RLE_MARKER = -1  # Special marker (will be shifted to valid range)
+    
+    def encode(self, tokens: Union[List[int], np.ndarray]) -> List[Tuple[int, int]]:
+        """
+        Encode tokens with run-length encoding.
+        
+        Args:
+            tokens: Token sequence
+            
+        Returns:
+            List of (token, count) tuples
+        """
+        if len(tokens) == 0:
+            return []
+        
+        tokens = list(tokens)
+        encoded = []
+        
+        i = 0
+        while i < len(tokens):
+            current = tokens[i]
+            count = 1
+            
+            # Count consecutive identical tokens
+            while i + count < len(tokens) and tokens[i + count] == current and count < self.max_count:
+                count += 1
+            
+            encoded.append((current, count))
+            i += count
+        
+        return encoded
+    
+    def decode(self, encoded: List[Tuple[int, int]]) -> List[int]:
+        """
+        Decode RLE-encoded sequence.
+        
+        Args:
+            encoded: List of (token, count) tuples
+            
+        Returns:
+            Original token sequence
+        """
+        tokens = []
+        for token, count in encoded:
+            tokens.extend([token] * count)
+        return tokens
+    
+    def compression_ratio(self, tokens: Union[List[int], np.ndarray]) -> float:
+        """Compute compression ratio."""
+        encoded = self.encode(tokens)
+        # Each (token, count) pair = 2 values vs count original tokens
+        compressed_size = len(encoded) * 2
+        original_size = len(tokens)
+        return original_size / compressed_size if compressed_size > 0 else 1.0
+    
+    def get_stats(self, tokens: Union[List[int], np.ndarray]) -> Dict:
+        """Get RLE statistics for a sequence."""
+        encoded = self.encode(tokens)
+        run_lengths = [count for _, count in encoded]
+        return {
+            "num_runs": len(encoded),
+            "avg_run_length": np.mean(run_lengths) if run_lengths else 0,
+            "max_run_length": max(run_lengths) if run_lengths else 0,
+            "compression_ratio": self.compression_ratio(tokens),
+        }
+
+
+class TokenCompressor:
+    """
+    Unified compression pipeline for DHB-Token sequences.
+    
+    Combines multiple compression methods for optimal results:
+    1. RLE for static segments (lossless, good for low-motion)
+    2. BPE for pattern merging (lossless, 2-4x on invariant patterns)
+    3. Entropy coding for final bitstream (lossless, 1.5-2x additional)
+    
+    Overall achievable: 3-8x compression on typical DHB-Token sequences.
+    
+    Example:
+        >>> compressor = TokenCompressor(vocab_size=256)
+        >>> compressor.fit(training_sequences)
+        >>> 
+        >>> # Compress
+        >>> compressed = compressor.compress(tokens)
+        >>> print(f"Ratio: {compressor.compression_ratio(tokens):.1f}x")
+        >>> 
+        >>> # Decompress (lossless)
+        >>> recovered = compressor.decompress(compressed)
+        >>> assert recovered == list(tokens)
+    """
+    
+    def __init__(
+        self,
+        vocab_size: int = 256,
+        use_rle: bool = True,
+        use_bpe: bool = True,
+        use_entropy: bool = True,
+        bpe_merges: int = 100,
+        rle_min_run: int = 3,
+    ):
+        """
+        Args:
+            vocab_size: VQ codebook size
+            use_rle: Enable run-length encoding
+            use_bpe: Enable byte-pair encoding
+            use_entropy: Enable entropy (Huffman) coding
+            bpe_merges: Number of BPE merges to learn
+            rle_min_run: Minimum run length for RLE
+        """
+        self.vocab_size = vocab_size
+        self.use_rle = use_rle
+        self.use_bpe = use_bpe
+        self.use_entropy = use_entropy
+        
+        self.rle = RLECompressor(min_run=rle_min_run) if use_rle else None
+        self.bpe = BPECompressor(vocab_size=vocab_size, num_merges=bpe_merges) if use_bpe else None
+        self.entropy = EntropyCompressor() if use_entropy else None
+        
+        self._fitted = False
+    
+    def fit(self, token_sequences: List[List[int]]) -> "TokenCompressor":
+        """
+        Fit all compression stages on training data.
+        
+        Args:
+            token_sequences: List of token sequences
+            
+        Returns:
+            self
+        """
+        current_sequences = [list(seq) for seq in token_sequences]
+        
+        # Stage 1: Learn BPE merges
+        if self.bpe:
+            self.bpe.fit(current_sequences)
+            current_sequences = [self.bpe.encode(seq) for seq in current_sequences]
+        
+        # Stage 2: Learn entropy codes (after BPE)
+        if self.entropy:
+            self.entropy.fit(current_sequences)
+        
+        self._fitted = True
+        return self
+    
+    def compress(self, tokens: Union[List[int], np.ndarray]) -> Dict:
+        """
+        Compress a token sequence.
+        
+        Args:
+            tokens: Token sequence
+            
+        Returns:
+            Dict with compressed data and metadata
+        """
+        if not self._fitted:
+            raise RuntimeError("TokenCompressor must be fitted before compressing")
+        
+        tokens = list(tokens)
+        original_len = len(tokens)
+        
+        result = {
+            "original_length": original_len,
+            "stages": {},
+        }
+        
+        # Stage 1: RLE (optional, applied first for static segments)
+        if self.rle:
+            rle_encoded = self.rle.encode(tokens)
+            # Flatten for next stage
+            tokens = []
+            for token, count in rle_encoded:
+                tokens.extend([token, count])
+            result["stages"]["rle"] = {
+                "length": len(rle_encoded),
+                "ratio": original_len / len(rle_encoded) if rle_encoded else 1.0,
+            }
+            result["rle_data"] = rle_encoded
+        
+        # Stage 2: BPE
+        if self.bpe:
+            bpe_encoded = self.bpe.encode(tokens)
+            result["stages"]["bpe"] = {
+                "length": len(bpe_encoded),
+                "ratio": len(tokens) / len(bpe_encoded) if bpe_encoded else 1.0,
+            }
+            tokens = bpe_encoded
+        
+        # Stage 3: Entropy coding
+        if self.entropy:
+            binary = self.entropy.encode(tokens)
+            result["stages"]["entropy"] = {
+                "bits": len(binary),
+                "bytes": len(binary) // 8 + (1 if len(binary) % 8 else 0),
+                "bits_per_original_token": len(binary) / original_len if original_len > 0 else 0,
+            }
+            result["binary"] = binary
+        else:
+            result["tokens"] = tokens
+        
+        # Compute overall ratio
+        if self.entropy:
+            compressed_bits = len(result["binary"])
+            original_bits = original_len * np.ceil(np.log2(self.vocab_size + 1))
+            result["overall_ratio"] = original_bits / compressed_bits if compressed_bits > 0 else 1.0
+        else:
+            result["overall_ratio"] = original_len / len(result.get("tokens", tokens))
+        
+        return result
+    
+    def decompress(self, compressed: Dict) -> List[int]:
+        """
+        Decompress back to original tokens.
+        
+        Args:
+            compressed: Output from compress()
+            
+        Returns:
+            Original token sequence
+        """
+        # Reverse entropy coding
+        if self.entropy and "binary" in compressed:
+            tokens = self.entropy.decode(compressed["binary"])
+        else:
+            tokens = compressed.get("tokens", [])
+        
+        # Reverse BPE
+        if self.bpe:
+            tokens = self.bpe.decode(tokens)
+        
+        # Reverse RLE
+        if self.rle and "rle_data" in compressed:
+            # Reconstruct from (token, count) pairs
+            tokens = self.rle.decode(compressed["rle_data"])
+        
+        return tokens
+    
+    def compression_ratio(self, tokens: Union[List[int], np.ndarray]) -> float:
+        """Compute overall compression ratio for a sequence."""
+        compressed = self.compress(tokens)
+        return compressed["overall_ratio"]
+    
+    def get_stats(self) -> Dict:
+        """Get overall compression statistics."""
+        stats = {
+            "fitted": self._fitted,
+            "stages": [],
+        }
+        if self.rle:
+            stats["stages"].append("RLE")
+        if self.bpe:
+            stats["stages"].append("BPE")
+            stats["bpe"] = self.bpe.get_stats()
+        if self.entropy:
+            stats["stages"].append("Entropy")
+            stats["entropy"] = self.entropy.get_stats()
+        return stats
+
+
+class TokenReuser:
+    """
+    Token reuse detector for inference acceleration.
+    
+    Inspired by FlashVLA (2025): Skip decoding when tokens are stable/repeated,
+    reusing previous outputs. Provides 2-5x effective speedup in long-horizon tasks.
+    
+    Works by detecting:
+    1. Exact token repeats (static segments)
+    2. Token sequences matching known patterns (from database)
+    3. Low-variance token regions (approximate reuse)
+    """
+    
+    def __init__(self, window_size: int = 5, similarity_threshold: float = 0.9):
+        """
+        Args:
+            window_size: Window for detecting stable regions
+            similarity_threshold: Threshold for approximate matching
+        """
+        self.window_size = window_size
+        self.similarity_threshold = similarity_threshold
+        self.pattern_cache: Dict[tuple, np.ndarray] = {}  # pattern -> cached output
+    
+    def detect_stable_regions(self, tokens: Union[List[int], np.ndarray]) -> List[Tuple[int, int, bool]]:
+        """
+        Detect regions where tokens are stable (can reuse previous decoding).
+        
+        Args:
+            tokens: Token sequence
+            
+        Returns:
+            List of (start, end, is_stable) tuples
+        """
+        tokens = np.array(tokens)
+        n = len(tokens)
+        
+        if n < self.window_size:
+            return [(0, n, False)]
+        
+        regions = []
+        i = 0
+        
+        while i < n:
+            # Check if next window_size tokens are identical
+            end = min(i + self.window_size, n)
+            window = tokens[i:end]
+            
+            if len(set(window)) == 1:  # All same
+                # Extend stable region
+                stable_start = i
+                while end < n and tokens[end] == tokens[i]:
+                    end += 1
+                regions.append((stable_start, end, True))
+                i = end
+            else:
+                regions.append((i, i + 1, False))
+                i += 1
+        
+        # Merge adjacent non-stable regions
+        merged = []
+        for start, end, is_stable in regions:
+            if merged and not merged[-1][2] and not is_stable:
+                merged[-1] = (merged[-1][0], end, False)
+            else:
+                merged.append((start, end, is_stable))
+        
+        return merged
+    
+    def compute_reuse_potential(self, tokens: Union[List[int], np.ndarray]) -> Dict:
+        """
+        Analyze reuse potential for a sequence.
+        
+        Args:
+            tokens: Token sequence
+            
+        Returns:
+            Statistics on reuse potential
+        """
+        regions = self.detect_stable_regions(tokens)
+        
+        total_len = len(tokens)
+        stable_len = sum(end - start for start, end, is_stable in regions if is_stable)
+        
+        return {
+            "total_tokens": total_len,
+            "stable_tokens": stable_len,
+            "reuse_fraction": stable_len / total_len if total_len > 0 else 0,
+            "num_regions": len(regions),
+            "num_stable_regions": sum(1 for _, _, s in regions if s),
+            "potential_speedup": total_len / (total_len - stable_len + len([r for r in regions if r[2]])) if total_len > stable_len else 1.0,
+        }
+
+
+# Convenience function
+def compress_token_sequence(
+    tokens: Union[List[int], np.ndarray],
+    vocab_size: int = 256,
+    method: str = "bpe",
+    **kwargs
+) -> Dict:
+    """
+    Compress a token sequence with specified method.
+    
+    Args:
+        tokens: Token sequence
+        vocab_size: VQ codebook size
+        method: "bpe", "entropy", "rle", or "full"
+        **kwargs: Additional arguments for compressor
+        
+    Returns:
+        Compression result dict
+    """
+    if method == "bpe":
+        compressor = BPECompressor(vocab_size=vocab_size, **kwargs)
+        compressor.fit([list(tokens)])  # Self-fit for single sequence
+        encoded = compressor.encode(tokens)
+        return {
+            "encoded": encoded,
+            "ratio": len(tokens) / len(encoded),
+            "method": "bpe",
+        }
+    elif method == "entropy":
+        compressor = EntropyCompressor()
+        compressor.fit([list(tokens)])
+        binary = compressor.encode(tokens)
+        return {
+            "binary": binary,
+            "bits_per_token": len(binary) / len(tokens),
+            "ratio": (len(tokens) * 8) / len(binary),  # vs 8-bit naive
+            "method": "entropy",
+        }
+    elif method == "rle":
+        compressor = RLECompressor(**kwargs)
+        encoded = compressor.encode(tokens)
+        return {
+            "encoded": encoded,
+            "ratio": compressor.compression_ratio(tokens),
+            "stats": compressor.get_stats(tokens),
+            "method": "rle",
+        }
+    elif method == "full":
+        compressor = TokenCompressor(vocab_size=vocab_size, **kwargs)
+        compressor.fit([list(tokens)])
+        return compressor.compress(tokens)
+    else:
+        raise ValueError(f"Unknown method: {method}")