Use the C++ implementation of Levenshtein instead of the Python one.

metrics_test.py running time reduced from 136.061s to 0.551s

DiarizationLM/diarizationlm/__init__.py

@@ -1,10 +1,8 @@
 """__init__ file."""
 
-from . import levenshtein
 from . import utils
 from . import metrics
 
-levenshtein_with_edits = levenshtein.levenshtein_with_edits
 PromptOptions = utils.PromptOptions
 transcript_preserving_speaker_transfer = (
     utils.transcript_preserving_speaker_transfer)

DiarizationLM/diarizationlm/levenshtein.py

@@ -1,4 +1,8 @@
-"""Function for the Levenshtein algorithm."""
+"""Function for the Levenshtein algorithm.
+
+Note: This Python implementation is very inefficient. Please use this C++
+implementation instead: https://github.com/wq2012/word_levenshtein
+"""
 import numpy as np
 from enum import Enum
 

DiarizationLM/diarizationlm/levenshtein_test.py

@@ -1,4 +1,4 @@
-"""Test levenshterin."""
+"""Test levenshtein."""
 import unittest
 from diarizationlm import levenshtein
 
@@ -21,6 +21,15 @@ def test_levenshtein_with_edits_2(self):
       [(0, 0), (1, -1), (2, 1), (-1, 2), (3, 3), (4, 4), (5, 5)],
       align[1])
 
+  def test_levenshtein_with_edit_multiple_spaces(self):
+    s1 = "     hello  good    morning how are you"
+    s2 = "hello morning hi how    are you  "
+    align = levenshtein.levenshtein_with_edits(s1, s2)
+    self.assertEqual(2, align[0])
+    self.assertListEqual(
+      [(0, 0), (1, -1), (2, 1), (-1, 2), (3, 3), (4, 4), (5, 5)],
+      align[1])
+
 
 if __name__ == "__main__":
   unittest.main()

DiarizationLM/diarizationlm/metrics.py

@@ -28,7 +28,7 @@
 from scipy import optimize
 import tqdm
 from diarizationlm import utils
-from diarizationlm import levenshtein
+import word_levenshtein as levenshtein
 
 
 @dataclasses.dataclass

DiarizationLM/diarizationlm/utils.py

@@ -9,7 +9,7 @@
 import numpy as np
 from scipy import optimize
 
-from diarizationlm import levenshtein
+import word_levenshtein as levenshtein
 
 PUNCTUATIONS = [",", ".", "_", "?", "!", "-", '"', "'"]
 

DiarizationLM/requirements.txt

@@ -6,3 +6,4 @@ openai
 datasets
 tqdm
 colortimelog
+word_levenshtein

DiarizationLM/setup.py

@@ -2,7 +2,7 @@
 
 import setuptools
 
-VERSION = "0.0.11"
+VERSION = "0.1.0"
 
 with open("README.md", "r") as file_object:
     LONG_DESCRIPTION = file_object.read()