[pybind] add pybind wrapper for int vector reader/writer. (#3833)

Author: csukuangfj

src/pybind/Makefile

@@ -106,6 +106,17 @@ LD_PRELOAD += $(MKL_LD_PRELOAD)
 export LD_PRELOAD
 endif
 
+# directories should be sorted in alphabetic order
+TEST_DIRS := \
+chain \
+cudamatrix \
+dlpack \
+feat \
+fst \
+matrix \
+nnet3 \
+tests
+
 .PHONY: all clean test
 
 all: $(LIBFILE)
@@ -124,15 +135,7 @@ clean:
 	-rm -f .depend.mk
 
 test: all
-	python3 tests/test_dlpack_subvector.py
-	python3 tests/test_kaldi_pybind.py
-	make -C chain test
-	make -C cudamatrix test
-	make -C dlpack test
-	make -C feat test
-	make -C fst test
-	make -C matrix test
-	make -C nnet3 test
+	for d in $(TEST_DIRS); do make -C $$d test; done
 
 # valgrind-python.supp is from http://svn.python.org/projects/python/trunk/Misc/valgrind-python.supp
 # since we do not compile Python from source, we follow the comment in valgrind-python.supp
@@ -145,7 +148,7 @@ depend:
 	rm -f .depend.mk
 	for f in $(CCFILES); do \
            $(CXX) -M -MT "$$(dirname $$f)/$$(basename -s .cc $$f).o" \
-           $(CXXFLAGS) $$f >> .depend.mk; \
+             $(CXXFLAGS) $$f >> .depend.mk; \
       done
 
 -include .depend.mk

src/pybind/fst/vector_fst_pybind_test.py

@@ -86,6 +86,9 @@ def test_std_vector_fst(self):
         print('read back compiled fst is:')
         print(read_back_compiled_fst)
 
+        os.remove(compiled_filename)
+        os.remove(fst_filename)
+
 
 if __name__ == '__main__':
     unittest.main()

src/pybind/kaldi/__init__.py

@@ -5,28 +5,6 @@
 
 from kaldi_pybind import *
 
+from pytorch_util import *
 from symbol_table import *
-from pytorch_util import PytorchToCuSubMatrix
-from pytorch_util import PytorchToCuSubVector
-from pytorch_util import PytorchToSubMatrix
-from pytorch_util import PytorchToSubVector
-
-from table import SequentialNnetChainExampleReader
-from table import RandomAccessNnetChainExampleReader
-from table import NnetChainExampleWriter
-
-from table import SequentialWaveReader
-from table import RandomAccessWaveReader
-
-from table import SequentialWaveInfoReader
-from table import RandomAccessWaveInfoReader
-
-from table import SequentialMatrixReader
-from table import RandomAccessMatrixReader
-from table import MatrixWriter
-
-from table import SequentialVectorReader
-from table import RandomAccessVectorReader
-from table import VectorWriter
-
-from table import CompressedMatrixWriter
+from table import *

src/pybind/kaldi/table.py

@@ -26,6 +26,9 @@
 import sys
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), os.pardir))
 
+import numpy as np
+
+import kaldi_pybind
 from kaldi_pybind.nnet3 import _SequentialNnetChainExampleReader
 from kaldi_pybind.nnet3 import _RandomAccessNnetChainExampleReader
 from kaldi_pybind.nnet3 import _NnetChainExampleWriter
@@ -45,6 +48,10 @@
 
 from kaldi_pybind import _CompressedMatrixWriter
 
+from kaldi_pybind import _SequentialInt32VectorReader
+from kaldi_pybind import _RandomAccessInt32VectorReader
+from kaldi_pybind import _Int32VectorWriter
+
 ################################################################################
 # Sequential Readers
 ################################################################################
@@ -199,6 +206,12 @@ class SequentialVectorReader(_SequentialReaderBase,
     pass
 
 
+class SequentialIntVectorReader(_SequentialReaderBase,
+                                _SequentialInt32VectorReader):
+    '''Sequential table reader for integer sequences.'''
+    pass
+
+
 ################################################################################
 # Random Access Readers
 ################################################################################
@@ -334,6 +347,12 @@ class RandomAccessVectorReader(_RandomAccessReaderBase,
     pass
 
 
+class RandomAccessIntVectorReader(_RandomAccessReaderBase,
+                                  _RandomAccessInt32VectorReader):
+    '''Random access table reader for integer sequences.'''
+    pass
+
+
 ################################################################################
 # Writers
 ################################################################################
@@ -431,19 +450,34 @@ class NnetChainExampleWriter(_WriterBase, _NnetChainExampleWriter):
 
 class MatrixWriter(_WriterBase, _BaseFloatMatrixWriter):
     '''Table writer for single precision matrices.'''
-    pass
+
+    def Write(self, key, value):
+        if isinstance(value, np.ndarray):
+            m = kaldi_pybind.FloatSubMatrix(value)
+            value = kaldi_pybind.FloatMatrix(m)
+        super().Write(key, value)
 
 
 class VectorWriter(_WriterBase, _BaseFloatVectorWriter):
     '''Table writer for single precision vectors.'''
-    pass
+
+    def Write(self, key, value):
+        if isinstance(value, np.ndarray):
+            v = kaldi_pybind.FloatSubVector(value)
+            value = kaldi_pybind.FloatVector(v)
+        super().Write(key, value)
 
 
 class CompressedMatrixWriter(_WriterBase, _CompressedMatrixWriter):
     '''Table writer for single precision compressed matrices.'''
     pass
 
 
+class IntVectorWriter(_WriterBase, _Int32VectorWriter):
+    '''Table writer for integer sequences.'''
+    pass
+
+
 if False:
     # TODO(fangjun): enable the following once other wrappers are added
 
@@ -526,11 +560,6 @@ class SequentialBoolReader(_SequentialReaderBase,
         '''Sequential table reader for Booleans.'''
         pass
 
-    class SequentialIntVectorReader(_SequentialReaderBase,
-                                    _kaldi_table.SequentialIntVectorReader):
-        '''Sequential table reader for integer sequences.'''
-        pass
-
     class SequentialIntVectorVectorReader(
             _SequentialReaderBase,
             _kaldi_table.SequentialIntVectorVectorReader):
@@ -623,11 +652,6 @@ class RandomAccessBoolReader(_RandomAccessReaderBase,
         '''Random access table reader for Booleans.'''
         pass
 
-    class RandomAccessIntVectorReader(_RandomAccessReaderBase,
-                                      _kaldi_table.RandomAccessIntVectorReader):
-        '''Random access table reader for integer sequences.'''
-        pass
-
     class RandomAccessIntVectorVectorReader(
             _RandomAccessReaderBase,
             _kaldi_table.RandomAccessIntVectorVectorReader):
@@ -886,10 +910,6 @@ class BoolWriter(_WriterBase, _kaldi_table.BoolWriter):
         '''Table writer for Booleans.'''
         pass
 
-    class IntVectorWriter(_WriterBase, _kaldi_table.IntVectorWriter):
-        '''Table writer for integer sequences.'''
-        pass
-
     class IntVectorVectorWriter(_WriterBase,
                                 _kaldi_table.IntVectorVectorWriter):
         '''Table writer for sequences of integer sequences.'''

src/pybind/matrix/kaldi_vector_pybind.cc

@@ -65,6 +65,8 @@ void pybind_kaldi_vector(py::module& m) {
       })
       .def(py::init<const MatrixIndexT, MatrixResizeType>(), py::arg("size"),
            py::arg("resize_type") = kSetZero)
+      .def(py::init<const VectorBase<float>&>(), py::arg("v"),
+           "Copy-constructor from base-class, needed to copy from SubVector.")
       .def("to_dlpack", [](py::object obj) { return VectorToDLPack(&obj); });
 
   py::class_<SubVector<float>, VectorBase<float>>(m, "FloatSubVector")

src/pybind/tests/Makefile

@@ -0,0 +1,5 @@
+
+test:
+	python3 ./test_dlpack_subvector.py
+	python3 ./test_kaldi_pybind.py
+	python3 ./test_table_types.py

src/pybind/tests/test_kaldi_pybind.py

@@ -25,30 +25,26 @@ def test_float_vector(self):
         self.assertTrue((kp_vector == gold).all())
 
     def test_float_matrix(self):
-        return
         # test FloatMatrix
         kp_matrix = kaldi.FloatMatrix(4, 5)
 
-        np_matrix = kp_matrix.numpy()
-
-        np_matrix[2][3] = 2.0
+        kp_matrix[2, 3] = 2.0
 
         gold = np.array([
             [0, 0, 0, 0, 0],
             [0, 0, 0, 0, 0],
             [0, 0, 0, 2, 0],
             [0, 0, 0, 0, 0],
         ])
-        self.assertTrue((kp_matrix == gold).all())
+        np.testing.assert_array_equal(kp_matrix.numpy(), gold)
 
     def test_matrix_reader_writer(self):
         kp_matrix = kaldi.FloatMatrix(2, 3)
         wspecifier = 'ark,t:test.ark'
         rspecifier = 'ark:test.ark'
         matrix_writer = kaldi.MatrixWriter(wspecifier)
 
-        np_matrix = kp_matrix.numpy()
-        np_matrix[0, 0] = 10
+        kp_matrix[0, 0] = 10
 
         matrix_writer.Write('id_1', kp_matrix)
         matrix_writer.Close()
@@ -59,7 +55,10 @@ def test_matrix_reader_writer(self):
 
         value = matrix_reader.Value()
         gold = np.array([[10, 0, 0], [0, 0, 0]])
-        self.assertTrue((np.array(value, copy=False) == gold).all())
+        np.testing.assert_array_equal(value.numpy(), gold)
+
+        matrix_reader.Close()
+        os.remove('test.ark')
 
     def test_matrix_reader_iterator(self):
         kp_matrix = kaldi.FloatMatrix(2, 3)
@@ -71,11 +70,13 @@ def test_matrix_reader_iterator(self):
 
         gold_key_list = ['id_1']
         gold_value_list = [np.array([[0, 0, 0], [0, 0, 0]])]
-        for (key, value), gold_key, gold_value in zip(
-                kaldi.SequentialMatrixReader(rspecifier), gold_key_list,
-                gold_value_list):
+        reader = kaldi.SequentialMatrixReader(rspecifier)
+        for (key, value), gold_key, gold_value in zip(reader, gold_key_list,
+                                                      gold_value_list):
             self.assertEqual(key, gold_key)
-            self.assertTrue((value == gold_value).all())
+            np.testing.assert_array_equal(value.numpy(), gold_value)
+        reader.Close()
+        os.remove('test.ark')
 
     def test_matrix_random_access_reader(self):
         kp_matrix = kaldi.FloatMatrix(2, 3)
@@ -88,8 +89,11 @@ def test_matrix_random_access_reader(self):
         reader = kaldi.RandomAccessMatrixReader(rspecifier)
         gold = np.array([[0, 0, 0], [0, 0, 0]])
         self.assertTrue('id_1' in reader)
-        self.assertTrue((np.array(reader['id_1']) == gold).all())
+
+        np.testing.assert_array_equal(reader['id_1'].numpy(), gold)
         self.assertFalse('id_2' in reader)
+        reader.Close()
+        os.remove('test.ark')
 
 
 if __name__ == '__main__':