feat: register keras custom objects to support load models

Author: nglehuy

.pre-commit-config.yaml

@@ -1,17 +1,25 @@
 repos:
   - repo: local
     hooks:
-      - id: black-formatter-check
-        name: black-formatter-check
-        entry: black --check
+      - id: black-formatter-fix
+        name: black-formatter-fix
+        entry: bash -c "for f in $@; do black --verbose $f; done"
+        language: system
+        types: [python]
+        stages: [pre-commit]
+        fail_fast: true
+        verbose: true
+      - id: isort-fix
+        name: isort-fix
+        entry: bash -c "for f in $@; do echo -e \"Organize import for file $f\" && isort $f; done"
         language: system
         types: [python]
         stages: [pre-commit]
         fail_fast: true
         verbose: true
       - id: pylint-check
         name: pylint-check
-        entry: pylint --rcfile=.pylintrc -rn -sn
+        entry: bash -c "for f in $@; do pylint --rcfile=.pylintrc -rn -sn $f; done"
         language: system
         types: [python]
         stages: [pre-commit]

.pylintrc

@@ -119,6 +119,7 @@ disable=too-few-public-methods,
         redefined-outer-name,
         consider-using-f-string,
         fixme,
+        unused-variable,
 
 # Enable the message, report, category or checker with the given id(s). You can
 # either give multiple identifier separated by comma (,) or put this option

.vscode/settings.json

@@ -1,37 +1,37 @@
 {
-    "editor.formatOnSave": true,
-    "python.analysis.fixAll": [
-        "source.unusedImports",
+    "[python]": {
+        "editor.defaultFormatter": "ms-python.black-formatter"
+    },
+    "autoDocstring.docstringFormat": "numpy",
+    "black-formatter.args": [
+        "--config",
+        "${workspaceFolder}/pyproject.toml"
+    ],
+    "black-formatter.path": [
+        "${interpreter}",
+        "-m",
+        "black"
     ],
     "editor.codeActionsOnSave": {
+        "source.fixAll": "explicit",
         "source.organizeImports": "explicit"
     },
+    "editor.formatOnSave": true,
     "isort.args": [
-        "--profile",
-        "black",
-        "--line-length",
-        "130"
-    ],
-    "autoDocstring.docstringFormat": "numpy",
-    "[python]": {
-        "editor.defaultFormatter": "ms-python.black-formatter",
-    },
-    "pylint.path": [
-        "${interpreter}",
-        "-m",
-        "pylint"
+        "--settings-file",
+        "${workspaceFolder}/pyproject.toml"
     ],
     "pylint.args": [
         "--rcfile=${workspaceFolder}/.pylintrc"
     ],
-    "black-formatter.path": [
+    "pylint.path": [
         "${interpreter}",
         "-m",
-        "black"
+        "pylint"
     ],
-    "isort.path": [
-        "${interpreter}",
-        "-m",
-        "isort"
-    ]
+    "python.analysis.fixAll": [
+        "source.unusedImports",
+        "source.convertImportFormat"
+    ],
+    "python.analysis.importFormat": "absolute"
 }

examples/configs/librispeech/data.yml.j2

@@ -6,6 +6,7 @@ data_config:
       - {{datadir}}/train-clean-100/transcripts.tsv
       - {{datadir}}/train-clean-360/transcripts.tsv
       - {{datadir}}/train-other-500/transcripts.tsv
+      - {{datadir}}/dev-clean/transcripts.tsv
     tfrecords_dir: {{datadir}}/tfrecords
     shuffle: True
     cache: False

examples/inferences/main.py

@@ -14,14 +14,12 @@
 
 import os
 
-import tensorflow as tf
-import keras
-
-from tensorflow_asr import schemas, tokenizers
-from tensorflow_asr.models import base_model
+from tensorflow_asr import keras, schemas, tf, tokenizers
 from tensorflow_asr.configs import Config
+from tensorflow_asr.models import base_model
 from tensorflow_asr.utils import cli_util, data_util, env_util, file_util
 
+env_util.setup_logging()
 logger = tf.get_logger()
 
 

examples/inferences/tflite.py

@@ -12,12 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import tensorflow as tf
 import tensorflow_text as tft
 from tensorflow.lite.python import interpreter
 
-from tensorflow_asr.utils import cli_util, data_util
+from tensorflow_asr import tf
+from tensorflow_asr.utils import cli_util, data_util, env_util
 
+env_util.setup_logging()
 logger = tf.get_logger()
 
 

examples/models/ctc/conformer/char-small.yml.j2

@@ -2,7 +2,7 @@
 {{decoder_config}}
 
 model_config:
-  class_name: tensorflow_asr.models.ctc>Conformer
+  class_name: tensorflow_asr.models.ctc.conformer>Conformer
   config:
     name: conformer
     speech_config:

examples/models/ctc/conformer/tiny.yml.j2

@@ -2,7 +2,7 @@
 {{decoder_config}}
 
 model_config:
-  class_name: tensorflow_asr.models.ctc>Conformer
+  class_name: tensorflow_asr.models.ctc.conformer>Conformer
   config:
     name: conformer
     speech_config:

examples/models/ctc/deepspeech2/base.yml.j2

@@ -2,26 +2,26 @@
 {{decoder_config}}
 
 model_config:
-  class_name: tensorflow_asr.models.ctc>DeepSpeech2
+  class_name: tensorflow_asr.models.ctc.deepspeech2>DeepSpeech2
   config:
     name: deepspeech2
     speech_config:
       sample_rate: 16000
       frame_ms: 25
       stride_ms: 10
-      num_feature_bins: 128
-      feature_type: spectrogram
-      normalize_zscore: True
+      nfft: 512
+      num_feature_bins: 80
+      feature_type: log_mel_spectrogram
       augmentation_config:
         feature_augment:
           time_masking:
-            prob: 0.5
+            prob: 1.0
             num_masks: 5
             mask_factor: -1 # whole utterance
             p_upperbound: 0.05
             mask_value: 0
           freq_masking:
-            prob: 0.5
+            prob: 1.0
             num_masks: 1
             mask_factor: 27
             mask_value: 0
@@ -100,9 +100,9 @@ learning_config:
       config:
         log_dir: {{modeldir}}/tensorboard
         histogram_freq: 0
-        write_graph: False
+        write_graph: True
         write_images: False
-        write_steps_per_second: False
+        write_steps_per_second: True
         update_freq: batch
         profile_batch: 0
 

examples/models/ctc/jasper/base.yml.j2

@@ -2,15 +2,17 @@
 {{decoder_config}}
 
 model_config:
-  class_name: tensorflow_asr.models.ctc>Jasper
+  class_name: tensorflow_asr.models.ctc.jasper>Jasper
   config:
     name: jasper
     speech_config:
       sample_rate: 16000
       frame_ms: 25
       stride_ms: 10
+      nfft: 512
       num_feature_bins: 80
       feature_type: log_mel_spectrogram
+      log_base: "10"
     dense: True
     first_additional_block_channels: 256
     first_additional_block_kernels: 11