add init for deepspeech workloads

mlcommons · Oct 2, 2023 · 0229fd8 · 0229fd8
1 parent 29c123e
commit 0229fd8
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/algorithmic_efficiency/workloads/librispeech_deepspeech/librispeech_jax/workload.py b/algorithmic_efficiency/workloads/librispeech_deepspeech/librispeech_jax/workload.py
@@ -17,6 +17,13 @@
 
 
 class LibriSpeechDeepSpeechWorkload(BaseDeepspeechLibrispeechWorkload):
+
+  def __init__(self,
+               tokenizer_vocab_path: Optional[str] = None,
+               use_specaug: bool = True) -> None:
+    super().__init__()
+    self.metrics_bundle = metrics.get_metrics_bundle(tokenizer_vocab_path)
+    self.use_specaug = use_specaug
 
   def init_model_fn(
       self,

diff --git a/algorithmic_efficiency/workloads/librispeech_deepspeech/librispeech_pytorch/workload.py b/algorithmic_efficiency/workloads/librispeech_deepspeech/librispeech_pytorch/workload.py
@@ -29,6 +29,13 @@
 
 class LibriSpeechDeepSpeechWorkload(BaseDeepspeechLibrispeechWorkload):
 
+  def __init__(self,
+               tokenizer_vocab_path: Optional[str] = None,
+               use_specaug: bool = True) -> None:
+    super().__init__()
+    self.tokenizer = metrics.load_tokenizer(tokenizer_vocab_path)
+    self.use_specaug = use_specaug
+
   def init_model_fn(
       self,
       rng: spec.RandomState,