Update models, remove decoder_main

models-ax630c/{base-decoder-loop.axmodel → base/base-decoder.axmodel}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4b12160aaa1ca31248a32ce05713fd72e273b16444389853c1f52990cf5130eb
-size 130364397
+oid sha256:51e595565f0121eb4dc9ee14172cb8a111a56bf280a927660fdee5fbffa9d52e
+size 184323085

models-ax630c/{base-encoder.axmodel → base/base-encoder.axmodel}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b9f89ed5bbe31bcf98aa0e479ced1699b39816db2d3e2e2ff84c6e887af2b79b
-size 56024079
+oid sha256:00ac3d4d0aa81f3910d4aa9c777e81fbf3b4bc22f26a9d9ac38f236392261603
+size 56706622

models-ax630c/base/base_config.json

@@ -0,0 +1,30 @@
+{
+    "model_type": "whisper-base",
+    "version": "1",
+    "maintainer": "k2-fsa",
+    "n_mels": 80,
+    "n_audio_ctx": 1500,
+    "n_audio_state": 512,
+    "n_audio_head": 8,
+    "n_audio_layer": 6,
+    "n_vocab": 51865,
+    "n_text_ctx": 448,
+    "n_text_state": 512,
+    "n_text_head": 8,
+    "n_text_layer": 6,
+    "sot_sequence": "50258,50259,50359",
+    "all_language_tokens": "50346,50356,50292,50319,50325,50330,50327,50328,50341,50331,50307,50353,50279,50320,50322,50326,50340,50267,50269,50297,50301,50316,50323,50287,50276,50342,50304,50277,50311,50350,50286,50278,50290,50309,50284,50268,50300,50321,50272,50291,50281,50266,50357,50333,50293,50299,50294,50337,50271,50263,50296,50264,50343,50315,50314,50270,50352,50317,50349,50348,50283,50265,50308,50305,50336,50261,50335,50262,50345,50344,50351,50310,50329,50332,50289,50274,50302,50259,50324,50282,50285,50313,50280,50334,50260,50303,50312,50318,50295,50273,50338,50298,50347,50288,50354,50355,50275,50306,50339",
+    "all_language_codes": "my,jw,bg,gl,yo,be,af,oc,tk,tg,et,ln,he,mr,si,so,ps,pt,pl,cy,lv,kk,km,ta,hi,nn,az,fi,is,as,hu,vi,ur,br,ro,tr,fa,pa,ar,hr,el,ja,su,gu,lt,te,la,uz,nl,ru,ml,ko,mt,bs,mn,ca,haw,sq,mg,tl,cs,fr,mk,sl,lo,de,yi,es,lb,sa,tt,eu,ka,sd,th,it,bn,en,sn,ms,da,ne,uk,am,zh,sr,hy,sw,mi,sv,fo,sk,bo,no,ha,ba,id,kn,ht",
+    "sot": 50258,
+    "sot_index": 0,
+    "eot": 50257,
+    "blank_id": 220,
+    "is_multilingual": 1,
+    "no_speech": 50362,
+    "non_speech_tokens": "1,2,7,8,9,10,14,25,26,27,28,29,31,58,59,60,61,62,63,90,91,92,93,359,503,522,542,873,893,902,918,922,931,1350,1853,1982,2460,2627,3246,3253,3268,3536,3846,3961,4183,4667,6585,6647,7273,9061,9383,10428,10929,11938,12033,12331,12562,13793,14157,14635,15265,15618,16553,16604,18362,18956,20075,21675,22520,26130,26161,26435,28279,29464,31650,32302,32470,36865,42863,47425,49870,50254",
+    "transcribe": 50359,
+    "translate": 50358,
+    "sot_prev": 50361,
+    "sot_lm": 50360,
+    "no_timestamps": 50363
+}

models-ax630c/{base-decoder-main.axmodel → small/small-decoder.axmodel}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:940f273d111e3aee53cdb692a384a29556981aa146afbb2f558f6aac262c0621
-size 135675471
+oid sha256:9d13b4c3693f72e280162a1fd78c52ffc8ecc9318d7d2bd56db9810945c88f1b
+size 345275786

models-ax630c/small/small_config.json

@@ -0,0 +1,30 @@
+{
+    "model_type": "whisper-small",
+    "version": "1",
+    "maintainer": "k2-fsa",
+    "n_mels": 80,
+    "n_audio_ctx": 1500,
+    "n_audio_state": 768,
+    "n_audio_head": 12,
+    "n_audio_layer": 12,
+    "n_vocab": 51865,
+    "n_text_ctx": 448,
+    "n_text_state": 768,
+    "n_text_head": 12,
+    "n_text_layer": 12,
+    "sot_sequence": "50258,50259,50359",
+    "all_language_tokens": "50349,50276,50334,50346,50335,50313,50353,50284,50280,50268,50266,50267,50330,50350,50259,50311,50336,50340,50288,50302,50291,50279,50299,50351,50270,50301,50314,50295,50331,50285,50303,50326,50342,50355,50348,50341,50354,50321,50272,50269,50357,50333,50283,50309,50271,50324,50323,50290,50327,50298,50319,50356,50282,50332,50275,50263,50294,50305,50293,50317,50338,50287,50292,50316,50343,50289,50260,50328,50312,50344,50325,50304,50339,50320,50308,50274,50262,50345,50278,50296,50337,50310,50329,50318,50347,50265,50307,50264,50352,50315,50273,50277,50300,50261,50286,50306,50322,50281,50297",
+    "all_language_codes": "mg,hi,am,my,yi,ne,ln,ro,uk,tr,ja,pt,be,as,en,is,lo,ps,no,bn,hr,he,te,tt,ca,lv,mn,mi,tg,da,sr,so,nn,ba,tl,tk,ha,pa,ar,pl,su,gu,cs,br,nl,sn,km,ur,af,sk,gl,jw,ms,sd,id,ru,la,sl,lt,sq,fo,ta,bg,kk,mt,th,zh,oc,hy,sa,yo,az,ht,mr,mk,it,es,lb,vi,ml,uz,eu,ka,sw,bo,fr,et,ko,haw,bs,sv,fi,fa,de,hu,kn,si,el,cy",
+    "sot": 50258,
+    "sot_index": 0,
+    "eot": 50257,
+    "blank_id": 220,
+    "is_multilingual": 1,
+    "no_speech": 50362,
+    "non_speech_tokens": "1,2,7,8,9,10,14,25,26,27,28,29,31,58,59,60,61,62,63,90,91,92,93,359,503,522,542,873,893,902,918,922,931,1350,1853,1982,2460,2627,3246,3253,3268,3536,3846,3961,4183,4667,6585,6647,7273,9061,9383,10428,10929,11938,12033,12331,12562,13793,14157,14635,15265,15618,16553,16604,18362,18956,20075,21675,22520,26130,26161,26435,28279,29464,31650,32302,32470,36865,42863,47425,49870,50254",
+    "transcribe": 50359,
+    "translate": 50358,
+    "sot_prev": 50361,
+    "sot_lm": 50360,
+    "no_timestamps": 50363
+}

models-ax630c/{base-positional_embedding.bin → tiny/tiny-decoder.axmodel}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:88fa1cdbf2b06f86b0ecb7be0fccfc39e906502986572b8cf5319c250e857169
-size 917504
+oid sha256:64a611a2575597fed3e705d9faf941df0b33d58bc10a733fa31f5e937fd58ec4
+size 129647157

models-ax650/small/small-positional_embedding.bin → models-ax630c/tiny/tiny-encoder.axmodel

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c10bc44f2bd94bdf1b7aa03581309fa536132b3fe79bfe22c9a6934a42cd8b58
-size 1376256
+oid sha256:0528e2d7e317668e43a5641695f93d0c80d9902e28f0e9f2fdef76470855efe7
+size 26853722

models-ax630c/tiny/tiny_config.json

@@ -0,0 +1,30 @@
+{
+    "model_type": "whisper-tiny",
+    "version": "1",
+    "maintainer": "k2-fsa",
+    "n_mels": 80,
+    "n_audio_ctx": 1500,
+    "n_audio_state": 384,
+    "n_audio_head": 6,
+    "n_audio_layer": 4,
+    "n_vocab": 51865,
+    "n_text_ctx": 448,
+    "n_text_state": 384,
+    "n_text_head": 6,
+    "n_text_layer": 4,
+    "sot_sequence": "50258,50259,50359",
+    "all_language_tokens": "50286,50307,50345,50299,50290,50265,50275,50294,50309,50262,50318,50331,50282,50349,50270,50326,50279,50272,50355,50287,50269,50315,50278,50283,50266,50260,50297,50348,50346,50296,50277,50334,50342,50313,50288,50322,50325,50259,50302,50332,50338,50344,50261,50330,50304,50357,50314,50340,50291,50352,50320,50271,50316,50336,50323,50293,50263,50308,50284,50273,50267,50312,50321,50328,50285,50298,50301,50327,50354,50303,50356,50351,50295,50339,50292,50319,50264,50310,50276,50335,50311,50341,50350,50268,50289,50281,50324,50333,50317,50343,50305,50274,50306,50353,50300,50347,50329,50337,50280",
+    "all_language_codes": "hu,et,lb,te,ur,fr,id,la,br,es,sw,tg,ms,mg,ca,so,he,ar,ba,ta,pl,bs,vi,cs,ja,zh,cy,tl,my,ml,fi,am,nn,ne,no,si,yo,en,bn,sd,fo,sa,de,be,az,su,mn,ps,hr,haw,mr,nl,kk,lo,km,lt,ru,mk,ro,sv,pt,hy,pa,oc,da,sk,lv,af,ha,sr,jw,tt,mi,ht,bg,gl,ko,eu,hi,yi,is,tk,as,tr,th,el,sn,gu,sq,mt,sl,it,kn,ln,fa,bo,ka,uz,uk",
+    "sot": 50258,
+    "sot_index": 0,
+    "eot": 50257,
+    "blank_id": 220,
+    "is_multilingual": 1,
+    "no_speech": 50362,
+    "non_speech_tokens": "1,2,7,8,9,10,14,25,26,27,28,29,31,58,59,60,61,62,63,90,91,92,93,359,503,522,542,873,893,902,918,922,931,1350,1853,1982,2460,2627,3246,3253,3268,3536,3846,3961,4183,4667,6585,6647,7273,9061,9383,10428,10929,11938,12033,12331,12562,13793,14157,14635,15265,15618,16553,16604,18362,18956,20075,21675,22520,26130,26161,26435,28279,29464,31650,32302,32470,36865,42863,47425,49870,50254",
+    "transcribe": 50359,
+    "translate": 50358,
+    "sot_prev": 50361,
+    "sot_lm": 50360,
+    "no_timestamps": 50363
+}

models-ax650/base/base-decoder-loop.axmodel

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:43fbc9a1672eabd705bb68fdfd6b0837c4d3bceec5e07c80cc829cf47417e11d
-size 183531172

models-ax650/base/base-decoder-main.axmodel

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:7d9167ae1c52ed1cb318fb27e45734c2ceb0560444078693bf831ce02f2c0331
-size 183985586

models-ax650/base/base-decoder.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a5212e05f8c0d2d4b2a319eaf87eb93253a9bb476dee8fa3d8a85f3137b61045
+size 184444593

models-ax650/base/base-encoder.axmodel

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:76899fc4d232fd6f458d3597e8a67a04719715e971bc82e286679014a929f5b6
-size 33082024
+oid sha256:c07d03194566f292d26cf1f0e104d5740e1a249a0dc92b23c5b15b9e96496c24
+size 33132600

models-ax650/small/small-decoder-loop.axmodel

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:b472a0f3539d17fece09e92bf6cd69ebf391928a6050896bbf86b558a25def22
-size 269002567

models-ax650/small/small-decoder-main.axmodel

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:f3bfc577f60c35192d8ce8cc24f9ca4aa84af72756ba11af9d178d337cb7eb1c
-size 285531695

models-ax650/small/small-decoder.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d4ed1e62122ed495624efa0275a0d4cb2450996b6f1a5e1d9ea9d48026d1bb66
+size 350609498

models-ax650/small/small-encoder.axmodel

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7b3bc8db9762f9b2dfe78bffbc8070fb877b2572c5288253573e49a8c7b37948
-size 139705612
+oid sha256:6e4e8462e3b6ac3ea9465d61560bc42e1398a06ea640d9ff5cdb82636ab73d47
+size 136275980

models-ax650/tiny/tiny-decoder-loop.axmodel

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:139aa1429a6f439b7a5d1e5f481cb761c673afd4718a25968ad979fccfdfecaf
-size 128541899

models-ax650/tiny/tiny-decoder-main.axmodel

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8c1d08ec9309a26103a955cb432e0ddc25476da706a6a0a94108d225a48385aa
-size 128909975

models-ax650/tiny/tiny-decoder.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ecd77406ae1c04883beb7517f3e9f0ca0c4b35e640467997f126efb539b96306
+size 129267343

models-ax650/tiny/tiny-encoder.axmodel

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1e38085752ab08b9eec10b89cb1374e1a57ede680cc0340292e8d7261399acae
-size 14102295
+oid sha256:698670eb7410eeb2a84e26cee8918a780c97d1232a9dc7f946f61949105299a9
+size 14085412

models-ax650/tiny/tiny-positional_embedding.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:c13450ae630323a0bdd39b1226f92a7ac251131a909c7efdb7d2f5516736eb83
-size 688128

models-ax650/turbo/turbo-decoder-loop.axmodel

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:1b50daea5f4776006bf0e635ca5224f118649a9b3c37b2b821ae4c321db096ec
-size 499257709

models-ax650/turbo/turbo-decoder-main.axmodel

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:56d34930691311002e401d8243abbae632593b57406e8ae780cd02c9076a783d
-size 500341239

models-ax650/turbo/turbo-decoder.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3d5a52c2c1f2c08fdb3b664201a79b42e3439184ffcc5f19e151eeb4a88cadd8
+size 501696186

models-ax650/turbo/turbo-encoder.axmodel

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d8962c9189a43f6fc0ea080df1bdcc43ca916b9337d16d0c6a7b30894c14e3ee
-size 892467653
+oid sha256:b2cbdf3941b8d739318148505c380167a41e7d710d1a3cfda1a257c2c7fe428f
+size 893420089

models-ax650/turbo/turbo-positional_embedding.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:a94790ac6719da6e134835255274ee7fc6066ad5e0f08a0f747c1e1cf6407dc3
-size 2293760

models-onnx/base/base-decoder-loop.onnx

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:1616a829b7d3d643616633551204b8d0f008fb7a7dc38919eda2e8c6c6ed9714
-size 194571088

models-onnx/base/base-decoder-main.onnx

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:1096b83590016bdbe74c66c7ccad1c0120abd6d37214560b1dfe4cd886a0e683
-size 205485892

models-onnx/base/base-decoder.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a2ba37d199ec7c1794facba8efc84a484b0224f6650e362fda7d2db75827023a
+size 195497242

models-onnx/base/base-encoder.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:dd4b51bd569e9b2b2d83a8ed56f3618811f0c593aa95c010069df675027b5f2b
-size 95026988
+oid sha256:1c75e7cca22000432ec4f8f50726299aa20db34a9c154646e8ac10c0ddb4699b
+size 95025778

models-onnx/base/base-positional_embedding.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:88fa1cdbf2b06f86b0ecb7be0fccfc39e906502986572b8cf5319c250e857169
-size 917504

models-onnx/small/small-decoder.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f696ce159cffc13416ef624b91dae341de9d3f7e720cfc832000fe987f6d50b4
+size 557821592

models-onnx/small/small-encoder.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f7dce3226784147e42021007dede9c2c144d33487711a1950b5935b9c4829f1b
+size 409408370

models-onnx/small/small-positional_embedding.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:c10bc44f2bd94bdf1b7aa03581309fa536132b3fe79bfe22c9a6934a42cd8b58
-size 1376256

models-onnx/tiny/tiny-decoder-loop.onnx

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:5cbb3533939e2dfdf567b27762b12cf0956b7d7982bfb915228d24789f483058
-size 112843354

models-onnx/tiny/tiny-decoder-main.onnx

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:59ced1cf4e9a6f2aef0a2457f64f846e5682033abb4b894ba7680a60c792ad73
-size 118301861

models-onnx/tiny/tiny-decoder.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0d1a16a5a70e9a5940b68559eb5142c7c53b8f595f2f5e4a01d7e168acca6eb1
+size 113537271

models-onnx/tiny/tiny-encoder.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a8030a6d1f3615b8a5e000995fee88357768c7dbaad05a79f853a4040c97087b
-size 37606186
+oid sha256:324e93c5ddc2e922273ebeb16bba8c453aa009c698e2262103ac5ce21df1c3ed
+size 37605342

models-onnx/tiny/tiny-positional_embedding.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:c13450ae630323a0bdd39b1226f92a7ac251131a909c7efdb7d2f5516736eb83
-size 688128

models-onnx/turbo/turbo-decoder.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:75b7d28ee244aa1d39d02d98514cd951e7697f6fb742c7f84f808b8aab9b1d2a
+size 635240242

models-ax650/base/base-positional_embedding.bin → models-onnx/turbo/turbo-encoder.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:88fa1cdbf2b06f86b0ecb7be0fccfc39e906502986572b8cf5319c250e857169
-size 917504
+oid sha256:883b1a459251b53f7431bdb81d2a573d0965c42be73179cbdeb0041154fb0a7d
+size 389433

python/languages.py

@@ -99,4 +99,4 @@ WHISPER_LANGUAGES = {
     "jw": "javanese",
     "su": "sundanese",
     "yue": "cantonese",
-}
+}

python/main.py

@@ -6,29 +6,49 @@ import time
 
 def get_args():
     parser = argparse.ArgumentParser(
-        prog="whisper",
-        description="Run Whisper on input audio file"
+        prog="whisper", description="Run Whisper on input audio file"
     )
     parser.add_argument("--wav", "-w", type=str, required=True, help="Input audio file")
-    parser.add_argument("--model_type", "-t", type=str, choices=["tiny", "base", "small", "large", "large-v3", "turbo"], required=True, help="model type, only support tiny, base and small currently")
-    parser.add_argument("--model_path", "-p", type=str, required=False, default="../models/models-ax650", help="model path for *.axmodel, tokens.txt, positional_embedding.bin")
-    parser.add_argument("--language", "-l", type=str, required=False, default="zh", help="Target language, support en, zh, ja, and others. See languages.py for more options.")
-    parser.add_argument("--task", type=str, required=False, choices=["translate", "transcribe"], default="transcribe")
-    parser.add_argument("--print_rtf", action="store_true", help="Print Real-Time Factor")
+    parser.add_argument(
+        "--model_type",
+        "-t",
+        type=str,
+        choices=["tiny", "base", "small", "large", "large-v3", "turbo"],
+        required=True,
+        help="model type, only support tiny, base and small currently",
+    )
+    parser.add_argument(
+        "--model_path",
+        "-p",
+        type=str,
+        required=False,
+        default="../models-ax650",
+        help="model path for *.axmodel, tokens.txt, positional_embedding.bin",
+    )
+    parser.add_argument(
+        "--language",
+        "-l",
+        type=str,
+        required=False,
+        default="zh",
+        help="Target language, support en, zh, ja, and others. See languages.py for more options.",
+    )
+    parser.add_argument(
+        "--task",
+        type=str,
+        required=False,
+        choices=["translate", "transcribe"],
+        default="transcribe",
+    )
+    parser.add_argument(
+        "--print_rtf", action="store_true", help="Print Real-Time Factor"
+    )
     return parser.parse_args()
 
 
-def print_args(args):
-    print(f"wav: {args.wav}")
-    print(f"model_type: {args.model_type}")
-    print(f"model_path: {args.model_path}")
-    print(f"language: {args.language}")
-    print(f"task: {args.task}")
-
-
 def main():
     args = get_args()
-    print_args(args)
+    print(vars(args))
 
     # Check wav existence
     wav_path = args.wav
@@ -36,19 +56,19 @@ def main():
 
     model = Whisper(args.model_type, args.model_path, args.language, args.task)
 
-    
-    
-    print("\n预测结果:")
+    print("ASR result:")
     start = time.time()
     print(model.run(wav_path))
     end = time.time()
 
     if args.print_rtf:
         import librosa
+
         samples, sr = librosa.load(wav_path, sr=16000)
         duration = len(samples) / sr
         process_time = end - start
         print(f"RTF: {process_time / duration}")
 
+
 if __name__ == "__main__":
     main()

python/test_wer.py

@@ -10,35 +10,35 @@ def setup_logging():
     # 获取脚本所在目录
     script_dir = os.path.dirname(os.path.abspath(__file__))
     log_file = os.path.join(script_dir, "test_wer.log")
-    
+
     # 配置日志格式
-    log_format = '%(asctime)s - %(levelname)s - %(message)s'
-    date_format = '%Y-%m-%d %H:%M:%S'
-    
+    log_format = "%(asctime)s - %(levelname)s - %(message)s"
+    date_format = "%Y-%m-%d %H:%M:%S"
+
     # 创建logger
     logger = logging.getLogger()
     logger.setLevel(logging.INFO)
-    
+
     # 清除现有的handler
     for handler in logger.handlers[:]:
         logger.removeHandler(handler)
-    
+
     # 创建文件handler
-    file_handler = logging.FileHandler(log_file, mode='a', encoding='utf-8')
+    file_handler = logging.FileHandler(log_file, mode="a", encoding="utf-8")
     file_handler.setLevel(logging.INFO)
     file_formatter = logging.Formatter(log_format, date_format)
     file_handler.setFormatter(file_formatter)
-    
+
     # 创建控制台handler
     console_handler = logging.StreamHandler()
     console_handler.setLevel(logging.INFO)
     console_formatter = logging.Formatter(log_format, date_format)
     console_handler.setFormatter(console_formatter)
-    
+
     # 添加handler到logger
     logger.addHandler(file_handler)
     logger.addHandler(console_handler)
-    
+
     return logger
 
 
@@ -46,21 +46,21 @@ class AIShellDataset:
     def __init__(self, gt_path: str):
         """
         初始化数据集
-        
+
         Args:
             json_path: voice.json文件的路径
         """
         self.gt_path = gt_path
         self.dataset_dir = os.path.dirname(gt_path)
         self.voice_dir = os.path.join(self.dataset_dir, "aishell_S0764")
-        
+
         # 检查必要文件和文件夹是否存在
         assert os.path.exists(gt_path), f"gt文件不存在: {gt_path}"
         assert os.path.exists(self.voice_dir), f"aishell_S0764文件夹不存在: {self.voice_dir}"
-        
+
         # 加载数据
         self.data = []
-        with open(gt_path, 'r', encoding='utf-8') as f:
+        with open(gt_path, "r", encoding="utf-8") as f:
             for line in f:
                 line = line.strip()
                 audio_path, gt = line.split(" ")
@@ -70,50 +70,50 @@ class AIShellDataset:
         # 使用logging而不是print
         logger = logging.getLogger()
         logger.info(f"加载了 {len(self.data)} 条数据")
-    
+
     def __iter__(self):
         """返回迭代器"""
         self.index = 0
         return self
-    
+
     def __next__(self):
         """返回下一个数据项"""
         if self.index >= len(self.data):
             raise StopIteration
-        
+
         item = self.data[self.index]
         audio_path = item["audio_path"]
         ground_truth = item["gt"]
-        
+
         self.index += 1
         return audio_path, ground_truth
-    
+
     def __len__(self):
         """返回数据集大小"""
         return len(self.data)
-    
+
 
 class CommonVoiceDataset:
     """Common Voice数据集解析器"""
-    
+
     def __init__(self, tsv_path: str):
         """
         初始化数据集
-        
+
         Args:
             json_path: voice.json文件的路径
         """
         self.tsv_path = tsv_path
         self.dataset_dir = os.path.dirname(tsv_path)
         self.voice_dir = os.path.join(self.dataset_dir, "clips")
-        
+
         # 检查必要文件和文件夹是否存在
         assert os.path.exists(tsv_path), f"{tsv_path}文件不存在: {tsv_path}"
         assert os.path.exists(self.voice_dir), f"voice文件夹不存在: {self.voice_dir}"
-        
+
         # 加载JSON数据
         self.data = []
-        with open(tsv_path, 'r', encoding='utf-8') as f:
+        with open(tsv_path, "r", encoding="utf-8") as f:
             f.readline()
             for line in f:
                 line = line.strip()
@@ -122,43 +122,77 @@ class CommonVoiceDataset:
                 gt = splits[2]
                 audio_path = os.path.join(self.voice_dir, audio_path)
                 self.data.append({"audio_path": audio_path, "gt": gt})
-        
+
         # 使用logging而不是print
         logger = logging.getLogger()
         logger.info(f"加载了 {len(self.data)} 条数据")
-    
+
     def __iter__(self):
         """返回迭代器"""
         self.index = 0
         return self
-    
+
     def __next__(self):
         """返回下一个数据项"""
         if self.index >= len(self.data):
             raise StopIteration
-        
+
         item = self.data[self.index]
         audio_path = item["audio_path"]
         ground_truth = item["gt"]
-        
+
         self.index += 1
         return audio_path, ground_truth
-    
+
     def __len__(self):
         """返回数据集大小"""
         return len(self.data)
 
+
 def get_args():
-    parser = argparse.ArgumentParser(
-        prog="whisper",
-        description="Test WER on dataset"
+    parser = argparse.ArgumentParser(prog="whisper", description="Test WER on dataset")
+    parser.add_argument(
+        "--dataset",
+        "-d",
+        type=str,
+        required=True,
+        choices=["aishell", "common_voice"],
+        help="Test dataset",
+    )
+    parser.add_argument(
+        "--gt_path",
+        "-g",
+        type=str,
+        required=True,
+        help="Test dataset ground truth file",
+    )
+    parser.add_argument(
+        "--max_num", type=int, default=-1, required=False, help="Maximum test data num"
+    )
+    parser.add_argument(
+        "--model_type",
+        "-t",
+        type=str,
+        choices=["tiny", "base", "small", "large", "large-v3", "turbo"],
+        required=True,
+        help="model type, only support tiny, base and small currently",
+    )
+    parser.add_argument(
+        "--model_path",
+        "-p",
+        type=str,
+        required=False,
+        default="../models/models-ax650",
+        help="model path for *.axmodel, tokens.txt, positional_embedding.bin",
+    )
+    parser.add_argument(
+        "--language",
+        "-l",
+        type=str,
+        required=False,
+        default="zh",
+        help="Target language, support en, zh, ja, and others. See languages.py for more options.",
     )
-    parser.add_argument("--dataset", "-d", type=str, required=True, choices=["aishell", "common_voice"], help="Test dataset")
-    parser.add_argument("--gt_path", "-g", type=str, required=True, help="Test dataset ground truth file")
-    parser.add_argument("--max_num", type=int, default=-1, required=False, help="Maximum test data num")
-    parser.add_argument("--model_type", "-t", type=str, choices=["tiny", "base", "small", "large", "large-v3", "turbo"], required=True, help="model type, only support tiny, base and small currently")
-    parser.add_argument("--model_path", "-p", type=str, required=False, default="../models/models-ax650", help="model path for *.axmodel, tokens.txt, positional_embedding.bin")
-    parser.add_argument("--language", "-l", type=str, required=False, default="zh", help="Target language, support en, zh, ja, and others. See languages.py for more options.")
     return parser.parse_args()
 
 
@@ -173,42 +207,42 @@ def print_args(args):
 
 
 def min_distance(word1: str, word2: str) -> int:
- 
+
     row = len(word1) + 1
     column = len(word2) + 1
- 
-    cache = [ [0]*column for i in range(row) ]
- 
+
+    cache = [[0] * column for i in range(row)]
+
     for i in range(row):
         for j in range(column):
- 
-            if i ==0 and j ==0:
+
+            if i == 0 and j == 0:
                 cache[i][j] = 0
-            elif i == 0 and j!=0:
+            elif i == 0 and j != 0:
                 cache[i][j] = j
-            elif j == 0 and i!=0:
+            elif j == 0 and i != 0:
                 cache[i][j] = i
             else:
-                if word1[i-1] == word2[j-1]:
-                    cache[i][j] = cache[i-1][j-1]
+                if word1[i - 1] == word2[j - 1]:
+                    cache[i][j] = cache[i - 1][j - 1]
                 else:
-                    replace = cache[i-1][j-1] + 1
-                    insert = cache[i][j-1] + 1
-                    remove = cache[i-1][j] + 1
- 
+                    replace = cache[i - 1][j - 1] + 1
+                    insert = cache[i][j - 1] + 1
+                    remove = cache[i - 1][j] + 1
+
                     cache[i][j] = min(replace, insert, remove)
- 
-    return cache[row-1][column-1]
+
+    return cache[row - 1][column - 1]
 
 
 def remove_punctuation(text):
     # 定义正则表达式模式，匹配所有标点符号
     # 这个模式包括常见的标点符号和中文标点
-    pattern = r'[^\w\s]|_'
-    
+    pattern = r"[^\w\s]|_"
+
     # 使用sub方法将所有匹配的标点符号替换为空字符串
-    cleaned_text = re.sub(pattern, '', text)
-    
+    cleaned_text = re.sub(pattern, "", text)
+
     return cleaned_text
 
 
@@ -254,7 +288,7 @@ def main():
 
         hyp.append(hypothesis)
         references.append(reference)
-        
+
         line_content = f"({n+1}/{max_data_num}) {os.path.basename(audio_path)}  gt: {reference}  predict: {hypothesis}  WER: {character_error_rate}%"
         wer_file.write(line_content + "\n")
         logger.info(line_content)
@@ -268,5 +302,6 @@ def main():
     wer_file.write(f"Total WER: {total_character_error_rate}%")
     wer_file.close()
 
+
 if __name__ == "__main__":
     main()

python/whisper.py

@@ -1,4 +1,4 @@
-import axengine as axe 
+import axengine as axe
 import numpy as np
 import librosa
 import os
@@ -9,27 +9,27 @@ from dataclasses import dataclass
 import zhconv
 
 
-NEG_INF = float("-inf")
-
 @dataclass
 class WhisperConfig:
-    n_mels          : int = 0
-    sample_rate     : int = 0
-    n_fft           : int = 0
-    hop_length      : int = 0
-
-    sot             : int = 0
-    eot             : int = 0
-    blank_id        : int = 0
-    no_timestamps   : int = 0
-    no_speech       : int = 0
-    translate       : int = 0
-    transcribe      : int = 0
-    n_vocab         : int = 0
-    n_text_ctx      : int = 0
-    n_text_state    : int = 0
-
-    sot_sequence    : np.ndarray = field(default_factory=lambda: np.array([0,0,0,0], dtype=np.int32))
+    n_mels: int = 0
+    sample_rate: int = 0
+    n_fft: int = 0
+    hop_length: int = 0
+
+    sot: int = 0
+    eot: int = 0
+    blank_id: int = 0
+    no_timestamps: int = 0
+    no_speech: int = 0
+    translate: int = 0
+    transcribe: int = 0
+    n_vocab: int = 0
+    n_text_ctx: int = 0
+    n_text_state: int = 0
+
+    sot_sequence: np.ndarray = field(
+        default_factory=lambda: np.array([0, 0, 0, 0], dtype=np.int32)
+    )
 
 
 class Whisper:
@@ -38,35 +38,41 @@ class Whisper:
 
         self.language = language
         self.task = task
-        self.encoder, self.decoder_main, self.decoder_loop, self.pe, self.tokenizer, model_config = \
-            self.load_model(model_type, model_path, language, task)
+        self.encoder, self.decoder, self.tokenizer, model_config = self.load_model(
+            model_type, model_path, language, task
+        )
         self.config = self.load_config(model_config)
 
-
     def load_model(self, model_type, model_path, language, task):
         encoder_path = f"{model_type}/{model_type}-encoder.axmodel"
-        decoder_main_path = f"{model_type}/{model_type}-decoder-main.axmodel"
-        decoder_loop_path = f"{model_type}/{model_type}-decoder-loop.axmodel"
-        pe_path = f"{model_type}/{model_type}-positional_embedding.bin"
+        decoder_path = f"{model_type}/{model_type}-decoder.axmodel"
         model_config_file = f"{model_type}/{model_type}_config.json"
+        token_file = f"{model_type}/{model_type}-tokens.txt"
 
-        required_files = [os.path.join(model_path, i) for i in (encoder_path, decoder_main_path, decoder_loop_path, pe_path, model_config_file)]
+        required_files = [
+            os.path.join(model_path, i)
+            for i in (encoder_path, decoder_path, model_config_file, token_file)
+        ]
         # Check file existence
         for i, file_path in enumerate(required_files):
             assert os.path.exists(file_path), f"{file_path} NOT exist"
 
         # Load encoder
-        encoder = axe.InferenceSession(required_files[0], providers=['AxEngineExecutionProvider'])
+        encoder = axe.InferenceSession(
+            required_files[0], providers=["AxEngineExecutionProvider"]
+        )
         # Load decoder main
-        decoder_main = axe.InferenceSession(required_files[1], providers=['AxEngineExecutionProvider'])
-        # Load decoder loop
-        decoder_loop = axe.InferenceSession(required_files[2], providers=['AxEngineExecutionProvider'])
-        # Load position embedding
-        pe = np.fromfile(required_files[3], dtype=np.float32)
+        decoder = axe.InferenceSession(
+            required_files[1], providers=["AxEngineExecutionProvider"]
+        )
         # Load tokens
-        model_config = json.load(open(required_files[4], "r"))
-        model_config["all_language_tokens"] = [int(i) for i in model_config["all_language_tokens"].split(",")]
-        model_config["all_language_codes"] = [i for i in model_config["all_language_codes"].split(",")]
+        model_config = json.load(open(required_files[2], "r"))
+        model_config["all_language_tokens"] = [
+            int(i) for i in model_config["all_language_tokens"].split(",")
+        ]
+        model_config["all_language_codes"] = [
+            i for i in model_config["all_language_codes"].split(",")
+        ]
         tokenizer = get_tokenizer(
             model_config["is_multilingual"],
             num_languages=len(model_config["all_language_codes"]),
@@ -74,8 +80,9 @@ class Whisper:
             task=task,
         )
 
-        return encoder, decoder_main, decoder_loop, pe, tokenizer, model_config
-    
+        self.id2token = self.load_tokens(required_files[3])
+
+        return encoder, decoder, tokenizer, model_config
 
     def load_config(self, model_config):
         config = WhisperConfig
@@ -94,34 +101,46 @@ class Whisper:
         config.n_vocab = model_config["n_vocab"]
         config.n_text_ctx = model_config["n_text_ctx"]
         config.n_text_state = model_config["n_text_state"]
+        config.n_text_layer = model_config["n_text_layer"]
 
-        lang_token = model_config["all_language_tokens"][model_config["all_language_codes"].index(self.language)]
-        task_token = config.transcribe if self.task == "transcribe" else config.translate
-        config.sot_sequence = np.array([config.sot, lang_token, task_token, config.no_timestamps], dtype=np.int32)
+        lang_token = model_config["all_language_tokens"][
+            model_config["all_language_codes"].index(self.language)
+        ]
+        task_token = (
+            config.transcribe if self.task == "transcribe" else config.translate
+        )
+        config.sot_sequence = np.array(
+            [config.sot, lang_token, task_token, config.no_timestamps], dtype=np.int32
+        )
 
         return config
-    
+
+    def load_tokens(self, filename):
+        tokens = dict()
+        with open(filename, "r") as f:
+            for line in f:
+                t, i = line.split()
+                tokens[int(i)] = t
+        return tokens
 
     def load_audio(self, audio: str):
         data, sample_rate = librosa.load(audio, sr=self.config.sample_rate)
         samples = np.ascontiguousarray(data)
         return samples, sample_rate
 
+    def compute_feature(self, audio: np.ndarray):
+        mel = librosa.feature.melspectrogram(
+            y=audio,
+            sr=self.config.sample_rate,
+            n_fft=self.config.n_fft,
+            hop_length=self.config.hop_length,
+            window="hann",
+            center=True,
+            pad_mode="reflect",
+            power=2.0,
+            n_mels=self.config.n_mels,
+        )
 
-    def compute_feature(self, audio: np.ndarray, padding = 480000):
-        if padding > 0:
-            audio = np.concatenate((audio, np.zeros((padding,), dtype=np.float32)), axis=-1)
-
-        mel = librosa.feature.melspectrogram(y=audio, 
-                                             sr=self.config.sample_rate, 
-                                             n_fft=self.config.n_fft, 
-                                             hop_length=self.config.hop_length, 
-                                             window="hann", 
-                                             center=True, 
-                                             pad_mode="reflect", 
-                                             power=2.0, 
-                                             n_mels=self.config.n_mels)
-        
         log_spec = np.log10(np.maximum(mel, 1e-10))
         log_spec = np.maximum(log_spec, log_spec.max() - 8.0)
         mel = (log_spec + 4.0) / 4.0
@@ -129,31 +148,71 @@ class Whisper:
         target = 3000
         if mel.shape[1] > target:
             # -50 so that there are some zero tail paddings.
-            mel = mel[:, : target]
+            mel = mel[:, :target]
             mel[:, -50:] = 0
 
         # We don't need to pad it to 30 seconds now!
         if mel.shape[1] < target:
-            mel = np.concatenate((mel, np.zeros((self.config.n_mels, target - mel.shape[1]), dtype=np.float32)), axis=-1)
-
-        return mel
-    
-
-    def supress_tokens(self, logits, is_initial):
-        if is_initial:
-            logits[self.config.eot] = NEG_INF
-            logits[self.config.blank_id] = NEG_INF
-
-        logits[self.config.no_timestamps] = NEG_INF
-        logits[self.config.sot] = NEG_INF
-        logits[self.config.no_speech] = NEG_INF
+            mel = np.concatenate(
+                (
+                    mel,
+                    np.zeros(
+                        (self.config.n_mels, target - mel.shape[1]), dtype=np.float32
+                    ),
+                ),
+                axis=-1,
+            )
+
+        return mel[np.newaxis, ...]
+
+    def run_encoder(
+        self,
+        mel: np.ndarray,
+    ) -> List[np.ndarray]:
+        cross_kv = self.encoder.run(
+            None,
+            {
+                self.encoder.get_inputs()[0].name: mel,
+            },
+        )
+        return cross_kv
 
-        if self.task == "transcribe":
-            logits[self.config.translate] = NEG_INF
-        else:
-            logits[self.config.transcribe] = NEG_INF
-        return logits
+    def run_decoder(self, inputs: List[np.ndarray]) -> List[np.ndarray]:
+        feed = {
+            self.decoder.get_inputs()[i].name: inputs[i] for i in range(len(inputs))
+        }
 
+        out = self.decoder.run(
+            None,
+            feed,
+        )
+        return out
+
+    def get_self_cache(self) -> List[np.ndarray]:
+        self_cache = []
+        batch_size = 1
+        for i in range(self.config.n_text_layer):
+            k = np.zeros(
+                (batch_size, self.config.n_text_ctx, self.config.n_text_state),
+                dtype=np.float32,
+            )
+            v = np.zeros(
+                (batch_size, self.config.n_text_ctx, self.config.n_text_state),
+                dtype=np.float32,
+            )
+            self_cache.extend([k, v])
+        return self_cache
+
+    def causal_mask_1d(self, n: int, L: int):
+        """
+        Returns a 1-D int mask of shape (L,) with:
+        0 -> allowed
+        1 -> masked (will be converted to -inf later)
+        """
+        mask = np.ones((L,), dtype=np.int32)
+        if n > 0:
+            mask[:n] = 0
+        return mask
 
     def run(self, audio: Union[str, np.ndarray]) -> str:
         if isinstance(audio, str):
@@ -161,64 +220,56 @@ class Whisper:
 
         mel = self.compute_feature(audio)
 
-        # Run encoder
-        x = self.encoder.run(None, input_feed={"mel": mel[None, ...]})
-        n_layer_cross_k, n_layer_cross_v = x
-
-        # Run decoder_main
-        x = self.decoder_main.run(None, input_feed={
-            "tokens": self.config.sot_sequence[None, ...],
-            "n_layer_cross_k": n_layer_cross_k,
-            "n_layer_cross_v": n_layer_cross_v
-        })
-        logits, n_layer_self_k_cache, n_layer_self_v_cache = x
-
-        # Decode token
-        logits = logits[0, -1, :]
-        logits = self.supress_tokens(logits, is_initial=True)
-        # logits.tofile("logits.bin")
-        max_token_id = np.argmax(logits)
-        output_tokens = []
-
-        # Position embedding offset
-        offset = self.config.sot_sequence.shape[0]
-
-        # Autoregressively run decoder until token meets EOT
-        for i in range(self.config.n_text_ctx - self.config.sot_sequence.shape[0]):
-            if max_token_id >= self.config.eot:
-                break
-
-            output_tokens.append(max_token_id)
-
-            mask = np.zeros((self.config.n_text_ctx,), dtype=np.float32)
-            mask[: self.config.n_text_ctx - offset - 1] = NEG_INF
-
-            # Run decoder_loop
-            x = self.decoder_loop.run(None, input_feed={
-                "tokens": np.array([[output_tokens[-1]]], dtype=np.int32),
-                "in_n_layer_self_k_cache": n_layer_self_k_cache,
-                "in_n_layer_self_v_cache": n_layer_self_v_cache,
-                "n_layer_cross_k": n_layer_cross_k,
-                "n_layer_cross_v": n_layer_cross_v,
-                "positional_embedding": self.pe[offset * self.config.n_text_state : (offset + 1) * self.config.n_text_state][None, ...],
-                "mask": mask
-            })
-            logits, n_layer_self_k_cache, n_layer_self_v_cache = x
-
-            # Decode token
+        cross_kv = self.run_encoder(mel)
+
+        self_kv = self.get_self_cache()
+
+        offset = np.array([0], dtype=np.int32)
+        for t in self.config.sot_sequence:
+            token = np.array([[t]], dtype=np.int32)  # sot
+            mask = self.causal_mask_1d(offset.item(), self.config.n_text_ctx)
+
+            out = self.run_decoder([token] + self_kv + cross_kv + [offset, mask])
+
+            for i in range(1, len(out)):
+                self_kv[i - 1][:, offset.item() : offset.item() + 1, :] = out[i]
+
+            offset += 1
+
+        idx = out[0][0, 0].argmax()
+
+        eot = self.config.eot
+
+        ans = []
+
+        while idx != eot and offset.item() < 100:
+            ans.append(idx)
+            token = np.array([[idx]], dtype=np.int32)
+
+            mask = self.causal_mask_1d(offset.item(), self.config.n_text_ctx)
+
+            out = self.run_decoder([token] + self_kv + cross_kv + [offset, mask])
+
+            for i in range(1, len(out)):
+                self_kv[i - 1][:, offset.item() : offset.item() + 1, :] = out[i]
+
             offset += 1
-            logits = self.supress_tokens(logits.flatten(), is_initial=False)
-            max_token_id = np.argmax(logits)
-        
-        text = self.tokenizer.decode(output_tokens)
+            idx = out[0][0, 0].argmax()
+
+        # print(ans)
+
+        s = b""
+        for i in ans:
+            if i in self.id2token:
+                s += base64.b64decode(self.id2token[i])
+
+        text = s.decode().strip()
 
         if self.language == "zh":
             try:
-                sim_zh = zhconv.convert(text, 'zh-hans')
+                sim_zh = zhconv.convert(text, "zh-hans")
                 return sim_zh
             except:
                 return text
-            
-        return text
 
-        
+        return text