|
import numpy as np |
|
import torch |
|
import ot |
|
from otfuncs import ( |
|
compute_distance_matrix_cosine, |
|
compute_distance_matrix_l2, |
|
compute_weights_norm, |
|
compute_weights_uniform, |
|
min_max_scaling |
|
) |
|
|
|
class Aligner: |
|
def __init__(self, ot_type, sinkhorn, dist_type, weight_type, distortion, thresh, tau, **kwargs): |
|
self.ot_type = ot_type |
|
self.sinkhorn = sinkhorn |
|
self.dist_type = dist_type |
|
self.weight_type = weight_type |
|
self.distotion = distortion |
|
self.thresh = thresh |
|
self.tau = tau |
|
self.epsilon = 0.1 |
|
self.stopThr = 1e-6 |
|
self.numItermax = 1000 |
|
self.div_type = kwargs['div_type'] |
|
|
|
self.dist_func = compute_distance_matrix_cosine if dist_type == 'cos' else compute_distance_matrix_l2 |
|
if weight_type == 'uniform': |
|
self.weight_func = compute_weights_uniform |
|
else: |
|
self.weight_func = compute_weights_norm |
|
|
|
def compute_alignment_matrixes(self, s1_word_embeddigs, s2_word_embeddigs): |
|
P, Cost, log, similarity_matrix = self.compute_optimal_transport(s1_word_embeddigs, s2_word_embeddigs) |
|
print(log.keys()) |
|
if torch.is_tensor(P): |
|
P = P.to('cpu').numpy() |
|
loss = log.get('cost', 'NotImplemented') |
|
|
|
return P, Cost, loss, similarity_matrix |
|
|
|
def compute_optimal_transport(self, s1_word_embeddigs, s2_word_embeddigs): |
|
s1_word_embeddigs = s1_word_embeddigs.to(torch.float64) |
|
s2_word_embeddigs = s2_word_embeddigs.to(torch.float64) |
|
|
|
C, similarity_matrix = self.dist_func(s1_word_embeddigs, s2_word_embeddigs, self.distotion) |
|
s1_weights, s2_weights = self.weight_func(s1_word_embeddigs, s2_word_embeddigs) |
|
|
|
if self.ot_type == 'ot': |
|
s1_weights = s1_weights / s1_weights.sum() |
|
s2_weights = s2_weights / s2_weights.sum() |
|
s1_weights, s2_weights, C = self.convert_to_numpy(s1_weights, s2_weights, C) |
|
|
|
if self.sinkhorn: |
|
P, log = ot.bregman.sinkhorn_log( |
|
s1_weights, s2_weights, C, |
|
reg=self.epsilon, stopThr=self.stopThr, |
|
numItermax=self.numItermax, log=True |
|
) |
|
else: |
|
P, log = ot.emd(s1_weights, s2_weights, C, log=True) |
|
|
|
P = min_max_scaling(P) |
|
|
|
elif self.ot_type == 'pot': |
|
s1_weights, s2_weights, C = self.convert_to_numpy(s1_weights, s2_weights, C) |
|
m = np.min((np.sum(s1_weights), np.sum(s2_weights))) * self.tau |
|
|
|
if self.sinkhorn: |
|
P, log = ot.partial.entropic_partial_wasserstein( |
|
s1_weights, s2_weights, C, |
|
reg=self.epsilon, |
|
m=m, stopThr=self.stopThr, numItermax=self.numItermax, log=True |
|
) |
|
else: |
|
|
|
P, log = ot.partial.partial_wasserstein(s1_weights, s2_weights, C, m=m, log=True) |
|
|
|
P = min_max_scaling(P) |
|
|
|
elif 'uot' in self.ot_type: |
|
tau = self.tau |
|
|
|
if self.ot_type == 'uot': |
|
P, log = ot.unbalanced.sinkhorn_stabilized_unbalanced( |
|
s1_weights, s2_weights, C, reg=self.epsilon, reg_m=tau, |
|
stopThr=self.stopThr, numItermax=self.numItermax, log=True |
|
) |
|
elif self.ot_type == 'uot-mm': |
|
P, log = ot.unbalanced.mm_unbalanced( |
|
s1_weights, s2_weights, C, reg_m=tau, div=self.div_type, |
|
stopThr=self.stopThr, numItermax=self.numItermax, log=True |
|
) |
|
|
|
P = min_max_scaling(P) |
|
|
|
elif self.ot_type == 'none': |
|
P = 1 - C |
|
|
|
return P, C, log, similarity_matrix |
|
|
|
def convert_to_numpy(self, s1_weights, s2_weights, C): |
|
if torch.is_tensor(s1_weights): |
|
s1_weights = s1_weights.to('cpu').numpy() |
|
s2_weights = s2_weights.to('cpu').numpy() |
|
if torch.is_tensor(C): |
|
C = C.to('cpu').numpy() |
|
|
|
return s1_weights, s2_weights, C |
|
|
