File size: 15,269 Bytes
2359bda |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 |
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)
# Prepare optimizers
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): # Cast an individual sentence to a list with length 1
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)
|