diff --git a/.gitignore b/.gitignore index 1ee3b86..d1683db 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,5 @@ venv venv/* -repos TensorRT-LLM */dataset cache diff --git a/configs/base.yaml b/configs/base.yaml index e8e9f04..8cb9225 100644 --- a/configs/base.yaml +++ b/configs/base.yaml @@ -17,7 +17,6 @@ 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 @@ -25,7 +24,7 @@ al: init_query_size: 10 query_size: 10 num_iterations: 5 - eval_zero_iteration: True + evaluate_zero_iteration: True subsample_size: -1 required_performance: rouge1: 0.5 @@ -40,9 +39,9 @@ model: assistant_response_start: "\n\n\n\n" peft: use: True - r: 32 - lora_alpha: 32 - lora_dropout: 0. + r: 64 + lora_alpha: 64 + lora_dropout: 0.1 bias: 'none' seed: ${seed} use_gradient_checkpointing: 'unsloth' @@ -51,8 +50,8 @@ model: training: dev_split_size: 0.2 hyperparameters: - num_epochs: 15 - train_batch_size: 8 + num_epochs: 5 + train_batch_size: 6 eval_batch_size: 4 gradient_accumulation_steps: 2 lr: 0.00003 @@ -63,7 +62,7 @@ training: model_max_length: ${model.model_max_length} dataset_num_proc: ${data.num_proc} packing: False - lr_scheduler_type: linear + lr_scheduler_type: cosine gradient_checkpointing: True @@ -76,7 +75,6 @@ inference: temperature: 0.6 model: ${model.checkpoint} num_display_generations: 5 - num_threads_for_bfcl: 32 evaluation: additional_metrics: [] @@ -90,3 +88,13 @@ 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/configs/data/hermes_tool.yaml b/configs/data/hermes_tool.yaml deleted file mode 100644 index 7affddd..0000000 --- a/configs/data/hermes_tool.yaml +++ /dev/null @@ -1,15 +0,0 @@ -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/configs/test.yaml b/configs/test.yaml index 08b0b03..2535b22 100644 --- a/configs/test.yaml +++ b/configs/test.yaml @@ -10,7 +10,6 @@ output_dir: outputs offline_mode: False cache_dir: cache seed: 42 -save_model: False al: init_query_size: 2 query_size: 2 @@ -65,7 +64,6 @@ inference: top_p: 0.1 model: ${model.checkpoint} num_display_generations: 0 - num_threads_for_bfcl: 32 evaluation: additional_metrics: [] diff --git a/pyproject.toml b/pyproject.toml index dd36ab8..bcfe10d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,6 @@ 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/requirements.txt b/requirements.txt index 02b6a42..a65f05f 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.6.0 +datasets==3.4.1 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.90.0 +openai==1.66.3 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.19.1 +transformers==4.52.4 +trl==0.15.2 torchmetrics==1.4.1 unsloth==2025.8.1 -# vllm==0.10.0 +vllm==0.8.1 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 1f22736..8ebef95 100644 --- a/src/atgen/labellers/api_labellers/openai_labeller.py +++ b/src/atgen/labellers/api_labellers/openai_labeller.py @@ -22,7 +22,6 @@ MAX_NUM_TRIES = 3 UPDATE_TIME_IN_SECONDS = 10 # update time when checking for the completion - class OpenAILabeller(BaseLabeler): def __init__( self, @@ -106,7 +105,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/labellers/get_labeller.py b/src/atgen/labellers/get_labeller.py index 18c794e..2f34d36 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 | float = 1_000_000, + budget: int = 1_000_000, workdir: str | Path = "tmp", **kwargs, ): diff --git a/src/atgen/metrics/compute_metrics.py b/src/atgen/metrics/compute_metrics.py index c4c9e54..2864372 100644 --- a/src/atgen/metrics/compute_metrics.py +++ b/src/atgen/metrics/compute_metrics.py @@ -81,24 +81,13 @@ 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) @@ -135,19 +124,13 @@ 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( @@ -212,14 +195,10 @@ 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: @@ -296,42 +275,31 @@ 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'(? 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) @@ -175,7 +177,7 @@ def run_active_learning(config, workdir: Union[str, Path]): lambda x: x["id"] not in set(labeled_ids) ) else: - query_ids: list[int] = al_strategy( + query_ids: list[str] = al_strategy( model=model, tokenizer=tokenizer, unlabeled_pool=unlabeled_data.remove_columns(output_column_name_train), @@ -219,20 +221,18 @@ def run_active_learning(config, workdir: Union[str, Path]): model_name=model_name, ) - if has_test and not config.data.use_test_benchmark: + if has_test: 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, + data_config=config.data, + split="test", + 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, @@ -253,46 +253,30 @@ def run_active_learning(config, workdir: Union[str, Path]): 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 - ( - is_performance_reached, - is_metrics_availability_checked, - available_metrics, - ) = check_performance_against_requirements( + is_performance_reached, is_metrics_availability_checked, available_metrics = ( + check_performance_against_requirements( + metrics=metrics, + required_performance_dict=required_performance_dict, + is_metrics_availability_checked=is_metrics_availability_checked, + available_metrics=available_metrics, + ) + ) + save_log_iter_results( + config=config, + workdir=workdir, + iter_dir=iter_dir, metrics=metrics, - required_performance_dict=required_performance_dict, - is_metrics_availability_checked=is_metrics_availability_checked, - available_metrics=available_metrics, + generations=generations, + al_iter=0, + train_result={}, + model=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 + 1` because we do not label data + # Start AL cycle. Use `num_al_iterations + 2` 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 for al_iter in range(start_iter, num_al_iterations + 1 + start_iter): print(f"Starting AL iteration #{al_iter}.") @@ -351,68 +335,46 @@ 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: 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, - ) + 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, + ) # Check required performance metrics - ( - is_performance_reached, - is_metrics_availability_checked, - available_metrics, - ) = check_performance_against_requirements( - metrics=metrics, - required_performance_dict=required_performance_dict, - is_metrics_availability_checked=is_metrics_availability_checked, - available_metrics=available_metrics, + is_performance_reached, is_metrics_availability_checked, available_metrics = ( + check_performance_against_requirements( + metrics=metrics, + required_performance_dict=required_performance_dict, + is_metrics_availability_checked=is_metrics_availability_checked, + available_metrics=available_metrics, + ) ) save_log_iter_results( config=config, @@ -423,7 +385,6 @@ 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 @@ -440,9 +401,7 @@ def run_active_learning(config, workdir: Union[str, Path]): ) query: Dataset = unlabeled_data.filter(lambda x: x["id"] in query_ids) - unlabeled_data: Dataset = unlabeled_data.filter( - lambda x: x["id"] not in query_ids - ) + unlabeled_data: Dataset = unlabeled_data.filter(lambda x: x["id"] not in query_ids) labeled_query: Dataset = labeller(query) if labeller.is_out_of_budget: print(f"Labeler ran out of budget at iteration {al_iter}.") diff --git a/src/atgen/run_scripts/run_subset_selection.py b/src/atgen/run_scripts/run_subset_selection.py deleted file mode 100644 index 59c451c..0000000 --- a/src/atgen/run_scripts/run_subset_selection.py +++ /dev/null @@ -1,381 +0,0 @@ -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 - budget = config.al.budget - if budget is None: - budget = 1e10 - - 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] - log.info(f"Loaded {len(labeled_ids)} labeled ids from {config.al.query_ids_path}") - 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/strategies/__init__.py b/src/atgen/strategies/__init__.py index 51aaf37..92470b5 100644 --- a/src/atgen/strategies/__init__.py +++ b/src/atgen/strategies/__init__.py @@ -19,12 +19,4 @@ "bleuvar": BLEUVarStrategy, "idds": IDDSStrategy, "dual": DualStrategy, - "kek": RandomStrategy, - "kek2": RandomStrategy, - "kek3": RandomStrategy, - "kek4": RandomStrategy, - "kek5": RandomStrategy, - "kek6": RandomStrategy, - "random1": RandomStrategy, - "random2": RandomStrategy, } diff --git a/src/atgen/strategies/base_strategy.py b/src/atgen/strategies/base_strategy.py index e81506f..eb96174 100644 --- a/src/atgen/strategies/base_strategy.py +++ b/src/atgen/strategies/base_strategy.py @@ -2,6 +2,10 @@ import random from datasets import Dataset +from transformers import ( + PreTrainedModel, + PreTrainedTokenizer, +) class BaseStrategy(ABC): @@ -11,10 +15,13 @@ 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/strategies/hadas.py b/src/atgen/strategies/hadas.py index 19e2e96..66801fb 100644 --- a/src/atgen/strategies/hadas.py +++ b/src/atgen/strategies/hadas.py @@ -102,7 +102,7 @@ def semantic_frame_score( batch_results = [] # Work-around for conversational data if isinstance(documents[0], list): - documents = [inst[-2]["content"].strip() for inst in documents] + documents = [inst[-2]['content'].strip() for inst in documents] for i_start in range(0, len(documents), batch_size): inputs = entailment_tokenizer( list( @@ -134,7 +134,7 @@ def discourse_score( we do not have access to during active learning. """ if isinstance(documents[0], list): - documents = [inst[-2]["content"].strip() for inst in documents] + documents = [inst[-2]['content'].strip() for inst in documents] data = convert_to_json(src_list=documents, output_list=summaries) scores = [ s["overall"] diff --git a/src/atgen/utils/check_performance_metrics.py b/src/atgen/utils/check_performance_metrics.py index be4e472..9702cc0 100644 --- a/src/atgen/utils/check_performance_metrics.py +++ b/src/atgen/utils/check_performance_metrics.py @@ -4,10 +4,7 @@ def check_performance_against_requirements( - metrics, - required_performance_dict, - is_metrics_availability_checked, - available_metrics, + metrics, required_performance_dict, is_metrics_availability_checked, available_metrics ): """ Check if the computed metrics meet the required performance thresholds. diff --git a/src/atgen/utils/constants.py b/src/atgen/utils/constants.py index bec6245..95fd7e9 100644 --- a/src/atgen/utils/constants.py +++ b/src/atgen/utils/constants.py @@ -13,9 +13,3 @@ 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 - -NUM_PROCS_FOR_DATASETS = 16 diff --git a/src/atgen/utils/data/load_data.py b/src/atgen/utils/data/load_data.py index 4f649ca..6a00b56 100644 --- a/src/atgen/utils/data/load_data.py +++ b/src/atgen/utils/data/load_data.py @@ -127,9 +127,7 @@ def load_data( phase=split, ) if data_config.task == "multi-choice-qa": - dataset = _preprocess_multi_choice_qa( - dataset=dataset, data_config=data_config, split=split - ) + dataset = _preprocess_multi_choice_qa(dataset=dataset, data_config=data_config, split=split) # Add `id` column to the dataset (practical use) or to train subset (benchmarking) dataset = _add_id_column(dataset) if subset_size is not None: @@ -137,10 +135,7 @@ def load_data( return dataset - -def _preprocess_multi_choice_qa( - dataset: Dataset, data_config: DictConfig, split: str -) -> 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"] @@ -148,21 +143,16 @@ def _preprocess_multi_choice_qa( 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"]}) @@ -171,8 +161,6 @@ def _preprocess_multi_choice_qa( 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 0ce2f1a..2bc07f5 100644 --- a/src/atgen/utils/data/prepare_conversational_data.py +++ b/src/atgen/utils/data/prepare_conversational_data.py @@ -19,18 +19,10 @@ 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 55099f7..c41e3a6 100644 --- a/src/atgen/utils/downloaders.py +++ b/src/atgen/utils/downloaders.py @@ -1,24 +1,19 @@ 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 dc0d844..96f355c 100644 --- a/src/atgen/utils/evaluate_bfcl.py +++ b/src/atgen/utils/evaluate_bfcl.py @@ -10,19 +10,14 @@ 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, -) +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__) @@ -34,7 +29,6 @@ "Irrelevance Detection", ] - def evaluate_bfcl( model_name: str, bfcl_results_dir: str | Path, @@ -45,7 +39,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 @@ -53,13 +47,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( @@ -68,14 +62,13 @@ 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, @@ -86,7 +79,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 @@ -94,7 +87,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 """ @@ -107,71 +100,55 @@ def _evaluate_bfcl( cwd = os.getcwd() os.chdir(bfcl_results_dir) model_name = save_dir - # 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}") 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)) @@ -196,57 +173,43 @@ 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.lower() 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 02f7f41..7bd7638 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,7 +258,9 @@ 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, @@ -278,7 +280,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 @@ -335,17 +337,12 @@ 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/installers.py b/src/atgen/utils/installers.py new file mode 100644 index 0000000..b540fb9 --- /dev/null +++ b/src/atgen/utils/installers.py @@ -0,0 +1,10 @@ +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/load_model_tokenizer.py b/src/atgen/utils/load_model_tokenizer.py index 8804311..2d33e85 100644 --- a/src/atgen/utils/load_model_tokenizer.py +++ b/src/atgen/utils/load_model_tokenizer.py @@ -1,5 +1,4 @@ from pathlib import Path -from time import sleep import torch @@ -23,9 +22,6 @@ class FastLanguageModel: ) -NUM_TRIES_LOAD_MODEL = 3 - - # def load_tokenizer( # model_config: DictConfig, # cache_dir: str, @@ -48,38 +44,32 @@ def load_model_tokenizer( cache_dir: str, ) -> tuple[FastLanguageModel, PreTrainedTokenizerFast]: dtype = getattr(torch, model_config.dtype, torch.bfloat16) - 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) + 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" + ) - 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 diff --git a/src/atgen/utils/main_decorator.py b/src/atgen/utils/main_decorator.py index 37c7d67..2e86397 100644 --- a/src/atgen/utils/main_decorator.py +++ b/src/atgen/utils/main_decorator.py @@ -8,7 +8,6 @@ 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" @@ -29,9 +28,6 @@ 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: @@ -43,7 +39,6 @@ def run_script(config): os.environ["HF_EVALUATE_OFFLINE"] = "1" os.environ["PYTHONHASHSEED"] = str(config.seed) - from transformers import set_seed set_seed(config.seed) diff --git a/src/atgen/utils/post_process_generations.py b/src/atgen/utils/post_process_generations.py index 4477191..fd5425a 100644 --- a/src/atgen/utils/post_process_generations.py +++ b/src/atgen/utils/post_process_generations.py @@ -3,21 +3,15 @@ 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 a659ef4..4d4959e 100644 --- a/src/atgen/utils/resolvers.py +++ b/src/atgen/utils/resolvers.py @@ -25,7 +25,6 @@ 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/save_log_iter_results.py b/src/atgen/utils/save_log_iter_results.py index 06a2062..39024a0 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, PreTrainedTokenizer +from transformers import PreTrainedModel from atgen.utils.combine_results import combine_results @@ -21,7 +21,6 @@ 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: @@ -33,6 +32,5 @@ 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") - tokenizer.save_pretrained(workdir / "tokenizer") + if config.save_model and model is not None: + model.save_pretrained(workdir / "model.bin") diff --git a/src/atgen/utils/training_utils.py b/src/atgen/utils/training_utils.py index eaf0298..4816e71 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 @@ -202,10 +202,10 @@ def _get_response_instruction_templates( """ if "gemma" in tokenizer.name_or_path.lower(): response_template = "model\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" + 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" 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" @@ -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,10 +287,10 @@ 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()) + texts.append(text) return {TEXT_FIELD: texts} diff --git a/src/atgen/utils/validate_and_fill_config.py b/src/atgen/utils/validate_and_fill_config.py index b767c75..9fd982e 100644 --- a/src/atgen/utils/validate_and_fill_config.py +++ b/src/atgen/utils/validate_and_fill_config.py @@ -66,18 +66,14 @@ 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, )