Speed up DSAlign

WORKSPACE

@@ -15,6 +15,32 @@ git_repository(
     tag = "2.0.0",
 )
 
+git_repository(
+    name = "cython",
+    remote = "https://github.com/cython/cython",
+    tag = "3.0a7",
+)
+
+# This is not robust. However, I don't know of a way to configure
+# bazel to find the correct python. It appears that tensorflow and
+# grpc have ways to do it, though.
+
+new_local_repository(
+    name = "python",
+    path = "/install/miniconda3/envs/100k-hours-lingvo-3",
+    build_file_content = """
+cc_library(
+    name = "python-lib",
+    hdrs = glob(["include/python3.7m/*.h"]) +
+           glob(["lib/python3.7/site-packages/numpy/core/include/numpy/*.h",
+                 "lib/python3.7/site-packages/numpy/core/include/numpy/random/*.h"]),
+    includes = ["include/python3.7m",
+                "lib/python3.7/site-packages/numpy/core/include"],
+    visibility = ["//visibility:public"]
+)
+    """
+)
+
 cc_tf_configure()
 
 lingvo_testonly_deps()

environment.yml

@@ -121,6 +121,7 @@ dependencies:
     - grpc-google-iam-v1==0.12.3
     - grpcio==1.30.0
     - grpcio-gcp==0.2.2
+    - gruut==1.2.2
     - gviz-api==1.9.0
     - h5py==2.10.0
     - hdfs==2.5.8
@@ -178,6 +179,7 @@ dependencies:
     - prompt-toolkit==3.0.5
     - protobuf==3.12.2
     - ptyprocess==0.6.0
+    - py4j==0.10.9
     - pyarrow==1.0.0
     - pyasn1==0.4.8
     - pyasn1-modules==0.2.8

galvasr2/BUILD

@@ -1,99 +1,3 @@
-py_binary(
-    name = "spark_forced_aligner",
-    main = "align/spark/align.py",
-    imports = ["align"],
-    srcs = [
-           "align/spark/align.py",
-           "align/spark/event_listener.py",
-           "align/spark/schemas.py",
-           ] +
-           glob(["align/*.py"]),
-    deps = [
-        "//lingvo/tools:audio_lib", # audio_lib.DecodeToWav
-        # Implicit pyspark dependency
-    ],
-    data = [
-         # trainer will be called as a subprocess
-         "//lingvo:trainer"
-    ],
-    visibility = ["//visibility:public"],
-)
-
-py_binary(
-    name = "cuda_decoder_forced_aligner",
-    main = "align/spark/align_cuda_decoder.py",
-    imports = ["align"],
-    srcs = [
-           "align/spark/align_cuda_decoder.py",
-           "align/spark/align_lib.py",
-           "align/spark/dsalign_lib.py",
-           "align/spark/event_listener.py",
-           "align/spark/schemas.py",
-           "align/spark/timeout.py",
-           ] +
-           glob(["align/*.py"]),
-    deps = [
-    ],
-    data = [
-    ],
-    visibility = ["//visibility:public"],
-)
-
-
-py_binary(
-    name = "spark_hours_per_language",
-    main = "align/spark/hours_per_language.py",
-    imports = ["align"],
-    srcs = [
-           "align/spark/hours_per_language.py",
-           "align/spark/align_lib.py",
-           "align/spark/event_listener.py",
-           "align/spark/schemas.py",
-           "align/spark/timeout.py",
-           ] +
-           glob(["align/*.py"]),
-    deps = [
-    ],
-    data = [
-    ],
-    visibility = ["//visibility:public"],
-)
-
-py_binary(
-    name = "spark_hours_per_license",
-    main = "align/spark/hours_per_license.py",
-    imports = ["align"],
-    srcs = [
-           "align/spark/hours_per_license.py",
-           "align/spark/align_lib.py",
-           "align/spark/event_listener.py",
-           "align/spark/schemas.py",
-           ] +
-           glob(["align/*.py"]),
-    deps = [
-    ],
-    data = [
-    ],
-    visibility = ["//visibility:public"],
-)
-
-py_binary(
-    name = "spark_sampling_rate",
-    main = "align/spark/sampling_rate.py",
-    imports = ["align"],
-    srcs = [
-           "align/spark/sampling_rate.py",
-           "align/spark/align_lib.py",
-           "align/spark/schemas.py",
-           ] +
-           glob(["align/*.py"]),
-    deps = [
-    ],
-    data = [
-    ],
-    visibility = ["//visibility:public"],
-)
-
 py_binary(
     name = "galvasr_tokenize_words",
     srcs = ["galvasr_tokenize_words.py"],
@@ -104,37 +8,30 @@ py_binary(
     "//lingvo/core:py_utils"
 ])
 
-# bazel run galvasr2:dump_cc_by_licenses
-# bazel build galvasr2:dump_cc_by_licenses
-# python dump_cc_by_licenses.py
-py_binary(
-    name = "dump_cc_by_licenses",
-    main = "dump_cc_by_licenses.py",
-    srcs = ["dump_cc_by_licenses.py",
-            "align/spark/schemas.py",],
-    python_version = "PY3",
-    srcs_version = "PY3",
-    visibility = ["//visibility:public"],
-)
-
-py_binary(
-    name = "ner_pipeline_transcript",
-    main = "ner_pipeline_transcript.py",
-    srcs = ["ner_pipeline_transcript.py",
-            "align/spark/align_lib.py",
-            "align/spark/schemas.py",],
-    python_version = "PY3",
-    srcs_version = "PY3",
-    visibility = ["//visibility:public"],
-)
-
-py_binary(
-    name = "zero_shot_classification_pipeline",
-    main = "zero_shot_classification_pipeline.py",
-    srcs = ["zero_shot_classification_pipeline.py",
-            "align/spark/align_lib.py",
-            "align/spark/schemas.py",],
-    python_version = "PY3",
-    srcs_version = "PY3",
-    visibility = ["//visibility:public"],
-)
+py_library(
+    name = "utils",
+    srcs = ["utils.py"],
+    visibility = ["//visibility:public"],
+)
+
+# py_binary(
+#     name = "ner_pipeline_transcript",
+#     main = "ner_pipeline_transcript.py",
+#     srcs = ["ner_pipeline_transcript.py",
+#             "align/spark/align_lib.py",
+#             "align/spark/schemas.py",],
+#     python_version = "PY3",
+#     srcs_version = "PY3",
+#     visibility = ["//visibility:public"],
+# )
+
+# py_binary(
+#     name = "zero_shot_classification_pipeline",
+#     main = "zero_shot_classification_pipeline.py",
+#     srcs = ["zero_shot_classification_pipeline.py",
+#             "align/spark/align_lib.py",
+#             "align/spark/schemas.py",],
+#     python_version = "PY3",
+#     srcs_version = "PY3",
+#     visibility = ["//visibility:public"],
+# )

galvasr2/align/BUILD

@@ -0,0 +1,15 @@
+load("//bazel:cython_rules.bzl", "pyx_library")
+
+pyx_library(name = 'smith_waterman',
+            srcs = ['smith_waterman.pyx'], # , 'a.pxd', 'b.py', 'pkg/__init__.py', 'pkg/c.pyx'],
+            deps = ["@python//:python-lib"],
+)
+
+
+py_library(
+    name = "dsalign_import",
+    imports =  ["."],
+    srcs = glob(["*.py"]),
+    deps = [":smith_waterman"],
+    visibility = ["//visibility:public"],
+)

galvasr2/align/search.py

@@ -1,6 +1,8 @@
 from collections import Counter
 from text import ngrams, similarity
 
+from galvasr2.align.smith_waterman import sw_align as sw_align_sped_up
+
 
 class FuzzySearch(object):
     def __init__(
@@ -20,6 +22,9 @@ def __init__(
         self.mismatch_score = mismatch_score
         self.gap_score = gap_score
         self.char_similarities = char_similarities
+        assert (
+            self.char_similarities is None
+        ), "Custom character similarities not supported at this time"
         self.ngrams = {}
         # build inverted index of ngram to where it occurs
         # character ngrams. Good.
@@ -38,11 +43,38 @@ def char_pair(a, b):
 
     def char_similarity(self, a, b):
         key = FuzzySearch.char_pair(a, b)
+        assert (
+            self.char_similarities is None
+        ), "Custom character similarities not supported at this time"
         if self.char_similarities and key in self.char_similarities:
             return self.char_similarities[key]
         return self.match_score if a == b else self.mismatch_score
 
     def sw_align(self, a, start, end):
+        align_start, align_end, score, substitutions = sw_align_sped_up(
+            a,
+            self.text,
+            start,
+            end,
+            self.gap_score,
+            self.match_score,
+            self.mismatch_score,
+        )
+        for new_start in range(align_start - 1, start - 1, -1):
+            if self.text[new_start] == " ":
+                align_start = new_start + 1
+                break
+        for new_end in range(align_end, end):
+            if self.text[new_end] == " ":
+                align_end = new_end
+                break
+
+        return align_start, align_end, score, substitutions
+
+    # unused. This was the original implementation of sw_align(),
+    # before I sped it up with cython. The cython implementation is
+    # approximately 10 times faster.
+    def sw_align_old(self, a, start, end):
         b = self.text[start:end]
         n, m = len(a), len(b)
         # building scoring matrix
@@ -94,6 +126,7 @@ def find_best(self, look_for, start=0, end=-1):
         end = len(self.text) if end < 0 else end
         if end - start < 2 * len(look_for):
             return self.sw_align(look_for, start, end)
+        # What happens in this case?
         window_size = len(look_for)
         windows = {}
         for i, ngram in enumerate(ngrams(" " + look_for + " ", 3)):

galvasr2/align/smith_waterman.pyx

@@ -0,0 +1,74 @@
+from collections import Counter
+
+cimport cython
+import numpy as np
+cimport numpy as np
+import time
+
+cdef char_pair(a, b):
+    if a < b:
+        return '' + a + b
+    else:
+        return '' + b + a
+
+
+@cython.boundscheck(False) # turn off bounds-checking for entire function
+@cython.wraparound(False)  # turn off negative index wrapping for entire function
+def sw_align(a, text, b_start, b_end, gap_score, match_score, mismatch_score):
+    b = text[b_start:b_end]
+    cdef int n = len(a)
+    cdef int m = len(b)
+    start_time = time.time()
+    cdef np.ndarray[np.int32_t, ndim=2] f = np.empty((n + 1, m + 1), np.int32)
+    f[0,0] = 0
+    cdef int i = 0
+    cdef int j = 0
+    for i in range(1, n + 1):
+        f[i,0] = gap_score * i
+    for j in range(1, m + 1):
+        f[0,j] = gap_score * j
+    cdef int max_score = 0
+    cdef int start_i = 0
+    cdef int start_j = 0
+    cdef int match = 0
+    cdef int insert = 0
+    cdef int delete = 0
+    cdef int score = 0
+    for i in range(1, n + 1):
+        for j in range(1, m + 1):
+            match = f[i - 1,j - 1] + (match_score if a[i-1] == b[j-1] 
+                                       else mismatch_score)
+            insert = f[i,j - 1] + gap_score
+            delete = f[i - 1,j] + gap_score
+            score = max(0, match, insert, delete)
+            # print(f"{i} {j} {score}")
+            f[i,j] = score
+            if score > max_score:
+                max_score = score
+                start_i = i
+                start_j = j
+    end_time = time.time()
+    # print(f"sw_align_sped_up compute matrix={end_time - start_time}")
+
+    start_time = time.time()
+    substitutions = Counter()
+    i = start_i
+    j = start_j
+    while (j > 0 or i > 0) and f[i,j] != 0:
+        if i > 0 and j > 0 and f[i,j] == (f[i-1,j-1] + (match_score if a[i-1] == b[j-1] 
+                                                        else mismatch_score)):
+            substitutions[char_pair(a[i-1], b[j-1])] += 1
+            i -= 1
+            j -= 1
+        elif i > 0 and f[i,j] == (f[i-1,j] + gap_score):
+            i -= 1
+        elif j > 0 and f[i,j] == (f[i,j-1] + gap_score):
+            j -= 1
+        else:
+            raise Exception('Switch-Waterman failure')
+    cdef int align_start = max(b_start, b_start + j - 1)
+    cdef int align_end = min(b_end, b_start + start_j)
+    cdef double final_score = f[start_i, start_j] / (match_score * max(align_end - align_start, n))
+    end_time = time.time()
+    # print(f"sw_align_sped_up backtrack={end_time - start_time}")
+    return align_start, align_end, final_score, substitutions