🚀 update demo examples and dependencies

Author: nglehuy

examples/demonstration/conformer.py

@@ -15,46 +15,26 @@
 import os
 import argparse
 from tensorflow_asr.utils import setup_environment, setup_devices
+from tensorflow_asr.utils.utils import get_reduced_length
 
 setup_environment()
 import tensorflow as tf
 
 parser = argparse.ArgumentParser(prog="Conformer non streaming")
 
-parser.add_argument("filename", metavar="FILENAME",
-                    help="audio file to be played back")
+parser.add_argument("filename", metavar="FILENAME", help="audio file to be played back")
 
-parser.add_argument("--config", type=str, default=None,
-                    help="Path to conformer config yaml")
+parser.add_argument("--config", type=str, default=None, help="Path to conformer config yaml")
 
-parser.add_argument("--saved", type=str, default=None,
-                    help="Path to conformer saved h5 weights")
-
-parser.add_argument("--blank", type=int, default=0,
-                    help="Path to conformer tflite")
+parser.add_argument("--saved", type=str, default=None, help="Path to conformer saved h5 weights")
 
 parser.add_argument("--beam_width", type=int, default=0, help="Beam width")
 
-parser.add_argument("--num_rnns", type=int, default=1,
-                    help="Number of RNN layers in prediction network")
-
-parser.add_argument("--nstates", type=int, default=2,
-                    help="Number of RNN states in prediction network (1 for GRU and 2 for LSTM)")
-
-parser.add_argument("--statesize", type=int, default=320,
-                    help="Size of RNN state in prediction network")
-
-parser.add_argument("--device", type=int, default=0,
-                    help="Device's id to run test on")
-
-parser.add_argument("--cpu", default=False, action="store_true",
-                    help="Whether to only use cpu")
+parser.add_argument("--device", type=int, default=0, help="Device's id to run test on")
 
-parser.add_argument("--subwords", type=str, default=None,
-                    help="Path to file that stores generated subwords")
+parser.add_argument("--cpu", default=False, action="store_true", help="Whether to only use cpu")
 
-parser.add_argument("--output_name", type=str, default="test",
-                    help="Result filename name prefix")
+parser.add_argument("--subwords", type=str, default=None, help="Path to file that stores generated subwords")
 
 args = parser.parse_args()
 
@@ -83,10 +63,12 @@
 conformer.add_featurizers(speech_featurizer, text_featurizer)
 
 signal = read_raw_audio(args.filename)
+features = speech_featurizer.tf_extract(signal)
+input_length = get_reduced_length(tf.shape(features)[0], conformer.time_reduction_factor)
 
 if (args.beam_width):
-    transcript = conformer.recognize_beam(signal[None, ...])
+    transcript = conformer.recognize_beam(features[None, ...], input_length[None, ...])
 else:
-    transcript = conformer.recognize(signal[None, ...])
+    transcript = conformer.recognize(features[None, ...], input_length[None, ...])
 
 tf.print("Transcript:", transcript[0])

examples/demonstration/streaming_tflite_conformer.py

@@ -93,7 +93,7 @@ def recognizer(Q):
 
     def recognize(signal, lastid, states):
         if signal.shape[0] < args.blocksize:
-            signal = np.pad(signal, [[0, args.blocksize - signal.shape[0]]])
+            signal = tf.pad(signal, [[0, args.blocksize - signal.shape[0]]])
         tflitemodel.set_tensor(input_details[0]["index"], signal)
         tflitemodel.set_tensor(input_details[1]["index"], lastid)
         tflitemodel.set_tensor(input_details[2]["index"], states)
@@ -104,8 +104,8 @@ def recognize(signal, lastid, states):
         text = "".join([chr(u) for u in upoints])
         return text, lastid, states
 
-    lastid = args.blank * np.ones(shape=[], dtype=np.int32)
-    states = np.zeros(shape=[args.num_rnns, args.nstates, 1, args.statesize], dtype=np.float32)
+    lastid = args.blank * tf.ones(shape=[], dtype=tf.int32)
+    states = tf.zeros(shape=[args.num_rnns, args.nstates, 1, args.statesize], dtype=tf.float32)
     transcript = ""
 
     while True:
@@ -122,51 +122,56 @@ def recognize(signal, lastid, states):
 tflite_process.start()
 
 
-def callback(outdata, frames, time, status):
-    assert frames == args.blocksize
-    if status.output_underflow:
-        print('Output underflow: increase blocksize?', file=sys.stderr)
-        raise sd.CallbackAbort
-    assert not status
+def send(q, Q, E):
+    def callback(outdata, frames, time, status):
+        assert frames == args.blocksize
+        if status.output_underflow:
+            print('Output underflow: increase blocksize?', file=sys.stderr)
+            raise sd.CallbackAbort
+        assert not status
+        try:
+            data = q.get_nowait()
+            Q.put(np.frombuffer(data, dtype=np.float32))
+        except queue.Empty as e:
+            print('Buffer is empty: increase buffersize?', file=sys.stderr)
+            raise sd.CallbackAbort from e
+        if len(data) < len(outdata):
+            outdata[:len(data)] = data
+            outdata[len(data):] = b'\x00' * (len(outdata) - len(data))
+            raise sd.CallbackStop
+        else:
+            outdata[:] = data
+
     try:
-        data = q.get_nowait()
-        Q.put(np.frombuffer(data, dtype=np.float32))
-    except queue.Empty as e:
-        print('Buffer is empty: increase buffersize?', file=sys.stderr)
-        raise sd.CallbackAbort from e
-    if len(data) < len(outdata):
-        outdata[:len(data)] = data
-        outdata[len(data):] = b'\x00' * (len(outdata) - len(data))
-        raise sd.CallbackStop
-    else:
-        outdata[:] = data
-
-
-try:
-    with sf.SoundFile(args.filename) as f:
-        for _ in range(args.buffersize):
-            data = f.buffer_read(args.blocksize, dtype='float32')
-            if not data:
-                break
-            q.put_nowait(data)  # Pre-fill queue
-        stream = sd.RawOutputStream(
-            samplerate=f.samplerate, blocksize=args.blocksize,
-            device=args.device, channels=f.channels, dtype='float32',
-            callback=callback, finished_callback=E.set)
-        with stream:
-            timeout = args.blocksize * args.buffersize / f.samplerate
-            while data:
+        with sf.SoundFile(args.filename) as f:
+            for _ in range(args.buffersize):
                 data = f.buffer_read(args.blocksize, dtype='float32')
-                q.put(data, timeout=timeout)
-            E.wait()
-
-except KeyboardInterrupt:
-    parser.exit('\nInterrupted by user')
-except queue.Full:
-    # A timeout occurred, i.e. there was an error in the callback
-    parser.exit(1)
-except Exception as e:
-    parser.exit(type(e).__name__ + ': ' + str(e))
-
-tflite_process.join()
-tflite_process.close()
+                if not data:
+                    break
+                q.put_nowait(data)  # Pre-fill queue
+            stream = sd.RawOutputStream(
+                samplerate=f.samplerate, blocksize=args.blocksize,
+                device=args.device, channels=f.channels, dtype='float32',
+                callback=callback, finished_callback=E.set)
+            with stream:
+                timeout = args.blocksize * args.buffersize / f.samplerate
+                while data:
+                    data = f.buffer_read(args.blocksize, dtype='float32')
+                    q.put(data, timeout=timeout)
+                E.wait()
+
+    except KeyboardInterrupt:
+        parser.exit('\nInterrupted by user')
+    except queue.Full:
+        # A timeout occurred, i.e. there was an error in the callback
+        parser.exit(1)
+    except Exception as e:
+        parser.exit(type(e).__name__ + ': ' + str(e))
+
+
+send_process = Process(target=send, args=[q, Q, E])
+send_process.start()
+send_process.join()
+send_process.close()
+
+tflite_process.terminate()

setup.py

@@ -19,18 +19,16 @@
 
 requirements = [
     "tensorflow-datasets>=3.2.1,<4.0.0",
-    "tensorflow-metadata>=0.26.0",
     "tensorflow-addons>=0.10.0",
     "setuptools>=47.1.1",
-    "librosa>=0.7.2",
+    "librosa>=0.8.0",
     "soundfile>=0.10.3",
     "PyYAML>=5.3.1",
     "matplotlib>=3.2.1",
-    "sox>=1.3.7",
-    "numba==0.49.1",
-    "tqdm>=4.51.0",
-    "colorama>=0.4.3",
-    "nlpaug>=1.0.1",
+    "sox>=1.4.1",
+    "tqdm>=4.54.1",
+    "colorama>=0.4.4",
+    "nlpaug>=1.1.1",
 ]
 
 setuptools.setup(

tensorflow_asr/featurizers/speech_featurizers.py

@@ -428,14 +428,12 @@ def tf_extract(self, signal: tf.Tensor) -> tf.Tensor:
         elif self.feature_type == "log_gammatone_spectrogram":
             features = self.compute_log_gammatone_spectrogram(signal)
         else:
-            raise ValueError("feature_type must be either 'mfcc',"
-                             "'log_mel_spectrogram' or 'spectrogram'")
+            raise ValueError("feature_type must be either 'mfcc', 'log_mel_spectrogram' or 'spectrogram'")
 
         features = tf.expand_dims(features, axis=-1)
 
         if self.normalize_feature:
-            features = tf_normalize_audio_features(
-                features, per_feature=self.normalize_per_feature)
+            features = tf_normalize_audio_features(features, per_feature=self.normalize_per_feature)
 
         return features
 

tensorflow_asr/models/ctc.py

@@ -46,11 +46,11 @@ def recognize(self, features: tf.Tensor, input_length: Optional[tf.Tensor]):
         logits = self(features, training=False)
         probs = tf.nn.softmax(logits)
 
-        def map_fn(prob): return tf.numpy_function(self.__perform_greedy, inp=[prob], Tout=tf.string)
+        def map_fn(prob): return tf.numpy_function(self._perform_greedy, inp=[prob], Tout=tf.string)
 
         return tf.map_fn(map_fn, probs, fn_output_signature=tf.TensorSpec([], dtype=tf.string))
 
-    def __perform_greedy(self, probs: np.ndarray):
+    def _perform_greedy(self, probs: np.ndarray):
         from ctc_decoders import ctc_greedy_decoder
         decoded = ctc_greedy_decoder(probs, vocabulary=self.text_featurizer.vocab_array)
         return tf.convert_to_tensor(decoded, dtype=tf.string)
@@ -85,11 +85,11 @@ def recognize_beam(self, features: tf.Tensor, input_length: Optional[tf.Tensor],
         logits = self(features, training=False)
         probs = tf.nn.softmax(logits)
 
-        def map_fn(prob): return tf.numpy_function(self.__perform_beam_search, inp=[prob, lm], Tout=tf.string)
+        def map_fn(prob): return tf.numpy_function(self._perform_beam_search, inp=[prob, lm], Tout=tf.string)
 
         return tf.map_fn(map_fn, probs, dtype=tf.string)
 
-    def __perform_beam_search(self, probs: np.ndarray, lm: bool = False):
+    def _perform_beam_search(self, probs: np.ndarray, lm: bool = False):
         from ctc_decoders import ctc_beam_search_decoder
         decoded = ctc_beam_search_decoder(
             probs_seq=probs,

tensorflow_asr/models/streaming_transducer.py

@@ -266,8 +266,8 @@ def recognize(self,
             self.encoder.get_initial_state(),
             input_length=input_length, with_batch=True
         )
-        return self.__perform_greedy_batch(encoded, input_length,
-                                           parallel_iterations=parallel_iterations, swap_memory=swap_memory)
+        return self._perform_greedy_batch(encoded, input_length,
+                                          parallel_iterations=parallel_iterations, swap_memory=swap_memory)
 
     def recognize_tflite(self, signal, predicted, encoder_states, prediction_states):
         """
@@ -286,7 +286,7 @@ def recognize_tflite(self, signal, predicted, encoder_states, prediction_states)
         """
         features = self.speech_featurizer.tf_extract(signal)
         encoded, new_encoder_states = self.encoder_inference(features, encoder_states)
-        hypothesis = self.__perform_greedy(encoded, tf.shape(encoded)[0], predicted, prediction_states)
+        hypothesis = self._perform_greedy(encoded, tf.shape(encoded)[0], predicted, prediction_states)
         transcript = self.text_featurizer.indices2upoints(hypothesis.prediction)
         return (
             transcript,
@@ -318,8 +318,8 @@ def recognize_beam(self,
             self.encoder.get_initial_state(),
             input_length=input_length, with_batch=True
         )
-        return self.__perform_beam_search_batch(encoded, input_length, lm,
-                                                parallel_iterations=parallel_iterations, swap_memory=swap_memory)
+        return self._perform_beam_search_batch(encoded, input_length, lm,
+                                               parallel_iterations=parallel_iterations, swap_memory=swap_memory)
 
     # -------------------------------- TFLITE -------------------------------------