offline_entropy_window_compress_v2.py

"""
new pipeline:
    1.prepare data: from original batch to unpack,sort,rerank windows and reconstruct indexes
    2.compress_segment_xxx(): use different compression algorithm
    3.reconstruct_result: reconstruct the window from compressed results and idx
    4.main()--- use prepare_segments and compress_seg_xx to produce and pass them to consumers
    5.write_consumer: get from compressed data, reconstruct and write the result
"""
import torch
import torch.nn.functional as F
from torch.utils.data import IterableDataset, Dataset, DataLoader
import json
import numpy as np
from pathlib import Path
from typing import Iterator, List, Dict, Any, Callable, Tuple, Optional
import logging
import argparse
import base64
import time
import math
import gc
from collections import defaultdict, Counter,deque 
from m1_compression.utils import *
from m1_compression.compressor import (
    load_m1_model_and_tokenizer,
    ALPHABET_SIZE,
)
import multiprocessing as mp
from m1_compression.enumerative_coder_simple import SimpleAdaptiveRankCodec
from m1_compression.batched_arithmetic_coder import BatchedArithmeticEncoder
from m1_compression.hybrid_arithmetic_coder import HybridArithmeticEncoder
from m1_compression.compressor import (
    load_m1_model_and_tokenizer,
    ALPHABET_SIZE,
    ARITHMETIC_CODER_BASE,
    ARITHMETIC_CODER_PRECISION,
)
from offline_entropy_window_split import unpack_windows

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

def pseudo_to_packed_bytes(lst: list[int]) -> bytes:
    out = bytearray()
    acc = bits = 0
    for v in lst:
        acc |= (v & 0x1FF) << bits
        bits += 9
        while bits >= 8:
            out.append(acc & 0xFF)
            acc >>= 8
            bits -= 8
    if bits:                        # flush tail
        out.append(acc)
    return bytes(out)

def packed_bytes_to_pseudo(b: bytes) -> list[int]:
    out, acc, bits = [], 0, 0
    for byte in b:
        acc |= byte << bits
        bits += 8
        while bits >= 9:
            out.append(acc & 0x1FF)
            acc >>= 9
            bits -= 9
    return out

def calculate_compression_ratio(original_bytes: List[bytes], compressed_segments: List[bytes]) -> float:
    if not compressed_segments or len(original_bytes) == 0:
        return 1.0
    
    total_compressed_length = sum(len(compressed_seg) for compressed_seg in compressed_segments)
    ratio = total_compressed_length / sum(len(orig_seg) for orig_seg in original_bytes)
    if ratio > 2.0:
        logger.warning(f"Unusual compression ratio: {ratio:.4f} (compressed larger than original)")
    
    return ratio

def collect_window_size_statistics(segmented_results: List[List[bytes]]) -> Dict[int, int]:
    window_size_counts = Counter()
    
    for segments in segmented_results:
        for segment in segments:
            window_size = len(segment)
            window_size_counts[window_size] += 1
    
    return dict(window_size_counts)

# def pad_batch(batch: List[bytes]):
#     batch_tensors = [torch.tensor(data, dtype=torch.int64) for data in batch]
#     lengths = torch.tensor([len(data) for data in batch], dtype=torch.int64)
#     padded_batch = torch.nn.utils.rnn.pad_sequence(
#         batch_tensors, 
#         batch_first=True, 
#         padding_value=0,
#         padding_side="right"
#     )
#     return padded_batch, lengths

def pad_batch(batch: List[bytes]):
   
    # fix 1: transfer bytes to (list(data))
    batch_tensors = [torch.tensor(list(data), dtype=torch.int64) for data in batch]
    lengths = torch.tensor([len(data) for data in batch], dtype=torch.int64)
    # fix 2: remove torch.nn.utils.rnn.pad_sequence 不支持的 'padding_side' 参数。
    # right padding（对于 batch_first=True）
    padded_batch = torch.nn.utils.rnn.pad_sequence(
        batch_tensors, 
        batch_first=True, 
        padding_value=0
    )
    return padded_batch, lengths

# control long seq with smaller batch
def get_batch_size_for_length(window_len, max_batch_size):
    """
    Determines the batch size for a given window length.
    VERY AGGRESSIVE reduction for long sequences to prevent OOM.
    """
    #  max_batch_size only for short len
    if window_len <= 128:
        return max_batch_size
    if window_len <= 256:
        return max(max_batch_size // 4, 1)
    if window_len <= 512:
        return max(max_batch_size // 16, 1)
    if window_len <= 1024:
        return max(max_batch_size // 64, 1)
    if window_len <= 2048:
        return 2 # 对于 1k-2k 的序列，最多处理 2 个
    # 对于超过 2048 的超长序列，一次只处理 1 个
    return 1

# 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 find_next_batch_range(all_windows, start_idx, max_m1_batch_size):
    M = len(all_windows)
    if start_idx >= M:
        return start_idx, start_idx

    first_window_len = len(all_windows[start_idx])
    base_batch_size = get_batch_size_for_length(first_window_len, max_m1_batch_size)

    low = start_idx
    high = min(start_idx + base_batch_size, M)
    high_batch_size = get_batch_size_for_length(len(all_windows[high - 1]), max_m1_batch_size)
    if high_batch_size == base_batch_size:
        return start_idx, high

    search_low = low
    search_high = high
    while search_low < search_high:
        mid = search_low + (search_high - search_low) // 2
        mid_window_len = len(all_windows[mid])
        if get_batch_size_for_length(mid_window_len, max_m1_batch_size) == base_batch_size:
            # This window is valid. The partition point must be to the right of it.
            # So, we continue searching in the range [mid + 1, high).
            search_low = mid + 1
        else:
            # This window is NOT valid. It might be the partition point itself,
            # or the point is to its left.
            # So, we continue searching in the range [low, mid).
            search_high = mid
    end_idx = search_low
    if end_idx == start_idx:
        return start_idx, start_idx + 1
    else:
        return start_idx, end_idx

class JsonlShardedDataset(Dataset):
    def __init__(
        self,
        file_path: str,
        current_proc_rank: int = 0,
        total_procs: int = 1,
    ) -> None:

        assert 0 <= current_proc_rank < total_procs, "rank must be in [0, world_size)"
        self.current_proc_rank = current_proc_rank
        self.total_procs = total_procs

        # -- load the whole file once (fast for < few-GB files) -------------
        with open(file_path, "r", encoding="utf-8") as f:
            full_data: List[Dict[str, Any]] = [json.loads(line) for line in f]

        # -- pick the slice that belongs to *this* process ------------------
        total = len(full_data)
        per_proc = math.ceil(total / total_procs)
        start = current_proc_rank * per_proc
        end = min(start + per_proc, total)
        self.data = full_data[start:end]

    def __len__(self) -> int:
        return len(self.data)

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        return self.data[idx]

class InterleavedJsonlDataset(IterableDataset):
    """
    An iterable-style dataset for reading a large JSONL file using an
    interleaving/striding pattern, without yielding state information.

    This is designed for multi-process data loading. Each process reads the
    entire file but only processes lines that match its rank (offset).
    For `N` total processes (world_size), process `r` (rank) will read
    lines r, r+N, r+2N, ... (0-indexed).

    This method ensures an even distribution of lines across processes.

    Args:
        file_path (str): Path to the JSONL file.
        rank (int): The rank of the current process, used as the offset.
        world_size (int): The total number of processes, used as the block_size/stride.
    """
    def __init__(
        self,
        file_path: str,
        rank: int,
        world_size: int,
    ) -> None:
        super().__init__()
        
        if not (0 <= rank < world_size):
            raise ValueError(f"Rank must be in [0, {world_size-1}], but got {rank}")

        self.file_path = file_path
        self.offset = rank
        self.block_size = world_size

    def __iter__(self) -> Iterator[Dict[str, Any]]:
        """
        The iterator method that yields the parsed JSON data for the assigned lines.
        """
        try:
            with open(self.file_path, "r", encoding="utf-8") as f:
                # We use a simple line counter to determine which lines to process.
                # The line_number is 0-indexed.
                for line_number, line in enumerate(f):
                    # Check if the current line number belongs to this process
                    if (line_number % self.block_size) == self.offset:
                        try:
                            # Yield the parsed JSON object
                            yield json.loads(line)
                        except json.JSONDecodeError:
                            # This line is malformed. We can either raise an error
                            # or, more robustly, just print a warning and skip it.
                            print(f"Warning: Rank {self.offset} could not decode JSON on line ~{line_number+1}. Skipping.")
                            continue
        except Exception as e:
            print(f"Error in worker {self.offset}: {e}")
            raise

def batched_m1_compress_predict_fn(model):
    def predict_fn(input_tensor: torch.Tensor, **kwargs) -> torch.Tensor:
        if input_tensor.dim() == 1:
            input_tensor = input_tensor.unsqueeze(0)
        with torch.no_grad():
            logits = model(input_tensor, **kwargs)
            logits = logits[..., :256]
            logits = logits.float()
            assert torch.isfinite(logits).all(), "Logits contain NaN or Inf values."
            probs = torch.softmax(logits, dim=-1)
        return probs
    
    return predict_fn

class CachingCompressorWrapper:
    def __init__(
        self,
        base_compression_fn: Callable, # add cache on base compressor, make module seperate
        cache_size: int = 819200, # default a big cache size
        cache_policy: str = 'fifo' # default fifo
    ):
        if cache_policy not in ['fifo']:
            raise ValueError(f"no caching policy: {cache_policy}.")
        self.base_compression_fn = base_compression_fn
        self.cache_size = cache_size
        self.cache_policy = cache_policy
        # self.cache 存储: raw_bytes -> compressed_pseudo_bytes (List[int])
        self.cache: Dict[bytes, List[int]] = {}
        self.fifo_queue: deque[bytes] = deque()
        logger.info(f"Create CachingCompressorWrapper '{self.base_compression_fn.__name__}'，"
                    f"Cache size: {self.cache_size}, policy: {self.cache_policy}")
         
    def compress(
        self,
        sorted_segments: List[bytes],
        *args, **kwargs
    ) -> List[List[int]]:
        """
        compressors with cache
        """
        if not sorted_segments:
            return []
        M = len(sorted_segments)
        # 1. unique data and indxes
        segment_to_indices = defaultdict(list)
        for i, seg in enumerate(sorted_segments):
            segment_to_indices[seg].append(i)
        unique_segments = list(segment_to_indices.keys())
        # 2. check in cache or not
        misses_data = []
        results_for_uniques: Dict[bytes, List[int]] = {}

        ## for each unique segment, check in cache or not
        for segment in unique_segments:
            if segment in self.cache:
                results_for_uniques[segment] = self.cache[segment]
            else:
                misses_data.append(segment)
        hit_count = len(unique_segments) - len(misses_data)
        logger.info(f"Cache checking: {len(unique_segments)} segments, "
                    f"Get {hit_count}, No caching {len(misses_data)} ")
        # 3. compress non-caching segments
        if misses_data:
            ## keep use original one
            newly_compressed = self.base_compression_fn(
                misses_data, *args, **kwargs
            )
            # 4.refresh cache and fill in result
            for i in range(len(misses_data)):
                raw_segment = misses_data[i]
                compressed_result = newly_compressed[i]
                results_for_uniques[raw_segment] = compressed_result
                # refresh cache
                if self.cache_size > 0 and raw_segment not in self.cache:
                    if len(self.cache) >= self.cache_size:
                        if self.cache_policy == 'fifo':
                            oldest_key = self.fifo_queue.popleft()
                            del self.cache[oldest_key]
                    self.cache[raw_segment] = compressed_result
                    self.fifo_queue.append(raw_segment)
        # 5. rebuild all results
        all_compressed_results = [None] * M
        for seg, indices in segment_to_indices.items():
            result = results_for_uniques[seg]
            for original_index in indices:
                all_compressed_results[original_index] = result
        return all_compressed_results
    
    def __call__(self, *args, **kwargs):
        return self.compress(*args, **kwargs)

def compress_segments_hybrid_arithmetic(
    sorted_segments: List[bytes],
    batched_predict_fn: Callable,
    first_byte_prob: torch.Tensor,
    max_m1_batch_size: int=4096,
    debug: bool = True
) -> List[List[int]]:
    """
    这个函数现在只处理它收到的数据，不需要关心缓存或去重。
    这些逻辑已经被外层的 CachingCompressorWrapper 处理了。
    """
    M = len(sorted_segments)
    if M == 0:
        return []
    # 注意：这里的 sorted_segments 已经是去重后、未命中缓存的数据了。
    logger.info(f"Hybrid AC 核心: 正在处理 {M} 个不重复、未命中缓存的段。")
    segment_to_compressed = {}
    ENCODING_BATCH_SIZE = 128
    encoder = HybridArithmeticEncoder(
        batched_predict_fn=batched_predict_fn,
        first_byte_prob=first_byte_prob
    )
    all_compressed_results = []
    for i in range(0, M, ENCODING_BATCH_SIZE):
        batch_start = i
        batch_end = min(i + ENCODING_BATCH_SIZE, M)
        batch_segments = sorted_segments[batch_start:batch_end]
        try:
            codes = encoder.batched_encode(batch_segments, return_num_padded_bits=False)
            # 对比压缩效果
            for seg, code in zip(batch_segments, codes):
                if len(code) < len(seg):
                    all_compressed_results.append(list(code))
                else:
                    all_compressed_results.append(list(seg)) # 压缩效果不好，用原始数据
        except Exception as e:
            logger.warning(f"Hybrid AC 核心: 批次 {batch_start}-{batch_end} 编码失败: {e}. 该批次使用原始字节。")
            for seg in batch_segments:
                all_compressed_results.append(list(seg))
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
    return all_compressed_results

def prepare_segments(batch: List[Dict[str, Any]])->Dict[str,Any]:
    """
    remove the unpack,sort and rerank methods from segment_pre..
    address unpack and judge the compressiable simultaneously
    """
    all_segments = []
    is_compressible_indicator = []
    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 segment,indicator in byte_windows:
            all_segments.append(segment)
            is_compressible = (indicator == 1 and len(segment) > 3)
            is_compressible_indicator.append(is_compressible)

            # record mapping for sample--segment
            sample_idx_to_list_segment_idx[sample_idx].append(segment_idx)
            segment_idx += 1
    
    effective_segments = {} # {index: window}
    raw_segments_map = {}   # {index: window}
    for i, (segment, is_comp) in enumerate(zip(all_segments, is_compressible_indicator)):
        if is_comp:
            effective_segments[i] = segment
        else:
            raw_segments_map[i] = segment

    
    # rerank by length, reduce padding
    sorted_indices_to_compress = sorted(
        effective_segments.keys(),
        key=lambda idx: len(effective_segments[idx])
    )
    sorted_segments_to_compress = [effective_segments[idx] for idx in sorted_indices_to_compress]

    # create reconstruct information --  
    # 1. mapping sample and segment idx -- in one time unpack
    # 2. mapping old and new idx
    sorted_to_original_idx_map = {new_idx: old_idx for new_idx, old_idx in enumerate(sorted_indices_to_compress)}
    
    reconstruction_info = {
        "sample_idx_to_list_segment_idx": sample_idx_to_list_segment_idx,
        "sorted_to_original_idx_map": sorted_to_original_idx_map,
        "raw_segments_map": raw_segments_map,
        "total_segments": len(all_segments),
        "batch_meta": batch, # meta data
        "effective_segments_map": effective_segments
    }

    return {
        "sorted_segments_to_compress": sorted_segments_to_compress,
        "reconstruction_info": reconstruction_info,
    }

# wrap it as the first processing function
def simple_rle_topk_compression(
    batch: List[bytes],
    predict_fn: Callable,
    first_byte_prob: torch.Tensor,
    max_m1_batch_size: int = 4096,
    debug: bool = True,
):
    """use language model to compress, return compressed bytes and padded bits
    
    Args:
        sliding_windows: List of byte sequences to compress
        predict_fn: Function that predicts next token probabilities
        return_num_padded_bits: Whether to return number of padded bits
        profile: Whether to print timing information for each major step
    """
    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)"

    # refactored batch output window AC:
    #### 1. pad the current batch
    batched_windows_np = [np.frombuffer(bytes(data), dtype=np.uint8) for data in batch]

    M = len(batched_windows_np)

    batched_repeat_probs = []
    batched_ranks = []
    batched_lengths = []
    if debug:
        batched_sorted_indices = []

    start_idx = 0
    while start_idx < M:
        # 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)
        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()

        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)
        ######## Use BatchArithmeticEncoder to replace address one by one ###########
        # we calculate two quantiles from prompt_probs
        # 1. the probability of the next byte
        # 2. the byte ids of the topk next bytes
        next_token_probs = torch.gather(
            prompt_probs,
            dim=-1,
            index=padded_batched_windows.unsqueeze(-1)
        ).squeeze(-1) # [B, L]
        sorted_indices = torch.argsort(prompt_probs, dim=-1, descending=True)
        rank_bitvector = padded_batched_windows.unsqueeze(-1) == sorted_indices
        ranks = torch.argmax(rank_bitvector.float(), dim=-1) # [B, L]
        start_idx = end_idx
        batched_repeat_probs.extend(next_token_probs.cpu().numpy().tolist())
        batched_ranks.extend(ranks.cpu().numpy().tolist())
        batched_lengths.extend(lengths.cpu().numpy().tolist())
        if debug:
            batched_sorted_indices.extend(sorted_indices.cpu().numpy().tolist())

    if debug:
        return batched_repeat_probs, batched_ranks, batched_lengths, batched_sorted_indices
    else:
        return batched_repeat_probs, batched_ranks, batched_lengths

def compress_segments_rank_based(
    sorted_segments: List[bytes],
    batched_predict_fn: Callable,
    first_byte_prob: torch.Tensor,
    max_m1_batch_size: int=4096,
    debug: bool = True
) -> List[List[int]]:
    """
     (SimpleAdaptiveRankCodec)。
    decompress GPU probs and CPU compression。
    """
    # --- GPU Stage : acquire probs and rank ---
    # use original simple_rle_topk_compression
    try:

        gpu_result = simple_rle_topk_compression(
            sorted_segments, 
            batched_predict_fn,
            first_byte_prob, 
            max_m1_batch_size=max_m1_batch_size,
            debug=debug,
        )
        if debug:
            batched_repeat_probs, batched_ranks, batched_lengths, batched_sorted_indices = gpu_result
        else:
            batched_repeat_probs, batched_ranks, batched_lengths = gpu_result
            batched_sorted_indices = None
    

        # --- CPU Stage: encoding one by one  ---

        if len(batched_lengths) != len(sorted_segments):
            logger.error(f"FATAL: Length mismatch after GPU stage. Expected {len(sorted_segments)}, got {len(batched_lengths)}. Falling back to raw data.")
            # 如果长度不匹配，说明上游出错了，直接返回原始数据
            return [list(seg) for seg in sorted_segments]

        M = len(batched_lengths)
        batched_compressed_bytes = [] 

        for i in range(M):
            lengths = batched_lengths[i]
            window_bytes = sorted_segments[i]
            repeat_probs = batched_repeat_probs[i][:lengths]
            ranks = batched_ranks[i][:lengths]

            codec = SimpleAdaptiveRankCodec(top_k=4)
            encoding = codec.encode_window(list(window_bytes), repeat_probs, ranks)
            compressed_bytes = codec.encoding_to_pseudo_bytes(encoding)

            #Add: compare compress result
            if len(compressed_bytes) >= len(window_bytes):
                # use raw bytes replace
                batched_compressed_bytes.append(list(window_bytes))
            else:
                # compress successfully
                batched_compressed_bytes.append(compressed_bytes)

            if debug:
                # make sure batched_compressed_bytes[i] is a list
                if batched_sorted_indices is None or batched_sorted_indices[i] is None:
                    logger.warning(f"Debug mode is on but sorted_indices for segment {i} is None. Skipping decode check.")
                    continue
                
                #Add: valid encode and decode rather than  pseudo_bytes -> encoding
                
                sorted_indices = batched_sorted_indices[i][:lengths]
                decoded = codec.decode_window(encoding, lengths, sorted_indices)
                assert bytes(decoded) == window_bytes, "decoded does not match window_bytes: \n{} and \n{}".format(decoded, window_bytes)
                if i < 10:
                    logger.info(f"Example input window bytes: {window_bytes}")
                    logger.info(f"Example encoding          : {encoding}")
                    logger.info(f"Example compressed bytes  : {compressed_bytes}")
            # batched_compressed_bytes.append(compressed_bytes) 这里重复会导致keyerror

        return batched_compressed_bytes

    except Exception as e:
        logger.error(f"Unhandled exception in compress_segments_rank_based: {e}. Falling back to raw data for the entire batch.", exc_info=True)
        # if any error back to original bytes
        return [list(seg) for seg in sorted_segments]


def compress_segments_arithmetic(
    sorted_segments: List[bytes],
    batched_predict_fn: Callable,
    first_byte_prob: torch.Tensor,
    max_m1_batch_size: int = 4096, 
    debug: bool = True
) -> List[List[int]]:
    """
    Final robust version for arithmetic compression.
    This version is inspired by successful production code and is designed to be stable.
    It compresses unique segments in small, manageable batches and handles failures gracefully.
    """
    device = first_byte_prob.device
    M = len(sorted_segments)
    
    if M == 0:
        return []

    # --- 1. unique  ---
    logger.info(f"Step 1: Identifying unique segments to compress.")
    
    # 创建从原始段到其所有出现位置的映射--这里是为unique准备
    # mapping original <-> position
    segment_to_indices = defaultdict(list)
    for i, seg in enumerate(sorted_segments):
        segment_to_indices[seg].append(i)
        
    # only valid segment
    unique_segments = [seg for seg in segment_to_indices.keys() if len(seg) > 2]
    logger.info(f"Found {len(unique_segments)} unique segments (len>2) out of {M} total segments.")

    # store only compressed result
    segment_to_compressed = {}
    
    # --- 2. safe encoding  ---
    ENCODING_BATCH_SIZE = 128 
    encoder = BatchedArithmeticEncoder(base=ARITHMETIC_CODER_BASE, precision=ARITHMETIC_CODER_PRECISION)

    logger.info(f"Step 2: Encoding unique segments in batches of size {ENCODING_BATCH_SIZE}.")
    
    for i in range(0, len(unique_segments), ENCODING_BATCH_SIZE):
        batch_start = i
        batch_end = min(i + ENCODING_BATCH_SIZE, len(unique_segments))
        batch_unique_segments = unique_segments[batch_start:batch_end]
        
        # for each small batch
        try:
            # prepare for bytes--padding
            batch_padded_segments, batch_lengths = pad_batch(batch_unique_segments)
            batch_padded_segments = batch_padded_segments.to(device)
            batch_lengths = batch_lengths.to(device)

            # batch_predict_fn
            with torch.no_grad():
                # 净化输入，防止模型侧的错误
                safe_padded_segments = batch_padded_segments.clamp(0, ALPHABET_SIZE - 1)
                probs = batched_predict_fn(safe_padded_segments)
            
            # # NOTE : normalize probs to avoid NaN/Inf
            # if not torch.isfinite(probs).all():
            #     logger.warning(f"NaN/Inf detected in model probabilities for batch {i//ENCODING_BATCH_SIZE}. Clamping.")
            #     probs = torch.nan_to_num(probs, nan=1e-9, posinf=1.0, neginf=1e-9)
            #     probs = probs / probs.sum(dim=-1, keepdim=True)

            final_probs = torch.cat([first_byte_prob.expand(probs.shape[0], -1, -1), probs[:, :-1, ...]], dim=1)
            normalized_probs = utils.batched_normalize_pdf_for_arithmetic_coding(final_probs)
            
            if not torch.isfinite(normalized_probs).all():
                raise ValueError("NaN or Inf in normalized probabilities after normalization.")

            # batch_encode
            codes, _ = encoder.batched_encode(
                normalized_probs,
                batch_padded_segments,
                lengths=batch_lengths,
                return_num_padded_bits=True
            )
            
            # store result
            for seg, code in zip(batch_unique_segments, codes):
                segment_to_compressed[seg] = list(code)

        except Exception as e:
            # fail，fallback to original data
            logger.warning(f"Batch encoding failed for unique segments {batch_start}-{batch_end}: {e}. Using raw bytes for this batch.")
            for seg in batch_unique_segments:
                segment_to_compressed[seg] = list(seg)
            
            # clean step
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

    # --- 3. reconstruct all result ---
    logger.info("Step 3: Reconstructing final list from unique compressed segments.")
    all_compressed_results = [None] * M
    
    for seg, indices in segment_to_indices.items():
        if len(seg) <= 2:
            # short segments
            result = list(seg)
        else:
            # from mapping to get results
            compressed_data = segment_to_compressed.get(seg, list(seg)) # not find,use original data
            if len(compressed_data) >= len(seg):
                result = list(seg) # back to origin
            else:
                result = compressed_data
        
        # fill back original data
        for original_index in indices:
            all_compressed_results[original_index] = result
            
    return all_compressed_results



def compress_segments_hybrid_arithmetic(
    sorted_segments: List[bytes],
    batched_predict_fn: Callable,
    first_byte_prob: torch.Tensor,
    max_m1_batch_size: int=4096,
    debug: bool = True
) -> List[List[int]]:
    """
    GPU and CPU hybrid version for arithmetic compression.
    """
    M = len(sorted_segments)
    if M == 0:
        return []
    logger.info("Step 1: Identifying unique segments to compress.")
    
    device = first_byte_prob.device
    segment_to_indices = defaultdict(list)
    for i, seg in enumerate(sorted_segments):
        segment_to_indices[seg].append(i)
    
    unique_segments = [seg for seg in segment_to_indices.keys() if len(seg) > 2]
    logger.info(f"Found {len(unique_segments)} unique segments (len>2) out of {M} total segments.")
    
    # store only compressed result
    segment_to_compressed = {}

    # --- 2. safe encoding  ---
    ENCODING_BATCH_SIZE = 128

    encoder = HybridArithmeticEncoder(
        batched_predict_fn=batched_predict_fn,
        first_byte_prob=first_byte_prob
    )
    logger.info(f"Step 2: Encoding unique segments in batches of size {ENCODING_BATCH_SIZE}.")
    
    for i in range(0, len(unique_segments), ENCODING_BATCH_SIZE):
        batch_start = i
        batch_end = min(i + ENCODING_BATCH_SIZE, len(unique_segments))
        batch_unique_segments = unique_segments[batch_start:batch_end]
        try:
            if debug:
                # in debug pattetn, get padded_bits for validation
                codes, padded_bits = encoder.batched_encode(batch_unique_segments, return_num_padded_bits=True)
                decoded_tensor = encoder.batched_decode(codes, padded_bits, batch_unique_segments)
                for j, original_seg_bytes in enumerate(batch_unique_segments):
                    original_len = len(original_seg_bytes)
                    decoded_bytes = bytes(decoded_tensor[j, :original_len].cpu().tolist())
                    assert decoded_bytes == original_seg_bytes, f"Hybrid decode mismatch for segment!"
            else:
                codes = encoder.batched_encode(batch_unique_segments, return_num_padded_bits=False)
            
            # store results
            for seg, code in zip(batch_unique_segments, codes):
                segment_to_compressed[seg] = list(code)
        
        except Exception as e:
            logger.warning(f"Batch encoding failed for unique segments {batch_start}-{batch_end}: {e}. Using raw bytes for this batch.")
            for seg in batch_unique_segments:
                segment_to_compressed[seg] = list(seg)
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
    
    logger.info("Step 3: Reconstructing final list from unique compressed segments.")
    all_compressed_results = [None] * M
    
    for seg, indices in segment_to_indices.items():
        if len(seg) <= 2:
            result = list(seg)
        else:
            compressed_data = segment_to_compressed.get(seg, list(seg))
            if len(compressed_data) >= len(seg):
                result = list(seg)
            else:
                result = compressed_data
        
        for original_index in indices:
            all_compressed_results[original_index] = result
    
    return all_compressed_results

    
def reconstruct_results(
    compressed_map: Dict[int, List[int]],
    reconstruction_info: Dict[str, Any],
    debug: bool = True
) -> List[Dict[str, Any]]:
    """
    Reconstruct the original results from the compressed results.
    Need and compressed ratio and assert the reconstruction is correct.
    """
    sample_idx_to_list_segment_idx = reconstruction_info["sample_idx_to_list_segment_idx"]
    raw_segments_map = reconstruction_info["raw_segments_map"]
    batch_meta = reconstruction_info["batch_meta"]

    #Add: valid reconstruction
    sorted_to_original_idx_map = reconstruction_info["sorted_to_original_idx_map"]
    # mapping original_idx -> compressed_data 
    original_idx_to_compressed_data = {
        v: compressed_map[k] 
        for k, v in sorted_to_original_idx_map.items() 
        if k in compressed_map
    }

    write_results = []
    ac_key = "m1_enumerative"

    # compute compress ratio
    total_original_bytes = 0
    total_compressed_pseudo_bytes = 0

    for sample_idx,item in enumerate(batch_meta):
        final_pseudo_bytes = []
        if debug:
            reconstructed_original_segments = []

        segment_indices_for_sample = sample_idx_to_list_segment_idx.get(sample_idx, [])
        for original_idx in segment_indices_for_sample:
            if original_idx in original_idx_to_compressed_data:
                # compressed bytes
                compressed_data = original_idx_to_compressed_data[original_idx]
                final_pseudo_bytes.extend(compressed_data)
                if debug:
                    #Add: assert encode and decode
                    total_compressed_pseudo_bytes += len(compressed_data)
                    original_segment_bytes = reconstruction_info["effective_segments_map"][original_idx]
                    reconstructed_original_segments.append(original_segment_bytes)
                    total_original_bytes += len(original_segment_bytes)

            elif original_idx in raw_segments_map:
                # raw bytes
                raw_data = raw_segments_map[original_idx]
                final_pseudo_bytes.extend(list(raw_data))
                
                if debug:
                    total_compressed_pseudo_bytes += len(raw_data)
                    reconstructed_original_segments.append(raw_data)
                    total_original_bytes += len(raw_data)
            
            else:
                # Case 3: wrong both not exist
                logger.error(f"FATAL LOGIC ERROR: Segment with original_idx {original_idx} does not exist in effective_segments_map or raw_segments_map!")
                # try from effective_segments_map 
                original_segment_bytes = reconstruction_info["effective_segments_map"].get(original_idx)
                if original_segment_bytes:
                    final_pseudo_bytes.extend(list(original_segment_bytes))


        packed_bytes = pseudo_to_packed_bytes(final_pseudo_bytes)
        result = {
            **item,
            "m1_compressed_data": base64.b64encode(packed_bytes).decode("ascii")
        }
        write_results.append(result)
        #Add: assert encode and decode
        if debug and reconstructed_original_segments:
            original_sample_bytes = item["text"].encode('utf-8')
            reconstructed_sample_bytes = b"".join(reconstructed_original_segments)
            assert reconstructed_sample_bytes == original_sample_bytes, \
                f"Sample {sample_idx} reconstruction failed!"
            # check pack and unpack
            unpacked_pseudo_bytes = packed_bytes_to_pseudo(packed_bytes)
            assert unpacked_pseudo_bytes == final_pseudo_bytes, \
                f"Pseudo-bytes packing/unpacking round-trip failed for sample {sample_idx}"

        #Add: copmute compress ratio
        if debug and total_original_bytes > 0:
            compression_ratio = total_compressed_pseudo_bytes / total_original_bytes
            logger.info(f"Batch compression stats: "
                        f"Original bytes: {total_original_bytes}, "
                        f"Compressed pseudo-bytes: {total_compressed_pseudo_bytes}, "
                        f"Ratio: {compression_ratio:.4f}")

    # this tab wrong
    return write_results


def writer_consumer(write_queue, output_file, buffer_size=100,debug=True):
    write_buf = []
    try:
        with open(output_file, 'w', encoding='utf-8') as f:
            while True:
                payload = write_queue.get()
                if payload is None:
                    break
                # get result from reconstruct results
                write_results = reconstruct_results(
                    payload["compressed_map"],
                    payload["reconstruction_info"],
                    debug=debug
                )
                write_buf.extend(write_results)
                # clean the complex expression of before segmentation
                
                # 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 = []

            # 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_processor_fn(
    batch: List[Dict[str, Any]],
    compression_fn: Callable, # <-- 传入一个压缩函数作为参数！
    predict_fn: Callable,
    first_byte_prob: torch.Tensor,
    max_m1_batch_size: int,
    debug: bool = True
):
    # 1. preparing data
    prep_data = prepare_segments(batch)
    sorted_segments = prep_data["sorted_segments_to_compress"]
    reconstruction_info = prep_data["reconstruction_info"]
    # 2. compress data
    if sorted_segments:
        #Add: Time consume
        start_time = time.time()
        compressed_pseudo_bytes = compression_fn(
            sorted_segments,
            predict_fn,
            first_byte_prob,
            max_m1_batch_size,
            debug
        )
        end_time = time.time()
        duration = end_time - start_time
        logger.info(
            f"Compressed {len(sorted_segments)} segments "
            f"in {duration:.4f} seconds ({len(sorted_segments)/duration if duration > 0 else float('inf'):.2f} segments/sec)."
        )
        # create a mapping from origal idx to compressed results
        #sorted_to_original_idx_map = reconstruction_info["sorted_to_original_idx_map"]
        
        compressed_map = {
            i: data 
            for i, data in enumerate(compressed_pseudo_bytes)
        }
        # compressed_map： key--sorted i; value--compressed bytes
        # but below sorted_to_original_idx_map[i] is original idx..
        # compressed_map = {
        #     sorted_to_original_idx_map[i]: data 
        #     for i, data in enumerate(compressed_pseudo_bytes)
        # }
    else:
        compressed_map = {}
    # 3. pack result to consumer
    payload = {
        "compressed_map": compressed_map,
        "reconstruction_info": reconstruction_info
    }
    return payload

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('--compressor', type=str, default='rank_based', 
                      choices=['rank_based', 'arithmetic','hybrid_arithmetic'],
                      help='Choose the compression algorithm.')
    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('--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=True,
                      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)')

    # adding cache parameters
    parser.add_argument('--use_global_cache', action='store_true', default=True,
                      help='Enable the global compression cache.')
    parser.add_argument('--cache_size', type=int, default=819200,
                      help='Size of the global compression cache.')
    
    
    args = parser.parse_args()

    # choose compression algorithm 
    if args.compressor == 'rank_based':
        compression_algorithm = compress_segments_rank_based
    elif args.compressor == 'arithmetic':
        compression_algorithm = compress_segments_arithmetic
    elif args.compressor == 'hybrid_arithmetic':
        compression_algorithm = compress_segments_hybrid_arithmetic
    else:
        raise ValueError(f"Unknown compressor: {args.compressor}")
    logger.info(f"Using compression algorithm: {compression_algorithm.__name__}")

    
    # use wrapper to make cache for each algorithm
    if args.use_global_cache:
        caching_wrapper = CachingCompressorWrapper(
            base_compression_fn=compression_algorithm,
            cache_size=args.cache_size
        )
        # use cache
        compression_algorithm_to_use = caching_wrapper
        logger.info("Global cache start....")
    else:
        # no cache
        compression_algorithm_to_use = compression_algorithm
        logger.info("No Global cache ...")
    
    
    mp.set_start_method('spawn', force=True)
    gc_freq = 100
    dump_freq = 25

    # 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...")

    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)
        )
        
        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):
            payload_for_writer = main_processor_fn(
                batch,
                compression_algorithm_to_use, # compressor with cache
                # compression_algorithm, # 把选择的算法传进去
                batched_predict_fn,
                first_byte_prob,
                args.max_compression_batch_size,
                args.debug,
            )
            logger.info(f"Processed batch {batch_idx}")
            write_queue.put(payload_for_writer)

            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()