import torch import torch.nn.functional as F from torch.utils.data import DataLoader import json import numpy as np from pathlib import Path from typing import List, Dict, Any, Callable, Tuple, Optional import logging import argparse import re import gc from collections import defaultdict, Counter from m1_compression.compressor import ( load_m1_model_and_tokenizer, ALPHABET_SIZE, ) import multiprocessing as mp from offline_utils import ( compress_windows_starts_lens, decompress_windows_starts_lens, unpack_windows, InterleavedJsonlDataset, batched_m1_compress_predict_fn, find_next_batch_range, ) logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') logger = logging.getLogger() MAX_LINE_LEN = 512 def print_windows(text: str, starts: list[int], lens: list[int], sample_idx: int = None, ): from rich.console import Console from rich.text import Text import io PALETTE = ( "#c6f6d5", "#bee3f8", "#fbb6ce", "#faf089", "#fed7e2", "#b2f5ea", ) string_io = io.StringIO() console = Console(record=True, force_terminal=True, color_system="truecolor", file=string_io) t = Text() last_end = 0 colour_idx = 0 for s, l in sorted(zip(starts, lens)): t.append(text[last_end:s]) if s == last_end: colour_idx = (colour_idx + 1) % len(PALETTE) t.append(text[s:s + l], style=f"on {PALETTE[colour_idx]} bold black") last_end = s + l t.append(text[last_end:]) console.print(t) # only save the last 100 samples save_idx = sample_idx % 100 return console.save_svg(f"window_visualize_{save_idx}.svg") def collect_lines(batched_bytes_data: List[bytes], max_len: int = 2048) -> Tuple[List[bytes], Dict[int, Tuple[int, int]]]: batched_lines = [] line_id_to_sample_offsets = {} line_idx = 0 for sample_idx, data_bytes in enumerate(batched_bytes_data): if len(data_bytes) == 0: continue # Find all lines with their consecutive newlines attached (handles \r\n, \r, \n) lines_with_positions = [] for match in re.finditer(b'[^\r\n]*(?:\r\n|\r|\n)*', data_bytes): if match.group(): # Skip empty matches lines_with_positions.append((match.group(), match.start())) for line, byte_offset in lines_with_positions: if len(line) > max_len: logger.info("Line too long with {} bytes, splitting into chunks...".format(len(line))) # Split long line into chunks of max_len for chunk_start in range(0, len(line), max_len): chunk_end = min(chunk_start + max_len, len(line)) batched_lines.append(line[chunk_start:chunk_end]) # Calculate the absolute byte offset for this chunk chunk_byte_offset = byte_offset + chunk_start line_id_to_sample_offsets[line_idx] = (sample_idx, chunk_byte_offset) line_idx += 1 else: batched_lines.append(line) line_id_to_sample_offsets[line_idx] = (sample_idx, byte_offset) line_idx += 1 return batched_lines, line_id_to_sample_offsets def calculate_skew(entropy: torch.Tensor) -> torch.Tensor: mean = torch.mean(entropy) diffs = entropy - mean var = torch.mean(torch.pow(diffs, 2.0)) std = torch.pow(var, 0.5) if std == 0.0: return torch.tensor(0.0) zscores = diffs / std skews = torch.mean(torch.pow(zscores, 3.0)) return skews def get_split_points( probs: torch.Tensor, next_bytes: torch.Tensor, lengths: torch.Tensor, base_global_quantile: float, base_monotonic_quantile: float, debug: bool = False, ): B, L = probs.shape[0], probs.shape[1] arange_ids = torch.arange(L, device=probs.device).unsqueeze(0) pad_mask = arange_ids < lengths.unsqueeze(1) padded_cross_entropy = F.cross_entropy( probs.transpose(1, 2), next_bytes, reduction="none" ) flattened_cross_entropy = padded_cross_entropy[pad_mask] assert flattened_cross_entropy.dim() == 1 skew_flattened_cross_entropy = calculate_skew(flattened_cross_entropy.float()) if skew_flattened_cross_entropy > 0.0: base_global_quantile = base_global_quantile - 0.04 * skew_flattened_cross_entropy.item() base_global_quantile = min(max(base_global_quantile, 0.0), 1.0) # entropy is a tensor of shape (B * L_b) threshold = torch.quantile(flattened_cross_entropy, base_global_quantile).clamp(0.1, 10.0) padded_cross_entropy_diff = torch.diff(padded_cross_entropy, dim=1) padded_cross_entropy_diff = torch.cat( [ torch.zeros(B, 1, device=padded_cross_entropy_diff.device), padded_cross_entropy_diff ], dim=1 ) flattened_cross_entropy_diff = padded_cross_entropy_diff[pad_mask] skew_flattened_cross_entropy_diff = calculate_skew(flattened_cross_entropy_diff.float()) if skew_flattened_cross_entropy_diff > 0.0: base_monotonic_quantile = base_monotonic_quantile - 0.04 * skew_flattened_cross_entropy_diff.item() base_monotonic_quantile = min(max(base_monotonic_quantile, 0.0), 1.0) diff_threshold = torch.quantile(flattened_cross_entropy_diff, base_monotonic_quantile).clamp(0.01, 10.0) split_points_mask = ((padded_cross_entropy > threshold) | (padded_cross_entropy_diff > diff_threshold)) & pad_mask if debug: logger.info(f"skew_flattened_cross_entropy: {skew_flattened_cross_entropy}") logger.info(f"skew_flattened_cross_entropy_diff: {skew_flattened_cross_entropy_diff}") logger.info(f"base_global_quantile: {base_global_quantile}") logger.info(f"base_monotonic_quantile: {base_monotonic_quantile}") logger.info(f"threshold: {threshold}") logger.info(f"diff_threshold: {diff_threshold}") return split_points_mask def get_batch_size_for_length(window_len, max_batch_size): """Determines the batch size for a given window length.""" BATCH_SIZE_TIERS = { 512: max_batch_size, 1024: max(max_batch_size // 2, 1), 2048: max(max_batch_size // 4, 1), } for max_len, batch_size in BATCH_SIZE_TIERS.items(): if window_len <= max_len: return batch_size return 1 def calculate_entropy_and_split_points_fn( batch: Dict[str, Any], # List, [{"text":"Hello world","id":"1"}] predict_fn: Callable, chunk_size: int = 512, base_global_quantile: float = 90.0, base_monotonic_quantile: float = 90.0, unigram_probs: Optional[torch.Tensor] = None, max_m1_batch_size: int = 2048, line_split: bool = False, debug: bool = False, ) -> List[Dict[str, Any]]: batched_bytes_data = [item["text"].encode('utf-8') for item in batch] # List bytes: [bytes, bytes,...] device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if unigram_probs is not None: unigram_probs = unigram_probs.to(device) # 2. batched segmentations all_split_point_masks = [] # 1. preprocess all samples, record each chunk if line_split: chunks, chunk_to_sample_and_offset = collect_lines(batched_bytes_data, max_len=MAX_LINE_LEN) sorted_chunks = sorted(enumerate(chunks), key=lambda x: len(x[1])) sorted_idx, sorted_chunks = zip(*sorted_chunks) sorted_chunks = list(sorted_chunks) # Convert tuple to list chunk_idx_map = { orig_idx: new_idx for new_idx, orig_idx in enumerate(sorted_idx) } chunks_np = [np.frombuffer(bytes(data), dtype=np.uint8) for data in sorted_chunks] num_chunks = len(sorted_chunks) start_idx = 0 while start_idx < num_chunks: # Use the new helper function to find the exact range for the next safe batch start_idx, end_idx = find_next_batch_range(chunks_np, start_idx, max_m1_batch_size, get_batch_size_for_length) batch_chunks_np = chunks_np[start_idx:end_idx] effective_batch_size = end_idx - start_idx lengths_pt = torch.tensor([len(chunk) for chunk in batch_chunks_np], dtype=torch.long, device=device) batch_chunks_pt = torch.zeros( (effective_batch_size, max(lengths_pt)), dtype=torch.long, device=device ) for i, chunk_np in enumerate(batch_chunks_np): batch_chunks_pt[i, :len(chunk_np)] = torch.tensor(chunk_np, dtype=torch.long, device=device) cur_batch = batch_chunks_pt[:effective_batch_size] cur_lengths = lengths_pt[:effective_batch_size] with torch.no_grad(): probs = predict_fn(cur_batch) # add unigram first_prob = unigram_probs.expand( effective_batch_size, 1, -1) final_probs = torch.cat([first_prob, probs[:, :-1, :]], dim=1) start_idx = end_idx # calculate (cross) entropy, # calculate dynamic threshold, # calculate split points split_points_mask = get_split_points( final_probs, cur_batch, cur_lengths, base_global_quantile, base_monotonic_quantile, debug, ) all_split_point_masks.append(split_points_mask) split_point_chunk_idx_lst = [] split_point_position_idx_lst = [] processed_chunks = 0 for mask in all_split_point_masks: split_point_chunk_idx, split_point_position_idx = mask.cpu().nonzero(as_tuple=True) split_point_chunk_idx_lst.append(split_point_chunk_idx + processed_chunks) split_point_position_idx_lst.append(split_point_position_idx) processed_chunks = processed_chunks + mask.shape[0] split_point_chunk_idx = torch.cat(split_point_chunk_idx_lst) split_point_position_idx = torch.cat(split_point_position_idx_lst) else: chunk_idx_map = None # 1. preprocess all samples, record each chunk chunks = [] chunk_to_sample_and_offset = {} chunk_idx = 0 for sample_idx, data_bytes in enumerate(batched_bytes_data): logger.debug(f"Processing sample {sample_idx+1} (bytes: {len(data_bytes)})") if len(data_bytes) == 0: continue byte_len = len(data_bytes) for i in range(0, byte_len, chunk_size): chunk_start = i chunk_end = min(i + chunk_size, byte_len) chunk = data_bytes[chunk_start:chunk_end] chunks.append(chunk) chunk_to_sample_and_offset[chunk_idx] = (sample_idx, chunk_start) # key: chunk_idx 被切分的块 -> (sample_idx, chunk_start)在原来某个样本中，起始位置 chunk_idx += 1 # 2. batched segmentations all_split_point_masks = [] batch_chunks_pt = torch.zeros( (max_m1_batch_size, chunk_size), dtype=torch.long, device=device ) lengths_pt = torch.zeros(max_m1_batch_size, dtype=torch.long, device=device) num_chunks = len(chunks) # batched get all segmentations for start_idx in range(0, num_chunks, max_m1_batch_size): end_idx = min(start_idx + max_m1_batch_size, num_chunks) batch_chunks = chunks[start_idx:end_idx] batch_chunks_np = [np.frombuffer(bytes(data), dtype=np.uint8) for data in batch_chunks] effective_batch_size = end_idx - start_idx # padding for i, chunk_np in enumerate(batch_chunks_np): batch_chunks_pt[i, :len(chunk_np)] = torch.tensor(chunk_np, dtype=torch.long, device=device) lengths_pt[i] = len(chunk_np) cur_batch = batch_chunks_pt[:effective_batch_size] cur_lengths = lengths_pt[:effective_batch_size] with torch.no_grad(): probs = predict_fn(cur_batch) # add unigram first_prob = unigram_probs.expand( effective_batch_size, 1, -1) final_probs = torch.cat([first_prob, probs[:, :-1, :]], dim=1) # calculate (cross) entropy, # calculate dynamic threshold, # calculate split points split_points_mask = get_split_points( final_probs, cur_batch, cur_lengths, base_global_quantile, base_monotonic_quantile, debug, ) all_split_point_masks.append(split_points_mask) all_split_point_masks = torch.cat(all_split_point_masks, dim=0) # （nums_chunks,chunk_size）-> 汇集所有chunks的切分点标记 all_split_points_tuple = all_split_point_masks.nonzero(as_tuple=True) # chunk_idx, position_idx包含两个一维张量，即第几个chunk的对应某个位置有切分点 # `(tensor([0, 0, 2]), tensor([15, 89, 412]))` 表示在第0个chunk的第15和89位置，以及第2个chunk的第412位置有切分点。 split_point_chunk_idx, split_point_position_idx = all_split_points_tuple[0].cpu(), all_split_points_tuple[1].cpu() sample_idx_to_split_positions = defaultdict(list) ## avoid scan each chunk -> dict -> 直接将每个split_point放到对应的chunk_idx中，用处理好的信息完成映射 # 1. transfer pytorch to numpy split_point_chunk_idx_np = split_point_chunk_idx.numpy() split_point_position_idx_np = split_point_position_idx.numpy() chunk_to_splits = defaultdict(list) # 2. scan all the split_point and put it into chunk-dicts 如果先遍历放到字典中时 split_points次 for i in range(len(split_point_chunk_idx_np)): chunk_idx = split_point_chunk_idx_np[i] position = split_point_position_idx_np[i] chunk_to_splits[chunk_idx].append(position) # 3. process all the chunk_to_split for chunk_idx in range(num_chunks): # chunk_idx: original index of chunk chunk = chunks[chunk_idx] sample_idx, chunk_start = chunk_to_sample_and_offset[chunk_idx] if line_split: sorted_chunk_idx = chunk_idx_map[chunk_idx] split_points = chunk_to_splits[sorted_chunk_idx] else: split_points = chunk_to_splits[chunk_idx] if len(split_points) == 0: split_points = [0] if split_points[0] != 0: split_points.insert(0, 0) split_points.append(len(chunk)) offset_split_points = [s + chunk_start for s in split_points] sample_idx_to_split_positions[sample_idx].extend(offset_split_points) # for chunk_idx in range(num_chunks): # 相当于遍历 chunk_size * split_points 次 # chunk = chunks[chunk_idx] # sample_idx, chunk_start = chunk_to_sample_and_offset[chunk_idx] # # nonzero() leads to a GPU-CPU sync point so we defer until all batches finished # split_points = split_point_position_idx[split_point_chunk_idx == chunk_idx].tolist() # if len(split_points) == 0: # split_points = [0] # if split_points[0] != 0: # split_points.insert(0, 0) # split_points.append(len(chunk)) # offset_split_points = [s + chunk_start for s in split_points] # sample_idx_to_split_positions[sample_idx].extend(offset_split_points) # 还原到原始样本的offset sample_idx_to_split_positions = {k: sorted(v) for k, v in sample_idx_to_split_positions.items()} write_results = [] min_window_size = 3 if debug: extreme_compression_results = [] for sample_idx, item in enumerate(batch): split_points = sample_idx_to_split_positions[sample_idx] split_windows_starts = [] split_windows_lens = [] cur_l = 0 cur_r = 0 for i in range(len(split_points) - 1): cur_l = split_points[i] cur_r = split_points[i+1] if cur_r - cur_l >= min_window_size: split_windows_starts.append(cur_l) split_windows_lens.append(cur_r - cur_l) # assert cur_r == len(item["text"].encode('utf-8')), f"last cur_r: {cur_r} != len(item['text']): {len(item['text'])}" compressed_windows_starts_lens_b64 = compress_windows_starts_lens(split_windows_starts, split_windows_lens) result = { **item, "windows_starts_lens_b64": compressed_windows_starts_lens_b64 } if debug: print_windows(item["text"], split_windows_starts, split_windows_lens, sample_idx=sample_idx) _debug_starts_lens = decompress_windows_starts_lens(compressed_windows_starts_lens_b64) _debug_starts, _debug_lens = _debug_starts_lens assert len(_debug_starts) == len(_debug_lens), f"Window starts and lens have different lengths: {len(_debug_starts)} != {len(_debug_lens)}" assert _debug_starts == split_windows_starts, f"Window starts do not match: {_debug_starts} != {split_windows_starts}" assert _debug_lens == split_windows_lens, f"Window lens do not match: {_debug_lens} != {split_windows_lens}" # calculate extreme compression rate: compress all windows into 1 byte debug_sample = item["text"].encode('utf-8') raw_bytes = len(debug_sample) - sum(_debug_lens) compressed_bytes = len(_debug_starts) extreme_compression_rate = (compressed_bytes + raw_bytes) / len(debug_sample) extreme_compression_results.append(extreme_compression_rate) logger.info(f"[Extreme compression rate] for sample idx {sample_idx}: {extreme_compression_rate:.4f}") debug_byte_windows = unpack_windows(debug_sample, compressed_windows_starts_lens_b64) debug_bytes_windows, debug_indicators = zip(*debug_byte_windows) assert b"".join(debug_bytes_windows) == debug_sample, f"Debug bytes windows do not match: {b''.join(debug_bytes_windows)} != {debug_sample}" debug_split_points = sample_idx_to_split_positions[sample_idx] logger.info(f"Original byte length: {len(debug_sample)}") logger.info(f"num split_points: {len(debug_split_points)}") _debug_compressed_windows = [x[0] for x in debug_byte_windows if x[1]] _debug_sorted_compressed_windows = sorted(_debug_compressed_windows, key=lambda x: len(x), reverse=True) _debug_raw_windows = [x[0] for x in debug_byte_windows if not x[1]] _debug_sorted_raw_windows = sorted(_debug_raw_windows, key=lambda x: len(x), reverse=True) for i, byte_window in enumerate(_debug_sorted_compressed_windows): logger.info(f"compressed byte_window[{i}]: {byte_window}") if i > 10: break for i, byte_window in enumerate(_debug_sorted_raw_windows): logger.info(f"raw byte_window[{i}]: {byte_window}") if i > 10: break write_results.append(result) if debug: logger.info(f"[Extreme compression rate] for all samples: {np.mean(extreme_compression_results):.4f}") return write_results def writer_consumer(write_queue, output_file, buffer_size=100): """ 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. """ write_buf = [] try: with open(output_file, 'w', encoding='utf-8') as f: while True: item = write_queue.get() if item is None: break write_buf.extend(item) # 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 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('--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('--apply_line_split', action='store_true', default=False, help='apply_line_split') 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)') args = parser.parse_args() 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 = JsonlShardedDataset( # args.input_file, # current_proc_rank=args.process_id, # total_procs=args.num_processes, # ) 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...") # Create queue and start writer process write_queue = mp.Queue(maxsize=200) writer_process = mp.Process( target=writer_consumer, args=(write_queue, output_file, dump_freq) ) writer_process.start() try: # Process each batch for batch_idx, batch in enumerate(dataloader): split_points_results = calculate_entropy_and_split_points_fn( batch, batched_predict_fn, chunk_size=args.chunk_size, base_global_quantile=args.base_global_quantile, base_monotonic_quantile=args.base_monotonic_quantile, unigram_probs=first_byte_prob, max_m1_batch_size=args.max_entropy_batch_size, line_split=args.apply_line_split, debug=args.debug, ) logger.info(f"Processed batch {batch_idx}") write_queue.put(split_points_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 writer process write_queue.put(None) except Exception as e: logger.error(f"Error during processing: {e}") # Try to terminate writer process cleanly try: write_queue.put(None) except: pass raise finally: # Wait for writer process to finish writer_process.join() if writer_process.exitcode != 0: logger.error(f"Writer process failed with exit code: {writer_process.exitcode}") logger.info(f"Completed processing successfully, output written to {output_file}") if __name__ == "__main__": main()