polos/models/estimators/estimator_base.py

# -*- 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