diff --git a/intx_sft.py b/intx_sft.py index 25890db..42460e3 100755 --- a/intx_sft.py +++ b/intx_sft.py @@ -79,6 +79,7 @@ def get_ds_prob(train_ds_len: list[int], total_len: int): for i, ds_len in enumerate(train_ds_len): if (ds_len / total_len) > 0.01: probs[i] = ds_len / res_total_len * res_probs + logger.debug(f"Dataset probabilities: {probs}") assert isclose(sum(probs), 1.0), ( f"Probs sum to {sum(probs)} ({probs}), expected 1.0" ) @@ -270,6 +271,7 @@ def main(): _get_tokenized_dataset = partial( get_tokenized_dataset, max_qas_len=ctx_args.max_qas_len, + max_qas_per_sample=ctx_args.max_qas_per_sample, base_model_max_len=model.base_model.config.max_position_embeddings, tokenizer=tokenizer, tokenizer_kwargs=tokenizer_kwargs, diff --git a/src/ctx_to_lora/configs.py b/src/ctx_to_lora/configs.py index d656aae..2d25970 100644 --- a/src/ctx_to_lora/configs.py +++ b/src/ctx_to_lora/configs.py @@ -322,6 +322,13 @@ class CtxTrainingArguments: "QA pairs that are longer than this value will be split up into multiple samples." }, ) + max_qas_per_sample: int = field( + default=-1, + metadata={ + "help": "Max QA pair per context. If a context has more QA pairs than this value, " + "they will be split up into multiple samples." + }, + ) max_packed_inp_len: int | None = field( default=2**14, metadata={"help": "Maximum packed input length for training."}, diff --git a/src/ctx_to_lora/data/definitions.py b/src/ctx_to_lora/data/definitions.py index 522fa0f..b3802b8 100644 --- a/src/ctx_to_lora/data/definitions.py +++ b/src/ctx_to_lora/data/definitions.py @@ -84,6 +84,20 @@ LONGBENCH_E_TASKS = [ ] DS_KWARGS = { + "context_numbers_2_10": dict( + train=dict( + path="json", + data_files=glob("data/raw_datasets/context_numbers_[2-9]/train.jsonl") + + ["data/raw_datasets/context_numbers_10/train.jsonl"], + split="train", + ), + validation=dict( + path="json", + data_files=glob("data/raw_datasets/context_numbers_[2-9]/val.jsonl") + + ["data/raw_datasets/context_numbers_10/val.jsonl"], + split="train", + ), + ), "hotpot_qa": dict( train=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="train[900:]"), validation=dict( @@ -631,6 +645,6 @@ EVAL_INTX_TEMPLATES = { for ds_name in LONGBENCH_E_TASKS: EVAL_INTX_TEMPLATES[ds_name] = EVAL_INTX_TEMPLATES[ds_name[:-2]] -for ds_name in DS_KWARGS: - if ds_name not in EVAL_INTX_TEMPLATES: - EVAL_INTX_TEMPLATES[ds_name] = "{input}" +# for ds_name in DS_KWARGS: +# if ds_name not in EVAL_INTX_TEMPLATES: +# EVAL_INTX_TEMPLATES[ds_name] = "{input}" diff --git a/src/ctx_to_lora/data/processing.py b/src/ctx_to_lora/data/processing.py index 84851fe..79b9ab5 100644 --- a/src/ctx_to_lora/data/processing.py +++ b/src/ctx_to_lora/data/processing.py @@ -42,7 +42,7 @@ def load_answers(ds_name, split): elif ds_name == "drop": def extract_ans(sample): - return {"answers": sample["answer_spans"]["spans"]} + return {"answers": sample["answers_spans"]["spans"]} ds_kwargs = get_ds_kwargs(ds_name, split) ds = load_dataset(**ds_kwargs, trust_remote_code=True) @@ -174,7 +174,7 @@ def get_preprocessing_fn( return { "context": sample["passage"], "prompt": prompt, - "response": ", ".join(set(sample["answers_spans"]["spans"])), + "response": sample["answers_spans"]["spans"][0], } elif ds_name == "ropes": @@ -284,10 +284,10 @@ def get_preprocessing_fn( def eval_intx_decorator(f): def g(sample): sample = f(sample) - prompt_field = "prompt" if "prompt" in sample else "prompts" - sample[prompt_field] = prompt_template.format( - input=sample[prompt_field] + assert "prompt" in sample, ( + f"Expected 'prompt' in sample, got {sample.keys()}" ) + sample["prompt"] = prompt_template.format(input=sample["prompt"]) return sample return g @@ -599,6 +599,7 @@ def get_tokenized_dataset( ds_name: str, split: str, max_qas_len: int, + max_qas_per_sample: int, base_model_max_len: int, tokenizer: PreTrainedTokenizerBase, tokenizer_kwargs: dict[str, Any], @@ -633,6 +634,7 @@ def get_tokenized_dataset( ) tokenize_kwargs = dict( max_qas_len=max_qas_len, + max_qas_per_sample=max_qas_per_sample, base_model_max_len=base_model_max_len, tokenizer_kwargs=tokenizer_kwargs, ctx_model_max_len=ctx_model_max_len, @@ -661,7 +663,7 @@ def get_tokenized_dataset( tokenized_ds = datasets.load_from_disk(ds_path) return tokenized_ds - num_proc = None if streaming and ("train" in split) else 8 + num_proc = None if streaming and ("train" in split) else 4 ds = load_and_process_dataset( **load_and_process_kwargs, num_proc=num_proc, @@ -676,8 +678,8 @@ def get_tokenized_dataset( ) if ("train" in split) and is_caching_enabled(): - tokenized_ds.save_to_disk(ds_path, num_proc=num_proc) tokenized_ds = tokenized_ds.shuffle() + tokenized_ds.save_to_disk(ds_path, num_proc=num_proc) # force reload from disk for fingerprint consistency tokenized_ds = datasets.load_from_disk(ds_path) return tokenized_ds @@ -685,6 +687,7 @@ def get_tokenized_dataset( def construct_and_tokenize_ctx_qa( max_qas_len, + max_qas_per_sample, base_model_max_len, tokenizer, tokenizer_kwargs, @@ -826,7 +829,10 @@ def construct_and_tokenize_ctx_qa( logging.debug(f"Split too long QAs with max length {max_qas_len}") tokenized_ds = tokenized_ds.map( split_too_long_qas, - fn_kwargs={"max_qas_len": max_qas_len}, + fn_kwargs={ + "max_qas_len": max_qas_len, + "max_qas_per_sample": max_qas_per_sample, + }, batched=True, batch_size=100_000, num_proc=4, @@ -1010,13 +1016,15 @@ def convert_ctx_prompt_response_to_messages( return dict(data=data) -def split_too_long_qas(samples: dict[str, any], max_qas_len: int): +def split_too_long_qas( + samples: dict[str, any], max_qas_len: int, max_qas_per_sample: int +): # samples keys: "input_ids", "attention_mask", "labels", "ctx_ids", "ctx_attn_mask" # split the qas into multiple samples if they are too long # e.g., if max_qas_len = 512, and qas is 1024 tokens long, # we split it such that each sample has at most 512 tokens # and the ctx_ids and ctx_attn_mask are the same for all samples - if max_qas_len < 0: + if max_qas_len < 0 and max_qas_per_sample < 0: return samples input_ids = samples["input_ids"] attention_mask = samples["attention_mask"] @@ -1029,7 +1037,10 @@ def split_too_long_qas(samples: dict[str, any], max_qas_len: int): longest_old_qas_len = max(total_lengths) if total_lengths else 0 # Early exit if no splitting needed - if all(length <= max_qas_len for length in total_lengths): + if (max_qas_len < 0 or all(length <= max_qas_len for length in total_lengths)) and ( + max_qas_per_sample < 0 + or all(len(seq) <= max_qas_per_sample for seq in input_ids) + ): logger.debug(f"Longest old qas len: {longest_old_qas_len}") logger.debug(f"Longest new qas len: {longest_old_qas_len}") return samples @@ -1047,8 +1058,11 @@ def split_too_long_qas(samples: dict[str, any], max_qas_len: int): out["ctx_attn_mask"].append(ctx_attn) for i, tot_inp_len in enumerate(total_lengths): - if tot_inp_len <= max_qas_len: + if (max_qas_len < 0 or tot_inp_len <= max_qas_len) and ( + max_qas_per_sample < 0 or len(input_ids[i]) <= max_qas_per_sample + ): # No need to split - add entire sample + # logger.debug(f"Sample {i} is within limits, adding as is.") for k in samples: out[k].append(samples[k][i]) continue @@ -1070,7 +1084,12 @@ def split_too_long_qas(samples: dict[str, any], max_qas_len: int): n_skip += 1 continue - if new_qas_len + inp_len <= max_qas_len: + # Check if we can add to current batch (both length and sample count limits) + can_add_to_current = ( + max_qas_len < 0 or new_qas_len + inp_len <= max_qas_len + ) and (max_qas_per_sample < 0 or len(new_input_ids) < max_qas_per_sample) + + if can_add_to_current: # Add to current batch new_qas_len += inp_len new_input_ids.append(inp_ids) @@ -1078,6 +1097,9 @@ def split_too_long_qas(samples: dict[str, any], max_qas_len: int): new_labels.append(label) else: # Current batch is full, save it and start new batch + # logger.debug( + # f"sample {i}: adding batch with {len(new_input_ids)} sequences" + # ) if new_input_ids: # Only add non-empty batches add_batch( new_input_ids, diff --git a/src/ctx_to_lora/eval_utils.py b/src/ctx_to_lora/eval_utils.py index 47203fa..0c6216b 100644 --- a/src/ctx_to_lora/eval_utils.py +++ b/src/ctx_to_lora/eval_utils.py @@ -98,7 +98,8 @@ def normalize_answer(s: str) -> str: def remove_punc(text: str) -> str: exclude = set(string.punctuation) - return "".join(ch for ch in text if ch not in exclude) + # return "".join(ch for ch in text if ch not in exclude) + return " ".join(re.split(f"[{string.punctuation}]", text)) def lower(text: str) -> str: return text.lower() @@ -737,6 +738,7 @@ def evaluate( _get_tokenized_dataset = partial( get_tokenized_dataset, max_qas_len=-1, + max_qas_per_sample=1, base_model_max_len=model.base_model.config.max_position_embeddings, tokenizer=tokenizer, tokenizer_kwargs=tokenizer_kwargs, diff --git a/src/ctx_to_lora/modeling/aggregator.py b/src/ctx_to_lora/modeling/aggregator.py index cd09885..7594bee 100644 --- a/src/ctx_to_lora/modeling/aggregator.py +++ b/src/ctx_to_lora/modeling/aggregator.py @@ -104,7 +104,6 @@ class Perceiver(nn.Module): self.per_rank_gen = per_rank_gen self.r = lora_r if self.per_rank_gen else 1 n_output_queries = num_layers * (num_modules * self.r + num_extra_modules) - self.tot_output_size = n_output_queries self.layer_to_layer = layer_to_layer_ctx_encoder if self.layer_to_layer: n_output_queries = num_modules * self.r + num_extra_modules diff --git a/src/ctx_to_lora/modeling/hypernet.py b/src/ctx_to_lora/modeling/hypernet.py index ab1aadc..9119500 100644 --- a/src/ctx_to_lora/modeling/hypernet.py +++ b/src/ctx_to_lora/modeling/hypernet.py @@ -1,7 +1,6 @@ import logging from collections.abc import Iterable from contextlib import contextmanager -from copy import deepcopy from dataclasses import dataclass from functools import partial from typing import Any @@ -23,7 +22,6 @@ from torch import Tensor, nn from transformers import ( PretrainedConfig, PreTrainedModel, - PreTrainedTokenizerBase, ) from transformers.modeling_outputs import ModelOutput from transformers.models.modernbert.modeling_modernbert import ModernBertModel @@ -1035,91 +1033,91 @@ class ModulatedPretrainedModel(nn.Module): return model_outputs -class ModulatedModelWithSharedInput(nn.Module): - def __init__( - self, - modulated_model: ModulatedPretrainedModel, - base_tokenizer: PreTrainedTokenizerBase, - ctx_tokenizer: PreTrainedTokenizerBase, - ctx_end_predicate: str | None = None, - remove_ctx_from_base_input: bool = False, - ): - super().__init__() - self.modulated_model = modulated_model - self.base_tokenizer = base_tokenizer - self.ctx_tokenizer = ctx_tokenizer - self.register_module("modulated_model", self.modulated_model) - self.ctx_end_predicate = ctx_end_predicate - self.remove_ctx_from_base_input = remove_ctx_from_base_input +# class ModulatedModelWithSharedInput(nn.Module): +# def __init__( +# self, +# modulated_model: ModulatedPretrainedModel, +# base_tokenizer: PreTrainedTokenizerBase, +# ctx_tokenizer: PreTrainedTokenizerBase, +# ctx_end_predicate: str | None = None, +# remove_ctx_from_base_input: bool = False, +# ): +# super().__init__() +# self.modulated_model = modulated_model +# self.base_tokenizer = base_tokenizer +# self.ctx_tokenizer = ctx_tokenizer +# self.register_module("modulated_model", self.modulated_model) +# self.ctx_end_predicate = ctx_end_predicate +# self.remove_ctx_from_base_input = remove_ctx_from_base_input - # delegate to self.modulated_model - @property - def config(self): - return self.modulated_model.config +# # delegate to self.modulated_model +# @property +# def config(self): +# return self.modulated_model.config - @property - def device(self): - return self.modulated_model.device +# @property +# def device(self): +# return self.modulated_model.device - def tie_weights(self): - self.modulated_model.base_model.tie_weights() +# def tie_weights(self): +# self.modulated_model.base_model.tie_weights() - def state_dict(self, *args, **kwargs): - return self.modulated_model.state_dict(*args, **kwargs) +# def state_dict(self, *args, **kwargs): +# return self.modulated_model.state_dict(*args, **kwargs) - def forward(self, *args, **kwargs): - input_ids = kwargs["input_ids"] if "input_ids" in kwargs else args[0] - input_txts = self.base_tokenizer.batch_decode( - input_ids, - skip_special_tokens=True, - clean_up_tokenization_spaces=True, - ) - ctx_txts = deepcopy(input_txts) - if self.ctx_end_predicate: - ctx_txts = [ - txt.split(self.ctx_end_predicate)[0].strip() for txt in input_txts - ] - ctx_inputs = self.ctx_tokenizer(ctx_txts, return_tensors="pt", padding=True).to( - self.device - ) - ctx_ids, ctx_attn_mask = ctx_inputs.input_ids, ctx_inputs.attention_mask - if self.remove_ctx_from_base_input: - raise NotImplementedError("Not implemented") - # input_txts = [ - # txt.split(self.ctx_end_predicate)[1].strip() for txt in input_txts - # ] - # inputs = self.base_tokenizer(input_txts, return_tensors="pt", padding=True) - # input_ids, attention_mask = inputs.input_ids, inputs.attention_mask - # kwargs["input_ids"] = input_ids - # kwargs["attention_mask"] = attention_mask - return self.modulated_model(ctx_ids, ctx_attn_mask, *args, **kwargs) +# def forward(self, *args, **kwargs): +# input_ids = kwargs["input_ids"] if "input_ids" in kwargs else args[0] +# input_txts = self.base_tokenizer.batch_decode( +# input_ids, +# skip_special_tokens=True, +# clean_up_tokenization_spaces=True, +# ) +# ctx_txts = deepcopy(input_txts) +# if self.ctx_end_predicate: +# ctx_txts = [ +# txt.split(self.ctx_end_predicate)[0].strip() for txt in input_txts +# ] +# ctx_inputs = self.ctx_tokenizer(ctx_txts, return_tensors="pt", padding=True).to( +# self.device +# ) +# ctx_ids, ctx_attn_mask = ctx_inputs.input_ids, ctx_inputs.attention_mask +# if self.remove_ctx_from_base_input: +# raise NotImplementedError("Not implemented") +# # input_txts = [ +# # txt.split(self.ctx_end_predicate)[1].strip() for txt in input_txts +# # ] +# # inputs = self.base_tokenizer(input_txts, return_tensors="pt", padding=True) +# # input_ids, attention_mask = inputs.input_ids, inputs.attention_mask +# # kwargs["input_ids"] = input_ids +# # kwargs["attention_mask"] = attention_mask +# return self.modulated_model(ctx_ids, ctx_attn_mask, *args, **kwargs) - def generate(self, *args, **kwargs): - input_ids = kwargs["input_ids"] if "input_ids" in kwargs else args[0] - input_txts = self.base_tokenizer.batch_decode( - input_ids, - skip_special_tokens=True, - clean_up_tokenization_spaces=True, - ) - ctx_txts = deepcopy(input_txts) - if self.ctx_end_predicate: - ctx_txts = [ - txt.split(self.ctx_end_predicate)[0].strip() for txt in input_txts - ] - ctx_inputs = self.ctx_tokenizer(ctx_txts, return_tensors="pt", padding=True).to( - self.device - ) - ctx_ids, ctx_attn_mask = ctx_inputs.input_ids, ctx_inputs.attention_mask - if self.remove_ctx_from_base_input: - raise NotImplementedError("Not implemented") - # input_txts = [ - # txt.split(self.ctx_end_predicate)[1].strip() for txt in input_txts - # ] - # inputs = self.base_tokenizer(input_txts, return_tensors="pt", padding=True) - # input_ids, attention_mask = inputs.input_ids, inputs.attention_mask - # kwargs["input_ids"] = input_ids - # kwargs["attention_mask"] = attention_mask - return self.modulated_model.generate(ctx_ids, ctx_attn_mask, *args, **kwargs) +# def generate(self, *args, **kwargs): +# input_ids = kwargs["input_ids"] if "input_ids" in kwargs else args[0] +# input_txts = self.base_tokenizer.batch_decode( +# input_ids, +# skip_special_tokens=True, +# clean_up_tokenization_spaces=True, +# ) +# ctx_txts = deepcopy(input_txts) +# if self.ctx_end_predicate: +# ctx_txts = [ +# txt.split(self.ctx_end_predicate)[0].strip() for txt in input_txts +# ] +# ctx_inputs = self.ctx_tokenizer(ctx_txts, return_tensors="pt", padding=True).to( +# self.device +# ) +# ctx_ids, ctx_attn_mask = ctx_inputs.input_ids, ctx_inputs.attention_mask +# if self.remove_ctx_from_base_input: +# raise NotImplementedError("Not implemented") +# # input_txts = [ +# # txt.split(self.ctx_end_predicate)[1].strip() for txt in input_txts +# # ] +# # inputs = self.base_tokenizer(input_txts, return_tensors="pt", padding=True) +# # input_ids, attention_mask = inputs.input_ids, inputs.attention_mask +# # kwargs["input_ids"] = input_ids +# # kwargs["attention_mask"] = attention_mask +# return self.modulated_model.generate(ctx_ids, ctx_attn_mask, *args, **kwargs) @contextmanager diff --git a/src/ctx_to_lora/modeling/idefics2.py b/src/ctx_to_lora/modeling/idefics2.py index 155c487..2da110a 100644 --- a/src/ctx_to_lora/modeling/idefics2.py +++ b/src/ctx_to_lora/modeling/idefics2.py @@ -32,6 +32,7 @@ from transformers.utils import ( ) if is_flash_attn_2_available(): + from flash_attn.bert_padding import unpad_input from transformers.modeling_flash_attention_utils import _flash_attention_forward logger = logging.get_logger(__name__) @@ -706,10 +707,11 @@ class Idefics2PerceiverResampler(Idefics2PreTrainedModel): # max_length_k = None if attention_mask is not None: logger.warning_once("Using attention mask for resampler") - seq_lens_k = attention_mask.sum(dim=-1, dtype=torch.int32) - cu_seq_lens_k = torch.cumsum(seq_lens_k, dim=0, dtype=torch.int32) - max_length_k = seq_lens_k.max().item() - cu_seq_lens_k = nn.functional.pad(cu_seq_lens_k, (1, 0)) + context, _, cu_seq_lens_k, max_length_k, _ = unpad_input( + context, attention_mask + ) + context = context.unsqueeze(0) + position_ids = True # goes down flash attn path that uses cu_seq_lens elif position_ids is not None: logger.warning_once("Using position ids for resampler") @@ -739,31 +741,26 @@ class Idefics2PerceiverResampler(Idefics2PreTrainedModel): else: raise ValueError("either position_ids or attention_mask is required") - x_attn_kwargs = dict( - attention_mask=attention_mask, + # attention_mask=attention_mask, position_ids=position_ids, cu_seq_lens_q=cu_seq_lens_q, cu_seq_lens_k=cu_seq_lens_k, max_length_q=max_length_q, max_length_k=max_length_k, ) - self_attn_mask = ( - torch.ones( - (bsz, self.n_latents), device=context.device, dtype=torch.float32 - ) - if attention_mask is not None - else None - ) - self_attn_position_ids = ( - torch.arange( - self.n_latents, device=position_ids.device, dtype=torch.int32 - ).repeat(1, bsz) - if position_ids is not None - else None - ) + # self_attn_mask = ( + # torch.ones( + # (bsz, self.n_latents), device=context.device, dtype=torch.float32 + # ) + # if attention_mask is not None + # else None + # ) + self_attn_position_ids = torch.arange( + self.n_latents, device=context.device, dtype=torch.int32 + ).repeat(1, bsz) self_attn_kwargs = dict( - attention_mask=self_attn_mask, + # attention_mask=self_attn_mask, position_ids=self_attn_position_ids, cu_seq_lens_q=cu_seq_lens_q, cu_seq_lens_k=cu_seq_lens_q, diff --git a/watcher.py b/watcher.py index a0742f9..771b11e 100644 --- a/watcher.py +++ b/watcher.py @@ -93,7 +93,8 @@ if __name__ == "__main__": generative=True, ) metrics.update(gen_metrics) - wandb.log(metrics, step=curstep) + for k in metrics: + wandb.log(metrics[k], step=curstep) wandb.finish() print(f"Logged metrics: {metrics}") print("=" * 80)