offline_entropy_window_compress_ac.py

import torch
import torch.nn.functional as F
from torch.utils.data import IterableDataset, DataLoader
import json
import numpy as np
from pathlib import Path
from typing import Iterator, List, Dict, Any, Tuple, Optional, Union, Callable
import logging
import argparse
import base64
import gc
from collections import defaultdict, Counter, deque
from m1_compression.batched_arithmetic_coder import (
    _pdf_to_cdf,
)
from m1_compression.hybrid_arithmetic_coder import CPUArithmeticEncoder
from m1_compression import utils
from m1_compression.compressor import (
    load_m1_model_and_tokenizer,
    load_m1_model_cpu,
    ALPHABET_SIZE,
    ARITHMETIC_CODER_BASE,
    ARITHMETIC_CODER_PRECISION,
)
import torch.multiprocessing as mp
from offline_utils import (
    unpack_windows, 
    pseudo_to_packed_bytes, 
    pad_batch, 
    find_next_batch_range,
    packed_bytes_to_pseudo,
    pseudo_to_packed_bytes,
    pad_batch,
    InterleavedJsonlDataset,
    batched_m1_compress_predict_fn,
)
MINIMUM_SEGMENT_SIZE = 3
COMPRESSION_OFFSET = 256
GC_FREQ = 10

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger()


class SegmentCache:
    """Cache for segments"""
    def __init__(self, cache_size: int = 819200, cache_desc: str = "Prediction"):
        self.cache_size = cache_size
        self.cache_desc = cache_desc
        self.cache: Dict[bytes, Union[torch.Tensor, List[int]]] = {}  # segment -> CDF tensor or compressed pseudo bytes
        logger.info(f"Created Cache with size: {cache_size}, type: {cache_desc}")
    
    def get_batch(self, segments: List[bytes]) -> Tuple[List[bytes], List[torch.Tensor], List[int]]:
        """
        Returns:
        - cache_misses: unique segments not in cache
        - cache_results: CDF tensors for segments in cache (in order of hit_indices)
        - hit_indices: indices of segments that were cache hits (in input order)
        """
        segment_to_indices = defaultdict(list)
        for idx, seg in enumerate(segments):
            segment_to_indices[seg].append(idx)
        unique_segments = list(segment_to_indices.keys())

        cache_results = {}
        cache_misses = []
        for seg in unique_segments:
            if seg in self.cache:
                cache_results[seg] = self.cache[seg]
            else:
                cache_misses.append((seg, segment_to_indices[seg]))

        hit_indices = []
        for seg, indices in segment_to_indices.items():
            if seg in cache_results:
                for idx in indices:
                    hit_indices.append(idx)

        logger.info(f"{self.cache_desc} cache: {len(unique_segments)} unique segments, {len(cache_results)} hits, {len(cache_misses)} misses, {len(segments)} total segments")
        return cache_misses, cache_results, hit_indices
    
    def put_batch(self, segments: List[bytes], values: List[Union[torch.Tensor, List[int]]]):
        """Store segment -> value mappings"""
        if self.cache_size <= 0:
            return
        for segment, value in zip(segments, values):
            if segment not in self.cache:
                if len(self.cache) < self.cache_size:
                    if isinstance(value, tuple):
                        assert len(value) == 2 or len(value) == 5, "value must be a tuple of length 2 or 5"
                        cloned_value = tuple(v.clone() if isinstance(v, torch.Tensor) else v for v in value)
                        self.cache[segment] = cloned_value
                    elif isinstance(value, torch.Tensor):
                        self.cache[segment] = value.clone()
                    else:
                        self.cache[segment] = value

def get_batch_size_for_length(window_len, max_batch_size):
    """Determines the batch size for a given window length."""
    BATCH_SIZE_TIERS = {
        128: max_batch_size,
        512: max(max_batch_size // 64, 1),
        1024: max(max_batch_size // 128, 1),
        2048: max(max_batch_size // 256, 1),
    }
    for max_len, batch_size in BATCH_SIZE_TIERS.items():
        if window_len <= max_len:
            return batch_size
    return 1

def segment_prediction_fn(
    batch: List[Dict[str, Any]], 
    max_m1_batch_size,
    batched_predict_fn,
    first_byte_prob,
    debug,
    prediction_cache: Optional[SegmentCache] = None
):
    """
    Consumer: reads from task_queue, compresses, puts result in result_queue.
    """
    all_segments = []
    compressed_or_raw_segments = []
    sample_idx_to_list_segment_idx = defaultdict(list)
    segment_idx = 0
    for sample_idx, item in enumerate(batch):
        assert "windows_starts_lens_b64" in item, "windows_starts_lens_b64 must be in item"
        sample_bytes = item["text"].encode('utf-8')
        byte_windows = unpack_windows(sample_bytes, item["windows_starts_lens_b64"])
        for byte_window_indicator in byte_windows:
            all_segments.append(byte_window_indicator[0])
            compressed_or_raw_segments.append(byte_window_indicator[1])
            sample_idx_to_list_segment_idx[sample_idx].append(segment_idx)
            segment_idx += 1

    effective_segments = []
    ineffective_segments = []
    for orig_idx, (segment, indicator) in enumerate(zip(all_segments, compressed_or_raw_segments)):
        if len(segment) > MINIMUM_SEGMENT_SIZE and indicator == 1:
            effective_segments.append((orig_idx, segment))
        else:
            ineffective_segments.append((orig_idx, segment))

    sorted_effective_segments = sorted(effective_segments, key=lambda x: len(x[1]))
    sorted_idx, sorted_segments = zip(*sorted_effective_segments)
    sorted_segments = list(sorted_segments)  # Convert tuple to list
    effective_segments_idx_map = {
        orig_idx: new_idx 
        for new_idx, orig_idx in enumerate(sorted_idx)
    }
    raw_idx, raw_segments = zip(*ineffective_segments)
    raw_segments = list(raw_segments)
    ineffective_segments_idx_map = {
        orig_idx: new_idx 
        for new_idx, orig_idx in enumerate(raw_idx)
    }
    if debug:
        start_event = torch.cuda.Event(enable_timing=True)
        end_event = torch.cuda.Event(enable_timing=True)
        start_event.record()
        torch.cuda.synchronize()## make sure all previous events are completed
        print("[Debug CUDA] time start", flush=True)

    assert first_byte_prob.shape == (1, 1, ALPHABET_SIZE), "first_byte_prob must be of shape (1, 1, ALPHABET_SIZE)"

    batched_windows_np = [np.frombuffer(bytes(data), dtype=np.uint8) for data in sorted_segments]

    M = len(batched_windows_np)
    batched_cdf_ends = [None] * M  # Pre-allocate to maintain order
    if debug:
        batched_pdfs = [None] * M
    else:
        batched_pdfs = None

    if prediction_cache is not None:
        cache_misses_tup, cache_results, hit_indices = prediction_cache.get_batch(sorted_segments)
        
        # Fill in cache hits
        for hit_idx in hit_indices:
            segment = sorted_segments[hit_idx]
            value = cache_results[segment]
            if debug:
                batched_cdf_ends[hit_idx] = value[0]
                batched_pdfs[hit_idx] = value[1]
            else:
                batched_cdf_ends[hit_idx] = value
        
        # Update miss_indices to only include segments not in cache
        cache_misses, cache_miss_indices = zip(*cache_misses_tup)
    else:
        cache_misses = sorted_segments
        cache_miss_indices = [[i] for i in range(M)]
        
    # Process cache misses
    if cache_misses:
        cache_miss_cdf_ends = []
        cache_miss_pdfs = [] if debug else None
        
        start_idx = 0
        batched_windows_np = [np.frombuffer(bytes(data), dtype=np.uint8) for data in cache_misses]
        miss_count = len(batched_windows_np)
        
        while start_idx < miss_count:
            # Use the new helper function to find the exact range for the next safe batch
            start_idx, end_idx = find_next_batch_range(batched_windows_np, start_idx, max_m1_batch_size, get_batch_size_for_length)
            windows_np_chunked = batched_windows_np[start_idx:end_idx]
            padded_batched_windows, lengths = pad_batch(windows_np_chunked)
            padded_batched_windows, lengths = padded_batched_windows.cuda(), lengths.cuda()
            with torch.no_grad():
                prompt_probs = batched_predict_fn(padded_batched_windows)
                prompt_probs = torch.cat(
                    [
                        first_byte_prob.expand(prompt_probs.shape[0], -1, -1), 
                        prompt_probs[:, :-1, ...]
                    ], 
                    dim=1
                )
            prompt_probs = utils.batched_normalize_pdf_for_arithmetic_coding(prompt_probs)
            cdfs_gpu = _pdf_to_cdf(prompt_probs)
            cdf_low = cdfs_gpu.gather(2, padded_batched_windows.unsqueeze(-1)).squeeze(-1)
            cdf_high = cdfs_gpu.gather(2, (padded_batched_windows + 1).unsqueeze(-1)).squeeze(-1)
            cdf_ends = torch.stack([cdf_low, cdf_high], dim=-1)

            start_idx = end_idx
            if debug:
                cache_miss_pdfs.extend(prompt_probs.cpu())
            cache_miss_cdf_ends.extend(cdf_ends.cpu())
        
        # each miss idx maps a list of original indices
        for idx, miss_indices in enumerate(cache_miss_indices):
            for orig_idx in miss_indices:
                batched_cdf_ends[orig_idx] = cache_miss_cdf_ends[idx]
                if debug:
                    batched_pdfs[orig_idx] = cache_miss_pdfs[idx]
        
        # Store new results in cache
        if prediction_cache is not None:
            if debug:
                prediction_cache.put_batch(cache_misses, zip(cache_miss_cdf_ends, cache_miss_pdfs))
            else:
                prediction_cache.put_batch(cache_misses, cache_miss_cdf_ends)

    return (
        batch,
        sorted_segments, 
        raw_segments, 
        effective_segments_idx_map, 
        ineffective_segments_idx_map, 
        sample_idx_to_list_segment_idx, 
        batched_cdf_ends,
        batched_pdfs,
    )

def segment_compression_fn(
    batch: List[Dict[str, Any]],
    sorted_segments: List[List[int]],
    raw_segments: List[List[int]],
    effective_segments_idx_map: Dict[int, int],
    ineffective_segments_idx_map: Dict[int, int],
    sample_idx_to_list_segment_idx: Dict[int, List[int]],
    batched_cdf_ends: List[torch.Tensor],
    batched_pdfs: List[torch.Tensor],
    output_window_size: int,
    escape_first_byte: bool,
    iterative_compress: bool,
    force_padding_to_threshold: bool,
    predict_fn: Callable,
    first_byte_prob: torch.Tensor,
    debug: bool = False,
    compression_cache: Optional[SegmentCache] = None
):
    ENCODING_BATCH_SIZE = 512 # 384
    if iterative_compress:
        assert not escape_first_byte, "iterative_compress does not support escape_first_byte"
    M = len(batched_cdf_ends)
    processed_batched_compressed_bytes = [None] * M
    if debug:
        batched_stop_steps = [None] * M
        batched_num_padded_bits = [None] * M
        batched_prompt_probs = [None] * M
        batched_lengths = [None] * M
    
    # Check cache for all segments
    if compression_cache is not None:
        cache_misses_tup, cache_results, hit_indices = compression_cache.get_batch(sorted_segments)
        
        # Fill in cache hits
        for hit_idx in hit_indices:
            segment = sorted_segments[hit_idx]
            if debug:
                assert len(cache_results[segment]) == 5, "cache_results must be a tuple of length 5"
                if isinstance(cache_results[segment][0], tuple):
                    processed_batched_compressed_bytes[hit_idx] = cache_results[segment][0][0]
                    batched_stop_steps[hit_idx] = None
                    batched_num_padded_bits[hit_idx] = None
                    batched_prompt_probs[hit_idx] = None
                    batched_lengths[hit_idx] = None
                else:
                    processed_batched_compressed_bytes[hit_idx] = cache_results[segment][0]
                    batched_stop_steps[hit_idx] = cache_results[segment][1]
                    batched_num_padded_bits[hit_idx] = cache_results[segment][2]
                    batched_prompt_probs[hit_idx] = cache_results[segment][3]
                    batched_lengths[hit_idx] = cache_results[segment][4]
            else:
                processed_batched_compressed_bytes[hit_idx] = cache_results[segment]
        
        # Update miss_indices to only include segments not in cache
        cache_misses, cache_miss_indices = zip(*cache_misses_tup)
    else:
        cache_misses = sorted_segments
        cache_miss_indices = [[i] for i in range(M)]

    # Process cache misses
    if cache_misses:
        cache_miss_compressed_bytes = []
        cache_miss_stop_steps = []
        cache_miss_num_padded_bits = []
        cache_miss_prompt_probs = []
        cache_miss_lengths = []
        ######## Use BatchArithmeticEncoder to replace address one by one ###########
        encoder = CPUArithmeticEncoder(
            base=ARITHMETIC_CODER_BASE,
            precision=ARITHMETIC_CODER_PRECISION
        )
        
        # Get CDF ends and segments for cache misses only
        miss_cdf_ends = [batched_cdf_ends[miss_indices[0]] for miss_indices in cache_miss_indices]
        if debug:
            miss_pdfs = [batched_pdfs[miss_indices[0]] for miss_indices in cache_miss_indices]
        else:
            miss_pdfs = None
        miss_count = len(cache_misses)
        
        cache_miss_compressed_results = []
        
        for chunk_idx in range(0, miss_count, ENCODING_BATCH_SIZE):
            chunk_start = chunk_idx
            chunk_end = min(chunk_idx + ENCODING_BATCH_SIZE, miss_count)
            chunk_size = chunk_end - chunk_start
            chunk_segments = cache_misses[chunk_start:chunk_end]
            chunk_cdf_ends = miss_cdf_ends[chunk_start:chunk_end]
            lengths = torch.tensor([len(segment) for segment in chunk_segments], dtype=torch.int64)
            padded_chunk_cdf_ends = torch.zeros(
                (chunk_size, lengths.max().item(), 2),
                device="cpu"
            )
            for idx, (cdf_end, length) in enumerate(zip(chunk_cdf_ends, lengths)):
                padded_chunk_cdf_ends[idx, :length, :] = cdf_end[:length, :]

            if escape_first_byte:
                chunked_compressed_bytes, chunked_stop_steps, chunked_num_padded_bits = encoder.incremental_batched_encode(
                    padded_chunk_cdf_ends[:, 1:, ...],
                    ALPHABET_SIZE,
                    lengths - 1,
                    bit_threshold=output_window_size,
                    force_padding_to_threshold=force_padding_to_threshold,
                    return_num_padded_bits=True
                )
                # if we escape the first byte, we need to add offset 1 to the stop step
                chunked_stop_steps = [step + 1 for step in chunked_stop_steps]
            else:
                chunked_compressed_bytes, chunked_stop_steps, chunked_num_padded_bits = encoder.incremental_batched_encode(
                    padded_chunk_cdf_ends,
                    ALPHABET_SIZE,
                    lengths,
                    bit_threshold=output_window_size,
                    force_padding_to_threshold=force_padding_to_threshold,
                    return_num_padded_bits=True
                )
            cache_miss_compressed_bytes.extend(chunked_compressed_bytes)
            cache_miss_stop_steps.extend(chunked_stop_steps)
            if debug:
                chunk_pdfs = miss_pdfs[chunk_start:chunk_end]
                padded_chunk_pdfs = torch.zeros(
                    (chunk_size, lengths.max().item(), ALPHABET_SIZE),
                    device="cpu"
                )
                for idx, (pdf, length) in enumerate(zip(chunk_pdfs, lengths)):
                    padded_chunk_pdfs[idx, :length, :] = pdf[:length, :]

                if escape_first_byte:
                    cache_miss_num_padded_bits.extend(chunked_num_padded_bits)
                    cache_miss_prompt_probs.extend(padded_chunk_pdfs[:, 1:, ...])
                    cache_miss_lengths.extend(lengths - 1)
                else:
                    cache_miss_num_padded_bits.extend(chunked_num_padded_bits)
                    cache_miss_prompt_probs.extend(padded_chunk_pdfs)
                    cache_miss_lengths.extend(lengths)

            for i in range(chunk_start, chunk_end):
                window_bytes = cache_misses[i]
                stop_step = cache_miss_stop_steps[i]
                _compressed_bytes = list(cache_miss_compressed_bytes[i])
                compressed_bytes = [COMPRESSION_OFFSET + b for b in _compressed_bytes]
                if escape_first_byte:
                    compressed_bytes = list(window_bytes[0:1]) + compressed_bytes
                if stop_step == -1 or stop_step >= len(window_bytes):
                    cache_miss_compressed_results.append(compressed_bytes)
                else:
                    remaining_raw_bytes = list(window_bytes[stop_step:])
                    if iterative_compress and len(remaining_raw_bytes) > MINIMUM_SEGMENT_SIZE:
                        cache_miss_compressed_results.append((remaining_raw_bytes, compressed_bytes))
                    else:
                        compressed_bytes = compressed_bytes + remaining_raw_bytes
                        cache_miss_compressed_results.append(compressed_bytes)

        if iterative_compress:
            incomplete_window_ids = []
            incomplete_window_remaining_bytes = []
            incomplete_window_compressed_bytes = []
            for i, compressed_bytes in enumerate(cache_miss_compressed_results):
                if isinstance(compressed_bytes, tuple):
                    incomplete_window_ids.append(i)
                    incomplete_window_remaining_bytes.append(compressed_bytes[0])
                    incomplete_window_compressed_bytes.append(compressed_bytes[1])

            remaining_compressed_bytes = iterative_compress_ac(
                incomplete_window_remaining_bytes,
                predict_fn,
                first_byte_prob,
                output_window_size,
                force_padding_to_threshold,
                ENCODING_BATCH_SIZE,
                debug
            )
            for i, remaining_compressed_b in enumerate(remaining_compressed_bytes):
                id_in_cache = incomplete_window_ids[i]
                final_compressed_bytes = incomplete_window_compressed_bytes[i] + remaining_compressed_b
                if debug:
                    cache_miss_compressed_results[id_in_cache] = (final_compressed_bytes, "skip_debug")
                else:
                    cache_miss_compressed_results[id_in_cache] = final_compressed_bytes
            logger.info(f"[DEBUG] total remaining windows: {len(incomplete_window_ids)}")
        # Fill in cache misses in the correct positions
        for idx, miss_indices in enumerate(cache_miss_indices):
            for orig_idx in miss_indices:
                if debug:
                    if isinstance(cache_miss_compressed_results[idx], tuple):
                        assert cache_miss_compressed_results[idx][1] == "skip_debug"
                        processed_batched_compressed_bytes[orig_idx] = cache_miss_compressed_results[idx][0]
                        batched_stop_steps[orig_idx] = None
                        batched_num_padded_bits[orig_idx] = None
                        batched_prompt_probs[orig_idx] = None
                        batched_lengths[orig_idx] = None
                    else:
                        processed_batched_compressed_bytes[orig_idx] = cache_miss_compressed_results[idx]
                        batched_stop_steps[orig_idx] = cache_miss_stop_steps[idx]
                        batched_num_padded_bits[orig_idx] = cache_miss_num_padded_bits[idx]
                        batched_prompt_probs[orig_idx] = cache_miss_prompt_probs[idx]
                        batched_lengths[orig_idx] = cache_miss_lengths[idx]
                else:
                    processed_batched_compressed_bytes[orig_idx] = cache_miss_compressed_results[idx]
        # Store new results in cache
        if compression_cache is not None:
            if debug:
                compression_cache.put_batch(
                    cache_misses, 
                    zip(
                        cache_miss_compressed_results, 
                        cache_miss_stop_steps, 
                        cache_miss_num_padded_bits, 
                        cache_miss_prompt_probs, 
                        cache_miss_lengths
                    )
                )
            else:
                compression_cache.put_batch(cache_misses, cache_miss_compressed_results)

    # 4.recompose all segmentations
    B = len(batch)
    #### fix: add pseudo length to split the compressed and raw bytes
    pseudo_lens_per_segment = [[] for _ in range(B)] 
    #### fix end

    compressed_bytes = [[] for _ in range(B)]
    original_bytes = [[] for _ in range(B)]

    for sample_idx, list_segment_idx in sample_idx_to_list_segment_idx.items():
        for segment_idx in list_segment_idx:
            if segment_idx in effective_segments_idx_map:
                compressed_idx = effective_segments_idx_map[segment_idx]
                compressed_byte = processed_batched_compressed_bytes[compressed_idx]
                
            else:
                raw_idx = ineffective_segments_idx_map[segment_idx]
                compressed_byte = raw_segments[raw_idx]
            
            #### fix: whatever the compressed or raw bytes windows,restore the pseudo bytes
            pseudo_lens_per_segment[sample_idx].append(len(compressed_byte))
            #### fix end
            compressed_bytes[sample_idx].extend(list(compressed_byte))
            if debug:
                if segment_idx in effective_segments_idx_map:
                    compressed_idx = effective_segments_idx_map[segment_idx]
                    original_byte = sorted_segments[compressed_idx]

                    _debug_prompt_probs = batched_prompt_probs[compressed_idx]
                    _debug_padded_bits = batched_num_padded_bits[compressed_idx]
                    _debug_lengths = batched_lengths[compressed_idx]
                    _debug_stop_step = batched_stop_steps[compressed_idx]

                    if _debug_prompt_probs is None:
                        original_bytes[sample_idx].append(original_byte)
                        continue

                    processed_compressed_byte = processed_batched_compressed_bytes[compressed_idx]
                    # de-postprocess the compressed byte
                    if escape_first_byte:
                        _debug_escaped_compressed_byte = processed_compressed_byte[1:]
                    else:
                        _debug_escaped_compressed_byte = processed_compressed_byte
                    if _debug_stop_step == -1 or _debug_stop_step >= len(original_byte):
                        _debug_compressed_byte = _debug_escaped_compressed_byte
                        _debug_raw_remaining_bytes = None
                        raw_bytes_len = None
                    else:
                        raw_bytes_len = len(original_byte[_debug_stop_step:])
                        _debug_compressed_byte = _debug_escaped_compressed_byte[:-raw_bytes_len]
                        _debug_raw_remaining_bytes = _debug_escaped_compressed_byte[-raw_bytes_len:]
                    _debug_compressed_byte = [b - COMPRESSION_OFFSET for b in _debug_compressed_byte]

                    print(f"#####  _debug_pdfs is {_debug_prompt_probs.shape}")
                    print(f"#####  _debug_padded is {_debug_padded_bits}")
                    print(f"#####  _debug_compressed is {_debug_compressed_byte}")
                    print(f"#####  _debug_lengths is {_debug_lengths}")
                    print(f"#####  _debug_stop_step is {_debug_stop_step}")
                    print(f"#####  _debug_raw_remaining_bytes is {_debug_raw_remaining_bytes}")
                    print(f"#####  raw_bytes_len is {raw_bytes_len}")
                    print(f"#####  original_byte len is {len(original_byte)}")
                    decoded = encoder.batched_decode(
                        _debug_prompt_probs.unsqueeze(0), 
                        [_debug_compressed_byte], 
                        [_debug_padded_bits], 
                        _debug_lengths.unsqueeze(0)
                    )[0, :_debug_lengths.item()].cpu().tolist()

                    print(f"#####  AC decoded is       {decoded}")
                    if escape_first_byte:
                        decoded = processed_compressed_byte[0:1] + decoded

                        if _debug_stop_step < (_debug_lengths.item() + 1):
                            decoded = decoded[:_debug_stop_step]
                    else:
                        if _debug_stop_step < _debug_lengths.item():
                            decoded = decoded[:_debug_stop_step]
                    print(f"#####  escape_first_byte decoded is {decoded}")
                    if _debug_raw_remaining_bytes:
                        decoded = decoded + _debug_raw_remaining_bytes
                    print(f"#####  decoded is       {decoded}")
                    print(f"#####  original_byte is {list(original_byte)}")
                    assert bytes(decoded) == original_byte, "roundtrip encoding/decoding failed \n{} and \n{}".format(bytes(decoded), original_byte)
                else:
                    raw_idx = ineffective_segments_idx_map[segment_idx]
                    original_byte = raw_segments[raw_idx]
                original_bytes[sample_idx].append(original_byte)

     # --- 关键：内部自验证测试 (仅在 debug 模式下运行) ---
    if debug:
        logger.info("Running internal self-verification test...")
        for i in range(B):
            item = batch[i]
            # 重新获取原始分段信息
            original_segments = unpack_windows(item["text"].encode('utf-8'), item["windows_starts_lens_b64"])
            generated_lens = pseudo_lens_per_segment[i]
            generated_pseudo_list = compressed_bytes[i]
            # 测试 1: 元数据列表的长度必须和原始分段数量一致
            assert len(original_segments) == len(generated_lens), \
                f"Metadata length mismatch for sample {i}: segments={len(original_segments)}, lens={len(generated_lens)}"
            # 测试 2: 使用元数据“走查”一遍生成的伪字节流
            test_ptr = 0
            for j in range(len(original_segments)):
                raw_chunk, indicator = original_segments[j]
                segment_len = generated_lens[j]
                pseudo_slice = generated_pseudo_list[test_ptr : test_ptr + segment_len]
                # 测试 2a: 对于“洞”，内容必须完全一致
                if indicator == 0:
                    assert list(raw_chunk) == pseudo_slice, \
                        f"Hole content mismatch for sample {i}, segment {j}"
                # 移动指针
                test_ptr += segment_len
            # 测试 3: 所有分段长度加起来必须等于总伪字节流长度
            assert test_ptr == len(generated_pseudo_list), \
                f"Total length mismatch for sample {i}: ptr_sum={test_ptr}, total_len={len(generated_pseudo_list)}"
        logger.info("✓ Internal self-verification test passed for all samples in the batch!")
    # --- 自验证测试结束 ---
    
    if debug:
        assert len(compressed_bytes) == len(batch)
        for sample_idx in range(len(batch)):
            assert b"".join(original_bytes[sample_idx]) == batch[sample_idx]["text"].encode('utf-8'), (
                "Assembled original bytes does not match the original batch: \n{} and \n{}".format(
                    b"".join(original_bytes[sample_idx]), batch[sample_idx]["text"].encode('utf-8')
                )
            )
        # window_size_stats = collect_window_size_statistics(original_bytes)
        # logger.info(f"Window size stats: {window_size_stats}")
        # logger.info(f"original_bytes: {original_bytes}")
        # logger.info(f"Finish compressing, Avg compress ratio is ..: {np.mean(compression_ratios):.4f}")
        logger.info(f"Example compressed bytes: {compressed_bytes[0]}")

    write_results = []
    ac_key = f"m1_ac_ow{output_window_size}_escapefb-{escape_first_byte}_iterative-{iterative_compress}_forcepadding-{force_padding_to_threshold}"
    for item, compressed_bytes_item in zip(batch, compressed_bytes):
        item = batch[i]
        compressed = pseudo_to_packed_bytes(compressed_bytes_item)

        result = {
            **item,
            ac_key: base64.b64encode(compressed).decode("ascii"),
            "pseudo_lens_per_segment": pseudo_lens_per_segment[i]
        }
        if debug:
            unpacked = packed_bytes_to_pseudo(compressed)
            assert unpacked == compressed_bytes_item, "Unpacked does not match compressed bytes item: \n{} and \n{}".format(unpacked, compressed_bytes_item)
            logger.info("✓ pseudo-bytes-enc-dec round-trip passes")
        write_results.append(result)
    orig_total_bytes = sum([len(data["text"].encode('utf-8')) for data in batch])
    compressed_total_bytes = sum([len(data) for data in compressed_bytes])
    compression_ratio = orig_total_bytes / compressed_total_bytes if compressed_total_bytes > 0 else 0
    logger.info(f"[DEBUG] original total bytes: {orig_total_bytes}, compressed total bytes: {compressed_total_bytes}, compression rate : {compression_ratio:.3f}")
    return write_results

def iterative_compress_ac(
    batch_windows: List[List[int]],
    predict_fn: Callable,
    first_byte_prob: torch.Tensor,
    output_window_size: int,
    force_padding_to_threshold: bool,
    max_m1_batch_size: int = 4096,
    debug: bool = False,
) -> List[bytes]:
    """
    Buffer-based compression pipeline that reads max_window_size from each file,
    performs batched compression, advances positions based on stop_steps, and repeats.
    """
    
    if debug:
        start_event = torch.cuda.Event(enable_timing=True)
        end_event = torch.cuda.Event(enable_timing=True)
        start_event.record()
        torch.cuda.synchronize()
        print("[Debug CUDA] time start", flush=True)
        original_total_bytes = sum([len(data) for data in batch_windows])
        print(f"[Debug] BufferBased->  Original total bytes: {original_total_bytes}", flush=True)
        print(f"[Debug] BufferBased->  Batch size: {len(batch_windows)}", flush=True)

    B = len(batch_windows)
    # Initialize buffers and positions for each file
    window_positions = [0] * B
    output_compressed_bytes = [[] for _ in range(B)]
    windows_done = [False] * B

    if debug:
        output_padded_bits = [[] for _ in range(B)]
        output_prompt_probs = [[] for _ in range(B)]
        output_lengths = [[] for _ in range(B)]
    
    iter_step = 0
    
    while not all(windows_done):
        iter_step += 1
        
        # Step 1: Read max_window_size bytes from each file to buffer
        current_windows = []
        active_file_indices = []
        
        for i in range(B):
            if windows_done[i]:
                continue
                
            # Read up to max_window_size bytes from current position
            start_pos = window_positions[i]
            end_pos = len(batch_windows[i])
            
            if start_pos >= len(batch_windows[i]) - MINIMUM_SEGMENT_SIZE:
                windows_done[i] = True
                continue
            window_bytes = batch_windows[i][start_pos:end_pos]
            current_windows.append(window_bytes)
            active_file_indices.append(i)
        
        if not current_windows:
            break

        start_idx = 0
        batched_windows_np = [np.array(data, dtype=np.uint8) for data in current_windows]
        current_windows_count = len(batched_windows_np)
        encoder = CPUArithmeticEncoder(
            base=ARITHMETIC_CODER_BASE,
            precision=ARITHMETIC_CODER_PRECISION
        )
        batched_compressed_bytes = []
        batched_stop_steps = []
        if debug:
            batched_num_padded_bits = []
            batched_pdfs = []
        _temp_cdf_ends = []
        _temp_lengths = []
        while start_idx < current_windows_count:
            # Use the new helper function to find the exact range for the next safe batch
            start_idx, end_idx = find_next_batch_range(batched_windows_np, start_idx, max_m1_batch_size, get_batch_size_for_length)
            windows_np_chunked = batched_windows_np[start_idx:end_idx]
            padded_batched_windows, lengths = pad_batch(windows_np_chunked)
            # NOTE: switch to GPU
            padded_batched_windows = padded_batched_windows.cuda()
            with torch.no_grad():
                prompt_probs = predict_fn(padded_batched_windows)
                prompt_probs = torch.cat(
                    [
                        first_byte_prob.expand(prompt_probs.shape[0], -1, -1), 
                        prompt_probs[:, :-1, ...]
                    ], 
                    dim=1
                )
            prompt_probs = utils.batched_normalize_pdf_for_arithmetic_coding(prompt_probs)
            cdfs_gpu = _pdf_to_cdf(prompt_probs)
            cdf_low = cdfs_gpu.gather(2, padded_batched_windows.unsqueeze(-1)).squeeze(-1)
            cdf_high = cdfs_gpu.gather(2, (padded_batched_windows + 1).unsqueeze(-1)).squeeze(-1)
            cdf_ends = torch.stack([cdf_low, cdf_high], dim=-1)

            start_idx = end_idx

            _temp_cdf_ends.append(cdf_ends.cpu())
            _temp_lengths.append(lengths)
            if debug:
                batched_pdfs.extend(prompt_probs.cpu())

        for cdf_ends, lengths in zip(_temp_cdf_ends, _temp_lengths):
            chunked_compressed_bytes, chunked_stop_steps, chunked_num_padded_bits = encoder.incremental_batched_encode(
                # NOTE: switch to CPU
                cdf_ends,
                ALPHABET_SIZE,
                lengths,
                bit_threshold=output_window_size,
                force_padding_to_threshold=force_padding_to_threshold,
                return_num_padded_bits=True
            )
            batched_compressed_bytes.extend(chunked_compressed_bytes)
            batched_stop_steps.extend(chunked_stop_steps)
            if debug:
                batched_num_padded_bits.extend(chunked_num_padded_bits)
        # NOTE: debug this function
        # Step 3: Process results and advance positions
        for window_idx, file_idx in enumerate(active_file_indices):
            compressed_bytes = batched_compressed_bytes[window_idx]
            stop_step = batched_stop_steps[window_idx]
            
            # Add compressed bytes to output
            output_compressed_bytes[file_idx].append(compressed_bytes)
            if debug:
                output_padded_bits[file_idx].append(batched_num_padded_bits[window_idx])
                output_prompt_probs[file_idx].append(batched_pdfs[window_idx])
                length = torch.tensor([stop_step], dtype=torch.long, device=batched_pdfs[window_idx].device)
                output_lengths[file_idx].append(length)

            window_positions[file_idx] += stop_step
            if window_positions[file_idx] >= len(batch_windows[file_idx]) - MINIMUM_SEGMENT_SIZE:
                windows_done[file_idx] = True

    # Concatenate all compressed bytes for each file
    final_compressed = []
    for i in range(B):
        _original_byte_window = batch_windows[i]
        _stopped_position = window_positions[i]

        _byte_array = b''.join(output_compressed_bytes[i])
        offset_compressed_bytes = [b + COMPRESSION_OFFSET for b in list(_byte_array)]
        if _stopped_position < len(_original_byte_window):
            raw_leftover_bytes = _original_byte_window[_stopped_position:]
            offset_compressed_bytes = offset_compressed_bytes + list(raw_leftover_bytes)
        
        final_compressed.append(offset_compressed_bytes)
    if debug:
        end_event.record()
        torch.cuda.synchronize()
        elapsed_time = start_event.elapsed_time(end_event)
        print(f"[Debug CUDA]  Elapsed time: {elapsed_time:.3f}ms", flush=True)
        encoder = CPUArithmeticEncoder(
            base=ARITHMETIC_CODER_BASE,
            precision=ARITHMETIC_CODER_PRECISION
        )
        for (
            output_compressed_bytes_item,
            output_padded_bits_item,
            output_prompt_probs_item,
            output_lengths_item,
            batch_windows_item,
            stopped_position
        ) in zip(output_compressed_bytes, output_padded_bits, output_prompt_probs, output_lengths, batch_windows, window_positions):
            original_bytes = batch_windows_item[:stopped_position]
            decoded_bytes = []
            for (
                _debug_compressed,
                _debug_padded,
                _debug_pdfs,
                _debug_lengths
            ) in zip(
                output_compressed_bytes_item,
                output_padded_bits_item, 
                output_prompt_probs_item, 
                output_lengths_item
            ):
                print(f"#####  _debug_pdfs is {_debug_pdfs.shape}")
                print(f"#####  _debug_padded is {_debug_padded}")
                print(f"#####  _debug_compressed is {_debug_compressed}")
                print(f"#####  _debug_lengths is {_debug_lengths}")
                print(f"#####  original_bytes is {original_bytes}")
                decoded = encoder.batched_decode(_debug_pdfs.unsqueeze(0), [_debug_compressed], [_debug_padded], _debug_lengths)
                decoded_bytes += decoded[0, :_debug_lengths.item()].cpu().tolist()
                print(f"#####  decoded is {bytes(decoded[0, :_debug_lengths.item()].cpu().tolist())}")
            assert decoded_bytes == original_bytes, "roundtrip encoding/decoding failed \n{} and \n{}".format(decoded_bytes, original_bytes)

    return final_compressed

def writer_consumer(
    write_queue, 
    output_file, 
    buffer_size=100, 
    debug=False, 
    output_window_size=16, 
    escape_first_byte=False, 
    compression_cache_size=819200,
    iterative_compress=False,
    force_padding_to_threshold=False,
    entropy_model_path=None,
    firstbyte_prob_path=None,
    num_workers=None,
    ):
    """
    Writer consumer: reads compressed results from write_queue and writes to file.
    Maintains its own buffer and writes when buffer is full or receives sentinel.
    """
    if num_workers is not None:
        num_threads = torch.get_num_threads()
        # new_num_threads = max(1, int(num_threads // 2 // num_workers))
        # TODO: HACK
        new_num_threads = 1 # max(1, int(num_threads // (num_workers + 1)))
        torch.set_num_threads(new_num_threads)
        logger.info(f"[Debug] Set num threads to {new_num_threads} for writer process {mp.current_process().name}")
    write_buf = []
    
    # Initialize compression cache for this worker
    compression_cache = SegmentCache(cache_size=compression_cache_size, cache_desc="Compression") if compression_cache_size > 0 else None

    if iterative_compress:
        model, _, _ = load_m1_model_and_tokenizer(entropy_model_path)
        predict_fn = batched_m1_compress_predict_fn(model)

        if firstbyte_prob_path is not None:
            with open(firstbyte_prob_path, 'r', encoding='utf-8') as f:
                first_byte_prob = json.load(f)
            print(first_byte_prob)
            first_byte_prob = torch.tensor(first_byte_prob, dtype=torch.float32, device="cuda").unsqueeze(0).unsqueeze(0)
        else:
            first_byte_prob = torch.ones((1, 1, ALPHABET_SIZE), dtype=torch.float32, device="cuda") / ALPHABET_SIZE
        # NOTE: use CPU
        # model = load_m1_model_cpu(entropy_model_path)
        # predict_fn = batched_m1_compress_predict_fn(model)

        # if firstbyte_prob_path is not None:
        #     with open(firstbyte_prob_path, 'r', encoding='utf-8') as f:
        #         first_byte_prob = json.load(f)
        #     print(first_byte_prob)
        #     first_byte_prob = torch.tensor(first_byte_prob, dtype=torch.float32, device="cpu").unsqueeze(0).unsqueeze(0)
        # else:
        #     first_byte_prob = torch.ones((1, 1, ALPHABET_SIZE), dtype=torch.float32, device="cpu") / ALPHABET_SIZE
    else:
        predict_fn = None
        first_byte_prob = None

    try:
        with open(output_file, 'w', encoding='utf-8') as f:
            while True:
                args = write_queue.get()
                if args is None:
                    break
                (
                    batch, 
                    sorted_segments, 
                    raw_segments, 
                    effective_segments_idx_map, 
                    ineffective_segments_idx_map, 
                    sample_idx_to_list_segment_idx, 
                    batched_cdf_ends,
                    batched_pdfs,
                ) = args
                write_results = segment_compression_fn(
                    batch,
                    sorted_segments,
                    raw_segments,
                    effective_segments_idx_map,
                    ineffective_segments_idx_map,
                    sample_idx_to_list_segment_idx,
                    batched_cdf_ends,
                    batched_pdfs,
                    output_window_size,
                    escape_first_byte,
                    iterative_compress,
                    force_padding_to_threshold,
                    predict_fn,
                    first_byte_prob,
                    debug=debug,
                    compression_cache=compression_cache
                )
                write_buf.extend(write_results)
                
                # Write buffer when it's full
                if len(write_buf) >= buffer_size:
                    logger.info(f"Writer: Dumping buffer of {len(write_buf)} items to {output_file}")
                    for buffered_item in write_buf:
                        f.write(json.dumps(buffered_item) + '\n')
                    f.flush()
                    write_buf = []

                    # Clean up GPU memory
                    gc.collect()
                    if torch.cuda.is_available():
                        torch.cuda.empty_cache()

            # Write remaining items in buffer
            if write_buf:
                logger.info(f"Writer: Dumping remaining {len(write_buf)} items to {output_file}")
                for buffered_item in write_buf:
                    f.write(json.dumps(buffered_item) + '\n')
                f.flush()
                    
    except Exception as e:
        logger.error(f"Writer process error: {e}")
        raise

def merge_output_files(output_file, writer_output_files):
    """Merge all writer output files into a single file"""
    logger.info(f"Merging {len(writer_output_files)} writer files into {output_file}")
    
    with open(output_file, 'w', encoding='utf-8') as outf:
        for writer_output_file in writer_output_files:
            if writer_output_file.exists():
                with open(writer_output_file, 'r', encoding='utf-8') as inf:
                    for line in inf:
                        outf.write(line)
                # Optionally remove the individual writer files
                writer_output_file.unlink()
                logger.info(f"Merged and removed writer file: {writer_output_file}")
    
    logger.info(f"Merged output written to: {output_file}")
    return output_file

def shutdown_writers(write_queue, writer_processes):
    """Send shutdown signals to shared queue and wait for all writers to complete"""
    # Send one sentinel per writer to ensure all writers get the shutdown signal
    for i in range(len(writer_processes)):
        write_queue.put(None)
        logger.info(f"Sent shutdown signal {i+1}/{len(writer_processes)}")
    
    # Wait for all writers to complete
    for i, writer_process in enumerate(writer_processes):
        writer_process.join()
        if writer_process.exitcode != 0:
            logger.error(f"Writer process {i} failed with exit code: {writer_process.exitcode}")
        else:
            logger.info(f"Writer process {i} completed successfully")


def main():
    # Set up argument parser
    parser = argparse.ArgumentParser(description='Process JSONL files using M1 arithmetic compression with buffer-based approach')
    parser.add_argument('--input_file', type=str, required=True,
                      help='Directory containing input JSONL files')
    parser.add_argument('--output_dir', type=str, required=True,
                      help='Directory to write compressed results')
    parser.add_argument('--entropy_model_path', type=str, required=True,
                      help='Path to the M1 model checkpoint')
    parser.add_argument('--compression_model_path', type=str, required=True,
                      help='Path to the M1 model checkpoint')
    parser.add_argument('--data_batch_size', type=int, default=512,
                      help='Size of batches for processing (default: 512)')
    parser.add_argument('--output_window_size', type=int, default=16,
                      help='Size of window for compression (default: 16)')
    parser.add_argument('--escape_first_byte', action='store_true', default=False,
                      help='Escape the first byte of each window (default: False)')
    parser.add_argument('--max_window_size', type=int, default=1024,
                      help='Maximum window size for reading from each file (default: 1024)')
    parser.add_argument('--max_entropy_batch_size', type=int, default=4096,
                      help='Size of max batch for compression (default: 4096)')
    parser.add_argument('--max_compression_batch_size', type=int, default=4096,
                      help='Size of max batch for compression (default: 4096)')
    parser.add_argument('--chunk_size', type=int, default=512,
                      help='Size of chunk for compression (default: 512)')
    parser.add_argument('--base_global_quantile', type=float, default=0.9,
                      help='Base global quantile for compression (default: 0.9)')
    parser.add_argument('--base_monotonic_quantile', type=float, default=0.9,
                      help='Base monotonic quantile for compression (default: 0.9)')
    parser.add_argument('--debug', action='store_true', default=False,
                      help='Debug mode (default: False)')
    parser.add_argument('--firstbyte_prob_path', type=str, default=None,
                      help='Probability path for the first word of each window (default : None)')
    parser.add_argument('--num_workers', type=int, default=1,
                      help='Number of workers for CPU jobs (default: 1)')
    parser.add_argument('--process_id', type=int, default=0,
                      help='Process ID for distributed processing (default: 0)')
    parser.add_argument('--num_processes', type=int, default=1,
                      help='Number of processes for distributed processing (default: 1)')
    parser.add_argument('--merge_output', action='store_true', default=False,
                      help='Merge all writer output files into a single file (default: False)')
    parser.add_argument('--prediction_cache_size', type=int, default=81920,
                      help='Size of prediction cache per process (default: 819200)')
    parser.add_argument('--compression_cache_size', type=int, default=81920,
                      help='Size of compression cache per worker (default: 819200)')
    parser.add_argument('--disable_caching', action='store_true', default=False,
                      help='Disable both prediction and compression caching (default: False)')
    parser.add_argument('--iterative_compress', action='store_true', default=False,
                      help='Iterative compression (default: False)')
    parser.add_argument('--force_padding_to_threshold', action='store_true', default=False,
                      help='Force padding to threshold (default: False)')

    args = parser.parse_args()

    num_threads = torch.get_num_threads()
    # new_num_threads = max(1, int(num_threads // 2))
    # TODO: HACK
    new_num_threads = 2 # max(1, int(num_threads // (args.num_workers + 1)))
    torch.set_num_threads(new_num_threads)
    logger.info(f"[Debug] Set num threads to {new_num_threads} for main process")

    mp.set_start_method('spawn', force=True)
    dump_freq = 100

    # Create output directory if it doesn't exist
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Load model and tokenizer
    model, _, _ = load_m1_model_and_tokenizer(args.entropy_model_path)
    batched_predict_fn = batched_m1_compress_predict_fn(model)

    if args.firstbyte_prob_path is not None:
        with open(args.firstbyte_prob_path, 'r', encoding='utf-8') as f:
            first_byte_prob = json.load(f)
        print(first_byte_prob)
        first_byte_prob = torch.tensor(first_byte_prob, dtype=torch.float32, device="cuda").unsqueeze(0).unsqueeze(0)
    else:
        first_byte_prob = torch.ones((1, 1, ALPHABET_SIZE), dtype=torch.float32, device="cuda") / ALPHABET_SIZE

    # Create dataset and dataloader
    dataset = InterleavedJsonlDataset(
        file_path=args.input_file,
        rank=args.process_id,
        world_size=args.num_processes,
    )
    dataloader = DataLoader(
        dataset,
        batch_size=args.data_batch_size,
        shuffle=False,
        collate_fn=lambda x: x
    )

    input_file = Path(args.input_file)
    logger.info(f"Processing file: {input_file}")
    
    output_file = output_dir / f"{input_file.stem}_out_{args.process_id}.jsonl"
    
    logger.info("Data loaded. Start processing...")

    # Initialize prediction cache for this process
    prediction_cache = None
    if not args.disable_caching and args.prediction_cache_size > 0:
        prediction_cache = SegmentCache(cache_size=args.prediction_cache_size, cache_desc="Prediction")
        logger.info(f"Prediction cache enabled with size: {args.prediction_cache_size}")
    else:
        logger.info("Prediction cache disabled")

    compression_cache_size = 0 if args.disable_caching else args.compression_cache_size
    if compression_cache_size > 0:
        logger.info(f"Compression cache enabled with size: {compression_cache_size} per worker")
    else:
        logger.info("Compression cache disabled")

    write_queue = mp.Queue(maxsize=200)
    writer_processes = []
    writer_output_files = []
    for i in range(args.num_workers):
        # Create unique output file for each writer
        output_path = Path(output_file)
        writer_output_file = output_path.parent / f"{output_path.stem}_writer_{i}.jsonl"
        writer_output_files.append(writer_output_file)
        writer_process = mp.Process(
            target=writer_consumer,
            args=(
                write_queue, 
                writer_output_file, 
                dump_freq, 
                args.debug, 
                args.output_window_size, 
                args.escape_first_byte, 
                compression_cache_size, 
                args.iterative_compress, 
                args.force_padding_to_threshold,
                args.entropy_model_path, 
                args.firstbyte_prob_path,
                args.num_workers,
            )
        )
        
        writer_processes.append(writer_process)
        writer_process.start()
        logger.info(f"Started writer process {i} for output file: {writer_output_file}")
    
    try:
        # Process each batch
        for batch_idx, batch in enumerate(dataloader):
            pred_results = segment_prediction_fn(
                batch,
                max_m1_batch_size=args.max_compression_batch_size,
                batched_predict_fn=batched_predict_fn,
                first_byte_prob=first_byte_prob,
                debug=args.debug,
                prediction_cache=prediction_cache
            )
            logger.info(f"Processed batch {batch_idx}")
            write_queue.put(pred_results)

            if batch_idx % GC_FREQ == 0:
                # Clean up GPU memory
                gc.collect()
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
        
        # Signal completion to all writer processes
        shutdown_writers(write_queue, writer_processes)
        
    except Exception as e:
        logger.error(f"Error during processing: {e}")
        # Try to terminate writer processes cleanly
        try:
            shutdown_writers(write_queue, writer_processes)
        except:
            pass
        raise

    if args.merge_output:
        final_output_file = merge_output_files(output_file, writer_output_files)
        logger.info(f"Completed processing successfully, merged output written to {final_output_file}")
    else:
        logger.info(f"Completed processing successfully, outputs written to {args.num_workers} separate files")

if __name__ == "__main__":
    main()