diff --git a/.gitignore b/.gitignore index d758e9e..9b6be1e 100644 --- a/.gitignore +++ b/.gitignore @@ -23,3 +23,5 @@ plots/ watcher_state.yaml ds_config.json eval_results/ +.github/ +*.code-workspace \ No newline at end of file diff --git a/README.md b/README.md index 72abfe5..af81de4 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/intx_sft.py b/intx_sft.py index 1cebfd6..e2e87c3 100755 --- a/intx_sft.py +++ b/intx_sft.py @@ -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 diff --git a/scripts/fw_qa_3/medium_gemma_per_rank_fac2_per_layer_4_gpus_gen_lora_reg.sh b/scripts/fw_qa_3/medium_gemma_per_rank_fac2_per_layer_4_gpus_gen_lora_reg.sh new file mode 100644 index 0000000..5193bf3 --- /dev/null +++ b/scripts/fw_qa_3/medium_gemma_per_rank_fac2_per_layer_4_gpus_gen_lora_reg.sh @@ -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 diff --git a/src/ctx_to_lora/configs.py b/src/ctx_to_lora/configs.py index 8245097..f72d90b 100644 --- a/src/ctx_to_lora/configs.py +++ b/src/ctx_to_lora/configs.py @@ -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( diff --git a/src/ctx_to_lora/data_definitions.py b/src/ctx_to_lora/data_definitions.py new file mode 100644 index 0000000..ebfb1b7 --- /dev/null +++ b/src/ctx_to_lora/data_definitions.py @@ -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}" diff --git a/src/ctx_to_lora/data_utils.py b/src/ctx_to_lora/data_utils.py index 99834d3..6492b2d 100644 --- a/src/ctx_to_lora/data_utils.py +++ b/src/ctx_to_lora/data_utils.py @@ -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, diff --git a/src/ctx_to_lora/eval.py b/src/ctx_to_lora/eval.py index da0821c..ac6b064 100644 --- a/src/ctx_to_lora/eval.py +++ b/src/ctx_to_lora/eval.py @@ -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 diff --git a/src/ctx_to_lora/modeling_utils.py b/src/ctx_to_lora/modeling_utils.py index 14ec4bc..173f626 100644 --- a/src/ctx_to_lora/modeling_utils.py +++ b/src/ctx_to_lora/modeling_utils.py @@ -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, diff --git a/src/ctx_to_lora/training_utils.py b/src/ctx_to_lora/training_utils.py index 2a11845..3e4ca2a 100644 --- a/src/ctx_to_lora/training_utils.py +++ b/src/ctx_to_lora/training_utils.py @@ -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) diff --git a/webui/app.py b/webui/app.py index 0265c2c..4ccb8c3 100644 --- a/webui/app.py +++ b/webui/app.py @@ -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