| |
|
| | from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoConfig |
| | import numpy as np |
| | import logging |
| | import os |
| | from typing import Dict, Type, Callable, List |
| | import transformers |
| | import torch |
| | from torch import nn |
| | from torch.optim import Optimizer |
| | from torch.utils.data import DataLoader |
| | from tqdm.autonotebook import tqdm, trange |
| | from .. import SentenceTransformer, util |
| | from ..evaluation import SentenceEvaluator |
| |
|
| |
|
| | logger = logging.getLogger(__name__) |
| |
|
| |
|
| | class CrossEncoder(): |
| | def __init__(self, model_name:str, num_labels:int = None, max_length:int = None, device:str = None, tokenizer_args:Dict = {}, |
| | automodel_args:Dict = {}, default_activation_function = None): |
| | """ |
| | A CrossEncoder takes exactly two sentences / texts as input and either predicts |
| | a score or label for this sentence pair. It can for example predict the similarity of the sentence pair |
| | on a scale of 0 ... 1. |
| | |
| | It does not yield a sentence embedding and does not work for individually sentences. |
| | |
| | :param model_name: Any model name from Huggingface Models Repository that can be loaded with AutoModel. We provide several pre-trained CrossEncoder models that can be used for common tasks |
| | :param num_labels: Number of labels of the classifier. If 1, the CrossEncoder is a regression model that outputs a continous score 0...1. If > 1, it output several scores that can be soft-maxed to get probability scores for the different classes. |
| | :param max_length: Max length for input sequences. Longer sequences will be truncated. If None, max length of the model will be used |
| | :param device: Device that should be used for the model. If None, it will use CUDA if available. |
| | :param tokenizer_args: Arguments passed to AutoTokenizer |
| | :param automodel_args: Arguments passed to AutoModelForSequenceClassification |
| | :param default_activation_function: Callable (like nn.Sigmoid) about the default activation function that should be used on-top of model.predict(). If None. nn.Sigmoid() will be used if num_labels=1, else nn.Identity() |
| | """ |
| |
|
| | self.config = AutoConfig.from_pretrained(model_name) |
| | classifier_trained = True |
| | if self.config.architectures is not None: |
| | classifier_trained = any([arch.endswith('ForSequenceClassification') for arch in self.config.architectures]) |
| |
|
| | if num_labels is None and not classifier_trained: |
| | num_labels = 1 |
| |
|
| | if num_labels is not None: |
| | self.config.num_labels = num_labels |
| |
|
| | self.model = AutoModelForSequenceClassification.from_pretrained(model_name, config=self.config, **automodel_args) |
| | self.tokenizer = AutoTokenizer.from_pretrained(model_name, **tokenizer_args) |
| | self.max_length = max_length |
| |
|
| | if device is None: |
| | device = "cuda" if torch.cuda.is_available() else "cpu" |
| | logger.info("Use pytorch device: {}".format(device)) |
| |
|
| | self._target_device = torch.device(device) |
| |
|
| | if default_activation_function is not None: |
| | self.default_activation_function = default_activation_function |
| | try: |
| | self.config.sbert_ce_default_activation_function = util.fullname(self.default_activation_function) |
| | except Exception as e: |
| | logger.warning("Was not able to update config about the default_activation_function: {}".format(str(e)) ) |
| | elif hasattr(self.config, 'sbert_ce_default_activation_function') and self.config.sbert_ce_default_activation_function is not None: |
| | self.default_activation_function = util.import_from_string(self.config.sbert_ce_default_activation_function)() |
| | else: |
| | self.default_activation_function = nn.Sigmoid() if self.config.num_labels == 1 else nn.Identity() |
| |
|
| | def smart_batching_collate(self, batch): |
| | texts = [[] for _ in range(len(batch[0].texts))] |
| | labels = [] |
| |
|
| | for example in batch: |
| | for idx, text in enumerate(example.texts): |
| | texts[idx].append(text.strip()) |
| |
|
| | labels.append(example.label) |
| |
|
| | tokenized = self.tokenizer(*texts, padding=True, truncation='longest_first', return_tensors="pt", max_length=self.max_length) |
| | labels = torch.tensor(labels, dtype=torch.float if self.config.num_labels == 1 else torch.long).to(self._target_device) |
| |
|
| | for name in tokenized: |
| | tokenized[name] = tokenized[name].to(self._target_device) |
| |
|
| | return tokenized, labels |
| |
|
| | def smart_batching_collate_text_only(self, batch): |
| | texts = [[] for _ in range(len(batch[0]))] |
| |
|
| | for example in batch: |
| | for idx, text in enumerate(example): |
| | texts[idx].append(text.strip()) |
| |
|
| | tokenized = self.tokenizer(*texts, padding=True, truncation='longest_first', return_tensors="pt", max_length=self.max_length) |
| |
|
| | for name in tokenized: |
| | tokenized[name] = tokenized[name].to(self._target_device) |
| |
|
| | return tokenized |
| |
|
| | def fit(self, |
| | train_dataloader: DataLoader, |
| | evaluator: SentenceEvaluator = None, |
| | epochs: int = 1, |
| | loss_fct = None, |
| | activation_fct = nn.Identity(), |
| | scheduler: str = 'WarmupLinear', |
| | warmup_steps: int = 10000, |
| | optimizer_class: Type[Optimizer] = torch.optim.AdamW, |
| | optimizer_params: Dict[str, object] = {'lr': 2e-5}, |
| | weight_decay: float = 0.01, |
| | evaluation_steps: int = 0, |
| | output_path: str = None, |
| | save_best_model: bool = True, |
| | max_grad_norm: float = 1, |
| | use_amp: bool = False, |
| | callback: Callable[[float, int, int], None] = None, |
| | show_progress_bar: bool = True |
| | ): |
| | """ |
| | Train the model with the given training objective |
| | Each training objective is sampled in turn for one batch. |
| | We sample only as many batches from each objective as there are in the smallest one |
| | to make sure of equal training with each dataset. |
| | |
| | :param train_dataloader: DataLoader with training InputExamples |
| | :param evaluator: An evaluator (sentence_transformers.evaluation) evaluates the model performance during training on held-out dev data. It is used to determine the best model that is saved to disc. |
| | :param epochs: Number of epochs for training |
| | :param loss_fct: Which loss function to use for training. If None, will use nn.BCEWithLogitsLoss() if self.config.num_labels == 1 else nn.CrossEntropyLoss() |
| | :param activation_fct: Activation function applied on top of logits output of model. |
| | :param scheduler: Learning rate scheduler. Available schedulers: constantlr, warmupconstant, warmuplinear, warmupcosine, warmupcosinewithhardrestarts |
| | :param warmup_steps: Behavior depends on the scheduler. For WarmupLinear (default), the learning rate is increased from o up to the maximal learning rate. After these many training steps, the learning rate is decreased linearly back to zero. |
| | :param optimizer_class: Optimizer |
| | :param optimizer_params: Optimizer parameters |
| | :param weight_decay: Weight decay for model parameters |
| | :param evaluation_steps: If > 0, evaluate the model using evaluator after each number of training steps |
| | :param output_path: Storage path for the model and evaluation files |
| | :param save_best_model: If true, the best model (according to evaluator) is stored at output_path |
| | :param max_grad_norm: Used for gradient normalization. |
| | :param use_amp: Use Automatic Mixed Precision (AMP). Only for Pytorch >= 1.6.0 |
| | :param callback: Callback function that is invoked after each evaluation. |
| | It must accept the following three parameters in this order: |
| | `score`, `epoch`, `steps` |
| | :param show_progress_bar: If True, output a tqdm progress bar |
| | """ |
| | train_dataloader.collate_fn = self.smart_batching_collate |
| |
|
| | if use_amp: |
| | from torch.cuda.amp import autocast |
| | scaler = torch.cuda.amp.GradScaler() |
| |
|
| | self.model.to(self._target_device) |
| |
|
| | if output_path is not None: |
| | os.makedirs(output_path, exist_ok=True) |
| |
|
| | self.best_score = -9999999 |
| | num_train_steps = int(len(train_dataloader) * epochs) |
| |
|
| | |
| | param_optimizer = list(self.model.named_parameters()) |
| |
|
| | no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight'] |
| | optimizer_grouped_parameters = [ |
| | {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': weight_decay}, |
| | {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0} |
| | ] |
| |
|
| | optimizer = optimizer_class(optimizer_grouped_parameters, **optimizer_params) |
| |
|
| | if isinstance(scheduler, str): |
| | scheduler = SentenceTransformer._get_scheduler(optimizer, scheduler=scheduler, warmup_steps=warmup_steps, t_total=num_train_steps) |
| |
|
| | if loss_fct is None: |
| | loss_fct = nn.BCEWithLogitsLoss() if self.config.num_labels == 1 else nn.CrossEntropyLoss() |
| |
|
| |
|
| | skip_scheduler = False |
| | for epoch in trange(epochs, desc="Epoch", disable=not show_progress_bar): |
| | training_steps = 0 |
| | self.model.zero_grad() |
| | self.model.train() |
| |
|
| | for features, labels in tqdm(train_dataloader, desc="Iteration", smoothing=0.05, disable=not show_progress_bar): |
| | if use_amp: |
| | with autocast(): |
| | model_predictions = self.model(**features, return_dict=True) |
| | logits = activation_fct(model_predictions.logits) |
| | if self.config.num_labels == 1: |
| | logits = logits.view(-1) |
| | loss_value = loss_fct(logits, labels) |
| |
|
| | scale_before_step = scaler.get_scale() |
| | scaler.scale(loss_value).backward() |
| | scaler.unscale_(optimizer) |
| | torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_grad_norm) |
| | scaler.step(optimizer) |
| | scaler.update() |
| |
|
| | skip_scheduler = scaler.get_scale() != scale_before_step |
| | else: |
| | model_predictions = self.model(**features, return_dict=True) |
| | logits = activation_fct(model_predictions.logits) |
| | if self.config.num_labels == 1: |
| | logits = logits.view(-1) |
| | loss_value = loss_fct(logits, labels) |
| | loss_value.backward() |
| | torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_grad_norm) |
| | optimizer.step() |
| |
|
| | optimizer.zero_grad() |
| |
|
| | if not skip_scheduler: |
| | scheduler.step() |
| |
|
| | training_steps += 1 |
| |
|
| | if evaluator is not None and evaluation_steps > 0 and training_steps % evaluation_steps == 0: |
| | self._eval_during_training(evaluator, output_path, save_best_model, epoch, training_steps, callback) |
| |
|
| | self.model.zero_grad() |
| | self.model.train() |
| |
|
| | if evaluator is not None: |
| | self._eval_during_training(evaluator, output_path, save_best_model, epoch, -1, callback) |
| |
|
| |
|
| |
|
| | def predict(self, sentences: List[List[str]], |
| | batch_size: int = 32, |
| | show_progress_bar: bool = None, |
| | num_workers: int = 0, |
| | activation_fct = None, |
| | apply_softmax = False, |
| | convert_to_numpy: bool = True, |
| | convert_to_tensor: bool = False |
| | ): |
| | """ |
| | Performs predicts with the CrossEncoder on the given sentence pairs. |
| | |
| | :param sentences: A list of sentence pairs [[Sent1, Sent2], [Sent3, Sent4]] |
| | :param batch_size: Batch size for encoding |
| | :param show_progress_bar: Output progress bar |
| | :param num_workers: Number of workers for tokenization |
| | :param activation_fct: Activation function applied on the logits output of the CrossEncoder. If None, nn.Sigmoid() will be used if num_labels=1, else nn.Identity |
| | :param convert_to_numpy: Convert the output to a numpy matrix. |
| | :param apply_softmax: If there are more than 2 dimensions and apply_softmax=True, applies softmax on the logits output |
| | :param convert_to_tensor: Conver the output to a tensor. |
| | :return: Predictions for the passed sentence pairs |
| | """ |
| | input_was_string = False |
| | if isinstance(sentences[0], str): |
| | sentences = [sentences] |
| | input_was_string = True |
| |
|
| | inp_dataloader = DataLoader(sentences, batch_size=batch_size, collate_fn=self.smart_batching_collate_text_only, num_workers=num_workers, shuffle=False) |
| |
|
| | if show_progress_bar is None: |
| | show_progress_bar = (logger.getEffectiveLevel() == logging.INFO or logger.getEffectiveLevel() == logging.DEBUG) |
| |
|
| | iterator = inp_dataloader |
| | if show_progress_bar: |
| | iterator = tqdm(inp_dataloader, desc="Batches") |
| |
|
| | if activation_fct is None: |
| | activation_fct = self.default_activation_function |
| |
|
| | pred_scores = [] |
| | self.model.eval() |
| | self.model.to(self._target_device) |
| | with torch.no_grad(): |
| | for features in iterator: |
| | model_predictions = self.model(**features, return_dict=True) |
| | logits = activation_fct(model_predictions.logits) |
| |
|
| | if apply_softmax and len(logits[0]) > 1: |
| | logits = torch.nn.functional.softmax(logits, dim=1) |
| | pred_scores.extend(logits) |
| |
|
| | if self.config.num_labels == 1: |
| | pred_scores = [score[0] for score in pred_scores] |
| |
|
| | if convert_to_tensor: |
| | pred_scores = torch.stack(pred_scores) |
| | elif convert_to_numpy: |
| | pred_scores = np.asarray([score.cpu().detach().numpy() for score in pred_scores]) |
| |
|
| | if input_was_string: |
| | pred_scores = pred_scores[0] |
| |
|
| | return pred_scores |
| |
|
| |
|
| | def _eval_during_training(self, evaluator, output_path, save_best_model, epoch, steps, callback): |
| | """Runs evaluation during the training""" |
| | if evaluator is not None: |
| | score = evaluator(self, output_path=output_path, epoch=epoch, steps=steps) |
| | if callback is not None: |
| | callback(score, epoch, steps) |
| | if score > self.best_score: |
| | self.best_score = score |
| | if save_best_model: |
| | self.save(output_path) |
| |
|
| | def save(self, path): |
| | """ |
| | Saves all model and tokenizer to path |
| | """ |
| | if path is None: |
| | return |
| |
|
| | logger.info("Save model to {}".format(path)) |
| | self.model.save_pretrained(path) |
| | self.tokenizer.save_pretrained(path) |
| |
|
| | def save_pretrained(self, path): |
| | """ |
| | Same function as save |
| | """ |
| | return self.save(path) |
| |
|
| |
|
