more fixes

This commit is contained in:
Wing Lian
2023-08-21 04:58:54 -04:00
parent 918e040601
commit 2455254b92
2 changed files with 46 additions and 52 deletions

View File

@@ -22,7 +22,7 @@ from transformers import (
from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR, IntervalStrategy from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR, IntervalStrategy
from axolotl.utils.bench import log_gpu_memory_usage from axolotl.utils.bench import log_gpu_memory_usage
from axolotl.utils.distributed import barrier, is_main_process, zero_first from axolotl.utils.distributed import is_main_process, zero_first
if TYPE_CHECKING: if TYPE_CHECKING:
from axolotl.utils.trainer import AxolotlTrainingArguments from axolotl.utils.trainer import AxolotlTrainingArguments
@@ -203,7 +203,6 @@ def bench_eval_callback_factory(trainer, tokenizer):
metrics: Dict[str, float], # pylint: disable=unused-argument metrics: Dict[str, float], # pylint: disable=unused-argument
**kwargs, # pylint: disable=unused-argument **kwargs, # pylint: disable=unused-argument
): ):
if is_main_process():
data_loader = trainer.get_eval_dataloader(bench_dataset) data_loader = trainer.get_eval_dataloader(bench_dataset)
source_max_len = trainer.data_collator.max_length source_max_len = trainer.data_collator.max_length
trainer.data_collator.max_length = args.bench_source_max_len trainer.data_collator.max_length = args.bench_source_max_len
@@ -231,9 +230,7 @@ def bench_eval_callback_factory(trainer, tokenizer):
results = {"bench_loss": loss_bench / len(data_loader)} results = {"bench_loss": loss_bench / len(data_loader)}
subject = bench_dataset["subject"] subject = bench_dataset["subject"]
subjects: dict = {s: {"refs": [], "preds": []} for s in set(subject)} subjects: dict = {s: {"refs": [], "preds": []} for s in set(subject)}
for s, p, r in zip( # pylint: disable=invalid-name for s, p, r in zip(subject, preds, refs): # pylint: disable=invalid-name
subject, preds, refs
):
subjects[s]["preds"].append(p) subjects[s]["preds"].append(p)
subjects[s]["refs"].append(r) subjects[s]["refs"].append(r)
subject_scores = [] subject_scores = []
@@ -243,13 +240,10 @@ def bench_eval_callback_factory(trainer, tokenizer):
predictions=subjects[subject]["preds"], predictions=subjects[subject]["preds"],
)["accuracy"] )["accuracy"]
if not pd.isna(subject_score): if not pd.isna(subject_score):
results[ results[f"bench_{bench_split}_accuracy_{subject}"] = subject_score
f"bench_{bench_split}_accuracy_{subject}"
] = subject_score
subject_scores.append(subject_score) subject_scores.append(subject_score)
results[f"bench_{bench_split}_accuracy"] = np.mean(subject_scores) results[f"bench_{bench_split}_accuracy"] = np.mean(subject_scores)
trainer.log(results) trainer.log(results)
trainer.data_collator.max_length = source_max_len trainer.data_collator.max_length = source_max_len
barrier()
return BenchEvalCallback return BenchEvalCallback

View File

@@ -147,7 +147,7 @@ class AxolotlTrainingArguments(TrainingArguments):
}, },
) )
bench_source_max_len: int = field( bench_source_max_len: int = field(
default=2048, metadata={"help": "Maximum source sequence length for mmlu."} default=2048, metadata={"help": "Maximum source sequence length for bench."}
) )
@@ -540,9 +540,9 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer, total_num_
) )
if cfg.do_bench_eval: if cfg.do_bench_eval:
training_arguments_kwargs["do_mmlu_eval"] = cfg.do_bench_eval training_arguments_kwargs["do_bench_eval"] = cfg.do_bench_eval
if cfg.bench_dataset: if cfg.bench_dataset:
training_arguments_kwargs["mmlu_dataset"] = cfg.bench_dataset training_arguments_kwargs["bench_dataset"] = cfg.bench_dataset
training_args = AxolotlTrainingArguments( # pylint: disable=unexpected-keyword-arg training_args = AxolotlTrainingArguments( # pylint: disable=unexpected-keyword-arg
max_steps=total_num_steps if cfg.max_steps else -1, max_steps=total_num_steps if cfg.max_steps else -1,