Spaces:
Running
Running
# -*- 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 | |