Spaces:
Sleeping
Sleeping
| # -*- coding: utf-8 -*- | |
| r""" | |
| Quality Estimator Model | |
| ======================= | |
| Quality Estimator Model estimates a quality score for the hyphotesis (e.g: the MT text) | |
| by looking only at source and MT. | |
| """ | |
| from argparse import Namespace | |
| from typing import Dict, List, Tuple, Union | |
| import pandas as pd | |
| import torch | |
| from polos.models.estimators import PolosEstimator, Estimator | |
| from polos.modules.feedforward import FeedForward | |
| from polos.modules.scalar_mix import ScalarMixWithDropout | |
| from torchnlp.utils import collate_tensors | |
| class QualityEstimator(PolosEstimator): | |
| """ | |
| Estimator class that uses a pretrained encoder to extract features from | |
| the sequences and then passes those features to a feed forward estimator. | |
| :param hparams: Namespace containing the hyperparameters. | |
| """ | |
| def __init__( | |
| self, | |
| hparams: Namespace, | |
| ) -> None: | |
| super().__init__(hparams) | |
| def _build_model(self) -> Estimator: | |
| """ | |
| Initializes the estimator architecture. | |
| """ | |
| super()._build_model() | |
| if self.hparams.encoder_model != "LASER": | |
| self.layer = ( | |
| int(self.hparams.layer) | |
| if self.hparams.layer != "mix" | |
| else self.hparams.layer | |
| ) | |
| self.scalar_mix = ( | |
| ScalarMixWithDropout( | |
| mixture_size=self.encoder.num_layers, | |
| dropout=self.hparams.scalar_mix_dropout, | |
| do_layer_norm=True, | |
| ) | |
| if self.layer == "mix" and self.hparams.pool != "default" | |
| else None | |
| ) | |
| self.ff = FeedForward( | |
| in_dim=self.encoder.output_units * 4, | |
| hidden_sizes=self.hparams.hidden_sizes, | |
| activations=self.hparams.activations, | |
| dropout=self.hparams.dropout, | |
| ) | |
| 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", "src", "score"]] | |
| df["mt"] = df["mt"].astype(str) | |
| df["src"] = df["src"].astype(str) | |
| df["score"] = df["score"].astype(float) | |
| return df.to_dict("records") | |
| def prepare_sample( | |
| self, sample: List[Dict[str, Union[str, float]]], inference: bool = False | |
| ) -> Union[ | |
| Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]], Dict[str, torch.Tensor] | |
| ]: | |
| """ | |
| Function that prepares a sample to input the model. | |
| :param sample: list of dictionaries. | |
| :param inference: If set to true prepares only the model inputs. | |
| :returns: Tuple with 2 dictionaries (model inputs and targets). | |
| If `inference=True` returns only the model inputs. | |
| """ | |
| sample = collate_tensors(sample) | |
| mt_inputs = self.encoder.prepare_sample(sample["mt"]) | |
| src_inputs = self.encoder.prepare_sample(sample["src"]) | |
| mt_inputs = {"mt_" + k: v for k, v in mt_inputs.items()} | |
| src_inputs = {"src_" + k: v for k, v in src_inputs.items()} | |
| inputs = {**mt_inputs, **src_inputs} | |
| if inference: | |
| return inputs | |
| targets = {"score": torch.tensor(sample["score"], dtype=torch.float)} | |
| return inputs, targets | |
| def forward( | |
| self, | |
| mt_tokens: torch.tensor, | |
| src_tokens: torch.tensor, | |
| mt_lengths: torch.tensor, | |
| src_lengths: torch.tensor, | |
| **kwargs | |
| ) -> Dict[str, torch.Tensor]: | |
| """ | |
| Function that encodes both Source, MT and returns a quality score. | |
| :param mt_tokens: MT sequences [batch_size x mt_seq_len] | |
| :param src_tokens: SRC sequences [batch_size x src_seq_len] | |
| :param mt_lengths: MT lengths [batch_size] | |
| :param src_lengths: SRC lengths [batch_size] | |
| :return: Dictionary with model outputs to be passed to the loss function. | |
| """ | |
| mt_sentemb = self.get_sentence_embedding(mt_tokens, mt_lengths) | |
| src_sentemb = self.get_sentence_embedding(src_tokens, src_lengths) | |
| diff_src = torch.abs(mt_sentemb - src_sentemb) | |
| prod_src = mt_sentemb * src_sentemb | |
| embedded_sequences = torch.cat( | |
| (mt_sentemb, src_sentemb, prod_src, diff_src), dim=1 | |
| ) | |
| return {"score": self.ff(embedded_sequences)} | |