mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
133 lines
4.6 KiB
Python
133 lines
4.6 KiB
Python
import os
|
|
from functools import partial
|
|
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
|
|
from ctx_to_lora.data.processing import get_tokenized_dataset, load_and_process_dataset
|
|
from ctx_to_lora.model_loading import get_tokenizer
|
|
|
|
tokenizer_kwargs = {"max_length": 2**13}
|
|
ctx_tokenizer_kwargs = {"max_length": 2**13} # not used for now
|
|
tokenizer = get_tokenizer("google/gemma-2-2b-it", train=True)
|
|
ctx_tokenizer = get_tokenizer("google/gemma-2-2b-it", train=True)
|
|
|
|
_get_tokenized_dataset = partial(
|
|
get_tokenized_dataset,
|
|
base_model_max_len=2**13,
|
|
tokenizer=tokenizer,
|
|
tokenizer_kwargs=tokenizer_kwargs,
|
|
ctx_model_max_len=2**13,
|
|
ctx_tokenizer=ctx_tokenizer,
|
|
ctx_tokenizer_kwargs=ctx_tokenizer_kwargs,
|
|
add_ctx_to_chat=False,
|
|
add_repeat_prompt=False,
|
|
add_negative_prompt=False,
|
|
use_kl_loss=False,
|
|
# streaming=data_args.streaming,
|
|
)
|
|
_load_and_process_dataset = partial(
|
|
load_and_process_dataset,
|
|
add_negative_prompt=False,
|
|
add_repeat_prompt=False,
|
|
repeat_prob=0,
|
|
is_pretrain=False,
|
|
streaming=False,
|
|
num_proc=8,
|
|
)
|
|
|
|
train_ds_names = [
|
|
# "fw_qa",
|
|
# "fw_qa_large",
|
|
# "ctx_qa",
|
|
# "pwc",
|
|
# "hotpot_qa",
|
|
# "squad",
|
|
# "drop",
|
|
# "narrativeqa",
|
|
# "quoref",
|
|
# "ropes",
|
|
# "synthetic_convqa",
|
|
# "booksum",
|
|
# "gov_report",
|
|
# "wikitext-2",
|
|
# "wikitext-103",
|
|
# "fw_qa_2"
|
|
# "fw_qa_xl",
|
|
# "fw_qa_3_medium",
|
|
"self_gen/google/gemma-2-2b-it_temp_0.0/fw_qa_v2_2k_len_level_0_tiny",
|
|
"self_gen/google/gemma-2-2b-it_temp_0.0/fw_qa_v2_2k_len_level_1_tiny",
|
|
"self_gen/google/gemma-2-2b-it_temp_0.0/fw_qa_v2_2k_len_level_2_tiny",
|
|
"self_gen/google/gemma-2-2b-it_temp_0.0/fw_qa_v2_2k_len_level_3_tiny",
|
|
]
|
|
loaded_dataset = dict()
|
|
val_ds_names = [] # ["fw_qa_large", "ctx_qa", "pwc", "hotpot_qa", "squad"]
|
|
for split, ds_names in zip(["train", "validation"], [train_ds_names, val_ds_names]):
|
|
loaded_dataset[split] = {
|
|
os.path.basename(ds_name): _load_and_process_dataset(ds_name, split)
|
|
for ds_name in ds_names
|
|
}
|
|
for ds_name in loaded_dataset[split]:
|
|
print(ds_name)
|
|
print(f"Num samples: {len(loaded_dataset[split][ds_name])}")
|
|
ds = loaded_dataset[split][ds_name]
|
|
sample = loaded_dataset[split][ds_name][0]
|
|
|
|
# TODO: let's investigate samples that are very long and very short
|
|
if "responses" in sample:
|
|
char_lens = []
|
|
responses = []
|
|
questions = []
|
|
ctxs = []
|
|
for sample in loaded_dataset[split][ds_name]:
|
|
char_lens += [len(res) for res in sample["responses"]]
|
|
responses += sample["responses"]
|
|
questions += sample["prompts"]
|
|
ctxs += [sample["context"]] * len(sample["responses"])
|
|
else:
|
|
responses = ds["response"]
|
|
questions = ds["prompt"]
|
|
ctxs = ds["context"]
|
|
|
|
char_lens = [
|
|
len(sample["response"]) for sample in loaded_dataset[split][ds_name]
|
|
]
|
|
|
|
print(f"Num avg. char_lens: {sum(char_lens) / len(char_lens)}")
|
|
print(f"Num max. char_lens: {max(char_lens)}")
|
|
print(f"Num min. char_lens: {min(char_lens)}")
|
|
print(f"Num std. char_lens: {np.std(char_lens)}")
|
|
plt.hist(char_lens, bins=100)
|
|
plt.savefig(f"tmp/n_token_hist_{split}_{ds_name}.png")
|
|
plt.close()
|
|
|
|
long_indices = [i for i, clen in enumerate(char_lens) if clen > 2000]
|
|
short_indices = [i for i, clen in enumerate(char_lens) if clen < 100]
|
|
breakpoint()
|
|
for i in long_indices:
|
|
print(f"Long sample {i}")
|
|
print(f"Response length: {len(responses[i])}")
|
|
print(f"Context: {ctxs[i]}")
|
|
print(f"Question: {questions[i]}")
|
|
print(f"Response: {responses[i]}")
|
|
print("=" * 80)
|
|
|
|
breakpoint()
|
|
for i in short_indices:
|
|
print(f"Short sample {i}")
|
|
print(f"Response length: {len(responses[i])}")
|
|
print(f"Context: {ctxs[i]}")
|
|
print(f"Question: {questions[i]}")
|
|
print(f"Response: {responses[i]}")
|
|
print("=" * 80)
|
|
|
|
ctx_tokens = [
|
|
len(sample["ctx_ids"]) for sample in loaded_dataset[split][ds_name]
|
|
]
|
|
print(f"Num avg. ctx char_lens: {sum(ctx_tokens) / len(ctx_tokens)}")
|
|
print(f"Num max. ctx char_lens: {max(ctx_tokens)}")
|
|
print(f"Num min. ctx char_lens: {min(ctx_tokens)}")
|
|
print(f"Num std. ctx char_lens: {np.std(ctx_tokens)}")
|
|
plt.hist(ctx_tokens, bins=100)
|
|
plt.savefig(f"tmp/ctx_n_token_hist_{split}_{ds_name}.png")
|
|
plt.close()
|