# -*- coding: utf-8 -*- r""" Estimator Base Model ======================= Abstract base class used to build new estimator models inside Polos. """ from argparse import Namespace from typing import Dict, List, Union import pandas as pd import torch import torch.nn as nn from tqdm import tqdm from polos.metrics import RegressionReport from polos.models.model_base import ModelBase from polos.models.utils import average_pooling, max_pooling, move_to_cpu, move_to_cuda class Estimator(ModelBase): """ Estimator base class that uses an Encoder to encode sequences. :param hparams: Namespace containing the hyperparameters. """ class ModelConfig(ModelBase.ModelConfig): """ Estimator ModelConfig: --------------------------- Encoder ----------------------------------------- :param encoder_learning_rate: Learning rate used for the encoder model. :param layerwise_decay: Decay for the layer wise learning rates. If 1.0 no decay is applied. :param layer: Layer that will be used to extract embeddings. If 'mix' embeddings from all layers are combined with a layer-wise attention mechanism :param scalar_mix_dropout: If layer='mix' we can regularize layer's importance by with a given probability setting that weight to - inf before softmax. ------------------------- Feed Forward --------------------------------------- :param loss: Loss function to be used (options: binary_xent, mse). :param hidden_sizes: String with size of the hidden layers in the feedforward. :param activations: Activation functions to be used in the feedforward :param dropout: Dropout probability to be applied to the feedforward :param final_activation: Activation function to be applied after getting the final regression score. Set to False if you wish to perform an 'unbounded' regression. """ encoder_learning_rate: float = 1e-06 layerwise_decay: float = 1.0 layer: str = "mix" scalar_mix_dropout: float = 0.0 loss: str = "mse" hidden_sizes: str = "1024" activations: str = "Tanh" dropout: float = 0.1 final_activation: str = "Sigmoid" def __init__(self, hparams: Namespace) -> None: super().__init__(hparams) def _build_model(self) -> ModelBase: """ Initializes the estimator architecture. """ super()._build_model() self.metrics = RegressionReport() def _build_loss(self): """ Initializes the loss function/s. """ super()._build_loss() if self.hparams.loss == "mse": self.loss = nn.MSELoss(reduction="sum") elif self.hparams.loss == "binary_xent": self.loss = nn.BCELoss(reduction="sum") else: raise Exception("{} is not a valid loss option.".format(self.hparams.loss)) def read_csv(self, path: str) -> List[dict]: """Reads a comma separated value file. :param path: path to a csv file. :return: List of records as dictionaries """ df = pd.read_csv(path) df = df[["mt","refs","score", "imgid"]] refs_list = [] for refs in df["refs"]: refs = eval(refs) refs_list.append(refs) df["refs"] = refs_list df["mt"] = df["mt"].astype(str) df["score"] = df["score"].astype(float) df["imgid"] = df["imgid"].astype(str) return df.to_dict("records") def compute_loss( self, model_out: Dict[str, torch.Tensor], targets: Dict[str, torch.Tensor] ) -> torch.Tensor: """ Computes Loss value according to a loss function. :param model_out: model specific output. Must contain a key 'score' with a tensor [batch_size x 1] with model predictions :param targets: Target score values [batch_size] """ assert torch.all((targets["score"] >= 0) & (targets["score"] <= 1)), f"gt({targets['score']}) is not in [0,1]!" # print("out",model_out["score"].view(-1).shape, targets["score"].shape) return self.loss(model_out["score"].view(-1), targets["score"]) def compute_metrics(self, outputs: List[Dict[str, torch.Tensor]]) -> dict: """ Private function that computes metrics of interest based on model predictions and respective targets. """ predictions = ( torch.cat([batch["val_prediction"]["score"].view(-1) for batch in outputs]) .cpu() .numpy() ) targets = ( torch.cat([batch["val_target"]["score"] for batch in outputs]).cpu().numpy() ) self.draw_histogram(predictions, targets) return self.metrics.compute(predictions, targets) def draw_histogram(self, predictions, targets): import matplotlib import matplotlib.pyplot as plt matplotlib.use('Agg') plt.figure(figsize=(10, 5)) plt.subplot(1, 2, 1) plt.hist(predictions, bins=20, alpha=0.5, color='blue', label='Predictions') plt.title('Predictions') plt.xlabel('Value') plt.ylabel('Frequency') plt.subplot(1, 2, 2) plt.hist(targets, bins=20, alpha=0.5, color='red', label='Targets') plt.title('Targets') plt.xlabel('Value') plt.ylabel('Frequency') plt.tight_layout() # plt.show() plt.savefig("histogram.png") def get_sentence_embedding( self, tokens: torch.Tensor, lengths: torch.Tensor, pooling=True ) -> torch.Tensor: """Auxiliar function that extracts sentence embeddings for a single sentence. :param tokens: sequences [batch_size x seq_len] :param lengths: lengths [batch_size] :return: torch.Tensor [batch_size x hidden_size] """ # When using just one GPU this should not change behavior # but when splitting batches across GPU the tokens have padding # from the entire original batch if self.trainer and self.trainer.use_dp and self.trainer.num_gpus > 1: tokens = tokens[:, : lengths.max()] encoder_out = self.encoder(tokens, lengths) if self.scalar_mix: embeddings = self.scalar_mix(encoder_out["all_layers"], encoder_out["mask"]) elif self.layer >= 0 and self.layer < self.encoder.num_layers: embeddings = encoder_out["all_layers"][self.layer] else: raise Exception("Invalid model layer {}.".format(self.layer)) if self.hparams.pool == "default": sentemb = encoder_out["sentemb"] elif self.hparams.pool == "max": sentemb = max_pooling( tokens, embeddings, self.encoder.tokenizer.padding_index ) elif self.hparams.pool == "avg": sentemb = average_pooling( tokens, embeddings, encoder_out["mask"], self.encoder.tokenizer.padding_index, ) # print("sentemb",sentemb[0,:]) return sentemb, embeddings, encoder_out["mask"], self.encoder.tokenizer.padding_index elif self.hparams.pool == "cls": sentemb = embeddings[:, 0, :] return sentemb, embeddings, encoder_out["mask"], self.encoder.tokenizer.padding_index elif self.hparams.pool == "cls+avg": cls_sentemb = embeddings[:, 0, :] avg_sentemb = average_pooling( tokens, embeddings, encoder_out["mask"], self.encoder.tokenizer.padding_index, ) sentemb = torch.cat((cls_sentemb, avg_sentemb), dim=1) else: raise Exception("Invalid pooling technique.") return sentemb def predict( self, samples: List[Dict[str, str]], cuda: bool = False, show_progress: bool = True, batch_size: int = -1, ) -> (Dict[str, Union[str, float]], List[float]): """Function that runs a model prediction, :param samples: List of dictionaries with 'mt' and 'ref' keys. :param cuda: Flag that runs inference using 1 single GPU. :param show_progress: Flag to show progress during inference of multiple examples. :para batch_size: Batch size used during inference. By default uses the same batch size used during training. :return: Dictionary with original samples, predicted scores and langid results for SRC and MT + list of predicted scores """ if self.training: self.eval() if cuda and torch.cuda.is_available(): self.to("cuda") batch_size = self.hparams.batch_size if batch_size < 1 else batch_size with torch.no_grad(): batches = [ samples[i : i + batch_size] for i in range(0, len(samples), batch_size) ] model_inputs = [] if show_progress: pbar = tqdm( total=len(batches), desc="Preparing batches...", dynamic_ncols=True, leave=None, ) for batch in batches: batch = self.prepare_sample(batch, inference=True) model_inputs.append(batch) if show_progress: pbar.update(1) if show_progress: pbar.close() if show_progress: pbar = tqdm( total=len(batches), desc="Scoring hypothesis...", dynamic_ncols=True, leave=None, ) scores = [] for model_input in model_inputs: if cuda and torch.cuda.is_available(): model_input = move_to_cuda(model_input) model_out = self.forward(**model_input) model_out = move_to_cpu(model_out) else: model_out = self.forward(**model_input) model_scores = model_out["score"].numpy().tolist() for i in range(len(model_scores)): scores.append(model_scores[i][0]) if show_progress: pbar.update(1) if show_progress: pbar.close() assert len(scores) == len(samples) for i in range(len(scores)): samples[i]["predicted_score"] = scores[i] return samples, scores def document_predict( self, documents: List[Dict[str, List[str]]], cuda: bool = False, show_progress: bool = False, ) -> (Dict[str, Union[str, float]], List[float]): """Function that scores entire documents by processing all segments in parallel. :param documents: List of dictionaries with 'mt', 'src' and 'ref' keys where each key is a list of segments. :param cuda: Flag that runs inference using 1 single GPU. :param show_progress: Flag to show progress during inference of multiple examples. :return: tuple with Dictionary with original samples and predicted document score, micro average scores, macro average scores. """ if self.training: self.eval() if cuda and torch.cuda.is_available(): self.to("cuda") inputs, lengths = [], [] for d in documents: d = [dict(zip(d, t)) for t in zip(*d.values())] # For very long documents we need to create chunks. # (64 sentences per chunk) if len(d) > 64: document_chunks, document_lengths = [], [] chunks = [d[i : i + 64] for i in range(0, len(d), 64)] for chunk in chunks: chunk = self.prepare_sample(chunk, inference=True) document_lengths.append(chunk["mt_lengths"]) if cuda and torch.cuda.is_available(): document_chunks.append(chunk) lengths.append(torch.cat(document_lengths, dim=0)) inputs.append(document_chunks) else: d_input = self.prepare_sample(d, inference=True) lengths.append(d_input["mt_lengths"]) if cuda and torch.cuda.is_available(): inputs.append(d_input) micro_average, average = [], [] for doc, seg_lengths in tqdm( zip(inputs, lengths), total=len(inputs), desc="Scoring Documents ...", dynamic_ncols=True, leave=None, ): if isinstance(doc, list): seg_scores = [] for chunk in doc: model_output = self.forward(**move_to_cuda(chunk)) seg_scores.append(move_to_cpu(model_output)["score"].view(1, -1)[0]) seg_scores = torch.cat(seg_scores, dim=0) else: model_output = self.forward(**move_to_cuda(doc)) seg_scores = move_to_cpu(model_output)["score"].view(1, -1)[0] # Invert segment-level scores for HTER # seg_scores = torch.ones_like(seg_scores) - seg_scores micro = (seg_scores * seg_lengths).sum() / seg_lengths.sum() macro = seg_scores.sum() / seg_scores.size()[0] micro_average.append(micro.item()) average.append(macro.item()) assert len(micro_average) == len(documents) for i in range(len(documents)): documents[i]["predicted_score"] = micro_average[i] return documents, micro_average, average