From 8cbbb45333188ba658896044438c83585887a1e3 Mon Sep 17 00:00:00 2001 From: Akim Tsvigun Date: Tue, 12 Aug 2025 18:32:02 +0000 Subject: [PATCH 1/4] merged --- .gitignore | 1 + configs/base.yaml | 8 +- configs/test.yaml | 2 + requirements.txt | 16 +-- src/atgen/labellers/get_labeller.py | 2 +- src/atgen/run_scripts/run_active_learning.py | 123 +++++++++++-------- src/atgen/strategies/base_strategy.py | 11 +- src/atgen/utils/constants.py | 4 + src/atgen/utils/installers.py | 10 -- src/atgen/utils/main_decorator.py | 6 +- src/atgen/utils/save_log_iter_results.py | 8 +- src/atgen/utils/training_utils.py | 10 +- 12 files changed, 111 insertions(+), 90 deletions(-) delete mode 100644 src/atgen/utils/installers.py diff --git a/.gitignore b/.gitignore index d1683db..1ee3b86 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ venv venv/* +repos TensorRT-LLM */dataset cache diff --git a/configs/base.yaml b/configs/base.yaml index 8cb9225..b501f13 100644 --- a/configs/base.yaml +++ b/configs/base.yaml @@ -17,6 +17,7 @@ name: base output_dir: outputs offline_mode: False cache_dir: cache +save_model: False seed: 42 al: # Strategy configuration now loaded from configs/al/.yaml @@ -39,9 +40,9 @@ model: assistant_response_start: "\n\n\n\n" peft: use: True - r: 64 - lora_alpha: 64 - lora_dropout: 0.1 + r: 32 + lora_alpha: 32 + lora_dropout: 0. bias: 'none' seed: ${seed} use_gradient_checkpointing: 'unsloth' @@ -75,6 +76,7 @@ inference: temperature: 0.6 model: ${model.checkpoint} num_display_generations: 5 + num_threads_for_bfcl: 32 evaluation: additional_metrics: [] diff --git a/configs/test.yaml b/configs/test.yaml index 2535b22..08b0b03 100644 --- a/configs/test.yaml +++ b/configs/test.yaml @@ -10,6 +10,7 @@ output_dir: outputs offline_mode: False cache_dir: cache seed: 42 +save_model: False al: init_query_size: 2 query_size: 2 @@ -64,6 +65,7 @@ inference: top_p: 0.1 model: ${model.checkpoint} num_display_generations: 0 + num_threads_for_bfcl: 32 evaluation: additional_metrics: [] diff --git a/requirements.txt b/requirements.txt index a65f05f..da23ecf 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,22 +1,22 @@ alignscore-SpeedOfMagic -accelerate==1.5.2 -anthropic==0.49.0 +accelerate>=1.5.2 +anthropic>=0.49.0 benepar==0.2.0 bert-score==0.3.13 bitsandbytes==0.45.3 ctc_score==0.1.3 -datasets==3.4.1 +datasets>=3.6.0 deepeval==2.5.5 evaluate==0.4.3 hydra-core==1.3.2 nltk==3.9.1 kaleido==0.2.1 omegaconf==2.3.0 -openai==1.66.3 +openai>=1.90.0 openpyxl==3.1.5 peft==0.14.0 plotly==5.23.0 -protobuf==3.20.3 +protobuf>=3.20.3 pytest==8.3.2 pytorch_lightning==2.5.0.post0 rake_nltk==1.0.6 @@ -28,9 +28,9 @@ spacy==3.7.5 streamlit==1.37.0 streamlit-authenticator==0.4.2 tabulate==0.9.0 -transformers==4.52.4 -trl==0.15.2 +transformers>=4.52.4 +trl==0.19.1 torchmetrics==1.4.1 unsloth==2025.8.1 -vllm==0.8.1 +vllm==0.10.0 xlrd==1.2.0 diff --git a/src/atgen/labellers/get_labeller.py b/src/atgen/labellers/get_labeller.py index 2f34d36..18c794e 100644 --- a/src/atgen/labellers/get_labeller.py +++ b/src/atgen/labellers/get_labeller.py @@ -12,7 +12,7 @@ def get_labeller( config: DictConfig, output_column_name: str = "output", - budget: int = 1_000_000, + budget: int | float = 1_000_000, workdir: str | Path = "tmp", **kwargs, ): diff --git a/src/atgen/run_scripts/run_active_learning.py b/src/atgen/run_scripts/run_active_learning.py index 65a8caf..ac713a6 100644 --- a/src/atgen/run_scripts/run_active_learning.py +++ b/src/atgen/run_scripts/run_active_learning.py @@ -32,7 +32,6 @@ def run_active_learning(config, workdir: Union[str, Path]): maybe_get_few_shot_examples, get_output_column_name_for_phase, ) - from atgen.utils.installers import install_spacy, install_nltk from atgen.utils.load_model_tokenizer import load_model_tokenizer from atgen.utils.prepare_model_for_training import prepare_model_for_training from atgen.utils.training_utils import get_trainer @@ -50,10 +49,7 @@ def run_active_learning(config, workdir: Union[str, Path]): from atgen.utils.check_performance_metrics import ( check_performance_against_requirements, ) - - # TODO Figure out how to stop downloading it every time - install_spacy() - install_nltk() + from atgen.utils.evaluate_bfcl import evaluate_bfcl seed = config.seed cache_dir = config.cache_dir @@ -65,7 +61,7 @@ def run_active_learning(config, workdir: Union[str, Path]): num_al_iterations = config.al.num_iterations required_performance_dict = check_required_performance( - config.al.required_performance + required_performance=config.al.required_performance ) budget = config.al.budget if budget is None: @@ -107,7 +103,7 @@ def run_active_learning(config, workdir: Union[str, Path]): cache_dir=config.cache_dir, seed=seed, ) - if has_test: + if has_test and not config.data.use_test_benchmark: test_data = load_data( data_config=config.data, split=TEST_DATA_SPLIT_DEFAULT_NAME, @@ -134,7 +130,7 @@ def run_active_learning(config, workdir: Union[str, Path]): print("Loading AL strategy.") al_strategy: BaseStrategy = get_strategy( - config.al.strategy, + strategy_name=config.al.strategy, subsample_size=config.al.subsample_size, unlabeled_pool=unlabeled_data[input_column_name], model=model, @@ -177,7 +173,7 @@ def run_active_learning(config, workdir: Union[str, Path]): lambda x: x["id"] not in set(labeled_ids) ) else: - query_ids: list[str] = al_strategy( + query_ids: list[int] = al_strategy( model=model, tokenizer=tokenizer, unlabeled_pool=unlabeled_data.remove_columns(output_column_name_train), @@ -231,27 +227,38 @@ def run_active_learning(config, workdir: Union[str, Path]): few_shot_examples=few_shot_examples, model_name=model_name, ) - # Evaluate the initial model before any training - if init_query_size_is_positive and config.al.evaluate_zero_iteration: - generations: list[str] = generate( - config.inference, - data=test_data, - model=model, - tokenizer=tokenizer, - save_dir=save_dir, - data_config=config.data, - model_config=config.model, - ) - if os.path.exists(save_dir): - rmtree(save_dir) - - metrics: dict[str, float] = compute_metrics( - generated_texts=generations, - reference_texts=test_data[output_column_name_test], - original_texts=test_data[input_column_name], - task=config.data.task, - config=config.evaluation, - cache_dir=cache_dir, + # Evaluate the initial model before any training + if init_query_size_is_positive and config.al.evaluate_zero_iteration: + generations: list[str] = generate( + config.inference, + data=test_data, + model=model, + tokenizer=tokenizer, + save_dir=save_dir, + data_config=config.data, + model_config=config.model, + ) + if os.path.exists(save_dir): + rmtree(save_dir) + + metrics: dict[str, float] = compute_metrics( + generated_texts=generations, + reference_texts=test_data[output_column_name_test], + original_texts=test_data[input_column_name], + task=config.data.task, + config=config.evaluation, + cache_dir=cache_dir, + ) + else: + if "bfcl" in config.data.test_split_name: + test_split_name = config.data.test_split_name.split("bfcl_")[1] + else: + raise NotImplementedError(f"Test split name {config.data.test_split_name} is not supported") + generations, metrics = evaluate_bfcl( + model_name=model_name, + bfcl_results_dir=iter_dir, + test_category=test_split_name, + num_threads=config.inference.num_threads_for_bfcl, ) # Check required performance metrics @@ -272,6 +279,7 @@ def run_active_learning(config, workdir: Union[str, Path]): al_iter=0, train_result={}, model=None, + tokenizer=None, ) # Start AL cycle. Use `num_al_iterations + 2` because we do not label data @@ -346,26 +354,40 @@ def run_active_learning(config, workdir: Union[str, Path]): if dev_split_size > 0: test_data = eval_data else: - generations: list[str] = generate( - config.inference, - data=test_data, - model=model, - tokenizer=tokenizer, - save_dir=save_dir, - data_config=config.data, - model_config=config.model, - ) - if os.path.exists(save_dir): - rmtree(save_dir) - - metrics: dict[str, float] = compute_metrics( - generated_texts=generations, - reference_texts=test_data[output_column_name_test], - original_texts=test_data[input_column_name], - task=config.data.task, - config=config.evaluation, - cache_dir=cache_dir, - ) + if not config.data.use_test_benchmark: + generations: list[str] = generate( + config.inference, + data=test_data, + model=model, + tokenizer=tokenizer, + save_dir=save_dir, + data_config=config.data, + model_config=config.model, + ) + if os.path.exists(save_dir): + rmtree(save_dir) + + metrics: dict[str, float] = compute_metrics( + generated_texts=generations, + reference_texts=test_data[output_column_name_test], + original_texts=test_data[input_column_name], + task=config.data.task, + config=config.evaluation, + cache_dir=cache_dir, + ) + else: + if "bfcl" in config.data.test_split_name: + test_split_name = config.data.test_split_name.split("bfcl_")[1] + else: + raise NotImplementedError(f"Test split name {config.data.test_split_name} is not supported") + generations, metrics = evaluate_bfcl( + model_name=model_name, + bfcl_results_dir=iter_dir, + model=model, + tokenizer=tokenizer, + test_category=test_split_name, + num_threads=config.inference.num_threads_for_bfcl, + ) # Check required performance metrics is_performance_reached, is_metrics_availability_checked, available_metrics = ( @@ -385,6 +407,7 @@ def run_active_learning(config, workdir: Union[str, Path]): al_iter=al_iter, train_result=train_result, model=model, + tokenizer=tokenizer, ) # Make AL query for the next round if we have not run out of iterations diff --git a/src/atgen/strategies/base_strategy.py b/src/atgen/strategies/base_strategy.py index eb96174..e81506f 100644 --- a/src/atgen/strategies/base_strategy.py +++ b/src/atgen/strategies/base_strategy.py @@ -2,10 +2,6 @@ import random from datasets import Dataset -from transformers import ( - PreTrainedModel, - PreTrainedTokenizer, -) class BaseStrategy(ABC): @@ -15,13 +11,10 @@ def __init__(self, subsample_size: int | float = -1): @abstractmethod def __call__( self, - model: PreTrainedModel, - tokenizer: PreTrainedTokenizer, unlabeled_pool: Dataset, - labeled_pool: Dataset, - input_column_name: str, - output_column_name: str, num_to_label: int, + *args, + **kwargs, ) -> list[int]: pass diff --git a/src/atgen/utils/constants.py b/src/atgen/utils/constants.py index 95fd7e9..afa22cd 100644 --- a/src/atgen/utils/constants.py +++ b/src/atgen/utils/constants.py @@ -13,3 +13,7 @@ OUTPUT_FIELD_PURPOSE_TRAIN = "train" OUTPUT_FIELD_PURPOSE_TEST = "test" + +DEFAULT_NUM_THREADS_BFCL = 32 +DEFAULT_GPU_MEMORY_UTILIZATION_BFCL = 0.75 +BFCL_NUM_RETRIES = 3 diff --git a/src/atgen/utils/installers.py b/src/atgen/utils/installers.py deleted file mode 100644 index b540fb9..0000000 --- a/src/atgen/utils/installers.py +++ /dev/null @@ -1,10 +0,0 @@ -import subprocess - - -def install_spacy(): - subprocess.run("python -m spacy download en_core_web_sm", shell=True) - -def install_nltk(): - import nltk - nltk.download('punkt') - nltk.download('punkt_tab') diff --git a/src/atgen/utils/main_decorator.py b/src/atgen/utils/main_decorator.py index 2e86397..40e9a6c 100644 --- a/src/atgen/utils/main_decorator.py +++ b/src/atgen/utils/main_decorator.py @@ -8,6 +8,7 @@ from .validate_and_fill_config import validate_and_fill_config from .resolvers import register_resolvers +from .downloaders import maybe_download_packages os.environ["WANDB_DISABLED"] = "true" @@ -28,6 +29,9 @@ def run_script(config): os.chdir(hydra.utils.get_original_cwd()) else: auto_generated_dir = config.output_dir + + maybe_download_packages(config.output_dir) + log.info(f"Work dir: {auto_generated_dir}") # Save config into yaml format with open(Path(auto_generated_dir) / "config.yaml", "w") as f: @@ -39,8 +43,8 @@ def run_script(config): os.environ["HF_EVALUATE_OFFLINE"] = "1" os.environ["PYTHONHASHSEED"] = str(config.seed) - from transformers import set_seed + from transformers import set_seed set_seed(config.seed) func(config, workdir=Path(auto_generated_dir)) diff --git a/src/atgen/utils/save_log_iter_results.py b/src/atgen/utils/save_log_iter_results.py index 39024a0..06a2062 100644 --- a/src/atgen/utils/save_log_iter_results.py +++ b/src/atgen/utils/save_log_iter_results.py @@ -4,7 +4,7 @@ from typing import Optional from omegaconf import DictConfig -from transformers import PreTrainedModel +from transformers import PreTrainedModel, PreTrainedTokenizer from atgen.utils.combine_results import combine_results @@ -21,6 +21,7 @@ def save_log_iter_results( al_iter: int, train_result: dict, model: Optional[PreTrainedModel] = None, + tokenizer: Optional[PreTrainedTokenizer] = None, ): log.info(metrics) with open(iter_dir / "train_result.json", "w") as f: @@ -32,5 +33,6 @@ def save_log_iter_results( combine_results(workdir, al_iter) log.info(f"Iteration {al_iter}: saving the trained model...") - if config.save_model and model is not None: - model.save_pretrained(workdir / "model.bin") + if config.save_model and (model is not None): + model.save_pretrained(workdir / "model") + tokenizer.save_pretrained(workdir / "tokenizer") diff --git a/src/atgen/utils/training_utils.py b/src/atgen/utils/training_utils.py index 4816e71..2216096 100644 --- a/src/atgen/utils/training_utils.py +++ b/src/atgen/utils/training_utils.py @@ -202,10 +202,10 @@ def _get_response_instruction_templates( """ if "gemma" in tokenizer.name_or_path.lower(): response_template = "model\n" - instruction_template = "\nuser\n" - elif "qwen" in tokenizer.name_or_path.lower(): - response_template = "<|im_start|>assistant\n" - instruction_template = "\n<|im_start|>user\n" + instruction_template = "user\n" + elif "qwen3" in tokenizer.name_or_path.lower(): + response_template = "<|im_start|>assistant\n" # don't include empty reasoning + instruction_template = "<|im_start|>user\n" elif "llama" in tokenizer.name_or_path.lower(): response_template = "<|start_header_id|>assistant<|end_header_id|>\n\n" instruction_template = "<|start_header_id|>user<|end_header_id|>\n\n" @@ -290,7 +290,7 @@ def _dataset_to_chat_template( if not found: text = "" - texts.append(text) + texts.append(text.strip()) return {TEXT_FIELD: texts} From a1f601f50b953ac72f0de080e4a2130ff651458a Mon Sep 17 00:00:00 2001 From: Akim Tsvigun Date: Tue, 12 Aug 2025 19:05:02 +0000 Subject: [PATCH 2/4] multiple bugs fixed --- configs/base.yaml | 10 ---------- src/atgen/run_scripts/run_active_learning.py | 8 ++++++-- src/atgen/utils/evaluate_bfcl.py | 4 ++++ 3 files changed, 10 insertions(+), 12 deletions(-) diff --git a/configs/base.yaml b/configs/base.yaml index b501f13..3554a2c 100644 --- a/configs/base.yaml +++ b/configs/base.yaml @@ -90,13 +90,3 @@ evaluation: deepeval_async_mode: True deepeval_verbose_mode: False deepeval_truths_extraction_limit: 10 - -data: - assistant_response_start: ${model.assistant_response_start} - # test_subset_size: 3 - # train_subset_size: 100 - dataset: SpeedOfMagic/gigaword_tiny - input_max_length: 100 - output_max_length: 20 - is_in_conversational_format: false - use_test_benchmark: false diff --git a/src/atgen/run_scripts/run_active_learning.py b/src/atgen/run_scripts/run_active_learning.py index ac713a6..0d9e580 100644 --- a/src/atgen/run_scripts/run_active_learning.py +++ b/src/atgen/run_scripts/run_active_learning.py @@ -343,12 +343,16 @@ def run_active_learning(config, workdir: Union[str, Path]): "No labeled training data available. Skipping training for this iteration." ) train_result = {"training_loss": 0.0, "skipped": True} - del trainer - rmtree(train_output_dir) + model = model.cpu() if config.model.save_in_fp_32: model = model.to(torch.float32) model = model.eval().merge_and_unload() + # Free up memory + del trainer + gc.collect() + torch.cuda.empty_cache() + rmtree(train_output_dir) if not has_test: if dev_split_size > 0: diff --git a/src/atgen/utils/evaluate_bfcl.py b/src/atgen/utils/evaluate_bfcl.py index 96f355c..29d6283 100644 --- a/src/atgen/utils/evaluate_bfcl.py +++ b/src/atgen/utils/evaluate_bfcl.py @@ -10,6 +10,7 @@ from shutil import rmtree from torch import cuda import gc +from torch import cuda from .constants import DEFAULT_NUM_THREADS_BFCL, DEFAULT_GPU_MEMORY_UTILIZATION_BFCL, BFCL_NUM_RETRIES @@ -100,6 +101,9 @@ def _evaluate_bfcl( cwd = os.getcwd() os.chdir(bfcl_results_dir) model_name = save_dir + del model, tokenizer + gc.collect() + cuda.empty_cache() logger.info(f"Starting BFCL evaluation for model: {model_name}") logger.info(f"Test category: {test_category}") From 47c1b4bc173bc4788e43c20a571fa25933980267 Mon Sep 17 00:00:00 2001 From: Akim Tsvigun Date: Sun, 17 Aug 2025 07:27:02 +0000 Subject: [PATCH 3/4] subset selection added --- configs/base.yaml | 9 +- pyproject.toml | 1 + src/atgen/run_scripts/run_active_learning.py | 94 +++-- src/atgen/run_scripts/run_subset_selection.py | 383 ++++++++++++++++++ src/atgen/utils/constants.py | 2 + src/atgen/utils/evaluate_bfcl.py | 11 +- src/atgen/utils/load_model_tokenizer.py | 60 +-- 7 files changed, 482 insertions(+), 78 deletions(-) create mode 100644 src/atgen/run_scripts/run_subset_selection.py diff --git a/configs/base.yaml b/configs/base.yaml index 3554a2c..2f54e5a 100644 --- a/configs/base.yaml +++ b/configs/base.yaml @@ -25,8 +25,9 @@ al: init_query_size: 10 query_size: 10 num_iterations: 5 - evaluate_zero_iteration: True + eval_zero_iteration: True subsample_size: -1 + query_ids_path: null required_performance: rouge1: 0.5 budget: @@ -51,8 +52,8 @@ model: training: dev_split_size: 0.2 hyperparameters: - num_epochs: 5 - train_batch_size: 6 + num_epochs: 15 + train_batch_size: 8 eval_batch_size: 4 gradient_accumulation_steps: 2 lr: 0.00003 @@ -63,7 +64,7 @@ training: model_max_length: ${model.model_max_length} dataset_num_proc: ${data.num_proc} packing: False - lr_scheduler_type: cosine + lr_scheduler_type: linear gradient_checkpointing: True diff --git a/pyproject.toml b/pyproject.toml index bcfe10d..dd36ab8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,7 @@ classifiers = [ [project.scripts] run-al = "atgen.run_scripts.run_active_learning:main" +run-ss = "atgen.run_scripts.run_subset_selection:main" [tool.setuptools.dynamic] dependencies = {file = ["requirements.txt"]} diff --git a/src/atgen/run_scripts/run_active_learning.py b/src/atgen/run_scripts/run_active_learning.py index 0d9e580..f18e462 100644 --- a/src/atgen/run_scripts/run_active_learning.py +++ b/src/atgen/run_scripts/run_active_learning.py @@ -2,8 +2,7 @@ import torch from shutil import rmtree import gc -import json - + import hydra from pathlib import Path from typing import Union @@ -13,8 +12,6 @@ DEFAULT_CONFIG_NAME, UNLABELED_DATA_SPLIT_DEFAULT_NAME, TEST_DATA_SPLIT_DEFAULT_NAME, - OUTPUT_FIELD_PURPOSE_TRAIN, - OUTPUT_FIELD_PURPOSE_TEST, ) log = logging.getLogger() @@ -30,7 +27,6 @@ def run_active_learning(config, workdir: Union[str, Path]): load_data, prepare_conversational_data, maybe_get_few_shot_examples, - get_output_column_name_for_phase, ) from atgen.utils.load_model_tokenizer import load_model_tokenizer from atgen.utils.prepare_model_for_training import prepare_model_for_training @@ -158,6 +154,12 @@ def run_active_learning(config, workdir: Union[str, Path]): init_query_size = config.al.init_query_size + config.data.few_shot.count init_query_size_is_positive = init_query_size > 0 if init_query_size_is_positive: + if config.al.init_query_size == 0: + raise ValueError( + "It's useless to duplicate selection process in the first iteration. " + "Either set `al.init_query_size` to a positive number or set `data.few_shot.count` to 0." + ) + iter_dir = workdir / "iter_0" iter_dir.mkdir(exist_ok=True) @@ -217,38 +219,40 @@ def run_active_learning(config, workdir: Union[str, Path]): model_name=model_name, ) - if has_test: + if has_test and not config.data.use_test_benchmark: print("Preparing test data") + test_data: Dataset = prepare_conversational_data( + dataset=test_data, + data_config=config.data, + split="test", + few_shot_examples=few_shot_examples, + model_name=model_name, + ) + # Evaluate the initial model before any training + # Don't need to evaluate if init_query_size is 0 because we'll get it + # in the first iteration + if init_query_size_is_positive and config.al.eval_zero_iteration: if not config.data.use_test_benchmark: - test_data: Dataset = prepare_conversational_data( - dataset=test_data, + generations: list[str] = generate( + config.inference, + data=test_data, + model=model, + tokenizer=tokenizer, + save_dir=save_dir, data_config=config.data, - split="test", - few_shot_examples=few_shot_examples, - model_name=model_name, + model_config=config.model, + ) + if os.path.exists(save_dir): + rmtree(save_dir) + + metrics: dict[str, float] = compute_metrics( + generated_texts=generations, + reference_texts=test_data[output_column_name_test], + original_texts=test_data[input_column_name], + task=config.data.task, + config=config.evaluation, + cache_dir=cache_dir, ) - # Evaluate the initial model before any training - if init_query_size_is_positive and config.al.evaluate_zero_iteration: - generations: list[str] = generate( - config.inference, - data=test_data, - model=model, - tokenizer=tokenizer, - save_dir=save_dir, - data_config=config.data, - model_config=config.model, - ) - if os.path.exists(save_dir): - rmtree(save_dir) - - metrics: dict[str, float] = compute_metrics( - generated_texts=generations, - reference_texts=test_data[output_column_name_test], - original_texts=test_data[input_column_name], - task=config.data.task, - config=config.evaluation, - cache_dir=cache_dir, - ) else: if "bfcl" in config.data.test_split_name: test_split_name = config.data.test_split_name.split("bfcl_")[1] @@ -270,19 +274,19 @@ def run_active_learning(config, workdir: Union[str, Path]): available_metrics=available_metrics, ) ) - save_log_iter_results( - config=config, - workdir=workdir, - iter_dir=iter_dir, - metrics=metrics, - generations=generations, - al_iter=0, - train_result={}, - model=None, - tokenizer=None, - ) + save_log_iter_results( + config=config, + workdir=workdir, + iter_dir=iter_dir, + metrics=metrics, + generations=generations, + al_iter=0, + train_result={}, + model=None, + tokenizer=None, + ) - # Start AL cycle. Use `num_al_iterations + 2` because we do not label data + # Start AL cycle. Use `num_al_iterations + 1` because we do not label data # but want to train the model on the last iteration. start_iter = 1 if init_query_size_is_positive else 0 diff --git a/src/atgen/run_scripts/run_subset_selection.py b/src/atgen/run_scripts/run_subset_selection.py new file mode 100644 index 0000000..5f12b0d --- /dev/null +++ b/src/atgen/run_scripts/run_subset_selection.py @@ -0,0 +1,383 @@ +import os +import torch +from shutil import rmtree +import gc +import json + +import hydra +from pathlib import Path +from typing import Union +import logging +from atgen.utils.main_decorator import main_decorator +from atgen.utils.constants import ( + DEFAULT_CONFIG_NAME, + UNLABELED_DATA_SPLIT_DEFAULT_NAME, + TEST_DATA_SPLIT_DEFAULT_NAME, + NUM_PROCS_FOR_DATASETS +) + +log = logging.getLogger() + + +@main_decorator +def run_subset_selection(config, workdir: Union[str, Path]): + from transformers import set_seed + from datasets import concatenate_datasets, Dataset + + from atgen.metrics.compute_metrics import compute_metrics + from atgen.utils.data import ( + load_data, + prepare_conversational_data, + maybe_get_few_shot_examples, + ) + from atgen.utils.load_model_tokenizer import load_model_tokenizer + from atgen.utils.prepare_model_for_training import prepare_model_for_training + from atgen.utils.training_utils import get_trainer + from atgen.strategies.get_strategy import get_strategy + from atgen.labellers import get_labeller + from atgen.utils.generate import generate + from atgen.utils.check_required_performance import check_required_performance + from atgen.utils.save_labeled_data import save_labeled_data + from atgen.utils.save_log_iter_results import save_log_iter_results + from atgen.strategies.base_strategy import BaseStrategy + from atgen.labellers.base_labeller import BaseLabeler + from atgen.utils.check_performance_metrics import ( + check_performance_against_requirements, + ) + from atgen.utils.evaluate_bfcl import evaluate_bfcl + + seed = config.seed + cache_dir = config.cache_dir + dev_split_size = config.training.dev_split_size + output_column_name_train = config.data.train_output_column_name + output_column_name_test = config.data.test_output_column_name + + model_name = config.model.checkpoint + + num_al_iterations = config.al.num_iterations + required_performance_dict = check_required_performance( + required_performance=config.al.required_performance + ) + budget = config.al.budget + if budget is None: + budget = 1e10 + + # Stopping criteria due to reaching required performance + is_performance_reached = False + + # Initialize variables for tracking available metrics + available_metrics = {} + is_metrics_availability_checked = False + + has_test = ( + config.data.test_split_name is not None and config.data.test_split_name != "" + ) + + print( + f"""Running Active Learning... +AL Strategy: {config.al.strategy} +Num Iterations: {num_al_iterations} +Query Size: {config.al.query_size} +Dataset: {config.data.dataset if isinstance(config.data.dataset, str) else 'custom'} +Seed: {seed} +Model: {model_name} +Config: {config.name} +Prompt:\n{config.data.system_prompt} +""" + ) + + if isinstance(workdir, str): + workdir = Path(workdir) + train_output_dir = workdir / "tmp" + save_dir = workdir / "tmp_best" + + print("Loading data.") + unlabeled_data = load_data( + data_config=config.data, + split=UNLABELED_DATA_SPLIT_DEFAULT_NAME, + cache_dir=config.cache_dir, + seed=seed, + ) + if has_test and not config.data.use_test_benchmark: + test_data = load_data( + data_config=config.data, + split=TEST_DATA_SPLIT_DEFAULT_NAME, + cache_dir=config.cache_dir, + seed=seed, + ) + # TODO: make better. Current workaround for multi-choice QA. + if config.data.get("processed_input_column_name", None) is not None: + config.data.input_column_name = config.data.processed_input_column_name + input_column_name = config.data.input_column_name + # After loading data, need to calculate the query size if it is proportional to the dataset size + if config.al.init_query_size is None: + config.al.init_query_size = int(len(unlabeled_data) * config.al.init_query_size) + log.info(f"Setting init query size to {config.al.init_query_size}") + if isinstance(config.al.query_size, float): + config.al.query_size = int(len(unlabeled_data) * config.al.query_size) + log.info(f"Setting query size to {config.al.query_size}") + al_query_size = config.al.query_size + + print("Initial iteration: loading model & tokenizer...") + model, tokenizer = load_model_tokenizer( + checkpoint=model_name, model_config=config.model, cache_dir=cache_dir + ) + + if config.al.query_ids_path: + with open(config.al.query_ids_path, "r") as f: + labeled_ids = json.load(f)[:al_query_size] + else: + print("Loading subset selection strategy...") + ss_strategy: BaseStrategy = get_strategy( + strategy_name=config.al.strategy, + subsample_size=config.al.subsample_size, + unlabeled_pool=unlabeled_data[input_column_name], + model=model, + tokenizer=tokenizer, + inference_config=config.inference, # for hadas + model_config=config.model, # for hadas + data_config=config.data, # for hadas + cache_dir=cache_dir, # for hadas, huds, graph_cut + seed=seed, + **config.al.strategy_kwargs, + ) + + print("Loading labeller...") + # TODO: unsure whether need to log here since may be confusing for a human labeller + labeller: BaseLabeler = get_labeller( + config.labeller, + output_column_name=output_column_name_train, + cache_dir=cache_dir, + budget=budget, + workdir=workdir, # if labeller is a human + data_config=config.data, # if labeller is a custom LLM on transformers + model_config=config.model, # if labeller is a custom LLM on transformers + ) + print("Calculating query_ids") + labeled_ids: list[int] = ss_strategy( + model=model, + tokenizer=tokenizer, + unlabeled_pool=unlabeled_data.remove_columns(output_column_name_train), + labeled_pool=None, + num_to_label=al_query_size, + batch_size=config.inference.batch_size, + max_new_tokens=config.inference.max_new_tokens, + ) + + query: Dataset = unlabeled_data.filter(lambda x: x["id"] in labeled_ids) + labeled_data: Dataset = labeller(query) + if labeller.is_out_of_budget: + labeled_data = labeled_data.filter(lambda x: x[labeller.output_column_name] != "", batched=False, num_proc=NUM_PROCS_FOR_DATASETS) + print(f"Labeler ran out of budget at iteration 0.") + + # Get the few-shot examples + few_shot_examples, labeled_data = maybe_get_few_shot_examples( + config=config, labeled_data=labeled_data, workdir=workdir + ) + + iter_dir = workdir / "iter_0" + iter_dir.mkdir(exist_ok=True) + + print(f"Saving labeled data...") + save_labeled_data( + labeled_data=labeled_data, + labeled_query=labeled_data, + workdir=workdir, + iter_dir=iter_dir, + labeled_ids=labeled_ids, + query_ids=labeled_ids, + ) + + if has_test and not config.data.use_test_benchmark: + print("Preparing test data") + test_data: Dataset = prepare_conversational_data( + dataset=test_data, + data_config=config.data, + split="test", + few_shot_examples=few_shot_examples, + model_name=model_name, + ) + # Evaluate the initial model before any training + if config.al.eval_zero_iteration: + if not config.data.use_test_benchmark: + generations: list[str] = generate( + config.inference, + data=test_data, + model=model, + tokenizer=tokenizer, + save_dir=save_dir, + data_config=config.data, + model_config=config.model, + ) + if os.path.exists(save_dir): + rmtree(save_dir) + + metrics: dict[str, float] = compute_metrics( + generated_texts=generations, + reference_texts=test_data[output_column_name_test], + original_texts=test_data[input_column_name], + task=config.data.task, + config=config.evaluation, + cache_dir=cache_dir, + ) + else: + if "bfcl" in config.data.test_split_name: + test_split_name = config.data.test_split_name.split("bfcl_")[1] + else: + raise NotImplementedError(f"Test split name {config.data.test_split_name} is not supported") + generations, metrics = evaluate_bfcl( + model_name=model_name, + bfcl_results_dir=iter_dir, + test_category=test_split_name, + num_threads=config.inference.num_threads_for_bfcl + ) + save_log_iter_results( + config=config, + workdir=workdir, + iter_dir=iter_dir, + metrics=metrics, + generations=generations, + al_iter=0, + train_result={}, + model=None, + tokenizer=None, + ) + + # Start AL cycle. Use `num_al_iterations + 2` because we do not label data + # but want to train the model on the last iteration. + + al_iter = 1 if config.al.eval_zero_iteration else 0 + iter_dir = workdir / ("iter_" + str(al_iter)) + iter_dir.mkdir(exist_ok=True) + + if not config.data.is_in_conversational_format: + train_eval_data = prepare_conversational_data( + dataset=labeled_data, + data_config=config.data, + split="train", + few_shot_examples=few_shot_examples, + model_name=model_name, + ) + else: + train_eval_data = labeled_data + + if dev_split_size > 0 and len(train_eval_data) > 1: + train_eval_data = train_eval_data.train_test_split( + test_size=dev_split_size, shuffle=True, seed=seed + ) + train_data = train_eval_data["train"] + eval_data = train_eval_data["test"] + else: + train_data = train_eval_data + eval_data = None + + model = prepare_model_for_training(model, config.model.peft) + + # Set seed for reproducibility + set_seed(seed) + trainer = get_trainer( + config=config, + model=model, + tokenizer=tokenizer, + train_data=train_data, + eval_data=eval_data, + output_dir=train_output_dir, + seed=seed, + ) + + # Launch training + if len(train_data) > 0: + train_result = trainer.train() + print(f"Training completed with {len(train_data)} examples") + else: + log.warning( + "No labeled training data available. Skipping training for this iteration." + ) + train_result = {"training_loss": 0.0, "skipped": True} + + model = model.cpu() + if config.model.save_in_fp_32: + model = model.to(torch.float32) + model = model.eval().merge_and_unload() + # Free up memory + del trainer + gc.collect() + torch.cuda.empty_cache() + rmtree(train_output_dir) + + if not has_test: + if dev_split_size > 0: + test_data = eval_data + else: + if not config.data.use_test_benchmark: + generations: list[str] = generate( + config.inference, + data=test_data, + model=model, + tokenizer=tokenizer, + save_dir=save_dir, + data_config=config.data, + model_config=config.model, + ) + if os.path.exists(save_dir): + rmtree(save_dir) + + metrics: dict[str, float] = compute_metrics( + generated_texts=generations, + reference_texts=test_data[output_column_name_test], + original_texts=test_data[input_column_name], + task=config.data.task, + config=config.evaluation, + cache_dir=cache_dir, + ) + else: + if "bfcl" in config.data.test_split_name: + test_split_name = config.data.test_split_name.split("bfcl_")[1] + else: + raise NotImplementedError(f"Test split name {config.data.test_split_name} is not supported") + generations, metrics = evaluate_bfcl( + model_name=model_name, + bfcl_results_dir=iter_dir, + model=model, + tokenizer=tokenizer, + test_category=test_split_name, + num_threads=config.inference.num_threads_for_bfcl, + ) + save_log_iter_results( + config=config, + workdir=workdir, + iter_dir=iter_dir, + metrics=metrics, + generations=generations, + al_iter=al_iter, + train_result=train_result, + model=model, + tokenizer=tokenizer, + ) + + + print("Subset selection is done.") + + +@hydra.main( + config_path=os.environ.get("HYDRA_CONFIG_PATH", os.getcwd() + "/configs/"), + config_name=os.environ.get("HYDRA_CONFIG_NAME", DEFAULT_CONFIG_NAME), + version_base="1.1", +) +def main(config): + if getattr(config, "debug", True): + try: + run_subset_selection(config) + except Exception as e: + print(e) + import pdb + import sys + + exc_type, exc_value, exc_traceback = sys.exc_info() + pdb.post_mortem(exc_traceback) + else: + run_subset_selection(config) + + +if __name__ == "__main__": + main() diff --git a/src/atgen/utils/constants.py b/src/atgen/utils/constants.py index afa22cd..bec6245 100644 --- a/src/atgen/utils/constants.py +++ b/src/atgen/utils/constants.py @@ -17,3 +17,5 @@ DEFAULT_NUM_THREADS_BFCL = 32 DEFAULT_GPU_MEMORY_UTILIZATION_BFCL = 0.75 BFCL_NUM_RETRIES = 3 + +NUM_PROCS_FOR_DATASETS = 16 diff --git a/src/atgen/utils/evaluate_bfcl.py b/src/atgen/utils/evaluate_bfcl.py index 29d6283..2825aab 100644 --- a/src/atgen/utils/evaluate_bfcl.py +++ b/src/atgen/utils/evaluate_bfcl.py @@ -101,9 +101,11 @@ def _evaluate_bfcl( cwd = os.getcwd() os.chdir(bfcl_results_dir) model_name = save_dir - del model, tokenizer - gc.collect() - cuda.empty_cache() + # Free up memory + del model, tokenizer + gc.collect() + cuda.empty_cache() + model = None logger.info(f"Starting BFCL evaluation for model: {model_name}") logger.info(f"Test category: {test_category}") @@ -207,7 +209,8 @@ def _extract_metrics(bfcl_results_dir: Path) -> dict[str, float]: [v for k, v in non_live_metrics.items() if "simple" in k.lower()] ) non_live_metrics["Non-live Overall"] = np.mean([ - v for k, v in non_live_metrics.items() if k.lower() in NON_LIVE_COLUMNS_FOR_OVERALL + v for k, v in non_live_metrics.items() + if k.strip("Non-live ") in NON_LIVE_COLUMNS_FOR_OVERALL ]) live_metrics = pd.read_csv(bfcl_results_dir / "score" / "data_live.csv").iloc[-1, 2:].dropna().to_dict() diff --git a/src/atgen/utils/load_model_tokenizer.py b/src/atgen/utils/load_model_tokenizer.py index 2d33e85..4aabaaf 100644 --- a/src/atgen/utils/load_model_tokenizer.py +++ b/src/atgen/utils/load_model_tokenizer.py @@ -1,4 +1,5 @@ from pathlib import Path +from time import sleep import torch @@ -22,6 +23,9 @@ class FastLanguageModel: ) +NUM_TRIES_LOAD_MODEL = 3 + + # def load_tokenizer( # model_config: DictConfig, # cache_dir: str, @@ -44,32 +48,38 @@ def load_model_tokenizer( cache_dir: str, ) -> tuple[FastLanguageModel, PreTrainedTokenizerFast]: dtype = getattr(torch, model_config.dtype, torch.bfloat16) - if torch.cuda.is_available(): - model, tokenizer = FastLanguageModel.from_pretrained( - model_name=checkpoint, - max_seq_length=model_config.model_max_length, - dtype=dtype, - load_in_4bit=model_config.quantize, - cache_dir=cache_dir, - trust_remote_code=True, - ) - tokenizer.model_max_length = model_config.model_max_length - tokenizer.padding_side = "left" - else: - kwargs = { - "cache_dir": cache_dir, - "torch_dtype": model_config.dtype, - "trust_remote_code": True, - } - model = AutoModelForCausalLM.from_pretrained(checkpoint, **kwargs) - - tokenizer = AutoTokenizer.from_pretrained( - model_config.checkpoint, - model_max_length=model_config.model_max_length, - cache_dir=cache_dir, - padding_side="left" - ) + n_tries = 0 + model = None + while (model is None) and (n_tries < NUM_TRIES_LOAD_MODEL): + sleep(10 * n_tries**2) + if torch.cuda.is_available(): + model, tokenizer = FastLanguageModel.from_pretrained( + model_name=checkpoint, + max_seq_length=model_config.model_max_length, + dtype=dtype, + load_in_4bit=model_config.quantize, + cache_dir=cache_dir, + trust_remote_code=True, + ) + tokenizer.model_max_length = model_config.model_max_length + tokenizer.padding_side = "left" + else: + kwargs = { + "cache_dir": cache_dir, + "torch_dtype": model_config.dtype, + "trust_remote_code": True, + } + model = AutoModelForCausalLM.from_pretrained(checkpoint, **kwargs) + tokenizer = AutoTokenizer.from_pretrained( + model_config.checkpoint, + model_max_length=model_config.model_max_length, + cache_dir=cache_dir, + padding_side="left" + ) + n_tries += 1 + if model is None: + raise RuntimeError(f"Failed to load model after {NUM_TRIES_LOAD_MODEL} tries") if tokenizer.pad_token is None: last_reserved_token = {v: k for k, v in tokenizer.vocab.items()}[ len(tokenizer) - 1 From f83b7534a7f19690d3726e7a315879159ede0327 Mon Sep 17 00:00:00 2001 From: Akim Tsvigun Date: Sun, 17 Aug 2025 08:17:51 +0000 Subject: [PATCH 4/4] code blacked; bugs fixed --- configs/base.yaml | 1 - configs/data/hermes_tool.yaml | 15 +++ requirements.txt | 2 +- .../api_labellers/openai_labeller.py | 3 +- src/atgen/metrics/compute_metrics.py | 102 +++++++++------ src/atgen/run_scripts/run_active_learning.py | 48 ++++--- src/atgen/run_scripts/run_subset_selection.py | 34 +++-- src/atgen/strategies/__init__.py | 8 ++ src/atgen/strategies/hadas.py | 4 +- src/atgen/utils/check_performance_metrics.py | 5 +- src/atgen/utils/data/load_data.py | 24 +++- .../utils/data/prepare_conversational_data.py | 14 ++- src/atgen/utils/downloaders.py | 9 +- src/atgen/utils/evaluate_bfcl.py | 118 +++++++++++------- src/atgen/utils/generate.py | 19 +-- src/atgen/utils/load_model_tokenizer.py | 2 +- src/atgen/utils/main_decorator.py | 1 + src/atgen/utils/post_process_generations.py | 18 ++- src/atgen/utils/resolvers.py | 1 + src/atgen/utils/training_utils.py | 16 +-- src/atgen/utils/validate_and_fill_config.py | 8 +- 21 files changed, 294 insertions(+), 158 deletions(-) create mode 100644 configs/data/hermes_tool.yaml diff --git a/configs/base.yaml b/configs/base.yaml index 2f54e5a..e8e9f04 100644 --- a/configs/base.yaml +++ b/configs/base.yaml @@ -27,7 +27,6 @@ al: num_iterations: 5 eval_zero_iteration: True subsample_size: -1 - query_ids_path: null required_performance: rouge1: 0.5 budget: diff --git a/configs/data/hermes_tool.yaml b/configs/data/hermes_tool.yaml new file mode 100644 index 0000000..7affddd --- /dev/null +++ b/configs/data/hermes_tool.yaml @@ -0,0 +1,15 @@ +dataset: 'Aktsvigun/xlam-function-calling-60k' +input_column_name: 'input' +output_column_name: 'answers' +unlabeled_data_split_name: train +test_split_name: bfcl_python +train_subset_size: null +test_subset_size: null +input_max_length: 4096 +output_max_length: 512 +fetch_kwargs: {} +is_in_conversational_format: false +system_prompt: "" +assistant_response_start: ${model.assistant_response_start} +use_test_benchmark: true +task: 'fc' diff --git a/requirements.txt b/requirements.txt index da23ecf..02b6a42 100644 --- a/requirements.txt +++ b/requirements.txt @@ -32,5 +32,5 @@ transformers>=4.52.4 trl==0.19.1 torchmetrics==1.4.1 unsloth==2025.8.1 -vllm==0.10.0 +# vllm==0.10.0 xlrd==1.2.0 diff --git a/src/atgen/labellers/api_labellers/openai_labeller.py b/src/atgen/labellers/api_labellers/openai_labeller.py index 8ebef95..1f22736 100644 --- a/src/atgen/labellers/api_labellers/openai_labeller.py +++ b/src/atgen/labellers/api_labellers/openai_labeller.py @@ -22,6 +22,7 @@ MAX_NUM_TRIES = 3 UPDATE_TIME_IN_SECONDS = 10 # update time when checking for the completion + class OpenAILabeller(BaseLabeler): def __init__( self, @@ -105,7 +106,7 @@ def _sync_call(self, dataset: Dataset) -> Dataset: # TODO: make a parameter for i, annotation in enumerate(annotations[:2]): print(f"Annotation of the {i+1}th instance: {annotation}") - print("-"*100) + print("-" * 100) return dataset def _batched_call(self, dataset: Dataset) -> Dataset: diff --git a/src/atgen/metrics/compute_metrics.py b/src/atgen/metrics/compute_metrics.py index 2864372..c4c9e54 100644 --- a/src/atgen/metrics/compute_metrics.py +++ b/src/atgen/metrics/compute_metrics.py @@ -81,13 +81,24 @@ def compute_metrics( elif task == "open-qa": metrics_to_calculate = ["exact_match"] + list(config.additional_metrics) elif task == "summarization": - metrics_to_calculate = ["exact_match", "sacrebleu", "bleu", "rouge", "word_length"] + list(config.additional_metrics) + metrics_to_calculate = [ + "exact_match", + "sacrebleu", + "bleu", + "rouge", + "word_length", + ] + list(config.additional_metrics) elif task == "translation": - metrics_to_calculate = ["exact_match", "sacrebleu", "bleu", "word_length"] + list(config.additional_metrics) + metrics_to_calculate = [ + "exact_match", + "sacrebleu", + "bleu", + "word_length", + ] + list(config.additional_metrics) elif task == "math": metrics_to_calculate = ["exact_match_math"] + list(config.additional_metrics) else: - raise NotImplementedError(f"Task {task} not implemented") + raise NotImplementedError(f"Task {task} not implemented") if "sacrebleu" in metrics_to_calculate: sacrebleu = load("sacrebleu", cache_dir=cache_dir) @@ -124,13 +135,19 @@ def compute_metrics( if isinstance(reference_texts[0], list): result["exact_match"] = np.array( [ - any(_preprocess_text(pred) == _preprocess_text(one_ref) for one_ref in ref) + any( + _preprocess_text(pred) == _preprocess_text(one_ref) + for one_ref in ref + ) for pred, ref in zip(generated_texts, reference_texts) ] ) else: result["exact_match"] = np.array( - [_preprocess_text(pred) == _preprocess_text(ref) for pred, ref in zip(generated_texts, reference_texts)] + [ + _preprocess_text(pred) == _preprocess_text(ref) + for pred, ref in zip(generated_texts, reference_texts) + ] ) if "exact_match_math" in metrics_to_calculate: # result["exact_match_math"] = np.array( @@ -195,10 +212,14 @@ def compute_metrics( ] ) else: - ref_word_lengths = np.array([len(ref.split()) for ref in reference_texts]) + ref_word_lengths = np.array( + [len(ref.split()) for ref in reference_texts] + ) # Avoid division by zero ref_word_lengths_safe = np.where(ref_word_lengths > 0, ref_word_lengths, 1) - result["word_length_rel"] = result["word_length_gen"] / ref_word_lengths_safe + result["word_length_rel"] = ( + result["word_length_gen"] / ref_word_lengths_safe + ) # AlignScore if "alignscore" in metrics_to_calculate and is_alignscore_available: @@ -275,31 +296,42 @@ def compute_metrics( return result -def _preprocess_text(text: str, do_lowercase: bool = True, do_remove_punctuation: bool = True, do_remove_extra_spaces: bool = True, do_remove_stopwords: bool = False, stopwords: Optional[list[str]] = None) -> str: - # Convert to lowercase - if do_lowercase: - text = text.lower() - - # Remove punctuation - if do_remove_punctuation: - # Keep hyphens within words, remove other punctuation - text = re.sub(r'(? str: + # Convert to lowercase + if do_lowercase: + text = text.lower() + + # Remove punctuation + if do_remove_punctuation: + # Keep hyphens within words, remove other punctuation + text = re.sub(r"(? Dataset: + +def _preprocess_multi_choice_qa( + dataset: Dataset, data_config: DictConfig, split: str +) -> Dataset: alphabet_titled = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" input_column_names = data_config.input_column_name options_column_name = input_column_names["options"] @@ -143,16 +148,21 @@ def _preprocess_multi_choice_qa(dataset: Dataset, data_config: DictConfig, split system_prompt = data_config.system_prompt messages = [] for inst in dataset: - preprocessed_options = '' + preprocessed_options = "" for option, letter in zip(inst[options_column_name], alphabet_titled): preprocessed_options += f"- {letter}. {option}\n" user_prompt_kwargs = { - key: inst[key] for key in input_column_names.keys() if key != options_column_name + key: inst[key] + for key in input_column_names.keys() + if key != options_column_name } user_prompt_kwargs["options"] = preprocessed_options inst_messages = [ {"role": "system", "content": system_prompt}, - {"role": "user", "content": user_prompt_template.format(**user_prompt_kwargs)}, + { + "role": "user", + "content": user_prompt_template.format(**user_prompt_kwargs), + }, ] if split == "train": inst_messages.append({"role": "assistant", "content": inst["answer"]}) @@ -161,6 +171,8 @@ def _preprocess_multi_choice_qa(dataset: Dataset, data_config: DictConfig, split dataset = dataset.remove_columns(["messages"]) dataset = dataset.add_column("messages", messages) # Can't directly update `processed_input_column_name` because test data is loaded separately - OmegaConf.update(data_config, "processed_input_column_name", "messages", force_add=True) + OmegaConf.update( + data_config, "processed_input_column_name", "messages", force_add=True + ) data_config.is_in_conversational_format = True return dataset diff --git a/src/atgen/utils/data/prepare_conversational_data.py b/src/atgen/utils/data/prepare_conversational_data.py index 2bc07f5..0ce2f1a 100644 --- a/src/atgen/utils/data/prepare_conversational_data.py +++ b/src/atgen/utils/data/prepare_conversational_data.py @@ -19,10 +19,18 @@ def prepare_conversational_data( few_shot_messages = [] if few_shot_examples: - for (fs_input, fs_output) in zip(few_shot_examples[input_column_name], few_shot_examples[output_column_name]): + for fs_input, fs_output in zip( + few_shot_examples[input_column_name], few_shot_examples[output_column_name] + ): few_shot_messages.append({"role": "user", "content": fs_input}) - response_start = data_config.assistant_response_start if data_config.assistant_response_start else "" - few_shot_messages.append({"role": "assistant", "content": response_start + fs_output}) + response_start = ( + data_config.assistant_response_start + if data_config.assistant_response_start + else "" + ) + few_shot_messages.append( + {"role": "assistant", "content": response_start + fs_output} + ) # Get appropriate preprocessing function based on all parameters preprocess_fn = get_preprocess_function( model_name=model_name, diff --git a/src/atgen/utils/downloaders.py b/src/atgen/utils/downloaders.py index c41e3a6..55099f7 100644 --- a/src/atgen/utils/downloaders.py +++ b/src/atgen/utils/downloaders.py @@ -1,19 +1,24 @@ import subprocess import os + def _download_spacy(): try: import spacy + spacy.load("en_core_web_sm") print("en_core_web_sm model already installed") except (ImportError, OSError): print("Installing en_core_web_sm model...") subprocess.run("python -m spacy download en_core_web_sm", shell=True) + def _download_nltk(): import nltk - nltk.download('punkt') - nltk.download('punkt_tab') + + nltk.download("punkt") + nltk.download("punkt_tab") + def maybe_download_packages(outputs_dir: str): if not os.path.exists(outputs_dir): diff --git a/src/atgen/utils/evaluate_bfcl.py b/src/atgen/utils/evaluate_bfcl.py index 2825aab..dc0d844 100644 --- a/src/atgen/utils/evaluate_bfcl.py +++ b/src/atgen/utils/evaluate_bfcl.py @@ -12,13 +12,17 @@ import gc from torch import cuda -from .constants import DEFAULT_NUM_THREADS_BFCL, DEFAULT_GPU_MEMORY_UTILIZATION_BFCL, BFCL_NUM_RETRIES +from .constants import ( + DEFAULT_NUM_THREADS_BFCL, + DEFAULT_GPU_MEMORY_UTILIZATION_BFCL, + BFCL_NUM_RETRIES, +) # Set up logging to output to stdout logging.basicConfig( level=logging.INFO, - format='%(asctime)s - %(levelname)s - %(message)s', - handlers=[logging.StreamHandler()] + format="%(asctime)s - %(levelname)s - %(message)s", + handlers=[logging.StreamHandler()], ) logger = logging.getLogger(__name__) @@ -30,6 +34,7 @@ "Irrelevance Detection", ] + def evaluate_bfcl( model_name: str, bfcl_results_dir: str | Path, @@ -40,7 +45,7 @@ def evaluate_bfcl( ) -> tuple[list[str], dict[str, float]]: """ Evaluate a model on the BFCL benchmark. - + Args: model_name: The name of the model to evaluate bfcl_results_dir: Directory to store the results of the evaluation @@ -48,13 +53,13 @@ def evaluate_bfcl( tokenizer: The tokenizer to use for the model test_category: The test category (default: "python") num_threads: Number of threads to use (default: DEFAULT_NUM_THREADS_BFCL) - + Returns: - Dictionary containing execution results + Dictionary containing execution results """ cuda.empty_cache() gc.collect() - + for _ in range(BFCL_NUM_RETRIES): try: return _evaluate_bfcl( @@ -63,13 +68,14 @@ def evaluate_bfcl( model=model, tokenizer=tokenizer, test_category=test_category, - num_threads=num_threads + num_threads=num_threads, ) except Exception as e: logger.error(f"Error evaluating BFCL: {e}") continue raise Exception("Failed to evaluate BFCL") + def _evaluate_bfcl( model_name: str, bfcl_results_dir: str | Path, @@ -80,7 +86,7 @@ def _evaluate_bfcl( ) -> tuple[list[str], dict[str, float]]: """ Evaluate a model on the BFCL benchmark. - + Args: model_name: The name of the model to evaluate bfcl_results_dir: Directory to store the results of the evaluation @@ -88,7 +94,7 @@ def _evaluate_bfcl( tokenizer: The tokenizer to use for the model test_category: The test category (default: "python") num_threads: Number of threads to use (default: DEFAULT_NUM_THREADS_BFCL) - + Returns: Dictionary containing execution results """ @@ -111,50 +117,61 @@ def _evaluate_bfcl( logger.info(f"Test category: {test_category}") logger.info(f"Number of threads: {num_threads}") logger.info(f"BFCL project root: {bfcl_results_dir}") - + # Set environment variables env = os.environ.copy() env["BFCL_PROJECT_ROOT"] = bfcl_results_dir env["TEST_CATEGORY"] = test_category env["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY", "kek") env["OPENAI_API_BASE"] = os.getenv("OPENAI_API_BASE", "http://localhost:8000") - + logger.info("Environment variables set successfully") - + # Run bfcl generate generate_cmd = [ - "bfcl", "generate", - "--model", model_name, - "--test-category", test_category, - "--num-threads", str(num_threads), - "--gpu-memory-utilization", str(DEFAULT_GPU_MEMORY_UTILIZATION_BFCL), + "bfcl", + "generate", + "--model", + model_name, + "--test-category", + test_category, + "--num-threads", + str(num_threads), + "--gpu-memory-utilization", + str(DEFAULT_GPU_MEMORY_UTILIZATION_BFCL), ] - + logger.info(f"Running BFCL generate command: {' '.join(generate_cmd)}") # generate_result = subprocess.run(generate_cmd, env=env, cwd=bfcl_results_dir) - with open(bfcl_results_dir / "generate_stdout.log", "w") as stdout_file, open(bfcl_results_dir / "generate_stderr.log", "w") as stderr_file: + with ( + open(bfcl_results_dir / "generate_stdout.log", "w") as stdout_file, + open(bfcl_results_dir / "generate_stderr.log", "w") as stderr_file, + ): generate_result = subprocess.run( - generate_cmd, - env=env, + generate_cmd, + env=env, cwd=bfcl_results_dir, stdout=stdout_file, - stderr=stderr_file + stderr=stderr_file, ) logger.info("BFCL generate completed") - + # Run bfcl evaluate evaluate_cmd = [ - "bfcl", "evaluate", - "--model", model_name, - "--test-category", test_category + "bfcl", + "evaluate", + "--model", + model_name, + "--test-category", + test_category, ] - + logger.info(f"Running BFCL evaluate command: {' '.join(evaluate_cmd)}") evaluate_result = subprocess.run(evaluate_cmd, env=env, cwd=bfcl_results_dir) logger.info("BFCL evaluate completed") - + logger.info("BFCL evaluation finished successfully") - + metrics = _extract_metrics(bfcl_results_dir) generations = _get_generations(model_name, Path(bfcl_results_dir)) @@ -179,44 +196,57 @@ def _get_generations(model_name: str, bfcl_results_dir: Path): def main(): """Example usage of the evaluate_bfcl function.""" import sys - + model_name = sys.argv[1] if len(sys.argv) > 1 else "default_model" bfcl_results_dir = sys.argv[2] if len(sys.argv) > 2 else "tmp" test_category = sys.argv[3] if len(sys.argv) > 3 else "python" num_threads = int(sys.argv[4]) if len(sys.argv) > 4 else 32 - + logger.info("Starting BFCL evaluation script") - + results = evaluate_bfcl( model_name=model_name, bfcl_results_dir=bfcl_results_dir, test_category=test_category, - num_threads=num_threads + num_threads=num_threads, ) - + logger.info("=== BFCL Evaluation Results ===") logger.info(results) + def _extract_metrics(bfcl_results_dir: Path) -> dict[str, float]: - non_live_metrics = pd.read_csv(bfcl_results_dir / "score" / "data_non_live.csv").iloc[-1, 2:].dropna().to_dict() + non_live_metrics = ( + pd.read_csv(bfcl_results_dir / "score" / "data_non_live.csv") + .iloc[-1, 2:] + .dropna() + .to_dict() + ) non_live_metrics = { "Non-live " + k: float(v.replace("%", "")) for k, v in non_live_metrics.items() - if 'overall' not in k.lower() + if "overall" not in k.lower() } if not "Simple AST" in non_live_metrics.keys(): non_live_metrics["Simple AST"] = np.mean( [v for k, v in non_live_metrics.items() if "simple" in k.lower()] ) - non_live_metrics["Non-live Overall"] = np.mean([ - v for k, v in non_live_metrics.items() - if k.strip("Non-live ") in NON_LIVE_COLUMNS_FOR_OVERALL - ]) + non_live_metrics["Non-live Overall"] = np.mean( + [ + v + for k, v in non_live_metrics.items() + if k.strip("Non-live ") in NON_LIVE_COLUMNS_FOR_OVERALL + ] + ) - live_metrics = pd.read_csv(bfcl_results_dir / "score" / "data_live.csv").iloc[-1, 2:].dropna().to_dict() + live_metrics = ( + pd.read_csv(bfcl_results_dir / "score" / "data_live.csv") + .iloc[-1, 2:] + .dropna() + .to_dict() + ) live_metrics = { - "Live " + k: float(v.replace("%", "")) - for k, v in live_metrics.items() + "Live " + k: float(v.replace("%", "")) for k, v in live_metrics.items() } non_live_metrics.update(live_metrics) return non_live_metrics diff --git a/src/atgen/utils/generate.py b/src/atgen/utils/generate.py index 7bd7638..02f7f41 100644 --- a/src/atgen/utils/generate.py +++ b/src/atgen/utils/generate.py @@ -109,7 +109,7 @@ def generate_vllm( generations=generations, data_config=data_config, model_name=llm_runner.llm_engine.model_config.model, - framework=VLLM_FRAMEWORK + framework=VLLM_FRAMEWORK, ) if delete_vllm_after_inference: del llm_runner @@ -210,7 +210,7 @@ def generate_response(s, messages): generations=generations, data_config=data_config, model_name=model_path, - framework=SGLANG_FRAMEWORK + framework=SGLANG_FRAMEWORK, ) # Clean up engine.shutdown() @@ -258,9 +258,7 @@ def generate_transformers( batch["input_ids"].to(model.device), attention_mask=batch["attention_mask"].to(model.device), max_new_tokens=inference_config.max_new_tokens, - temperature=inference_config.get( - "temperature", DEFAULT_TEMPERATURE - ), + temperature=inference_config.get("temperature", DEFAULT_TEMPERATURE), top_p=inference_config.get("top_p", DEFAULT_TOP_P), return_dict_in_generate=True, output_scores=True, @@ -280,7 +278,7 @@ def generate_transformers( generations=generations, data_config=data_config, model_name=model.name_or_path, - framework=TRANSFORMERS_FRAMEWORK + framework=TRANSFORMERS_FRAMEWORK, ) _maybe_display_generations(generations, inference_config.num_display_generations) return generations @@ -337,12 +335,17 @@ def tokenize_conversational_example( example: dict[str, Any], tokenizer: PreTrainedTokenizer, data_config: DictConfig ) -> dict[str, list[int]]: if data_config.assistant_response_start: - input_ids = tokenizer.apply_chat_template(example["messages"], continue_final_message=True) + input_ids = tokenizer.apply_chat_template( + example["messages"], continue_final_message=True + ) else: - input_ids = tokenizer.apply_chat_template(example["messages"], add_generation_prompt=True) + input_ids = tokenizer.apply_chat_template( + example["messages"], add_generation_prompt=True + ) attention_mask = [1 for _ in range(len(input_ids))] return {"input_ids": input_ids, "attention_mask": attention_mask} + def _maybe_display_generations(generations: list[str], num_display_gens: int): if num_display_gens: log.info("Displaying the first {} generations:".format(num_display_gens)) diff --git a/src/atgen/utils/load_model_tokenizer.py b/src/atgen/utils/load_model_tokenizer.py index 4aabaaf..8804311 100644 --- a/src/atgen/utils/load_model_tokenizer.py +++ b/src/atgen/utils/load_model_tokenizer.py @@ -75,7 +75,7 @@ def load_model_tokenizer( model_config.checkpoint, model_max_length=model_config.model_max_length, cache_dir=cache_dir, - padding_side="left" + padding_side="left", ) n_tries += 1 if model is None: diff --git a/src/atgen/utils/main_decorator.py b/src/atgen/utils/main_decorator.py index 40e9a6c..37c7d67 100644 --- a/src/atgen/utils/main_decorator.py +++ b/src/atgen/utils/main_decorator.py @@ -45,6 +45,7 @@ def run_script(config): os.environ["PYTHONHASHSEED"] = str(config.seed) from transformers import set_seed + set_seed(config.seed) func(config, workdir=Path(auto_generated_dir)) diff --git a/src/atgen/utils/post_process_generations.py b/src/atgen/utils/post_process_generations.py index fd5425a..4477191 100644 --- a/src/atgen/utils/post_process_generations.py +++ b/src/atgen/utils/post_process_generations.py @@ -3,15 +3,21 @@ from .constants import REASONING_END_TOKEN + def post_process_generations( - generations: list[str], - data_config: DictConfig, - model_name: Optional[str] = None, - framework: Literal["vllm", "sglang", "transformers"] = "vllm" + generations: list[str], + data_config: DictConfig, + model_name: Optional[str] = None, + framework: Literal["vllm", "sglang", "transformers"] = "vllm", ) -> list[str]: # Remove assistant response start from transformers generations - if framework == "transformers" and (ass_resp_start := data_config.assistant_response_start): - return [ass_resp_start.join(gen.split(ass_resp_start)[1:]).strip() for gen in generations] + if framework == "transformers" and ( + ass_resp_start := data_config.assistant_response_start + ): + return [ + ass_resp_start.join(gen.split(ass_resp_start)[1:]).strip() + for gen in generations + ] elif model_name and "deepseek-r1" in model_name: return [_remove_thinking_part(gen) for gen in generations] return generations diff --git a/src/atgen/utils/resolvers.py b/src/atgen/utils/resolvers.py index 4d4959e..a659ef4 100644 --- a/src/atgen/utils/resolvers.py +++ b/src/atgen/utils/resolvers.py @@ -25,6 +25,7 @@ def to_string(model_name: str): """ return model_name.replace("/", "__") + def register_resolvers() -> None: """Register all custom resolvers with OmegaConf""" # Register resolvers only if they are not already registered diff --git a/src/atgen/utils/training_utils.py b/src/atgen/utils/training_utils.py index 2216096..eaf0298 100644 --- a/src/atgen/utils/training_utils.py +++ b/src/atgen/utils/training_utils.py @@ -117,7 +117,7 @@ def torch_call(self, examples): # Ensure sequence ends with instruction_template + last assistant response last_response_idx = response_token_ids_idxs[-1] - + # Find the user message before the last assistant response preceding_human_idxs = [ idx for idx in human_token_ids_idxs if idx < last_response_idx @@ -130,32 +130,32 @@ def torch_call(self, examples): ) batch["labels"][i, :] = self.ignore_index continue - + # Find any content after the last assistant response next_human_idxs = [ idx for idx in human_token_ids_idxs if idx > last_response_idx ] - + # If there's content after the last assistant response, truncate it if next_human_idxs: end_idx = next_human_idxs[0] # Truncate the sequence batch["input_ids"][i, end_idx:] = self.tokenizer.pad_token_id batch["attention_mask"][i, end_idx:] = 0 - + # Set all labels to ignore_index as default batch["labels"][i, :] = self.ignore_index # Only unmask the last assistant response content_start_idx = last_response_idx + len(self.response_token_ids) end_idx = batch["input_ids"].shape[1] - + # Find the actual end of the response (before padding) if self.tokenizer.pad_token_id is not None: padding_mask = batch["input_ids"][i] == self.tokenizer.pad_token_id if padding_mask.any(): end_idx = padding_mask.nonzero()[0].item() - + # Unmask only the content after the template and before the end batch["labels"][i, content_start_idx:end_idx] = batch["input_ids"][ i, content_start_idx:end_idx @@ -279,7 +279,7 @@ def _dataset_to_chat_template( # This ensures we only keep examples that will work with the collator tokenized = tokenizer(text, truncation=True, return_tensors="pt") input_ids = tokenized["input_ids"][0].tolist() - + # Check if response_token_ids exist in the tokenized input response_token_ids = data_collator.response_token_ids found = False @@ -287,7 +287,7 @@ def _dataset_to_chat_template( if input_ids[i : i + len(response_token_ids)] == response_token_ids: found = True break - + if not found: text = "" texts.append(text.strip()) diff --git a/src/atgen/utils/validate_and_fill_config.py b/src/atgen/utils/validate_and_fill_config.py index 9fd982e..b767c75 100644 --- a/src/atgen/utils/validate_and_fill_config.py +++ b/src/atgen/utils/validate_and_fill_config.py @@ -66,14 +66,18 @@ def validate_and_fill_config(config: DictConfig) -> DictConfig: OmegaConf.update( config, "data.train_output_column_name", - get_output_column_name_for_phase(config.data.output_column_name, OUTPUT_FIELD_PURPOSE_TRAIN), + get_output_column_name_for_phase( + config.data.output_column_name, OUTPUT_FIELD_PURPOSE_TRAIN + ), force_add=True, ) if "test_output_column_name" not in config.data: OmegaConf.update( config, "data.test_output_column_name", - get_output_column_name_for_phase(config.data.output_column_name, OUTPUT_FIELD_PURPOSE_TEST), + get_output_column_name_for_phase( + config.data.output_column_name, OUTPUT_FIELD_PURPOSE_TEST + ), force_add=True, )