Spaces:
Sleeping
Sleeping
File size: 4,438 Bytes
03f6091 |
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 |
# -*- 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)}
|