From 7eb957e9939ecea1968da2e51a3217aa3c45ac74 Mon Sep 17 00:00:00 2001
From: Yifan Yang <yifanyeung@qq.com>
Date: Fri, 2 Jun 2023 12:42:54 +0800
Subject: [PATCH] Fix typo and reformatted zipformer

---
 egs/librispeech/ASR/zipformer/decoder.py      |  1 -
 egs/librispeech/ASR/zipformer/export.py       |  2 +-
 .../ASR/zipformer/generate_averaged_model.py  |  6 +----
 egs/librispeech/ASR/zipformer/model.py        |  2 +-
 egs/librispeech/ASR/zipformer/pretrained.py   |  3 ++-
 egs/librispeech/ASR/zipformer/profile.py      | 14 +++++------
 egs/librispeech/ASR/zipformer/scaling.py      | 11 ++++----
 egs/librispeech/ASR/zipformer/subsampling.py  |  4 +--
 egs/librispeech/ASR/zipformer/zipformer.py    | 25 +++++++++++--------
 9 files changed, 35 insertions(+), 33 deletions(-)

diff --git a/egs/librispeech/ASR/zipformer/decoder.py b/egs/librispeech/ASR/zipformer/decoder.py
index 78d22a878..b54aa3294 100644
--- a/egs/librispeech/ASR/zipformer/decoder.py
+++ b/egs/librispeech/ASR/zipformer/decoder.py
@@ -17,7 +17,6 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-
 from scaling import Balancer
 
 
diff --git a/egs/librispeech/ASR/zipformer/export.py b/egs/librispeech/ASR/zipformer/export.py
index d1152372a..22f0756e2 100755
--- a/egs/librispeech/ASR/zipformer/export.py
+++ b/egs/librispeech/ASR/zipformer/export.py
@@ -160,6 +160,7 @@ from typing import List, Tuple
 
 import sentencepiece as spm
 import torch
+from scaling_converter import convert_scaled_to_non_scaled
 from torch import Tensor, nn
 from train import add_model_arguments, get_params, get_transducer_model
 
@@ -170,7 +171,6 @@ from icefall.checkpoint import (
     load_checkpoint,
 )
 from icefall.utils import make_pad_mask, str2bool
-from scaling_converter import convert_scaled_to_non_scaled
 
 
 def get_parser():
diff --git a/egs/librispeech/ASR/zipformer/generate_averaged_model.py b/egs/librispeech/ASR/zipformer/generate_averaged_model.py
index fe29355f2..5df4bb436 100755
--- a/egs/librispeech/ASR/zipformer/generate_averaged_model.py
+++ b/egs/librispeech/ASR/zipformer/generate_averaged_model.py
@@ -43,13 +43,9 @@ from pathlib import Path
 import sentencepiece as spm
 import torch
 from asr_datamodule import LibriSpeechAsrDataModule
-
 from train import add_model_arguments, get_params, get_transducer_model
 
-from icefall.checkpoint import (
-    average_checkpoints_with_averaged_model,
-    find_checkpoints,
-)
+from icefall.checkpoint import average_checkpoints_with_averaged_model, find_checkpoints
 
 
 def get_parser():
diff --git a/egs/librispeech/ASR/zipformer/model.py b/egs/librispeech/ASR/zipformer/model.py
index 7fcab04ae..5db9e2f82 100644
--- a/egs/librispeech/ASR/zipformer/model.py
+++ b/egs/librispeech/ASR/zipformer/model.py
@@ -19,9 +19,9 @@ import k2
 import torch
 import torch.nn as nn
 from encoder_interface import EncoderInterface
+from scaling import ScaledLinear
 
 from icefall.utils import add_sos, make_pad_mask
-from scaling import ScaledLinear
 
 
 class Transducer(nn.Module):
diff --git a/egs/librispeech/ASR/zipformer/pretrained.py b/egs/librispeech/ASR/zipformer/pretrained.py
index 1efc2d8b7..33e6944ab 100755
--- a/egs/librispeech/ASR/zipformer/pretrained.py
+++ b/egs/librispeech/ASR/zipformer/pretrained.py
@@ -120,10 +120,11 @@ from beam_search import (
     greedy_search_batch,
     modified_beam_search,
 )
-from icefall.utils import make_pad_mask
 from torch.nn.utils.rnn import pad_sequence
 from train import add_model_arguments, get_params, get_transducer_model
 
+from icefall.utils import make_pad_mask
+
 
 def get_parser():
     parser = argparse.ArgumentParser(
diff --git a/egs/librispeech/ASR/zipformer/profile.py b/egs/librispeech/ASR/zipformer/profile.py
index b460b5338..b56ce08be 100755
--- a/egs/librispeech/ASR/zipformer/profile.py
+++ b/egs/librispeech/ASR/zipformer/profile.py
@@ -22,24 +22,24 @@ Usage: ./zipformer/profile.py
 
 import argparse
 import logging
+from typing import Tuple
+
 import sentencepiece as spm
 import torch
-
-from typing import Tuple
-from torch import Tensor, nn
-
-from icefall.utils import make_pad_mask
-from icefall.profiler import get_model_profile
 from scaling import BiasNorm
+from torch import Tensor, nn
 from train import (
+    add_model_arguments,
     get_encoder_embed,
     get_encoder_model,
     get_joiner_model,
-    add_model_arguments,
     get_params,
 )
 from zipformer import BypassModule
 
+from icefall.profiler import get_model_profile
+from icefall.utils import make_pad_mask
+
 
 def get_parser():
     parser = argparse.ArgumentParser(
diff --git a/egs/librispeech/ASR/zipformer/scaling.py b/egs/librispeech/ASR/zipformer/scaling.py
index 23b23bfd8..e0ed144c8 100644
--- a/egs/librispeech/ASR/zipformer/scaling.py
+++ b/egs/librispeech/ASR/zipformer/scaling.py
@@ -15,15 +15,16 @@
 # limitations under the License.
 
 
-from typing import Optional, Tuple, Union
 import logging
-import k2
-from torch.cuda.amp import custom_fwd, custom_bwd
-import random
-import torch
 import math
+import random
+from typing import Optional, Tuple, Union
+
+import k2
+import torch
 import torch.nn as nn
 from torch import Tensor
+from torch.cuda.amp import custom_bwd, custom_fwd
 
 
 class PiecewiseLinear(object):
diff --git a/egs/librispeech/ASR/zipformer/subsampling.py b/egs/librispeech/ASR/zipformer/subsampling.py
index 47403f13c..c0c1c1963 100644
--- a/egs/librispeech/ASR/zipformer/subsampling.py
+++ b/egs/librispeech/ASR/zipformer/subsampling.py
@@ -16,11 +16,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Tuple
 import warnings
+from typing import Tuple
 
 import torch
-from torch import Tensor, nn
 from scaling import (
     Balancer,
     BiasNorm,
@@ -34,6 +33,7 @@ from scaling import (
     SwooshR,
     Whiten,
 )
+from torch import Tensor, nn
 
 
 class ConvNeXt(nn.Module):
diff --git a/egs/librispeech/ASR/zipformer/zipformer.py b/egs/librispeech/ASR/zipformer/zipformer.py
index 7aae6ae22..3bebe6868 100644
--- a/egs/librispeech/ASR/zipformer/zipformer.py
+++ b/egs/librispeech/ASR/zipformer/zipformer.py
@@ -17,28 +17,33 @@
 # limitations under the License.
 
 import copy
+import logging
 import math
+import random
 import warnings
 from typing import List, Optional, Tuple, Union
-import logging
+
 import torch
-import random
 from encoder_interface import EncoderInterface
 from scaling import (
+    Identity,  # more friendly to backward hooks than nn.Identity(), for diagnostic reasons.
+)
+from scaling import (
+    ScaledLinear,  # not as in other dirs.. just scales down initial parameter values.
+)
+from scaling import (
+    ActivationDropoutAndLinear,
     Balancer,
     BiasNorm,
-    Dropout2,
     ChunkCausalDepthwiseConv1d,
-    ActivationDropoutAndLinear,
-    ScaledLinear,  # not as in other dirs.. just scales down initial parameter values.
+    Dropout2,
+    FloatLike,
+    ScheduledFloat,
     Whiten,
-    Identity,  # more friendly to backward hooks than nn.Identity(), for diagnostic reasons.
+    convert_num_channels,
+    limit_param_value,
     penalize_abs_values_gt,
     softmax,
-    ScheduledFloat,
-    FloatLike,
-    limit_param_value,
-    convert_num_channels,
 )
 from torch import Tensor, nn