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

webui chat + result display + eval datasets + preprocessing eval + refactoring data_utils.py to data_definition.py

Browse source
This commit is contained in:
51616 2025-05-20 06:06:20 +00:00
parent cdfd42cf43
commit 0e4c1444f9
11 changed files with 891 additions and 424 deletions
Download patch file Download diff file Expand all files Collapse all files

2
.gitignore vendored
View file

@ -23,3 +23,5 @@ plots/
watcher_state.yaml
ds_config.json
eval_results/
.github/
*.code-workspace

2
README.md
View file

@ -51,7 +51,7 @@ run python intx_sft.py configs/...yaml ... --from_pretrained_checkpoint=train_ou
LongBench
```bash
# generative
run python src/ctx_to_lora/eval.py --checkpoint_path train_outputs/runs/May08_13-56-31_slurm0-a3nodeset-5_59383_906acb28/checkpoint-105000/pytorch_model.bin
run python src/ctx_to_lora/eval.py --checkpoint_path train_outputs/runs/May08_13-56-31_slurm0-a3nodeset-5_59383_906acb28/checkpoint-105000/pytorch_model.bin
# benchmark
cd LongBench/LongBench

15
intx_sft.py
View file

@ -308,7 +308,7 @@ def main():
model, tokenizer = get_model_and_tokenizer(
**vars(model_args),
train=True,
requires_grad=ctx_args.exp_setup == ExperimentSetup.FULL_FINETUNE,
requires_grad=False, # ctx_args.exp_setup == ExperimentSetup.FULL_FINETUNE,
peft_config=get_lora_config(model_name, **vars(lora_args)),
)
ctx_name = ctx_encoder_args.ctx_encoder_model_name_or_path
@ -324,7 +324,7 @@ def main():
ctx_encoder_model_config = model.config
ctx_tokenizer = tokenizer
if ctx_args.exp_setup == ExperimentSetup.HYPER_LORA:
if ctx_args.exp_setup == ExperimentSetup.HYPERLORA:
# TODO: handle only extra_modules case (no target_modules)
logger.info("Using HyperLoRA")
if not ctx_args.from_pretrained_checkpoint:
@ -384,15 +384,17 @@ def main():
logger.info("Loading dataset...")
add_ctx_to_chat = not isinstance(model, ModulatedPretrainedModel)
tokenizer_kwargs = {"max_length": ctx_args.max_base_len}
ctx_tokenizer_kwargs = {"max_length": ctx_args.max_ctx_len} # not used for now
# tokenizer_kwargs = {"max_length": ctx_args.max_base_len}
# ctx_tokenizer_kwargs = {"max_length": ctx_args.max_ctx_len} # not used for now
_get_tokenized_dataset = partial(
get_tokenized_dataset,
base_model_max_len=model.base_model.config.max_position_embeddings,
tokenizer=tokenizer,
tokenizer_kwargs=tokenizer_kwargs,
tokenizer_kwargs={},
ctx_model_max_len=model.ctx_encoder.config.max_position_embeddings,
ctx_tokenizer=ctx_tokenizer,
ctx_tokenizer_kwargs=ctx_tokenizer_kwargs,
ctx_tokenizer_kwargs={},
add_ctx_to_chat=add_ctx_to_chat,
add_repeat_prompt=ctx_args.add_repeat_prompt,
add_negative_prompt=ctx_args.add_negative_prompt,
@ -415,6 +417,7 @@ def main():
train_ds = tokenized_ds["train"]
logging.info(f"train_ds: {train_ds}")
val_ds = dict()
if "validation" in tokenized_ds:
n_val_samples = data_args.max_val_samples_per_ds

34
scripts/fw_qa_3/medium_gemma_per_rank_fac2_per_layer_4_gpus_gen_lora_reg.sh Normal file
View file

@ -0,0 +1,34 @@
#!/bin/bash
#SBATCH --job-name=ctxlora_medium
#SBATCH --partition=a3
#SBATCH --nodes=1
#SBATCH --gpus=4
#SBATCH --output=outputs/%x-%j.out
#SBATCH --error=outputs/%x-%j.out
# module load
# module load cuda/12.1
# module load cudnn/8.9.7
# module load nccl/cuda-12.1/2.18.3
# module load hpcx/2.20
# export OMP_NUM_THREADS=24
# export TRITON_CACHE_DIR=/tmp/.triton/
. ~/miniconda3/etc/profile.d/conda.sh
conda activate /home/rujikorn_sakana_ai/.conda/envs/ctx-to-lora
# eval "$@"
accelerate launch --num_processes=4 --gradient_accumulation_steps=8 --gradient_clipping=1.0 \
--gpu_ids all --main_process_port 29564 intx_sft.py configs/pretrain_all_xl_3_medium.yaml \
--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=32 \
--gradient_accumulation_steps=8 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \
--target_modules=down_proj \
--num_self_attends_per_block=8 --num_latent_factor=2 \
--lora_r=8 \
--eval_steps=5000 --save_steps=5000 --learning_rate=4e-5 --lora_dropout=0.0 \
--neftune_noise_alpha=5 --use_light_weight_lora=False \
--load_best_model_at_end=True --metric_for_best_model=pwc_loss --add_negative_prompt=False \
--add_repeat_prompt=False \
--use_sequence_packing=True --per_rank_gen=True \
--per_layer_processing=True \
--gen_lora_l1_reg_coef=0.1

8
src/ctx_to_lora/configs.py
View file

@ -122,9 +122,9 @@ class ArgumentParser(HfArgumentParser):
class ExperimentSetup(str, Enum):
LORA = "lora"
HYPER_LORA = "hyper_lora"
FULL_FINETUNE = "full_finetune"
# LORA = "lora"
HYPERLORA = "hyper_lora"
# FULL_FINETUNE = "full_finetune"
@dataclass
@ -277,7 +277,7 @@ class LoRAArguments:
@dataclass
class CtxTrainingArguments:
exp_setup: ExperimentSetup = field(
default=ExperimentSetup.LORA,
default=ExperimentSetup.HYPERLORA,
metadata={"help": "Experiment setup - LoRA, HyperLoRA, or full finetuning"},
)
from_pretrained_checkpoint: str = field(

365
src/ctx_to_lora/data_definitions.py Normal file
View file

@ -0,0 +1,365 @@
from glob import glob
IGNORE_INDEX = -100
FW_QA_PATHS = [
f"data/raw_datasets/fw_qa/{i:05d}.parquet" for i in [0, 1, 6, 7, 8, 10, 22, 30, 35]
]
TRANSFORMED_DATA_DIR = "data/processed_datasets"
# approximate length of the datasets (train split)
# needed for streaming datasets
DS_LEN = {
"fw_qa_3_mini": 100_000,
"fw_qa_3_medium": 121_000_000,
"fw_qa_3": 270_000_000,
"fw_qa_xl": 27_000_000,
"ctx_qa": 300_000,
"pwc": 240_000,
"hotpot_qa": 100_000,
"squad": 90_000,
"drop": 77_000,
"narrativeqa": 40_000,
"quoref": 11_000,
"ropes": 11_000,
"synthetic_convqa": 40_000,
}
LONGBENCH_TASKS = [
"longbench/narrativeqa",
"longbench/qasper",
"longbench/multifieldqa_en",
"longbench/multifieldqa_zh",
"longbench/hotpotqa",
"longbench/2wikimqa",
"longbench/musique",
"longbench/dureader",
"longbench/gov_report",
"longbench/qmsum",
"longbench/multi_news",
"longbench/vcsum",
]
LONGBENCH_E_TASKS = [
"longbench/qasper_e",
"longbench/multifieldqa_en_e",
"longbench/hotpotqa_e",
"longbench/2wikimqa_e",
"longbench/gov_report_e",
"longbench/multi_news_e",
]
DS_KWARGS = {
"hotpot_qa": dict(
train=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="train[900:]"),
validation=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="train[:900]"),
test=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="validation"),
),
"hotpot_qa_tiny": dict(
train=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="train[900:2000]"),
validation=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="train[:900]"),
),
"pwc": dict(
train=dict(path="sggetao/PwC", split="train[900:]"),
validation=dict(path="sggetao/PwC", split="train[:900]"),
test=dict(path="sggetao/PwC", split="test"),
),
"pwc_tiny": dict(
train=dict(path="sggetao/PwC", split="train[900:2000]"),
validation=dict(path="sggetao/PwC", split="train[:900]"),
),
"squad": dict(
train=dict(path="rajpurkar/squad", split="train[900:]"),
validation=dict(path="rajpurkar/squad", split="train[:900]"),
test=dict(path="rajpurkar/squad", split="validation"),
),
# "fw_qa_tiny": dict(
# train=dict(
# path="parquet",
# data_files="data/raw_datasets/fw_qa/00000.parquet",
# split="train",
# ),
# validation=dict(
# path="parquet",
# data_files="data/raw_datasets/fw_qa/00000_val.parquet",
# split="train",
# ),
# ),
# "fw_qa": dict(
# train=dict(
# path="parquet",
# data_files=FW_QA_PATHS,
# split="train",
# ),
# validation=dict(
# path="parquet",
# data_files="data/raw_datasets/fw_qa/*_val.parquet",
# split="train",
# ),
# ),
# "fw_qa_large": dict(
# train=dict(
# path="parquet",
# data_files=glob("data/raw_datasets/fw_qa_large/*res.parquet"),
# split="train",
# ),
# validation=dict(
# path="parquet",
# data_files=glob("data/raw_datasets/fw_qa_large/*val.parquet"),
# split="train",
# ),
# ),
"fw_qa_xl": dict(
train=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_xl/*[!val].parquet"),
split="train",
),
validation=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_xl/*val.parquet"),
split="train",
),
# there is no explcit test for this ds
# used only for sanity check
test=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_xl/*val.parquet"),
split="train",
),
),
# "fw_qa_2": dict(
# train=dict(
# path="parquet",
# data_files=glob("data/raw_datasets/fw_qa_2/*[!val].parquet"),
# split="train",
# ),
# validation=dict(
# path="parquet",
# data_files=glob("data/raw_datasets/fw_qa_2/*val.parquet"),
# split="train",
# ),
# ),
"fw_qa_3_mini": dict(
train=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_3/000_00001.parquet"),
split="train[:100000]",
),
),
"fw_qa_3_medium": dict(
train=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_3/00[0-5]*[!val].parquet"),
split="train",
),
),
"fw_qa_3": dict(
train=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_3/*[!val].parquet"),
split="train",
),
validation=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_3/*val.parquet"),
split="train",
),
),
"ctx_qa": dict(
train=dict(
path="parquet",
data_files=glob("data/raw_datasets/ctx_qa/*res.parquet"),
split="train",
),
validation=dict(
path="parquet",
data_files=glob("data/raw_datasets/ctx_qa/*val.parquet"),
split="train",
),
# there is no explcit test for this ds
# used only for sanity check
test=dict(
path="parquet",
data_files=glob("data/raw_datasets/ctx_qa/*val.parquet"),
split="train",
),
),
"drop": dict(
train=dict(
path="ucinlp/drop",
split="train",
),
),
"narrativeqa": dict(
train=dict(
path="nvidia/ChatQA-Training-Data",
name="narrativeqa",
split="train",
),
),
"quoref": dict(
train=dict(
path="nvidia/ChatQA-Training-Data",
name="quoref",
split="train",
),
),
"ropes": dict(
train=dict(
path="nvidia/ChatQA-Training-Data",
name="ropes",
split="train",
),
),
"synthetic_convqa": dict(
train=dict(
path="nvidia/ChatQA-Training-Data",
name="synthetic_convqa",
split="train",
),
),
"booksum": dict(
train=dict(
path="kmfoda/booksum",
split="train",
),
validation=dict(
path="kmfoda/booksum",
split="validation",
),
test=dict(
path="kmfoda/booksum",
split="test",
),
),
"gov_report": dict(
train=dict(
path="ccdv/govreport-summarization",
split="train",
),
validation=dict(
path="ccdv/govreport-summarization",
split="validation",
),
test=dict(
path="ccdv/govreport-summarization",
split="test",
),
),
# "wikitext-2": dict(
# train=dict(
# path="EleutherAI/wikitext_document_level",
# name="wikitext-2-raw-v1",
# split="train",
# ),
# ),
# "wikitext-103": dict(
# train=dict(
# path="EleutherAI/wikitext_document_level",
# name="wikitext-103-raw-v1",
# split="train",
# ),
# ),
"gsm8k": dict(
train=dict(
path="openai/gsm8k",
name="main",
split="train[100:]",
),
validation=dict(
path="openai/gsm8k",
name="main",
split="train[:100]",
),
),
"openmathintx-2": dict(
train=dict(
path="nvidia/OpenMathInstruct-2",
split="train_1M[:100000]",
),
),
"opencoder-edu": dict(
train=dict(
path="OpenCoder-LLM/opc-sft-stage2",
name="educational_instruct",
split="train[900:]",
),
validation=dict(
path="OpenCoder-LLM/opc-sft-stage2",
name="educational_instruct",
split="train[:900]",
),
),
"triviaqa_retrieved": dict(
test=dict(
path="json",
data_files="data/eval/triviaqa/triviaqa_retrieved.jsonl",
split="train",
),
),
# TODO: for negative training and evaluating
"natural_questions": ...,
}
# LongBench kwargs
for ds_name in LONGBENCH_TASKS + LONGBENCH_E_TASKS:
DS_KWARGS[ds_name] = dict(
test=dict(
path="THUDM/LongBench",
name=ds_name.split("/")[-1],
split="test",
)
)
# for training closed qa datasets, e.g., hotpot_qa, squad, etc.
CLOSED_QA_INTX_TEMPLATES = [
"Answer the question based on the given passages. Only give me the answer and do not output any other words.\n\nQuestion: {input}",
"Answer without any explanation.\n\nQuestion: {input}",
"Based on the provided text, what is the answer to the following question? Provide only the answer.\n\nQuestion: {input}",
"Extract the answer to the question from the text. Be concise. Do not explain.\n\nQuestion: {input}",
"What is the answer to this question, based on the context? Respond with the answer only.\n\nQuestion: {input}",
"Provide a direct answer to the question using the given passages. Do not give any explanation.\n\nQuestion: {input}",
"Answer the question using only information from the provided text. No extra words.\n\nQuestion: {input}",
"From the passages, answer the question. Just the answer, please.\n\nQuestion: {input}",
"Give the answer to the question. Do not include any other text.\n\nQuestion: {input}",
"The answer to the question is in the text. Find it and state it clearly. No need for explanation.\n\nQuestion: {input}",
"Concisely answer the question based on the text provided. Don't include any other words. Just the answer.\n\nQuestion: {input}",
"Read the passages and answer the question with the minimal necessary words.\n\nQuestion: {input}",
"What is the direct response to the question, according to the text? Avoid explanation.\n\nQuestion: {input}",
"Please provide only the answer to the question, derived from the text.\n\nQuestion: {input}",
"Using the provided context, answer the question. Output the answer and nothing else.\n\nQuestion: {input}",
"Identify the answer in the text and present it without elaboration.\n\nQuestion: {input}",
"Answer the following question based on the text. Your answer should be brief and to the point. No explanation.\n\nQuestion: {input}",
"Based on the information given, what is the answer to the question? Only state the answer.\n\nQuestion: {input}",
"Find the answer to the question in the provided passages and write it down. No explanations.\n\nQuestion: {input}",
"The question is: {input}. Provide the answer based on the text, and nothing more.",
"Question: {input}\nAnswer directly based on the text provided. No extra words.",
"Question: {input}\nPlease provide the answer based on the text. No explanation is needed.",
]
EVAL_INTX_TEMPLATES = {
"triviaqa_retrieved": "Only give me the answer and do not output any other words.\n\nQuestion: {input}",
"hotpot_qa": "Answer the question based on the given passages. Only give me the answer and do not output any other words.\n\nQuestion: {input}",
"squad": "Answer the question based on the given passages. Only give me the answer and do not output any other words.\n\nQuestion: {input}",
"longbench/qasper": 'You are given a scientific article and a question. Answer the question as concisely as you can, using a single phrase or sentence if possible.\nIf the question cannot be answered based on the information in the article, write "unanswerable".\nIf the question is a yes/no question, answer "yes", "no", or "unanswerable". Do not provide any explanation.\n\nQuestion: {input}',
"longbench/narrativeqa": "Answer the question as concisely as you can, using a single phrase if possible. Do not provide any explanation.\n\nQuestion: {input}",
"longbench/multifieldqa_en": "Answer the following question based on the above text, only give me the answer and do not output any other words.\n\nQuestion: {input}",
"longbench/multifieldqa_zh": "阅读以下文字并用中文简短回答：\n\n现在请基于上面的文章回答下面的问题，只告诉我答案，不要输出任何其他字词。\n\n问题：{input}",
"longbench/hotpotqa": "Answer the question based on the given passages. Only give me the answer and do not output any other words.\n\nQuestion: {input}",
"longbench/2wikimqa": "Answer the question based on the given passages. Only give me the answer and do not output any other words.\n\nQuestion: {input}",
"longbench/musique": "Answer the question based on the given passages. Only give me the answer and do not output any other words.\n\nQuestion: {input}",
"longbench/dureader": "请基于给定的文章回答下述问题\n\n请基于上述文章回答下面的问题。\n\n问题：{input}",
"longbench/gov_report": "You are given a report by a government agency. Write a one-page summary of the report.",
"longbench/qmsum": "You are given a meeting transcript and a query containing a question or instruction. Answer the query based on the above meeting transcript.\n\nQuery: {input}",
"longbench/multi_news": "You are given several news passages. Write a one-page summary of all the news.",
"longbench/vcsum": "下面有一段会议记录，请你阅读后，写一段总结，总结会议的内容。",
}
for ds_name in LONGBENCH_E_TASKS:
EVAL_INTX_TEMPLATES[ds_name] = EVAL_INTX_TEMPLATES[ds_name[:-2]]
for ds_name in DS_KWARGS:
if ds_name not in EVAL_INTX_TEMPLATES:
EVAL_INTX_TEMPLATES[ds_name] = "{input}"

699
src/ctx_to_lora/data_utils.py
View file

@ -1,266 +1,236 @@
import logging
import numpy as np
import hashlib
import json
import random
from functools import partial
from os import path
from glob import glob
from typing import Any, Callable, Iterator, Optional
import datasets
import numpy as np
from datasets import load_dataset, IterableDataset
from transformers import PreTrainedTokenizerBase
from ctx_to_lora.training_utils import TRAINING_TASK
IGNORE_INDEX = -100
from ctx_to_lora.data_definitions import (
DS_KWARGS,
EVAL_INTX_TEMPLATES,
TRANSFORMED_DATA_DIR,
IGNORE_INDEX,
CLOSED_QA_INTX_TEMPLATES,
)
logger = logging.getLogger()
FW_QA_PATHS = [
f"data/raw_datasets/fw_qa/{i:05d}.parquet" for i in [0, 1, 6, 7, 8, 10, 22, 30, 35]
]
TRANSFORMED_DATA_DIR = "data/processed_datasets"
# TODO: add continued pre-training pipeline
# approximate length of the datasets
# needed for streaming datasets
DS_LEN = {
"fw_qa_3_mini": 100_000,
"fw_qa_3_medium": 121_000_000,
"fw_qa_3": 270_000_000,
"fw_qa_xl": 27_000_000,
"ctx_qa": 300_000,
"pwc": 240_000,
"hotpot_qa": 100_000,
"squad": 90_000,
"drop": 77_000,
"narrativeqa": 40_000,
"quoref": 11_000,
"ropes": 11_000,
"synthetic_convqa": 40_000,
}
DS_KWARGS = {
"hotpot_qa": dict(
train=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="train[900:]"),
validation=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="train[:900]"),
test=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="validation"),
),
"hotpot_qa_tiny": dict(
train=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="train[900:2000]"),
validation=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="train[:900]"),
),
"pwc": dict(
train=dict(path="sggetao/PwC", split="train[900:]"),
validation=dict(path="sggetao/PwC", split="train[:900]"),
test=dict(path="sggetao/PwC", split="test"),
),
"pwc_tiny": dict(
train=dict(path="sggetao/PwC", split="train[900:2000]"),
validation=dict(path="sggetao/PwC", split="train[:900]"),
),
"squad": dict(
train=dict(path="rajpurkar/squad", split="train[900:]"),
validation=dict(path="rajpurkar/squad", split="train[:900]"),
test=dict(path="rajpurkar/squad", split="validation"),
),
"fw_qa_tiny": dict(
train=dict(
path="parquet",
data_files="data/raw_datasets/fw_qa/00000.parquet",
split="train",
),
validation=dict(
path="parquet",
data_files="data/raw_datasets/fw_qa/00000_val.parquet",
split="train",
),
),
"fw_qa": dict(
train=dict(
path="parquet",
data_files=FW_QA_PATHS,
split="train",
),
validation=dict(
path="parquet",
data_files="data/raw_datasets/fw_qa/*_val.parquet",
split="train",
),
),
"fw_qa_large": dict(
train=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_large/*res.parquet"),
split="train",
),
validation=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_large/*val.parquet"),
split="train",
),
),
"fw_qa_xl": dict(
train=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_xl/*[!val].parquet"),
split="train",
),
validation=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_xl/*val.parquet"),
split="train",
),
),
"fw_qa_2": dict(
train=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_2/*[!val].parquet"),
split="train",
),
validation=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_2/*val.parquet"),
split="train",
),
),
"fw_qa_3_mini": dict(
train=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_3/000_00001.parquet"),
split="train[:100000]",
),
),
"fw_qa_3_medium": dict(
train=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_3/00[0-5]*[!val].parquet"),
split="train",
),
),
"fw_qa_3": dict(
train=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_3/*[!val].parquet"),
split="train",
),
validation=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_3/*val.parquet"),
split="train",
),
),
"ctx_qa": dict(
train=dict(
path="parquet",
data_files=glob("data/raw_datasets/ctx_qa/*res.parquet"),
split="train",
),
validation=dict(
path="parquet",
data_files=glob("data/raw_datasets/ctx_qa/*val.parquet"),
split="train",
),
),
"drop": dict(
train=dict(
path="ucinlp/drop",
split="train",
),
),
"narrativeqa": dict(
train=dict(
path="nvidia/ChatQA-Training-Data",
name="narrativeqa",
split="train",
),
),
"quoref": dict(
train=dict(
path="nvidia/ChatQA-Training-Data",
name="quoref",
split="train",
),
),
"ropes": dict(
train=dict(
path="nvidia/ChatQA-Training-Data",
name="ropes",
split="train",
),
),
"synthetic_convqa": dict(
train=dict(
path="nvidia/ChatQA-Training-Data",
name="synthetic_convqa",
split="train",
),
),
"booksum": dict(
train=dict(
path="kmfoda/booksum",
split="train+validation+test",
)
),
"gov_report": dict(
train=dict(
path="ccdv/govreport-summarization",
split="train",
),
validation=dict(
path="ccdv/govreport-summarization",
split="validation",
),
test=dict(
path="ccdv/govreport-summarization",
split="test",
),
),
# "wikitext-2": dict(
# train=dict(
# path="EleutherAI/wikitext_document_level",
# name="wikitext-2-raw-v1",
# split="train",
# ),
# ),
# "wikitext-103": dict(
# train=dict(
# path="EleutherAI/wikitext_document_level",
# name="wikitext-103-raw-v1",
# split="train",
# ),
# ),
"gsm8k": dict(
train=dict(
path="openai/gsm8k",
name="main",
split="train[100:]",
),
validation=dict(
path="openai/gsm8k",
name="main",
split="train[:100]",
),
),
"openmathintx-2": dict(
train=dict(
path="nvidia/OpenMathInstruct-2",
split="train_1M[:100000]",
),
),
"opencoder-edu": dict(
train=dict(
path="OpenCoder-LLM/opc-sft-stage2",
name="educational_instruct",
split="train[900:]",
),
validation=dict(
path="OpenCoder-LLM/opc-sft-stage2",
name="educational_instruct",
split="train[:900]",
),
),
}
def closed_qa_prompting(prompt: str):
template = random.choice(CLOSED_QA_INTX_TEMPLATES)
return template.format(input=prompt)
def get_preprocessing_fn(
ds_name: str, is_eval: bool
) -> Callable[[dict[str, Any]], dict[str, Any]]:
"""
Get preprocessing function for a specific dataset.
Args:
ds_name: Name of the dataset
Returns:
A preprocessing function that takes and returns a dictionary
"""
f = lambda x: x
if ds_name.startswith("pwc"):
def f(sample):
return {
"context": sample["input"],
"prompt": sample["prompt"],
"response": sample["answer"],
}
elif ds_name.startswith("hotpot_qa"):
def f(sample):
txt = ""
for p in sample["context"]["sentences"]:
txt += " " + "".join(p)
q = sample["question"]
prompt = closed_qa_prompting(q) if not is_eval else q
return {
"context": txt.strip(),
"prompt": prompt,
"response": sample["answer"],
}
elif "squad" in ds_name:
def f(sample):
q = sample["question"]
prompt = closed_qa_prompting(q) if not is_eval else q
return {
"context": sample["context"],
"prompt": prompt,
"response": sample["answers"]["text"][0],
}
elif ds_name == "drop":
def f(sample):
q = sample["question"]
prompt = closed_qa_prompting(q) if not is_eval else q
return {
"context": sample["passage"],
"prompt": prompt,
"response": ", ".join(sample["answers_spans"]["spans"]),
}
elif ds_name in ["narrativeqa", "quoref", "ropes"]:
def f(sample):
response = sample["answers"][0]
if isinstance(response, list):
response = response[0]
q = sample["messages"][-1]["content"]
prompt = closed_qa_prompting(q) if not is_eval else q
return {
"context": sample["document"],
"prompt": prompt,
"response": response,
}
elif ds_name == "synthetic_convqa":
def f(sample):
response = sample["answers"][0]
if isinstance(response, list):
response = response[0]
return {
"context": sample["document"],
"prompt": sample["messages"][-1]["content"],
"response": response,
}
elif ds_name == "booksum":
prompt_templates = [
"Summarization the provided text.",
"# Summary",
"### Summary",
"Summary of the text",
]
# TODO: use these templates for training
# analysis_templates = []
# summary_templates = []
def f(sample):
return {
"context": sample["chapter"],
"prompt": "Summarize the provided text.",
"response": sample["summary_text"],
}
elif ds_name == "gov_report":
def f(sample):
return {
"context": sample["report"],
"prompt": "Summarize the provided text.",
"response": sample["summary"],
}
elif "wikitext" in ds_name:
def f(sample):
return {
"context": sample["page"],
"prompt": "PLAECHOLDER",
"response": "PLAECHOLDER",
}
elif ds_name == "openmathintx-2":
def f(sample):
return {
"context": sample["problem"],
"prompt": sample["problem"],
"response": sample["generated_solution"],
}
elif "gsm8k" in ds_name:
def f(sample):
return {
"context": sample["question"],
"prompt": sample["question"],
"response": sample["answer"],
}
elif "opencoder-edu" in ds_name:
def f(sample):
return {
"context": sample["instruction"],
"prompt": sample["instruction"],
"response": "```python\n" + sample["code"].strip() + "\n```",
}
elif ds_name.startswith("longbench"):
def f(sample):
return {
"context": sample["context"],
"prompt": sample["input"],
"response": sample["answers"][0],
}
elif "natural_questions" in ds_name:
pass
# ctx = " ".join(
# [
# token
# for is_html, token in zip(
# ds[0]["document"]["tokens"]["is_html"],
# ds[0]["document"]["tokens"]["token"],
# )
# if not is_html
# ]
# )
elif ds_name == "triviaqa_retrieved":
# maybe needed for training (negative sample for training)
# def f(sample):
# ctx = sample["entity_page"]["wiki_context"]
# if not ctx:
# return None
def f(sample):
return {
"context": sample["context"],
"prompt": sample["prompt"],
"response": sample["answer"],
}
if is_eval and ds_name in EVAL_INTX_TEMPLATES:
prompt_template = EVAL_INTX_TEMPLATES[ds_name]
def eval_intx_decorator(f):
def g(sample):
sample = f(sample)
sample["prompt"] = prompt_template.format(input=sample["prompt"])
return sample
return g
return eval_intx_decorator(f) if is_eval else f
def get_ds_kwargs(ds_name: str, split: str) -> dict[str, Any]:
@ -425,7 +395,7 @@ def _load_and_process_dataset(
and ds_kwargs["path"] == "parquet"
)
ds = load_dataset(
**get_ds_kwargs(ds_name, split),
**ds_kwargs,
trust_remote_code=True,
streaming=load_as_stream,
)
@ -437,24 +407,27 @@ def _load_and_process_dataset(
if take is not None:
ds = ds.take(take)
except ValueError as e:
logger.info(
logger.warning(
f"Failed to load dataset {ds_name} with split {split}. Error: {e}\nSkipping..."
)
return None
cols_to_remove = [
col for col in ds.column_names if col not in ["context", "prompt", "response"]
]
is_eval = split != "train"
ds = ds.map(
get_preprocessing_fn(ds_name),
get_preprocessing_fn(ds_name, is_eval),
remove_columns=cols_to_remove,
num_proc=num_proc,
)
# ds = ds.remove_columns(cols_to_remove)
ds = ds.filter(filter_none, batched=True, num_proc=num_proc)
ds = ds.filter(filter_long_samples, batched=True, num_proc=num_proc)
if split == "train":
# TODO: drop samples or split into multiple contexts
ds = ds.filter(filter_long_samples, batched=True, num_proc=num_proc)
if add_negative_prompt:
# TODO: make explicit dataset for this
ds = ds.map(
add_negative_prompt_fn,
batched=True,
@ -462,21 +435,24 @@ def _load_and_process_dataset(
num_proc=num_proc,
)
if add_repeat_prompt and "context_numbers" not in ds_name:
# TODO: convert to AE/LM
ds = ds.map(
add_repeat_prompt_fn,
batched=True,
batch_size=100_000,
num_proc=num_proc,
)
ds.save_to_disk(ds_path, num_proc=num_proc)
ds.save_to_disk(ds_path, num_proc=num_proc)
return ds
def get_tokenized_dataset(
ds_name: str,
split: str,
base_model_max_len: int,
tokenizer: PreTrainedTokenizerBase,
tokenizer_kwargs: dict[str, Any],
ctx_model_max_len: int,
ctx_tokenizer: PreTrainedTokenizerBase,
ctx_tokenizer_kwargs: dict[str, Any],
add_ctx_to_chat: bool,
@ -486,6 +462,8 @@ def get_tokenized_dataset(
set_format: Optional[str] = None,
streaming: bool = False,
) -> dict[str, Any]:
# TODO: update hash function to include max_len
# TODO: max_len could be included in the kwargs
assert not use_kl_loss, "KL loss is deprecated"
logger.debug(f"Loading dataset {ds_name} with split {split}...")
need_ctx_ids = not add_ctx_to_chat
@ -502,27 +480,32 @@ def get_tokenized_dataset(
ds_hash = hashlib.sha256(json.dumps(load_and_process_kwargs).encode()).hexdigest()
ds_path = f"{TRANSFORMED_DATA_DIR}/{ds_hash}"
if path.exists(ds_path):
if path.exists(ds_path) and split == "train":
# load the cached ds
logger.info(f"Loaded processed dataset from {ds_path}")
else:
logger.info(f"Loading dataset {ds_name} with split {split}...")
_load_and_process_dataset(
ds = _load_and_process_dataset(
**load_and_process_kwargs,
ds_path=ds_path,
num_proc=num_proc,
)
ds = datasets.load_from_disk(ds_path)
if split == "train":
# force loading the cached version for newly created ds
ds = datasets.load_from_disk(ds_path)
tokenized_ds = construct_and_tokenize_ctx_qa(
base_model_max_len,
tokenizer,
tokenizer_kwargs,
ctx_model_max_len,
ctx_tokenizer,
ctx_tokenizer_kwargs,
add_ctx_to_chat,
use_kl_loss,
need_ctx_ids,
ds,
split,
set_format,
num_proc,
)
@ -530,17 +513,23 @@ def get_tokenized_dataset(
def construct_and_tokenize_ctx_qa(
base_model_max_len,
tokenizer,
tokenizer_kwargs,
ctx_model_max_len,
ctx_tokenizer,
ctx_tokenizer_kwargs,
add_ctx_to_chat,
use_kl_loss,
need_ctx_ids,
ds,
split,
set_format=None,
num_proc=None,
):
is_eval = split != "train"
# TODO: update hash function to include max_len
# TODO: max_len could be included in the kwargs
kwargs = dict(
tokenizer=repr(tokenizer),
tokenizer_kwargs=json.dumps(tokenizer_kwargs),
@ -550,13 +539,14 @@ def construct_and_tokenize_ctx_qa(
use_kl_loss=use_kl_loss,
need_ctx_ids=need_ctx_ids,
ds=ds._fingerprint,
# TODO: add split (after moving to continued pre-training pipeline)
set_format=set_format,
)
kwargs_str = json.dumps(kwargs)
logger.debug(f"Tokenizing dataset with kwargs: {kwargs_str}")
ds_hash = hashlib.sha256(kwargs_str.encode()).hexdigest()
ds_path = f"{TRANSFORMED_DATA_DIR}/{ds_hash}"
if path.exists(ds_path):
if path.exists(ds_path) and split == "train":
# load the cached ds
logger.info(f"Loaded tokenized dataset from {ds_path}")
ds = datasets.load_from_disk(ds_path)
@ -576,12 +566,14 @@ def construct_and_tokenize_ctx_qa(
num_proc=num_proc,
)
# tokenize the chat + mask the assistant inputs
ds = ds.filter(
filter_long_chat,
batched=True,
batch_size=100_000,
num_proc=num_proc,
)
# TODO: drop samples or split into multiple contexts
if split == "train":
ds = ds.filter(
filter_long_chat,
batched=True,
batch_size=100_000,
num_proc=num_proc,
)
# add "input_ids", "attention_mask", "labels"
tokenized_ds = ds.map(
@ -593,6 +585,21 @@ def construct_and_tokenize_ctx_qa(
},
num_proc=num_proc,
)
if split == "train":
# drop?
# currently the above function would raise if the chat is too long
...
else:
tokenized_ds = tokenized_ds.map(
truncate_middle_if_too_long,
fn_kwargs={
"max_length": base_model_max_len,
"columns": ["input_ids", "attention_mask"],
},
batched=True,
batch_size=100_000,
num_proc=num_proc,
)
# for use_kl_loss, we need "chat_ids" and "chat_attn_mask"
if use_kl_loss:
@ -625,7 +632,23 @@ def construct_and_tokenize_ctx_qa(
batch_size=100_000,
num_proc=num_proc,
)
if split == "train":
# TODO: do something if ctx length is longer than the ctx model length
# e.g., drop or split for multi-lora training
...
else:
# truncate in the middle
tokenized_ds = tokenized_ds.map(
truncate_middle_if_too_long,
fn_kwargs={
"max_length": ctx_model_max_len,
"columns": ["ctx_ids", "ctx_attn_mask"],
},
batched=True,
batch_size=100_000,
num_proc=num_proc,
)
tokenized_ds = tokenized_ds.remove_columns(
["messages", "chat", "context", "prompt", "response"]
)
@ -635,9 +658,11 @@ def construct_and_tokenize_ctx_qa(
# # the columns are unknown when using streaming dataset
# tokenized_ds = tokenized_ds._resolve_features()
# validate_columns(tokenized_ds)
tokenized_ds.save_to_disk(ds_path, num_proc=num_proc)
del tokenized_ds
return datasets.load_from_disk(ds_path)
if split == "train":
tokenized_ds.save_to_disk(ds_path, num_proc=num_proc)
del tokenized_ds
return datasets.load_from_disk(ds_path)
return tokenized_ds
def get_sft_prompt_formatting_fn(
@ -734,127 +759,6 @@ def convert_ctx_prompt_response_to_messages(
return dict(messages=messages)
def get_preprocessing_fn(ds_name: str) -> Callable[[dict[str, Any]], dict[str, Any]]:
"""
Get preprocessing function for a specific dataset.
Args:
ds_name: Name of the dataset
Returns:
A preprocessing function that takes and returns a dictionary
"""
f = lambda x: x
if ds_name.startswith("pwc"):
def f(sample):
return {
"context": sample["input"],
"prompt": sample["prompt"],
"response": sample["answer"],
}
elif ds_name.startswith("hotpot_qa"):
def f(sample):
txt = ""
for p in sample["context"]["sentences"]:
txt += " " + "".join(p)
return {
"context": txt.strip(),
"prompt": sample["question"],
"response": sample["answer"],
}
elif "squad" in ds_name:
def f(sample):
return {
"context": sample["context"],
"prompt": sample["question"],
"response": sample["answers"]["text"][0],
}
elif ds_name == "drop":
def f(sample):
return {
"context": sample["passage"],
"prompt": sample["question"],
"response": ", ".join(sample["answers_spans"]["spans"]),
}
elif ds_name in ["narrativeqa", "quoref", "ropes", "synthetic_convqa"]:
def f(sample):
response = sample["answers"][0]
if isinstance(response, list):
response = response[0]
return {
"context": sample["document"],
"prompt": sample["messages"][0]["content"],
"response": response,
}
elif ds_name == "booksum":
def f(sample):
return {
"context": sample["chapter"],
"prompt": "Summarize the provided text.",
"response": sample["summary_text"],
}
elif ds_name == "gov_report":
def f(sample):
return {
"context": sample["report"],
"prompt": "Summarize the provided text.",
"response": sample["summary"],
}
elif "wikitext" in ds_name:
def f(sample):
return {
"context": sample["page"],
"prompt": "PLAECHOLDER",
"response": "PLAECHOLDER",
}
elif ds_name == "openmathintx-2":
def f(sample):
return {
"context": sample["problem"],
"prompt": sample["problem"],
"response": sample["generated_solution"],
}
elif "gsm8k" in ds_name:
def f(sample):
return {
"context": sample["question"],
"prompt": sample["question"],
"response": sample["answer"],
}
elif "opencoder-edu" in ds_name:
def f(sample):
return {
"context": sample["instruction"],
"prompt": sample["instruction"],
"response": "```python\n" + sample["code"].strip() + "\n```",
}
return f
# adapted from https://github.com/huggingface/trl/issues/632#issuecomment-1972630547
def get_assistant_start_end_indices(
messages: list[dict[str, str]],
@ -985,6 +889,31 @@ def tokenize_chat_messages(
return conversation_ids
def truncate_middle_if_too_long(
samples: dict[str, any],
max_length: int,
columns: list[str],
) -> dict[str, any]:
"""
Truncate the middle of a list of tokens to fit within a maximum length.
Args:
tokens: List of token IDs
max_length: Maximum length for the truncated tokens
Returns:
List of truncated token IDs
"""
half = max_length // 2
for col in columns:
# if len(samples[col]) <= max_length:
# return samples[col]
samples[col] = [
t[:half] + t[-half:] if len(t) > max_length else t for t in samples[col]
]
return samples
def tokenize_ctx_text(
samples: dict[str, Any],
tokenizer: PreTrainedTokenizerBase,

123
src/ctx_to_lora/eval.py
View file

@ -1,18 +1,18 @@
import re
import string
import yaml
import sys
import gc
import json
import os
import random
import shutil
import logging
from argparse import Namespace
from dataclasses import fields
from functools import partial
from collections import defaultdict
from collections import Counter, defaultdict
from typing import Callable
import numpy as np
import torch
import numpy as np
from datasets import disable_caching
from peft import PeftModel
from rouge_score import rouge_scorer
@ -28,10 +28,79 @@ from transformers import (
from transformers.utils import is_liger_kernel_available
from ctx_to_lora.data_definitions import LONGBENCH_E_TASKS, LONGBENCH_TASKS
from ctx_to_lora.data_utils import get_tokenized_dataset
from ctx_to_lora.modeling_utils import ModulatedPretrainedModel
from ctx_to_lora.model_loading import get_tokenizer, get_model
logger = logging.getLogger()
# TODO: add negative samples eval (per-sample?)
# TODO: add length info for more fine grain eval
# longbench metrics
def normalize_answer(s):
"""Lower text and remove punctuation, articles and extra whitespace."""
def remove_articles(text):
return re.sub(r"\b(a|an|the)\b", " ", text)
def white_space_fix(text):
return " ".join(text.split())
def remove_punc(text):
exclude = set(string.punctuation)
return "".join(ch for ch in text if ch not in exclude)
def lower(text):
return text.lower()
return white_space_fix(remove_articles(remove_punc(lower(s))))
def f1_score(prediction: str, ground_truth: str):
common = Counter(prediction) & Counter(ground_truth)
num_same = sum(common.values())
if num_same == 0:
return 0
precision = 1.0 * num_same / len(prediction)
recall = 1.0 * num_same / len(ground_truth)
f1 = (2 * precision * recall) / (precision + recall)
return f1
def compute_qa_f1_score(pred_texts: list[str], label_texts: list[str]):
"""
Word-level F1 score for evaluating question answering systems.
Order of the words does not matter.
"""
res = []
for prediction, label in zip(pred_texts, label_texts):
normalized_prediction = normalize_answer(prediction)
normalized_label = normalize_answer(label)
prediction_words = normalized_prediction.split()
label_words = normalized_label.split()
res.append(f1_score(prediction_words, label_words))
return dict(qa_f1=np.mean(res))
CLOSED_QA_DATASETS = {
"longbench/narrativeqa",
"longbench/qasper",
"longbench/multifieldqa_en",
"longbench/hotpotqa",
"longbench/2wikimqa",
"longbench/musique",
"squad",
"triviaqa_retrieved",
}
for ds_name in list(CLOSED_QA_DATASETS):
CLOSED_QA_DATASETS.add(f"{ds_name}_e")
def clear_gpu():
gc.collect()
@ -40,6 +109,15 @@ def clear_gpu():
torch.cuda.reset_max_memory_cached()
def add_longbench_tasks(ds_names: list[str]):
if "longbench" in ds_names:
ds_names.remove("longbench")
ds_names += LONGBENCH_TASKS
if "longbench_e" in ds_names:
ds_names.remove("longbench_e")
ds_names += LONGBENCH_E_TASKS
def compute_rouge(pred_texts, label_texts):
out = defaultdict(list)
scorer = rouge_scorer.RougeScorer(["rouge1", "rougeL"], use_stemmer=False)
@ -142,6 +220,11 @@ def compute_metrics(
def save_generated_text(samples, output_dir, split):
os.makedirs(output_dir, exist_ok=True)
# Create any necessary subdirectories if split contains path separators
if "/" in split:
split_dir = os.path.join(output_dir, os.path.dirname(split))
os.makedirs(split_dir, exist_ok=True)
with open(f"{output_dir}/{split}_generated_text.jsonl", "w") as f:
for sample in samples:
f.write(json.dumps(sample) + "\n")
@ -160,7 +243,7 @@ def decode_test_result(test_dataset, test_result, tokenizer, ctx_tokenizer):
# replace them with the pad token id
label_toks = np.where(label_toks == -100, tokenizer.pad_token_id, label_toks)
label_text = tokenizer.decode(label_toks, skip_special_tokens=True)
d["label"] = label_text
d["label"] = label_text.strip()
# remove the label part
input_toks = sample["input_ids"][:start_idx]
@ -170,7 +253,7 @@ def decode_test_result(test_dataset, test_result, tokenizer, ctx_tokenizer):
gen_toks = np.where(gen_toks == -100, tokenizer.pad_token_id, gen_toks)
d["input"] = tokenizer.decode(input_toks, skip_special_tokens=True)
d["generated"] = tokenizer.decode(gen_toks, skip_special_tokens=True)
d["generated"] = tokenizer.decode(gen_toks, skip_special_tokens=True).strip()
if "ctx_ids" in sample:
d["context"] = ctx_tokenizer.decode(
sample["ctx_ids"], skip_special_tokens=True
@ -186,6 +269,7 @@ def eval_generation(eval_trainer, tokenizer, ctx_tokenizer, datasets, split, gen
datasets = {"": datasets}
out = {}
for ds_name, ds in datasets.items():
print(f"Evaluating: {ds_name}")
split_name = f"{split}_{ds_name}" if ds_name else split
eval_result = eval_trainer.predict(
ds,
@ -193,12 +277,21 @@ def eval_generation(eval_trainer, tokenizer, ctx_tokenizer, datasets, split, gen
**gen_kwargs,
)
decoded_txts = decode_test_result(ds, eval_result, tokenizer, ctx_tokenizer)
rouge_metrics = compute_rouge(
[txt["generated"] for txt in decoded_txts],
[txt["label"] for txt in decoded_txts],
)
pred_texts = [txt["generated"] for txt in decoded_txts]
label_texts = [txt["label"] for txt in decoded_txts]
rouge_metrics = compute_rouge(pred_texts, label_texts)
for k, v in rouge_metrics.items():
eval_result.metrics[f"{split_name}_{k}"] = v
if ds_name in CLOSED_QA_DATASETS:
print("Computing QA F1 Score")
qa_f1_metric = compute_qa_f1_score(
[txt["generated"] for txt in decoded_txts],
[txt["label"] for txt in decoded_txts],
)
for k, v in qa_f1_metric.items():
eval_result.metrics[f"{split_name}_{k}"] = v
save_generated_text(
decoded_txts,
split=split_name,
@ -346,8 +439,8 @@ def evaluate(
train=False,
use_flash_attn=True,
)
if generative:
model = model.to(torch.bfloat16)
# if generative:
# model = model.to(torch.bfloat16)
else:
model = get_model(model_name_or_path, train=False, requires_grad=False)
# NOTE: there is still some randomness in the eval result
@ -392,8 +485,10 @@ def evaluate(
_get_tokenized_dataset = partial(
get_tokenized_dataset,
base_model_max_len=model.base_model.config.max_position_embeddings,
tokenizer=tokenizer,
tokenizer_kwargs=tokenizer_kwargs,
ctx_model_max_len=model.ctx_encoder.config.max_position_embeddings,
ctx_tokenizer=ctx_tokenizer,
ctx_tokenizer_kwargs=ctx_tokenizer_kwargs,
add_ctx_to_chat=add_ctx_to_chat,
@ -405,6 +500,7 @@ def evaluate(
datasets = dict()
ds_names = args.val_ds_names if split == "validation" else args.test_ds_names
add_longbench_tasks(ds_names)
for ds_name in ds_names:
datasets[ds_name] = _get_tokenized_dataset(ds_name, split)
print(datasets)
@ -526,6 +622,7 @@ if __name__ == "__main__":
)
cli_args = parser.parse_args()
# setup_logging(output_dir, debug=os.getenv("DEBUG", False))
assert bool(cli_args.model_name_or_path) ^ bool(
cli_args.checkpoint_path

12
src/ctx_to_lora/modeling_utils.py
View file

@ -982,6 +982,7 @@ class HyperLoRA(nn.Module):
# hidden_size=self.config.base_hidden_size,
# )
@torch.autocast(device_type="cuda", dtype=torch.bfloat16)
def _to_lora_dict(
self, flat_loras: Float[Tensor, "bs n_layers n_modules r max_io_dim"]
) -> dict[str, dict[str, Float[Tensor, "bs n_layers r _"]]]:
@ -1035,6 +1036,7 @@ class HyperLoRA(nn.Module):
)
return {k: v for k, v in zip(self.extra_modules, layernorms)}
@torch.autocast(device_type="cuda", dtype=torch.bfloat16)
def forward(
self,
features: Float[Tensor, "bs seq_len feature_dim"],
@ -1065,6 +1067,7 @@ class HyperLoRA(nn.Module):
return flat_loras, flat_layernorms
@torch.autocast(device_type="cuda", dtype=torch.bfloat16)
def generate_weights(
self,
features: Float[Tensor, "bs seq_len feature_dim"],
@ -1486,9 +1489,7 @@ class ModulatedPretrainedModel(nn.Module):
*model_inputs_args: Any,
**model_inputs_kwargs: dict[str, Any],
):
# TODO: generate LoRA for each sample in ctx_ids
# apply all LoRAs to the base model
# the lora should be applied until removed
# TODO: make this persistent
generated_loras, _ = self.generate_weights(
ctx_ids, ctx_attn_mask, ctx_position_ids
)
@ -1672,11 +1673,6 @@ class ModulatedPretrainedModel(nn.Module):
return model_outputs
# def combine_loras(loras: list[LoRA]) -> LoRA:
# # TODO: implement
# ...
class ModulatedModelWithSharedInput(nn.Module):
def __init__(
self,

10
src/ctx_to_lora/training_utils.py
View file

@ -114,7 +114,7 @@ def train_model(
logger.info(f"Resuming from the checkpoint: {checkpoint}")
is_modulated_model = isinstance(model, ModulatedPretrainedModel)
trainer_cls = Trainer # if not is_modulated_model else ModulatedModelTrainer
trainer_cls = Trainer if not is_modulated_model else ModulatedModelTrainer
trainer_kwargs = dict(
model=model,
args=training_args,
@ -123,10 +123,10 @@ def train_model(
data_collator=train_collator,
compute_metrics=compute_metrics,
)
# if is_modulated_model:
# logger.info(f"Training with modulated model. Using CustomTrainer.")
# trainer_kwargs["gen_lora_l1_reg_coef"] = training_args.gen_lora_l1_reg_coef
# del training_args.gen_lora_l1_reg_coef
if is_modulated_model:
logger.info(f"Training with modulated model. Using CustomTrainer.")
trainer_kwargs["gen_lora_l1_reg_coef"] = training_args.gen_lora_l1_reg_coef
del training_args.gen_lora_l1_reg_coef
trainer = trainer_cls(**trainer_kwargs)

45
webui/app.py
View file

@ -21,9 +21,14 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
modulated_model = None
# TODO: grab the results on-deman
# TODO: cache the grabed results
def get_run_data(run_path):
"""
Loads and processes data from all_results.json and *_results.json for a given run.
Also recursively searches for results files in subfolders.
Args:
run_path: Path to the directory containing the results files.
@ -38,7 +43,7 @@ def get_run_data(run_path):
except FileNotFoundError:
data = {}
# Load data from *_results.json files
# Load data from *_results.json files in the current directory
for results_json_file in glob(os.path.join(run_path, "*_results.json")):
try:
with open(results_json_file) as f:
@ -55,6 +60,24 @@ def get_run_data(run_path):
except FileNotFoundError:
print(f"Warning: {results_json_file} not found.")
# Recursively search for *_results.json files in subdirectories
for subdir in [d for d in glob(os.path.join(run_path, "*")) if os.path.isdir(d)]:
subdir_name = os.path.basename(subdir)
# Look for *_results.json files in the subdirectory
for results_json_file in glob(os.path.join(subdir, "*_results.json")):
try:
with open(results_json_file) as f:
split_data = json.load(f)
# Use a prefix that includes the subfolder name
file_prefix = f"{subdir_name}_{os.path.basename(results_json_file).replace('_results.json', '')}"
for key, value in split_data.items():
if isinstance(value, float):
value = round(value, 4)
new_key = f"{file_prefix}_{key}"
data[new_key] = value
except FileNotFoundError:
print(f"Warning: {results_json_file} not found.")
grouped_data = {}
for key, value in data.items():
if isinstance(value, dict):
@ -96,6 +119,7 @@ def get_run_data(run_path):
def get_generated_text_data(run_path):
"""
Loads generated text data from all *_generated_text.jsonl files.
Also recursively searches for generated text files in subfolders.
Args:
run_path: Path to the directory containing the .jsonl files.
@ -106,7 +130,7 @@ def get_generated_text_data(run_path):
generated_data = {}
files = sorted(glob(f"{run_path}/*_generated_text.jsonl"))
print(files)
# Find all *_generated_text.jsonl files
# Find all *_generated_text.jsonl files in the current directory
for filename in files:
# Extract split name from filename (remove _generated_text.jsonl)
split = filename.split("/")[-1].split("_generated_text.jsonl")[0]
@ -117,6 +141,23 @@ def get_generated_text_data(run_path):
generated_data[split] = data
except FileNotFoundError:
generated_data[split] = None
# Recursively search for *_generated_text.jsonl files in subdirectories
for subdir in [d for d in glob(os.path.join(run_path, "*")) if os.path.isdir(d)]:
subdir_name = os.path.basename(subdir)
subdir_files = sorted(glob(f"{subdir}/*_generated_text.jsonl"))
for filename in subdir_files:
# Extract split name from filename and include subfolder name
base_split = filename.split("/")[-1].split("_generated_text.jsonl")[0]
split = f"{subdir_name}_{base_split}"
try:
with open(filename) as f:
lines = f.readlines()
data = [json.loads(line) for line in lines]
generated_data[split] = data
except FileNotFoundError:
generated_data[split] = None
return generated_data