[update] Change the behavior of WordAlign class

- Restore 'align' method
m-yoshinaka · Oct 5, 2020 · 611f819 · 611f819
1 parent c79fb96
commit 611f819
Showing 1 changed file with 10 additions and 7 deletions.
diff --git a/sapphire/word_alignment.py b/sapphire/word_alignment.py
@@ -52,13 +52,7 @@ def __init__(self, lambda_, use_hungarian):
         self.use_hungarian = use_hungarian
 
     def __call__(self, sim_matrix):
-        if self.use_hungarian:
-            alignments = self._hungarian_assign(sim_matrix)
-        else:
-            alignments = self._grow_diag_final(sim_matrix)
-
-        return [(s + 1, t + 1) for s, t in alignments
-                if sim_matrix[s][t] >= self.lambda_]
+        return self.align(sim_matrix=sim_matrix)
 
     @staticmethod
     def _hungarian_assign(sim_matrix):
@@ -120,3 +114,12 @@ def _final(matrix):
                 alignments.append((s, t))
 
         return alignments
+
+    def align(self, sim_matrix):
+        if self.use_hungarian:
+            alignments = self._hungarian_assign(sim_matrix)
+        else:
+            alignments = self._grow_diag_final(sim_matrix)
+
+        return [(s + 1, t + 1) for s, t in alignments
+                if sim_matrix[s][t] >= self.lambda_]  # 1-index alignment