# based on # https://github.com/MeetKai/functionary/blob/aa3dbdd65f7e388f2386622606bdfeec95c2b863/functionary/train/packing/packed_dataset.py import logging import numpy as np from ctx_to_lora.utils import check_is_iterable, concat_list, setup_logging logger = logging.getLogger() def pack_data_points_by_length( lens: list[list[int]], ctx_lens: list[int], n_queries: list[int], max_packed_inp_len: int, max_packed_ctx_len: int, max_size: int = -1, ) -> tuple[list[int], list[int]]: if not lens: return [] len_arr = np.array([sum(l) for l in lens], dtype=np.long) ctx_len_arr = np.array(ctx_lens, dtype=np.long) n = len(len_arr) assert len(ctx_len_arr) == n, "Length of ctx_len_arr must match length of lens" assert len(n_queries) == n, "Length of n_queries must match length of lens" if n == 1: return ( [0] if len_arr[0] <= max_packed_inp_len and ctx_len_arr[0] <= max_packed_ctx_len else [] ) # Create cumulative sum arrays for efficient range queries cumsum_inp_len = np.cumsum(len_arr) cumsum_ctx_len = np.cumsum(ctx_len_arr) # indices of flattened input_ids cumsum_n_queries = np.cumsum(n_queries) # start_indices_inp = [0] # start_indices_ctx = [0] idx_pairs = [] i = 0 while i < n: # Find the maximum j such that sum(lens[i:j+1]) <= max_packed_inp_len start_sum_inp = cumsum_inp_len[i - 1] if i > 0 else 0 valid_ends_inp = (cumsum_inp_len[i:] - start_sum_inp) <= max_packed_inp_len start_sum_ctx = cumsum_ctx_len[i - 1] if i > 0 else 0 valid_ends_ctx = (cumsum_ctx_len[i:] - start_sum_ctx) <= max_packed_ctx_len valid_ends = valid_ends_inp & valid_ends_ctx if not np.any(valid_ends): # Single item exceeds max_packed_inp_len, skip it logging.debug( f"Skipping item {i} with input length {len_arr[i]} and context length {ctx_len_arr[i]}" ) i += 1 continue # Find the last valid index max_valid_idx = i + np.where(valid_ends)[0][-1] # Apply max_size constraint if max_size > 0: max_valid_idx = min(max_valid_idx, i + max_size - 1) idx_pairs.append((i, max_valid_idx + 1)) i = max_valid_idx + 1 return idx_pairs def pack_data_points_FA( batch: dict[str, any], need_flatten: bool, ) -> dict[str, np.ndarray]: if not batch: raise ValueError("Batch is empty") # Pre-allocate lists with known sizes total_ctx_len = sum(len(x) for x in batch["ctx_ids"]) if not need_flatten: total_inp_len = sum(len(x) for x in batch["input_ids"]) else: total_inp_len = sum(len(y) for x in batch["input_ids"] for y in x) ctx_ids = np.empty(total_ctx_len, dtype=np.long) ctx_position_ids = np.empty(total_ctx_len, dtype=np.long) input_ids = np.empty(total_inp_len, dtype=np.long) position_ids = np.empty(total_inp_len, dtype=np.long) labels = np.empty(total_inp_len, dtype=np.long) offset = ctx_offset = 0 for input_ids_b, labels_b in zip(batch["input_ids"], batch["labels"]): inp_start = offset if need_flatten: # compute position_ids for each sub-list in input_ids_b local_start = inp_start for ids_b in input_ids_b: position_ids[local_start : local_start + len(ids_b)] = np.arange( len(ids_b), dtype=np.long ) local_start += len(ids_b) input_ids_b = concat_list(input_ids_b) labels_b = concat_list(labels_b) inp_len = len(input_ids_b) inp_end = offset + inp_len if not need_flatten: position_ids[inp_start:inp_end] = np.arange(inp_len, dtype=np.long) input_ids[inp_start:inp_end] = input_ids_b labels[inp_start:inp_end] = labels_b offset += inp_len for ctx_ids_b in batch["ctx_ids"]: ctx_len = len(ctx_ids_b) ctx_start, ctx_end = ctx_offset, ctx_offset + ctx_len ctx_ids[ctx_start:ctx_end] = ctx_ids_b ctx_position_ids[ctx_start:ctx_end] = np.arange(ctx_len, dtype=np.long) ctx_offset += ctx_len return { "ctx_ids": ctx_ids, "ctx_position_ids": ctx_position_ids, "input_ids": input_ids, "position_ids": position_ids, "labels": labels, } def pack_batch( batch: dict[str, any], max_packed_inp_len: int, max_packed_ctx_len: int, max_packed_size: int = -1, ) -> dict[str, any]: need_flatten = check_is_iterable(batch["input_ids"][0][0]) if need_flatten: n_queries = [len(x) for x in batch["input_ids"]] inp_lens = [[len(y) for y in x] for x in batch["input_ids"]] # batch = {k: concat_list([x for x in batch[k]]) for k in keys} # for k in ["input_ids", "attention_mask", "labels"]: # batch[k] = concat_list([x for x in batch[k]]) else: n_queries = [1] * len(batch["input_ids"]) inp_lens = [[len(x)] for x in batch["input_ids"]] inp_count = len(inp_lens) # total_original_tokens = sum(inp_lens) if "ctx_ids" not in batch: raise ValueError("Batch must contain 'ctx_ids' and 'labels' keys") # we do not pad so we can just take the length of the tokens ctx_lens = [len(x) for x in batch["ctx_ids"]] # Group indices idx_pairs = pack_data_points_by_length( inp_lens, ctx_lens, n_queries, max_packed_inp_len, max_packed_ctx_len, max_packed_size, ) # Pack groups # n_samples = len(boundaries) - 1 packed_batch = { "ctx_ids": [], "ctx_position_ids": [], "input_ids": [], "position_ids": [], "labels": [], "n_queries": [], } packing_efficiency_ratios = [] ctx_packing_efficiency_ratios = [] for idx_pair in idx_pairs: start_idx, end_idx = idx_pair[0], idx_pair[1] # start_idx = start_indices[i] # end_idx = start_indices[i + 1] group_items = { "ctx_ids": batch["ctx_ids"][start_idx:end_idx], "input_ids": batch["input_ids"][start_idx:end_idx], "labels": batch["labels"][start_idx:end_idx], } packed_item = pack_data_points_FA(group_items, need_flatten) packed_batch["ctx_ids"].append(packed_item["ctx_ids"]) packed_batch["ctx_position_ids"].append(packed_item["ctx_position_ids"]) packed_batch["input_ids"].append(packed_item["input_ids"]) packed_batch["position_ids"].append(packed_item["position_ids"]) packed_batch["labels"].append(packed_item["labels"]) packed_batch["n_queries"].append(n_queries[start_idx:end_idx]) if max_packed_inp_len > 0: inp_efficiency = len(packed_item["input_ids"]) / max_packed_inp_len packing_efficiency_ratios.append(inp_efficiency) if max_packed_ctx_len > 0: ctx_efficiency = len(packed_item["ctx_ids"]) / max_packed_ctx_len ctx_packing_efficiency_ratios.append(ctx_efficiency) # # Calculate total original tokens # total_original_inp_tokens = sum(inp_lens) # total_original_ctx_tokens = sum(ctx_lens) # total_original_tokens = total_original_inp_tokens + total_original_ctx_tokens # # Calculate length statistics packed_inp_lens_arr = np.array([len(x) for x in packed_batch["input_ids"]]) packed_ctx_lens_arr = np.array([len(x) for x in packed_batch["ctx_ids"]]) # # Log performance statistics avg_inp_packing_efficiency = ( np.mean(packing_efficiency_ratios) if packing_efficiency_ratios else 0 ) avg_ctx_packing_efficiency = ( np.mean(ctx_packing_efficiency_ratios) if ctx_packing_efficiency_ratios else 0 ) logging.debug( f"Packing stats - Original samples: {inp_count}\n" f"# Packed samples: {len(idx_pairs)}\n" f"Avg inp packing efficiency: {avg_inp_packing_efficiency:.3f}\n" f"Avg ctx packing efficiency: {avg_ctx_packing_efficiency:.3f}\n\n" f"Input IDs length stats:\n\n" f" Avg: {np.mean(packed_inp_lens_arr):.1f}, Std: {np.std(packed_inp_lens_arr):.1f}, " f"Min: {np.min(packed_inp_lens_arr)}, Max: {np.max(packed_inp_lens_arr)}\n" f"Context IDs length stats:\n\n" f" Avg: {np.mean(packed_ctx_lens_arr):.1f}, Std: {np.std(packed_ctx_lens_arr):.1f}, " f"Min: {np.min(packed_ctx_lens_arr)}, Max: {np.max(packed_ctx_lens_arr)}" ) return packed_batch if __name__ == "__main__": from datasets import disable_caching from ctx_to_lora.data.processing import get_tokenized_dataset from ctx_to_lora.model_loading import get_model_and_tokenizer disable_caching() setup_logging("tmp/packing_debug.log", debug=True) logger.info("Starting packing script...") model, tokenizer = get_model_and_tokenizer( "google/gemma-2-2b-it", train=True, requires_grad=False, ) base_model_max_len = model.base_model.config.max_position_embeddings tokenizer_kwargs = {"max_length": base_model_max_len} # not used ctx_tokenizer_kwargs = {"max_length": base_model_max_len} # not used for now ds = get_tokenized_dataset( ds_name="self_gen/google/gemma-2-2b-it_temp_0.0/fw_qa_v2_2k_len_level_2_tiny", split="train", max_qas_len=2048, base_model_max_len=model.base_model.config.max_position_embeddings, tokenizer=tokenizer, tokenizer_kwargs={}, ctx_model_max_len=model.base_model.config.max_position_embeddings, ctx_tokenizer=tokenizer, ctx_tokenizer_kwargs={}, add_ctx_to_chat=False, add_repeat_prompt=False, repeat_prob=0.0, add_negative_prompt=False, use_kl_loss=False, ) # ds.set_format("torch") print(ds) packed_ds = ds.map( pack_batch, fn_kwargs={ "max_packed_inp_len": 2**13, # 8k "max_packed_ctx_len": 2**14, # 16k }, batched=True, batch_size=100_000, remove_columns=ds.column_names, num_proc=8, ) print(packed_ds) # packed_ds = PackedDataset( # dataset=ds, # tokenizer=tokenizer, # max_input_length=base_model_max_len, # max_packed_length=base_model_max_len * 2, # # max_packed_size=3, # maximum number of data points being packed # ) # packed_ds.stat() # packed_ds = Dataset.from_list(packed_ds) # packed_ds.set_format("torch") # print(packed_ds) # print(ds[0]) # print(packed_ds[0]) orig_seq = ds[0]["ctx_ids"] packed_seq = packed_ds[0]["ctx_ids"][: len(orig_seq)] print(orig_seq, packed_seq) print(orig_seq == packed_seq)