From ac5ddf7ec3ae49cd240db55450428d4fa81988ef Mon Sep 17 00:00:00 2001 From: calpt Date: Mon, 13 Jun 2022 14:22:42 +0200 Subject: [PATCH 1/3] Infer label names for training for flex head models --- src/transformers/adapters/heads/base.py | 6 +++++ src/transformers/adapters/trainer.py | 29 ++----------------------- 2 files changed, 8 insertions(+), 27 deletions(-) diff --git a/src/transformers/adapters/heads/base.py b/src/transformers/adapters/heads/base.py index 87dcf9772..3ad8f3e5a 100644 --- a/src/transformers/adapters/heads/base.py +++ b/src/transformers/adapters/heads/base.py @@ -91,6 +91,9 @@ def build(self, model): def get_output_embeddings(self): return None # override for heads with output embeddings + def get_label_names(self): + return ["labels"] + class ClassificationHead(PredictionHead): def __init__( @@ -405,6 +408,9 @@ def forward(self, outputs, cls_output=None, attention_mask=None, return_dict=Fal outputs = (total_loss,) + outputs return outputs + def get_label_names(self): + return ["start_positions", "end_positions"] + class ModelWithFlexibleHeadsAdaptersMixin(ModelWithHeadsAdaptersMixin): """ diff --git a/src/transformers/adapters/trainer.py b/src/transformers/adapters/trainer.py index 39374c551..a74f2b857 100644 --- a/src/transformers/adapters/trainer.py +++ b/src/transformers/adapters/trainer.py @@ -86,6 +86,8 @@ def __init__( "Expected a model with an active adapter setup." "If you want to fully finetune the model use the Trainer class." ) + if (self.label_names is None or len(self.label_names) < 1) and self.active_head is not None: + self.label_names = model.heads[model.active_head].get_label_names() def create_optimizer(self): """ @@ -215,33 +217,6 @@ def _load_heads(self, resume_from_checkpoint): ): self.model.load_head(os.path.join(resume_from_checkpoint, file_name)) - def _remove_unused_columns(self, dataset: "datasets.Dataset", description: Optional[str] = None): - if not self.args.remove_unused_columns: - return dataset - if self._signature_columns is None: - # Inspect model forward signature to keep only the arguments it accepts. - signature = inspect.signature(self.model.forward) - self._signature_columns = list(signature.parameters.keys()) - # Labels may be named label or label_ids, the default data collator handles that. - self._signature_columns += ["label", "label_ids"] - self._signature_columns += self.label_names - columns = [k for k in self._signature_columns if k in dataset.column_names] - ignored_columns = list(set(dataset.column_names) - set(self._signature_columns)) - if len(ignored_columns) > 0: - dset_description = "" if description is None else f"in the {description} set " - logger.info( - f"The following columns {dset_description} don't have a corresponding argument in " - f"`{self.model.__class__.__name__}.forward` and have been ignored: {', '.join(ignored_columns)}." - ) - - if version.parse(datasets.__version__) < version.parse("1.4.0"): - dataset.set_format( - type=dataset.format["type"], columns=columns, format_kwargs=dataset.format["format_kwargs"] - ) - return dataset - else: - return dataset.remove_columns(ignored_columns) - class AdapterTrainerCallback(TrainerCallback): def __init__(self, trainer): From 5bd08d1b66e3a3504d41d8c21fc02dc280e91430 Mon Sep 17 00:00:00 2001 From: calpt Date: Tue, 14 Jun 2022 10:21:37 +0200 Subject: [PATCH 2/3] Fixes --- src/transformers/adapters/trainer.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/transformers/adapters/trainer.py b/src/transformers/adapters/trainer.py index a74f2b857..e42407f92 100644 --- a/src/transformers/adapters/trainer.py +++ b/src/transformers/adapters/trainer.py @@ -1,11 +1,8 @@ -import inspect import os import re from typing import Callable, Dict, List, Optional, Tuple, Union -import datasets import torch -from packaging import version from torch import nn from torch.utils.data.dataset import Dataset @@ -86,7 +83,7 @@ def __init__( "Expected a model with an active adapter setup." "If you want to fully finetune the model use the Trainer class." ) - if (self.label_names is None or len(self.label_names) < 1) and self.active_head is not None: + if (self.label_names is None or len(self.label_names) < 1) and model.active_head is not None: self.label_names = model.heads[model.active_head].get_label_names() def create_optimizer(self): @@ -250,6 +247,7 @@ def on_train_end(self, args: TrainingArguments, state: TrainerState, control: Tr fusion_dir = os.path.join(state.best_model_checkpoint, fusion) if os.path.exists(fusion_dir): model.load_adapter_fusion(fusion_dir) + model.to(args.device) def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): # apply adapter fusion weight regularization on the value matrix From 5a7c4052bcc8d341b1d5bf24dcece3b06b70298f Mon Sep 17 00:00:00 2001 From: calpt Date: Tue, 14 Jun 2022 10:57:21 +0200 Subject: [PATCH 3/3] fix --- src/transformers/adapters/trainer.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/transformers/adapters/trainer.py b/src/transformers/adapters/trainer.py index e42407f92..f5e5a4204 100644 --- a/src/transformers/adapters/trainer.py +++ b/src/transformers/adapters/trainer.py @@ -84,7 +84,10 @@ def __init__( "If you want to fully finetune the model use the Trainer class." ) if (self.label_names is None or len(self.label_names) < 1) and model.active_head is not None: - self.label_names = model.heads[model.active_head].get_label_names() + all_label_names = set() + for head in model._active_heads: + all_label_names |= set(model.heads[head].get_label_names()) + self.label_names = list(all_label_names) def create_optimizer(self): """