|
from sentence_transformers.evaluation import SentenceEvaluator |
|
from sentence_transformers.util import batch_to_device |
|
from sentence_transformers import SentenceTransformer |
|
from typing import List, Tuple, Dict |
|
import torch |
|
import numpy as np |
|
import logging |
|
import os |
|
import csv |
|
|
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
class MSEEvaluatorFromDataFrame(SentenceEvaluator): |
|
""" |
|
Computes the mean squared error (x100) between the computed sentence embedding |
|
and some target sentence embedding. |
|
:param dataframe: |
|
It must have the following format. Rows contains different, parallel sentences. Columns are the respective language codes |
|
[{'en': 'My sentence', 'es': 'Sentence in Spanisch', 'fr': 'Sentence in French'...}, |
|
{'en': 'My second sentence', ....] |
|
:param combinations: |
|
Must be of the format [('en', 'es'), ('en', 'fr'), ...] |
|
First entry in a tuple is the source language. The sentence in the respective language will be fetched from the dataframe and passed to the teacher model. |
|
Second entry in a tuple the the target language. Sentence will be fetched from the dataframe and passed to the student model |
|
""" |
|
def __init__(self, dataframe: List[Dict[str, str]], teacher_model: SentenceTransformer, combinations: List[Tuple[str, str]], batch_size: int = 8, name='', write_csv: bool = True): |
|
|
|
self.combinations = combinations |
|
self.name = name |
|
self.batch_size = batch_size |
|
|
|
|
|
if name: |
|
name = "_"+name |
|
|
|
self.csv_file = "mse_evaluation" + name + "_results.csv" |
|
self.csv_headers = ["epoch", "steps"] |
|
self.write_csv = write_csv |
|
self.data = {} |
|
|
|
logger.info("Compute teacher embeddings") |
|
all_source_sentences = set() |
|
for src_lang, trg_lang in self.combinations: |
|
src_sentences = [] |
|
trg_sentences = [] |
|
|
|
for row in dataframe: |
|
if row[src_lang].strip() != "" and row[trg_lang].strip() != "": |
|
all_source_sentences.add(row[src_lang]) |
|
src_sentences.append(row[src_lang]) |
|
trg_sentences.append(row[trg_lang]) |
|
|
|
self.data[(src_lang, trg_lang)] = (src_sentences, trg_sentences) |
|
self.csv_headers.append("{}-{}".format(src_lang, trg_lang)) |
|
|
|
all_source_sentences = list(all_source_sentences) |
|
all_src_embeddings = teacher_model.encode(all_source_sentences, batch_size=self.batch_size) |
|
self.teacher_embeddings = {sent: emb for sent, emb in zip(all_source_sentences, all_src_embeddings)} |
|
|
|
def __call__(self, model, output_path: str = None, epoch: int = -1, steps: int = -1): |
|
model.eval() |
|
|
|
mse_scores = [] |
|
for src_lang, trg_lang in self.combinations: |
|
src_sentences, trg_sentences = self.data[(src_lang, trg_lang)] |
|
|
|
src_embeddings = np.asarray([self.teacher_embeddings[sent] for sent in src_sentences]) |
|
trg_embeddings = np.asarray(model.encode(trg_sentences, batch_size=self.batch_size)) |
|
|
|
mse = ((src_embeddings - trg_embeddings) ** 2).mean() |
|
mse *= 100 |
|
mse_scores.append(mse) |
|
|
|
logger.info("MSE evaluation on {} dataset - {}-{}:".format(self.name, src_lang, trg_lang)) |
|
logger.info("MSE (*100):\t{:4f}".format(mse)) |
|
|
|
if output_path is not None and self.write_csv: |
|
csv_path = os.path.join(output_path, self.csv_file) |
|
output_file_exists = os.path.isfile(csv_path) |
|
with open(csv_path, newline='', mode="a" if output_file_exists else 'w', encoding="utf-8") as f: |
|
writer = csv.writer(f) |
|
if not output_file_exists: |
|
writer.writerow(self.csv_headers) |
|
|
|
writer.writerow([epoch, steps]+mse_scores) |
|
|
|
return -np.mean(mse_scores) |
|
|
|
|
