apunkt/doc-to-lora
1
0
Fork 0
mirror of https://github.com/SakanaAI/doc-to-lora.git synced 2026-07-23 17:01:04 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

scaling gen lora + q gen round for cd + gsm8k

Browse source
This commit is contained in:
51616 2025-09-18 16:02:22 +09:00
parent 2af651fa41
commit 76bf6981ba
7 changed files with 183 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

23
run_eval.py
View file

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

31
src/ctx_to_lora/eval_utils.py
View file

@ -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)}")

29
src/ctx_to_lora/metrics.py
View file

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

9
src/ctx_to_lora/model_loading.py
View file

@ -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:

47
src/ctx_to_lora/modeling/context_distillation.py
View file

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

27
src/ctx_to_lora/modeling/hypernet.py
View file

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

23
src/ctx_to_lora/modeling/llm_lingua.py
View file

@ -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",