mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
113 lines
3.5 KiB
Python
113 lines
3.5 KiB
Python
from collections import defaultdict
|
|
|
|
import torch
|
|
|
|
# Store activations
|
|
activations = defaultdict(list)
|
|
|
|
|
|
def apply_forward_hook(module, name):
|
|
def hook_fn(m, input_args, output_tensor):
|
|
activations[name].append(output_tensor.cpu().detach())
|
|
|
|
module.register_forward_hook(hook_fn)
|
|
|
|
|
|
def last_token_pool(
|
|
outputs: dict[str, torch.Tensor], attention_mask: torch.Tensor
|
|
) -> torch.Tensor:
|
|
last_hidden_states = (
|
|
outputs["hidden_states"][-1].detach()
|
|
if "hidden_states" in outputs
|
|
else outputs["last_hidden_state"].detach()
|
|
)
|
|
left_padding = attention_mask[:, -1].sum() == attention_mask.shape[0]
|
|
if left_padding:
|
|
return last_hidden_states[:, -1]
|
|
else:
|
|
sequence_lengths = attention_mask.sum(dim=1) - 1
|
|
batch_size = last_hidden_states.shape[0]
|
|
return last_hidden_states[
|
|
torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths
|
|
]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
from functools import partial
|
|
|
|
import torch
|
|
from peft import get_peft_model
|
|
from torch.utils.data import DataLoader
|
|
from tqdm import tqdm
|
|
|
|
from ctx_to_lora.data.collator import generation_collator
|
|
from ctx_to_lora.data.processing import get_tokenized_dataset
|
|
from ctx_to_lora.model_loading import get_lora_config, get_model_and_tokenizer
|
|
from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
|
|
|
|
model_name_or_path = "google/gemma-2-2b-it"
|
|
base_model, tokenizer = get_model_and_tokenizer(
|
|
model_name_or_path, train=False, requires_grad=False
|
|
)
|
|
ds = get_tokenized_dataset(
|
|
ds_name="squad",
|
|
split="test",
|
|
max_qas_len=-1,
|
|
max_qas_per_sample=1,
|
|
base_model_max_len=8192,
|
|
tokenizer=tokenizer,
|
|
ctx_model_max_len=8192,
|
|
ctx_tokenizer=tokenizer,
|
|
max_ctx_chunk_len=-1,
|
|
min_ctx_chunk_len=-1,
|
|
num_chunk_probs=None,
|
|
max_ctx_chunk_num=None,
|
|
add_ctx_to_chat=False,
|
|
use_kl_loss=True,
|
|
max_new_tokens=0,
|
|
add_self_distill_template=False,
|
|
set_format="pt",
|
|
)
|
|
ds = ds.take(100)
|
|
|
|
peft_config = get_lora_config(
|
|
model_name_or_path,
|
|
lora_r=8,
|
|
lora_dropout=0,
|
|
target_modules=["down_proj"],
|
|
)
|
|
peft_config.lora_alpha = 16
|
|
peft_model = get_peft_model(base_model, peft_config)
|
|
|
|
checkpoint_path = "train_outputs/runs/Sep06_12-27-02_slurm0-a3nodeset-11_84523_bbcb67ca/checkpoint-40000/pytorch_model.bin"
|
|
model = ModulatedPretrainedModel.from_state_dict(
|
|
state_dict=torch.load(checkpoint_path, weights_only=False),
|
|
train=False,
|
|
base_model_kwargs=dict(attn_implementation="flash_attention_2"),
|
|
use_flash_attn=True,
|
|
use_sequence_packing=False, # for generation
|
|
)
|
|
dataloader = DataLoader(
|
|
ds,
|
|
batch_size=1,
|
|
shuffle=False,
|
|
sampler=None,
|
|
batch_sampler=None,
|
|
num_workers=0,
|
|
collate_fn=partial(generation_collator, tokenizer=tokenizer),
|
|
pin_memory=False,
|
|
drop_last=False,
|
|
timeout=0,
|
|
persistent_workers=False,
|
|
)
|
|
|
|
apply_forward_hook(model.ctx_encoder, "ctx_encoder")
|
|
apply_forward_hook(model.layers, "hypernet_layers")
|
|
|
|
for sample_idx, sample in tqdm(enumerate(dataloader)):
|
|
for k, v in sample.items():
|
|
sample[k] = v.to(model.device)
|
|
with torch.no_grad():
|
|
lora_state, _ = model.generate_weights(
|
|
sample["ctx_ids"], sample["ctx_attn_mask"]
|
|
)
|