search confidence works

wenet-e2e · robin1001 · Oct 30, 2023 · Oct 16, 2023 · Oct 20, 2023 · Oct 23, 2023
commit cd9c65975116f530e5382dc09ad37fe4c2b098d9
diff --git a/wenet/cli/paraformer_model.py b/wenet/cli/paraformer_model.py
@@ -34,9 +34,25 @@ def transcribe(self, audio_file: str):
         decoder_out, token_num = self.model.forward_paraformer(
             feats, feats_lens)
 
-        results = paraformer_greedy_search(decoder_out, token_num)
-        hyp = [self.char_dict[x] for x in results[0].tokens]
-
-        # TODO(Mddct): deal with '@@' and 'eos'
-        result = ''.join(hyp)
+        res = paraformer_greedy_search(decoder_out, token_num)[0]
+
+        tokens_info = True
+        result = {}
+        result['confidence'] = res.confidence
+        # # TODO(Mddct): deal with '@@' and 'eos'
+        result['rec'] = "".join([self.char_dict[x] for x in res.tokens])
+
+        if tokens_info:
+            tokens_info = []
+            for i, x in enumerate(res.tokens):
+                tokens_info.append({
+                    'token': self.char_dict[x],
+                    # TODO(Mddct): support times
+                    # 'start': 0,
+                    # 'end': 0,
+                    'confidence': res.tokens_confidence[i]
+                })
+            result['tokens'] = tokens_info
+
+        # result = ''.join(hyp)
         return result
diff --git a/wenet/paraformer/ali_paraformer/model.py b/wenet/paraformer/ali_paraformer/model.py
@@ -312,7 +312,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.feed_forward(self.norm1(x))
 
 
-class SANMDecoderLayer(DecoderLayer):
+class SanmDecoderLayer(DecoderLayer):
 
     def __init__(self,
                  size: int,
@@ -403,7 +403,7 @@ def __init__(
         del self.embed
         del self.decoders
         self.decoders = torch.nn.ModuleList([
-            SANMDecoderLayer(
+            SanmDecoderLayer(
                 encoder_output_size,
                 DummyMultiHeadSANM(attention_heads, encoder_output_size,
                                    encoder_output_size, dropout_rate,
@@ -493,7 +493,7 @@ def forward_paraformer(
         # decoder
         decoder_out, _, _ = self.decoder(encoder_out, encoder_out_mask,
                                          acoustic_embed, token_num)
-        # decoder_out = decoder_out.log_softmax(dim=-1)
+        decoder_out = decoder_out.log_softmax(dim=-1)
         return decoder_out, token_num
 
     def decode(self,

diff --git a/wenet/paraformer/search.py b/wenet/paraformer/search.py
@@ -1,3 +1,4 @@
+import math
 from typing import List, Tuple
 import torch
 
@@ -13,12 +14,22 @@ def paraformer_greedy_search(
     maxlen = decoder_out.size(1)
     topk_prob, topk_index = decoder_out.topk(1, dim=2)
     topk_index = topk_index.view(batch_size, maxlen)  # (B, maxlen)
+    topk_prob = topk_prob.view(batch_size, maxlen)
     results = []
-    topk_index = topk_index.cpu()
-    decoder_out_lens = decoder_out_lens.cpu()
-    # TODO(Mddct): scores, times etc
-    for (i, hyp) in enumerate(topk_index.tolist()):
-        r = DecodeResult(hyp[:decoder_out_lens.numpy()[i]])
+    topk_index = topk_index.cpu().tolist()
+    topk_prob = topk_prob.cpu().tolist()
+    decoder_out_lens = decoder_out_lens.cpu().numpy()
+    # TODO(Mddct) times
+    for (i, hyp) in enumerate(topk_index):
+        confidence = 0.0
+        tokens_confidence = []
+        lens = decoder_out_lens[i]
+        for logp in topk_prob[i][:lens]:
+            tokens_confidence.append(math.exp(logp))
+            confidence += logp
+        r = DecodeResult(hyp[:lens],
+                         tokens_confidence=tokens_confidence,
+                         confidence=math.exp(confidence / lens))
         results.append(r)
     return results