diff --git a/run_eval.py b/run_eval.py index 8b9cb46..5e3f390 100644 --- a/run_eval.py +++ b/run_eval.py @@ -63,7 +63,7 @@ if __name__ == "__main__": parser.add_argument( "--max_test_samples_per_ds", type=int, - default=1000, + default=500, help=( "Maximum number of validation samples per dataset. " "If -1, uses values from checkpoint config." @@ -102,6 +102,12 @@ if __name__ == "__main__": action="store_true", help="Use generated queries for context distillation training.", ) + parser.add_argument( + "--q_gen_rounds", + type=int, + default=4, + help="Number of rounds of query generation for context distillation.", + ) parser.add_argument( "--use_iterative_mode", action="store_true", @@ -128,6 +134,21 @@ if __name__ == "__main__": action="store_true", help="Add ctx to base model's input", ) + parser.add_argument( + "--truncate_if_too_long_inp", + action="store_true", + help="Truncate input sequences that are too long", + ) + parser.add_argument( + "--truncate_if_too_long_ctx", + action="store_true", + help="Truncate ctx sequences that are too long", + ) + parser.add_argument( + "--gen_lora_scaling", + type=float, + default=1.0, + ) cli_args = vars(parser.parse_args()) # setup_logging(output_dir, debug=os.getenv("DEBUG", False)) diff --git a/src/ctx_to_lora/eval_utils.py b/src/ctx_to_lora/eval_utils.py index bd2430e..00e64b6 100644 --- a/src/ctx_to_lora/eval_utils.py +++ b/src/ctx_to_lora/eval_utils.py @@ -29,6 +29,7 @@ from ctx_to_lora.data.collator import eval_collator, generation_collator from ctx_to_lora.data.definitions import ( CLOSED_QA_DATASETS, CTX_AFFIXES, + GSM8K_DATASETS, LONGBENCH_E_TASKS, LONGBENCH_TASKS, MULTI_ANSWER_DATASETS, @@ -38,6 +39,7 @@ from ctx_to_lora.data.self_gen_template import SELF_QA_INTX from ctx_to_lora.metrics import ( LENGTH_BINS, Evaluator, + compute_gsm8k_acc, compute_metrics, compute_per_token_acc, compute_perplexity, @@ -614,6 +616,13 @@ def eval_generation( for k, v in qa_f1_metric.items(): eval_result.metrics[f"{split_name}_{k}"] = v eval_result.metrics[f"{split_name}_num_samples_{k}"] = n + elif ds_name in GSM8K_DATASETS: + print("Computing GSM8K Accuracy") + gsm8k_acc_metric, per_sample_metric = compute_gsm8k_acc( + pred_texts, label_texts + ) + for k, v in gsm8k_acc_metric.items(): + eval_result.metrics[f"{split_name}_{k}"] = v else: rouge_metrics, per_sample_metric = compute_rouge(pred_texts, label_texts) for k, v in rouge_metrics.items(): @@ -774,13 +783,21 @@ def evaluate( except FileNotFoundError: raise FileNotFoundError(f"Checkpoint {checkpoint_path} not found. ") ctx_name = state_dict["ctx_encoder_args"].ctx_encoder_model_name_or_path + model = ModulatedPretrainedModel.from_state_dict( state_dict, train=False, base_model_kwargs=model_kwargs, use_flash_attn=True, use_sequence_packing=False, # for generation + user_defined_scaling=args.gen_lora_scaling, ) + if getattr(args, "use_llmlingua", False): + print("Using LLMLingua-2 for compressing inp") + inp_compressor = LLMLinguaModel( + model.base_model, tokenizer, args.llmlingua_compression_rate + ) + model.inp_compressor = inp_compressor ctx_model_max_len = model.ctx_encoder.config.max_position_embeddings model.enable_iterative_mode(args.use_iterative_mode) add_tracker(model.base_model.generate, "generate") @@ -832,6 +849,7 @@ def evaluate( ) ctx_distill_kwargs["q_model"] = q_model ctx_distill_kwargs["q_tokenizer"] = q_tokenizer + ctx_distill_kwargs["q_gen_rounds"] = args.q_gen_rounds model = CtxDistillModel(peft_model, **ctx_distill_kwargs) add_tracker(model._distill_context, "distill_context") @@ -892,6 +910,8 @@ def evaluate( max_new_tokens=max_new_tokens, set_format="pt", add_self_distill_template=use_cd, # only for eval + truncate_if_too_long_inp=args.truncate_if_too_long_inp, # only for eval + truncate_if_too_long_ctx=args.truncate_if_too_long_ctx, # only for eval ) datasets = dict() @@ -1046,11 +1066,15 @@ def run_eval( use_cd: bool = False, cd_update_iterations: int = 10, cd_use_gen_q: bool = False, + q_gen_rounds: int = 4, use_iterative_mode: bool = False, use_llmlingua: bool = False, llmlingua_compression_rate: float = 0.9, use_t2l: bool = False, add_ctx_to_input: bool = False, + truncate_if_too_long_inp: bool = False, + truncate_if_too_long_ctx: bool = False, + gen_lora_scaling: float = 1, ) -> None: """Run evaluation with the specified parameters.""" assert bool(model_name_or_path) ^ bool(checkpoint_path), ( @@ -1100,6 +1124,9 @@ def run_eval( # but remove_context has to be false for correct file naming args.remove_context = False args.use_iterative_mode = use_iterative_mode + if use_llmlingua: + args.use_llmlingua = use_llmlingua + args.llmlingua_compression_rate = llmlingua_compression_rate else: args = Namespace( model_name_or_path=model_name_or_path, @@ -1114,6 +1141,7 @@ def run_eval( args.use_cd = use_cd args.cd_update_iterations = cd_update_iterations args.cd_use_gen_q = cd_use_gen_q + args.q_gen_rounds = q_gen_rounds if use_llmlingua: args.use_llmlingua = use_llmlingua args.llmlingua_compression_rate = llmlingua_compression_rate @@ -1124,6 +1152,9 @@ def run_eval( if max_test_samples_per_ds > 0: args.max_test_samples_per_ds = max_test_samples_per_ds args.add_ctx_to_input = add_ctx_to_input + args.gen_lora_scaling = gen_lora_scaling + args.truncate_if_too_long_inp = truncate_if_too_long_inp + args.truncate_if_too_long_ctx = truncate_if_too_long_ctx setup_logging(args.logging_dir) logger.debug(f"CMD: {' '.join(os.sys.argv)}") diff --git a/src/ctx_to_lora/metrics.py b/src/ctx_to_lora/metrics.py index ab163c4..00c4d3d 100644 --- a/src/ctx_to_lora/metrics.py +++ b/src/ctx_to_lora/metrics.py @@ -1,3 +1,4 @@ +import re from collections import defaultdict from collections.abc import Callable @@ -31,6 +32,34 @@ def get_length_bin(length: int): return (start, end) +def compute_gsm8k_acc(pred_texts, label_texts): + out = defaultdict(list) + for pred_text, label_text in zip(pred_texts, label_texts): + if "Answer" not in pred_text: + out["gsm8k_acc"].append(0) + continue + + label_ans = float(label_text.split("####")[-1].strip().replace(",", "")) + # find all the numbers (including decimals) in the string + numbers = re.findall(r"\d+\.?\d*", pred_text) + + if not numbers: + out["gsm8k_acc"].append(0) + continue + + prediction = float(numbers[-1].rstrip(".").strip().replace(",", "")) + + if label_ans == prediction: + out["gsm8k_acc"].append(1) + else: + out["gsm8k_acc"].append(0) + + out_mean = dict() + for k in out: + out_mean[k] = np.mean(out[k]) + return out_mean, out + + def compute_rouge(pred_texts, label_texts): out = defaultdict(list) scorer = rouge_scorer.RougeScorer(["rougeL"], use_stemmer=True) diff --git a/src/ctx_to_lora/model_loading.py b/src/ctx_to_lora/model_loading.py index ed68bc8..6d607df 100644 --- a/src/ctx_to_lora/model_loading.py +++ b/src/ctx_to_lora/model_loading.py @@ -82,7 +82,12 @@ def get_tokenizer( tokenizer.pad_token_id = tokenizer.eos_token_id template_path = f"chat_templates/{model_name_or_path}.jinja" - assert os.path.exists(template_path), f"Chat template not found at {template_path}." + # assert os.path.exists(template_path), f"Chat template not found at {template_path}." + if not os.path.exists(template_path): + logger.warning( + f"Chat template not found at {template_path}. Using default template." + ) + return tokenizer logger.info(f"Using chat template from {template_path}") chat_template = open(template_path).read() @@ -108,7 +113,7 @@ def get_model( torch_dtype=dtype, trust_remote_code=True, attn_implementation="eager", - use_cache=False, + use_cache=None, ) is_vision_model = check_is_vision_model(model_name_or_path) if model_kwargs is not None: diff --git a/src/ctx_to_lora/modeling/context_distillation.py b/src/ctx_to_lora/modeling/context_distillation.py index 2f0077d..b3bb510 100644 --- a/src/ctx_to_lora/modeling/context_distillation.py +++ b/src/ctx_to_lora/modeling/context_distillation.py @@ -1,5 +1,7 @@ import gc +import json import logging +import os import re from typing import Any @@ -189,6 +191,9 @@ class CtxDistillModel(nn.Module): q_model: PreTrainedModel | None = None, q_tokenizer=None, reprompt_ctx: bool = False, + lora_save_dir: str | None = None, + save_after_distill: bool = True, + q_gen_rounds: int = 4, ): super().__init__() self.register_module("base_model", base_model) @@ -203,7 +208,10 @@ class CtxDistillModel(nn.Module): self.reset = reset self.device = base_model.device self.to(self.device) - self.q_gen_rounds = 4 + self.q_gen_rounds = q_gen_rounds + # New save options + self.lora_save_dir = lora_save_dir + self.save_after_distill = save_after_distill self.peft_config = base_model.peft_config["default"] self.adapter_name = "default" @@ -236,6 +244,26 @@ class CtxDistillModel(nn.Module): layer.reset_lora_parameters(self.adapter_name, init_lora_weights=True) self._init_optim() + def save_lora(self): + """ + Save current LoRA adapter in PEFT format plus a lightweight JSON summary + for easy human inspection/manipulation. + """ + if self.lora_save_dir is None: + return + os.makedirs(self.lora_save_dir, exist_ok=True) + # Standard PEFT save (produces adapter_config.json + adapter_model.bin / safetensors) + self.base_model.save_pretrained(self.lora_save_dir) + # Human-readable summary of LoRA parameter shapes + summary = { + name: list(p.shape) + for name, p in self.base_model.named_parameters() + if ("lora_A" in name or "lora_B" in name) and p.requires_grad + } + with open(os.path.join(self.lora_save_dir, "lora_summary.json"), "w") as f: + json.dump(summary, f, indent=2) + print(f"LoRA adapter saved to {self.lora_save_dir}") + @torch.enable_grad() def _distill_context( self, @@ -312,9 +340,21 @@ class CtxDistillModel(nn.Module): # rename for separate timing return self.base_model.generate(*args, **kwargs) + def get_lora_state(self, clone: bool = True): + """ + Return a dict of current LoRA parameter tensors. + clone=True returns detached cloned tensors (safe to store). + """ + return { + name: (p.detach().clone() if clone else p) + for name, p in self.base_model.named_parameters() + if ("lora_A" in name or "lora_B" in name) + } + def generate( self, *model_inputs_args: Any, + distill_only: bool = False, **model_inputs_kwargs: dict[str, Any], ): if self.reset: @@ -450,6 +490,9 @@ class CtxDistillModel(nn.Module): inp_res_attention_mask, student_labels, ) + # Save LoRA after distillation if requested + if distill_only: + return self.get_lora_state() # _, inp_ids = ctx_inp_split( # ctx_inp_ids, @@ -535,6 +578,8 @@ if __name__ == "__main__": q_tokenizer=q_tokenizer, tokenizer=tokenizer, reprompt_ctx=False, + lora_save_dir="./saved_lora_adapter", # example save path + save_after_distill=True, ) with torch.no_grad(): diff --git a/src/ctx_to_lora/modeling/hypernet.py b/src/ctx_to_lora/modeling/hypernet.py index 287f633..6625e23 100644 --- a/src/ctx_to_lora/modeling/hypernet.py +++ b/src/ctx_to_lora/modeling/hypernet.py @@ -710,6 +710,8 @@ class ModulatedPretrainedModel(nn.Module): use_base_input_as_ctx: bool = False, # need non-packed inputs for generation use_sequence_packing: bool = True, + user_defined_scaling: float = 1, + inp_compressor=None, ): assert not use_base_input_as_ctx super().__init__() @@ -719,6 +721,8 @@ class ModulatedPretrainedModel(nn.Module): self.ctx_encoder_args = ctx_encoder_args self.use_base_input_as_ctx = use_base_input_as_ctx self.use_sequence_packing = use_sequence_packing + self.user_defined_scaling = user_defined_scaling + self.inp_compressor = inp_compressor self.model_accepts_loss_kwargs = True self.active_adapters = [] @@ -959,10 +963,18 @@ class ModulatedPretrainedModel(nn.Module): if isinstance(self.ctx_encoder.base_model, ModernBertModel): ctx_features = ctx_features.unsqueeze(0) + if self.user_defined_scaling == 1: + return self.hypernet.generate_weights( + ctx_features, ctx_attn_mask, ctx_position_ids + ) - return self.hypernet.generate_weights( + lora_dict, _ = self.hypernet.generate_weights( ctx_features, ctx_attn_mask, ctx_position_ids ) + for module in lora_dict: + lora_dict[module]["A"] = lora_dict[module]["A"] * self.user_defined_scaling + lora_dict[module]["B"] = lora_dict[module]["B"] * self.user_defined_scaling + return lora_dict, None def enable_iterative_mode(self, x: bool): self.hypernet.enable_iterative_mode(x) @@ -1218,6 +1230,19 @@ class ModulatedPretrainedModel(nn.Module): n_queries, position_ids, ) + if self.inp_compressor is not None: + print("Using inp compressor") + input_ids = model_inputs_kwargs["input_ids"] + print(f"input_ids before: {input_ids.shape}") + model_inputs_kwargs["input_ids"] = self.inp_compressor.compress_tokens( + input_ids, + query_text="Answer the following question. Only output the answer and do not output any other words.", + ).to(self.base_model.device) + # batch size always 1 + model_inputs_kwargs["attention_mask"] = torch.ones_like( + model_inputs_kwargs["input_ids"] + ) + print(f"input_ids after: {model_inputs_kwargs['input_ids'].shape}") model_outputs = self.base_model.generate( *model_inputs_args, **model_inputs_kwargs ) diff --git a/src/ctx_to_lora/modeling/llm_lingua.py b/src/ctx_to_lora/modeling/llm_lingua.py index f2b33d6..99e9b17 100644 --- a/src/ctx_to_lora/modeling/llm_lingua.py +++ b/src/ctx_to_lora/modeling/llm_lingua.py @@ -30,6 +30,28 @@ class LLMLinguaModel(nn.Module): prompt_txt, rate=rate, force_tokens=["\n", "?"] ) + def compress_tokens(self, input_ids, query_text): + bs = input_ids.shape[0] + txt = self.tokenizer.batch_decode( + input_ids[:, self.len_prefix : -self.len_suffix] + ) + q_start_idx = txt[0].rfind(query_text) + ctx_txt = txt[0][:q_start_idx] + q_txt = txt[0][q_start_idx:] + compressed_txt = self.compress(ctx_txt, rate=self.compression_rate) + compressed_ids = self.tokenizer( + compressed_txt["compressed_prompt"] + "\n\n" + q_txt, + return_attention_mask=False, + add_special_tokens=False, + return_tensors="pt", + )["input_ids"].to(self.base_model.device) + + out = torch.cat( + [self.prefix.expand(bs, -1), compressed_ids, self.suffix.expand(bs, -1)], + dim=-1, + ) + return out + def generate(self, *args, **kwargs): # take ctx_ids # strip prefix and suffix @@ -37,7 +59,6 @@ class LLMLinguaModel(nn.Module): ctx_ids = kwargs["ctx_ids"][:, self.len_prefix : -self.len_suffix] # decode ctx_ids to ctx_txt ctx_txt = self.tokenizer.batch_decode(ctx_ids) - # 4x compression compressed_ctx_txt = self.compress(ctx_txt, rate=self.compression_rate) compressed_ctx_ids = self.tokenizer( compressed_ctx_txt["compressed_prompt"] + "\n\n",