Update hyperparams.yaml

hyperparams.yaml

@@ -1,7 +1,7 @@
 # Hparams NEEDED
 HPARAMS_NEEDED: ["wav2vec_output_dim", "emb_size", "dec_neurons", "dec_layers", "output_neurons", "log_softmax", "tokenizer"]
 # Modules Needed
-MODULES_NEEDED: ["encoder_w2v2", "embedding", "ctc_lin", "seq_lin"]
+MODULES_NEEDED: ["encoder_w2v2", "embedding", "ctc_lin", "seq_lin", "lm_model"]
 
 # Pretrain folder (HuggingFace)
 output_folder: !ref output_folder_seq2seq_cv_podcast_arhiv_augmentation
@@ -25,7 +25,6 @@ bos_index: 0
 eos_index: 0
 unk_index: 0
 
-# Decoding parameters
 # Decoding parameters
 min_decode_ratio: 0.0
 max_decode_ratio: 1.0
@@ -33,13 +32,28 @@ valid_beam_size: 10
 test_beam_size: 20
 using_eos_threshold: True
 eos_threshold: 1.5
-using_max_attn_shift: True
+using_max_attn_shift: False
 max_attn_shift: 700
 length_normalization: True
 temperature: 1.0
+temperature_lm: 1.4
 # Scoring parameters
 coverage_penalty: 1.5
-
+lm_weight: 0.4
+
+
+# This is the RNNLM that is used according to the Huggingface repository
+# NB: It has to match the pre-trained RNNLM!!
+lm_model: !new:speechbrain.lobes.models.RNNLM.RNNLM
+  output_neurons: !ref <output_neurons>
+  embedding_dim: !ref <emb_size>
+  activation: !name:torch.nn.LeakyReLU
+  dropout: 0.0
+  rnn_layers: 2
+  rnn_neurons: 2048
+  dnn_blocks: 1
+  dnn_neurons: 512
+  return_hidden: True  # For inference
 
 # Wav2vec2 encoder
 encoder_w2v2: !new:speechbrain.lobes.models.huggingface_transformers.wav2vec2.Wav2Vec2
@@ -89,6 +103,7 @@ modules:
    decoder: !ref <decoder>
    ctc_lin: !ref <ctc_lin>
    seq_lin: !ref <seq_lin>
+   lm_model: !ref <lm_model>
 
 model: !new:torch.nn.ModuleList
    - [!ref <encoder_w2v2>, !ref <embedding>, !ref <decoder>, !ref <ctc_lin>, !ref <seq_lin>]
@@ -97,11 +112,22 @@ model: !new:torch.nn.ModuleList
 coverage_scorer: !new:speechbrain.decoders.scorer.CoverageScorer
   vocab_size: !ref <output_neurons>
 
+rnnlm_scorer: !new:speechbrain.decoders.scorer.RNNLMScorer
+  language_model: !ref <lm_model>
+  temperature: !ref <temperature_lm>
+
 scorer: !new:speechbrain.decoders.scorer.ScorerBuilder
   full_scorers: [!ref <coverage_scorer>]
   weights:
      coverage: !ref <coverage_penalty>
 
+scorer_lm: !new:speechbrain.decoders.scorer.ScorerBuilder
+  full_scorers: [!ref <rnnlm_scorer>,
+                 !ref <coverage_scorer>]
+  weights:
+     rnnlm: !ref <lm_weight>
+     coverage: !ref <coverage_penalty>
+
 test_search: !new:speechbrain.decoders.S2SRNNBeamSearcher
     embedding: !ref <embedding>
     decoder: !ref <decoder>
@@ -115,7 +141,7 @@ test_search: !new:speechbrain.decoders.S2SRNNBeamSearcher
     using_max_attn_shift: !ref <using_max_attn_shift>
     max_attn_shift: !ref <max_attn_shift>
     temperature: !ref <temperature>
-    scorer: !ref <scorer>
+    scorer: !ref <scorer_lm>
 
 
 ############################## Logging and Pretrainer ##########################
@@ -123,5 +149,7 @@ test_search: !new:speechbrain.decoders.S2SRNNBeamSearcher
 pretrainer: !new:speechbrain.utils.parameter_transfer.Pretrainer
     loadables:
         model: !ref <model>
+        lm: !ref <lm_model>
     paths:
         model: !ref <pretrained_path>/model.ckpt
+        lm: lm.ckpt