Aluode/PerceptionLabPortable
0
1import logging2import os3from pathlib import Path4from time import sleep5from typing import Callable, Optional, Union6 7import numpy as np8import tensorflow as tf9from huggingface_hub import Repository, create_repo10from packaging.version import parse11 12from . import IntervalStrategy, PreTrainedTokenizerBase13from .modelcard import TrainingSummary14from .modeling_tf_utils import keras15 16 17logger = logging.getLogger(__name__)18 19 20class KerasMetricCallback(keras.callbacks.Callback):21 """22 Callback to compute metrics at the end of every epoch. Unlike normal Keras metrics, these do not need to be23 compilable by TF. It is particularly useful for common NLP metrics like BLEU and ROUGE that require string24 operations or generation loops that cannot be compiled. Predictions (or generations) will be computed on the25 `eval_dataset` before being passed to the `metric_fn` in `np.ndarray` format. The `metric_fn` should compute26 metrics and return a dict mapping metric names to metric values.27 28 We provide an example of a suitable metric_fn that computes ROUGE scores for a summarization model below. Note that29 this example skips some post-processing for readability and simplicity, and should probably not be used as-is!30 31 ```py32 from datasets import load_metric33 34 rouge_metric = load_metric("rouge")35 36 37 def rouge_fn(predictions, labels):38 decoded_predictions = tokenizer.batch_decode(predictions, skip_special_tokens=True)39 decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)40 result = rouge_metric.compute(predictions=decoded_predictions, references=decoded_labels)41 return {key: value.mid.fmeasure * 100 for key, value in result.items()}42 ```43 44 The above function will return a dict containing values which will be logged like any other Keras metric:45 46 ```47 {'rouge1': 37.4199, 'rouge2': 13.9768, 'rougeL': 34.361, 'rougeLsum': 35.078148 ```49 50 Args:51 metric_fn (`Callable`):52 Metric function provided by the user. It will be called with two arguments - `predictions` and `labels`.53 These contain the model's outputs and matching labels from the dataset. It should return a dict mapping54 metric names to numerical values.55 eval_dataset (`tf.data.Dataset` or `dict` or `tuple` or `np.ndarray` or `tf.Tensor`):56 Validation data to be used to generate predictions for the `metric_fn`.57 output_cols (`list[str], *optional*):58 A list of columns to be retained from the model output as the predictions. Defaults to all.59 label_cols ('`list[str]`, *optional*'):60 A list of columns to be retained from the input dataset as the labels. Will be autodetected if this is not61 supplied.62 batch_size (`int`, *optional*):63 Batch size. Only used when the data is not a pre-batched `tf.data.Dataset`.64 predict_with_generate (`bool`, *optional*, defaults to `False`):65 Whether we should use `model.generate()` to get outputs for the model.66 use_xla_generation (`bool`, *optional*, defaults to `False`):67 If we're generating, whether to compile model generation with XLA. This can massively increase the speed of68 generation (up to 100X speedup) but will require a new XLA compilation for each input shape. When using XLA69 generation, it's a good idea to pad your inputs to the same size, or to use the `pad_to_multiple_of`70 argument in your `tokenizer` or `DataCollator`, which will reduce the number of unique input shapes and71 save a lot of compilation time. This option has no effect is `predict_with_generate` is `False`.72 generate_kwargs (`dict`, *optional*):73 Keyword arguments to pass to `model.generate()` when generating. Has no effect if `predict_with_generate`74 is `False`.75 76 """77 78 def __init__(79 self,80 metric_fn: Callable,81 eval_dataset: Union[tf.data.Dataset, np.ndarray, tf.Tensor, tuple, dict],82 output_cols: Optional[list[str]] = None,83 label_cols: Optional[list[str]] = None,84 batch_size: Optional[int] = None,85 predict_with_generate: bool = False,86 use_xla_generation: bool = False,87 generate_kwargs: Optional[dict] = None,88 ):89 super().__init__()90 self.metric_fn = metric_fn91 self.batch_size = batch_size92 if not isinstance(eval_dataset, tf.data.Dataset):93 if batch_size is None:94 raise ValueError(95 "When passing data to KerasMetricCallback that is not a pre-batched tf.data.Dataset "96 "the batch_size argument must be set."97 )98 # Wrap a tf.data.Dataset around it99 eval_dataset = tf.data.Dataset.from_tensor_slices(eval_dataset).batch(batch_size, drop_remainder=False)100 self.eval_dataset = eval_dataset101 self.predict_with_generate = predict_with_generate102 self.output_cols = output_cols103 104 # This next block attempts to parse out which elements of the dataset should be appended to the labels list105 # that is passed to the metric_fn106 if isinstance(eval_dataset.element_spec, tuple) and len(eval_dataset.element_spec) == 2:107 input_spec, label_spec = eval_dataset.element_spec108 else:109 input_spec = eval_dataset.element_spec110 label_spec = None111 if label_cols is not None:112 for label in label_cols:113 if label not in input_spec:114 raise ValueError(f"Label {label} is in label_cols but could not be found in the dataset inputs!")115 self.label_cols = label_cols116 self.use_keras_label = False117 elif label_spec is not None:118 # If the dataset inputs are split into a 2-tuple of inputs and labels,119 # assume the second element is the labels120 self.label_cols = None121 self.use_keras_label = True122 elif "labels" in input_spec:123 self.label_cols = ["labels"]124 self.use_keras_label = False125 logging.warning("No label_cols specified for KerasMetricCallback, assuming you want the 'labels' key.")126 elif "start_positions" in input_spec and "end_positions" in input_spec:127 self.label_cols = ["start_positions", "end_positions"]128 self.use_keras_label = False129 logging.warning(130 "No label_cols specified for KerasMetricCallback, assuming you want the "131 "start_positions and end_positions keys."132 )133 else:134 raise ValueError("Could not autodetect label_cols for KerasMetricCallback, please specify them!")135 if parse(tf.__version__) < parse("2.7"):136 logging.warning("TF versions less than 2.7 may encounter issues with KerasMetricCallback!")137 138 self.use_xla_generation = use_xla_generation139 self.generate_kwargs = {} if generate_kwargs is None else generate_kwargs140 141 self.generation_function = None142 143 @staticmethod144 def _concatenate_batches(batches, padding_index=-100):145 # If all batches are unidimensional or same length, do a simple concatenation146 if batches[0].ndim == 1 or all(batch.shape[1] == batches[0].shape[1] for batch in batches):147 return np.concatenate(batches, axis=0)148 149 # Welp, they're not the same length. Let's do some padding150 max_len = max([batch.shape[1] for batch in batches])151 num_samples = sum([batch.shape[0] for batch in batches])152 output = np.full_like(153 batches[0], fill_value=padding_index, shape=[num_samples, max_len] + list(batches[0].shape[2:])154 )155 # i keeps track of which part of the concatenated array we're writing the next batch to156 i = 0157 for batch in batches:158 output[i : i + len(batch), : batch.shape[1]] = batch159 i += len(batch)160 return output161 162 def _postprocess_predictions_or_labels(self, inputs):163 if isinstance(inputs[0], dict):164 outputs = {}165 for key in inputs[0]:166 outputs[key] = self._concatenate_batches([batch[key] for batch in inputs])167 # If it's a dict with only one key, just return the array168 if len(outputs) == 1:169 outputs = list(outputs.values())[0]170 elif isinstance(inputs[0], (tuple, list)):171 outputs = []172 for input_list in zip(*inputs):173 outputs.append(self._concatenate_batches(input_list))174 if len(outputs) == 1:175 outputs = outputs[0] # If it's a list with only one element, just return the array176 elif isinstance(inputs[0], np.ndarray):177 outputs = self._concatenate_batches(inputs)178 elif isinstance(inputs[0], tf.Tensor):179 outputs = self._concatenate_batches([tensor.numpy() for tensor in inputs])180 else:181 raise TypeError(f"Couldn't handle batch of type {type(inputs[0])}!")182 return outputs183 184 def on_epoch_end(self, epoch, logs=None):185 if hasattr(self.model, "config"):186 ignore_keys = getattr(self.model.config, "keys_to_ignore_at_inference", [])187 else:188 ignore_keys = []189 190 main_input_name = None191 if self.predict_with_generate:192 # This dense conditional recognizes the case where we have an encoder-decoder model, but193 # avoids getting tangled up when we just have a model with a layer called 'encoder'194 if hasattr(self.model, "encoder") and hasattr(self.model.encoder, "main_input_name"):195 main_input_name = self.model.encoder.main_input_name196 else:197 main_input_name = getattr(self.model, "main_input_name", "input_ids")198 199 if self.use_xla_generation and self.generation_function is None:200 201 def generation_function(inputs, attention_mask):202 return self.model.generate(inputs, attention_mask=attention_mask, **self.generate_kwargs)203 204 self.generation_function = tf.function(generation_function, jit_compile=True)205 206 prediction_list = []207 label_list = []208 209 # The whole predict/generate loop is handled inside this method210 for batch in self.eval_dataset:211 if isinstance(batch, tuple):212 batch, labels = batch213 else:214 labels = None215 if self.predict_with_generate:216 if isinstance(batch, dict):217 generation_inputs = batch[main_input_name]218 attention_mask = batch.get("attention_mask", None)219 else:220 generation_inputs = batch221 attention_mask = None222 if self.use_xla_generation:223 predictions = self.generation_function(generation_inputs, attention_mask=attention_mask)224 else:225 predictions = self.model.generate(226 generation_inputs, attention_mask=attention_mask, **self.generate_kwargs227 )228 else:229 predictions = self.model.predict_on_batch(batch)230 if isinstance(predictions, dict):231 # This converts any dict-subclass to a regular dict232 # Keras REALLY doesn't like it when we pass around a BatchEncoding or other derived class233 predictions = dict(predictions)234 if self.output_cols is not None:235 predictions = {key: predictions[key] for key in self.output_cols}236 else:237 predictions = {238 key: val for key, val in predictions.items() if key not in ignore_keys + ["loss"]239 }240 prediction_list.append(predictions)241 if not self.use_keras_label:242 labels = {key: batch[key].numpy() for key in self.label_cols}243 elif isinstance(labels, dict):244 labels = {key: array.numpy() for key, array in labels.items()}245 elif isinstance(labels, (list, tuple)):246 labels = [array.numpy() for array in labels]247 elif isinstance(labels, tf.Tensor):248 labels = labels.numpy()249 else:250 raise TypeError(f"Confused by labels of type {type(labels)}")251 label_list.append(labels)252 253 all_preds = self._postprocess_predictions_or_labels(prediction_list)254 all_labels = self._postprocess_predictions_or_labels(label_list)255 256 metric_output = self.metric_fn((all_preds, all_labels))257 if not isinstance(metric_output, dict):258 raise TypeError(259 f"metric_fn should return a dict mapping metric names to values but instead returned {metric_output}"260 )261 # This is the critical bit - Keras passes a dict containing the loss and standard metric values for this epoch262 # in the logs argument. Ordinarily, this is so the callback can read them, but in this case we write a bunch of263 # new keys in there, which will then get read by the History callback and treated like any other metric value.264 # I promise that I have it in writing from Chollet that this is okay.265 logs.update(metric_output)266 267 268class PushToHubCallback(keras.callbacks.Callback):269 """270 Callback that will save and push the model to the Hub regularly. By default, it pushes once per epoch, but this can271 be changed with the `save_strategy` argument. Pushed models can be accessed like any other model on the hub, such272 as with the `from_pretrained` method.273 274 ```py275 from transformers.keras_callbacks import PushToHubCallback276 277 push_to_hub_callback = PushToHubCallback(278 output_dir="./model_save",279 tokenizer=tokenizer,280 hub_model_id="gpt5-7xlarge",281 )282 283 model.fit(train_dataset, callbacks=[push_to_hub_callback])284 ```285 286 Args:287 output_dir (`str`):288 The output directory where the model predictions and checkpoints will be written and synced with the289 repository on the Hub.290 save_strategy (`str` or [`~trainer_utils.IntervalStrategy`], *optional*, defaults to `"epoch"`):291 The checkpoint save strategy to adopt during training. Possible values are:292 293 - `"no"`: Save is done at the end of training.294 - `"epoch"`: Save is done at the end of each epoch.295 - `"steps"`: Save is done every `save_steps`296 save_steps (`int`, *optional*):297 The number of steps between saves when using the "steps" `save_strategy`.298 tokenizer (`PreTrainedTokenizerBase`, *optional*):299 The tokenizer used by the model. If supplied, will be uploaded to the repo alongside the weights.300 hub_model_id (`str`, *optional*):301 The name of the repository to keep in sync with the local `output_dir`. It can be a simple model ID in302 which case the model will be pushed in your namespace. Otherwise it should be the whole repository name,303 for instance `"user_name/model"`, which allows you to push to an organization you are a member of with304 `"organization_name/model"`.305 306 Will default to the name of `output_dir`.307 hub_token (`str`, *optional*):308 The token to use to push the model to the Hub. Will default to the token in the cache folder obtained with309 `hf auth login`.310 checkpoint (`bool`, *optional*, defaults to `False`):311 Whether to save full training checkpoints (including epoch and optimizer state) to allow training to be312 resumed. Only usable when `save_strategy` is `"epoch"`.313 """314 315 def __init__(316 self,317 output_dir: Union[str, Path],318 save_strategy: Union[str, IntervalStrategy] = "epoch",319 save_steps: Optional[int] = None,320 tokenizer: Optional[PreTrainedTokenizerBase] = None,321 hub_model_id: Optional[str] = None,322 hub_token: Optional[str] = None,323 checkpoint: bool = False,324 **model_card_args,325 ):326 super().__init__()327 if checkpoint and save_strategy != "epoch":328 raise ValueError("Cannot save checkpoints when save_strategy is not 'epoch'!")329 if isinstance(save_strategy, str):330 save_strategy = IntervalStrategy(save_strategy.lower())331 self.save_strategy = save_strategy332 if self.save_strategy == IntervalStrategy.STEPS and (not isinstance(save_steps, int) or save_steps <= 0):333 raise ValueError("Please supply a positive integer argument for save_steps when save_strategy == 'steps'!")334 self.save_steps = save_steps335 output_dir = Path(output_dir)336 337 # Create repo and retrieve repo_id338 if hub_model_id is None:339 hub_model_id = output_dir.absolute().name340 self.hub_model_id = create_repo(repo_id=hub_model_id, exist_ok=True, token=hub_token).repo_id341 342 self.output_dir = output_dir343 self.repo = Repository(str(self.output_dir), clone_from=self.hub_model_id, token=hub_token)344 345 self.tokenizer = tokenizer346 self.last_job = None347 self.checkpoint = checkpoint348 self.training_history = None349 self.model_card_args = model_card_args350 351 def on_train_begin(self, logs=None):352 # Although we can access model.history, we have no guarantees that the History callback will fire before this353 # one, so we keep track of it here too354 self.training_history = []355 356 def on_train_batch_end(self, batch, logs=None):357 if self.save_strategy == IntervalStrategy.STEPS and (batch + 1) % self.save_steps == 0:358 if self.last_job is not None and not self.last_job.is_done:359 return # The last upload is still running, don't start another360 self.model.save_pretrained(self.output_dir)361 if self.tokenizer is not None:362 self.tokenizer.save_pretrained(self.output_dir)363 _, self.last_job = self.repo.push_to_hub(364 commit_message=f"Training in progress steps {batch}", blocking=False365 )366 367 def on_epoch_end(self, epoch, logs=None):368 logs = logs.copy() # Don't accidentally write things that Keras will read later369 if "epoch" not in logs:370 logs["epoch"] = epoch371 self.training_history.append(logs)372 if self.save_strategy == IntervalStrategy.EPOCH:373 if self.last_job is not None and not self.last_job.is_done:374 return # The last upload is still running, don't start another375 self.model.save_pretrained(self.output_dir)376 if self.tokenizer is not None:377 self.tokenizer.save_pretrained(self.output_dir)378 if self.checkpoint:379 checkpoint_dir = os.path.join(self.output_dir, "checkpoint")380 self.model._save_checkpoint(checkpoint_dir, epoch)381 train_summary = TrainingSummary.from_keras(382 model=self.model,383 model_name=self.hub_model_id,384 keras_history=self.training_history,385 **self.model_card_args,386 )387 model_card = train_summary.to_model_card()388 with (self.output_dir / "README.md").open("w") as f:389 f.write(model_card)390 _, self.last_job = self.repo.push_to_hub(391 commit_message=f"Training in progress epoch {epoch}", blocking=False392 )393 394 def on_train_end(self, logs=None):395 # Makes sure the latest version of the model is uploaded396 if self.last_job is not None and not self.last_job.is_done:397 logging.info("Pushing the last epoch to the Hub, this may take a while...")398 while not self.last_job.is_done:399 sleep(1)400 else:401 self.model.save_pretrained(self.output_dir)402 if self.tokenizer is not None:403 self.tokenizer.save_pretrained(self.output_dir)404 train_summary = TrainingSummary.from_keras(405 model=self.model,406 model_name=self.hub_model_id,407 keras_history=self.training_history,408 **self.model_card_args,409 )410 model_card = train_summary.to_model_card()411 with (self.output_dir / "README.md").open("w") as f:412 f.write(model_card)413 self.repo.push_to_hub(commit_message="End of training", blocking=True)414