fix

aligner.py

@@ -30,19 +30,19 @@ class Aligner:
             self.weight_func = compute_weights_norm
 
     def compute_alignment_matrixes(self, s1_word_embeddigs, s2_word_embeddigs):
-        P, Cost, log, similarity_matrix, relative_distance = self.compute_optimal_transport(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, relative_distance
+        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, relative_distance = self.dist_func(s1_word_embeddigs, s2_word_embeddigs, self.distotion)
+        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':
@@ -96,7 +96,7 @@ class Aligner:
         elif self.ot_type == 'none':
             P = 1 - C
 
-        return P, C, log, similarity_matrix, relative_distance
+        return P, C, log, similarity_matrix
 
     def convert_to_numpy(self, s1_weights, s2_weights, C):
         if torch.is_tensor(s1_weights):

app.py

@@ -8,8 +8,6 @@ from nltk.tokenize import word_tokenize
 from transformers import AutoModel, AutoTokenizer
 
 from aligner import Aligner
-
-# from utils import align_matrix_heatmap, plot_align_matrix_heatmap
 from plotools import (
     plot_align_matrix_heatmap_plotly,
     plot_similarity_matrix_heatmap_plotly,
@@ -45,8 +43,8 @@ def init_aligner(
         dist_type="cos",
         weight_type="uniform",
         distortion=distortion,
-        thresh=threshhold,  # 0.25252525252525254
-        tau=tau,  # 0.9803921568627451
+        thresh=threshhold,
+        tau=tau,
         div_type="--",
     )
 
@@ -86,14 +84,14 @@ def main():
         1.0,
         value=0.98,
         help="fraction of fertility to be aligned (fraction of mass to be transported) / penalties",
-    )  # with 0.02 interva
+    )
     threshhold = st.sidebar.slider(
         "threshhold: $\lambda$",
         0.0,
         1.0,
         value=0.22,
         help="sparsity of alignment matrix",
-    )  # with 0.01 interval
+    )
     show_assignments = st.sidebar.checkbox("show assignments", value=True)
     if show_assignments:
         n_neighbors = st.sidebar.slider(

otfuncs.py

@@ -12,11 +12,11 @@ def compute_distance_matrix_cosine(
         torch.matmul(F.normalize(s1_word_embeddigs), F.normalize(s2_word_embeddigs).t())
         + 1.0
     ) / 2  # Range 0-1
-    C, relative_distance = apply_distortion(sim_matrix, distortion_ratio)
+    C = apply_distortion(sim_matrix, distortion_ratio)
     C = min_max_scaling(C)  # Range 0-1
     C = 1.0 - C  # Convert to distance
 
-    return C, sim_matrix, relative_distance
+    return C, sim_matrix
 
 
 def compute_distance_matrix_l2(s1_word_embeddigs, s2_word_embeddigs, distortion_ratio):
@@ -48,7 +48,7 @@ def apply_distortion(sim_matrix, ratio):
 
     sim_matrix = torch.mul(sim_matrix, distortion_mask)
 
-    return sim_matrix, relative_distance
+    return sim_matrix
 
 
 def compute_weights_norm(s1_word_embeddigs, s2_word_embeddigs):