Merge pull request #238 from pzelasko/feature/sampler-updates

Complimentary sampler update following Lhotse's changes

Author: danpovey

egs/aishell/asr/simple_v1/mmi_att_transformer_decode.py

@@ -6,17 +6,18 @@
 # Apache 2.0
 
 import argparse
-import k2
 import logging
-import numpy as np
 import os
-import torch
-from k2 import Fsa, SymbolTable
-from kaldialign import edit_distance
 from pathlib import Path
 from typing import List
 from typing import Union
 
+import k2
+import numpy as np
+import torch
+from k2 import Fsa, SymbolTable
+from kaldialign import edit_distance
+
 from snowfall.common import average_checkpoint
 from snowfall.common import find_first_disambig_symbol
 from snowfall.common import get_texts
@@ -25,16 +26,20 @@
 from snowfall.data import AishellAsrDataModule
 from snowfall.decoding.graph import compile_LG
 from snowfall.models import AcousticModel
-from snowfall.models.transformer import Transformer
 from snowfall.models.conformer import Conformer
+from snowfall.models.transformer import Transformer
 from snowfall.training.ctc_graph import build_ctc_topo
 from snowfall.training.mmi_graph import create_bigram_phone_lm
 from snowfall.training.mmi_graph import get_phone_symbols
 
 
 def decode(dataloader: torch.utils.data.DataLoader, model: AcousticModel,
            device: Union[str, torch.device], LG: Fsa, symbols: SymbolTable):
-    tot_num_cuts = len(dataloader.dataset.cuts)
+    num_batches = None
+    try:
+        num_batches = len(dataloader)
+    except TypeError:
+        pass
     num_cuts = 0
     results = []  # a list of pair (ref_words, hyp_words)
     for batch_idx, batch in enumerate(dataloader):
@@ -83,10 +88,8 @@ def decode(dataloader: torch.utils.data.DataLoader, model: AcousticModel,
             results.append((ref_words, hyp_words))
 
         if batch_idx % 10 == 0:
-            logging.info(
-                'batch {}, cuts processed until now is {}/{} ({:.6f}%)'.format(
-                    batch_idx, num_cuts, tot_num_cuts,
-                    float(num_cuts) / tot_num_cuts * 100))
+            batch_str = f"{batch_idx}" if num_batches is None else f"{batch_idx}/{num_batches}"
+            logging.info(f"batch {batch_str}, number of cuts processed until now is {num_cuts}")
 
         num_cuts += len(texts)