mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
python api + remove generate_multi_lora + rename config path + test eval scripts
This commit is contained in:
51616
parent
530fcd46cc
commit
3b798c670c
27 changed files with 345 additions and 233 deletions
93
README.md
93
README.md
|
|
@ -1,5 +1,5 @@
|
|||
<div align="center">
|
||||
<h1>Ctx-to-LoRA</h1>
|
||||
<h1>Doc-to-LoRA</h1>
|
||||
<br>
|
||||
<img height="500px" src="assets/cover.png" />
|
||||
</div>
|
||||
|
|
@ -7,56 +7,57 @@
|
|||
|
||||
---
|
||||
|
||||
## 🚀 API Usage [WIP]
|
||||
## 🚀 Python API Usage
|
||||
```python
|
||||
from ctx_to_lora.modeling import ModulatedPretrainedModel
|
||||
model = ModulatedPretrainedModel.from_state_dict(...)
|
||||
# caveat: this interface only supports non-batched inputs
|
||||
# for batched inference please see `src/ctx_to_lora/modeling/hypernet.py`
|
||||
import torch
|
||||
|
||||
ctx_info = "..."
|
||||
query = "..."
|
||||
from ctx_to_lora.model_loading import get_tokenizer
|
||||
from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
|
||||
|
||||
ctx_ids = model.ctx_encoder.tokenize(ctx_info)
|
||||
input_ids = model.tokenize(query)
|
||||
outputs = model.generate(ctx_ids, input_ids)
|
||||
print(model.decode(outputs))
|
||||
# model loading
|
||||
checkpoint_path = ...
|
||||
state_dict = torch.load(checkpoint_path, weights_only=False)
|
||||
model = ModulatedPretrainedModel.from_state_dict(
|
||||
state_dict, train=False, use_sequence_packing=False
|
||||
)
|
||||
model.reset()
|
||||
tokenizer = get_tokenizer(model.base_model.name_or_path)
|
||||
|
||||
# prepare data
|
||||
doc = open("data/sakana_wiki.txt", "r").read()
|
||||
chat = [{"role": "user", "content": "Summarize what Sakana AI does."}]
|
||||
chat_ids = tokenizer.apply_chat_template(
|
||||
chat,
|
||||
add_special_tokens=False,
|
||||
return_attention_mask=False,
|
||||
add_generation_prompt=True,
|
||||
return_tensors="pt",
|
||||
).to(model.device)
|
||||
|
||||
|
||||
# calls after internalization will be influenced by internalized info
|
||||
model.internalize(doc)
|
||||
|
||||
outputs = model.generate(input_ids=chat_ids, max_new_tokens=256)
|
||||
print(tokenizer.decode(outputs[0]))
|
||||
|
||||
|
||||
# remove internalized info
|
||||
model.reset()
|
||||
|
||||
outputs = model.generate(input_ids=chat_ids, max_new_tokens=256)
|
||||
print(tokenizer.decode(outputs[0]))
|
||||
```
|
||||
|
||||
|
||||
***Generate data from scratch***
|
||||
# 0. download fineweb_edu to `data/raw_datasets/fineweb_edu
|
||||
uv run data/download_fineweb_edu.py
|
||||
```
|
||||
|
||||
1. Recursively generate more data! (depends on step 0)
|
||||
### 🎮 Interactive Demo [WIP]
|
||||
```bash
|
||||
# run from 000 to 0013
|
||||
run uv run data/generate_fw_edu_qa_v2.py --shard_pattern "000_00000" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it --max_length=2000 --max_model_length=2048;
|
||||
run uv run data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_0*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it;
|
||||
uv run webui/demo.py
|
||||
```
|
||||
|
||||
2. Self-generated response QA data (depends on step 0 and 1)
|
||||
```bash
|
||||
# Example commands using gemma-2-2b-it
|
||||
# self-gen data for fw_qa_v2
|
||||
uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_pattern 'data/raw_datasets/fw_qa_v2/min_0_to_2000/013*_level_3*' --closed_qa_prob 1.0 # or 0.0
|
||||
|
||||
# val split
|
||||
uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_pattern 'data/raw_datasets/fw_qa_v2/min_0_to_2000/*_level_0_val.parquet'
|
||||
|
||||
# self-gen data for other ds listed in qa_short_ctx_self_gen_no_fw_qa.yaml
|
||||
uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --config configs/qa_short_ctx_self_gen_no_fw_qa.yaml
|
||||
```
|
||||
|
||||
### Evaluation
|
||||
See [eval scripts](scripts/eval/).
|
||||
```bash
|
||||
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/pytorch_model.bin --datasets squad --split test --max_ctx_chunk_len 8192 --eval_batch_size_gen 4
|
||||
|
||||
|
||||
# base model
|
||||
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets negative_nq triviaqa_retrieved squad longbench_e --split test --eval_batch_size 2
|
||||
|
||||
# base model w/o context
|
||||
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets negative_nq triviaqa_retrieved squad longbench_e --split test --remove_context
|
||||
|
||||
|
||||
### 🧪 Experimental Scripts
|
||||
| Experiment | Data prep | Training | Evaluation | Notes |
|
||||
| ------------------------------------ | ----------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| [Main experiment](scripts/main_exp/) | `uv run bash scripts/main_exp/0-download_data.sh`<br>`uv run bash scripts/main_exp/gen_data.sh` *(optional, rebuilds from scratch)* | `uv run bash scripts/main_exp/1-train.sh` | `uv run bash scripts/main_exp/eval/<script>.sh` (pick from `base_model.sh`, `cd.sh`, `cd_oracle.sh`, `d2l.sh`, `llmlingua.sh`, `t2l.sh`) | Downloading data is fastest; regenerate only if you need fresh synthetic data. Evaluation scripts reproduce the main paper metrics. |
|
||||
| [NIAH](scripts/niah/) | `uv run bash scripts/niah/0-gen_data.sh` | `uv run bash scripts/niah/1-train.sh` | `uv run bash scripts/niah/2-eval.sh` | Run the scripts in order; data generation only needs to happen once |
|
||||
|
|
|
|||
32
configs/main_exp/self_gen_lv1_closed_qa_1_and_lv3_l2l.yaml
Normal file
32
configs/main_exp/self_gen_lv1_closed_qa_1_and_lv3_l2l.yaml
Normal file
|
|
@ -0,0 +1,32 @@
|
|||
# LoRA
|
||||
lora_r: 8
|
||||
lora_dropout: 0.0
|
||||
target_modules:
|
||||
- down_proj
|
||||
|
||||
use_kl_loss: true
|
||||
|
||||
ctx_encoder_type: per_layer_activations
|
||||
n_latent_queries: 8
|
||||
num_blocks: 9
|
||||
num_self_attn_per_block: 0
|
||||
|
||||
gradient_accumulation_steps: 11
|
||||
max_packed_inp_len: 6144
|
||||
max_packed_ctx_len: 6144
|
||||
|
||||
# data
|
||||
train_ds_names:
|
||||
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_1*.parquet
|
||||
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_3*.parquet
|
||||
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact
|
||||
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/squad_compact
|
||||
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/ropes_compact
|
||||
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/drop_compact
|
||||
|
||||
val_ds_names:
|
||||
- squad
|
||||
- pwc
|
||||
- drop
|
||||
- ropes
|
||||
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet
|
||||
|
|
@ -15,9 +15,6 @@ gradient_accumulation_steps: 11
|
|||
max_packed_inp_len: 6144
|
||||
max_packed_ctx_len: 6144
|
||||
|
||||
|
||||
|
||||
|
||||
# data
|
||||
train_ds_names:
|
||||
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_1*.parquet
|
||||
|
|
@ -31,11 +28,4 @@ val_ds_names:
|
|||
- pwc
|
||||
- drop
|
||||
- ropes
|
||||
# - hotpot_qa
|
||||
# - negative_nq
|
||||
# - booksum
|
||||
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet
|
||||
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet
|
||||
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_qa/validation/*.parquet
|
||||
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/facts/validation/*.parquet
|
||||
|
||||
|
|
|
|||
8
data/sakana_wiki.txt
Normal file
8
data/sakana_wiki.txt
Normal file
|
|
@ -0,0 +1,8 @@
|
|||
Sakana AI Co, Ltd. is a Japanese artificial intelligence company based in Tokyo.
|
||||
Overview
|
||||
|
||||
Sakana AI's main research fields are evolution and collective intelligence of AI. The company's name is derived from the Japanese word さかな (sakana), which means fish. This represents the idea of a school of fish coming together and forming a coherent entity from simple rules, which is an analogy of collective intelligence.[2]
|
||||
|
||||
The company was founded by David Ha, Llion Jones and Ren Ito. Llion Jones co-authored the famous paper "Attention Is All You Need" when he was working for Google in 2017. The company raised $30M in its seed funding round from Lux Capital and Khosla Ventures.[3] The company raised approximately $200M from companies such as Mitsubishi UFJ, SMBC, Mizuho, Itochu, KDDI, Nomura and Nvidia in its series A funding round in 2024.[4]
|
||||
|
||||
In January 2024, Sakana AI developed a method to build new AI models by 'breeding' multiple existing models, which it sees as a means to democratise AI development, as this process does not require large computational resources.[5] Sakana AI is also developing a model called the AI Scientist, which automates the entire process of scientific research.[6] The Nikkei estimated the company's value at 19 billion yen in 2024.[7]
|
||||
|
|
@ -186,21 +186,10 @@ def self_generate(
|
|||
raise ValueError(
|
||||
f"Multiple sources of truth for closed_qa_prob: CLI arg --closed_qa_prob={args.closed_qa_prob} and dataset name contains closed_qa_prob specification."
|
||||
)
|
||||
if "_summ_prob_" in ds_name and args.summ_prob != 0.0:
|
||||
raise ValueError(
|
||||
f"Multiple sources of truth for summ_prob: CLI arg --summ_prob={args.summ_prob} and dataset name contains summ_prob specification."
|
||||
)
|
||||
if "_struct_prob_" in ds_name and args.struct_prob != 0.0:
|
||||
raise ValueError(
|
||||
f"Multiple sources of truth for struct_prob: CLI arg --struct_prob={args.struct_prob} and dataset name contains struct_prob specification."
|
||||
)
|
||||
|
||||
# Base values from args
|
||||
temp = args.temp
|
||||
closed_qa_prob = args.closed_qa_prob
|
||||
summ_prob = args.summ_prob
|
||||
struct_prob = args.struct_prob
|
||||
cot_prob = args.cot_prob
|
||||
|
||||
# Overrides from ds_name pattern if present
|
||||
if ds_name is not None:
|
||||
|
|
@ -212,21 +201,10 @@ def self_generate(
|
|||
m = re.search(r"_closed_qa_prob_([\d.]+)", ds_name)
|
||||
if m:
|
||||
closed_qa_prob = float(m.group(1))
|
||||
if "_summ_prob_" in ds_name:
|
||||
m = re.search(r"_summ_prob_([\d.]+)", ds_name)
|
||||
if m:
|
||||
summ_prob = float(m.group(1))
|
||||
if "_struct_prob_" in ds_name:
|
||||
m = re.search(r"_struct_prob_([\d.]+)", ds_name)
|
||||
if m:
|
||||
struct_prob = float(m.group(1))
|
||||
|
||||
print(f"Processing dataset: {ds_name}, split: {split}")
|
||||
print(f"Using temperature: {temp}")
|
||||
print(f"Using closed QA prompt probability: {closed_qa_prob}")
|
||||
print(f"Using chain-of-thought prompt probability: {cot_prob}")
|
||||
print(f"Using summarization augmentation probability: {summ_prob}")
|
||||
print(f"Using structuring augmentation probability: {struct_prob}")
|
||||
|
||||
if parquet_file:
|
||||
print(f"Loading dataset from parquet file: {parquet_file}")
|
||||
|
|
@ -249,7 +227,7 @@ def self_generate(
|
|||
print(f"Loaded dataset: {ds_name} with split: {split}")
|
||||
|
||||
if args.debug:
|
||||
ds = ds.take(50)
|
||||
ds = ds.take(10)
|
||||
|
||||
ds = ds.filter(filter_none, batched=False, num_proc=8)
|
||||
|
||||
|
|
@ -445,22 +423,6 @@ def execute_qa_generation(
|
|||
# bos + question + eos + start model turn + response + eos
|
||||
input_ids = all_ids[:sys_start] + all_ids[q_start:res_end]
|
||||
|
||||
# TODO: save prompt logprobs with qa_start
|
||||
# TODO: also save the prompt tokens for output alignment during training
|
||||
# prompt_logp = completions[c + i].prompt_logprobs
|
||||
# q_len = res_start - q_start
|
||||
# prompt_logp_indices = np.empty((q_len, k), dtype=np.int32)
|
||||
# # float-16 is better for this range
|
||||
# prompt_logp_vals = np.empty((q_len, k), dtype=np.float16)
|
||||
# for li, info_d in enumerate(prompt_logp[q_start:res_start]):
|
||||
# for j, (idx, tok_info) in enumerate(info_d.items()):
|
||||
# # vllm returns logprob for gt token first,
|
||||
# # which means that we might have k + 1 logprobs
|
||||
# if j >= k:
|
||||
# continue
|
||||
# prompt_logp_indices[li, j] = idx
|
||||
# prompt_logp_vals[li, j] = tok_info.logprob
|
||||
|
||||
# relative to the input_ids
|
||||
res_start = res_start - q_start + sys_start
|
||||
res_end = res_start + n_response_tokens
|
||||
|
|
@ -472,10 +434,6 @@ def execute_qa_generation(
|
|||
self_gen_data[ctx]["response_start_end"].append((res_start, res_end))
|
||||
self_gen_data[ctx]["logprobs_vals"].append(logp_vals)
|
||||
self_gen_data[ctx]["logprobs_indices"].append(logp_indices)
|
||||
# NOTE: also have to shift back 1 for training
|
||||
# self_gen_data[ctx]["prompt_start_end"].append((sys_start, res_start))
|
||||
# self_gen_data[ctx]["prompt_logprobs_vals"].append(prompt_logp_vals)
|
||||
# self_gen_data[ctx]["prompt_logprobs_indices"].append(prompt_logp_indices)
|
||||
|
||||
c += i + 1
|
||||
|
||||
|
|
@ -516,20 +474,22 @@ def execute_qa_generation(
|
|||
# random.shuffle(samples)
|
||||
|
||||
# Save results
|
||||
if not args.debug:
|
||||
# df = pd.DataFrame(samples)
|
||||
# ds_out = Dataset.from_pandas(df)
|
||||
ds_out = Dataset.from_list(samples)
|
||||
# fpath = f"{SELF_GEN_DATA_DIR}/{args.vllm_model}_temp_{temp}_closed_qa_prob_{closed_qa_prob}/{ds_name}/{split}/ds{shard_name}"
|
||||
# df = pd.DataFrame(samples)
|
||||
# ds_out = Dataset.from_pandas(df)
|
||||
ds_out = Dataset.from_list(samples)
|
||||
# fpath = f"{SELF_GEN_DATA_DIR}/{args.vllm_model}_temp_{temp}_closed_qa_prob_{closed_qa_prob}/{ds_name}/{split}/ds{shard_name}"
|
||||
|
||||
os.makedirs(os.path.dirname(fpath), exist_ok=True)
|
||||
fpath = f"{fpath}.parquet"
|
||||
ds_out.to_parquet(fpath)
|
||||
print(f"Saved to {fpath}")
|
||||
if args.debug:
|
||||
fpath += "_debug"
|
||||
os.makedirs(os.path.dirname(fpath), exist_ok=True)
|
||||
|
||||
# Cleanup
|
||||
del samples, ds_out, completions, messages, ctxs, questions
|
||||
clear_gpu()
|
||||
fpath = f"{fpath}.parquet"
|
||||
ds_out.to_parquet(fpath)
|
||||
print(f"Saved to {fpath}")
|
||||
|
||||
# Cleanup
|
||||
del samples, ds_out, completions, messages, ctxs, questions
|
||||
clear_gpu()
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
|
|
|
|||
43
examples/python_api.py
Normal file
43
examples/python_api.py
Normal file
|
|
@ -0,0 +1,43 @@
|
|||
# caveat: this interface only supports non-batched inputs
|
||||
# for batched inference please see `src/ctx_to_lora/modeling/hypernet.py`
|
||||
import torch
|
||||
|
||||
from ctx_to_lora.model_loading import get_tokenizer
|
||||
from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
|
||||
|
||||
checkpoint_path = "../../ctx-to-lora/train_outputs/runs/Sep29_14-42-46_slurm0-a3nodeset-9_88483_1e7bb34e/checkpoint-40000/pytorch_model.bin"
|
||||
state_dict = torch.load(checkpoint_path, weights_only=False)
|
||||
model = ModulatedPretrainedModel.from_state_dict(
|
||||
state_dict, train=False, use_sequence_packing=False
|
||||
)
|
||||
model.reset()
|
||||
|
||||
tokenizer = get_tokenizer(model.base_model.name_or_path)
|
||||
|
||||
doc = open("data/sakana_wiki.txt").read()
|
||||
chat = [{"role": "user", "content": "Summarize what Sakana AI does."}]
|
||||
chat_ids = tokenizer.apply_chat_template(
|
||||
chat,
|
||||
add_special_tokens=False,
|
||||
return_attention_mask=False,
|
||||
add_generation_prompt=True,
|
||||
return_tensors="pt",
|
||||
).to(model.device)
|
||||
|
||||
|
||||
outputs = model.generate(input_ids=chat_ids, max_new_tokens=256)
|
||||
print(tokenizer.decode(outputs[0]))
|
||||
|
||||
|
||||
# calls after internalization will be influenced by internalized info
|
||||
model.internalize(doc)
|
||||
|
||||
outputs = model.generate(input_ids=chat_ids, max_new_tokens=256)
|
||||
print(tokenizer.decode(outputs[0]))
|
||||
|
||||
|
||||
# remove internalized info
|
||||
model.reset()
|
||||
|
||||
outputs = model.generate(input_ids=chat_ids, max_new_tokens=256)
|
||||
print(tokenizer.decode(outputs[0]))
|
||||
|
|
@ -2,6 +2,7 @@
|
|||
# module load cuda/12.4
|
||||
# module load cudnn/9.1.0
|
||||
# module load hpcx/v2.21
|
||||
curl -LsSf https://astral.sh/uv/install.sh | sh
|
||||
uv self update
|
||||
uv venv --python 3.10 --seed
|
||||
uv pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --torch-backend=cu124
|
||||
|
|
@ -16,12 +17,11 @@ uv run data/build_pwc_compact.py
|
|||
uv run data/build_ropes_compact.py
|
||||
uv run data/build_squad_compact.py
|
||||
|
||||
# [Optional]
|
||||
# needed for gated models
|
||||
# optional: needed for gated models
|
||||
# uv run huggingface-cli login
|
||||
|
||||
# needed for logging with wandb
|
||||
# optional: needed for logging with wandb
|
||||
# wandb login
|
||||
|
||||
# dev
|
||||
# optional: dev
|
||||
# uv run pre-commit install
|
||||
|
|
|
|||
11
scripts/main_exp/0-download_data.py
Normal file
11
scripts/main_exp/0-download_data.py
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
# TO BE ADDED: see which split is better
|
||||
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
if __name__ == "__main__":
|
||||
fw_dir = "./data/raw_datasets/self_gen/"
|
||||
snapshot_download(
|
||||
"SakanaAI/self_gen_qa_fw_qa",
|
||||
repo_type="dataset",
|
||||
local_dir=fw_dir,
|
||||
)
|
||||
25
scripts/main_exp/README.md
Normal file
25
scripts/main_exp/README.md
Normal file
|
|
@ -0,0 +1,25 @@
|
|||
# D2L pipeline
|
||||
### Data
|
||||
You can either download the generated data (recommended) or generate them by youself.
|
||||
```bash
|
||||
# download training data (recommended)
|
||||
uv run bash scripts/main_exp/0-download_data.sh
|
||||
```
|
||||
|
||||
Generating data from scratch can take very long if not parallelized across multiple gpus.
|
||||
```bash
|
||||
# generate training data (takes very long if not parallelized across multiple gpus)
|
||||
# optional: use the command below for generating data from scratch
|
||||
# uv run bash scripts/main_exp/gen_data.sh
|
||||
```
|
||||
|
||||
|
||||
### Training
|
||||
Simply run the training script.
|
||||
```bash
|
||||
# train
|
||||
uv run bash scripts/main_exp/1-train.sh
|
||||
```
|
||||
|
||||
### Evaluation
|
||||
All evaluation scripts for reproducing the main results in the paper are included in [eval](eval/) directory.
|
||||
|
|
@ -2,7 +2,7 @@
|
|||
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --eval_batch_size_gen 1
|
||||
|
||||
# w/ truncation
|
||||
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --eval_batch_size_gen 1 --truncate_if_too_long_inp
|
||||
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --eval_batch_size_gen 1 --truncate_if_too_long_inp
|
||||
|
||||
# no context
|
||||
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --eval_batch_size_gen 1 --remove_context
|
||||
|
|
|
|||
8
scripts/main_exp/eval/base_model_test.sh
Executable file
8
scripts/main_exp/eval/base_model_test.sh
Executable file
|
|
@ -0,0 +1,8 @@
|
|||
# no truncation
|
||||
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --eval_batch_size_gen 1 --max_test_samples_per_ds 10
|
||||
|
||||
# w/ truncation
|
||||
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --eval_batch_size_gen 1 --truncate_if_too_long_inp --max_test_samples_per_ds 10
|
||||
|
||||
# no context
|
||||
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --eval_batch_size_gen 1 --remove_context --max_test_samples_per_ds 10
|
||||
|
|
@ -1,6 +1,6 @@
|
|||
# qa
|
||||
uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp --cd_use_gen_q --q_gen_rounds=4
|
||||
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp --cd_use_gen_q --q_gen_rounds=4
|
||||
|
||||
|
||||
# longbench
|
||||
uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/multifieldqa_en_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp --cd_use_gen_q --q_gen_rounds=1
|
||||
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/multifieldqa_en_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp --cd_use_gen_q --q_gen_rounds=1
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
# qa
|
||||
uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp
|
||||
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp
|
||||
|
||||
|
||||
# longbench
|
||||
uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/multifieldqa_en_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp
|
||||
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/multifieldqa_en_e longbench/2wikimqa_e longbench/qasper_e --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp
|
||||
|
|
|
|||
6
scripts/main_exp/eval/cd_oracle_test.sh
Executable file
6
scripts/main_exp/eval/cd_oracle_test.sh
Executable file
|
|
@ -0,0 +1,6 @@
|
|||
# qa
|
||||
uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp --max_test_samples_per_ds 10
|
||||
|
||||
|
||||
# longbench
|
||||
uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/multifieldqa_en_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp --max_test_samples_per_ds 10
|
||||
6
scripts/main_exp/eval/cd_test.sh
Executable file
6
scripts/main_exp/eval/cd_test.sh
Executable file
|
|
@ -0,0 +1,6 @@
|
|||
# qa
|
||||
uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp --cd_use_gen_q --q_gen_rounds=4 --max_test_samples_per_ds 10
|
||||
|
||||
|
||||
# longbench
|
||||
uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/multifieldqa_en_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp --cd_use_gen_q --q_gen_rounds=1 --max_test_samples_per_ds 10
|
||||
|
|
@ -1,13 +1,13 @@
|
|||
# main results
|
||||
# batched
|
||||
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/pytorch_model.bin --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --max_ctx_chunk_len 8192 --eval_batch_size_gen 1
|
||||
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/checkpoint-$step/pytorch_model.bin --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --max_ctx_chunk_len 8192 --eval_batch_size_gen 1
|
||||
|
||||
# iterative
|
||||
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/pytorch_model.bin --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --max_ctx_chunk_len 8192 --eval_batch_size_gen 1 --use_iterative_mode
|
||||
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/checkpoint-$step/pytorch_model.bin --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --max_ctx_chunk_len 8192 --eval_batch_size_gen 1 --use_iterative_mode
|
||||
|
||||
# query internalization
|
||||
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/pytorch_model.bin --datasets squad --split test --eval_batch_size_gen=1 --flip_ctx_inp
|
||||
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/checkpoint-$step/pytorch_model.bin --datasets squad --split test --eval_batch_size_gen=1 --flip_ctx_inp
|
||||
|
||||
# replaced squad context
|
||||
WANDB_MODE=disabled uv run python run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/pytorch_model.bin --datasets squad_assistant_ctx_no_passage --split test
|
||||
WANDB_MODE=disabled uv run python run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/pytorch_model.bin --datasets squad_negative_no_passage --split test
|
||||
WANDB_MODE=disabled uv run python run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/checkpoint-$step/pytorch_model.bin --datasets squad_assistant_ctx_no_passage --split test
|
||||
WANDB_MODE=disabled uv run python run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/checkpoint-$step/pytorch_model.bin --datasets squad_negative_no_passage --split test
|
||||
|
|
|
|||
13
scripts/main_exp/eval/d2l_test.sh
Executable file
13
scripts/main_exp/eval/d2l_test.sh
Executable file
|
|
@ -0,0 +1,13 @@
|
|||
# main results
|
||||
# batched
|
||||
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path ../../ctx-to-lora/train_outputs/runs/Sep29_14-42-46_slurm0-a3nodeset-9_88483_1e7bb34e/checkpoint-40000/pytorch_model.bin --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --max_ctx_chunk_len 8192 --eval_batch_size_gen 1 --max_test_samples_per_ds 10
|
||||
|
||||
# iterative
|
||||
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path ../../ctx-to-lora/train_outputs/runs/Sep29_14-42-46_slurm0-a3nodeset-9_88483_1e7bb34e/checkpoint-40000/pytorch_model.bin --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --max_ctx_chunk_len 8192 --eval_batch_size_gen 1 --use_iterative_mode --max_test_samples_per_ds 10
|
||||
|
||||
# query internalization
|
||||
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path ../../ctx-to-lora/train_outputs/runs/Sep29_14-42-46_slurm0-a3nodeset-9_88483_1e7bb34e/checkpoint-40000/pytorch_model.bin --datasets squad --split test --eval_batch_size_gen=1 --flip_ctx_inp --max_test_samples_per_ds 10
|
||||
|
||||
# replaced squad context
|
||||
WANDB_MODE=disabled uv run python run_eval.py --checkpoint_path ../../ctx-to-lora/train_outputs/runs/Sep29_14-42-46_slurm0-a3nodeset-9_88483_1e7bb34e/checkpoint-40000/pytorch_model.bin --datasets squad_assistant_ctx_no_passage --split test --max_test_samples_per_ds 10
|
||||
WANDB_MODE=disabled uv run python run_eval.py --checkpoint_path ../../ctx-to-lora/train_outputs/runs/Sep29_14-42-46_slurm0-a3nodeset-9_88483_1e7bb34e/checkpoint-40000/pytorch_model.bin --datasets squad_negative_no_passage --split test --max_test_samples_per_ds 10
|
||||
5
scripts/main_exp/eval/llmlingua_test.sh
Executable file
5
scripts/main_exp/eval/llmlingua_test.sh
Executable file
|
|
@ -0,0 +1,5 @@
|
|||
for dataset in squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e; do
|
||||
for rate in 0.9 0.8 0.6 0.4 0.2 0.1; do
|
||||
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets "$dataset" --split test --eval_batch_size_gen=1 --use_llmlingua --llmlingua_compression_rate "$rate" --truncate_if_too_long_ctx --max_test_samples_per_ds 10
|
||||
done
|
||||
done
|
||||
4
scripts/main_exp/eval/t2l_test.sh
Executable file
4
scripts/main_exp/eval/t2l_test.sh
Executable file
|
|
@ -0,0 +1,4 @@
|
|||
# download t2l checkpoint
|
||||
uv run huggingface-cli download SakanaAI/text-to-lora --local-dir . --include "trained_t2l/gemma_2b_t2l"
|
||||
|
||||
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --eval_batch_size_gen=1 --use_t2l --max_test_samples_per_ds 10
|
||||
|
|
@ -4,8 +4,8 @@ uv run data/download_fineweb_edu.py
|
|||
# generate qa data
|
||||
# run from 000 to 013
|
||||
for shard_id in $(seq -f "%03g" 0 13); do
|
||||
uv run data/generate_fw_edu_qa_v2.py --shard_pattern "${shard_id}_00000" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it --max_length=2000 --max_model_length=2048;
|
||||
uv run data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/${shard_id}*level_0*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it;
|
||||
uv run data/generate_fw_edu_qa_v2.py --shard_pattern "${shard_id}_00000" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it --max_length=2000 --max_model_length=2048
|
||||
uv run data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/${shard_id}*level_0" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it
|
||||
|
||||
# self-generated response QA data
|
||||
uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_pattern "data/raw_datasets/fw_qa_v2/min_0_to_2000/${shard_id}*_level_1*" --closed_qa_prob 1.0
|
||||
|
|
|
|||
|
|
@ -4,17 +4,17 @@ uv run data/download_fineweb_edu.py
|
|||
# generate qa data
|
||||
# run from 000 to 013
|
||||
for shard_id in $(seq -f "%03g" 0 1); do
|
||||
uv run data/generate_fw_edu_qa_v2.py --shard_pattern "${shard_id}_00000" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it --max_length=2000 --max_model_length=2048;
|
||||
uv run data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/${shard_id}*level_0*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it;
|
||||
uv run data/generate_fw_edu_qa_v2.py --shard_pattern "${shard_id}_00000" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it --max_length=2000 --max_model_length=2048 --debug
|
||||
uv run data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/${shard_id}*level_0" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it --debug
|
||||
|
||||
# self-generated response QA data
|
||||
uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_pattern "data/raw_datasets/fw_qa_v2/min_0_to_2000/${shard_id}*_level_1*" --closed_qa_prob 1.0
|
||||
uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_pattern "data/raw_datasets/fw_qa_v2/min_0_to_2000/${shard_id}*_level_1*" --closed_qa_prob 1.0 --debug
|
||||
done
|
||||
|
||||
|
||||
# val split
|
||||
uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_pattern 'data/raw_datasets/fw_qa_v2/min_0_to_2000/*_level_0_val.parquet'
|
||||
uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_pattern 'data/raw_datasets/fw_qa_v2/min_0_to_2000/*_level_0_val.parquet' --debug
|
||||
|
||||
# self-gen data for other ds
|
||||
uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --ds_names squad_compact ropes_compact drop_compact --closed_qa_prob 1.0
|
||||
uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --ds_names pwc_compact --closed_qa_prob 0.0
|
||||
uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --ds_names squad_compact ropes_compact drop_compact --split train --closed_qa_prob 1.0 --debug
|
||||
uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --ds_names pwc_compact --split train --closed_qa_prob 0.0 --debug
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ from collections.abc import Callable
|
|||
from typing import Any
|
||||
|
||||
from ctx_to_lora.data.definitions import CLOSED_QA_INTX_TEMPLATES, EVAL_INTX_TEMPLATES
|
||||
from ctx_to_lora.utils import concat_list
|
||||
|
||||
logger = logging.getLogger()
|
||||
|
||||
|
|
@ -42,6 +43,29 @@ def get_preprocessing_fn(
|
|||
# already processed data, do nothing
|
||||
return f
|
||||
|
||||
if "fw_qa_v2" in ds_name:
|
||||
|
||||
def f(sample):
|
||||
# get questions/answers from all levels in the ds
|
||||
q_cols = [col for col in sample.keys() if col.startswith("prompts_level")]
|
||||
r_cols = [col for col in sample.keys() if col.startswith("responses_level")]
|
||||
questions = concat_list([sample[col] for col in q_cols])
|
||||
responses = concat_list([sample[col] for col in r_cols])
|
||||
min_len = min(len(questions), len(responses))
|
||||
|
||||
if min_len == 0:
|
||||
return {
|
||||
"context": None,
|
||||
"prompts": None,
|
||||
"responses": None,
|
||||
}
|
||||
|
||||
return {
|
||||
"context": sample["context"],
|
||||
"prompts": questions[:min_len],
|
||||
"responses": responses[:min_len],
|
||||
}
|
||||
|
||||
elif ds_name.startswith("longbench"):
|
||||
|
||||
def f(sample):
|
||||
|
|
|
|||
|
|
@ -941,9 +941,6 @@ def evaluate(
|
|||
|
||||
collator = generation_collator if generative else eval_collator
|
||||
|
||||
if max_ctx_chunk_len > 0:
|
||||
model.generate = model.generate_with_multi_loras
|
||||
|
||||
# if isinstance(model, CtxDistillModel):
|
||||
|
||||
# model.generate = partial(
|
||||
|
|
|
|||
|
|
@ -31,8 +31,10 @@ from ctx_to_lora.configs import (
|
|||
CtxEncoderArguments,
|
||||
HypernetArguments,
|
||||
)
|
||||
from ctx_to_lora.data.processing import tokenize_ctx_text
|
||||
from ctx_to_lora.model_loading import (
|
||||
get_model,
|
||||
get_tokenizer,
|
||||
)
|
||||
from ctx_to_lora.modeling.aggregator import (
|
||||
AGGREGATOR_CLS,
|
||||
|
|
@ -458,7 +460,7 @@ class ModulatedPretrainedModel(nn.Module):
|
|||
self.user_defined_scaling = user_defined_scaling
|
||||
self.inp_compressor = inp_compressor
|
||||
self.model_accepts_loss_kwargs = True
|
||||
self.active_adapters = []
|
||||
self.generated_loras = None
|
||||
|
||||
self.register_module("base_model", base_model)
|
||||
self._init_model()
|
||||
|
|
@ -496,6 +498,28 @@ class ModulatedPretrainedModel(nn.Module):
|
|||
model.load_state_dict(state_dict)
|
||||
return model
|
||||
|
||||
def patch_lora_forward(self):
|
||||
layers = get_layers(self.base_model)
|
||||
|
||||
lora_forward_fn = (
|
||||
lora_forward_packed if self.use_sequence_packing else lora_forward
|
||||
)
|
||||
for layer_idx in self.hypernet.layer_indices:
|
||||
for module_info in get_peft_modules(layers[layer_idx], self.peft_config):
|
||||
name = module_info["name"]
|
||||
module = module_info["module"]
|
||||
if getattr(module, "patched_forward", False):
|
||||
continue
|
||||
logger.debug(f"Applying LoRA forward to {name}")
|
||||
module.forward_orig = module.forward
|
||||
module.patched_forward = True
|
||||
module.forward = partial(
|
||||
lora_forward_fn,
|
||||
self=module,
|
||||
lora_dropout_p=self.peft_config.lora_dropout,
|
||||
scaling=self.peft_config.lora_alpha,
|
||||
)
|
||||
|
||||
def _init_model(self):
|
||||
# disable adapter of the base model
|
||||
# this only works with LoRA(?)
|
||||
|
|
@ -506,21 +530,7 @@ class ModulatedPretrainedModel(nn.Module):
|
|||
HyperLoRA(self.hypernet_config).to(self.device).to(torch.float32)
|
||||
)
|
||||
|
||||
layers = get_layers(self.base_model)
|
||||
lora_forward_fn = (
|
||||
lora_forward_packed if self.use_sequence_packing else lora_forward
|
||||
)
|
||||
for layer_idx in self.hypernet.layer_indices:
|
||||
for module_info in get_peft_modules(layers[layer_idx], self.peft_config):
|
||||
name = module_info["name"]
|
||||
module = module_info["module"]
|
||||
logger.debug(f"Applying LoRA forward to {name}")
|
||||
module.forward = partial(
|
||||
lora_forward_fn,
|
||||
self=module,
|
||||
lora_dropout_p=self.peft_config.lora_dropout,
|
||||
scaling=self.peft_config.lora_alpha,
|
||||
)
|
||||
self.patch_lora_forward()
|
||||
|
||||
ctx_model_name = self.ctx_encoder_args.ctx_encoder_model_name_or_path
|
||||
if ctx_model_name is None:
|
||||
|
|
@ -762,58 +772,6 @@ class ModulatedPretrainedModel(nn.Module):
|
|||
else:
|
||||
return model_outputs
|
||||
|
||||
@torch.inference_mode()
|
||||
def generate_with_multi_loras(
|
||||
self,
|
||||
ctx_ids: Integer[Tensor, "bs ctx_length"],
|
||||
ctx_attn_mask: Integer[Tensor, "bs ctx_length"] | None = None,
|
||||
ctx_position_ids: Integer[Tensor, "bs ctx_length"] | None = None,
|
||||
n_ctx_chunks: Integer[Tensor, "n_ctx"] | None = None,
|
||||
n_queries: Integer[Tensor, "n_ctx"] | None = None,
|
||||
scalers: Float[Tensor, "n_ctx"] | None = None,
|
||||
bias_scaler: float | None = None,
|
||||
*model_inputs_args: Any,
|
||||
**model_inputs_kwargs: dict[str, Any],
|
||||
):
|
||||
generated_loras, _ = self.generate_weights(
|
||||
ctx_ids, ctx_attn_mask, ctx_position_ids
|
||||
)
|
||||
|
||||
generated_loras = combine_lora(
|
||||
generated_loras,
|
||||
n_ctx_chunks,
|
||||
lora_bias=self.hypernet.get_head_bias()
|
||||
if self.hypernet.config.use_bias
|
||||
else None,
|
||||
scalers=scalers,
|
||||
bias_scaler=bias_scaler,
|
||||
)
|
||||
|
||||
# apply lora hook to the base model
|
||||
position_ids = (
|
||||
model_inputs_kwargs["position_ids"]
|
||||
if "position_ids" in model_inputs_kwargs
|
||||
else None
|
||||
)
|
||||
if n_queries is None:
|
||||
# assumes each group has one q
|
||||
n_queries = torch.ones(
|
||||
len(n_ctx_chunks), dtype=torch.int32, device=self.device
|
||||
)
|
||||
|
||||
apply_lora_to_layers(
|
||||
self.base_model,
|
||||
self.hypernet.layer_indices,
|
||||
generated_loras,
|
||||
n_queries,
|
||||
position_ids, # TODO: remove (cant be used with generation?)
|
||||
)
|
||||
|
||||
model_outputs = self.base_model.generate(
|
||||
*model_inputs_args, **model_inputs_kwargs
|
||||
)
|
||||
return model_outputs
|
||||
|
||||
def combine_lora(self, *args, **kwargs):
|
||||
# for timing
|
||||
return combine_lora(*args, **kwargs)
|
||||
|
|
@ -822,6 +780,38 @@ class ModulatedPretrainedModel(nn.Module):
|
|||
# for timing
|
||||
return apply_lora_to_layers(*args, **kwargs)
|
||||
|
||||
# for simple api usage
|
||||
def internalize(self, ctx_str: str):
|
||||
ctx_tokenizer = get_tokenizer(self.ctx_encoder.base_model.name_or_path)
|
||||
ctx_ids = tokenize_ctx_text(dict(context=[ctx_str]), ctx_tokenizer)["ctx_ids"]
|
||||
return self._internalize_from_ids(torch.tensor(ctx_ids, device=self.device))
|
||||
|
||||
def _internalize_from_ids(
|
||||
self,
|
||||
ctx_ids: Integer[Tensor, "n_ctx ctx_len"] | None = None,
|
||||
ctx_attn_mask: Integer[Tensor, "n_ctx ctx_len"] | None = None,
|
||||
ctx_position_ids: Integer[Tensor, "n_ctx ctx_len"] | None = None,
|
||||
):
|
||||
self.patch_lora_forward()
|
||||
if ctx_attn_mask is None and ctx_position_ids is None:
|
||||
assert ctx_ids.shape[0] == 1
|
||||
ctx_attn_mask = torch.ones_like(ctx_ids)
|
||||
generated_loras, generated_layernorms = self.generate_weights(
|
||||
ctx_ids, ctx_attn_mask, ctx_position_ids
|
||||
)
|
||||
self.generated_loras = generated_loras
|
||||
|
||||
def reset(self):
|
||||
self.generated_loras = None
|
||||
layers = get_layers(self.base_model)
|
||||
for layer_idx in self.hypernet.layer_indices:
|
||||
for module_info in get_peft_modules(layers[layer_idx], self.peft_config):
|
||||
name = module_info["name"]
|
||||
module = module_info["module"]
|
||||
logger.debug(f"Resetting forward for {name}")
|
||||
module.forward = module.forward_orig
|
||||
module.patched_forward = False
|
||||
|
||||
@torch.inference_mode()
|
||||
def generate(
|
||||
self,
|
||||
|
|
@ -837,23 +827,22 @@ class ModulatedPretrainedModel(nn.Module):
|
|||
generated_layernorms = None
|
||||
if (
|
||||
ctx_ids is None
|
||||
and not self.active_adapters
|
||||
and not self.generated_loras
|
||||
and not self.use_base_input_as_ctx
|
||||
):
|
||||
logger.warning(
|
||||
(
|
||||
"*" * 100,
|
||||
"\n\nNo ctx_ids provided, using the base model for generation\n\n",
|
||||
"*" * 100,
|
||||
)
|
||||
print(
|
||||
"*" * 100
|
||||
+ "\n\nNo ctx_ids provided, using the base model for generation\n\n"
|
||||
+ "*" * 100
|
||||
)
|
||||
if ctx_ids is None and self.active_adapters:
|
||||
logger.info(
|
||||
(
|
||||
"*" * 100,
|
||||
"\n\nUsing active LoRAs for generation\n\n",
|
||||
"*" * 100,
|
||||
)
|
||||
elif ctx_ids is None and self.generated_loras:
|
||||
generated_loras = self.generated_loras
|
||||
if n_ctx_chunks is None:
|
||||
n_ctx_chunks = torch.tensor((1,), device=self.device)
|
||||
print(
|
||||
"*" * 100
|
||||
+ "\n\nUsing internalized LoRAs for generation\n\n"
|
||||
+ "*" * 100
|
||||
)
|
||||
else:
|
||||
if self.use_base_input_as_ctx:
|
||||
|
|
@ -895,7 +884,9 @@ class ModulatedPretrainedModel(nn.Module):
|
|||
if n_queries is None:
|
||||
if ctx_position_ids is None:
|
||||
n_queries = torch.ones(
|
||||
ctx_ids.shape[0], dtype=torch.int32, device=self.device
|
||||
model_inputs_kwargs["input_ids"].shape[0],
|
||||
dtype=torch.int32,
|
||||
device=self.device,
|
||||
)
|
||||
else:
|
||||
# quite redundant (we do cu_seqlens many places)
|
||||
|
|
@ -913,19 +904,7 @@ 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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -675,7 +675,7 @@ def chat():
|
|||
# do_sample=False,
|
||||
# )
|
||||
|
||||
outputs = modulated_model.generate_with_multi_loras(
|
||||
outputs = modulated_model.generate(
|
||||
ctx_ids=ctx_ids,
|
||||
ctx_attn_mask=ctx_attn_mask,
|
||||
n_ctx_chunks=torch.tensor(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue