roberta_embeddings.py

import torch
import ot
from sklearn import metrics


roberta = torch.hub.load('pytorch/fairseq', 'roberta.large')
roberta.eval()  # disable dropout (or leave in train mode to finetune)


def get_roberta_features(text):
    return roberta.extract_features(roberta.encode(text)).detach().numpy()[0]


def get_optimal_transport_distance(v, w, sinkhorn_method='sinkhorn', entropy_regularization=0.05):
    dists = metrics.pairwise.euclidean_distances(v, w) ** 2
    transport = ot.sinkhorn([], [], dists, reg=entropy_regularization, method=sinkhorn_method, numItermax=100)
    cost = (dists * transport).sum()
    return cost