Spaces:
Sleeping
Sleeping
| # -*- coding: utf-8 -*- | |
| r""" | |
| Translation Ranking Base Model | |
| ============================== | |
| Abstract base class used to build new ranking systems inside Polos. | |
| This task consists of ranking "good" translations above "worse" ones. | |
| """ | |
| from argparse import Namespace | |
| from typing import List | |
| import pandas as pd | |
| import torch | |
| import torch.nn as nn | |
| from polos.models.model_base import ModelBase | |
| from polos.models.utils import average_pooling, max_pooling | |
| from polos.modules.scalar_mix import ScalarMixWithDropout | |
| from polos.metrics import WMTKendall | |
| class RankingBase(ModelBase): | |
| """ | |
| Ranking Model base class used to fine-tune pretrained models such as XLM-R | |
| to produce better sentence embeddings by optmizing Triplet Margin Loss. | |
| :param hparams: Namespace containing the hyperparameters. | |
| """ | |
| def __init__( | |
| self, | |
| hparams: Namespace, | |
| ) -> None: | |
| super().__init__(hparams) | |
| 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[["src", "ref", "pos", "neg"]] | |
| df["src"] = df["src"].astype(str) | |
| df["ref"] = df["ref"].astype(str) | |
| df["pos"] = df["pos"].astype(str) | |
| df["neg"] = df["neg"].astype(str) | |
| return df.to_dict("records") | |
| def _build_loss(self): | |
| """ Initializes the loss function/s. """ | |
| self.loss = nn.TripletMarginLoss(margin=1.0, p=2) | |
| def _build_model(self) -> ModelBase: | |
| """ | |
| Initializes the ranking model architecture. | |
| """ | |
| super()._build_model() | |
| self.metrics = WMTKendall() | |
| 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 | |
| ) | |
| def get_sentence_embedding( | |
| self, tokens: torch.Tensor, lengths: torch.Tensor | |
| ) -> 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) | |
| # for LASER we dont care about the word embeddings | |
| if self.hparams.encoder_model == "LASER": | |
| pass | |
| elif 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" or self.hparams.encoder_model == "LASER": | |
| 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, | |
| ) | |
| elif self.hparams.pool == "cls": | |
| sentemb = embeddings[:, 0, :] | |
| else: | |
| raise Exception("Invalid pooling technique.") | |
| return sentemb | |