From 861147c7308b91068ffa02724fdf74ee623a909e Mon Sep 17 00:00:00 2001 From: zhifu gao Date: Wed, 24 Apr 2024 16:03:38 +0800 Subject: [PATCH] Dev gzf exp (#1654) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * sensevoice finetune * sensevoice finetune * sensevoice finetune * sensevoice finetune * sensevoice finetune * sensevoice finetune * sensevoice finetune * sensevoice finetune * sensevoice finetune * sensevoice finetune * bugfix * update with main (#1631) * update seaco finetune * v1.0.24 --------- Co-authored-by: 维石 * sensevoice * sensevoice * sensevoice * update with main (#1638) * update seaco finetune * v1.0.24 * update rwkv template --------- Co-authored-by: 维石 * sensevoice * sensevoice * sensevoice * sensevoice * sensevoice * sensevoice * sensevoice * sensevoice * sensevoice * sensevoice * sensevoice * sensevoice * sensevoice * sensevoice * sensevoice * sense voice * sense voice * sense voice * sense voice * sense voice * sense voice * sense voice * sense voice * sense voice * sense voice * sense voice * sense voice * sense voice * sense voice * sense voice * sense voice * sense voice * sense voice * sense voice * sense voice * whisper * whisper * update style * update style --------- Co-authored-by: 维石 --- .pre-commit-config.yaml | 6 + docs/conf.py | 18 +- docs/m2met2/conf.py | 21 +- .../aishell/paraformer/utils/compute_wer.py | 164 +- .../paraformer/utils/extract_embeds.py | 16 +- .../paraformer/utils/postprocess_text_zh.py | 39 +- .../aishell/paraformer/utils/text2token.py | 8 +- .../aishell/paraformer/utils/text_tokenize.py | 18 +- .../aishell/paraformer/utils/textnorm_zh.py | 447 +++-- .../common_voice/whisper_lid/demo_funasr.py | 2 +- .../whisper_lid/demo_modelscope.py | 8 +- .../bicif_paraformer/demo.py | 17 +- .../bicif_paraformer/export.py | 12 +- .../campplus_sv/demo.py | 9 +- .../conformer/demo.py | 5 +- .../contextual_paraformer/demo.py | 8 +- .../ct_transformer/demo.py | 10 +- .../ct_transformer/export.py | 11 +- .../ct_transformer_streaming/demo.py | 2 +- .../ct_transformer_streaming/export.py | 11 +- .../emotion2vec/demo.py | 17 +- .../fsmn_vad_streaming/demo.py | 23 +- .../fsmn_vad_streaming/export.py | 4 +- .../lcbnet/compute_wer_details.py | 221 ++- .../lcbnet/demo.py | 12 +- .../monotonic_aligner/demo.py | 13 +- .../paraformer-zh-spk/demo.py | 19 +- .../paraformer/demo.py | 29 +- .../paraformer/export.py | 10 +- .../paraformer_streaming/demo.py | 38 +- .../paraformer_streaming/export.py | 10 +- .../qwen_audio/demo_chat.py | 3 +- .../qwen_audio/demo_chat_from_local.py | 14 +- .../qwen_audio/demo_from_local.py | 7 +- .../industrial_data_pretraining/scama/demo.py | 42 +- .../seaco_paraformer/demo.py | 28 +- .../sense_voice/demo.py | 23 +- .../transducer/demo.py | 11 +- .../uniasr/demo.py | 12 +- .../whisper/demo.py | 30 +- .../whisper/demo_from_openai.py | 28 +- .../inverse_text_normalization/de/__init__.py | 10 +- .../de/taggers/cardinal.py | 29 +- .../de/taggers/date.py | 21 +- .../de/taggers/decimal.py | 18 +- .../de/taggers/electronic.py | 12 +- .../de/taggers/fraction.py | 30 +- .../de/taggers/measure.py | 30 +- .../de/taggers/money.py | 38 +- .../de/taggers/ordinal.py | 12 +- .../de/taggers/telephone.py | 13 +- .../de/taggers/time.py | 4 +- .../de/taggers/tokenize_and_classify.py | 74 +- .../de/taggers/whitelist.py | 3 +- .../de/verbalizers/__init__.py | 2 - .../de/verbalizers/cardinal.py | 5 +- .../de/verbalizers/decimal.py | 15 +- .../de/verbalizers/measure.py | 7 +- .../de/verbalizers/money.py | 5 +- .../de/verbalizers/time.py | 31 +- .../de/verbalizers/verbalize.py | 11 +- .../de/verbalizers/verbalize_final.py | 9 +- .../inverse_text_normalization/en/__init__.py | 10 +- .../en/clean_eval_data.py | 116 +- .../en/taggers/cardinal.py | 47 +- .../en/taggers/date.py | 31 +- .../en/taggers/decimal.py | 50 +- .../en/taggers/electronic.py | 52 +- .../en/taggers/fraction.py | 1 - .../en/taggers/measure.py | 40 +- .../en/taggers/money.py | 41 +- .../en/taggers/ordinal.py | 3 +- .../en/taggers/punctuation.py | 5 +- .../en/taggers/telephone.py | 71 +- .../en/taggers/time.py | 52 +- .../en/taggers/tokenize_and_classify.py | 13 +- .../en/taggers/whitelist.py | 3 +- .../en/taggers/word.py | 3 +- .../inverse_text_normalization/en/utils.py | 8 +- .../en/verbalizers/cardinal.py | 16 +- .../en/verbalizers/date.py | 26 +- .../en/verbalizers/decimal.py | 22 +- .../en/verbalizers/electronic.py | 20 +- .../en/verbalizers/fraction.py | 3 +- .../en/verbalizers/measure.py | 7 +- .../en/verbalizers/money.py | 5 +- .../en/verbalizers/ordinal.py | 12 +- .../en/verbalizers/telephone.py | 15 +- .../en/verbalizers/time.py | 21 +- .../en/verbalizers/verbalize.py | 3 +- .../en/verbalizers/verbalize_final.py | 9 +- .../en/verbalizers/whitelist.py | 15 +- .../en/verbalizers/word.py | 18 +- .../inverse_text_normalization/es/__init__.py | 9 +- .../es/data/__init__.py | 2 - .../es/data/electronic/__init__.py | 2 - .../es/data/numbers/__init__.py | 2 - .../es/data/ordinals/__init__.py | 2 - .../es/data/time/__init__.py | 2 - .../es/taggers/__init__.py | 2 - .../es/taggers/cardinal.py | 47 +- .../es/taggers/date.py | 17 +- .../es/taggers/decimal.py | 47 +- .../es/taggers/electronic.py | 44 +- .../es/taggers/measure.py | 42 +- .../es/taggers/money.py | 46 +- .../es/taggers/ordinal.py | 47 +- .../es/taggers/punctuation.py | 5 +- .../es/taggers/telephone.py | 19 +- .../es/taggers/time.py | 55 +- .../es/taggers/tokenize_and_classify.py | 13 +- .../es/taggers/whitelist.py | 3 +- .../es/taggers/word.py | 3 +- .../inverse_text_normalization/es/utils.py | 5 +- .../es/verbalizers/__init__.py | 2 - .../es/verbalizers/cardinal.py | 15 +- .../es/verbalizers/date.py | 13 +- .../es/verbalizers/decimal.py | 25 +- .../es/verbalizers/electronic.py | 19 +- .../es/verbalizers/measure.py | 7 +- .../es/verbalizers/money.py | 5 +- .../es/verbalizers/ordinal.py | 17 +- .../es/verbalizers/telephone.py | 9 +- .../es/verbalizers/time.py | 21 +- .../es/verbalizers/verbalize.py | 3 +- .../es/verbalizers/verbalize_final.py | 9 +- .../es/verbalizers/whitelist.py | 15 +- .../es/verbalizers/word.py | 18 +- .../export_models.py | 137 +- .../inverse_text_normalization/fr/__init__.py | 9 +- .../fr/data/__init__.py | 2 - .../fr/data/electronic/__init__.py | 2 - .../fr/data/measurements/__init__.py | 2 - .../fr/data/money/__init__.py | 2 - .../fr/data/numbers/__init__.py | 2 - .../fr/data/ordinals/__init__.py | 2 - .../fr/data/time/__init__.py | 2 - .../fr/graph_utils.py | 37 +- .../fr/taggers/__init__.py | 2 - .../fr/taggers/cardinal.py | 78 +- .../fr/taggers/date.py | 18 +- .../fr/taggers/decimal.py | 41 +- .../fr/taggers/electronic.py | 29 +- .../fr/taggers/fraction.py | 23 +- .../fr/taggers/measure.py | 20 +- .../fr/taggers/money.py | 35 +- .../fr/taggers/ordinal.py | 19 +- .../fr/taggers/punctuation.py | 5 +- .../fr/taggers/telephone.py | 13 +- .../fr/taggers/time.py | 23 +- .../fr/taggers/tokenize_and_classify.py | 13 +- .../fr/taggers/whitelist.py | 4 +- .../fr/taggers/word.py | 3 +- .../inverse_text_normalization/fr/utils.py | 5 +- .../fr/verbalizers/__init__.py | 2 - .../fr/verbalizers/cardinal.py | 15 +- .../fr/verbalizers/date.py | 19 +- .../fr/verbalizers/decimal.py | 24 +- .../fr/verbalizers/electronic.py | 15 +- .../fr/verbalizers/fraction.py | 27 +- .../fr/verbalizers/measure.py | 23 +- .../fr/verbalizers/money.py | 5 +- .../fr/verbalizers/ordinal.py | 21 +- .../fr/verbalizers/telephone.py | 9 +- .../fr/verbalizers/time.py | 24 +- .../fr/verbalizers/verbalize.py | 5 +- .../fr/verbalizers/verbalize_final.py | 9 +- .../fr/verbalizers/whitelist.py | 8 +- .../fr/verbalizers/word.py | 11 +- .../inverse_text_normalization/id/__init__.py | 8 +- .../id/clean_eval_data.py | 114 +- .../id/graph_utils.py | 60 +- .../id/taggers/cardinal.py | 78 +- .../id/taggers/date.py | 30 +- .../id/taggers/decimal.py | 48 +- .../id/taggers/electronic.py | 57 +- .../id/taggers/measure.py | 39 +- .../id/taggers/money.py | 40 +- .../id/taggers/ordinal.py | 4 +- .../id/taggers/punctuation.py | 4 +- .../id/taggers/telephone.py | 70 +- .../id/taggers/time.py | 65 +- .../id/taggers/tokenize_and_classify.py | 12 +- .../id/taggers/whitelist.py | 2 +- .../id/taggers/word.py | 2 +- .../inverse_text_normalization/id/utils.py | 6 +- .../id/verbalizers/cardinal.py | 14 +- .../id/verbalizers/date.py | 24 +- .../id/verbalizers/decimal.py | 20 +- .../id/verbalizers/electronic.py | 18 +- .../id/verbalizers/fraction.py | 2 +- .../id/verbalizers/measure.py | 12 +- .../id/verbalizers/money.py | 10 +- .../id/verbalizers/ordinal.py | 11 +- .../id/verbalizers/telephone.py | 14 +- .../id/verbalizers/time.py | 20 +- .../id/verbalizers/verbalize.py | 2 +- .../id/verbalizers/verbalize_final.py | 8 +- .../id/verbalizers/whitelist.py | 13 +- .../id/verbalizers/word.py | 17 +- .../inverse_normalize.py | 128 +- .../inverse_text_normalization/ja/__init__.py | 8 +- .../ja/clean_eval_data.py | 114 +- .../ja/graph_utils.py | 12 +- .../ja/taggers/cardinal.py | 240 ++- .../ja/taggers/date.py | 30 +- .../ja/taggers/decimal.py | 51 +- .../ja/taggers/electronic.py | 57 +- .../ja/taggers/fraction.py | 16 +- .../ja/taggers/measure.py | 46 +- .../ja/taggers/money.py | 41 +- .../ja/taggers/ordinal.py | 34 +- .../ja/taggers/preprocessor.py | 11 +- .../ja/taggers/punctuation.py | 4 +- .../ja/taggers/telephone.py | 74 +- .../ja/taggers/time.py | 61 +- .../ja/taggers/tokenize_and_classify.py | 26 +- .../ja/taggers/whitelist.py | 2 +- .../ja/taggers/word.py | 8 +- .../inverse_text_normalization/ja/utils.py | 6 +- .../ja/verbalizers/cardinal.py | 16 +- .../ja/verbalizers/date.py | 26 +- .../ja/verbalizers/decimal.py | 22 +- .../ja/verbalizers/electronic.py | 20 +- .../ja/verbalizers/fraction.py | 13 +- .../ja/verbalizers/measure.py | 13 +- .../ja/verbalizers/money.py | 11 +- .../ja/verbalizers/ordinal.py | 13 +- .../ja/verbalizers/telephone.py | 15 +- .../ja/verbalizers/time.py | 21 +- .../ja/verbalizers/verbalize.py | 3 +- .../ja/verbalizers/verbalize_final.py | 9 +- .../ja/verbalizers/whitelist.py | 15 +- .../ja/verbalizers/word.py | 18 +- .../inverse_text_normalization/ko/__init__.py | 9 +- .../ko/clean_eval_data.py | 115 +- .../ko/graph_utils.py | 61 +- .../ko/taggers/cardinal.py | 123 +- .../ko/taggers/date.py | 45 +- .../ko/taggers/decimal.py | 54 +- .../ko/taggers/electronic.py | 58 +- .../ko/taggers/fraction.py | 23 +- .../ko/taggers/measure.py | 19 +- .../ko/taggers/money.py | 7 +- .../ko/taggers/punctuation.py | 5 +- .../ko/taggers/telephone.py | 27 +- .../ko/taggers/time.py | 43 +- .../ko/taggers/tokenize_and_classify.py | 13 +- .../ko/taggers/whitelist.py | 3 +- .../ko/taggers/word.py | 3 +- .../inverse_text_normalization/ko/utils.py | 11 +- .../ko/verbalizers/cardinal.py | 15 +- .../ko/verbalizers/date.py | 48 +- .../ko/verbalizers/decimal.py | 21 +- .../ko/verbalizers/electronic.py | 19 +- .../ko/verbalizers/fraction.py | 22 +- .../ko/verbalizers/measure.py | 13 +- .../ko/verbalizers/money.py | 11 +- .../ko/verbalizers/ordinal.py | 12 +- .../ko/verbalizers/telephone.py | 15 +- .../ko/verbalizers/time.py | 58 +- .../ko/verbalizers/verbalize.py | 3 +- .../ko/verbalizers/verbalize_final.py | 9 +- .../ko/verbalizers/whitelist.py | 14 +- .../ko/verbalizers/word.py | 18 +- .../inverse_text_normalization/pt/__init__.py | 11 +- .../pt/data/__init__.py | 1 - .../pt/data/electronic/__init__.py | 1 - .../pt/data/numbers/__init__.py | 1 - .../pt/data/ordinals/__init__.py | 1 - .../pt/data/time/__init__.py | 1 - .../pt/taggers/__init__.py | 1 - .../pt/taggers/cardinal.py | 74 +- .../pt/taggers/date.py | 25 +- .../pt/taggers/decimal.py | 45 +- .../pt/taggers/electronic.py | 34 +- .../pt/taggers/measure.py | 50 +- .../pt/taggers/money.py | 48 +- .../pt/taggers/ordinal.py | 33 +- .../pt/taggers/punctuation.py | 7 +- .../pt/taggers/telephone.py | 17 +- .../pt/taggers/time.py | 96 +- .../pt/taggers/tokenize_and_classify.py | 13 +- .../pt/taggers/whitelist.py | 5 +- .../pt/taggers/word.py | 5 +- .../inverse_text_normalization/pt/utils.py | 5 +- .../pt/verbalizers/__init__.py | 1 - .../pt/verbalizers/cardinal.py | 17 +- .../pt/verbalizers/date.py | 21 +- .../pt/verbalizers/decimal.py | 27 +- .../pt/verbalizers/electronic.py | 21 +- .../pt/verbalizers/measure.py | 9 +- .../pt/verbalizers/money.py | 14 +- .../pt/verbalizers/ordinal.py | 17 +- .../pt/verbalizers/telephone.py | 9 +- .../pt/verbalizers/time.py | 23 +- .../pt/verbalizers/verbalize.py | 3 - .../pt/verbalizers/verbalize_final.py | 9 +- .../pt/verbalizers/whitelist.py | 14 +- .../pt/verbalizers/word.py | 20 +- .../inverse_text_normalization/ru/__init__.py | 2 - .../ru/taggers/cardinal.py | 8 +- .../ru/taggers/date.py | 4 +- .../ru/taggers/decimals.py | 26 +- .../ru/taggers/electronic.py | 4 +- .../ru/taggers/measure.py | 4 +- .../ru/taggers/money.py | 4 +- .../ru/taggers/ordinal.py | 4 +- .../ru/taggers/telephone.py | 4 +- .../ru/taggers/time.py | 22 +- .../ru/taggers/tokenize_and_classify.py | 23 +- .../ru/taggers/whitelist.py | 4 +- .../ru/verbalizers/cardinal.py | 16 +- .../ru/verbalizers/date.py | 3 +- .../ru/verbalizers/decimal.py | 16 +- .../ru/verbalizers/electronic.py | 5 +- .../ru/verbalizers/measure.py | 11 +- .../ru/verbalizers/money.py | 7 +- .../ru/verbalizers/ordinal.py | 3 +- .../ru/verbalizers/telephone.py | 7 +- .../ru/verbalizers/time.py | 20 +- .../ru/verbalizers/verbalize.py | 3 +- .../ru/verbalizers/verbalize_final.py | 9 +- .../run_evaluate.py | 49 +- .../inverse_text_normalization/tl/__init__.py | 9 +- .../tl/clean_eval_data.py | 115 +- .../tl/graph_utils.py | 61 +- .../tl/taggers/cardinal.py | 92 +- .../tl/taggers/date.py | 31 +- .../tl/taggers/decimal.py | 48 +- .../tl/taggers/electronic.py | 58 +- .../tl/taggers/fraction.py | 1 - .../tl/taggers/measure.py | 40 +- .../tl/taggers/money.py | 41 +- .../tl/taggers/ordinal.py | 3 +- .../tl/taggers/punctuation.py | 5 +- .../tl/taggers/telephone.py | 71 +- .../tl/taggers/time.py | 51 +- .../tl/taggers/tokenize_and_classify.py | 13 +- .../tl/taggers/whitelist.py | 3 +- .../tl/taggers/word.py | 3 +- .../inverse_text_normalization/tl/utils.py | 7 +- .../tl/verbalizers/cardinal.py | 16 +- .../tl/verbalizers/date.py | 25 +- .../tl/verbalizers/decimal.py | 22 +- .../tl/verbalizers/electronic.py | 19 +- .../tl/verbalizers/fraction.py | 3 +- .../tl/verbalizers/measure.py | 13 +- .../tl/verbalizers/money.py | 11 +- .../tl/verbalizers/ordinal.py | 12 +- .../tl/verbalizers/telephone.py | 15 +- .../tl/verbalizers/time.py | 21 +- .../tl/verbalizers/verbalize.py | 3 +- .../tl/verbalizers/verbalize_final.py | 9 +- .../tl/verbalizers/whitelist.py | 14 +- .../tl/verbalizers/word.py | 18 +- .../inverse_text_normalization/vi/__init__.py | 10 +- .../vi/data/math/__init__.py | 2 - .../vi/graph_utils.py | 13 +- .../vi/taggers/cardinal.py | 35 +- .../vi/taggers/date.py | 45 +- .../vi/taggers/decimal.py | 31 +- .../vi/taggers/electronic.py | 25 +- .../vi/taggers/fraction.py | 19 +- .../vi/taggers/measure.py | 18 +- .../vi/taggers/money.py | 18 +- .../vi/taggers/ordinal.py | 9 +- .../vi/taggers/punctuation.py | 1 - .../vi/taggers/telephone.py | 5 +- .../vi/taggers/time.py | 6 +- .../vi/taggers/tokenize_and_classify.py | 11 +- .../vi/taggers/whitelist.py | 1 - .../vi/taggers/word.py | 1 - .../inverse_text_normalization/vi/utils.py | 2 - .../vi/verbalizers/cardinal.py | 8 +- .../vi/verbalizers/date.py | 14 +- .../vi/verbalizers/decimal.py | 8 +- .../vi/verbalizers/electronic.py | 8 +- .../vi/verbalizers/fraction.py | 11 +- .../vi/verbalizers/measure.py | 1 - .../vi/verbalizers/money.py | 7 +- .../vi/verbalizers/ordinal.py | 7 +- .../vi/verbalizers/telephone.py | 7 +- .../vi/verbalizers/time.py | 15 +- .../vi/verbalizers/verbalize.py | 1 - .../vi/verbalizers/verbalize_final.py | 7 +- .../vi/verbalizers/whitelist.py | 4 +- .../vi/verbalizers/word.py | 11 +- .../inverse_text_normalization/zh/__init__.py | 10 +- .../zh/graph_utils.py | 64 +- .../zh/taggers/cardinal.py | 110 +- .../zh/taggers/date.py | 27 +- .../zh/taggers/decimal.py | 15 +- .../zh/taggers/electronic.py | 59 +- .../zh/taggers/fraction.py | 21 +- .../zh/taggers/measure.py | 31 +- .../zh/taggers/money.py | 13 +- .../zh/taggers/punctuation.py | 5 +- .../zh/taggers/telephone.py | 47 +- .../zh/taggers/time.py | 17 +- .../zh/taggers/tokenize_and_classify.py | 12 +- .../zh/taggers/whitelist.py | 3 +- .../zh/taggers/word.py | 3 +- .../inverse_text_normalization/zh/utils.py | 19 +- .../zh/verbalizers/cardinal.py | 17 +- .../zh/verbalizers/date.py | 40 +- .../zh/verbalizers/decimal.py | 22 +- .../zh/verbalizers/electronic.py | 19 +- .../zh/verbalizers/fraction.py | 20 +- .../zh/verbalizers/measure.py | 13 +- .../zh/verbalizers/money.py | 11 +- .../zh/verbalizers/telephone.py | 15 +- .../zh/verbalizers/time.py | 27 +- .../zh/verbalizers/verbalize.py | 3 +- .../zh/verbalizers/verbalize_final.py | 9 +- .../zh/verbalizers/whitelist.py | 15 +- .../zh/verbalizers/word.py | 18 +- fun_text_processing/num2words/.coveragerc | 5 - fun_text_processing/num2words/.gitignore | 7 - fun_text_processing/num2words/.travis.yml | 20 - fun_text_processing/num2words/MANIFEST.in | 1 - .../num2words/docker-compose.yml | 7 - .../num2words/num2words/__init__.py | 147 -- .../num2words/num2words/base.py | 300 ---- .../num2words/num2words/compat.py | 14 - .../num2words/num2words/currency.py | 35 - .../num2words/num2words/lang_AR.py | 352 ---- .../num2words/num2words/lang_AZ.py | 89 - .../num2words/num2words/lang_BG.py | 91 - .../num2words/num2words/lang_BN.py | 108 -- .../num2words/num2words/lang_CA.py | 89 - .../num2words/num2words/lang_CZ.py | 142 -- .../num2words/num2words/lang_DE.py | 142 -- .../num2words/num2words/lang_DK.py | 117 -- .../num2words/num2words/lang_EL.py | 115 -- .../num2words/num2words/lang_EN.py | 95 - .../num2words/num2words/lang_EN_IN.py | 11 - .../num2words/num2words/lang_ES.py | 351 ---- .../num2words/num2words/lang_ES_CO.py | 14 - .../num2words/num2words/lang_ES_NI.py | 18 - .../num2words/num2words/lang_ES_VE.py | 17 - .../num2words/num2words/lang_EU.py | 76 - .../num2words/num2words/lang_FI.py | 721 -------- .../num2words/num2words/lang_FR.py | 93 - .../num2words/num2words/lang_FR_BE.py | 34 - .../num2words/num2words/lang_FR_CH.py | 33 - .../num2words/num2words/lang_FR_DZ.py | 18 - .../num2words/num2words/lang_HE.py | 144 -- .../num2words/num2words/lang_HU.py | 150 -- .../num2words/num2words/lang_ID.py | 189 -- .../num2words/num2words/lang_IT.py | 208 --- .../num2words/num2words/lang_IW.py | 91 - .../num2words/num2words/lang_JA.py | 574 ------ .../num2words/num2words/lang_KH.py | 136 -- .../num2words/num2words/lang_KN.py | 167 -- .../num2words/num2words/lang_KO.py | 133 -- .../num2words/num2words/lang_KZ.py | 118 -- .../num2words/num2words/lang_LO.py | 110 -- .../num2words/num2words/lang_LT.py | 167 -- .../num2words/num2words/lang_LV.py | 174 -- .../num2words/num2words/lang_MN.py | 98 - .../num2words/num2words/lang_MY.py | 97 - .../num2words/num2words/lang_NE.py | 250 --- .../num2words/num2words/lang_NL.py | 140 -- .../num2words/num2words/lang_NO.py | 89 - .../num2words/num2words/lang_PL.py | 245 --- .../num2words/num2words/lang_PT.py | 226 --- .../num2words/num2words/lang_PT_BR.py | 98 - .../num2words/num2words/lang_RO.py | 158 -- .../num2words/num2words/lang_RU.py | 231 --- .../num2words/num2words/lang_SI.py | 108 -- .../num2words/num2words/lang_SL.py | 169 -- .../num2words/num2words/lang_SR.py | 204 --- .../num2words/num2words/lang_SV.py | 102 -- .../num2words/num2words/lang_SW.py | 108 -- .../num2words/num2words/lang_TE.py | 169 -- .../num2words/num2words/lang_TH.py | 168 -- .../num2words/num2words/lang_TR.py | 846 --------- .../num2words/num2words/lang_UK.py | 837 --------- .../num2words/num2words/lang_UR.py | 232 --- .../num2words/num2words/lang_VI.py | 86 - .../num2words/num2words/utils.py | 20 - .../num2words/requirements-test.txt | 6 - fun_text_processing/num2words/setup.py | 64 - fun_text_processing/num2words/tox.ini | 32 - .../text_normalization/data_loader_utils.py | 83 +- .../text_normalization/de/__init__.py | 2 - .../text_normalization/de/data/__init__.py | 2 - .../de/data/electronic/__init__.py | 2 - .../de/data/measure/__init__.py | 2 - .../de/data/money/__init__.py | 2 - .../de/data/months/__init__.py | 2 - .../de/data/numbers/__init__.py | 2 - .../de/data/ordinals/__init__.py | 2 - .../de/data/time/__init__.py | 2 - .../text_normalization/de/taggers/__init__.py | 2 - .../text_normalization/de/taggers/cardinal.py | 69 +- .../text_normalization/de/taggers/date.py | 22 +- .../text_normalization/de/taggers/decimal.py | 37 +- .../de/taggers/electronic.py | 32 +- .../text_normalization/de/taggers/fraction.py | 11 +- .../text_normalization/de/taggers/measure.py | 77 +- .../text_normalization/de/taggers/money.py | 61 +- .../text_normalization/de/taggers/ordinal.py | 9 +- .../de/taggers/telephone.py | 25 +- .../text_normalization/de/taggers/time.py | 38 +- .../de/taggers/tokenize_and_classify.py | 30 +- .../de/taggers/whitelist.py | 10 +- .../text_normalization/de/taggers/word.py | 4 +- .../text_normalization/de/utils.py | 5 +- .../de/verbalizers/__init__.py | 2 - .../de/verbalizers/cardinal.py | 5 +- .../text_normalization/de/verbalizers/date.py | 21 +- .../de/verbalizers/decimal.py | 28 +- .../de/verbalizers/electronic.py | 17 +- .../de/verbalizers/fraction.py | 32 +- .../de/verbalizers/measure.py | 9 +- .../de/verbalizers/money.py | 33 +- .../de/verbalizers/ordinal.py | 16 +- .../de/verbalizers/telephone.py | 21 +- .../text_normalization/de/verbalizers/time.py | 47 +- .../de/verbalizers/verbalize.py | 7 +- .../de/verbalizers/verbalize_final.py | 13 +- .../text_normalization/en/__init__.py | 2 - .../text_normalization/en/clean_eval_data.py | 116 +- .../en/data/address/__init__.py | 2 - .../en/data/electronic/__init__.py | 1 - .../en/data/telephone/__init__.py | 1 - .../text_normalization/en/graph_utils.py | 61 +- .../en/taggers/abbreviation.py | 12 +- .../text_normalization/en/taggers/cardinal.py | 57 +- .../text_normalization/en/taggers/date.py | 108 +- .../text_normalization/en/taggers/decimal.py | 69 +- .../en/taggers/electronic.py | 46 +- .../text_normalization/en/taggers/fraction.py | 16 +- .../text_normalization/en/taggers/measure.py | 110 +- .../text_normalization/en/taggers/money.py | 87 +- .../text_normalization/en/taggers/ordinal.py | 7 +- .../en/taggers/punctuation.py | 17 +- .../text_normalization/en/taggers/range.py | 42 +- .../text_normalization/en/taggers/roman.py | 34 +- .../text_normalization/en/taggers/serial.py | 34 +- .../en/taggers/telephone.py | 73 +- .../text_normalization/en/taggers/time.py | 31 +- .../en/taggers/tokenize_and_classify.py | 97 +- .../en/taggers/tokenize_and_classify_lm.py | 62 +- .../tokenize_and_classify_with_audio.py | 59 +- .../en/taggers/whitelist.py | 40 +- .../text_normalization/en/taggers/word.py | 16 +- .../text_normalization/en/utils.py | 8 +- .../en/verbalizers/abbreviation.py | 4 +- .../en/verbalizers/cardinal.py | 14 +- .../text_normalization/en/verbalizers/date.py | 33 +- .../en/verbalizers/decimal.py | 21 +- .../en/verbalizers/electronic.py | 32 +- .../en/verbalizers/fraction.py | 46 +- .../en/verbalizers/measure.py | 49 +- .../en/verbalizers/money.py | 22 +- .../en/verbalizers/ordinal.py | 19 +- .../en/verbalizers/post_processing.py | 101 +- .../en/verbalizers/roman.py | 24 +- .../en/verbalizers/telephone.py | 21 +- .../text_normalization/en/verbalizers/time.py | 22 +- .../en/verbalizers/verbalize.py | 8 +- .../en/verbalizers/verbalize_final.py | 12 +- .../en/verbalizers/whitelist.py | 14 +- .../text_normalization/en/verbalizers/word.py | 18 +- .../text_normalization/es/__init__.py | 3 - .../text_normalization/es/data/__init__.py | 1 - .../es/data/dates/__init__.py | 1 - .../es/data/electronic/__init__.py | 1 - .../es/data/fractions/__init__.py | 1 - .../es/data/measures/__init__.py | 1 - .../es/data/money/__init__.py | 1 - .../es/data/numbers/__init__.py | 1 - .../es/data/ordinals/__init__.py | 1 - .../es/data/roman/__init__.py | 1 - .../es/data/time/__init__.py | 1 - .../text_normalization/es/graph_utils.py | 42 +- .../text_normalization/es/taggers/__init__.py | 1 - .../text_normalization/es/taggers/cardinal.py | 66 +- .../text_normalization/es/taggers/date.py | 33 +- .../text_normalization/es/taggers/decimals.py | 36 +- .../es/taggers/electronic.py | 31 +- .../text_normalization/es/taggers/fraction.py | 75 +- .../text_normalization/es/taggers/measure.py | 70 +- .../text_normalization/es/taggers/money.py | 85 +- .../text_normalization/es/taggers/ordinal.py | 43 +- .../es/taggers/telephone.py | 8 +- .../text_normalization/es/taggers/time.py | 79 +- .../es/taggers/tokenize_and_classify.py | 30 +- .../es/taggers/whitelist.py | 11 +- .../text_normalization/es/taggers/word.py | 4 +- .../text_normalization/es/utils.py | 5 +- .../es/verbalizers/__init__.py | 1 - .../es/verbalizers/cardinal.py | 18 +- .../text_normalization/es/verbalizers/date.py | 27 +- .../es/verbalizers/decimals.py | 44 +- .../es/verbalizers/electronic.py | 17 +- .../es/verbalizers/fraction.py | 70 +- .../es/verbalizers/measure.py | 41 +- .../es/verbalizers/money.py | 62 +- .../es/verbalizers/ordinal.py | 34 +- .../es/verbalizers/telephone.py | 8 +- .../text_normalization/es/verbalizers/time.py | 36 +- .../es/verbalizers/verbalize.py | 6 +- .../es/verbalizers/verbalize_final.py | 13 +- .../text_normalization/export_models.py | 84 +- .../text_normalization/normalize.py | 137 +- .../normalize_with_audio.py | 112 +- .../text_normalization/ru/alphabet.py | 11 +- .../text_normalization/ru/taggers/cardinal.py | 81 +- .../text_normalization/ru/taggers/date.py | 62 +- .../text_normalization/ru/taggers/decimals.py | 58 +- .../ru/taggers/electronic.py | 53 +- .../text_normalization/ru/taggers/measure.py | 94 +- .../text_normalization/ru/taggers/money.py | 43 +- .../ru/taggers/number_names.py | 60 +- .../text_normalization/ru/taggers/ordinal.py | 11 +- .../ru/taggers/telephone.py | 36 +- .../text_normalization/ru/taggers/time.py | 63 +- .../ru/taggers/tokenize_and_classify.py | 43 +- .../ru/taggers/whitelist.py | 8 +- .../text_normalization/ru/taggers/word.py | 4 +- .../text_normalization/ru/utils.py | 8 +- .../ru/verbalizers/cardinal.py | 11 +- .../text_normalization/ru/verbalizers/date.py | 4 +- .../ru/verbalizers/decimal.py | 23 +- .../ru/verbalizers/electronic.py | 4 +- .../ru/verbalizers/measure.py | 8 +- .../ru/verbalizers/money.py | 6 +- .../ru/verbalizers/ordinal.py | 4 +- .../ru/verbalizers/telephone.py | 8 +- .../text_normalization/ru/verbalizers/time.py | 24 +- .../ru/verbalizers/verbalize.py | 4 +- .../ru/verbalizers/verbalize_final.py | 13 +- .../text_normalization/run_evaluate.py | 49 +- .../text_normalization/token_parser.py | 20 +- .../text_normalization/zh/__init__.py | 1 - .../text_normalization/zh/data/__init__.py | 1 - .../zh/data/char/__init__.py | 1 - .../zh/data/date/__init__.py | 1 - .../zh/data/denylist/__init__.py | 1 - .../zh/data/erhua/__init__.py | 1 - .../zh/data/math/__init__.py | 1 - .../zh/data/measure/__init__.py | 1 - .../zh/data/money/__init__.py | 1 - .../zh/data/number/__init__.py | 1 - .../zh/data/time/__init__.py | 1 - .../zh/data/whitelist/__init__.py | 1 - .../text_normalization/zh/graph_utils.py | 22 +- .../text_normalization/zh/taggers/__init__.py | 1 - .../text_normalization/zh/taggers/cardinal.py | 75 +- .../text_normalization/zh/taggers/char.py | 10 +- .../text_normalization/zh/taggers/date.py | 67 +- .../text_normalization/zh/taggers/fraction.py | 18 +- .../zh/taggers/math_symbol.py | 16 +- .../text_normalization/zh/taggers/measure.py | 20 +- .../text_normalization/zh/taggers/money.py | 16 +- .../zh/taggers/preprocessor.py | 38 +- .../text_normalization/zh/taggers/time.py | 38 +- .../zh/taggers/tokenize_and_classify.py | 14 +- .../zh/taggers/whitelist.py | 14 +- .../text_normalization/zh/utils.py | 16 +- .../zh/verbalizers/__init__.py | 1 - .../zh/verbalizers/cardinal.py | 8 +- .../text_normalization/zh/verbalizers/char.py | 10 +- .../text_normalization/zh/verbalizers/date.py | 38 +- .../zh/verbalizers/fraction.py | 14 +- .../zh/verbalizers/math_symbol.py | 10 +- .../zh/verbalizers/measure.py | 26 +- .../zh/verbalizers/money.py | 27 +- .../zh/verbalizers/postprocessor.py | 42 +- .../text_normalization/zh/verbalizers/time.py | 50 +- .../zh/verbalizers/verbalize.py | 4 +- .../zh/verbalizers/verbalize_final.py | 33 +- .../zh/verbalizers/whitelist.py | 12 +- funasr/__init__.py | 6 +- funasr/auto/auto_frontend.py | 41 +- funasr/auto/auto_model.py | 277 +-- funasr/auto/auto_tokenizer.py | 13 +- funasr/bin/compute_audio_cmvn.py | 77 +- funasr/bin/export.py | 27 +- funasr/bin/inference.py | 10 +- funasr/bin/tokenize_text.py | 10 +- funasr/bin/train.py | 115 +- funasr/datasets/audio_datasets/datasets.py | 175 +- .../audio_datasets/espnet_samplers.py | 57 +- funasr/datasets/audio_datasets/index_ds.py | 71 +- funasr/datasets/audio_datasets/jsonl2scp.py | 25 +- .../datasets/audio_datasets/preprocessor.py | 60 +- funasr/datasets/audio_datasets/samplers.py | 229 +-- funasr/datasets/audio_datasets/scp2jsonl.py | 61 +- funasr/datasets/dataloader_entry.py | 189 +- .../large_datasets/build_dataloader.py | 32 +- funasr/datasets/large_datasets/collate_fn.py | 54 +- .../large_datasets/datapipes/batch.py | 14 +- .../large_datasets/datapipes/filter.py | 5 +- .../datasets/large_datasets/datapipes/map.py | 4 +- funasr/datasets/large_datasets/dataset.py | 145 +- .../datasets/large_datasets/utils/clipping.py | 8 +- .../datasets/large_datasets/utils/filter.py | 21 +- .../large_datasets/utils/hotword_utils.py | 71 +- .../large_datasets/utils/low_frame_rate.py | 4 +- .../datasets/large_datasets/utils/padding.py | 40 +- .../datasets/large_datasets/utils/tokenize.py | 30 +- funasr/datasets/llm_datasets/datasets.py | 302 ++-- funasr/datasets/llm_datasets/preprocessor.py | 31 +- .../llm_datasets_qwenaudio/datasets.py | 105 +- .../datasets/llm_datasets_vicuna/datasets.py | 105 +- .../datasets/sense_voice_datasets/datasets.py | 83 +- funasr/download/download_dataset_from_hub.py | 19 +- funasr/download/download_from_hub.py | 68 +- funasr/download/file.py | 96 +- funasr/download/name_maps_from_hub.py | 6 +- funasr/download/runtime_sdk_download_tool.py | 39 +- funasr/frontends/default.py | 126 +- funasr/frontends/eend_ola_feature.py | 24 +- funasr/frontends/fused.py | 14 +- funasr/frontends/s3prl.py | 20 +- funasr/frontends/utils/beamformer.py | 4 +- funasr/frontends/utils/complex_utils.py | 29 +- funasr/frontends/utils/dnn_beamformer.py | 25 +- funasr/frontends/utils/feature_transform.py | 8 +- funasr/frontends/utils/log_mel.py | 8 +- funasr/frontends/utils/mask_estimator.py | 4 +- funasr/frontends/utils/stft.py | 25 +- funasr/frontends/wav_frontend.py | 344 ++-- funasr/frontends/whisper_frontend.py | 37 +- funasr/frontends/windowing.py | 2 +- funasr/losses/label_smoothing_loss.py | 10 +- funasr/metrics/common.py | 14 +- funasr/metrics/compute_acc.py | 8 +- funasr/metrics/compute_eer.py | 4 +- funasr/metrics/compute_min_dcf.py | 101 +- funasr/metrics/wer.py | 377 ++-- funasr/models/bat/model.py | 2 +- .../models/bicif_paraformer/cif_predictor.py | 270 +-- funasr/models/bicif_paraformer/export_meta.py | 39 +- funasr/models/bicif_paraformer/model.py | 248 +-- funasr/models/branchformer/encoder.py | 50 +- funasr/models/branchformer/fastformer.py | 30 +- funasr/models/branchformer/model.py | 1 + funasr/models/campplus/cluster_backend.py | 36 +- funasr/models/campplus/components.py | 244 ++- funasr/models/campplus/model.py | 135 +- funasr/models/campplus/utils.py | 168 +- funasr/models/conformer/encoder.py | 167 +- funasr/models/conformer/model.py | 2 +- funasr/models/conformer_rwkv/decoder.py | 180 +- funasr/models/conformer_rwkv/model.py | 2 +- .../models/contextual_paraformer/decoder.py | 128 +- .../contextual_paraformer/export_meta.py | 73 +- funasr/models/contextual_paraformer/model.py | 292 +-- funasr/models/ct_transformer/export_meta.py | 34 +- funasr/models/ct_transformer/model.py | 152 +- funasr/models/ct_transformer/utils.py | 29 +- .../ct_transformer_streaming/attention.py | 2 - .../ct_transformer_streaming/encoder.py | 141 +- .../ct_transformer_streaming/export_meta.py | 47 +- .../models/ct_transformer_streaming/model.py | 70 +- funasr/models/ctc/ctc.py | 4 +- funasr/models/data2vec/data2vec.py | 36 +- funasr/models/data2vec/data2vec_encoder.py | 201 +-- funasr/models/data2vec/data_utils.py | 32 +- funasr/models/data2vec/ema_module.py | 12 +- funasr/models/data2vec/multihead_attention.py | 232 +-- funasr/models/data2vec/quant_noise.py | 20 +- funasr/models/data2vec/utils.py | 4 +- funasr/models/data2vec/wav2vec2.py | 116 +- funasr/models/e_branchformer/encoder.py | 14 +- funasr/models/e_branchformer/model.py | 1 + funasr/models/eend/e2e_diar_eend_ola.py | 160 +- funasr/models/eend/eend_ola_dataloader.py | 24 +- funasr/models/eend/encoder.py | 51 +- .../models/eend/encoder_decoder_attractor.py | 30 +- funasr/models/eend/utils/feature.py | 120 +- funasr/models/eend/utils/kaldi_data.py | 80 +- funasr/models/eend/utils/losses.py | 19 +- funasr/models/eend/utils/power.py | 39 +- funasr/models/eend/utils/report.py | 125 +- funasr/models/emotion2vec/audio.py | 13 +- funasr/models/emotion2vec/base.py | 70 +- funasr/models/emotion2vec/fairseq_modules.py | 264 ++- funasr/models/emotion2vec/model.py | 61 +- funasr/models/emotion2vec/modules.py | 14 +- funasr/models/emotion2vec/timm_modules.py | 42 +- funasr/models/eres2net/eres2net.py | 167 +- funasr/models/eres2net/eres2net_aug.py | 103 +- funasr/models/eres2net/fusion.py | 3 +- funasr/models/fsmn_vad_streaming/encoder.py | 147 +- .../models/fsmn_vad_streaming/export_meta.py | 30 +- funasr/models/fsmn_vad_streaming/model.py | 1595 +++++++++-------- funasr/models/language_model/rnn/argument.py | 16 +- .../models/language_model/rnn/attentions.py | 162 +- funasr/models/language_model/rnn/decoders.py | 245 +-- funasr/models/language_model/rnn/encoders.py | 24 +- funasr/models/language_model/seq_rnn_lm.py | 9 +- .../language_model/transformer_encoder.py | 36 +- .../models/language_model/transformer_lm.py | 7 +- funasr/models/lcbnet/attention.py | 3 +- funasr/models/lcbnet/encoder.py | 120 +- funasr/models/lcbnet/model.py | 181 +- funasr/models/llm_asr/adaptor.py | 19 +- funasr/models/llm_asr/model.py | 131 +- funasr/models/llm_asr_nar/adaptor.py | 3 +- funasr/models/llm_asr_nar/model.py | 351 ++-- funasr/models/lora/layers.py | 89 +- funasr/models/lora/utils.py | 35 +- funasr/models/mfcca/e2e_asr_mfcca.py | 118 +- funasr/models/mfcca/encoder_layer_mfcca.py | 67 +- funasr/models/mfcca/mfcca_encoder.py | 92 +- funasr/models/monotonic_aligner/model.py | 128 +- funasr/models/mossformer/e2e_ss.py | 24 +- funasr/models/mossformer/mossformer.py | 206 ++- .../models/mossformer/mossformer_decoder.py | 5 +- .../models/mossformer/mossformer_encoder.py | 123 +- funasr/models/normalize/global_mvn.py | 3 +- funasr/models/normalize/utterance_mvn.py | 1 + funasr/models/paraformer/cif_predictor.py | 256 ++- funasr/models/paraformer/decoder.py | 405 ++--- funasr/models/paraformer/export_meta.py | 78 +- funasr/models/paraformer/model.py | 259 +-- funasr/models/paraformer/search.py | 29 +- .../paraformer_streaming/export_meta.py | 98 +- funasr/models/paraformer_streaming/model.py | 503 +++--- funasr/models/qwen_audio/audio.py | 118 +- funasr/models/qwen_audio/model.py | 100 +- funasr/models/rwkv_bat/rwkv.py | 1 + funasr/models/rwkv_bat/rwkv_attention.py | 53 +- funasr/models/rwkv_bat/rwkv_encoder.py | 18 +- funasr/models/rwkv_bat/rwkv_feed_forward.py | 8 +- funasr/models/rwkv_bat/rwkv_subsampling.py | 128 +- funasr/models/sa_asr/attention.py | 13 +- funasr/models/sa_asr/beam_search_sa_asr.py | 65 +- funasr/models/sa_asr/e2e_sa_asr.py | 185 +- funasr/models/sa_asr/transformer_decoder.py | 353 ++-- funasr/models/sanm/attention.py | 181 +- funasr/models/sanm/decoder.py | 104 +- funasr/models/sanm/encoder.py | 124 +- funasr/models/sanm/model.py | 1 + funasr/models/sanm/multihead_att.py | 24 +- .../models/sanm/positionwise_feed_forward.py | 1 - funasr/models/scama/beam_search.py | 121 +- funasr/models/scama/chunk_utilis.py | 383 ++-- funasr/models/scama/decoder.py | 140 +- funasr/models/scama/encoder.py | 157 +- funasr/models/scama/model.py | 433 +++-- funasr/models/scama/utils.py | 39 +- funasr/models/seaco_paraformer/export_meta.py | 274 +-- funasr/models/seaco_paraformer/model.py | 403 +++-- funasr/models/sense_voice/decoder.py | 509 +++--- funasr/models/sense_voice/encoder.py | 100 +- funasr/models/sense_voice/model.py | 559 +++--- funasr/models/sense_voice/rwkv_v6.py | 784 ++++---- .../sense_voice/whisper_lib/__init__.py | 8 +- .../models/sense_voice/whisper_lib/audio.py | 6 +- .../sense_voice/whisper_lib/decoding.py | 199 +- .../models/sense_voice/whisper_lib/model.py | 73 +- .../whisper_lib/normalizers/basic.py | 33 +- .../whisper_lib/normalizers/english.py | 10 +- .../models/sense_voice/whisper_lib/timing.py | 46 +- .../sense_voice/whisper_lib/tokenizer.py | 20 +- .../sense_voice/whisper_lib/transcribe.py | 62 +- .../sense_voice/whisper_lib/triton_ops.py | 13 +- .../models/sense_voice/whisper_lib/utils.py | 63 +- funasr/models/sond/attention.py | 34 +- funasr/models/sond/e2e_diar_sond.py | 158 +- funasr/models/sond/encoder/ci_scorers.py | 12 +- funasr/models/sond/encoder/conv_encoder.py | 65 +- .../models/sond/encoder/ecapa_tdnn_encoder.py | 64 +- funasr/models/sond/encoder/fsmn_encoder.py | 79 +- .../models/sond/encoder/resnet34_encoder.py | 235 +-- .../sond/encoder/self_attention_encoder.py | 57 +- funasr/models/sond/label_aggregation.py | 10 +- funasr/models/sond/pooling/pooling_layers.py | 25 +- .../models/sond/pooling/statistic_pooling.py | 41 +- funasr/models/sond/sv_decoder.py | 20 +- funasr/models/specaug/mask_along_axis.py | 10 +- funasr/models/specaug/profileaug.py | 67 +- funasr/models/specaug/specaug.py | 15 +- funasr/models/specaug/time_warp.py | 1 + .../transducer/beam_search_transducer.py | 34 +- funasr/models/transducer/joint_network.py | 4 +- funasr/models/transducer/model.py | 203 ++- funasr/models/transducer/rnn_decoder.py | 37 +- funasr/models/transducer/rnnt_decoder.py | 27 +- funasr/models/transformer/attention.py | 80 +- funasr/models/transformer/decoder.py | 299 ++- funasr/models/transformer/embedding.py | 88 +- funasr/models/transformer/encoder.py | 81 +- funasr/models/transformer/layer_norm.py | 25 +- funasr/models/transformer/model.py | 183 +- .../transformer/positionwise_feed_forward.py | 21 +- funasr/models/transformer/scorers/ctc.py | 4 +- .../transformer/scorers/ctc_prefix_score.py | 42 +- .../transformer/scorers/length_bonus.py | 5 +- .../transformer/scorers/scorer_interface.py | 4 +- funasr/models/transformer/search.py | 27 +- .../models/transformer/utils/dynamic_conv.py | 6 +- .../transformer/utils/dynamic_conv2d.py | 10 +- funasr/models/transformer/utils/lightconv.py | 8 +- .../models/transformer/utils/lightconv2d.py | 14 +- funasr/models/transformer/utils/mask.py | 6 +- .../transformer/utils/multi_layer_conv.py | 2 +- funasr/models/transformer/utils/nets_utils.py | 45 +- .../models/transformer/utils/subsampling.py | 117 +- funasr/models/transformer/utils/vgg2l.py | 8 +- funasr/models/uniasr/beam_search.py | 55 +- funasr/models/uniasr/model.py | 433 +++-- funasr/models/whisper/model.py | 55 +- funasr/models/whisper_lid/decoder.py | 31 +- funasr/models/whisper_lid/encoder.py | 42 +- funasr/models/whisper_lid/eres2net/ResNet.py | 171 +- funasr/models/whisper_lid/eres2net/fusion.py | 3 +- .../whisper_lid/eres2net/pooling_layers.py | 29 +- .../models/whisper_lid/eres2net/simple_avg.py | 3 +- funasr/models/whisper_lid/lid_predictor.py | 44 +- funasr/models/whisper_lid/model.py | 307 ++-- funasr/models/xvector/e2e_sv.py | 61 +- funasr/optimizers/__init__.py | 2 +- funasr/optimizers/fairseq_adam.py | 22 +- funasr/register.py | 37 +- funasr/schedulers/lambdalr_cus.py | 13 +- funasr/schedulers/noam_lr.py | 5 +- funasr/schedulers/tri_stage_scheduler.py | 16 +- funasr/schedulers/warmup_lr.py | 5 +- funasr/tokenizer/abs_tokenizer.py | 49 +- funasr/tokenizer/build_tokenizer.py | 4 +- funasr/tokenizer/char_tokenizer.py | 191 +- funasr/tokenizer/cleaner.py | 1 + funasr/tokenizer/hf_tokenizer.py | 19 +- funasr/tokenizer/korean_cleaner.py | 7 +- funasr/tokenizer/phoneme_tokenizer.py | 4 +- funasr/tokenizer/sentencepiece_tokenizer.py | 7 +- funasr/tokenizer/token_id_converter.py | 4 +- funasr/tokenizer/whisper_tokenizer.py | 75 +- funasr/tokenizer/word_tokenizer.py | 3 +- funasr/train_utils/average_nbest_models.py | 18 +- funasr/train_utils/device_funcs.py | 13 +- funasr/train_utils/initialize.py | 6 +- funasr/train_utils/load_pretrained_model.py | 175 +- funasr/train_utils/model_summary.py | 4 +- funasr/train_utils/recursive_op.py | 1 + funasr/train_utils/trainer.py | 407 +++-- funasr/utils/datadir_writer.py | 4 +- funasr/utils/export_utils.py | 134 +- funasr/utils/install_model_requirements.py | 9 +- funasr/utils/load_utils.py | 100 +- funasr/utils/misc.py | 37 +- funasr/utils/postprocess_utils.py | 129 +- funasr/utils/speaker_utils.py | 108 +- funasr/utils/timestamp_tools.py | 145 +- funasr/utils/torch_function.py | 20 +- funasr/utils/vad_utils.py | 23 +- runtime/html5/h5Server.py | 82 +- .../cmake/cmake_extension.py | 9 +- runtime/python/grpc/grpc_main_client.py | 125 +- runtime/python/http/client.py | 33 +- runtime/python/http/server.py | 160 +- runtime/python/libtorch/demo.py | 4 +- .../libtorch/funasr_torch/paraformer_bin.py | 120 +- .../funasr_torch/utils/compute_wer.py | 160 +- .../libtorch/funasr_torch/utils/frontend.py | 82 +- .../funasr_torch/utils/postprocess_utils.py | 90 +- .../funasr_torch/utils/timestamp_utils.py | 37 +- .../libtorch/funasr_torch/utils/utils.py | 30 +- runtime/python/libtorch/setup.py | 42 +- .../onnxruntime/demo_contextual_paraformer.py | 4 +- .../onnxruntime/demo_paraformer_offline.py | 2 +- .../onnxruntime/demo_paraformer_online.py | 15 +- .../python/onnxruntime/demo_punc_offline.py | 4 +- .../python/onnxruntime/demo_punc_online.py | 4 +- .../onnxruntime/demo_seaco_paraformer.py | 4 +- .../python/onnxruntime/demo_vad_offline.py | 4 +- runtime/python/onnxruntime/demo_vad_online.py | 15 +- .../python/onnxruntime/funasr_client_http.py | 27 +- .../onnxruntime/funasr_onnx/paraformer_bin.py | 272 +-- .../funasr_onnx/paraformer_online_bin.py | 153 +- .../onnxruntime/funasr_onnx/punc_bin.py | 168 +- .../onnxruntime/funasr_onnx/utils/e2e_vad.py | 260 ++- .../onnxruntime/funasr_onnx/utils/frontend.py | 202 ++- .../funasr_onnx/utils/postprocess_utils.py | 123 +- .../funasr_onnx/utils/timestamp_utils.py | 35 +- .../onnxruntime/funasr_onnx/utils/utils.py | 135 +- .../python/onnxruntime/funasr_onnx/vad_bin.py | 612 ++++--- .../python/onnxruntime/funasr_server_http.py | 31 +- runtime/python/onnxruntime/setup.py | 43 +- runtime/python/utils/compute_wer.py | 160 +- runtime/python/utils/proce_text.py | 39 +- runtime/python/utils/test_cer.py | 40 +- runtime/python/utils/test_rtf.py | 58 +- runtime/python/utils/test_rtf_gpu.py | 60 +- runtime/python/websocket/funasr_client_api.py | 211 +-- runtime/python/websocket/funasr_wss_client.py | 315 ++-- runtime/python/websocket/funasr_wss_server.py | 547 +++--- runtime/tools/fst/ctc_token_fst.py | 23 +- runtime/tools/fst/generate_lexicon.py | 16 +- .../client/decode_manifest_triton.py | 43 +- .../client/decode_manifest_triton_wo_cuts.py | 85 +- runtime/triton_gpu/client/speech_client.py | 4 +- .../feature_extractor/1/model.py | 150 +- .../scoring/1/model.py | 46 +- .../cif_search/1/model.py | 65 +- .../feature_extractor/1/model.py | 53 +- .../lfr_cmvn_pe/export_lfr_cmvn_pe_onnx.py | 87 +- setup.py | 32 +- tests/run_test.py | 18 +- tests/test_asr_inference_pipeline.py | 387 ++-- tests/test_asr_vad_punc_inference_pipeline.py | 13 +- tests/test_lm_pipeline.py | 6 +- tests/test_punctuation_pipeline.py | 8 +- tests/test_sv_inference_pipeline.py | 38 +- tests/test_tp_pipeline.py | 42 +- tests/test_vad_inference_pipeline.py | 66 +- 1015 files changed, 25663 insertions(+), 30644 deletions(-) create mode 100644 .pre-commit-config.yaml delete mode 100644 fun_text_processing/num2words/.coveragerc delete mode 100644 fun_text_processing/num2words/.gitignore delete mode 100644 fun_text_processing/num2words/.travis.yml delete mode 100644 fun_text_processing/num2words/MANIFEST.in delete mode 100644 fun_text_processing/num2words/docker-compose.yml delete mode 100644 fun_text_processing/num2words/num2words/__init__.py delete mode 100644 fun_text_processing/num2words/num2words/base.py delete mode 100644 fun_text_processing/num2words/num2words/compat.py delete mode 100644 fun_text_processing/num2words/num2words/currency.py delete mode 100644 fun_text_processing/num2words/num2words/lang_AR.py delete mode 100644 fun_text_processing/num2words/num2words/lang_AZ.py delete mode 100644 fun_text_processing/num2words/num2words/lang_BG.py delete mode 100644 fun_text_processing/num2words/num2words/lang_BN.py delete mode 100644 fun_text_processing/num2words/num2words/lang_CA.py delete mode 100644 fun_text_processing/num2words/num2words/lang_CZ.py delete mode 100644 fun_text_processing/num2words/num2words/lang_DE.py delete mode 100644 fun_text_processing/num2words/num2words/lang_DK.py delete mode 100644 fun_text_processing/num2words/num2words/lang_EL.py delete mode 100644 fun_text_processing/num2words/num2words/lang_EN.py delete mode 100644 fun_text_processing/num2words/num2words/lang_EN_IN.py delete mode 100644 fun_text_processing/num2words/num2words/lang_ES.py delete mode 100644 fun_text_processing/num2words/num2words/lang_ES_CO.py delete mode 100644 fun_text_processing/num2words/num2words/lang_ES_NI.py delete mode 100644 fun_text_processing/num2words/num2words/lang_ES_VE.py delete mode 100644 fun_text_processing/num2words/num2words/lang_EU.py delete mode 100644 fun_text_processing/num2words/num2words/lang_FI.py delete mode 100644 fun_text_processing/num2words/num2words/lang_FR.py delete mode 100644 fun_text_processing/num2words/num2words/lang_FR_BE.py delete mode 100644 fun_text_processing/num2words/num2words/lang_FR_CH.py delete mode 100644 fun_text_processing/num2words/num2words/lang_FR_DZ.py delete mode 100644 fun_text_processing/num2words/num2words/lang_HE.py delete mode 100644 fun_text_processing/num2words/num2words/lang_HU.py delete mode 100644 fun_text_processing/num2words/num2words/lang_ID.py delete mode 100644 fun_text_processing/num2words/num2words/lang_IT.py delete mode 100644 fun_text_processing/num2words/num2words/lang_IW.py delete mode 100644 fun_text_processing/num2words/num2words/lang_JA.py delete mode 100644 fun_text_processing/num2words/num2words/lang_KH.py delete mode 100644 fun_text_processing/num2words/num2words/lang_KN.py delete mode 100644 fun_text_processing/num2words/num2words/lang_KO.py delete mode 100644 fun_text_processing/num2words/num2words/lang_KZ.py delete mode 100644 fun_text_processing/num2words/num2words/lang_LO.py delete mode 100644 fun_text_processing/num2words/num2words/lang_LT.py delete mode 100644 fun_text_processing/num2words/num2words/lang_LV.py delete mode 100644 fun_text_processing/num2words/num2words/lang_MN.py delete mode 100644 fun_text_processing/num2words/num2words/lang_MY.py delete mode 100644 fun_text_processing/num2words/num2words/lang_NE.py delete mode 100644 fun_text_processing/num2words/num2words/lang_NL.py delete mode 100644 fun_text_processing/num2words/num2words/lang_NO.py delete mode 100644 fun_text_processing/num2words/num2words/lang_PL.py delete mode 100644 fun_text_processing/num2words/num2words/lang_PT.py delete mode 100644 fun_text_processing/num2words/num2words/lang_PT_BR.py delete mode 100644 fun_text_processing/num2words/num2words/lang_RO.py delete mode 100644 fun_text_processing/num2words/num2words/lang_RU.py delete mode 100644 fun_text_processing/num2words/num2words/lang_SI.py delete mode 100644 fun_text_processing/num2words/num2words/lang_SL.py delete mode 100644 fun_text_processing/num2words/num2words/lang_SR.py delete mode 100644 fun_text_processing/num2words/num2words/lang_SV.py delete mode 100644 fun_text_processing/num2words/num2words/lang_SW.py delete mode 100644 fun_text_processing/num2words/num2words/lang_TE.py delete mode 100644 fun_text_processing/num2words/num2words/lang_TH.py delete mode 100644 fun_text_processing/num2words/num2words/lang_TR.py delete mode 100644 fun_text_processing/num2words/num2words/lang_UK.py delete mode 100644 fun_text_processing/num2words/num2words/lang_UR.py delete mode 100644 fun_text_processing/num2words/num2words/lang_VI.py delete mode 100644 fun_text_processing/num2words/num2words/utils.py delete mode 100644 fun_text_processing/num2words/requirements-test.txt delete mode 100644 fun_text_processing/num2words/setup.py delete mode 100644 fun_text_processing/num2words/tox.ini diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 000000000..d0c684ab4 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,6 @@ +repos: + - repo: https://github.com/psf/black + rev: 24.4.0 + hooks: + - id: black + args: ['--line-length=100'] # 示例参数，black默认使用4个空格缩进 diff --git a/docs/conf.py b/docs/conf.py index 9a9012b80..a1bdcd7c4 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -17,9 +17,9 @@ # -- Project information ----------------------------------------------------- -project = 'FunASR' -copyright = '2022, Speech Lab, Alibaba Group' -author = 'Speech Lab, Alibaba Grou' +project = "FunASR" +copyright = "2022, Speech Lab, Alibaba Group" +author = "Speech Lab, Alibaba Group" # -- General configuration --------------------------------------------------- @@ -30,18 +30,18 @@ author = 'Speech Lab, Alibaba Grou' extensions = [ "nbsphinx", "sphinx.ext.autodoc", - 'sphinx.ext.napoleon', - 'sphinx.ext.viewcode', + "sphinx.ext.napoleon", + "sphinx.ext.viewcode", "sphinx.ext.mathjax", "sphinx.ext.todo", # "sphinxarg.ext", "sphinx_markdown_tables", - 'recommonmark', - 'sphinx_rtd_theme', + "recommonmark", + "sphinx_rtd_theme", ] # Add any paths that contain templates here, relative to this directory. -templates_path = ['_templates'] +templates_path = ["_templates"] source_suffix = [".rst", ".md"] @@ -64,4 +64,4 @@ html_theme = "sphinx_rtd_theme" # Add any paths that contain custom static files (such as style sheets) here, # relative to this directory. They are copied after the builtin static files, # so a file named "default.css" will overwrite the builtin "default.css". -# html_static_path = ['_static'] \ No newline at end of file +# html_static_path = ['_static'] diff --git a/docs/m2met2/conf.py b/docs/m2met2/conf.py index a9f738249..6edc1b458 100644 --- a/docs/m2met2/conf.py +++ b/docs/m2met2/conf.py @@ -6,23 +6,23 @@ # -- Project information ----------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information -project = 'MULTI-PARTY MEETING TRANSCRIPTION CHALLENGE 2.0' -copyright = '2023, Speech Lab, Alibaba Group; ASLP Group, Northwestern Polytechnical University' -author = 'Speech Lab, Alibaba Group; Audio, Speech and Language Processing Group, Northwestern Polytechnical University' +project = "MULTI-PARTY MEETING TRANSCRIPTION CHALLENGE 2.0" +copyright = "2023, Speech Lab, Alibaba Group; ASLP Group, Northwestern Polytechnical University" +author = "Speech Lab, Alibaba Group; Audio, Speech and Language Processing Group, Northwestern Polytechnical University" extensions = [ "nbsphinx", "sphinx.ext.autodoc", - 'sphinx.ext.napoleon', - 'sphinx.ext.viewcode', + "sphinx.ext.napoleon", + "sphinx.ext.viewcode", "sphinx.ext.mathjax", "sphinx.ext.todo", # "sphinxarg.ext", "sphinx_markdown_tables", # 'recommonmark', - 'sphinx_rtd_theme', - 'myst_parser', + "sphinx_rtd_theme", + "myst_parser", ] myst_enable_extensions = [ @@ -33,13 +33,12 @@ myst_enable_extensions = [ ] myst_heading_anchors = 2 -myst_highlight_code_blocks=True -myst_update_mathjax=False +myst_highlight_code_blocks = True +myst_update_mathjax = False -templates_path = ['_templates'] +templates_path = ["_templates"] source_suffix = [".rst", ".md"] pygments_style = "sphinx" html_theme = "sphinx_rtd_theme" - diff --git a/examples/aishell/paraformer/utils/compute_wer.py b/examples/aishell/paraformer/utils/compute_wer.py index 26a9f491f..3d00010ca 100755 --- a/examples/aishell/paraformer/utils/compute_wer.py +++ b/examples/aishell/paraformer/utils/compute_wer.py @@ -2,66 +2,98 @@ import os import numpy as np import sys -def compute_wer(ref_file, - hyp_file, - cer_detail_file): + +def compute_wer(ref_file, hyp_file, cer_detail_file): rst = { - 'Wrd': 0, - 'Corr': 0, - 'Ins': 0, - 'Del': 0, - 'Sub': 0, - 'Snt': 0, - 'Err': 0.0, - 'S.Err': 0.0, - 'wrong_words': 0, - 'wrong_sentences': 0 + "Wrd": 0, + "Corr": 0, + "Ins": 0, + "Del": 0, + "Sub": 0, + "Snt": 0, + "Err": 0.0, + "S.Err": 0.0, + "wrong_words": 0, + "wrong_sentences": 0, } hyp_dict = {} ref_dict = {} - with open(hyp_file, 'r') as hyp_reader: + with open(hyp_file, "r") as hyp_reader: for line in hyp_reader: key = line.strip().split()[0] value = line.strip().split()[1:] hyp_dict[key] = value - with open(ref_file, 'r') as ref_reader: + with open(ref_file, "r") as ref_reader: for line in ref_reader: key = line.strip().split()[0] value = line.strip().split()[1:] ref_dict[key] = value - cer_detail_writer = open(cer_detail_file, 'w') + cer_detail_writer = open(cer_detail_file, "w") for hyp_key in hyp_dict: if hyp_key in ref_dict: - out_item = compute_wer_by_line(hyp_dict[hyp_key], ref_dict[hyp_key]) - rst['Wrd'] += out_item['nwords'] - rst['Corr'] += out_item['cor'] - rst['wrong_words'] += out_item['wrong'] - rst['Ins'] += out_item['ins'] - rst['Del'] += out_item['del'] - rst['Sub'] += out_item['sub'] - rst['Snt'] += 1 - if out_item['wrong'] > 0: - rst['wrong_sentences'] += 1 - cer_detail_writer.write(hyp_key + print_cer_detail(out_item) + '\n') - cer_detail_writer.write("ref:" + '\t' + " ".join(list(map(lambda x: x.lower(), ref_dict[hyp_key]))) + '\n') - cer_detail_writer.write("hyp:" + '\t' + " ".join(list(map(lambda x: x.lower(), hyp_dict[hyp_key]))) + '\n') + out_item = compute_wer_by_line(hyp_dict[hyp_key], ref_dict[hyp_key]) + rst["Wrd"] += out_item["nwords"] + rst["Corr"] += out_item["cor"] + rst["wrong_words"] += out_item["wrong"] + rst["Ins"] += out_item["ins"] + rst["Del"] += out_item["del"] + rst["Sub"] += out_item["sub"] + rst["Snt"] += 1 + if out_item["wrong"] > 0: + rst["wrong_sentences"] += 1 + cer_detail_writer.write(hyp_key + print_cer_detail(out_item) + "\n") + cer_detail_writer.write( + "ref:" + "\t" + " ".join(list(map(lambda x: x.lower(), ref_dict[hyp_key]))) + "\n" + ) + cer_detail_writer.write( + "hyp:" + "\t" + " ".join(list(map(lambda x: x.lower(), hyp_dict[hyp_key]))) + "\n" + ) - if rst['Wrd'] > 0: - rst['Err'] = round(rst['wrong_words'] * 100 / rst['Wrd'], 2) - if rst['Snt'] > 0: - rst['S.Err'] = round(rst['wrong_sentences'] * 100 / rst['Snt'], 2) + if rst["Wrd"] > 0: + rst["Err"] = round(rst["wrong_words"] * 100 / rst["Wrd"], 2) + if rst["Snt"] > 0: + rst["S.Err"] = round(rst["wrong_sentences"] * 100 / rst["Snt"], 2) - cer_detail_writer.write('\n') - cer_detail_writer.write("%WER " + str(rst['Err']) + " [ " + str(rst['wrong_words'])+ " / " + str(rst['Wrd']) + - ", " + str(rst['Ins']) + " ins, " + str(rst['Del']) + " del, " + str(rst['Sub']) + " sub ]" + '\n') - cer_detail_writer.write("%SER " + str(rst['S.Err']) + " [ " + str(rst['wrong_sentences']) + " / " + str(rst['Snt']) + " ]" + '\n') - cer_detail_writer.write("Scored " + str(len(hyp_dict)) + " sentences, " + str(len(hyp_dict) - rst['Snt']) + " not present in hyp." + '\n') + cer_detail_writer.write("\n") + cer_detail_writer.write( + "%WER " + + str(rst["Err"]) + + " [ " + + str(rst["wrong_words"]) + + " / " + + str(rst["Wrd"]) + + ", " + + str(rst["Ins"]) + + " ins, " + + str(rst["Del"]) + + " del, " + + str(rst["Sub"]) + + " sub ]" + + "\n" + ) + cer_detail_writer.write( + "%SER " + + str(rst["S.Err"]) + + " [ " + + str(rst["wrong_sentences"]) + + " / " + + str(rst["Snt"]) + + " ]" + + "\n" + ) + cer_detail_writer.write( + "Scored " + + str(len(hyp_dict)) + + " sentences, " + + str(len(hyp_dict) - rst["Snt"]) + + " not present in hyp." + + "\n" + ) - -def compute_wer_by_line(hyp, - ref): + +def compute_wer_by_line(hyp, ref): hyp = list(map(lambda x: x.lower(), hyp)) ref = list(map(lambda x: x.lower(), ref)) @@ -96,14 +128,7 @@ def compute_wer_by_line(hyp, match_idx = [] i = len_hyp j = len_ref - rst = { - 'nwords': len_ref, - 'cor': 0, - 'wrong': 0, - 'ins': 0, - 'del': 0, - 'sub': 0 - } + rst = {"nwords": len_ref, "cor": 0, "wrong": 0, "ins": 0, "del": 0, "sub": 0} while i >= 0 or j >= 0: i_idx = max(0, i) j_idx = max(0, j) @@ -111,42 +136,57 @@ def compute_wer_by_line(hyp, if ops_matrix[i_idx][j_idx] == 0: # correct if i - 1 >= 0 and j - 1 >= 0: match_idx.append((j - 1, i - 1)) - rst['cor'] += 1 + rst["cor"] += 1 i -= 1 j -= 1 elif ops_matrix[i_idx][j_idx] == 2: # insert i -= 1 - rst['ins'] += 1 + rst["ins"] += 1 elif ops_matrix[i_idx][j_idx] == 3: # delete j -= 1 - rst['del'] += 1 + rst["del"] += 1 elif ops_matrix[i_idx][j_idx] == 1: # substitute i -= 1 j -= 1 - rst['sub'] += 1 + rst["sub"] += 1 if i < 0 and j >= 0: - rst['del'] += 1 + rst["del"] += 1 elif j < 0 and i >= 0: - rst['ins'] += 1 + rst["ins"] += 1 match_idx.reverse() wrong_cnt = cost_matrix[len_hyp][len_ref] - rst['wrong'] = wrong_cnt + rst["wrong"] = wrong_cnt return rst -def print_cer_detail(rst): - return ("(" + "nwords=" + str(rst['nwords']) + ",cor=" + str(rst['cor']) - + ",ins=" + str(rst['ins']) + ",del=" + str(rst['del']) + ",sub=" - + str(rst['sub']) + ") corr:" + '{:.2%}'.format(rst['cor']/rst['nwords']) - + ",cer:" + '{:.2%}'.format(rst['wrong']/rst['nwords'])) -if __name__ == '__main__': +def print_cer_detail(rst): + return ( + "(" + + "nwords=" + + str(rst["nwords"]) + + ",cor=" + + str(rst["cor"]) + + ",ins=" + + str(rst["ins"]) + + ",del=" + + str(rst["del"]) + + ",sub=" + + str(rst["sub"]) + + ") corr:" + + "{:.2%}".format(rst["cor"] / rst["nwords"]) + + ",cer:" + + "{:.2%}".format(rst["wrong"] / rst["nwords"]) + ) + + +if __name__ == "__main__": if len(sys.argv) != 4: print("usage : python compute-wer.py test.ref test.hyp test.wer") sys.exit(0) diff --git a/examples/aishell/paraformer/utils/extract_embeds.py b/examples/aishell/paraformer/utils/extract_embeds.py index 7b817d8ca..e0cf98dae 100755 --- a/examples/aishell/paraformer/utils/extract_embeds.py +++ b/examples/aishell/paraformer/utils/extract_embeds.py @@ -5,6 +5,7 @@ import os import torch from kaldiio import WriteHelper import re + text_file_json = sys.argv[1] out_ark = sys.argv[2] out_scp = sys.argv[3] @@ -16,17 +17,17 @@ model = AutoModel.from_pretrained(model_path) tokenizer = AutoTokenizer.from_pretrained(model_path) extractor = pipeline(task="feature-extraction", model=model, tokenizer=tokenizer, device=device) -with open(text_file_json, 'r') as f: +with open(text_file_json, "r") as f: js = f.readlines() f_shape = open(out_shape, "w") -with WriteHelper('ark,scp:{},{}'.format(out_ark, out_scp)) as writer: +with WriteHelper("ark,scp:{},{}".format(out_ark, out_scp)) as writer: with torch.no_grad(): for idx, line in enumerate(js): id, tokens = line.strip().split(" ", 1) tokens = re.sub(" ", "", tokens.strip()) - tokens = ' '.join([j for j in tokens]) + tokens = " ".join([j for j in tokens]) token_num = len(tokens.split(" ")) outputs = extractor(tokens) outputs = np.array(outputs) @@ -38,10 +39,11 @@ with WriteHelper('ark,scp:{},{}'.format(out_ark, out_scp)) as writer: shape_line = "{} {},{}\n".format(id, token_num_embeds, dim) f_shape.write(shape_line) else: - print("{}, size has changed, {}, {}, {}".format(id, token_num, token_num_embeds, tokens)) - + print( + "{}, size has changed, {}, {}, {}".format( + id, token_num, token_num_embeds, tokens + ) + ) f_shape.close() - - diff --git a/examples/aishell/paraformer/utils/postprocess_text_zh.py b/examples/aishell/paraformer/utils/postprocess_text_zh.py index 9e517a4e1..d03febd96 100755 --- a/examples/aishell/paraformer/utils/postprocess_text_zh.py +++ b/examples/aishell/paraformer/utils/postprocess_text_zh.py @@ -1,4 +1,3 @@ - import sys import re @@ -7,25 +6,25 @@ out_f = sys.argv[2] with open(in_f, "r", encoding="utf-8") as f: - lines = f.readlines() + lines = f.readlines() with open(out_f, "w", encoding="utf-8") as f: - for line in lines: - outs = line.strip().split(" ", 1) - if len(outs) == 2: - idx, text = outs - text = re.sub("", "", text) - text = re.sub("", "", text) - text = re.sub("@@", "", text) - text = re.sub("@", "", text) - text = re.sub("", "", text) - text = re.sub(" ", "", text) - text = text.lower() - else: - idx = outs[0] - text = " " + for line in lines: + outs = line.strip().split(" ", 1) + if len(outs) == 2: + idx, text = outs + text = re.sub("", "", text) + text = re.sub("", "", text) + text = re.sub("@@", "", text) + text = re.sub("@", "", text) + text = re.sub("", "", text) + text = re.sub(" ", "", text) + text = text.lower() + else: + idx = outs[0] + text = " " - text = [x for x in text] - text = " ".join(text) - out = "{} {}\n".format(idx, text) - f.write(out) + text = [x for x in text] + text = " ".join(text) + out = "{} {}\n".format(idx, text) + f.write(out) diff --git a/examples/aishell/paraformer/utils/text2token.py b/examples/aishell/paraformer/utils/text2token.py index 56c39138f..263cb9d1a 100755 --- a/examples/aishell/paraformer/utils/text2token.py +++ b/examples/aishell/paraformer/utils/text2token.py @@ -37,9 +37,7 @@ def get_parser(): help="number of characters to split, i.e., \ aabb -> a a b b with -n 1 and aa bb with -n 2", ) - parser.add_argument( - "--skip-ncols", "-s", default=0, type=int, help="skip first n columns" - ) + parser.add_argument("--skip-ncols", "-s", default=0, type=int, help="skip first n columns") parser.add_argument("--space", default="", type=str, help="space symbol") parser.add_argument( "--non-lang-syms", @@ -80,9 +78,7 @@ def main(): else: f = codecs.getreader("utf-8")(sys.stdin if is_python2 else sys.stdin.buffer) - sys.stdout = codecs.getwriter("utf-8")( - sys.stdout if is_python2 else sys.stdout.buffer - ) + sys.stdout = codecs.getwriter("utf-8")(sys.stdout if is_python2 else sys.stdout.buffer) line = f.readline() n = args.nchar while line: diff --git a/examples/aishell/paraformer/utils/text_tokenize.py b/examples/aishell/paraformer/utils/text_tokenize.py index 962ea11bc..31500f191 100755 --- a/examples/aishell/paraformer/utils/text_tokenize.py +++ b/examples/aishell/paraformer/utils/text_tokenize.py @@ -4,7 +4,7 @@ import argparse def load_dict(seg_file): seg_dict = {} - with open(seg_file, 'r') as infile: + with open(seg_file, "r") as infile: for line in infile: s = line.strip().split() key = s[0] @@ -28,8 +28,7 @@ def forward_segment(text, dic): return word_list -def tokenize(txt, - seg_dict): +def tokenize(txt, seg_dict): out_txt = "" pattern = re.compile(r"([\u4E00-\u9FA5A-Za-z0-9])") for word in txt: @@ -87,20 +86,19 @@ def main(): parser = get_parser() args = parser.parse_args() - txt_writer = open("{}/text.{}.txt".format(args.output_dir, args.txt_index), 'w') - shape_writer = open("{}/len.{}".format(args.output_dir, args.txt_index), 'w') + txt_writer = open("{}/text.{}.txt".format(args.output_dir, args.txt_index), "w") + shape_writer = open("{}/len.{}".format(args.output_dir, args.txt_index), "w") seg_dict = load_dict(args.seg_file) - with open(args.text_file, 'r') as infile: + with open(args.text_file, "r") as infile: for line in infile: s = line.strip().split() text_id = s[0] text_list = forward_segment("".join(s[1:]).lower(), seg_dict) text = tokenize(text_list, seg_dict) lens = len(text.strip().split()) - txt_writer.write(text_id + " " + text + '\n') - shape_writer.write(text_id + " " + str(lens) + '\n') + txt_writer.write(text_id + " " + text + "\n") + shape_writer.write(text_id + " " + str(lens) + "\n") -if __name__ == '__main__': +if __name__ == "__main__": main() - diff --git a/examples/aishell/paraformer/utils/textnorm_zh.py b/examples/aishell/paraformer/utils/textnorm_zh.py index 79feb83fd..9de8e81ef 100755 --- a/examples/aishell/paraformer/utils/textnorm_zh.py +++ b/examples/aishell/paraformer/utils/textnorm_zh.py @@ -14,50 +14,59 @@ import sys, os, argparse, codecs, string, re # ================================================================================ # # basic constant # ================================================================================ # -CHINESE_DIGIS = u'零一二三四五六七八九' -BIG_CHINESE_DIGIS_SIMPLIFIED = u'零壹贰叁肆伍陆柒捌玖' -BIG_CHINESE_DIGIS_TRADITIONAL = u'零壹貳參肆伍陸柒捌玖' -SMALLER_BIG_CHINESE_UNITS_SIMPLIFIED = u'十百千万' -SMALLER_BIG_CHINESE_UNITS_TRADITIONAL = u'拾佰仟萬' -LARGER_CHINESE_NUMERING_UNITS_SIMPLIFIED = u'亿兆京垓秭穰沟涧正载' -LARGER_CHINESE_NUMERING_UNITS_TRADITIONAL = u'億兆京垓秭穰溝澗正載' -SMALLER_CHINESE_NUMERING_UNITS_SIMPLIFIED = u'十百千万' -SMALLER_CHINESE_NUMERING_UNITS_TRADITIONAL = u'拾佰仟萬' +CHINESE_DIGIS = "零一二三四五六七八九" +BIG_CHINESE_DIGIS_SIMPLIFIED = "零壹贰叁肆伍陆柒捌玖" +BIG_CHINESE_DIGIS_TRADITIONAL = "零壹貳參肆伍陸柒捌玖" +SMALLER_BIG_CHINESE_UNITS_SIMPLIFIED = "十百千万" +SMALLER_BIG_CHINESE_UNITS_TRADITIONAL = "拾佰仟萬" +LARGER_CHINESE_NUMERING_UNITS_SIMPLIFIED = "亿兆京垓秭穰沟涧正载" +LARGER_CHINESE_NUMERING_UNITS_TRADITIONAL = "億兆京垓秭穰溝澗正載" +SMALLER_CHINESE_NUMERING_UNITS_SIMPLIFIED = "十百千万" +SMALLER_CHINESE_NUMERING_UNITS_TRADITIONAL = "拾佰仟萬" -ZERO_ALT = u'〇' -ONE_ALT = u'幺' -TWO_ALTS = [u'两', u'兩'] +ZERO_ALT = "〇" +ONE_ALT = "幺" +TWO_ALTS = ["两", "兩"] -POSITIVE = [u'正', u'正'] -NEGATIVE = [u'负', u'負'] -POINT = [u'点', u'點'] +POSITIVE = ["正", "正"] +NEGATIVE = ["负", "負"] +POINT = ["点", "點"] # PLUS = [u'加', u'加'] # SIL = [u'杠', u'槓'] -FILLER_CHARS = ['呃', '啊'] -ER_WHITELIST = '(儿女|儿子|儿孙|女儿|儿媳|妻儿|' \ - '胎儿|婴儿|新生儿|婴幼儿|幼儿|少儿|小儿|儿歌|儿童|儿科|托儿所|孤儿|' \ - '儿戏|儿化|台儿庄|鹿儿岛|正儿八经|吊儿郎当|生儿育女|托儿带女|养儿防老|痴儿呆女|' \ - '佳儿佳妇|儿怜兽扰|儿无常父|儿不嫌母丑|儿行千里母担忧|儿大不由爷|苏乞儿)' +FILLER_CHARS = ["呃", "啊"] +ER_WHITELIST = ( + "(儿女|儿子|儿孙|女儿|儿媳|妻儿|" + "胎儿|婴儿|新生儿|婴幼儿|幼儿|少儿|小儿|儿歌|儿童|儿科|托儿所|孤儿|" + "儿戏|儿化|台儿庄|鹿儿岛|正儿八经|吊儿郎当|生儿育女|托儿带女|养儿防老|痴儿呆女|" + "佳儿佳妇|儿怜兽扰|儿无常父|儿不嫌母丑|儿行千里母担忧|儿大不由爷|苏乞儿)" +) # 中文数字系统类型 -NUMBERING_TYPES = ['low', 'mid', 'high'] +NUMBERING_TYPES = ["low", "mid", "high"] -CURRENCY_NAMES = '(人民币|美元|日元|英镑|欧元|马克|法郎|加拿大元|澳元|港币|先令|芬兰马克|爱尔兰镑|' \ - '里拉|荷兰盾|埃斯库多|比塞塔|印尼盾|林吉特|新西兰元|比索|卢布|新加坡元|韩元|泰铢)' -CURRENCY_UNITS = '((亿|千万|百万|万|千|百)|(亿|千万|百万|万|千|百|)元|(亿|千万|百万|万|千|百|)块|角|毛|分)' -COM_QUANTIFIERS = '(匹|张|座|回|场|尾|条|个|首|阙|阵|网|炮|顶|丘|棵|只|支|袭|辆|挑|担|颗|壳|窠|曲|墙|群|腔|' \ - '砣|座|客|贯|扎|捆|刀|令|打|手|罗|坡|山|岭|江|溪|钟|队|单|双|对|出|口|头|脚|板|跳|枝|件|贴|' \ - '针|线|管|名|位|身|堂|课|本|页|家|户|层|丝|毫|厘|分|钱|两|斤|担|铢|石|钧|锱|忽|(千|毫|微)克|' \ - '毫|厘|分|寸|尺|丈|里|寻|常|铺|程|(千|分|厘|毫|微)米|撮|勺|合|升|斗|石|盘|碗|碟|叠|桶|笼|盆|' \ - '盒|杯|钟|斛|锅|簋|篮|盘|桶|罐|瓶|壶|卮|盏|箩|箱|煲|啖|袋|钵|年|月|日|季|刻|时|周|天|秒|分|旬|' \ - '纪|岁|世|更|夜|春|夏|秋|冬|代|伏|辈|丸|泡|粒|颗|幢|堆|条|根|支|道|面|片|张|颗|块)' +CURRENCY_NAMES = ( + "(人民币|美元|日元|英镑|欧元|马克|法郎|加拿大元|澳元|港币|先令|芬兰马克|爱尔兰镑|" + "里拉|荷兰盾|埃斯库多|比塞塔|印尼盾|林吉特|新西兰元|比索|卢布|新加坡元|韩元|泰铢)" +) +CURRENCY_UNITS = ( + "((亿|千万|百万|万|千|百)|(亿|千万|百万|万|千|百|)元|(亿|千万|百万|万|千|百|)块|角|毛|分)" +) +COM_QUANTIFIERS = ( + "(匹|张|座|回|场|尾|条|个|首|阙|阵|网|炮|顶|丘|棵|只|支|袭|辆|挑|担|颗|壳|窠|曲|墙|群|腔|" + "砣|座|客|贯|扎|捆|刀|令|打|手|罗|坡|山|岭|江|溪|钟|队|单|双|对|出|口|头|脚|板|跳|枝|件|贴|" + "针|线|管|名|位|身|堂|课|本|页|家|户|层|丝|毫|厘|分|钱|两|斤|担|铢|石|钧|锱|忽|(千|毫|微)克|" + "毫|厘|分|寸|尺|丈|里|寻|常|铺|程|(千|分|厘|毫|微)米|撮|勺|合|升|斗|石|盘|碗|碟|叠|桶|笼|盆|" + "盒|杯|钟|斛|锅|簋|篮|盘|桶|罐|瓶|壶|卮|盏|箩|箱|煲|啖|袋|钵|年|月|日|季|刻|时|周|天|秒|分|旬|" + "纪|岁|世|更|夜|春|夏|秋|冬|代|伏|辈|丸|泡|粒|颗|幢|堆|条|根|支|道|面|片|张|颗|块)" +) # punctuation information are based on Zhon project (https://github.com/tsroten/zhon.git) -CHINESE_PUNC_STOP = '！？｡。' -CHINESE_PUNC_NON_STOP = '＂＃＄％＆＇（）＊＋，－／：；＜＝＞＠［＼］＾＿｀｛｜｝～｟｠｢｣､、〃》「」『』【】〔〕〖〗〘〙〚〛〜〝〞〟〰〾〿–—‘’‛“”„‟…‧﹏' +CHINESE_PUNC_STOP = "！？｡。" +CHINESE_PUNC_NON_STOP = "＂＃＄％＆＇（）＊＋，－／：；＜＝＞＠［＼］＾＿｀｛｜｝～｟｠｢｣､、〃》「」『』【】〔〕〖〗〘〙〚〛〜〝〞〟〰〾〿–—‘’‛“”„‟…‧﹏" CHINESE_PUNC_LIST = CHINESE_PUNC_STOP + CHINESE_PUNC_NON_STOP + # ================================================================================ # # basic class # ================================================================================ # @@ -72,7 +81,7 @@ class ChineseChar(object): def __init__(self, simplified, traditional): self.simplified = simplified self.traditional = traditional - #self.__repr__ = self.__str__ + # self.__repr__ = self.__str__ def __str__(self): return self.simplified or self.traditional or None @@ -95,26 +104,49 @@ class ChineseNumberUnit(ChineseChar): self.big_t = big_t def __str__(self): - return '10^{}'.format(self.power) + return "10^{}".format(self.power) @classmethod def create(cls, index, value, numbering_type=NUMBERING_TYPES[1], small_unit=False): if small_unit: - return ChineseNumberUnit(power=index + 1, - simplified=value[0], traditional=value[1], big_s=value[1], big_t=value[1]) + return ChineseNumberUnit( + power=index + 1, + simplified=value[0], + traditional=value[1], + big_s=value[1], + big_t=value[1], + ) elif numbering_type == NUMBERING_TYPES[0]: - return ChineseNumberUnit(power=index + 8, - simplified=value[0], traditional=value[1], big_s=value[0], big_t=value[1]) + return ChineseNumberUnit( + power=index + 8, + simplified=value[0], + traditional=value[1], + big_s=value[0], + big_t=value[1], + ) elif numbering_type == NUMBERING_TYPES[1]: - return ChineseNumberUnit(power=(index + 2) * 4, - simplified=value[0], traditional=value[1], big_s=value[0], big_t=value[1]) + return ChineseNumberUnit( + power=(index + 2) * 4, + simplified=value[0], + traditional=value[1], + big_s=value[0], + big_t=value[1], + ) elif numbering_type == NUMBERING_TYPES[2]: - return ChineseNumberUnit(power=pow(2, index + 3), - simplified=value[0], traditional=value[1], big_s=value[0], big_t=value[1]) + return ChineseNumberUnit( + power=pow(2, index + 3), + simplified=value[0], + traditional=value[1], + big_s=value[0], + big_t=value[1], + ) else: raise ValueError( - 'Counting type should be in {0} ({1} provided).'.format(NUMBERING_TYPES, numbering_type)) + "Counting type should be in {0} ({1} provided).".format( + NUMBERING_TYPES, numbering_type + ) + ) class ChineseNumberDigit(ChineseChar): @@ -158,6 +190,7 @@ class NumberSystem(object): """ 中文数字系统 """ + pass @@ -207,27 +240,27 @@ def create_system(numbering_type=NUMBERING_TYPES[1]): # chinese number units of '亿' and larger all_larger_units = zip( - LARGER_CHINESE_NUMERING_UNITS_SIMPLIFIED, LARGER_CHINESE_NUMERING_UNITS_TRADITIONAL) - larger_units = [CNU.create(i, v, numbering_type, False) - for i, v in enumerate(all_larger_units)] + LARGER_CHINESE_NUMERING_UNITS_SIMPLIFIED, LARGER_CHINESE_NUMERING_UNITS_TRADITIONAL + ) + larger_units = [CNU.create(i, v, numbering_type, False) for i, v in enumerate(all_larger_units)] # chinese number units of '十, 百, 千, 万' all_smaller_units = zip( - SMALLER_CHINESE_NUMERING_UNITS_SIMPLIFIED, SMALLER_CHINESE_NUMERING_UNITS_TRADITIONAL) - smaller_units = [CNU.create(i, v, small_unit=True) - for i, v in enumerate(all_smaller_units)] + SMALLER_CHINESE_NUMERING_UNITS_SIMPLIFIED, SMALLER_CHINESE_NUMERING_UNITS_TRADITIONAL + ) + smaller_units = [CNU.create(i, v, small_unit=True) for i, v in enumerate(all_smaller_units)] # digis - chinese_digis = zip(CHINESE_DIGIS, CHINESE_DIGIS, - BIG_CHINESE_DIGIS_SIMPLIFIED, BIG_CHINESE_DIGIS_TRADITIONAL) + chinese_digis = zip( + CHINESE_DIGIS, CHINESE_DIGIS, BIG_CHINESE_DIGIS_SIMPLIFIED, BIG_CHINESE_DIGIS_TRADITIONAL + ) digits = [CND.create(i, v) for i, v in enumerate(chinese_digis)] digits[0].alt_s, digits[0].alt_t = ZERO_ALT, ZERO_ALT digits[1].alt_s, digits[1].alt_t = ONE_ALT, ONE_ALT digits[2].alt_s, digits[2].alt_t = TWO_ALTS[0], TWO_ALTS[1] # symbols - positive_cn = CM(POSITIVE[0], POSITIVE[1], '+', lambda x: x) - negative_cn = CM(NEGATIVE[0], NEGATIVE[1], '-', lambda x: -x) - point_cn = CM(POINT[0], POINT[1], '.', lambda x, - y: float(str(x) + '.' + str(y))) + positive_cn = CM(POSITIVE[0], POSITIVE[1], "+", lambda x: x) + negative_cn = CM(NEGATIVE[0], NEGATIVE[1], "-", lambda x: -x) + point_cn = CM(POINT[0], POINT[1], ".", lambda x, y: float(str(x) + "." + str(y))) # sil_cn = CM(SIL[0], SIL[1], '-', lambda x, y: float(str(x) + '-' + str(y))) system = NumberSystem() system.units = smaller_units + larger_units @@ -251,13 +284,14 @@ def chn2num(chinese_string, numbering_type=NUMBERING_TYPES[1]): return m def string2symbols(chinese_string, system): - int_string, dec_string = chinese_string, '' + int_string, dec_string = chinese_string, "" for p in [system.math.point.simplified, system.math.point.traditional]: if p in chinese_string: int_string, dec_string = chinese_string.split(p) break - return [get_symbol(c, system) for c in int_string], \ - [get_symbol(c, system) for c in dec_string] + return [get_symbol(c, system) for c in int_string], [ + get_symbol(c, system) for c in dec_string + ] def correct_symbols(integer_symbols, system): """ @@ -271,8 +305,7 @@ def chn2num(chinese_string, numbering_type=NUMBERING_TYPES[1]): if len(integer_symbols) > 1: if isinstance(integer_symbols[-1], CND) and isinstance(integer_symbols[-2], CNU): - integer_symbols.append( - CNU(integer_symbols[-2].power - 1, None, None, None, None)) + integer_symbols.append(CNU(integer_symbols[-2].power - 1, None, None, None, None)) result = [] unit_count = 0 @@ -288,9 +321,13 @@ def chn2num(chinese_string, numbering_type=NUMBERING_TYPES[1]): result.append(current_unit) elif unit_count > 1: for i in range(len(result)): - if isinstance(result[-i - 1], CNU) and result[-i - 1].power < current_unit.power: - result[-i - 1] = CNU(result[-i - 1].power + - current_unit.power, None, None, None, None) + if ( + isinstance(result[-i - 1], CNU) + and result[-i - 1].power < current_unit.power + ): + result[-i - 1] = CNU( + result[-i - 1].power + current_unit.power, None, None, None, None + ) return result def compute_value(integer_symbols): @@ -307,8 +344,7 @@ def chn2num(chinese_string, numbering_type=NUMBERING_TYPES[1]): elif isinstance(s, CNU): value[-1] *= pow(10, s.power) if s.power > last_power: - value[:-1] = list(map(lambda v: v * - pow(10, s.power), value[:-1])) + value[:-1] = list(map(lambda v: v * pow(10, s.power), value[:-1])) last_power = s.power value.append(0) return sum(value) @@ -317,20 +353,28 @@ def chn2num(chinese_string, numbering_type=NUMBERING_TYPES[1]): int_part, dec_part = string2symbols(chinese_string, system) int_part = correct_symbols(int_part, system) int_str = str(compute_value(int_part)) - dec_str = ''.join([str(d.value) for d in dec_part]) + dec_str = "".join([str(d.value) for d in dec_part]) if dec_part: - return '{0}.{1}'.format(int_str, dec_str) + return "{0}.{1}".format(int_str, dec_str) else: return int_str -def num2chn(number_string, numbering_type=NUMBERING_TYPES[1], big=False, - traditional=False, alt_zero=False, alt_one=False, alt_two=True, - use_zeros=True, use_units=True): +def num2chn( + number_string, + numbering_type=NUMBERING_TYPES[1], + big=False, + traditional=False, + alt_zero=False, + alt_one=False, + alt_two=True, + use_zeros=True, + use_units=True, +): def get_value(value_string, use_zeros=True): - striped_string = value_string.lstrip('0') + striped_string = value_string.lstrip("0") # record nothing if all zeros if not striped_string: @@ -345,14 +389,17 @@ def num2chn(number_string, numbering_type=NUMBERING_TYPES[1], big=False, # recursively record multiple digits else: - result_unit = next(u for u in reversed( - system.units) if u.power < len(striped_string)) - result_string = value_string[:-result_unit.power] - return get_value(result_string) + [result_unit] + get_value(striped_string[-result_unit.power:]) + result_unit = next(u for u in reversed(system.units) if u.power < len(striped_string)) + result_string = value_string[: -result_unit.power] + return ( + get_value(result_string) + + [result_unit] + + get_value(striped_string[-result_unit.power :]) + ) system = create_system(numbering_type) - int_dec = number_string.split('.') + int_dec = number_string.split(".") if len(int_dec) == 1: int_string = int_dec[0] dec_string = "" @@ -361,7 +408,8 @@ def num2chn(number_string, numbering_type=NUMBERING_TYPES[1], big=False, dec_string = int_dec[1] else: raise ValueError( - "invalid input num string with more than one dot: {}".format(number_string)) + "invalid input num string with more than one dot: {}".format(number_string) + ) if use_units and len(int_string) > 1: result_symbols = get_value(int_string) @@ -372,51 +420,62 @@ def num2chn(number_string, numbering_type=NUMBERING_TYPES[1], big=False, result_symbols += [system.math.point] + dec_symbols if alt_two: - liang = CND(2, system.digits[2].alt_s, system.digits[2].alt_t, - system.digits[2].big_s, system.digits[2].big_t) + liang = CND( + 2, + system.digits[2].alt_s, + system.digits[2].alt_t, + system.digits[2].big_s, + system.digits[2].big_t, + ) for i, v in enumerate(result_symbols): if isinstance(v, CND) and v.value == 2: - next_symbol = result_symbols[i + - 1] if i < len(result_symbols) - 1 else None + next_symbol = result_symbols[i + 1] if i < len(result_symbols) - 1 else None previous_symbol = result_symbols[i - 1] if i > 0 else None if isinstance(next_symbol, CNU) and isinstance(previous_symbol, (CNU, type(None))): - if next_symbol.power != 1 and ((previous_symbol is None) or (previous_symbol.power != 1)): + if next_symbol.power != 1 and ( + (previous_symbol is None) or (previous_symbol.power != 1) + ): result_symbols[i] = liang # if big is True, '两' will not be used and `alt_two` has no impact on output if big: - attr_name = 'big_' + attr_name = "big_" if traditional: - attr_name += 't' + attr_name += "t" else: - attr_name += 's' + attr_name += "s" else: if traditional: - attr_name = 'traditional' + attr_name = "traditional" else: - attr_name = 'simplified' + attr_name = "simplified" - result = ''.join([getattr(s, attr_name) for s in result_symbols]) + result = "".join([getattr(s, attr_name) for s in result_symbols]) # if not use_zeros: # result = result.strip(getattr(system.digits[0], attr_name)) if alt_zero: - result = result.replace( - getattr(system.digits[0], attr_name), system.digits[0].alt_s) + result = result.replace(getattr(system.digits[0], attr_name), system.digits[0].alt_s) if alt_one: - result = result.replace( - getattr(system.digits[1], attr_name), system.digits[1].alt_s) + result = result.replace(getattr(system.digits[1], attr_name), system.digits[1].alt_s) for i, p in enumerate(POINT): if result.startswith(p): return CHINESE_DIGIS[0] + result # ^10, 11, .., 19 - if len(result) >= 2 and result[1] in [SMALLER_CHINESE_NUMERING_UNITS_SIMPLIFIED[0], - SMALLER_CHINESE_NUMERING_UNITS_TRADITIONAL[0]] and \ - result[0] in [CHINESE_DIGIS[1], BIG_CHINESE_DIGIS_SIMPLIFIED[1], BIG_CHINESE_DIGIS_TRADITIONAL[1]]: + if ( + len(result) >= 2 + and result[1] + in [ + SMALLER_CHINESE_NUMERING_UNITS_SIMPLIFIED[0], + SMALLER_CHINESE_NUMERING_UNITS_TRADITIONAL[0], + ] + and result[0] + in [CHINESE_DIGIS[1], BIG_CHINESE_DIGIS_SIMPLIFIED[1], BIG_CHINESE_DIGIS_TRADITIONAL[1]] + ): result = result[1:] return result @@ -440,6 +499,7 @@ class Cardinal: def cardinal2chntext(self): return num2chn(self.cardinal) + class Digit: """ DIGIT类 @@ -476,17 +536,17 @@ class TelePhone: def telephone2chntext(self, fixed=False): if fixed: - sil_parts = self.telephone.split('-') - self.raw_chntext = ''.join([ - num2chn(part, alt_two=False, use_units=False) for part in sil_parts - ]) - self.chntext = self.raw_chntext.replace('', '') + sil_parts = self.telephone.split("-") + self.raw_chntext = "".join( + [num2chn(part, alt_two=False, use_units=False) for part in sil_parts] + ) + self.chntext = self.raw_chntext.replace("", "") else: - sp_parts = self.telephone.strip('+').split() - self.raw_chntext = ''.join([ - num2chn(part, alt_two=False, use_units=False) for part in sp_parts - ]) - self.chntext = self.raw_chntext.replace('', '') + sp_parts = self.telephone.strip("+").split() + self.raw_chntext = "".join( + [num2chn(part, alt_two=False, use_units=False) for part in sp_parts] + ) + self.chntext = self.raw_chntext.replace("", "") return self.chntext @@ -500,12 +560,12 @@ class Fraction: self.chntext = chntext def chntext2fraction(self): - denominator, numerator = self.chntext.split('分之') - return chn2num(numerator) + '/' + chn2num(denominator) + denominator, numerator = self.chntext.split("分之") + return chn2num(numerator) + "/" + chn2num(denominator) def fraction2chntext(self): - numerator, denominator = self.fraction.split('/') - return num2chn(denominator) + '分之' + num2chn(numerator) + numerator, denominator = self.fraction.split("/") + return num2chn(denominator) + "分之" + num2chn(numerator) class Date: @@ -544,23 +604,23 @@ class Date: def date2chntext(self): date = self.date try: - year, other = date.strip().split('年', 1) - year = Digit(digit=year).digit2chntext() + '年' + year, other = date.strip().split("年", 1) + year = Digit(digit=year).digit2chntext() + "年" except ValueError: other = date - year = '' + year = "" if other: try: - month, day = other.strip().split('月', 1) - month = Cardinal(cardinal=month).cardinal2chntext() + '月' + month, day = other.strip().split("月", 1) + month = Cardinal(cardinal=month).cardinal2chntext() + "月" except ValueError: day = date - month = '' + month = "" if day: day = Cardinal(cardinal=day[:-1]).cardinal2chntext() + day[-1] else: - month = '' - day = '' + month = "" + day = "" chntext = year + month + day self.chntext = chntext return self.chntext @@ -580,7 +640,7 @@ class Money: def money2chntext(self): money = self.money - pattern = re.compile(r'(\d+(\.\d+)?)') + pattern = re.compile(r"(\d+(\.\d+)?)") matchers = pattern.findall(money) if matchers: for matcher in matchers: @@ -599,10 +659,10 @@ class Percentage: self.chntext = chntext def chntext2percentage(self): - return chn2num(self.chntext.strip().strip('百分之')) + '%' + return chn2num(self.chntext.strip().strip("百分之")) + "%" def percentage2chntext(self): - return '百分之' + num2chn(self.percentage.strip().strip('%')) + return "百分之" + num2chn(self.percentage.strip().strip("%")) def remove_erhua(text, er_whitelist): @@ -612,9 +672,9 @@ def remove_erhua(text, er_whitelist): """ er_pattern = re.compile(er_whitelist) - new_str='' - while re.search('儿',text): - a = re.search('儿',text).span() + new_str = "" + while re.search("儿", text): + a = re.search("儿", text).span() remove_er_flag = 0 if er_pattern.search(text): @@ -622,23 +682,24 @@ def remove_erhua(text, er_whitelist): if b[0] <= a[0]: remove_er_flag = 1 - if remove_er_flag == 0 : - new_str = new_str + text[0:a[0]] - text = text[a[1]:] + if remove_er_flag == 0: + new_str = new_str + text[0 : a[0]] + text = text[a[1] :] else: - new_str = new_str + text[0:b[1]] - text = text[b[1]:] + new_str = new_str + text[0 : b[1]] + text = text[b[1] :] text = new_str + text return text + # ================================================================================ # # NSW Normalizer # ================================================================================ # class NSWNormalizer: def __init__(self, raw_text): - self.raw_text = '^' + raw_text + '$' - self.norm_text = '' + self.raw_text = "^" + raw_text + "$" + self.norm_text = "" def _particular(self): text = self.norm_text @@ -647,7 +708,7 @@ class NSWNormalizer: if matchers: # print('particular') for matcher in matchers: - text = text.replace(matcher[0], matcher[1]+'2'+matcher[2], 1) + text = text.replace(matcher[0], matcher[1] + "2" + matcher[2], 1) self.norm_text = text return self.norm_text @@ -658,15 +719,17 @@ class NSWNormalizer: pattern = re.compile(r"\D+((([089]\d|(19|20)\d{2})年)?(\d{1,2}月(\d{1,2}[日号])?)?)") matchers = pattern.findall(text) if matchers: - #print('date') + # print('date') for matcher in matchers: text = text.replace(matcher[0], Date(date=matcher[0]).date2chntext(), 1) # 规范化金钱 - pattern = re.compile(r"\D+((\d+(\.\d+)?)[多余几]?" + CURRENCY_UNITS + r"(\d" + CURRENCY_UNITS + r"?)?)") + pattern = re.compile( + r"\D+((\d+(\.\d+)?)[多余几]?" + CURRENCY_UNITS + r"(\d" + CURRENCY_UNITS + r"?)?)" + ) matchers = pattern.findall(text) if matchers: - #print('money') + # print('money') for matcher in matchers: text = text.replace(matcher[0], Money(money=matcher[0]).money2chntext(), 1) @@ -679,39 +742,45 @@ class NSWNormalizer: pattern = re.compile(r"\D((\+?86 ?)?1([38]\d|5[0-35-9]|7[678]|9[89])\d{8})\D") matchers = pattern.findall(text) if matchers: - #print('telephone') + # print('telephone') for matcher in matchers: - text = text.replace(matcher[0], TelePhone(telephone=matcher[0]).telephone2chntext(), 1) + text = text.replace( + matcher[0], TelePhone(telephone=matcher[0]).telephone2chntext(), 1 + ) # 固话 pattern = re.compile(r"\D((0(10|2[1-3]|[3-9]\d{2})-?)?[1-9]\d{6,7})\D") matchers = pattern.findall(text) if matchers: # print('fixed telephone') for matcher in matchers: - text = text.replace(matcher[0], TelePhone(telephone=matcher[0]).telephone2chntext(fixed=True), 1) + text = text.replace( + matcher[0], TelePhone(telephone=matcher[0]).telephone2chntext(fixed=True), 1 + ) # 规范化分数 pattern = re.compile(r"(\d+/\d+)") matchers = pattern.findall(text) if matchers: - #print('fraction') + # print('fraction') for matcher in matchers: text = text.replace(matcher, Fraction(fraction=matcher).fraction2chntext(), 1) # 规范化百分数 - text = text.replace('％', '%') + text = text.replace("％", "%") pattern = re.compile(r"(\d+(\.\d+)?%)") matchers = pattern.findall(text) if matchers: - #print('percentage') + # print('percentage') for matcher in matchers: - text = text.replace(matcher[0], Percentage(percentage=matcher[0]).percentage2chntext(), 1) + text = text.replace( + matcher[0], Percentage(percentage=matcher[0]).percentage2chntext(), 1 + ) # 规范化纯数+量词 pattern = re.compile(r"(\d+(\.\d+)?)[多余几]?" + COM_QUANTIFIERS) matchers = pattern.findall(text) if matchers: - #print('cardinal+quantifier') + # print('cardinal+quantifier') for matcher in matchers: text = text.replace(matcher[0], Cardinal(cardinal=matcher[0]).cardinal2chntext(), 1) @@ -719,7 +788,7 @@ class NSWNormalizer: pattern = re.compile(r"(\d{4,32})") matchers = pattern.findall(text) if matchers: - #print('digit') + # print('digit') for matcher in matchers: text = text.replace(matcher, Digit(digit=matcher).digit2chntext(), 1) @@ -727,74 +796,82 @@ class NSWNormalizer: pattern = re.compile(r"(\d+(\.\d+)?)") matchers = pattern.findall(text) if matchers: - #print('cardinal') + # print('cardinal') for matcher in matchers: text = text.replace(matcher[0], Cardinal(cardinal=matcher[0]).cardinal2chntext(), 1) self.norm_text = text self._particular() - return self.norm_text.lstrip('^').rstrip('$') + return self.norm_text.lstrip("^").rstrip("$") def nsw_test_case(raw_text): - print('I:' + raw_text) - print('O:' + NSWNormalizer(raw_text).normalize()) - print('') + print("I:" + raw_text) + print("O:" + NSWNormalizer(raw_text).normalize()) + print("") def nsw_test(): - nsw_test_case('固话：0595-23865596或23880880。') - nsw_test_case('固话：0595-23865596或23880880。') - nsw_test_case('手机：+86 19859213959或15659451527。') - nsw_test_case('分数：32477/76391。') - nsw_test_case('百分数：80.03%。') - nsw_test_case('编号：31520181154418。') - nsw_test_case('纯数：2983.07克或12345.60米。') - nsw_test_case('日期：1999年2月20日或09年3月15号。') - nsw_test_case('金钱：12块5，34.5元，20.1万') - nsw_test_case('特殊：O2O或B2C。') - nsw_test_case('3456万吨') - nsw_test_case('2938个') - nsw_test_case('938') - nsw_test_case('今天吃了115个小笼包231个馒头') - nsw_test_case('有62％的概率') + nsw_test_case("固话：0595-23865596或23880880。") + nsw_test_case("固话：0595-23865596或23880880。") + nsw_test_case("手机：+86 19859213959或15659451527。") + nsw_test_case("分数：32477/76391。") + nsw_test_case("百分数：80.03%。") + nsw_test_case("编号：31520181154418。") + nsw_test_case("纯数：2983.07克或12345.60米。") + nsw_test_case("日期：1999年2月20日或09年3月15号。") + nsw_test_case("金钱：12块5，34.5元，20.1万") + nsw_test_case("特殊：O2O或B2C。") + nsw_test_case("3456万吨") + nsw_test_case("2938个") + nsw_test_case("938") + nsw_test_case("今天吃了115个小笼包231个馒头") + nsw_test_case("有62％的概率") -if __name__ == '__main__': - #nsw_test() +if __name__ == "__main__": + # nsw_test() p = argparse.ArgumentParser() - p.add_argument('ifile', help='input filename, assume utf-8 encoding') - p.add_argument('ofile', help='output filename') - p.add_argument('--to_upper', action='store_true', help='convert to upper case') - p.add_argument('--to_lower', action='store_true', help='convert to lower case') - p.add_argument('--has_key', action='store_true', help="input text has Kaldi's key as first field.") - p.add_argument('--remove_fillers', type=bool, default=True, help='remove filler chars such as "呃, 啊"') - p.add_argument('--remove_erhua', type=bool, default=True, help='remove erhua chars such as "这儿"') - p.add_argument('--log_interval', type=int, default=10000, help='log interval in number of processed lines') + p.add_argument("ifile", help="input filename, assume utf-8 encoding") + p.add_argument("ofile", help="output filename") + p.add_argument("--to_upper", action="store_true", help="convert to upper case") + p.add_argument("--to_lower", action="store_true", help="convert to lower case") + p.add_argument( + "--has_key", action="store_true", help="input text has Kaldi's key as first field." + ) + p.add_argument( + "--remove_fillers", type=bool, default=True, help='remove filler chars such as "呃, 啊"' + ) + p.add_argument( + "--remove_erhua", type=bool, default=True, help='remove erhua chars such as "这儿"' + ) + p.add_argument( + "--log_interval", type=int, default=10000, help="log interval in number of processed lines" + ) args = p.parse_args() - ifile = codecs.open(args.ifile, 'r', 'utf8') - ofile = codecs.open(args.ofile, 'w+', 'utf8') + ifile = codecs.open(args.ifile, "r", "utf8") + ofile = codecs.open(args.ofile, "w+", "utf8") n = 0 for l in ifile: - key = '' - text = '' + key = "" + text = "" if args.has_key: cols = l.split(maxsplit=1) key = cols[0] if len(cols) == 2: text = cols[1].strip() else: - text = '' + text = "" else: text = l.strip() # cases if args.to_upper and args.to_lower: - sys.stderr.write('text norm: to_upper OR to_lower?') + sys.stderr.write("text norm: to_upper OR to_lower?") exit(1) if args.to_upper: text = text.upper() @@ -804,7 +881,7 @@ if __name__ == '__main__': # Filler chars removal if args.remove_fillers: for ch in FILLER_CHARS: - text = text.replace(ch, '') + text = text.replace(ch, "") if args.remove_erhua: text = remove_erhua(text, ER_WHITELIST) @@ -813,16 +890,16 @@ if __name__ == '__main__': text = NSWNormalizer(text).normalize() # Punctuations removal - old_chars = CHINESE_PUNC_LIST + string.punctuation # includes all CN and EN punctuations - new_chars = ' ' * len(old_chars) - del_chars = '' + old_chars = CHINESE_PUNC_LIST + string.punctuation # includes all CN and EN punctuations + new_chars = " " * len(old_chars) + del_chars = "" text = text.translate(str.maketrans(old_chars, new_chars, del_chars)) # if args.has_key: - ofile.write(key + '\t' + text + '\n') + ofile.write(key + "\t" + text + "\n") else: - ofile.write(text + '\n') + ofile.write(text + "\n") n += 1 if n % args.log_interval == 0: diff --git a/examples/common_voice/whisper_lid/demo_funasr.py b/examples/common_voice/whisper_lid/demo_funasr.py index 5818f4cc8..f771b615a 100644 --- a/examples/common_voice/whisper_lid/demo_funasr.py +++ b/examples/common_voice/whisper_lid/demo_funasr.py @@ -16,4 +16,4 @@ model = AutoModel(model="iic/speech_whisper-large_lid_multilingual_pytorch") for wav_id in multilingual_wavs: wav_file = f"{model.model_path}/examples/{wav_id}" res = model.generate(input=wav_file, data_type="sound", inference_clip_length=250) - print("detect sample {}: {}".format(wav_id, res)) \ No newline at end of file + print("detect sample {}: {}".format(wav_id, res)) diff --git a/examples/common_voice/whisper_lid/demo_modelscope.py b/examples/common_voice/whisper_lid/demo_modelscope.py index f2cc5eac3..b0bc22bcb 100644 --- a/examples/common_voice/whisper_lid/demo_modelscope.py +++ b/examples/common_voice/whisper_lid/demo_modelscope.py @@ -6,7 +6,7 @@ from modelscope.pipelines import pipeline from modelscope.utils.constant import Tasks -multilingual_wavs=[ +multilingual_wavs = [ "https://www.modelscope.cn/api/v1/models/iic/speech_whisper-large_lid_multilingual_pytorch/repo?Revision=master&FilePath=examples/example_zh-CN.mp3", "https://www.modelscope.cn/api/v1/models/iic/speech_whisper-large_lid_multilingual_pytorch/repo?Revision=master&FilePath=examples/example_en.mp3", "https://www.modelscope.cn/api/v1/models/iic/speech_whisper-large_lid_multilingual_pytorch/repo?Revision=master&FilePath=examples/example_ja.mp3", @@ -14,9 +14,9 @@ multilingual_wavs=[ ] inference_pipeline = pipeline( - task=Tasks.auto_speech_recognition, - model='iic/speech_whisper-large_lid_multilingual_pytorch') + task=Tasks.auto_speech_recognition, model="iic/speech_whisper-large_lid_multilingual_pytorch" +) for wav in multilingual_wavs: rec_result = inference_pipeline(input=wav, inference_clip_length=250) - print(rec_result) \ No newline at end of file + print(rec_result) diff --git a/examples/industrial_data_pretraining/bicif_paraformer/demo.py b/examples/industrial_data_pretraining/bicif_paraformer/demo.py index 03b4479fd..57c1c6c0c 100644 --- a/examples/industrial_data_pretraining/bicif_paraformer/demo.py +++ b/examples/industrial_data_pretraining/bicif_paraformer/demo.py @@ -5,11 +5,16 @@ from funasr import AutoModel -model = AutoModel(model="iic/speech_paraformer-large-vad-punc_asr_nat-zh-cn-16k-common-vocab8404-pytorch", - vad_model="iic/speech_fsmn_vad_zh-cn-16k-common-pytorch", - punc_model="iic/punc_ct-transformer_cn-en-common-vocab471067-large", - # spk_model="iic/speech_campplus_sv_zh-cn_16k-common", - ) +model = AutoModel( + model="iic/speech_paraformer-large-vad-punc_asr_nat-zh-cn-16k-common-vocab8404-pytorch", + vad_model="iic/speech_fsmn_vad_zh-cn-16k-common-pytorch", + punc_model="iic/punc_ct-transformer_cn-en-common-vocab471067-large", + # spk_model="iic/speech_campplus_sv_zh-cn_16k-common", +) -res = model.generate(input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_vad_punc_example.wav", batch_size_s=300, batch_size_threshold_s=60) +res = model.generate( + input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_vad_punc_example.wav", + batch_size_s=300, + batch_size_threshold_s=60, +) print(res) diff --git a/examples/industrial_data_pretraining/bicif_paraformer/export.py b/examples/industrial_data_pretraining/bicif_paraformer/export.py index d1b1ea756..31098d2fd 100644 --- a/examples/industrial_data_pretraining/bicif_paraformer/export.py +++ b/examples/industrial_data_pretraining/bicif_paraformer/export.py @@ -7,7 +7,10 @@ from funasr import AutoModel -model = AutoModel(model="iic/speech_paraformer-large-vad-punc_asr_nat-zh-cn-16k-common-vocab8404-pytorch", device="cpu") +model = AutoModel( + model="iic/speech_paraformer-large-vad-punc_asr_nat-zh-cn-16k-common-vocab8404-pytorch", + device="cpu", +) res = model.export(type="onnx", quantize=False) print(res) @@ -16,7 +19,10 @@ print(res) # method2, inference from local path from funasr import AutoModel -model = AutoModel(model="/Users/zhifu/.cache/modelscope/hub/iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch", device="cpu") +model = AutoModel( + model="/Users/zhifu/.cache/modelscope/hub/iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch", + device="cpu", +) res = model.export(type="onnx", quantize=False) -print(res) \ No newline at end of file +print(res) diff --git a/examples/industrial_data_pretraining/campplus_sv/demo.py b/examples/industrial_data_pretraining/campplus_sv/demo.py index c3e939819..ba3356be9 100644 --- a/examples/industrial_data_pretraining/campplus_sv/demo.py +++ b/examples/industrial_data_pretraining/campplus_sv/demo.py @@ -5,8 +5,9 @@ from funasr import AutoModel -model = AutoModel(model="iic/speech_campplus_sv_zh-cn_16k-common" - ) +model = AutoModel(model="iic/speech_campplus_sv_zh-cn_16k-common") -res = model.generate(input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav") -print(res) \ No newline at end of file +res = model.generate( + input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav" +) +print(res) diff --git a/examples/industrial_data_pretraining/conformer/demo.py b/examples/industrial_data_pretraining/conformer/demo.py index 1b07c3004..43cf67de6 100644 --- a/examples/industrial_data_pretraining/conformer/demo.py +++ b/examples/industrial_data_pretraining/conformer/demo.py @@ -7,6 +7,7 @@ from funasr import AutoModel model = AutoModel(model="iic/speech_conformer_asr_nat-zh-cn-16k-aishell2-vocab5212-pytorch") -res = model.generate(input="https://modelscope.oss-cn-beijing.aliyuncs.com/test/audios/asr_example.wav") +res = model.generate( + input="https://modelscope.oss-cn-beijing.aliyuncs.com/test/audios/asr_example.wav" +) print(res) - diff --git a/examples/industrial_data_pretraining/contextual_paraformer/demo.py b/examples/industrial_data_pretraining/contextual_paraformer/demo.py index e7930dfca..d7952a7ce 100755 --- a/examples/industrial_data_pretraining/contextual_paraformer/demo.py +++ b/examples/industrial_data_pretraining/contextual_paraformer/demo.py @@ -7,6 +7,8 @@ from funasr import AutoModel model = AutoModel(model="iic/speech_paraformer-large-contextual_asr_nat-zh-cn-16k-common-vocab8404") -res = model.generate(input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav", - hotword='达摩院 魔搭') -print(res) \ No newline at end of file +res = model.generate( + input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav", + hotword="达摩院 魔搭", +) +print(res) diff --git a/examples/industrial_data_pretraining/ct_transformer/demo.py b/examples/industrial_data_pretraining/ct_transformer/demo.py index 23f3dd756..c19af2b2a 100644 --- a/examples/industrial_data_pretraining/ct_transformer/demo.py +++ b/examples/industrial_data_pretraining/ct_transformer/demo.py @@ -7,7 +7,9 @@ from funasr import AutoModel model = AutoModel(model="iic/punc_ct-transformer_zh-cn-common-vocab272727-pytorch") -res = model.generate(input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_text/punc_example.txt") +res = model.generate( + input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_text/punc_example.txt" +) print(res) @@ -15,5 +17,7 @@ from funasr import AutoModel model = AutoModel(model="iic/punc_ct-transformer_cn-en-common-vocab471067-large") -res = model.generate(input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_text/punc_example.txt") -print(res) \ No newline at end of file +res = model.generate( + input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_text/punc_example.txt" +) +print(res) diff --git a/examples/industrial_data_pretraining/ct_transformer/export.py b/examples/industrial_data_pretraining/ct_transformer/export.py index c429a2ffe..7115728c6 100644 --- a/examples/industrial_data_pretraining/ct_transformer/export.py +++ b/examples/industrial_data_pretraining/ct_transformer/export.py @@ -7,8 +7,9 @@ from funasr import AutoModel -model = AutoModel(model="iic/punc_ct-transformer_zh-cn-common-vocab272727-pytorch", - ) +model = AutoModel( + model="iic/punc_ct-transformer_zh-cn-common-vocab272727-pytorch", +) res = model.export(type="onnx", quantize=False) print(res) @@ -17,7 +18,9 @@ print(res) # method2, inference from local path from funasr import AutoModel -model = AutoModel(model="/Users/zhifu/.cache/modelscope/hub/iic/punc_ct-transformer_zh-cn-common-vocab272727-pytorch") +model = AutoModel( + model="/Users/zhifu/.cache/modelscope/hub/iic/punc_ct-transformer_zh-cn-common-vocab272727-pytorch" +) res = model.export(type="onnx", quantize=False) -print(res) \ No newline at end of file +print(res) diff --git a/examples/industrial_data_pretraining/ct_transformer_streaming/demo.py b/examples/industrial_data_pretraining/ct_transformer_streaming/demo.py index c27a73b54..42e0b9bd9 100644 --- a/examples/industrial_data_pretraining/ct_transformer_streaming/demo.py +++ b/examples/industrial_data_pretraining/ct_transformer_streaming/demo.py @@ -13,6 +13,6 @@ rec_result_all = "outputs: " cache = {} for vad in vads: rec_result = model.generate(input=vad, cache=cache) - rec_result_all += rec_result[0]['text'] + rec_result_all += rec_result[0]["text"] print(rec_result_all) diff --git a/examples/industrial_data_pretraining/ct_transformer_streaming/export.py b/examples/industrial_data_pretraining/ct_transformer_streaming/export.py index ec1557bb1..afab1ae6b 100644 --- a/examples/industrial_data_pretraining/ct_transformer_streaming/export.py +++ b/examples/industrial_data_pretraining/ct_transformer_streaming/export.py @@ -7,8 +7,9 @@ from funasr import AutoModel -model = AutoModel(model="iic/punc_ct-transformer_zh-cn-common-vad_realtime-vocab272727", - ) +model = AutoModel( + model="iic/punc_ct-transformer_zh-cn-common-vad_realtime-vocab272727", +) res = model.export(type="onnx", quantize=False) print(res) @@ -17,7 +18,9 @@ print(res) # method2, inference from local path from funasr import AutoModel -model = AutoModel(model="/Users/zhifu/.cache/modelscope/hub/iic/punc_ct-transformer_zh-cn-common-vad_realtime-vocab272727") +model = AutoModel( + model="/Users/zhifu/.cache/modelscope/hub/iic/punc_ct-transformer_zh-cn-common-vad_realtime-vocab272727" +) res = model.export(type="onnx", quantize=False) -print(res) \ No newline at end of file +print(res) diff --git a/examples/industrial_data_pretraining/emotion2vec/demo.py b/examples/industrial_data_pretraining/emotion2vec/demo.py index 9e926cada..f33dfee1f 100644 --- a/examples/industrial_data_pretraining/emotion2vec/demo.py +++ b/examples/industrial_data_pretraining/emotion2vec/demo.py @@ -6,12 +6,15 @@ from funasr import AutoModel # model="iic/emotion2vec_base" -model = AutoModel(model="iic/emotion2vec_base_finetuned", - # vad_model="iic/speech_fsmn_vad_zh-cn-16k-common-pytorch", - # vad_model_revision="master", - # vad_kwargs={"max_single_segment_time": 2000}, - ) +model = AutoModel( + model="iic/emotion2vec_base_finetuned", + # vad_model="iic/speech_fsmn_vad_zh-cn-16k-common-pytorch", + # vad_model_revision="master", + # vad_kwargs={"max_single_segment_time": 2000}, +) wav_file = f"{model.model_path}/example/test.wav" -res = model.generate(wav_file, output_dir="./outputs", granularity="utterance", extract_embedding=False) -print(res) \ No newline at end of file +res = model.generate( + wav_file, output_dir="./outputs", granularity="utterance", extract_embedding=False +) +print(res) diff --git a/examples/industrial_data_pretraining/fsmn_vad_streaming/demo.py b/examples/industrial_data_pretraining/fsmn_vad_streaming/demo.py index f131c8c5f..b21348fc5 100644 --- a/examples/industrial_data_pretraining/fsmn_vad_streaming/demo.py +++ b/examples/industrial_data_pretraining/fsmn_vad_streaming/demo.py @@ -4,6 +4,7 @@ # MIT License (https://opensource.org/licenses/MIT) from funasr import AutoModel + wav_file = "https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/vad_example.wav" model = AutoModel(model="iic/speech_fsmn_vad_zh-cn-16k-common-pytorch") @@ -14,28 +15,28 @@ print(res) # beg/end: ms - import soundfile import os wav_file = os.path.join(model.model_path, "example/vad_example.wav") speech, sample_rate = soundfile.read(wav_file) -chunk_size = 200 # ms +chunk_size = 200 # ms chunk_stride = int(chunk_size * sample_rate / 1000) cache = {} -total_chunk_num = int(len((speech)-1)/chunk_stride+1) +total_chunk_num = int(len((speech) - 1) / chunk_stride + 1) for i in range(total_chunk_num): - speech_chunk = speech[i*chunk_stride:(i+1)*chunk_stride] + speech_chunk = speech[i * chunk_stride : (i + 1) * chunk_stride] is_final = i == total_chunk_num - 1 - res = model.generate(input=speech_chunk, - cache=cache, - is_final=is_final, - chunk_size=chunk_size, - disable_pbar=True, - ) + res = model.generate( + input=speech_chunk, + cache=cache, + is_final=is_final, + chunk_size=chunk_size, + disable_pbar=True, + ) # print(res) if len(res[0]["value"]): print(res) @@ -44,4 +45,4 @@ for i in range(total_chunk_num): # 1. [[beg1, end1], [beg2, end2], .., [begN, endN]]; [[beg, end]]; [[beg1, end1], [beg2, end2]] # 2. [[beg, -1]] # 3. [[-1, end]] -# beg/end: ms \ No newline at end of file +# beg/end: ms diff --git a/examples/industrial_data_pretraining/fsmn_vad_streaming/export.py b/examples/industrial_data_pretraining/fsmn_vad_streaming/export.py index a3546e19c..38a1798f9 100644 --- a/examples/industrial_data_pretraining/fsmn_vad_streaming/export.py +++ b/examples/industrial_data_pretraining/fsmn_vad_streaming/export.py @@ -17,7 +17,9 @@ print(res) from funasr import AutoModel -model = AutoModel(model="/Users/zhifu/.cache/modelscope/hub/iic/speech_fsmn_vad_zh-cn-16k-common-pytorch") +model = AutoModel( + model="/Users/zhifu/.cache/modelscope/hub/iic/speech_fsmn_vad_zh-cn-16k-common-pytorch" +) res = model.export(type="onnx", quantize=False) print(res) diff --git a/examples/industrial_data_pretraining/lcbnet/compute_wer_details.py b/examples/industrial_data_pretraining/lcbnet/compute_wer_details.py index e72d87155..f672cc85c 100755 --- a/examples/industrial_data_pretraining/lcbnet/compute_wer_details.py +++ b/examples/industrial_data_pretraining/lcbnet/compute_wer_details.py @@ -9,6 +9,7 @@ import argparse from tqdm import tqdm import os import pdb + remove_tag = False spacelist = [" ", "\t", "\r", "\n"] puncts = [ @@ -51,9 +52,9 @@ class WordError(object): def get_wer(self): assert self.ref_words != 0 errors = ( - self.errors[Code.substitution] - + self.errors[Code.insertion] - + self.errors[Code.deletion] + self.errors[Code.substitution] + + self.errors[Code.insertion] + + self.errors[Code.deletion] ) return 100.0 * errors / self.ref_words @@ -299,30 +300,30 @@ def default_cluster(word): for i in reversed(range(len(unicode_names))): if unicode_names[i].startswith("DIGIT"): # 1 unicode_names[i] = "Number" # 'DIGIT' - elif unicode_names[i].startswith("CJK UNIFIED IDEOGRAPH") or unicode_names[ - i - ].startswith("CJK COMPATIBILITY IDEOGRAPH"): + elif unicode_names[i].startswith("CJK UNIFIED IDEOGRAPH") or unicode_names[i].startswith( + "CJK COMPATIBILITY IDEOGRAPH" + ): # 明 / 郎 unicode_names[i] = "Mandarin" # 'CJK IDEOGRAPH' - elif unicode_names[i].startswith("LATIN CAPITAL LETTER") or unicode_names[ - i - ].startswith("LATIN SMALL LETTER"): + elif unicode_names[i].startswith("LATIN CAPITAL LETTER") or unicode_names[i].startswith( + "LATIN SMALL LETTER" + ): # A / a unicode_names[i] = "English" # 'LATIN LETTER' elif unicode_names[i].startswith("HIRAGANA LETTER"): # は こ め unicode_names[i] = "Japanese" # 'GANA LETTER' elif ( - unicode_names[i].startswith("AMPERSAND") - or unicode_names[i].startswith("APOSTROPHE") - or unicode_names[i].startswith("COMMERCIAL AT") - or unicode_names[i].startswith("DEGREE CELSIUS") - or unicode_names[i].startswith("EQUALS SIGN") - or unicode_names[i].startswith("FULL STOP") - or unicode_names[i].startswith("HYPHEN-MINUS") - or unicode_names[i].startswith("LOW LINE") - or unicode_names[i].startswith("NUMBER SIGN") - or unicode_names[i].startswith("PLUS SIGN") - or unicode_names[i].startswith("SEMICOLON") + unicode_names[i].startswith("AMPERSAND") + or unicode_names[i].startswith("APOSTROPHE") + or unicode_names[i].startswith("COMMERCIAL AT") + or unicode_names[i].startswith("DEGREE CELSIUS") + or unicode_names[i].startswith("EQUALS SIGN") + or unicode_names[i].startswith("FULL STOP") + or unicode_names[i].startswith("HYPHEN-MINUS") + or unicode_names[i].startswith("LOW LINE") + or unicode_names[i].startswith("NUMBER SIGN") + or unicode_names[i].startswith("PLUS SIGN") + or unicode_names[i].startswith("SEMICOLON") ): # & / ' / @ / ℃ / = / . / - / _ / # / + / ; del unicode_names[i] @@ -411,11 +412,13 @@ def main(args): if len(array) == 0: continue fid = array[0] - rec_sets[rec_names[i]][fid] = normalize(array[1:], ignore_words, case_sensitive, split) + rec_sets[rec_names[i]][fid] = normalize( + array[1:], ignore_words, case_sensitive, split + ) calculators_dict[rec_names[i]] = Calculator() ub_wer_dict[rec_names[i]] = {"u_wer": WordError(), "b_wer": WordError(), "wer": WordError()} - hotwords_related_dict[rec_names[i]] = {'tp': 0, 'tn': 0, 'fp': 0, 'fn': 0} + hotwords_related_dict[rec_names[i]] = {"tp": 0, "tn": 0, "fp": 0, "fn": 0} # tp: 热词在label里，同时在rec里 # tn: 热词不在label里，同时不在rec里 # fp: 热词不在label里，但是在rec里 @@ -431,21 +434,22 @@ def main(args): _file_total_len = int(pipe.read().strip()) # compute error rate on the interaction of reference file and hyp file - for line in tqdm(open(ref_file, 'r', encoding='utf-8'), total=_file_total_len): + for line in tqdm(open(ref_file, "r", encoding="utf-8"), total=_file_total_len): if tochar: array = characterize(line) else: - array = line.rstrip('\n').split() - if len(array) == 0: continue + array = line.rstrip("\n").split() + if len(array) == 0: + continue fid = array[0] lab = normalize(array[1:], ignore_words, case_sensitive, split) if verbose: - print('\nutt: %s' % fid) + print("\nutt: %s" % fid) ocr_text = ref_ocr_dict[fid] ocr_set = set(ocr_text) - print('ocr: {}'.format(" ".join(ocr_text))) + print("ocr: {}".format(" ".join(ocr_text))) list_match = [] # 指label里面在ocr里面的内容 list_not_mathch = [] tmp_error = 0 @@ -458,7 +462,7 @@ def main(args): else: tmp_match += 1 list_match.append(lab[index]) - print('label in ocr: {}'.format(" ".join(list_match))) + print("label in ocr: {}".format(" ".join(list_match))) # for each reco file base_wrong_ocr_wer = None @@ -482,33 +486,44 @@ def main(args): result = calculators_dict[rec_name].calculate(lab.copy(), rec.copy()) if verbose: - if result['all'] != 0: - wer = float(result['ins'] + result['sub'] + result['del']) * 100.0 / result['all'] + if result["all"] != 0: + wer = ( + float(result["ins"] + result["sub"] + result["del"]) * 100.0 / result["all"] + ) else: wer = 0.0 - print('WER(%s): %4.2f %%' % (rec_name, wer), end=' ') - print('N=%d C=%d S=%d D=%d I=%d' % - (result['all'], result['cor'], result['sub'], result['del'], result['ins'])) - + print("WER(%s): %4.2f %%" % (rec_name, wer), end=" ") + print( + "N=%d C=%d S=%d D=%d I=%d" + % (result["all"], result["cor"], result["sub"], result["del"], result["ins"]) + ) # print(result['rec']) wrong_rec_but_in_ocr = [] - for idx in range(len(result['lab'])): - if result['lab'][idx] != "": - if result['lab'][idx] != result['rec'][idx].replace("", ""): - if result['lab'][idx] in list_match: - wrong_rec_but_in_ocr.append(result['lab'][idx]) + for idx in range(len(result["lab"])): + if result["lab"][idx] != "": + if result["lab"][idx] != result["rec"][idx].replace("", ""): + if result["lab"][idx] in list_match: + wrong_rec_but_in_ocr.append(result["lab"][idx]) wrong_rec_but_in_ocr_dict[rec_name] += 1 - print('wrong_rec_but_in_ocr: {}'.format(" ".join(wrong_rec_but_in_ocr))) + print("wrong_rec_but_in_ocr: {}".format(" ".join(wrong_rec_but_in_ocr))) if rec_name == "base": base_wrong_ocr_wer = len(wrong_rec_but_in_ocr) if "ocr" in rec_name or "hot" in rec_name: ocr_wrong_ocr_wer = len(wrong_rec_but_in_ocr) if ocr_wrong_ocr_wer < base_wrong_ocr_wer: - print("{} {} helps, {} -> {}".format(fid, rec_name, base_wrong_ocr_wer, ocr_wrong_ocr_wer)) + print( + "{} {} helps, {} -> {}".format( + fid, rec_name, base_wrong_ocr_wer, ocr_wrong_ocr_wer + ) + ) elif ocr_wrong_ocr_wer > base_wrong_ocr_wer: - print("{} {} hurts, {} -> {}".format(fid, rec_name, base_wrong_ocr_wer, ocr_wrong_ocr_wer)) + print( + "{} {} hurts, {} -> {}".format( + fid, rec_name, base_wrong_ocr_wer, ocr_wrong_ocr_wer + ) + ) # recall = 0 # false_alarm = 0 @@ -537,11 +552,11 @@ def main(args): # if badhotword == word: # count += 1 if count == 0: - hotwords_related_dict[rec_name]['tn'] += 1 + hotwords_related_dict[rec_name]["tn"] += 1 _tn += 1 # fp: 0 else: - hotwords_related_dict[rec_name]['fp'] += count + hotwords_related_dict[rec_name]["fp"] += count _fp += count # tn: 0 # if badhotword in _rec_list: @@ -553,23 +568,30 @@ def main(args): rec_count = len([word for word in _rec_list if hotword == word]) # print(f"good {hotword} true_count: {true_count}, rec_count: {rec_count}") if rec_count == true_count: - hotwords_related_dict[rec_name]['tp'] += true_count + hotwords_related_dict[rec_name]["tp"] += true_count _tp += true_count elif rec_count > true_count: - hotwords_related_dict[rec_name]['tp'] += true_count + hotwords_related_dict[rec_name]["tp"] += true_count # fp: 不在label里，但是在rec里 - hotwords_related_dict[rec_name]['fp'] += rec_count - true_count + hotwords_related_dict[rec_name]["fp"] += rec_count - true_count _tp += true_count _fp += rec_count - true_count else: - hotwords_related_dict[rec_name]['tp'] += rec_count + hotwords_related_dict[rec_name]["tp"] += rec_count # fn: 热词在label里，但是不在rec里 - hotwords_related_dict[rec_name]['fn'] += true_count - rec_count + hotwords_related_dict[rec_name]["fn"] += true_count - rec_count _tp += rec_count _fn += true_count - rec_count - print("hotword: tp: {}, tn: {}, fp: {}, fn: {}, all: {}, recall: {:.2f}%".format( - _tp, _tn, _fp, _fn, sum([_tp, _tn, _fp, _fn]), _tp / (_tp + _fn) * 100 if (_tp + _fn) != 0 else 0 - )) + print( + "hotword: tp: {}, tn: {}, fp: {}, fn: {}, all: {}, recall: {:.2f}%".format( + _tp, + _tn, + _fp, + _fn, + sum([_tp, _tn, _fp, _fn]), + _tp / (_tp + _fn) * 100 if (_tp + _fn) != 0 else 0, + ) + ) # if hotword in _rec_list: # hotwords_related_dict[rec_name]['tp'] += 1 @@ -612,77 +634,89 @@ def main(args): ub_wer_dict[rec_name]["u_wer"].errors[Code.insertion] += 1 space = {} - space['lab'] = [] - space['rec'] = [] - for idx in range(len(result['lab'])): - len_lab = width(result['lab'][idx]) - len_rec = width(result['rec'][idx]) + space["lab"] = [] + space["rec"] = [] + for idx in range(len(result["lab"])): + len_lab = width(result["lab"][idx]) + len_rec = width(result["rec"][idx]) length = max(len_lab, len_rec) - space['lab'].append(length - len_lab) - space['rec'].append(length - len_rec) - upper_lab = len(result['lab']) - upper_rec = len(result['rec']) + space["lab"].append(length - len_lab) + space["rec"].append(length - len_rec) + upper_lab = len(result["lab"]) + upper_rec = len(result["rec"]) lab1, rec1 = 0, 0 while lab1 < upper_lab or rec1 < upper_rec: if verbose > 1: - print('lab(%s):' % fid.encode('utf-8'), end=' ') + print("lab(%s):" % fid.encode("utf-8"), end=" ") else: - print('lab:', end=' ') + print("lab:", end=" ") lab2 = min(upper_lab, lab1 + max_words_per_line) for idx in range(lab1, lab2): - token = result['lab'][idx] - print('{token}'.format(token=token), end='') - for n in range(space['lab'][idx]): - print(padding_symbol, end='') - print(' ', end='') + token = result["lab"][idx] + print("{token}".format(token=token), end="") + for n in range(space["lab"][idx]): + print(padding_symbol, end="") + print(" ", end="") print() if verbose > 1: - print('rec(%s):' % fid.encode('utf-8'), end=' ') + print("rec(%s):" % fid.encode("utf-8"), end=" ") else: - print('rec:', end=' ') + print("rec:", end=" ") rec2 = min(upper_rec, rec1 + max_words_per_line) for idx in range(rec1, rec2): - token = result['rec'][idx] - print('{token}'.format(token=token), end='') - for n in range(space['rec'][idx]): - print(padding_symbol, end='') - print(' ', end='') + token = result["rec"][idx] + print("{token}".format(token=token), end="") + for n in range(space["rec"][idx]): + print(padding_symbol, end="") + print(" ", end="") print() # print('\n', end='\n') lab1 = lab2 rec1 = rec2 - print('\n', end='\n') + print("\n", end="\n") # break if verbose: - print('===========================================================================') + print("===========================================================================") print() print(wrong_rec_but_in_ocr_dict) for rec_name in rec_names: result = calculators_dict[rec_name].overall() - if result['all'] != 0: - wer = float(result['ins'] + result['sub'] + result['del']) * 100.0 / result['all'] + if result["all"] != 0: + wer = float(result["ins"] + result["sub"] + result["del"]) * 100.0 / result["all"] else: wer = 0.0 - print('{} Overall -> {:4.2f} %'.format(rec_name, wer), end=' ') - print('N=%d C=%d S=%d D=%d I=%d' % - (result['all'], result['cor'], result['sub'], result['del'], result['ins'])) + print("{} Overall -> {:4.2f} %".format(rec_name, wer), end=" ") + print( + "N=%d C=%d S=%d D=%d I=%d" + % (result["all"], result["cor"], result["sub"], result["del"], result["ins"]) + ) print(f"WER: {ub_wer_dict[rec_name]['wer'].get_result_string()}") print(f"U-WER: {ub_wer_dict[rec_name]['u_wer'].get_result_string()}") print(f"B-WER: {ub_wer_dict[rec_name]['b_wer'].get_result_string()}") - print('hotword: tp: {}, tn: {}, fp: {}, fn: {}, all: {}, recall: {:.2f}%'.format( - hotwords_related_dict[rec_name]['tp'], - hotwords_related_dict[rec_name]['tn'], - hotwords_related_dict[rec_name]['fp'], - hotwords_related_dict[rec_name]['fn'], - sum([v for k, v in hotwords_related_dict[rec_name].items()]), - hotwords_related_dict[rec_name]['tp'] / ( - hotwords_related_dict[rec_name]['tp'] + hotwords_related_dict[rec_name]['fn'] - ) * 100 if hotwords_related_dict[rec_name]['tp'] + hotwords_related_dict[rec_name]['fn'] != 0 else 0 - )) + print( + "hotword: tp: {}, tn: {}, fp: {}, fn: {}, all: {}, recall: {:.2f}%".format( + hotwords_related_dict[rec_name]["tp"], + hotwords_related_dict[rec_name]["tn"], + hotwords_related_dict[rec_name]["fp"], + hotwords_related_dict[rec_name]["fn"], + sum([v for k, v in hotwords_related_dict[rec_name].items()]), + ( + hotwords_related_dict[rec_name]["tp"] + / ( + hotwords_related_dict[rec_name]["tp"] + + hotwords_related_dict[rec_name]["fn"] + ) + * 100 + if hotwords_related_dict[rec_name]["tp"] + hotwords_related_dict[rec_name]["fn"] + != 0 + else 0 + ), + ) + ) # tp: 热词在label里，同时在rec里 # tn: 热词不在label里，同时不在rec里 @@ -695,8 +729,7 @@ def main(args): if __name__ == "__main__": args = get_args() - + # print("") print(args) main(args) - diff --git a/examples/industrial_data_pretraining/lcbnet/demo.py b/examples/industrial_data_pretraining/lcbnet/demo.py index 8b3d493e6..39e11eeb6 100755 --- a/examples/industrial_data_pretraining/lcbnet/demo.py +++ b/examples/industrial_data_pretraining/lcbnet/demo.py @@ -5,11 +5,15 @@ from funasr import AutoModel -model = AutoModel(model="iic/LCB-NET", - model_revision="v1.0.0") +model = AutoModel(model="iic/LCB-NET", model_revision="v1.0.0") - -res = model.generate(input=("https://www.modelscope.cn/api/v1/models/iic/LCB-NET/repo?Revision=master&FilePath=example/asr_example.wav","https://www.modelscope.cn/api/v1/models/iic/LCB-NET/repo?Revision=master&FilePath=example/ocr.txt"),data_type=("sound", "text")) +res = model.generate( + input=( + "https://www.modelscope.cn/api/v1/models/iic/LCB-NET/repo?Revision=master&FilePath=example/asr_example.wav", + "https://www.modelscope.cn/api/v1/models/iic/LCB-NET/repo?Revision=master&FilePath=example/ocr.txt", + ), + data_type=("sound", "text"), +) print(res) diff --git a/examples/industrial_data_pretraining/monotonic_aligner/demo.py b/examples/industrial_data_pretraining/monotonic_aligner/demo.py index 345076ec0..35ddb899c 100644 --- a/examples/industrial_data_pretraining/monotonic_aligner/demo.py +++ b/examples/industrial_data_pretraining/monotonic_aligner/demo.py @@ -7,9 +7,12 @@ from funasr import AutoModel model = AutoModel(model="iic/speech_timestamp_prediction-v1-16k-offline") -res = model.generate(input=("https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav", - "欢迎大家来到魔搭社区进行体验"), - data_type=("sound", "text"), - batch_size=2, - ) +res = model.generate( + input=( + "https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav", + "欢迎大家来到魔搭社区进行体验", + ), + data_type=("sound", "text"), + batch_size=2, +) print(res) diff --git a/examples/industrial_data_pretraining/paraformer-zh-spk/demo.py b/examples/industrial_data_pretraining/paraformer-zh-spk/demo.py index cbc9d8baa..9afd801f3 100644 --- a/examples/industrial_data_pretraining/paraformer-zh-spk/demo.py +++ b/examples/industrial_data_pretraining/paraformer-zh-spk/demo.py @@ -5,12 +5,15 @@ from funasr import AutoModel -model = AutoModel(model="iic/speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch", - vad_model="iic/speech_fsmn_vad_zh-cn-16k-common-pytorch", - punc_model="iic/punc_ct-transformer_zh-cn-common-vocab272727-pytorch", - spk_model="iic/speech_campplus_sv_zh-cn_16k-common", - ) +model = AutoModel( + model="iic/speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch", + vad_model="iic/speech_fsmn_vad_zh-cn-16k-common-pytorch", + punc_model="iic/punc_ct-transformer_zh-cn-common-vocab272727-pytorch", + spk_model="iic/speech_campplus_sv_zh-cn_16k-common", +) -res = model.generate(input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav", - hotword='达摩院 磨搭') -print(res) \ No newline at end of file +res = model.generate( + input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav", + hotword="达摩院 磨搭", +) +print(res) diff --git a/examples/industrial_data_pretraining/paraformer/demo.py b/examples/industrial_data_pretraining/paraformer/demo.py index 3c409a5c3..f6f4c7530 100644 --- a/examples/industrial_data_pretraining/paraformer/demo.py +++ b/examples/industrial_data_pretraining/paraformer/demo.py @@ -5,22 +5,29 @@ from funasr import AutoModel -model = AutoModel(model="iic/speech_paraformer-large-vad-punc_asr_nat-zh-cn-16k-common-vocab8404-pytorch", - vad_model="iic/speech_fsmn_vad_zh-cn-16k-common-pytorch", - vad_kwargs={"max_single_segment_time": 60000}, - punc_model="iic/punc_ct-transformer_zh-cn-common-vocab272727-pytorch", - # spk_model="iic/speech_campplus_sv_zh-cn_16k-common", - ) +model = AutoModel( + model="iic/speech_paraformer-large-vad-punc_asr_nat-zh-cn-16k-common-vocab8404-pytorch", + vad_model="iic/speech_fsmn_vad_zh-cn-16k-common-pytorch", + vad_kwargs={"max_single_segment_time": 60000}, + punc_model="iic/punc_ct-transformer_zh-cn-common-vocab272727-pytorch", + # spk_model="iic/speech_campplus_sv_zh-cn_16k-common", +) -res = model.generate(input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav") -res = model.generate(input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav") +res = model.generate( + input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav" +) +res = model.generate( + input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav" +) -res = model.generate(input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav") +res = model.generate( + input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav" +) print(res) -''' can not use currently +""" can not use currently from funasr import AutoFrontend frontend = AutoFrontend(model="iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch") @@ -30,4 +37,4 @@ fbanks = frontend(input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ for batch_idx, fbank_dict in enumerate(fbanks): res = model.generate(**fbank_dict) print(res) -''' \ No newline at end of file +""" diff --git a/examples/industrial_data_pretraining/paraformer/export.py b/examples/industrial_data_pretraining/paraformer/export.py index 1c86927c3..19512c1f3 100644 --- a/examples/industrial_data_pretraining/paraformer/export.py +++ b/examples/industrial_data_pretraining/paraformer/export.py @@ -9,7 +9,9 @@ from funasr import AutoModel -model = AutoModel(model="iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch",) +model = AutoModel( + model="iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch", +) res = model.export(type="onnx", quantize=False) print(res) @@ -18,7 +20,9 @@ print(res) # method2, inference from local path from funasr import AutoModel -model = AutoModel(model="/Users/zhifu/.cache/modelscope/hub/iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch") +model = AutoModel( + model="/Users/zhifu/.cache/modelscope/hub/iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch" +) res = model.export(type="onnx", quantize=False) -print(res) \ No newline at end of file +print(res) diff --git a/examples/industrial_data_pretraining/paraformer_streaming/demo.py b/examples/industrial_data_pretraining/paraformer_streaming/demo.py index 57356b82b..a1e988266 100644 --- a/examples/industrial_data_pretraining/paraformer_streaming/demo.py +++ b/examples/industrial_data_pretraining/paraformer_streaming/demo.py @@ -7,17 +7,18 @@ import os from funasr import AutoModel -chunk_size = [0, 10, 5] #[0, 10, 5] 600ms, [0, 8, 4] 480ms -encoder_chunk_look_back = 4 #number of chunks to lookback for encoder self-attention -decoder_chunk_look_back = 1 #number of encoder chunks to lookback for decoder cross-attention +chunk_size = [0, 10, 5] # [0, 10, 5] 600ms, [0, 8, 4] 480ms +encoder_chunk_look_back = 4 # number of chunks to lookback for encoder self-attention +decoder_chunk_look_back = 1 # number of encoder chunks to lookback for decoder cross-attention model = AutoModel(model="iic/speech_paraformer_asr_nat-zh-cn-16k-common-vocab8404-online") wav_file = os.path.join(model.model_path, "example/asr_example.wav") -res = model.generate(input=wav_file, - chunk_size=chunk_size, - encoder_chunk_look_back=encoder_chunk_look_back, - decoder_chunk_look_back=decoder_chunk_look_back, - ) +res = model.generate( + input=wav_file, + chunk_size=chunk_size, + encoder_chunk_look_back=encoder_chunk_look_back, + decoder_chunk_look_back=decoder_chunk_look_back, +) print(res) @@ -27,18 +28,19 @@ import soundfile wav_file = os.path.join(model.model_path, "example/asr_example.wav") speech, sample_rate = soundfile.read(wav_file) -chunk_stride = chunk_size[1] * 960 # 600ms、480ms +chunk_stride = chunk_size[1] * 960 # 600ms、480ms cache = {} -total_chunk_num = int(len((speech)-1)/chunk_stride+1) +total_chunk_num = int(len((speech) - 1) / chunk_stride + 1) for i in range(total_chunk_num): - speech_chunk = speech[i*chunk_stride:(i+1)*chunk_stride] + speech_chunk = speech[i * chunk_stride : (i + 1) * chunk_stride] is_final = i == total_chunk_num - 1 - res = model.generate(input=speech_chunk, - cache=cache, - is_final=is_final, - chunk_size=chunk_size, - encoder_chunk_look_back=encoder_chunk_look_back, - decoder_chunk_look_back=decoder_chunk_look_back, - ) + res = model.generate( + input=speech_chunk, + cache=cache, + is_final=is_final, + chunk_size=chunk_size, + encoder_chunk_look_back=encoder_chunk_look_back, + decoder_chunk_look_back=decoder_chunk_look_back, + ) print(res) diff --git a/examples/industrial_data_pretraining/paraformer_streaming/export.py b/examples/industrial_data_pretraining/paraformer_streaming/export.py index 95cc56b10..588fd0838 100644 --- a/examples/industrial_data_pretraining/paraformer_streaming/export.py +++ b/examples/industrial_data_pretraining/paraformer_streaming/export.py @@ -9,7 +9,9 @@ from funasr import AutoModel -model = AutoModel(model="iic/speech_paraformer_asr_nat-zh-cn-16k-common-vocab8404-online", ) +model = AutoModel( + model="iic/speech_paraformer_asr_nat-zh-cn-16k-common-vocab8404-online", +) res = model.export(type="onnx", quantize=False) print(res) @@ -19,7 +21,9 @@ print(res) from funasr import AutoModel -model = AutoModel(model="/Users/zhifu/.cache/modelscope/hub/iic/speech_paraformer_asr_nat-zh-cn-16k-common-vocab8404-online") +model = AutoModel( + model="/Users/zhifu/.cache/modelscope/hub/iic/speech_paraformer_asr_nat-zh-cn-16k-common-vocab8404-online" +) res = model.export(type="onnx", quantize=False) -print(res) \ No newline at end of file +print(res) diff --git a/examples/industrial_data_pretraining/qwen_audio/demo_chat.py b/examples/industrial_data_pretraining/qwen_audio/demo_chat.py index a2fb05be0..6a69fa1cd 100644 --- a/examples/industrial_data_pretraining/qwen_audio/demo_chat.py +++ b/examples/industrial_data_pretraining/qwen_audio/demo_chat.py @@ -12,7 +12,7 @@ model = AutoModel(model="Qwen-Audio-Chat") audio_in = "https://github.com/QwenLM/Qwen-Audio/raw/main/assets/audio/1272-128104-0000.flac" # 1st dialogue turn -prompt = 'what does the person say?' +prompt = "what does the person say?" cache = {"history": None} res = model.generate(input=audio_in, prompt=prompt, cache=cache) print(res) @@ -22,4 +22,3 @@ print(res) prompt = 'Find the start time and end time of the word "middle classes"' res = model.generate(input=None, prompt=prompt, cache=cache) print(res) - diff --git a/examples/industrial_data_pretraining/qwen_audio/demo_chat_from_local.py b/examples/industrial_data_pretraining/qwen_audio/demo_chat_from_local.py index 233321765..99d4caf55 100644 --- a/examples/industrial_data_pretraining/qwen_audio/demo_chat_from_local.py +++ b/examples/industrial_data_pretraining/qwen_audio/demo_chat_from_local.py @@ -7,14 +7,17 @@ from funasr import AutoModel -model = AutoModel(model="Qwen-Audio-Chat", - model_path="/nfs/zhifu.gzf/init_model/qwen/Qwen-Audio-Chat", - ) +model = AutoModel( + model="Qwen-Audio-Chat", + model_path="/nfs/zhifu.gzf/init_model/qwen/Qwen-Audio-Chat", +) -audio_in = "https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav" +audio_in = ( + "https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav" +) # 1st dialogue turn -prompt = 'what does the person say?' +prompt = "what does the person say?" cache = {"history": None} res = model.generate(input=audio_in, prompt=prompt, cache=cache) print(res) @@ -24,4 +27,3 @@ print(res) prompt = 'Find the start time and end time of the word "middle classes"' res = model.generate(input=None, prompt=prompt, cache=cache) print(res) - diff --git a/examples/industrial_data_pretraining/qwen_audio/demo_from_local.py b/examples/industrial_data_pretraining/qwen_audio/demo_from_local.py index 40cf4b941..5ade20518 100644 --- a/examples/industrial_data_pretraining/qwen_audio/demo_from_local.py +++ b/examples/industrial_data_pretraining/qwen_audio/demo_from_local.py @@ -7,9 +7,10 @@ from funasr import AutoModel -model = AutoModel(model="Qwen-Audio", - model_path="/nfs/zhifu.gzf/init_model/qwen/Qwen-Audio", - ) +model = AutoModel( + model="Qwen-Audio", + model_path="/nfs/zhifu.gzf/init_model/qwen/Qwen-Audio", +) audio_in = "https://github.com/QwenLM/Qwen-Audio/raw/main/assets/audio/1272-128104-0000.flac" prompt = "<|startoftranscription|><|en|><|transcribe|><|en|><|notimestamps|><|wo_itn|>" diff --git a/examples/industrial_data_pretraining/scama/demo.py b/examples/industrial_data_pretraining/scama/demo.py index 9d458ab94..00d85ae2d 100644 --- a/examples/industrial_data_pretraining/scama/demo.py +++ b/examples/industrial_data_pretraining/scama/demo.py @@ -5,17 +5,20 @@ from funasr import AutoModel -chunk_size = [5, 10, 5] #[0, 10, 5] 600ms, [0, 8, 4] 480ms -encoder_chunk_look_back = 0 #number of chunks to lookback for encoder self-attention -decoder_chunk_look_back = 0 #number of encoder chunks to lookback for decoder cross-attention +chunk_size = [5, 10, 5] # [0, 10, 5] 600ms, [0, 8, 4] 480ms +encoder_chunk_look_back = 0 # number of chunks to lookback for encoder self-attention +decoder_chunk_look_back = 0 # number of encoder chunks to lookback for decoder cross-attention -model = AutoModel(model="/Users/zhifu/Downloads/modelscope_models/speech_SCAMA_asr-zh-cn-16k-common-vocab8358-streaming") +model = AutoModel( + model="/Users/zhifu/Downloads/modelscope_models/speech_SCAMA_asr-zh-cn-16k-common-vocab8358-streaming" +) cache = {} -res = model.generate(input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav", - chunk_size=chunk_size, - encoder_chunk_look_back=encoder_chunk_look_back, - decoder_chunk_look_back=decoder_chunk_look_back, - ) +res = model.generate( + input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav", + chunk_size=chunk_size, + encoder_chunk_look_back=encoder_chunk_look_back, + decoder_chunk_look_back=decoder_chunk_look_back, +) print(res) @@ -25,18 +28,19 @@ import os wav_file = os.path.join(model.model_path, "example/asr_example.wav") speech, sample_rate = soundfile.read(wav_file) -chunk_stride = chunk_size[1] * 960 # 600ms、480ms +chunk_stride = chunk_size[1] * 960 # 600ms、480ms cache = {} -total_chunk_num = int(len((speech)-1)/chunk_stride+1) +total_chunk_num = int(len((speech) - 1) / chunk_stride + 1) for i in range(total_chunk_num): - speech_chunk = speech[i*chunk_stride:(i+1)*chunk_stride] + speech_chunk = speech[i * chunk_stride : (i + 1) * chunk_stride] is_final = i == total_chunk_num - 1 - res = model.generate(input=speech_chunk, - cache=cache, - is_final=is_final, - chunk_size=chunk_size, - encoder_chunk_look_back=encoder_chunk_look_back, - decoder_chunk_look_back=decoder_chunk_look_back, - ) + res = model.generate( + input=speech_chunk, + cache=cache, + is_final=is_final, + chunk_size=chunk_size, + encoder_chunk_look_back=encoder_chunk_look_back, + decoder_chunk_look_back=decoder_chunk_look_back, + ) print(res) diff --git a/examples/industrial_data_pretraining/seaco_paraformer/demo.py b/examples/industrial_data_pretraining/seaco_paraformer/demo.py index c7f78d3f4..a88e880df 100644 --- a/examples/industrial_data_pretraining/seaco_paraformer/demo.py +++ b/examples/industrial_data_pretraining/seaco_paraformer/demo.py @@ -5,24 +5,26 @@ from funasr import AutoModel -model = AutoModel(model="iic/speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch", - # vad_model="iic/speech_fsmn_vad_zh-cn-16k-common-pytorch", - # punc_model="iic/punc_ct-transformer_zh-cn-common-vocab272727-pytorch", - # spk_model="iic/speech_campplus_sv_zh-cn_16k-common", - ) +model = AutoModel( + model="iic/speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch", + # vad_model="iic/speech_fsmn_vad_zh-cn-16k-common-pytorch", + # punc_model="iic/punc_ct-transformer_zh-cn-common-vocab272727-pytorch", + # spk_model="iic/speech_campplus_sv_zh-cn_16k-common", +) # example1 -res = model.generate(input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav", - hotword='达摩院 魔搭', - # return_raw_text=True, # return raw text recognition results splited by space of equal length with timestamp - # preset_spk_num=2, # preset speaker num for speaker cluster model - # sentence_timestamp=True, # return sentence level information when spk_model is not given - ) +res = model.generate( + input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav", + hotword="达摩院 魔搭", + # return_raw_text=True, # return raw text recognition results splited by space of equal length with timestamp + # preset_spk_num=2, # preset speaker num for speaker cluster model + # sentence_timestamp=True, # return sentence level information when spk_model is not given +) print(res) -''' +""" # tensor or numpy as input # example2 import torchaudio @@ -39,4 +41,4 @@ import soundfile wav_file = os.path.join(model.model_path, "example/asr_example.wav") speech, sample_rate = soundfile.read(wav_file) res = model.generate(input=[speech], batch_size_s=300, is_final=True) -''' +""" diff --git a/examples/industrial_data_pretraining/sense_voice/demo.py b/examples/industrial_data_pretraining/sense_voice/demo.py index 0d8ef9756..ed583f0f8 100644 --- a/examples/industrial_data_pretraining/sense_voice/demo.py +++ b/examples/industrial_data_pretraining/sense_voice/demo.py @@ -5,20 +5,23 @@ from funasr import AutoModel -model = AutoModel(model="/Users/zhifu/Downloads/modelscope_models/SenseVoiceModelscope", - vad_model="iic/speech_fsmn_vad_zh-cn-16k-common-pytorch", - vad_kwargs={"max_single_segment_time": 30000}, - ) +model = AutoModel( + model="/Users/zhifu/Downloads/modelscope_models/SenseVoiceModelscope", + vad_model="iic/speech_fsmn_vad_zh-cn-16k-common-pytorch", + vad_kwargs={"max_single_segment_time": 30000}, +) -input_wav = "https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav" +input_wav = ( + "https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav" +) DecodingOptions = { - "task": ("ASR", "AED", "SER"), - "language": "auto", - "fp16": True, - "gain_event": True, - } + "task": ("ASR", "AED", "SER"), + "language": "auto", + "fp16": True, + "gain_event": True, +} res = model.generate(input=input_wav, batch_size_s=0, DecodingOptions=DecodingOptions) print(res) diff --git a/examples/industrial_data_pretraining/transducer/demo.py b/examples/industrial_data_pretraining/transducer/demo.py index 5a4aa4240..0b3d63856 100644 --- a/examples/industrial_data_pretraining/transducer/demo.py +++ b/examples/industrial_data_pretraining/transducer/demo.py @@ -7,8 +7,11 @@ from funasr import AutoModel # Transducer, BAT and RWKV_BAT models are just same to use, use the correct model_revision # https://modelscope.cn/models?name=transducer&page=1&tasks=auto-speech-recognition&type=audio -model = AutoModel(model="iic/speech_bat_asr-zh-cn-16k-aishell1-vocab4234-pytorch", - ) +model = AutoModel( + model="iic/speech_bat_asr-zh-cn-16k-aishell1-vocab4234-pytorch", +) -res = model.generate(input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav") -print(res) \ No newline at end of file +res = model.generate( + input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav" +) +print(res) diff --git a/examples/industrial_data_pretraining/uniasr/demo.py b/examples/industrial_data_pretraining/uniasr/demo.py index 774f42fe5..d3741b8c2 100644 --- a/examples/industrial_data_pretraining/uniasr/demo.py +++ b/examples/industrial_data_pretraining/uniasr/demo.py @@ -6,14 +6,18 @@ from funasr import AutoModel -model = AutoModel(model="iic/speech_UniASR-large_asr_2pass-zh-cn-16k-common-vocab8358-tensorflow1-offline",) +model = AutoModel( + model="iic/speech_UniASR-large_asr_2pass-zh-cn-16k-common-vocab8358-tensorflow1-offline", +) -res = model.generate(input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav") +res = model.generate( + input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav" +) print(res) -''' can not use currently +""" can not use currently from funasr import AutoFrontend frontend = AutoFrontend(model="iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch") @@ -23,4 +27,4 @@ fbanks = frontend(input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ for batch_idx, fbank_dict in enumerate(fbanks): res = model.generate(**fbank_dict) print(res) -''' \ No newline at end of file +""" diff --git a/examples/industrial_data_pretraining/whisper/demo.py b/examples/industrial_data_pretraining/whisper/demo.py index c94678fe2..e9c2a8320 100644 --- a/examples/industrial_data_pretraining/whisper/demo.py +++ b/examples/industrial_data_pretraining/whisper/demo.py @@ -7,22 +7,24 @@ from funasr import AutoModel -model = AutoModel(model="iic/Whisper-large-v3", - vad_model="iic/speech_fsmn_vad_zh-cn-16k-common-pytorch", - vad_kwargs={"max_single_segment_time": 30000}, - ) +model = AutoModel( + model="iic/Whisper-large-v3", + vad_model="iic/speech_fsmn_vad_zh-cn-16k-common-pytorch", + vad_kwargs={"max_single_segment_time": 30000}, +) DecodingOptions = { - "task": "transcribe", - "language": None, - "beam_size": None, - "fp16": True, - "without_timestamps": False, - "prompt": None, - } + "task": "transcribe", + "language": None, + "beam_size": None, + "fp16": True, + "without_timestamps": False, + "prompt": None, +} res = model.generate( - DecodingOptions=DecodingOptions, - batch_size_s=0, - input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav") + DecodingOptions=DecodingOptions, + batch_size_s=0, + input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav", +) print(res) diff --git a/examples/industrial_data_pretraining/whisper/demo_from_openai.py b/examples/industrial_data_pretraining/whisper/demo_from_openai.py index a257bc820..097e942e8 100644 --- a/examples/industrial_data_pretraining/whisper/demo_from_openai.py +++ b/examples/industrial_data_pretraining/whisper/demo_from_openai.py @@ -10,15 +10,25 @@ from funasr import AutoModel # model = AutoModel(model="Whisper-small", hub="openai") # model = AutoModel(model="Whisper-medium", hub="openai") # model = AutoModel(model="Whisper-large-v2", hub="openai") -model = AutoModel(model="Whisper-large-v3", - vad_model="iic/speech_fsmn_vad_zh-cn-16k-common-pytorch", - vad_kwargs={"max_single_segment_time": 30000}, - hub="openai", - ) +model = AutoModel( + model="Whisper-large-v3", + vad_model="iic/speech_fsmn_vad_zh-cn-16k-common-pytorch", + vad_kwargs={"max_single_segment_time": 30000}, + hub="openai", +) +DecodingOptions = { + "task": "transcribe", + "language": None, + "beam_size": None, + "fp16": True, + "without_timestamps": False, + "prompt": None, +} res = model.generate( - language=None, - task="transcribe", - batch_size_s=0, - input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav") + DecodingOptions=DecodingOptions, + batch_size_s=0, + input="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav", +) + print(res) diff --git a/fun_text_processing/inverse_text_normalization/de/__init__.py b/fun_text_processing/inverse_text_normalization/de/__init__.py index 68b10481a..4f73ff2c3 100644 --- a/fun_text_processing/inverse_text_normalization/de/__init__.py +++ b/fun_text_processing/inverse_text_normalization/de/__init__.py @@ -1,5 +1,7 @@ - - -from fun_text_processing.inverse_text_normalization.en.taggers.tokenize_and_classify import ClassifyFst +from fun_text_processing.inverse_text_normalization.en.taggers.tokenize_and_classify import ( + ClassifyFst, +) from fun_text_processing.inverse_text_normalization.en.verbalizers.verbalize import VerbalizeFst -from fun_text_processing.inverse_text_normalization.en.verbalizers.verbalize_final import VerbalizeFinalFst +from fun_text_processing.inverse_text_normalization.en.verbalizers.verbalize_final import ( + VerbalizeFinalFst, +) diff --git a/fun_text_processing/inverse_text_normalization/de/taggers/cardinal.py b/fun_text_processing/inverse_text_normalization/de/taggers/cardinal.py index 61058df94..1aadad81a 100644 --- a/fun_text_processing/inverse_text_normalization/de/taggers/cardinal.py +++ b/fun_text_processing/inverse_text_normalization/de/taggers/cardinal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_SIGMA, GraphFst from pynini.lib import pynutil @@ -6,7 +5,7 @@ from pynini.lib import pynutil class CardinalFst(GraphFst): """ - Finite state transducer for classifying cardinals. Numbers below ten are not converted. + Finite state transducer for classifying cardinals. Numbers below ten are not converted. Allows both compound numeral strings or separated by whitespace. "und" (en: "and") can be inserted between "hundert" and following number or "tausend" and following single or double digit number. @@ -18,7 +17,7 @@ class CardinalFst(GraphFst): e.g. ein tausend -> cardinal { integer: "1000" } } e.g. eintausend -> cardinal { integer: "1000" } } e.g. ein tausend zwanzig -> cardinal { integer: "1020" } } - + Args: tn_cardinal_tagger: TN cardinal tagger """ @@ -31,24 +30,36 @@ class CardinalFst(GraphFst): graph = (tn_cardinal_tagger.graph @ optional_delete_space).invert().optimize() self.graph_hundred_component_at_least_one_none_zero_digit = ( - (tn_cardinal_tagger.graph_hundred_component_at_least_one_none_zero_digit @ optional_delete_space) + ( + tn_cardinal_tagger.graph_hundred_component_at_least_one_none_zero_digit + @ optional_delete_space + ) .invert() .optimize() ) - self.graph_ties = (tn_cardinal_tagger.two_digit_non_zero @ optional_delete_space).invert().optimize() + self.graph_ties = ( + (tn_cardinal_tagger.two_digit_non_zero @ optional_delete_space).invert().optimize() + ) # this is to make sure if there is an ambiguity with decimal, decimal is chosen, e.g. 1000000 vs. 1 million graph = pynutil.add_weight(graph, weight=0.001) self.graph_no_exception = graph - self.digit = pynini.arcmap(tn_cardinal_tagger.digit, map_type="rmweight").invert().optimize() - graph_exception = pynini.project(self.digit, 'input') + self.digit = ( + pynini.arcmap(tn_cardinal_tagger.digit, map_type="rmweight").invert().optimize() + ) + graph_exception = pynini.project(self.digit, "input") self.graph = (pynini.project(graph, "input") - graph_exception.arcsort()) @ graph self.optional_minus_graph = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("minus ", "\"-\" "), 0, 1 + pynutil.insert("negative: ") + pynini.cross("minus ", '"-" '), 0, 1 ) - final_graph = self.optional_minus_graph + pynutil.insert("integer: \"") + self.graph + pynutil.insert("\"") + final_graph = ( + self.optional_minus_graph + + pynutil.insert('integer: "') + + self.graph + + pynutil.insert('"') + ) final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/de/taggers/date.py b/fun_text_processing/inverse_text_normalization/de/taggers/date.py index e36a13a83..31aa614e1 100644 --- a/fun_text_processing/inverse_text_normalization/de/taggers/date.py +++ b/fun_text_processing/inverse_text_normalization/de/taggers/date.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_DIGIT, @@ -35,20 +34,28 @@ class DateFst(GraphFst): ): super().__init__(name="date", kind="classify", deterministic=deterministic) - add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | (pynutil.insert("0") + DAMO_DIGIT) + add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | ( + pynutil.insert("0") + DAMO_DIGIT + ) optional_delete_space = pynini.closure(DAMO_SIGMA | pynutil.delete(" ", weight=0.0001)) tagger = tn_date_verbalizer.graph.invert().optimize() delete_day_marker = ( - pynutil.delete("day: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + pynutil.delete('day: "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') ) @ itn_cardinal_tagger.graph_no_exception - month_as_number = pynutil.delete("month: \"") + itn_cardinal_tagger.graph_no_exception + pynutil.delete("\"") - month_as_string = pynutil.delete("month: \"") + tn_date_tagger.month_abbr.invert() + pynutil.delete("\"") + month_as_number = ( + pynutil.delete('month: "') + + itn_cardinal_tagger.graph_no_exception + + pynutil.delete('"') + ) + month_as_string = ( + pynutil.delete('month: "') + tn_date_tagger.month_abbr.invert() + pynutil.delete('"') + ) convert_year = (tn_date_tagger.year @ optional_delete_space).invert().optimize() delete_year_marker = ( - pynutil.delete("year: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + pynutil.delete('year: "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') ) @ convert_year # day. month as string (year) @@ -73,5 +80,5 @@ class DateFst(GraphFst): final_graph = tagger @ verbalizer - graph = pynutil.insert("name: \"") + convert_space(final_graph) + pynutil.insert("\"") + graph = pynutil.insert('name: "') + convert_space(final_graph) + pynutil.insert('"') self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/de/taggers/decimal.py b/fun_text_processing/inverse_text_normalization/de/taggers/decimal.py index a037a9775..e2dba5a65 100644 --- a/fun_text_processing/inverse_text_normalization/de/taggers/decimal.py +++ b/fun_text_processing/inverse_text_normalization/de/taggers/decimal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.de.taggers.decimal import get_quantity, quantities from fun_text_processing.text_normalization.en.graph_utils import DAMO_SIGMA, GraphFst @@ -15,18 +14,24 @@ class DecimalFst(GraphFst): tn_decimal_tagger: TN decimal tagger """ - def __init__(self, itn_cardinal_tagger: GraphFst, tn_decimal_tagger: GraphFst, deterministic: bool = True): + def __init__( + self, itn_cardinal_tagger: GraphFst, tn_decimal_tagger: GraphFst, deterministic: bool = True + ): super().__init__(name="decimal", kind="classify", deterministic=deterministic) self.graph = tn_decimal_tagger.graph.invert().optimize() delete_point = pynutil.delete(" komma") - allow_spelling = pynini.cdrewrite(pynini.cross("eine ", "eins ") + quantities, "[BOS]", "[EOS]", DAMO_SIGMA) + allow_spelling = pynini.cdrewrite( + pynini.cross("eine ", "eins ") + quantities, "[BOS]", "[EOS]", DAMO_SIGMA + ) - graph_fractional = pynutil.insert("fractional_part: \"") + self.graph + pynutil.insert("\"") + graph_fractional = pynutil.insert('fractional_part: "') + self.graph + pynutil.insert('"') graph_integer = ( - pynutil.insert("integer_part: \"") + itn_cardinal_tagger.graph_no_exception + pynutil.insert("\"") + pynutil.insert('integer_part: "') + + itn_cardinal_tagger.graph_no_exception + + pynutil.insert('"') ) final_graph_wo_sign = graph_integer + delete_point + pynini.accep(" ") + graph_fractional @@ -35,7 +40,8 @@ class DecimalFst(GraphFst): @ ( final_graph_wo_sign | get_quantity( - final_graph_wo_sign, itn_cardinal_tagger.graph_hundred_component_at_least_one_none_zero_digit + final_graph_wo_sign, + itn_cardinal_tagger.graph_hundred_component_at_least_one_none_zero_digit, ) ).optimize() ) diff --git a/fun_text_processing/inverse_text_normalization/de/taggers/electronic.py b/fun_text_processing/inverse_text_normalization/de/taggers/electronic.py index 185a52411..823188a2d 100644 --- a/fun_text_processing/inverse_text_normalization/de/taggers/electronic.py +++ b/fun_text_processing/inverse_text_normalization/de/taggers/electronic.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import GraphFst from pynini.lib import pynutil @@ -8,18 +7,23 @@ class ElectronicFst(GraphFst): """ Finite state transducer for classifying electronic: email addresses, etc. e.g. c d f eins at a b c punkt e d u -> tokens { name: "cdf1.abc.edu" } - + Args: tn_electronic_tagger: TN eletronic tagger tn_electronic_verbalizer: TN eletronic verbalizer """ - def __init__(self, tn_electronic_tagger: GraphFst, tn_electronic_verbalizer: GraphFst, deterministic: bool = True): + def __init__( + self, + tn_electronic_tagger: GraphFst, + tn_electronic_verbalizer: GraphFst, + deterministic: bool = True, + ): super().__init__(name="electronic", kind="classify", deterministic=deterministic) tagger = pynini.invert(tn_electronic_verbalizer.graph).optimize() verbalizer = pynini.invert(tn_electronic_tagger.graph).optimize() final_graph = tagger @ verbalizer - graph = pynutil.insert("name: \"") + final_graph + pynutil.insert("\"") + graph = pynutil.insert('name: "') + final_graph + pynutil.insert('"') self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/de/taggers/fraction.py b/fun_text_processing/inverse_text_normalization/de/taggers/fraction.py index 94fff82ef..dd307ba2d 100644 --- a/fun_text_processing/inverse_text_normalization/de/taggers/fraction.py +++ b/fun_text_processing/inverse_text_normalization/de/taggers/fraction.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_NOT_QUOTE, @@ -15,28 +14,41 @@ class FractionFst(GraphFst): e.g. ein halb -> tokens { name: "1/2" } e.g. ein ein halb -> tokens { name: "1 1/2" } e.g. drei zwei ein hundertstel -> tokens { name: "3 2/100" } - + Args: itn_cardinal_tagger: ITN cardinal tagger tn_fraction_verbalizer: TN fraction verbalizer """ - def __init__(self, itn_cardinal_tagger: GraphFst, tn_fraction_verbalizer: GraphFst, deterministic: bool = True): + def __init__( + self, + itn_cardinal_tagger: GraphFst, + tn_fraction_verbalizer: GraphFst, + deterministic: bool = True, + ): super().__init__(name="fraction", kind="classify", deterministic=deterministic) tagger = tn_fraction_verbalizer.graph.invert().optimize() - delete_optional_sign = pynini.closure(pynutil.delete("negative: ") + pynini.cross("\"true\" ", "-"), 0, 1) + delete_optional_sign = pynini.closure( + pynutil.delete("negative: ") + pynini.cross('"true" ', "-"), 0, 1 + ) delete_integer_marker = ( - pynutil.delete("integer_part: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + pynutil.delete('integer_part: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') ) @ itn_cardinal_tagger.graph_no_exception delete_numerator_marker = ( - pynutil.delete("numerator: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + pynutil.delete('numerator: "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') ) @ itn_cardinal_tagger.graph_no_exception delete_denominator_marker = ( - pynutil.insert('/') - + (pynutil.delete("denominator: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"")) + pynutil.insert("/") + + ( + pynutil.delete('denominator: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + ) @ itn_cardinal_tagger.graph_no_exception ) @@ -50,5 +62,5 @@ class FractionFst(GraphFst): self.graph = tagger @ verbalizer - graph = pynutil.insert("name: \"") + convert_space(self.graph) + pynutil.insert("\"") + graph = pynutil.insert('name: "') + convert_space(self.graph) + pynutil.insert('"') self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/de/taggers/measure.py b/fun_text_processing/inverse_text_normalization/de/taggers/measure.py index 681f1575b..cd3c55798 100644 --- a/fun_text_processing/inverse_text_normalization/de/taggers/measure.py +++ b/fun_text_processing/inverse_text_normalization/de/taggers/measure.py @@ -1,6 +1,8 @@ - import pynini -from fun_text_processing.text_normalization.de.taggers.measure import singular_to_plural, unit_singular +from fun_text_processing.text_normalization.de.taggers.measure import ( + singular_to_plural, + unit_singular, +) from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_SIGMA, GraphFst, @@ -35,25 +37,31 @@ class MeasureFst(GraphFst): super().__init__(name="measure", kind="classify", deterministic=deterministic) cardinal_graph = ( - pynini.cdrewrite(pynini.cross(pynini.union("ein", "eine"), "eins"), "[BOS]", "[EOS]", DAMO_SIGMA) + pynini.cdrewrite( + pynini.cross(pynini.union("ein", "eine"), "eins"), "[BOS]", "[EOS]", DAMO_SIGMA + ) @ itn_cardinal_tagger.graph_no_exception ) graph_unit_singular = pynini.invert(unit_singular) # singular -> abbr - unit = (pynini.invert(singular_to_plural()) @ graph_unit_singular) | graph_unit_singular # plural -> abbr + unit = ( + pynini.invert(singular_to_plural()) @ graph_unit_singular + ) | graph_unit_singular # plural -> abbr unit = convert_space(unit) graph_unit_singular = convert_space(graph_unit_singular) optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("minus", "\"true\"") + delete_extra_space, 0, 1 + pynutil.insert("negative: ") + pynini.cross("minus", '"true"') + delete_extra_space, + 0, + 1, ) unit_misc = pynutil.insert("/") + pynutil.delete("pro") + delete_space + graph_unit_singular unit = ( - pynutil.insert("units: \"") + pynutil.insert('units: "') + (unit | unit_misc | pynutil.add_weight(unit + delete_space + unit_misc, 0.01)) - + pynutil.insert("\"") + + pynutil.insert('"') ) subgraph_decimal = ( @@ -68,9 +76,9 @@ class MeasureFst(GraphFst): subgraph_fraction = ( pynutil.insert("decimal { ") + optional_graph_negative - + pynutil.insert("integer_part: \"") + + pynutil.insert('integer_part: "') + itn_fraction_tagger.graph - + pynutil.insert("\" }") + + pynutil.insert('" }') + delete_extra_space + unit ) @@ -78,9 +86,9 @@ class MeasureFst(GraphFst): subgraph_cardinal = ( pynutil.insert("cardinal { ") + optional_graph_negative - + pynutil.insert("integer: \"") + + pynutil.insert('integer: "') + cardinal_graph - + pynutil.insert("\"") + + pynutil.insert('"') + pynutil.insert(" }") + delete_extra_space + unit diff --git a/fun_text_processing/inverse_text_normalization/de/taggers/money.py b/fun_text_processing/inverse_text_normalization/de/taggers/money.py index 091bde0c2..041adc127 100644 --- a/fun_text_processing/inverse_text_normalization/de/taggers/money.py +++ b/fun_text_processing/inverse_text_normalization/de/taggers/money.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.text_normalization.de.taggers.money import maj_singular, min_plural, min_singular +from fun_text_processing.text_normalization.de.taggers.money import ( + maj_singular, + min_plural, + min_singular, +) from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_DIGIT, DAMO_SIGMA, @@ -23,26 +26,35 @@ class MoneyFst(GraphFst): itn_decimal_tagger: ITN Decimal Tagger """ - def __init__(self, itn_cardinal_tagger: GraphFst, itn_decimal_tagger: GraphFst, deterministic: bool = True): + def __init__( + self, + itn_cardinal_tagger: GraphFst, + itn_decimal_tagger: GraphFst, + deterministic: bool = True, + ): super().__init__(name="money", kind="classify", deterministic=deterministic) cardinal_graph = ( - pynini.cdrewrite(pynini.cross(pynini.union("ein", "eine"), "eins"), "[BOS]", "[EOS]", DAMO_SIGMA) + pynini.cdrewrite( + pynini.cross(pynini.union("ein", "eine"), "eins"), "[BOS]", "[EOS]", DAMO_SIGMA + ) @ itn_cardinal_tagger.graph_no_exception ) graph_decimal_final = itn_decimal_tagger.final_graph_wo_negative graph_unit = pynini.invert(maj_singular) - graph_unit = pynutil.insert("currency: \"") + convert_space(graph_unit) + pynutil.insert("\"") + graph_unit = pynutil.insert('currency: "') + convert_space(graph_unit) + pynutil.insert('"') - add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | (pynutil.insert("0") + DAMO_DIGIT) + add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | ( + pynutil.insert("0") + DAMO_DIGIT + ) min_unit = pynini.project(min_singular | min_plural, "output") # elf euro (und) vier cent, vier cent cents_standalone = ( - pynutil.insert("fractional_part: \"") + pynutil.insert('fractional_part: "') + cardinal_graph @ add_leading_zero_to_double_digit + delete_space + pynutil.delete(min_unit) - + pynutil.insert("\"") + + pynutil.insert('"') ) optional_cents_standalone = pynini.closure( @@ -56,23 +68,23 @@ class MoneyFst(GraphFst): # elf euro vierzig, only after integer optional_cents_suffix = pynini.closure( delete_extra_space - + pynutil.insert("fractional_part: \"") + + pynutil.insert('fractional_part: "') + pynutil.add_weight(cardinal_graph @ add_leading_zero_to_double_digit, -0.7) - + pynutil.insert("\""), + + pynutil.insert('"'), 0, 1, ) graph_integer = ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + cardinal_graph - + pynutil.insert("\"") + + pynutil.insert('"') + delete_extra_space + graph_unit + (optional_cents_standalone | optional_cents_suffix) ) graph_decimal = graph_decimal_final + delete_extra_space + graph_unit - graph_decimal |= pynutil.insert("currency: \"€\" integer_part: \"0\" ") + cents_standalone + graph_decimal |= pynutil.insert('currency: "€" integer_part: "0" ') + cents_standalone final_graph = graph_integer | graph_decimal final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/de/taggers/ordinal.py b/fun_text_processing/inverse_text_normalization/de/taggers/ordinal.py index bc4dd8404..d2c159077 100644 --- a/fun_text_processing/inverse_text_normalization/de/taggers/ordinal.py +++ b/fun_text_processing/inverse_text_normalization/de/taggers/ordinal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst from pynini.lib import pynutil @@ -14,16 +13,21 @@ class OrdinalFst(GraphFst): tn_ordinal_verbalizer: TN Ordinal Verbalizer """ - def __init__(self, itn_cardinal_tagger: GraphFst, tn_ordinal_verbalizer: GraphFst, deterministic: bool = True): + def __init__( + self, + itn_cardinal_tagger: GraphFst, + tn_ordinal_verbalizer: GraphFst, + deterministic: bool = True, + ): super().__init__(name="ordinal", kind="classify", deterministic=deterministic) tagger = tn_ordinal_verbalizer.graph.invert().optimize() graph = ( - pynutil.delete("integer: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + pynutil.delete('integer: "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') ) @ itn_cardinal_tagger.graph final_graph = tagger @ graph + pynutil.insert(".") - graph = pynutil.insert("name: \"") + final_graph + pynutil.insert("\"") + graph = pynutil.insert('name: "') + final_graph + pynutil.insert('"') self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/de/taggers/telephone.py b/fun_text_processing/inverse_text_normalization/de/taggers/telephone.py index a0bfa205a..1dace7beb 100644 --- a/fun_text_processing/inverse_text_normalization/de/taggers/telephone.py +++ b/fun_text_processing/inverse_text_normalization/de/taggers/telephone.py @@ -1,14 +1,17 @@ - import pynini -from fun_text_processing.text_normalization.en.graph_utils import GraphFst, convert_space, insert_space +from fun_text_processing.text_normalization.en.graph_utils import ( + GraphFst, + convert_space, + insert_space, +) from pynini.lib import pynutil class TelephoneFst(GraphFst): """ - Finite state transducer for classifying telephone numbers, e.g. + Finite state transducer for classifying telephone numbers, e.g. null vier eins eins eins zwei drei vier eins zwei drei vier -> tokens { name: "(0411) 1234-1234" } - + Args: tn_cardinal_tagger: TN Cardinal Tagger """ @@ -35,6 +38,6 @@ class TelephoneFst(GraphFst): + digit ) graph = convert_space(pynini.invert(number_part)) - final_graph = pynutil.insert("name: \"") + graph + pynutil.insert("\"") + final_graph = pynutil.insert('name: "') + graph + pynutil.insert('"') self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/de/taggers/time.py b/fun_text_processing/inverse_text_normalization/de/taggers/time.py index 9a30a9a93..f9b05c9f5 100644 --- a/fun_text_processing/inverse_text_normalization/de/taggers/time.py +++ b/fun_text_processing/inverse_text_normalization/de/taggers/time.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_SIGMA, GraphFst from pynini.lib import pynutil @@ -17,7 +15,7 @@ class TimeFst(GraphFst): e.g. drei vor zwölf -> time { minutes: "57" hours: "11" } e.g. drei nach zwölf -> time { minutes: "3" hours: "12" } e.g. drei uhr zehn minuten zehn sekunden -> time { hours: "3" hours: "10" sekunden: "10"} - + Args: tn_time_verbalizer: TN time verbalizer """ diff --git a/fun_text_processing/inverse_text_normalization/de/taggers/tokenize_and_classify.py b/fun_text_processing/inverse_text_normalization/de/taggers/tokenize_and_classify.py index c8c029632..13d64b185 100644 --- a/fun_text_processing/inverse_text_normalization/de/taggers/tokenize_and_classify.py +++ b/fun_text_processing/inverse_text_normalization/de/taggers/tokenize_and_classify.py @@ -1,4 +1,3 @@ - import os import pynini @@ -15,15 +14,27 @@ from fun_text_processing.inverse_text_normalization.de.taggers.time import TimeF from fun_text_processing.inverse_text_normalization.de.taggers.whitelist import WhiteListFst from fun_text_processing.inverse_text_normalization.en.taggers.punctuation import PunctuationFst from fun_text_processing.inverse_text_normalization.en.taggers.word import WordFst -from fun_text_processing.text_normalization.de.taggers.cardinal import CardinalFst as TNCardinalTagger +from fun_text_processing.text_normalization.de.taggers.cardinal import ( + CardinalFst as TNCardinalTagger, +) from fun_text_processing.text_normalization.de.taggers.date import DateFst as TNDateTagger from fun_text_processing.text_normalization.de.taggers.decimal import DecimalFst as TNDecimalTagger -from fun_text_processing.text_normalization.de.taggers.electronic import ElectronicFst as TNElectronicTagger -from fun_text_processing.text_normalization.de.taggers.whitelist import WhiteListFst as TNWhitelistTagger +from fun_text_processing.text_normalization.de.taggers.electronic import ( + ElectronicFst as TNElectronicTagger, +) +from fun_text_processing.text_normalization.de.taggers.whitelist import ( + WhiteListFst as TNWhitelistTagger, +) from fun_text_processing.text_normalization.de.verbalizers.date import DateFst as TNDateVerbalizer -from fun_text_processing.text_normalization.de.verbalizers.electronic import ElectronicFst as TNElectronicVerbalizer -from fun_text_processing.text_normalization.de.verbalizers.fraction import FractionFst as TNFractionVerbalizer -from fun_text_processing.text_normalization.de.verbalizers.ordinal import OrdinalFst as TNOrdinalVerbalizer +from fun_text_processing.text_normalization.de.verbalizers.electronic import ( + ElectronicFst as TNElectronicVerbalizer, +) +from fun_text_processing.text_normalization.de.verbalizers.fraction import ( + FractionFst as TNFractionVerbalizer, +) +from fun_text_processing.text_normalization.de.verbalizers.ordinal import ( + OrdinalFst as TNOrdinalVerbalizer, +) from fun_text_processing.text_normalization.de.verbalizers.time import TimeFst as TNTimeVerbalizer from fun_text_processing.text_normalization.en.graph_utils import ( GraphFst, @@ -39,7 +50,7 @@ import logging class ClassifyFst(GraphFst): """ Final class that composes all other classification grammars. This class can process an entire sentence, that is lower cased. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. Args: @@ -47,11 +58,13 @@ class ClassifyFst(GraphFst): overwrite_cache: set to True to overwrite .far files """ - def __init__(self, cache_dir: str = None, overwrite_cache: bool = False, deterministic: bool = True): + def __init__( + self, cache_dir: str = None, overwrite_cache: bool = False, deterministic: bool = True + ): super().__init__(name="tokenize_and_classify", kind="classify", deterministic=deterministic) far_file = None - if cache_dir is not None and cache_dir != 'None': + if cache_dir is not None and cache_dir != "None": os.makedirs(cache_dir, exist_ok=True) far_file = os.path.join(cache_dir, "_de_itn.far") if not overwrite_cache and far_file and os.path.exists(far_file): @@ -63,9 +76,15 @@ class ClassifyFst(GraphFst): tn_date_tagger = TNDateTagger(cardinal=tn_cardinal_tagger, deterministic=False) tn_decimal_tagger = TNDecimalTagger(cardinal=tn_cardinal_tagger, deterministic=False) tn_ordinal_verbalizer = TNOrdinalVerbalizer(deterministic=False) - tn_fraction_verbalizer = TNFractionVerbalizer(ordinal=tn_ordinal_verbalizer, deterministic=False) - tn_time_verbalizer = TNTimeVerbalizer(cardinal_tagger=tn_cardinal_tagger, deterministic=False) - tn_date_verbalizer = TNDateVerbalizer(ordinal=tn_ordinal_verbalizer, deterministic=False) + tn_fraction_verbalizer = TNFractionVerbalizer( + ordinal=tn_ordinal_verbalizer, deterministic=False + ) + tn_time_verbalizer = TNTimeVerbalizer( + cardinal_tagger=tn_cardinal_tagger, deterministic=False + ) + tn_date_verbalizer = TNDateVerbalizer( + ordinal=tn_ordinal_verbalizer, deterministic=False + ) tn_electronic_tagger = TNElectronicTagger(deterministic=False) tn_electronic_verbalizer = TNElectronicVerbalizer(deterministic=False) tn_whitelist_tagger = TNWhitelistTagger(input_case="cased", deterministic=False) @@ -73,19 +92,27 @@ class ClassifyFst(GraphFst): cardinal = CardinalFst(tn_cardinal_tagger=tn_cardinal_tagger) cardinal_graph = cardinal.fst - ordinal = OrdinalFst(itn_cardinal_tagger=cardinal, tn_ordinal_verbalizer=tn_ordinal_verbalizer) + ordinal = OrdinalFst( + itn_cardinal_tagger=cardinal, tn_ordinal_verbalizer=tn_ordinal_verbalizer + ) ordinal_graph = ordinal.fst decimal = DecimalFst(itn_cardinal_tagger=cardinal, tn_decimal_tagger=tn_decimal_tagger) decimal_graph = decimal.fst - fraction = FractionFst(itn_cardinal_tagger=cardinal, tn_fraction_verbalizer=tn_fraction_verbalizer) + fraction = FractionFst( + itn_cardinal_tagger=cardinal, tn_fraction_verbalizer=tn_fraction_verbalizer + ) fraction_graph = fraction.fst measure_graph = MeasureFst( - itn_cardinal_tagger=cardinal, itn_decimal_tagger=decimal, itn_fraction_tagger=fraction + itn_cardinal_tagger=cardinal, + itn_decimal_tagger=decimal, + itn_fraction_tagger=fraction, ).fst date_graph = DateFst( - itn_cardinal_tagger=cardinal, tn_date_verbalizer=tn_date_verbalizer, tn_date_tagger=tn_date_tagger + itn_cardinal_tagger=cardinal, + tn_date_verbalizer=tn_date_verbalizer, + tn_date_tagger=tn_date_tagger, ).fst word_graph = WordFst().fst time_graph = TimeFst(tn_time_verbalizer=tn_time_verbalizer).fst @@ -93,7 +120,8 @@ class ClassifyFst(GraphFst): whitelist_graph = WhiteListFst(tn_whitelist_tagger=tn_whitelist_tagger).fst punct_graph = PunctuationFst().fst electronic_graph = ElectronicFst( - tn_electronic_tagger=tn_electronic_tagger, tn_electronic_verbalizer=tn_electronic_verbalizer + tn_electronic_tagger=tn_electronic_tagger, + tn_electronic_verbalizer=tn_electronic_verbalizer, ).fst telephone_graph = TelephoneFst(tn_cardinal_tagger=tn_cardinal_tagger).fst @@ -112,10 +140,16 @@ class ClassifyFst(GraphFst): | pynutil.add_weight(word_graph, 100) ) - punct = pynutil.insert("tokens { ") + pynutil.add_weight(punct_graph, weight=1.1) + pynutil.insert(" }") + punct = ( + pynutil.insert("tokens { ") + + pynutil.add_weight(punct_graph, weight=1.1) + + pynutil.insert(" }") + ) token = pynutil.insert("tokens { ") + classify + pynutil.insert(" }") token_plus_punct = ( - pynini.closure(punct + pynutil.insert(" ")) + token + pynini.closure(pynutil.insert(" ") + punct) + pynini.closure(punct + pynutil.insert(" ")) + + token + + pynini.closure(pynutil.insert(" ") + punct) ) graph = token_plus_punct + pynini.closure(delete_extra_space + token_plus_punct) diff --git a/fun_text_processing/inverse_text_normalization/de/taggers/whitelist.py b/fun_text_processing/inverse_text_normalization/de/taggers/whitelist.py index ab853d04a..884c2f0d6 100644 --- a/fun_text_processing/inverse_text_normalization/de/taggers/whitelist.py +++ b/fun_text_processing/inverse_text_normalization/de/taggers/whitelist.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import GraphFst, convert_space from pynini.lib import pynutil @@ -16,5 +15,5 @@ class WhiteListFst(GraphFst): super().__init__(name="whitelist", kind="classify", deterministic=deterministic) whitelist = pynini.invert(tn_whitelist_tagger.graph) - graph = pynutil.insert("name: \"") + convert_space(whitelist) + pynutil.insert("\"") + graph = pynutil.insert('name: "') + convert_space(whitelist) + pynutil.insert('"') self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/de/verbalizers/__init__.py b/fun_text_processing/inverse_text_normalization/de/verbalizers/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/de/verbalizers/__init__.py +++ b/fun_text_processing/inverse_text_normalization/de/verbalizers/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/inverse_text_normalization/de/verbalizers/cardinal.py b/fun_text_processing/inverse_text_normalization/de/verbalizers/cardinal.py index b98fd489c..a9f84261e 100644 --- a/fun_text_processing/inverse_text_normalization/de/verbalizers/cardinal.py +++ b/fun_text_processing/inverse_text_normalization/de/verbalizers/cardinal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst from pynini.lib import pynutil @@ -16,7 +15,9 @@ class CardinalFst(GraphFst): def __init__(self, tn_cardinal_verbalizer: GraphFst, deterministic: bool = True): super().__init__(name="cardinal", kind="verbalize", deterministic=deterministic) self.numbers = tn_cardinal_verbalizer.numbers - optional_sign = pynini.closure(pynutil.delete("negative: \"") + DAMO_NOT_QUOTE + pynutil.delete("\" "), 0, 1) + optional_sign = pynini.closure( + pynutil.delete('negative: "') + DAMO_NOT_QUOTE + pynutil.delete('" '), 0, 1 + ) graph = optional_sign + self.numbers delete_tokens = self.delete_tokens(graph) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/inverse_text_normalization/de/verbalizers/decimal.py b/fun_text_processing/inverse_text_normalization/de/verbalizers/decimal.py index 3f65387a4..1c8042f6b 100644 --- a/fun_text_processing/inverse_text_normalization/de/verbalizers/decimal.py +++ b/fun_text_processing/inverse_text_normalization/de/verbalizers/decimal.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_preserve_order +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_preserve_order, +) from pynini.lib import pynutil @@ -17,13 +20,17 @@ class DecimalFst(GraphFst): super().__init__(name="decimal", kind="verbalize", deterministic=deterministic) delete_space = pynutil.delete(" ") optional_sign = pynini.closure( - pynutil.delete("negative: \"") + DAMO_NOT_QUOTE + pynutil.delete("\"") + delete_space, 0, 1 + pynutil.delete('negative: "') + DAMO_NOT_QUOTE + pynutil.delete('"') + delete_space, + 0, + 1, ) optional_integer = pynini.closure(tn_decimal_verbalizer.integer, 0, 1) optional_fractional = pynini.closure( delete_space + pynutil.insert(",") + tn_decimal_verbalizer.fractional_default, 0, 1 ) - graph = (optional_integer + optional_fractional + tn_decimal_verbalizer.optional_quantity).optimize() + graph = ( + optional_integer + optional_fractional + tn_decimal_verbalizer.optional_quantity + ).optimize() self.numbers = optional_sign + graph graph = self.numbers + delete_preserve_order delete_tokens = self.delete_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/de/verbalizers/measure.py b/fun_text_processing/inverse_text_normalization/de/verbalizers/measure.py index 2ba8f94c6..2b9480b23 100644 --- a/fun_text_processing/inverse_text_normalization/de/verbalizers/measure.py +++ b/fun_text_processing/inverse_text_normalization/de/verbalizers/measure.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_CHAR, GraphFst, delete_space from pynini.lib import pynutil @@ -18,13 +17,13 @@ class MeasureFst(GraphFst): def __init__(self, decimal: GraphFst, cardinal: GraphFst, deterministic: bool = True): super().__init__(name="measure", kind="verbalize", deterministic=deterministic) - optional_sign = pynini.closure(pynini.cross("negative: \"true\"", "-"), 0, 1) + optional_sign = pynini.closure(pynini.cross('negative: "true"', "-"), 0, 1) unit = ( pynutil.delete("units:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space ) graph_decimal = ( diff --git a/fun_text_processing/inverse_text_normalization/de/verbalizers/money.py b/fun_text_processing/inverse_text_normalization/de/verbalizers/money.py index ce4a2b9f2..bca749296 100644 --- a/fun_text_processing/inverse_text_normalization/de/verbalizers/money.py +++ b/fun_text_processing/inverse_text_normalization/de/verbalizers/money.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_CHAR, GraphFst, delete_space from pynini.lib import pynutil @@ -18,9 +17,9 @@ class MoneyFst(GraphFst): unit = ( pynutil.delete("currency:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) graph = unit + delete_space + decimal.numbers delete_tokens = self.delete_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/de/verbalizers/time.py b/fun_text_processing/inverse_text_normalization/de/verbalizers/time.py index 306e89611..429563fa7 100644 --- a/fun_text_processing/inverse_text_normalization/de/verbalizers/time.py +++ b/fun_text_processing/inverse_text_normalization/de/verbalizers/time.py @@ -1,6 +1,10 @@ - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_ALPHA, DAMO_DIGIT, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_ALPHA, + DAMO_DIGIT, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -9,26 +13,35 @@ class TimeFst(GraphFst): Finite state transducer for verbalizing time, e.g. time { hours: "8" minutes: "30" zone: "e s t" } -> 08:30 Uhr est time { hours: "8" } -> 8 Uhr - time { hours: "8" minutes: "30" seconds: "10" } -> 08:30:10 Uhr + time { hours: "8" minutes: "30" seconds: "10" } -> 08:30:10 Uhr """ def __init__(self, deterministic: bool = True): super().__init__(name="time", kind="verbalize", deterministic=deterministic) - add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | (pynutil.insert("0") + DAMO_DIGIT) - hour = pynutil.delete("hours: \"") + pynini.closure(DAMO_DIGIT, 1) + pynutil.delete("\"") - minute = pynutil.delete("minutes: \"") + pynini.closure(DAMO_DIGIT, 1) + pynutil.delete("\"") + add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | ( + pynutil.insert("0") + DAMO_DIGIT + ) + hour = pynutil.delete('hours: "') + pynini.closure(DAMO_DIGIT, 1) + pynutil.delete('"') + minute = pynutil.delete('minutes: "') + pynini.closure(DAMO_DIGIT, 1) + pynutil.delete('"') - second = pynutil.delete("seconds: \"") + pynini.closure(DAMO_DIGIT, 1) + pynutil.delete("\"") + second = pynutil.delete('seconds: "') + pynini.closure(DAMO_DIGIT, 1) + pynutil.delete('"') zone = ( - pynutil.delete("zone: \"") + pynini.closure(DAMO_ALPHA + delete_space) + DAMO_ALPHA + pynutil.delete("\"") + pynutil.delete('zone: "') + + pynini.closure(DAMO_ALPHA + delete_space) + + DAMO_ALPHA + + pynutil.delete('"') ) optional_zone = pynini.closure(pynini.accep(" ") + zone, 0, 1) graph = ( delete_space + pynutil.insert(":") + (minute @ add_leading_zero_to_double_digit) - + pynini.closure(delete_space + pynutil.insert(":") + (second @ add_leading_zero_to_double_digit), 0, 1) + + pynini.closure( + delete_space + pynutil.insert(":") + (second @ add_leading_zero_to_double_digit), + 0, + 1, + ) + pynutil.insert(" Uhr") + optional_zone ) diff --git a/fun_text_processing/inverse_text_normalization/de/verbalizers/verbalize.py b/fun_text_processing/inverse_text_normalization/de/verbalizers/verbalize.py index 9d8a88e72..42937726d 100644 --- a/fun_text_processing/inverse_text_normalization/de/verbalizers/verbalize.py +++ b/fun_text_processing/inverse_text_normalization/de/verbalizers/verbalize.py @@ -1,18 +1,21 @@ - from fun_text_processing.inverse_text_normalization.de.verbalizers.cardinal import CardinalFst from fun_text_processing.inverse_text_normalization.de.verbalizers.decimal import DecimalFst from fun_text_processing.inverse_text_normalization.de.verbalizers.measure import MeasureFst from fun_text_processing.inverse_text_normalization.de.verbalizers.money import MoneyFst from fun_text_processing.inverse_text_normalization.de.verbalizers.time import TimeFst -from fun_text_processing.text_normalization.de.verbalizers.cardinal import CardinalFst as TNCardinalVerbalizer -from fun_text_processing.text_normalization.de.verbalizers.decimal import DecimalFst as TNDecimalVerbalizer +from fun_text_processing.text_normalization.de.verbalizers.cardinal import ( + CardinalFst as TNCardinalVerbalizer, +) +from fun_text_processing.text_normalization.de.verbalizers.decimal import ( + DecimalFst as TNDecimalVerbalizer, +) from fun_text_processing.text_normalization.en.graph_utils import GraphFst class VerbalizeFst(GraphFst): """ Composes other verbalizer grammars. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. """ diff --git a/fun_text_processing/inverse_text_normalization/de/verbalizers/verbalize_final.py b/fun_text_processing/inverse_text_normalization/de/verbalizers/verbalize_final.py index 78f293856..470fa27b8 100644 --- a/fun_text_processing/inverse_text_normalization/de/verbalizers/verbalize_final.py +++ b/fun_text_processing/inverse_text_normalization/de/verbalizers/verbalize_final.py @@ -1,14 +1,17 @@ - import pynini from fun_text_processing.inverse_text_normalization.de.verbalizers.verbalize import VerbalizeFst from fun_text_processing.inverse_text_normalization.en.verbalizers.word import WordFst -from fun_text_processing.text_normalization.en.graph_utils import GraphFst, delete_extra_space, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + GraphFst, + delete_extra_space, + delete_space, +) from pynini.lib import pynutil class VerbalizeFinalFst(GraphFst): """ - Finite state transducer that verbalizes an entire sentence, e.g. + Finite state transducer that verbalizes an entire sentence, e.g. tokens { name: "jetzt" } tokens { name: "ist" } tokens { time { hours: "12" minutes: "30" } } -> jetzt ist 12:30 Uhr """ diff --git a/fun_text_processing/inverse_text_normalization/en/__init__.py b/fun_text_processing/inverse_text_normalization/en/__init__.py index 68b10481a..4f73ff2c3 100644 --- a/fun_text_processing/inverse_text_normalization/en/__init__.py +++ b/fun_text_processing/inverse_text_normalization/en/__init__.py @@ -1,5 +1,7 @@ - - -from fun_text_processing.inverse_text_normalization.en.taggers.tokenize_and_classify import ClassifyFst +from fun_text_processing.inverse_text_normalization.en.taggers.tokenize_and_classify import ( + ClassifyFst, +) from fun_text_processing.inverse_text_normalization.en.verbalizers.verbalize import VerbalizeFst -from fun_text_processing.inverse_text_normalization.en.verbalizers.verbalize_final import VerbalizeFinalFst +from fun_text_processing.inverse_text_normalization.en.verbalizers.verbalize_final import ( + VerbalizeFinalFst, +) diff --git a/fun_text_processing/inverse_text_normalization/en/clean_eval_data.py b/fun_text_processing/inverse_text_normalization/en/clean_eval_data.py index 0ea70c772..e975bf5ec 100644 --- a/fun_text_processing/inverse_text_normalization/en/clean_eval_data.py +++ b/fun_text_processing/inverse_text_normalization/en/clean_eval_data.py @@ -1,5 +1,3 @@ - - from argparse import ArgumentParser from typing import List @@ -55,7 +53,7 @@ class Filter: Args: processes given instance with process function - + Returns: processed instance if instance belongs to expected class type or original instance """ if instance.token_type != self.class_type: @@ -73,7 +71,9 @@ def process_cardinal_1(instance: Instance) -> Instance: normalized = instance.normalized un_normalized = re.sub(r"[^0-9]", "", un_normalized) normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_ordinal_1(instance: Instance) -> bool: @@ -86,7 +86,9 @@ def process_ordinal_1(instance: Instance) -> Instance: normalized = instance.normalized un_normalized = re.sub(r"[,\s]", "", un_normalized) normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_decimal_1(instance: Instance) -> bool: @@ -99,7 +101,9 @@ def process_decimal_1(instance: Instance) -> Instance: un_normalized = re.sub(r",", "", un_normalized) normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_measure_1(instance: Instance) -> bool: @@ -116,7 +120,9 @@ def process_measure_1(instance: Instance) -> Instance: normalized = re.sub(r"[^a-z\s]", "", normalized) normalized = re.sub(r"per ([a-z\s]*)s$", r"per \1", normalized) normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_money_1(instance: Instance) -> bool: @@ -133,7 +139,9 @@ def process_money_1(instance: Instance) -> Instance: un_normalized = re.sub(r"(\d)m\s*$", r"\1 million", un_normalized) un_normalized = re.sub(r"(\d)bn?\s*$", r"\1 billion", un_normalized) normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_time_1(instance: Instance) -> bool: @@ -148,7 +156,9 @@ def process_time_1(instance: Instance) -> Instance: un_normalized = re.sub(r"(\d)\s?p\s?m\s?", r"\1 p.m.", un_normalized) normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_plain_1(instance: Instance) -> bool: @@ -159,7 +169,9 @@ def filter_plain_1(instance: Instance) -> bool: def process_plain_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_punct_1(instance: Instance) -> bool: @@ -170,7 +182,9 @@ def filter_punct_1(instance: Instance) -> bool: def process_punct_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_date_1(instance: Instance) -> bool: @@ -183,7 +197,9 @@ def process_date_1(instance: Instance) -> Instance: un_normalized = re.sub(r",", "", un_normalized) normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_letters_1(instance: Instance) -> bool: @@ -195,7 +211,9 @@ def process_letters_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_verbatim_1(instance: Instance) -> bool: @@ -206,7 +224,9 @@ def filter_verbatim_1(instance: Instance) -> bool: def process_verbatim_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_digit_1(instance: Instance) -> bool: @@ -218,7 +238,9 @@ def process_digit_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_telephone_1(instance: Instance) -> bool: @@ -230,7 +252,9 @@ def process_telephone_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_electronic_1(instance: Instance) -> bool: @@ -242,7 +266,9 @@ def process_electronic_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_fraction_1(instance: Instance) -> bool: @@ -254,7 +280,9 @@ def process_fraction_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_address_1(instance: Instance) -> bool: @@ -266,27 +294,51 @@ def process_address_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) filters = [] -filters.append(Filter(class_type="CARDINAL", process_func=process_cardinal_1, filter_func=filter_cardinal_1)) -filters.append(Filter(class_type="ORDINAL", process_func=process_ordinal_1, filter_func=filter_ordinal_1)) -filters.append(Filter(class_type="DECIMAL", process_func=process_decimal_1, filter_func=filter_decimal_1)) -filters.append(Filter(class_type="MEASURE", process_func=process_measure_1, filter_func=filter_measure_1)) +filters.append( + Filter(class_type="CARDINAL", process_func=process_cardinal_1, filter_func=filter_cardinal_1) +) +filters.append( + Filter(class_type="ORDINAL", process_func=process_ordinal_1, filter_func=filter_ordinal_1) +) +filters.append( + Filter(class_type="DECIMAL", process_func=process_decimal_1, filter_func=filter_decimal_1) +) +filters.append( + Filter(class_type="MEASURE", process_func=process_measure_1, filter_func=filter_measure_1) +) filters.append(Filter(class_type="MONEY", process_func=process_money_1, filter_func=filter_money_1)) filters.append(Filter(class_type="TIME", process_func=process_time_1, filter_func=filter_time_1)) filters.append(Filter(class_type="DATE", process_func=process_date_1, filter_func=filter_date_1)) filters.append(Filter(class_type="PLAIN", process_func=process_plain_1, filter_func=filter_plain_1)) filters.append(Filter(class_type="PUNCT", process_func=process_punct_1, filter_func=filter_punct_1)) -filters.append(Filter(class_type="LETTERS", process_func=process_letters_1, filter_func=filter_letters_1)) -filters.append(Filter(class_type="VERBATIM", process_func=process_verbatim_1, filter_func=filter_verbatim_1)) +filters.append( + Filter(class_type="LETTERS", process_func=process_letters_1, filter_func=filter_letters_1) +) +filters.append( + Filter(class_type="VERBATIM", process_func=process_verbatim_1, filter_func=filter_verbatim_1) +) filters.append(Filter(class_type="DIGIT", process_func=process_digit_1, filter_func=filter_digit_1)) -filters.append(Filter(class_type="TELEPHONE", process_func=process_telephone_1, filter_func=filter_telephone_1)) -filters.append(Filter(class_type="ELECTRONIC", process_func=process_electronic_1, filter_func=filter_electronic_1)) -filters.append(Filter(class_type="FRACTION", process_func=process_fraction_1, filter_func=filter_fraction_1)) -filters.append(Filter(class_type="ADDRESS", process_func=process_address_1, filter_func=filter_address_1)) +filters.append( + Filter(class_type="TELEPHONE", process_func=process_telephone_1, filter_func=filter_telephone_1) +) +filters.append( + Filter( + class_type="ELECTRONIC", process_func=process_electronic_1, filter_func=filter_electronic_1 + ) +) +filters.append( + Filter(class_type="FRACTION", process_func=process_fraction_1, filter_func=filter_fraction_1) +) +filters.append( + Filter(class_type="ADDRESS", process_func=process_address_1, filter_func=filter_address_1) +) filters.append(Filter(class_type=EOS_TYPE, process_func=lambda x: x, filter_func=lambda x: True)) @@ -315,8 +367,10 @@ def filter_loaded_data(data: List[Instance], verbose: bool = False) -> List[Inst def parse_args(): parser = ArgumentParser() - parser.add_argument("--input", help="input file path", type=str, default='./en_with_types/output-00001-of-00100') - parser.add_argument("--verbose", help="print filtered instances", action='store_true') + parser.add_argument( + "--input", help="input file path", type=str, default="./en_with_types/output-00001-of-00100" + ) + parser.add_argument("--verbose", help="print filtered instances", action="store_true") return parser.parse_args() diff --git a/fun_text_processing/inverse_text_normalization/en/taggers/cardinal.py b/fun_text_processing/inverse_text_normalization/en/taggers/cardinal.py index d2e138cc6..5246d0eed 100644 --- a/fun_text_processing/inverse_text_normalization/en/taggers/cardinal.py +++ b/fun_text_processing/inverse_text_normalization/en/taggers/cardinal.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.inverse_text_normalization.en.utils import get_abs_path, num_to_word from fun_text_processing.text_normalization.en.graph_utils import ( @@ -17,7 +15,7 @@ class CardinalFst(GraphFst): """ Finite state transducer for classifying cardinals e.g. minus twenty three -> cardinal { integer: "23" negative: "-" } } - Numbers below thirteen are not converted. + Numbers below thirteen are not converted. """ def __init__(self): @@ -29,7 +27,9 @@ class CardinalFst(GraphFst): graph_hundred = pynini.cross("hundred", "") - graph_hundred_component = pynini.union(graph_digit + delete_space + graph_hundred, pynutil.insert("0")) + graph_hundred_component = pynini.union( + graph_digit + delete_space + graph_hundred, pynutil.insert("0") + ) graph_hundred_component += delete_space graph_hundred_component += pynini.union( graph_teen | pynutil.insert("00"), @@ -44,32 +44,46 @@ class CardinalFst(GraphFst): ) graph_thousands = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("thousand"), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + pynutil.delete("thousand"), pynutil.insert("000", weight=0.1), ) graph_million = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("million"), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + pynutil.delete("million"), pynutil.insert("000", weight=0.1), ) graph_billion = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("billion"), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + pynutil.delete("billion"), pynutil.insert("000", weight=0.1), ) graph_trillion = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("trillion"), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + pynutil.delete("trillion"), pynutil.insert("000", weight=0.1), ) graph_quadrillion = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("quadrillion"), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + pynutil.delete("quadrillion"), pynutil.insert("000", weight=0.1), ) graph_quintillion = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("quintillion"), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + pynutil.delete("quintillion"), pynutil.insert("000", weight=0.1), ) graph_sextillion = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("sextillion"), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + pynutil.delete("sextillion"), pynutil.insert("000", weight=0.1), ) @@ -93,7 +107,10 @@ class CardinalFst(GraphFst): ) graph = graph @ pynini.union( - pynutil.delete(pynini.closure("0")) + pynini.difference(DAMO_DIGIT, "0") + pynini.closure(DAMO_DIGIT), "0" + pynutil.delete(pynini.closure("0")) + + pynini.difference(DAMO_DIGIT, "0") + + pynini.closure(DAMO_DIGIT), + "0", ) labels_exception = [num_to_word(x) for x in range(0, 13)] @@ -110,10 +127,12 @@ class CardinalFst(GraphFst): self.graph = (pynini.project(graph, "input") - graph_exception.arcsort()) @ graph optional_minus_graph = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("minus", "\"-\"") + DAMO_SPACE, 0, 1 + pynutil.insert("negative: ") + pynini.cross("minus", '"-"') + DAMO_SPACE, 0, 1 ) - final_graph = optional_minus_graph + pynutil.insert("integer: \"") + self.graph + pynutil.insert("\"") + final_graph = ( + optional_minus_graph + pynutil.insert('integer: "') + self.graph + pynutil.insert('"') + ) final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/en/taggers/date.py b/fun_text_processing/inverse_text_normalization/en/taggers/date.py index 05178f066..82fe2a7e1 100644 --- a/fun_text_processing/inverse_text_normalization/en/taggers/date.py +++ b/fun_text_processing/inverse_text_normalization/en/taggers/date.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.en.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import ( @@ -63,7 +62,9 @@ def _get_year_graph(): def _get_thousands_graph(): graph_ties = _get_ties_graph() - graph_hundred_component = (graph_digit + delete_space + pynutil.delete("hundred")) | pynutil.insert("0") + graph_hundred_component = ( + graph_digit + delete_space + pynutil.delete("hundred") + ) | pynutil.insert("0") graph = ( graph_digit + delete_space @@ -90,7 +91,7 @@ def _get_year_graph(): class DateFst(GraphFst): """ - Finite state transducer for classifying date, + Finite state transducer for classifying date, e.g. january fifth twenty twelve -> date { month: "january" day: "5" year: "2012" preserve_order: true } e.g. the fifth of january twenty twelve -> date { day: "5" month: "january" year: "2012" preserve_order: true } e.g. twenty twenty -> date { year: "2012" preserve_order: true } @@ -108,18 +109,26 @@ class DateFst(GraphFst): year_graph = pynutil.add_weight(year_graph, YEAR_WEIGHT) month_graph = _get_month_graph() - month_graph = pynutil.insert("month: \"") + month_graph + pynutil.insert("\"") + month_graph = pynutil.insert('month: "') + month_graph + pynutil.insert('"') - day_graph = pynutil.insert("day: \"") + pynutil.add_weight(ordinal_graph, -0.7) + pynutil.insert("\"") + day_graph = ( + pynutil.insert('day: "') + pynutil.add_weight(ordinal_graph, -0.7) + pynutil.insert('"') + ) graph_year = ( delete_extra_space - + pynutil.insert("year: \"") + + pynutil.insert('year: "') + pynutil.add_weight(year_graph, -YEAR_WEIGHT) - + pynutil.insert("\"") + + pynutil.insert('"') + ) + optional_graph_year = pynini.closure( + graph_year, + 0, + 1, ) - optional_graph_year = pynini.closure(graph_year, 0, 1,) graph_mdy = month_graph + ( - (delete_extra_space + day_graph) | graph_year | (delete_extra_space + day_graph + graph_year) + (delete_extra_space + day_graph) + | graph_year + | (delete_extra_space + day_graph + graph_year) ) graph_dmy = ( pynutil.delete("the") @@ -131,7 +140,9 @@ class DateFst(GraphFst): + month_graph + optional_graph_year ) - graph_year = pynutil.insert("year: \"") + (year_graph | _get_range_graph()) + pynutil.insert("\"") + graph_year = ( + pynutil.insert('year: "') + (year_graph | _get_range_graph()) + pynutil.insert('"') + ) final_graph = graph_mdy | graph_dmy | graph_year final_graph += pynutil.insert(" preserve_order: true") diff --git a/fun_text_processing/inverse_text_normalization/en/taggers/decimal.py b/fun_text_processing/inverse_text_normalization/en/taggers/decimal.py index 196f3fc69..987af43d5 100644 --- a/fun_text_processing/inverse_text_normalization/en/taggers/decimal.py +++ b/fun_text_processing/inverse_text_normalization/en/taggers/decimal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.en.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import ( @@ -10,30 +9,42 @@ from fun_text_processing.text_normalization.en.graph_utils import ( from pynini.lib import pynutil -def get_quantity(decimal: 'pynini.FstLike', cardinal_up_to_hundred: 'pynini.FstLike') -> 'pynini.FstLike': +def get_quantity( + decimal: "pynini.FstLike", cardinal_up_to_hundred: "pynini.FstLike" +) -> "pynini.FstLike": """ Returns FST that transforms either a cardinal or decimal followed by a quantity into a numeral, e.g. one million -> integer_part: "1" quantity: "million" e.g. one point five million -> integer_part: "1" fractional_part: "5" quantity: "million" - Args: + Args: decimal: decimal FST cardinal_up_to_hundred: cardinal FST """ numbers = cardinal_up_to_hundred @ ( - pynutil.delete(pynini.closure("0")) + pynini.difference(DAMO_DIGIT, "0") + pynini.closure(DAMO_DIGIT) + pynutil.delete(pynini.closure("0")) + + pynini.difference(DAMO_DIGIT, "0") + + pynini.closure(DAMO_DIGIT) + ) + suffix = pynini.union( + "million", "billion", "trillion", "quadrillion", "quintillion", "sextillion" ) - suffix = pynini.union("million", "billion", "trillion", "quadrillion", "quintillion", "sextillion") res = ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + numbers - + pynutil.insert("\"") + + pynutil.insert('"') + delete_extra_space - + pynutil.insert("quantity: \"") + + pynutil.insert('quantity: "') + suffix - + pynutil.insert("\"") + + pynutil.insert('"') + ) + res |= ( + decimal + + delete_extra_space + + pynutil.insert('quantity: "') + + (suffix | "thousand") + + pynutil.insert('"') ) - res |= decimal + delete_extra_space + pynutil.insert("quantity: \"") + (suffix | "thousand") + pynutil.insert("\"") return res @@ -52,7 +63,9 @@ class DecimalFst(GraphFst): cardinal_graph = cardinal.graph_no_exception graph_decimal = pynini.string_file(get_abs_path("data/numbers/digit.tsv")) - graph_decimal |= pynini.string_file(get_abs_path("data/numbers/zero.tsv")) | pynini.cross("o", "0") + graph_decimal |= pynini.string_file(get_abs_path("data/numbers/zero.tsv")) | pynini.cross( + "o", "0" + ) graph_decimal = pynini.closure(graph_decimal + delete_space) + graph_decimal self.graph = graph_decimal @@ -60,13 +73,20 @@ class DecimalFst(GraphFst): point = pynutil.delete("point") optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("minus", "\"true\"") + delete_extra_space, 0, 1 + pynutil.insert("negative: ") + pynini.cross("minus", '"true"') + delete_extra_space, + 0, + 1, ) - graph_fractional = pynutil.insert("fractional_part: \"") + graph_decimal + pynutil.insert("\"") - graph_integer = pynutil.insert("integer_part: \"") + cardinal_graph + pynutil.insert("\"") + graph_fractional = ( + pynutil.insert('fractional_part: "') + graph_decimal + pynutil.insert('"') + ) + graph_integer = pynutil.insert('integer_part: "') + cardinal_graph + pynutil.insert('"') final_graph_wo_sign = ( - pynini.closure(graph_integer + delete_extra_space, 0, 1) + point + delete_extra_space + graph_fractional + pynini.closure(graph_integer + delete_extra_space, 0, 1) + + point + + delete_extra_space + + graph_fractional ) final_graph = optional_graph_negative + final_graph_wo_sign diff --git a/fun_text_processing/inverse_text_normalization/en/taggers/electronic.py b/fun_text_processing/inverse_text_normalization/en/taggers/electronic.py index d11917628..49292fcd8 100644 --- a/fun_text_processing/inverse_text_normalization/en/taggers/electronic.py +++ b/fun_text_processing/inverse_text_normalization/en/taggers/electronic.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.en.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import DAMO_ALPHA, GraphFst, insert_space @@ -25,35 +24,50 @@ class ElectronicFst(GraphFst): accepted_username = alpha_num | symbols process_dot = pynini.cross("dot", ".") - username = (alpha_num + pynini.closure(delete_extra_space + accepted_username)) | pynutil.add_weight( - pynini.closure(DAMO_ALPHA, 1), weight=0.0001 - ) - username = pynutil.insert("username: \"") + username + pynutil.insert("\"") + username = ( + alpha_num + pynini.closure(delete_extra_space + accepted_username) + ) | pynutil.add_weight(pynini.closure(DAMO_ALPHA, 1), weight=0.0001) + username = pynutil.insert('username: "') + username + pynutil.insert('"') single_alphanum = pynini.closure(alpha_num + delete_extra_space) + alpha_num - server = single_alphanum | pynini.string_file(get_abs_path("data/electronic/server_name.tsv")) + server = single_alphanum | pynini.string_file( + get_abs_path("data/electronic/server_name.tsv") + ) domain = single_alphanum | pynini.string_file(get_abs_path("data/electronic/domain.tsv")) domain_graph = ( - pynutil.insert("domain: \"") + pynutil.insert('domain: "') + server + delete_extra_space + process_dot + delete_extra_space + domain - + pynutil.insert("\"") + + pynutil.insert('"') + ) + graph = ( + username + + delete_extra_space + + pynutil.delete("at") + + insert_space + + delete_extra_space + + domain_graph ) - graph = username + delete_extra_space + pynutil.delete("at") + insert_space + delete_extra_space + domain_graph ############# url ### protocol_end = pynini.cross(pynini.union("w w w", "www"), "www") - protocol_start = (pynini.cross("h t t p", "http") | pynini.cross("h t t p s", "https")) + pynini.cross( - " colon slash slash ", "://" - ) + protocol_start = ( + pynini.cross("h t t p", "http") | pynini.cross("h t t p s", "https") + ) + pynini.cross(" colon slash slash ", "://") # .com, ending = ( delete_extra_space + symbols + delete_extra_space - + (domain | pynini.closure(accepted_username + delete_extra_space,) + accepted_username) + + ( + domain + | pynini.closure( + accepted_username + delete_extra_space, + ) + + accepted_username + ) ) protocol_default = ( @@ -64,12 +78,18 @@ class ElectronicFst(GraphFst): + pynini.closure(ending, 1) ).optimize() protocol = ( - pynini.closure(protocol_start, 0, 1) + protocol_end + delete_extra_space + process_dot + protocol_default + pynini.closure(protocol_start, 0, 1) + + protocol_end + + delete_extra_space + + process_dot + + protocol_default ).optimize() - protocol |= pynini.closure(protocol_end + delete_extra_space + process_dot, 0, 1) + protocol_default + protocol |= ( + pynini.closure(protocol_end + delete_extra_space + process_dot, 0, 1) + protocol_default + ) - protocol = pynutil.insert("protocol: \"") + protocol.optimize() + pynutil.insert("\"") + protocol = pynutil.insert('protocol: "') + protocol.optimize() + pynutil.insert('"') graph |= protocol ######## diff --git a/fun_text_processing/inverse_text_normalization/en/taggers/fraction.py b/fun_text_processing/inverse_text_normalization/en/taggers/fraction.py index 6d4d4b9d6..2829e9e31 100644 --- a/fun_text_processing/inverse_text_normalization/en/taggers/fraction.py +++ b/fun_text_processing/inverse_text_normalization/en/taggers/fraction.py @@ -1,4 +1,3 @@ - from fun_text_processing.text_normalization.en.graph_utils import GraphFst diff --git a/fun_text_processing/inverse_text_normalization/en/taggers/measure.py b/fun_text_processing/inverse_text_normalization/en/taggers/measure.py index 8aa9aca09..3d581e3e4 100644 --- a/fun_text_processing/inverse_text_normalization/en/taggers/measure.py +++ b/fun_text_processing/inverse_text_normalization/en/taggers/measure.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.en.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import ( @@ -32,22 +31,37 @@ class MeasureFst(GraphFst): graph_unit_plural = get_singulars(graph_unit_singular) # plural -> abbr optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("minus", "\"true\"") + delete_extra_space, 0, 1 + pynutil.insert("negative: ") + pynini.cross("minus", '"true"') + delete_extra_space, + 0, + 1, ) unit_singular = convert_space(graph_unit_singular) unit_plural = convert_space(graph_unit_plural) - unit_misc = pynutil.insert("/") + pynutil.delete("per") + delete_space + convert_space(graph_unit_singular) + unit_misc = ( + pynutil.insert("/") + + pynutil.delete("per") + + delete_space + + convert_space(graph_unit_singular) + ) unit_singular = ( - pynutil.insert("units: \"") - + (unit_singular | unit_misc | pynutil.add_weight(unit_singular + delete_space + unit_misc, 0.01)) - + pynutil.insert("\"") + pynutil.insert('units: "') + + ( + unit_singular + | unit_misc + | pynutil.add_weight(unit_singular + delete_space + unit_misc, 0.01) + ) + + pynutil.insert('"') ) unit_plural = ( - pynutil.insert("units: \"") - + (unit_plural | unit_misc | pynutil.add_weight(unit_plural + delete_space + unit_misc, 0.01)) - + pynutil.insert("\"") + pynutil.insert('units: "') + + ( + unit_plural + | unit_misc + | pynutil.add_weight(unit_plural + delete_space + unit_misc, 0.01) + ) + + pynutil.insert('"') ) subgraph_decimal = ( @@ -61,9 +75,9 @@ class MeasureFst(GraphFst): subgraph_cardinal = ( pynutil.insert("cardinal { ") + optional_graph_negative - + pynutil.insert("integer: \"") + + pynutil.insert('integer: "') + ((DAMO_SIGMA - "one") @ cardinal_graph) - + pynutil.insert("\"") + + pynutil.insert('"') + pynutil.insert(" }") + delete_extra_space + unit_plural @@ -71,9 +85,9 @@ class MeasureFst(GraphFst): subgraph_cardinal |= ( pynutil.insert("cardinal { ") + optional_graph_negative - + pynutil.insert("integer: \"") + + pynutil.insert('integer: "') + pynini.cross("one", "1") - + pynutil.insert("\"") + + pynutil.insert('"') + pynutil.insert(" }") + delete_extra_space + unit_singular diff --git a/fun_text_processing/inverse_text_normalization/en/taggers/money.py b/fun_text_processing/inverse_text_normalization/en/taggers/money.py index df6812db8..df219755b 100644 --- a/fun_text_processing/inverse_text_normalization/en/taggers/money.py +++ b/fun_text_processing/inverse_text_normalization/en/taggers/money.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.en.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import ( @@ -33,8 +32,11 @@ class MoneyFst(GraphFst): # add support for missing hundred (only for 3 digit numbers) # "one fifty" -> "one hundred fifty" with_hundred = pynini.compose( - pynini.closure(DAMO_NOT_SPACE) + pynini.accep(" ") + pynutil.insert("hundred ") + DAMO_SIGMA, - pynini.compose(cardinal_graph, DAMO_DIGIT ** 3), + pynini.closure(DAMO_NOT_SPACE) + + pynini.accep(" ") + + pynutil.insert("hundred ") + + DAMO_SIGMA, + pynini.compose(cardinal_graph, DAMO_DIGIT**3), ) cardinal_graph |= with_hundred graph_decimal_final = decimal.final_graph_wo_negative @@ -43,20 +45,27 @@ class MoneyFst(GraphFst): unit_singular = pynini.invert(unit) unit_plural = get_singulars(unit_singular) - graph_unit_singular = pynutil.insert("currency: \"") + convert_space(unit_singular) + pynutil.insert("\"") - graph_unit_plural = pynutil.insert("currency: \"") + convert_space(unit_plural) + pynutil.insert("\"") + graph_unit_singular = ( + pynutil.insert('currency: "') + convert_space(unit_singular) + pynutil.insert('"') + ) + graph_unit_plural = ( + pynutil.insert('currency: "') + convert_space(unit_plural) + pynutil.insert('"') + ) - add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | (pynutil.insert("0") + DAMO_DIGIT) + add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | ( + pynutil.insert("0") + DAMO_DIGIT + ) # twelve dollars (and) fifty cents, zero cents cents_standalone = ( - pynutil.insert("fractional_part: \"") + pynutil.insert('fractional_part: "') + pynini.union( - pynutil.add_weight(((DAMO_SIGMA - "one") @ cardinal_graph), -0.7) @ add_leading_zero_to_double_digit + pynutil.add_weight(((DAMO_SIGMA - "one") @ cardinal_graph), -0.7) + @ add_leading_zero_to_double_digit + delete_space + (pynutil.delete("cents") | pynutil.delete("cent")), pynini.cross("one", "01") + delete_space + pynutil.delete("cent"), ) - + pynutil.insert("\"") + + pynutil.insert('"') ) optional_cents_standalone = pynini.closure( @@ -70,31 +79,31 @@ class MoneyFst(GraphFst): # twelve dollars fifty, only after integer optional_cents_suffix = pynini.closure( delete_extra_space - + pynutil.insert("fractional_part: \"") + + pynutil.insert('fractional_part: "') + pynutil.add_weight(cardinal_graph @ add_leading_zero_to_double_digit, -0.7) - + pynutil.insert("\""), + + pynutil.insert('"'), 0, 1, ) graph_integer = ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + ((DAMO_SIGMA - "one") @ cardinal_graph) - + pynutil.insert("\"") + + pynutil.insert('"') + delete_extra_space + graph_unit_plural + (optional_cents_standalone | optional_cents_suffix) ) graph_integer |= ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + pynini.cross("one", "1") - + pynutil.insert("\"") + + pynutil.insert('"') + delete_extra_space + graph_unit_singular + (optional_cents_standalone | optional_cents_suffix) ) graph_decimal = graph_decimal_final + delete_extra_space + graph_unit_plural - graph_decimal |= pynutil.insert("currency: \"$\" integer_part: \"0\" ") + cents_standalone + graph_decimal |= pynutil.insert('currency: "$" integer_part: "0" ') + cents_standalone final_graph = graph_integer | graph_decimal final_graph = self.add_tokens(final_graph) diff --git a/fun_text_processing/inverse_text_normalization/en/taggers/ordinal.py b/fun_text_processing/inverse_text_normalization/en/taggers/ordinal.py index e57aaae31..2c6f276f8 100644 --- a/fun_text_processing/inverse_text_normalization/en/taggers/ordinal.py +++ b/fun_text_processing/inverse_text_normalization/en/taggers/ordinal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.en.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import DAMO_CHAR, GraphFst @@ -25,6 +24,6 @@ class OrdinalFst(GraphFst): ) self.graph = graph @ cardinal_graph - final_graph = pynutil.insert("integer: \"") + self.graph + pynutil.insert("\"") + final_graph = pynutil.insert('integer: "') + self.graph + pynutil.insert('"') final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/en/taggers/punctuation.py b/fun_text_processing/inverse_text_normalization/en/taggers/punctuation.py index 73eb283b0..dfe16facf 100644 --- a/fun_text_processing/inverse_text_normalization/en/taggers/punctuation.py +++ b/fun_text_processing/inverse_text_normalization/en/taggers/punctuation.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import GraphFst from pynini.lib import pynutil @@ -13,9 +12,9 @@ class PunctuationFst(GraphFst): def __init__(self): super().__init__(name="punctuation", kind="classify") - s = "!#$%&\'()*+,-./:;<=>?@^_`{|}~" + s = "!#$%&'()*+,-./:;<=>?@^_`{|}~" punct = pynini.union(*s) - graph = pynutil.insert("name: \"") + punct + pynutil.insert("\"") + graph = pynutil.insert('name: "') + punct + pynutil.insert('"') self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/en/taggers/telephone.py b/fun_text_processing/inverse_text_normalization/en/taggers/telephone.py index 2865c71ad..280ff88d9 100644 --- a/fun_text_processing/inverse_text_normalization/en/taggers/telephone.py +++ b/fun_text_processing/inverse_text_normalization/en/taggers/telephone.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.en.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import ( @@ -24,7 +23,7 @@ def get_serial_number(cardinal): class TelephoneFst(GraphFst): """ - Finite state transducer for classifying telephone numbers, e.g. + Finite state transducer for classifying telephone numbers, e.g. one two three one two three five six seven eight -> { number_part: "123-123-5678" } This class also support card number and IP format. @@ -61,61 +60,89 @@ class TelephoneFst(GraphFst): double_digit.invert() # to handle cases like "one twenty three" - two_digit_cardinal = pynini.compose(cardinal.graph_no_exception, DAMO_DIGIT ** 2) + two_digit_cardinal = pynini.compose(cardinal.graph_no_exception, DAMO_DIGIT**2) double_digit_to_digit = ( - pynini.compose(double_digit, str_to_digit + pynutil.delete(" ") + str_to_digit) | two_digit_cardinal + pynini.compose(double_digit, str_to_digit + pynutil.delete(" ") + str_to_digit) + | two_digit_cardinal ) - single_or_double_digit = (pynutil.add_weight(double_digit_to_digit, -0.0001) | str_to_digit).optimize() + single_or_double_digit = ( + pynutil.add_weight(double_digit_to_digit, -0.0001) | str_to_digit + ).optimize() single_or_double_digit |= ( single_or_double_digit - + pynini.closure(pynutil.add_weight(pynutil.delete(" ") + single_or_double_digit, 0.0001)) + + pynini.closure( + pynutil.add_weight(pynutil.delete(" ") + single_or_double_digit, 0.0001) + ) ).optimize() number_part = pynini.compose( single_or_double_digit, - DAMO_DIGIT ** 3 + pynutil.insert("-") + DAMO_DIGIT ** 3 + pynutil.insert("-") + DAMO_DIGIT ** 4, + DAMO_DIGIT**3 + + pynutil.insert("-") + + DAMO_DIGIT**3 + + pynutil.insert("-") + + DAMO_DIGIT**4, ).optimize() - number_part = pynutil.insert("number_part: \"") + number_part.optimize() + pynutil.insert("\"") + number_part = ( + pynutil.insert('number_part: "') + number_part.optimize() + pynutil.insert('"') + ) cardinal_option = pynini.compose(single_or_double_digit, DAMO_DIGIT ** (2, 3)) country_code = ( - pynutil.insert("country_code: \"") + pynutil.insert('country_code: "') + pynini.closure(pynini.cross("plus ", "+"), 0, 1) - + ((pynini.closure(str_to_digit + pynutil.delete(" "), 0, 2) + str_to_digit) | cardinal_option) - + pynutil.insert("\"") + + ( + (pynini.closure(str_to_digit + pynutil.delete(" "), 0, 2) + str_to_digit) + | cardinal_option + ) + + pynutil.insert('"') ) - optional_country_code = pynini.closure(country_code + pynutil.delete(" ") + insert_space, 0, 1).optimize() + optional_country_code = pynini.closure( + country_code + pynutil.delete(" ") + insert_space, 0, 1 + ).optimize() graph = optional_country_code + number_part # credit card number - space_four_digits = insert_space + DAMO_DIGIT ** 4 - credit_card_graph = pynini.compose(single_or_double_digit, DAMO_DIGIT ** 4 + space_four_digits ** 3).optimize() - graph |= pynutil.insert("number_part: \"") + credit_card_graph.optimize() + pynutil.insert("\"") + space_four_digits = insert_space + DAMO_DIGIT**4 + credit_card_graph = pynini.compose( + single_or_double_digit, DAMO_DIGIT**4 + space_four_digits**3 + ).optimize() + graph |= ( + pynutil.insert('number_part: "') + credit_card_graph.optimize() + pynutil.insert('"') + ) # SSN ssn_graph = pynini.compose( single_or_double_digit, - DAMO_DIGIT ** 3 + pynutil.insert("-") + DAMO_DIGIT ** 2 + pynutil.insert("-") + DAMO_DIGIT ** 4, + DAMO_DIGIT**3 + + pynutil.insert("-") + + DAMO_DIGIT**2 + + pynutil.insert("-") + + DAMO_DIGIT**4, ).optimize() - graph |= pynutil.insert("number_part: \"") + ssn_graph.optimize() + pynutil.insert("\"") + graph |= pynutil.insert('number_part: "') + ssn_graph.optimize() + pynutil.insert('"') # ip - digit_or_double = pynini.closure(str_to_digit + pynutil.delete(" "), 0, 1) + double_digit_to_digit - digit_or_double |= double_digit_to_digit + pynini.closure(pynutil.delete(" ") + str_to_digit, 0, 1) + digit_or_double = ( + pynini.closure(str_to_digit + pynutil.delete(" "), 0, 1) + double_digit_to_digit + ) + digit_or_double |= double_digit_to_digit + pynini.closure( + pynutil.delete(" ") + str_to_digit, 0, 1 + ) digit_or_double |= str_to_digit + (pynutil.delete(" ") + str_to_digit) ** (0, 2) digit_or_double |= cardinal_option digit_or_double = digit_or_double.optimize() ip_graph = digit_or_double + (pynini.cross(" dot ", ".") + digit_or_double) ** 3 - graph |= pynutil.insert("number_part: \"") + ip_graph.optimize() + pynutil.insert("\"") + graph |= pynutil.insert('number_part: "') + ip_graph.optimize() + pynutil.insert('"') graph |= ( - pynutil.insert("number_part: \"") + pynutil.insert('number_part: "') + pynutil.add_weight(get_serial_number(cardinal=cardinal), weight=0.0001) - + pynutil.insert("\"") + + pynutil.insert('"') ) final_graph = self.add_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/en/taggers/time.py b/fun_text_processing/inverse_text_normalization/en/taggers/time.py index 07a07737c..72cfe541b 100644 --- a/fun_text_processing/inverse_text_normalization/en/taggers/time.py +++ b/fun_text_processing/inverse_text_normalization/en/taggers/time.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.inverse_text_normalization.en.taggers.cardinal import CardinalFst from fun_text_processing.inverse_text_normalization.en.utils import get_abs_path, num_to_word @@ -47,21 +45,23 @@ class TimeFst(GraphFst): graph_minute_verbose = pynini.cross("half", "30") | pynini.cross("quarter", "15") oclock = pynini.cross(pynini.union("o' clock", "o clock", "o'clock", "oclock"), "") - final_graph_hour = pynutil.insert("hours: \"") + graph_hour + pynutil.insert("\"") + final_graph_hour = pynutil.insert('hours: "') + graph_hour + pynutil.insert('"') graph_minute = ( oclock + pynutil.insert("00") | pynutil.delete("o") + delete_space + graph_minute_single | graph_minute_double ) - final_suffix = pynutil.insert("suffix: \"") + convert_space(suffix_graph) + pynutil.insert("\"") + final_suffix = ( + pynutil.insert('suffix: "') + convert_space(suffix_graph) + pynutil.insert('"') + ) final_suffix = delete_space + insert_space + final_suffix final_suffix_optional = pynini.closure(final_suffix, 0, 1) final_time_zone_optional = pynini.closure( delete_space + insert_space - + pynutil.insert("zone: \"") + + pynutil.insert('zone: "') + convert_space(time_zone_graph) - + pynutil.insert("\""), + + pynutil.insert('"'), 0, 1, ) @@ -70,13 +70,17 @@ class TimeFst(GraphFst): # two o eight, two thirty five (am/pm) # two pm/am graph_hm = ( - final_graph_hour + delete_extra_space + pynutil.insert("minutes: \"") + graph_minute + pynutil.insert("\"") + final_graph_hour + + delete_extra_space + + pynutil.insert('minutes: "') + + graph_minute + + pynutil.insert('"') ) # 10 past four, quarter past four, half past four graph_m_past_h = ( - pynutil.insert("minutes: \"") + pynutil.insert('minutes: "') + pynini.union(graph_minute_single, graph_minute_double, graph_minute_verbose) - + pynutil.insert("\"") + + pynutil.insert('"') + delete_space + pynutil.delete("past") + delete_extra_space @@ -84,42 +88,48 @@ class TimeFst(GraphFst): ) graph_quarter_time = ( - pynutil.insert("minutes: \"") + pynutil.insert('minutes: "') + pynini.cross("quarter", "45") - + pynutil.insert("\"") + + pynutil.insert('"') + delete_space + pynutil.delete(pynini.union("to", "till")) + delete_extra_space - + pynutil.insert("hours: \"") + + pynutil.insert('hours: "') + to_hour_graph - + pynutil.insert("\"") + + pynutil.insert('"') ) graph_m_to_h_suffix_time = ( - pynutil.insert("minutes: \"") + pynutil.insert('minutes: "') + ((graph_minute_single | graph_minute_double).optimize() @ minute_to_graph) - + pynutil.insert("\"") - + pynini.closure(delete_space + pynutil.delete(pynini.union("min", "mins", "minute", "minutes")), 0, 1) + + pynutil.insert('"') + + pynini.closure( + delete_space + pynutil.delete(pynini.union("min", "mins", "minute", "minutes")), + 0, + 1, + ) + delete_space + pynutil.delete(pynini.union("to", "till")) + delete_extra_space - + pynutil.insert("hours: \"") + + pynutil.insert('hours: "') + to_hour_graph - + pynutil.insert("\"") + + pynutil.insert('"') + final_suffix ) graph_h = ( final_graph_hour + delete_extra_space - + pynutil.insert("minutes: \"") + + pynutil.insert('minutes: "') + (pynutil.insert("00") | graph_minute) - + pynutil.insert("\"") + + pynutil.insert('"') + final_suffix + final_time_zone_optional ) final_graph = ( - (graph_hm | graph_m_past_h | graph_quarter_time) + final_suffix_optional + final_time_zone_optional + (graph_hm | graph_m_past_h | graph_quarter_time) + + final_suffix_optional + + final_time_zone_optional ) final_graph |= graph_h final_graph |= graph_m_to_h_suffix_time diff --git a/fun_text_processing/inverse_text_normalization/en/taggers/tokenize_and_classify.py b/fun_text_processing/inverse_text_normalization/en/taggers/tokenize_and_classify.py index d2bddec4e..756ea9ed2 100644 --- a/fun_text_processing/inverse_text_normalization/en/taggers/tokenize_and_classify.py +++ b/fun_text_processing/inverse_text_normalization/en/taggers/tokenize_and_classify.py @@ -1,4 +1,3 @@ - import os import pynini @@ -28,7 +27,7 @@ import logging class ClassifyFst(GraphFst): """ Final class that composes all other classification grammars. This class can process an entire sentence, that is lower cased. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. Args: @@ -81,10 +80,16 @@ class ClassifyFst(GraphFst): | pynutil.add_weight(word_graph, 100) ) - punct = pynutil.insert("tokens { ") + pynutil.add_weight(punct_graph, weight=1.1) + pynutil.insert(" }") + punct = ( + pynutil.insert("tokens { ") + + pynutil.add_weight(punct_graph, weight=1.1) + + pynutil.insert(" }") + ) token = pynutil.insert("tokens { ") + classify + pynutil.insert(" }") token_plus_punct = ( - pynini.closure(punct + pynutil.insert(" ")) + token + pynini.closure(pynutil.insert(" ") + punct) + pynini.closure(punct + pynutil.insert(" ")) + + token + + pynini.closure(pynutil.insert(" ") + punct) ) graph = token_plus_punct + pynini.closure(delete_extra_space + token_plus_punct) diff --git a/fun_text_processing/inverse_text_normalization/en/taggers/whitelist.py b/fun_text_processing/inverse_text_normalization/en/taggers/whitelist.py index f33309867..c7443ac26 100644 --- a/fun_text_processing/inverse_text_normalization/en/taggers/whitelist.py +++ b/fun_text_processing/inverse_text_normalization/en/taggers/whitelist.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.en.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import GraphFst, convert_space @@ -16,5 +15,5 @@ class WhiteListFst(GraphFst): super().__init__(name="whitelist", kind="classify") whitelist = pynini.string_file(get_abs_path("data/whitelist.tsv")).invert() - graph = pynutil.insert("name: \"") + convert_space(whitelist) + pynutil.insert("\"") + graph = pynutil.insert('name: "') + convert_space(whitelist) + pynutil.insert('"') self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/en/taggers/word.py b/fun_text_processing/inverse_text_normalization/en/taggers/word.py index 37775e4bb..9758f6885 100644 --- a/fun_text_processing/inverse_text_normalization/en/taggers/word.py +++ b/fun_text_processing/inverse_text_normalization/en/taggers/word.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_SPACE, GraphFst from pynini.lib import pynutil @@ -12,5 +11,5 @@ class WordFst(GraphFst): def __init__(self): super().__init__(name="word", kind="classify") - word = pynutil.insert("name: \"") + pynini.closure(DAMO_NOT_SPACE, 1) + pynutil.insert("\"") + word = pynutil.insert('name: "') + pynini.closure(DAMO_NOT_SPACE, 1) + pynutil.insert('"') self.fst = word.optimize() diff --git a/fun_text_processing/inverse_text_normalization/en/utils.py b/fun_text_processing/inverse_text_normalization/en/utils.py index 5c0c8b570..38a95bac5 100644 --- a/fun_text_processing/inverse_text_normalization/en/utils.py +++ b/fun_text_processing/inverse_text_normalization/en/utils.py @@ -1,5 +1,3 @@ - - import os from typing import Union @@ -15,7 +13,7 @@ def num_to_word(x: Union[str, int]): Args x: integer - Returns: spoken representation + Returns: spoken representation """ if isinstance(x, int): x = str(x) @@ -29,7 +27,7 @@ def get_abs_path(rel_path): Args: rel_path: relative path to this file - + Returns absolute path """ - return os.path.dirname(os.path.abspath(__file__)) + '/' + rel_path + return os.path.dirname(os.path.abspath(__file__)) + "/" + rel_path diff --git a/fun_text_processing/inverse_text_normalization/en/verbalizers/cardinal.py b/fun_text_processing/inverse_text_normalization/en/verbalizers/cardinal.py index 4c3f0ebcc..939c40680 100644 --- a/fun_text_processing/inverse_text_normalization/en/verbalizers/cardinal.py +++ b/fun_text_processing/inverse_text_normalization/en/verbalizers/cardinal.py @@ -1,7 +1,9 @@ - - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -16,9 +18,9 @@ class CardinalFst(GraphFst): optional_sign = pynini.closure( pynutil.delete("negative:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + DAMO_NOT_QUOTE - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space, 0, 1, @@ -26,9 +28,9 @@ class CardinalFst(GraphFst): graph = ( pynutil.delete("integer:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) self.numbers = graph graph = optional_sign + graph diff --git a/fun_text_processing/inverse_text_normalization/en/verbalizers/date.py b/fun_text_processing/inverse_text_normalization/en/verbalizers/date.py index 0e5d4617c..d8cdfb099 100644 --- a/fun_text_processing/inverse_text_normalization/en/verbalizers/date.py +++ b/fun_text_processing/inverse_text_normalization/en/verbalizers/date.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_NOT_QUOTE, @@ -22,43 +20,47 @@ class DateFst(GraphFst): month = ( pynutil.delete("month:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) day = ( pynutil.delete("day:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) year = ( pynutil.delete("year:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') ) # month (day) year graph_mdy = ( - month + pynini.closure(delete_extra_space + day, 0, 1) + pynini.closure(delete_extra_space + year, 0, 1) + month + + pynini.closure(delete_extra_space + day, 0, 1) + + pynini.closure(delete_extra_space + year, 0, 1) ) # (day) month year graph_dmy = ( - pynini.closure(day + delete_extra_space, 0, 1) + month + pynini.closure(delete_extra_space + year, 0, 1) + pynini.closure(day + delete_extra_space, 0, 1) + + month + + pynini.closure(delete_extra_space + year, 0, 1) ) optional_preserve_order = pynini.closure( pynutil.delete("preserve_order:") + delete_space + pynutil.delete("true") + delete_space | pynutil.delete("field_order:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + DAMO_NOT_QUOTE - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space ) diff --git a/fun_text_processing/inverse_text_normalization/en/verbalizers/decimal.py b/fun_text_processing/inverse_text_normalization/en/verbalizers/decimal.py index 9a5041a79..0b7b637dc 100644 --- a/fun_text_processing/inverse_text_normalization/en/verbalizers/decimal.py +++ b/fun_text_processing/inverse_text_normalization/en/verbalizers/decimal.py @@ -1,7 +1,9 @@ - - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -13,30 +15,30 @@ class DecimalFst(GraphFst): def __init__(self): super().__init__(name="decimal", kind="verbalize") - optionl_sign = pynini.closure(pynini.cross("negative: \"true\"", "-") + delete_space, 0, 1) + optionl_sign = pynini.closure(pynini.cross('negative: "true"', "-") + delete_space, 0, 1) integer = ( pynutil.delete("integer_part:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_integer = pynini.closure(integer + delete_space, 0, 1) fractional = ( pynutil.insert(".") + pynutil.delete("fractional_part:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_fractional = pynini.closure(fractional + delete_space, 0, 1) quantity = ( pynutil.delete("quantity:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_quantity = pynini.closure(pynutil.insert(" ") + quantity + delete_space, 0, 1) graph = optional_integer + optional_fractional + optional_quantity diff --git a/fun_text_processing/inverse_text_normalization/en/verbalizers/electronic.py b/fun_text_processing/inverse_text_normalization/en/verbalizers/electronic.py index 95bffb42e..a88607407 100644 --- a/fun_text_processing/inverse_text_normalization/en/verbalizers/electronic.py +++ b/fun_text_processing/inverse_text_normalization/en/verbalizers/electronic.py @@ -1,7 +1,9 @@ - - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -16,24 +18,24 @@ class ElectronicFst(GraphFst): user_name = ( pynutil.delete("username:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) domain = ( pynutil.delete("domain:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) protocol = ( pynutil.delete("protocol:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) graph = user_name + delete_space + pynutil.insert("@") + domain diff --git a/fun_text_processing/inverse_text_normalization/en/verbalizers/fraction.py b/fun_text_processing/inverse_text_normalization/en/verbalizers/fraction.py index fc2f13cc9..b546ad1a6 100644 --- a/fun_text_processing/inverse_text_normalization/en/verbalizers/fraction.py +++ b/fun_text_processing/inverse_text_normalization/en/verbalizers/fraction.py @@ -1,10 +1,9 @@ - from fun_text_processing.text_normalization.en.graph_utils import GraphFst class FractionFst(GraphFst): """ - Finite state transducer for verbalizing fraction, + Finite state transducer for verbalizing fraction, """ def __init__(self): diff --git a/fun_text_processing/inverse_text_normalization/en/verbalizers/measure.py b/fun_text_processing/inverse_text_normalization/en/verbalizers/measure.py index b3fe93214..8f6a13448 100644 --- a/fun_text_processing/inverse_text_normalization/en/verbalizers/measure.py +++ b/fun_text_processing/inverse_text_normalization/en/verbalizers/measure.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_CHAR, GraphFst, delete_space from pynini.lib import pynutil @@ -16,13 +15,13 @@ class MeasureFst(GraphFst): def __init__(self, decimal: GraphFst, cardinal: GraphFst): super().__init__(name="measure", kind="verbalize") - optional_sign = pynini.closure(pynini.cross("negative: \"true\"", "-"), 0, 1) + optional_sign = pynini.closure(pynini.cross('negative: "true"', "-"), 0, 1) unit = ( pynutil.delete("units:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space ) graph_decimal = ( diff --git a/fun_text_processing/inverse_text_normalization/en/verbalizers/money.py b/fun_text_processing/inverse_text_normalization/en/verbalizers/money.py index 89beb93dc..99747f4e3 100644 --- a/fun_text_processing/inverse_text_normalization/en/verbalizers/money.py +++ b/fun_text_processing/inverse_text_normalization/en/verbalizers/money.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_CHAR, GraphFst, delete_space from pynini.lib import pynutil @@ -18,9 +17,9 @@ class MoneyFst(GraphFst): unit = ( pynutil.delete("currency:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) graph = unit + delete_space + decimal.numbers delete_tokens = self.delete_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/en/verbalizers/ordinal.py b/fun_text_processing/inverse_text_normalization/en/verbalizers/ordinal.py index f2c709f12..29f4f530e 100644 --- a/fun_text_processing/inverse_text_normalization/en/verbalizers/ordinal.py +++ b/fun_text_processing/inverse_text_normalization/en/verbalizers/ordinal.py @@ -1,6 +1,10 @@ - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, DAMO_SIGMA, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + DAMO_SIGMA, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -15,9 +19,9 @@ class OrdinalFst(GraphFst): graph = ( pynutil.delete("integer:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) convert_eleven = pynini.cross("11", "11th") convert_twelve = pynini.cross("12", "12th") diff --git a/fun_text_processing/inverse_text_normalization/en/verbalizers/telephone.py b/fun_text_processing/inverse_text_normalization/en/verbalizers/telephone.py index a154aacb1..7f7b0b6ce 100644 --- a/fun_text_processing/inverse_text_normalization/en/verbalizers/telephone.py +++ b/fun_text_processing/inverse_text_normalization/en/verbalizers/telephone.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst from pynini.lib import pynutil @@ -8,17 +7,21 @@ class TelephoneFst(GraphFst): """ Finite state transducer for verbalizing telephone, e.g. telephone { number_part: "123-123-5678" } - -> 123-123-5678 + -> 123-123-5678 """ def __init__(self): super().__init__(name="telephone", kind="verbalize") - number_part = pynutil.delete("number_part: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") - optional_country_code = pynini.closure( - pynutil.delete("country_code: \"") + number_part = ( + pynutil.delete('number_part: "') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') + ) + optional_country_code = pynini.closure( + pynutil.delete('country_code: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + pynini.accep(" "), 0, 1, diff --git a/fun_text_processing/inverse_text_normalization/en/verbalizers/time.py b/fun_text_processing/inverse_text_normalization/en/verbalizers/time.py index 49175f689..c85f259dd 100644 --- a/fun_text_processing/inverse_text_normalization/en/verbalizers/time.py +++ b/fun_text_processing/inverse_text_normalization/en/verbalizers/time.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_CHAR, @@ -20,29 +19,31 @@ class TimeFst(GraphFst): def __init__(self): super().__init__(name="time", kind="verbalize") - add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | (pynutil.insert("0") + DAMO_DIGIT) + add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | ( + pynutil.insert("0") + DAMO_DIGIT + ) hour = ( pynutil.delete("hours:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_DIGIT, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) minute = ( pynutil.delete("minutes:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_DIGIT, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) suffix = ( delete_space + insert_space + pynutil.delete("suffix:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_suffix = pynini.closure(suffix, 0, 1) zone = ( @@ -50,9 +51,9 @@ class TimeFst(GraphFst): + insert_space + pynutil.delete("zone:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_zone = pynini.closure(zone, 0, 1) graph = ( diff --git a/fun_text_processing/inverse_text_normalization/en/verbalizers/verbalize.py b/fun_text_processing/inverse_text_normalization/en/verbalizers/verbalize.py index 2f7eca56a..c9ce9d5c4 100644 --- a/fun_text_processing/inverse_text_normalization/en/verbalizers/verbalize.py +++ b/fun_text_processing/inverse_text_normalization/en/verbalizers/verbalize.py @@ -1,4 +1,3 @@ - from fun_text_processing.inverse_text_normalization.en.verbalizers.cardinal import CardinalFst from fun_text_processing.inverse_text_normalization.en.verbalizers.date import DateFst from fun_text_processing.inverse_text_normalization.en.verbalizers.decimal import DecimalFst @@ -15,7 +14,7 @@ from fun_text_processing.text_normalization.en.graph_utils import GraphFst class VerbalizeFst(GraphFst): """ Composes other verbalizer grammars. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. """ diff --git a/fun_text_processing/inverse_text_normalization/en/verbalizers/verbalize_final.py b/fun_text_processing/inverse_text_normalization/en/verbalizers/verbalize_final.py index 02ed2f99e..f0284281b 100644 --- a/fun_text_processing/inverse_text_normalization/en/verbalizers/verbalize_final.py +++ b/fun_text_processing/inverse_text_normalization/en/verbalizers/verbalize_final.py @@ -1,14 +1,17 @@ - import pynini from fun_text_processing.inverse_text_normalization.en.verbalizers.verbalize import VerbalizeFst from fun_text_processing.inverse_text_normalization.en.verbalizers.word import WordFst -from fun_text_processing.text_normalization.en.graph_utils import GraphFst, delete_extra_space, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + GraphFst, + delete_extra_space, + delete_space, +) from pynini.lib import pynutil class VerbalizeFinalFst(GraphFst): """ - Finite state transducer that verbalizes an entire sentence, e.g. + Finite state transducer that verbalizes an entire sentence, e.g. tokens { name: "its" } tokens { time { hours: "12" minutes: "30" } } tokens { name: "now" } -> its 12:30 now """ diff --git a/fun_text_processing/inverse_text_normalization/en/verbalizers/whitelist.py b/fun_text_processing/inverse_text_normalization/en/verbalizers/whitelist.py index f5a5f129f..cd57a4298 100644 --- a/fun_text_processing/inverse_text_normalization/en/verbalizers/whitelist.py +++ b/fun_text_processing/inverse_text_normalization/en/verbalizers/whitelist.py @@ -1,7 +1,10 @@ - - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_CHAR, DAMO_SIGMA, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_CHAR, + DAMO_SIGMA, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -16,9 +19,9 @@ class WhiteListFst(GraphFst): graph = ( pynutil.delete("name:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) - graph = graph @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + graph = graph @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/en/verbalizers/word.py b/fun_text_processing/inverse_text_normalization/en/verbalizers/word.py index 99861dc31..1d33cd156 100644 --- a/fun_text_processing/inverse_text_normalization/en/verbalizers/word.py +++ b/fun_text_processing/inverse_text_normalization/en/verbalizers/word.py @@ -1,6 +1,10 @@ - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_CHAR, DAMO_SIGMA, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_CHAR, + DAMO_SIGMA, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -13,7 +17,13 @@ class WordFst(GraphFst): def __init__(self): super().__init__(name="word", kind="verbalize") chars = pynini.closure(DAMO_CHAR - " ", 1) - char = pynutil.delete("name:") + delete_space + pynutil.delete("\"") + chars + pynutil.delete("\"") - graph = char @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + char = ( + pynutil.delete("name:") + + delete_space + + pynutil.delete('"') + + chars + + pynutil.delete('"') + ) + graph = char @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/es/__init__.py b/fun_text_processing/inverse_text_normalization/es/__init__.py index b95c71ea4..19bb6ec0f 100644 --- a/fun_text_processing/inverse_text_normalization/es/__init__.py +++ b/fun_text_processing/inverse_text_normalization/es/__init__.py @@ -1,4 +1,7 @@ - -from fun_text_processing.inverse_text_normalization.es.taggers.tokenize_and_classify import ClassifyFst +from fun_text_processing.inverse_text_normalization.es.taggers.tokenize_and_classify import ( + ClassifyFst, +) from fun_text_processing.inverse_text_normalization.es.verbalizers.verbalize import VerbalizeFst -from fun_text_processing.inverse_text_normalization.es.verbalizers.verbalize_final import VerbalizeFinalFst +from fun_text_processing.inverse_text_normalization.es.verbalizers.verbalize_final import ( + VerbalizeFinalFst, +) diff --git a/fun_text_processing/inverse_text_normalization/es/data/__init__.py b/fun_text_processing/inverse_text_normalization/es/data/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/es/data/__init__.py +++ b/fun_text_processing/inverse_text_normalization/es/data/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/inverse_text_normalization/es/data/electronic/__init__.py b/fun_text_processing/inverse_text_normalization/es/data/electronic/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/es/data/electronic/__init__.py +++ b/fun_text_processing/inverse_text_normalization/es/data/electronic/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/inverse_text_normalization/es/data/numbers/__init__.py b/fun_text_processing/inverse_text_normalization/es/data/numbers/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/es/data/numbers/__init__.py +++ b/fun_text_processing/inverse_text_normalization/es/data/numbers/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/inverse_text_normalization/es/data/ordinals/__init__.py b/fun_text_processing/inverse_text_normalization/es/data/ordinals/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/es/data/ordinals/__init__.py +++ b/fun_text_processing/inverse_text_normalization/es/data/ordinals/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/inverse_text_normalization/es/data/time/__init__.py b/fun_text_processing/inverse_text_normalization/es/data/time/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/es/data/time/__init__.py +++ b/fun_text_processing/inverse_text_normalization/es/data/time/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/inverse_text_normalization/es/taggers/__init__.py b/fun_text_processing/inverse_text_normalization/es/taggers/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/es/taggers/__init__.py +++ b/fun_text_processing/inverse_text_normalization/es/taggers/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/inverse_text_normalization/es/taggers/cardinal.py b/fun_text_processing/inverse_text_normalization/es/taggers/cardinal.py index 425ec8fd4..bb9ffd177 100644 --- a/fun_text_processing/inverse_text_normalization/es/taggers/cardinal.py +++ b/fun_text_processing/inverse_text_normalization/es/taggers/cardinal.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.inverse_text_normalization.es.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import ( @@ -16,10 +14,10 @@ from pynini.lib import pynutil class CardinalFst(GraphFst): """ Finite state transducer for classifying cardinals - e.g. menos veintitrés -> cardinal { negative: "-" integer: "23"} + e.g. menos veintitrés -> cardinal { negative: "-" integer: "23"} This class converts cardinals up to (but not including) "un cuatrillón", i.e up to "one septillion" in English (10^{24}). - Cardinals below ten are not converted (in order to avoid + Cardinals below ten are not converted (in order to avoid "vivo en una casa" --> "vivo en 1 casa" and any other odd conversions.) Although technically Spanish grammar requires that "y" only comes after @@ -54,7 +52,9 @@ class CardinalFst(GraphFst): ) graph_thousands = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("mil"), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + pynutil.delete("mil"), pynutil.insert("001") + pynutil.delete("mil"), # because we say 'mil', not 'un mil' pynutil.insert("000", weight=0.1), ) @@ -67,7 +67,9 @@ class CardinalFst(GraphFst): ) graph_mil_millones = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("mil"), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + pynutil.delete("mil"), pynutil.insert("001") + pynutil.delete("mil"), # because we say 'mil', not 'un mil' ) graph_mil_millones += delete_space + ( @@ -89,7 +91,9 @@ class CardinalFst(GraphFst): ) graph_mil_billones = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("mil"), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + pynutil.delete("mil"), pynutil.insert("001") + pynutil.delete("mil"), # because we say 'mil', not 'un mil' ) graph_mil_billones += delete_space + ( @@ -104,7 +108,9 @@ class CardinalFst(GraphFst): ) graph_mil_trillones = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("mil"), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + pynutil.delete("mil"), pynutil.insert("001") + pynutil.delete("mil"), # because we say 'mil', not 'un mil' ) graph_mil_trillones += delete_space + ( @@ -131,7 +137,10 @@ class CardinalFst(GraphFst): ) graph = graph @ pynini.union( - pynutil.delete(pynini.closure("0")) + pynini.difference(DAMO_DIGIT, "0") + pynini.closure(DAMO_DIGIT), "0" + pynutil.delete(pynini.closure("0")) + + pynini.difference(DAMO_DIGIT, "0") + + pynini.closure(DAMO_DIGIT), + "0", ) # ignore "y" inside cardinal numbers @@ -144,32 +153,34 @@ class CardinalFst(GraphFst): self.graph_no_exception = graph # save self.numbers_up_to_thousand for use in DecimalFst - digits_up_to_thousand = DAMO_DIGIT | (DAMO_DIGIT ** 2) | (DAMO_DIGIT ** 3) + digits_up_to_thousand = DAMO_DIGIT | (DAMO_DIGIT**2) | (DAMO_DIGIT**3) numbers_up_to_thousand = pynini.compose(graph, digits_up_to_thousand).optimize() self.numbers_up_to_thousand = numbers_up_to_thousand # save self.numbers_up_to_million for use in DecimalFst digits_up_to_million = ( DAMO_DIGIT - | (DAMO_DIGIT ** 2) - | (DAMO_DIGIT ** 3) - | (DAMO_DIGIT ** 4) - | (DAMO_DIGIT ** 5) - | (DAMO_DIGIT ** 6) + | (DAMO_DIGIT**2) + | (DAMO_DIGIT**3) + | (DAMO_DIGIT**4) + | (DAMO_DIGIT**5) + | (DAMO_DIGIT**6) ) numbers_up_to_million = pynini.compose(graph, digits_up_to_million).optimize() self.numbers_up_to_million = numbers_up_to_million # don't convert cardinals from zero to nine inclusive - graph_exception = pynini.project(pynini.union(graph_digit, graph_zero), 'input') + graph_exception = pynini.project(pynini.union(graph_digit, graph_zero), "input") self.graph = (pynini.project(graph, "input") - graph_exception.arcsort()) @ graph optional_minus_graph = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("menos", "\"-\"") + DAMO_SPACE, 0, 1 + pynutil.insert("negative: ") + pynini.cross("menos", '"-"') + DAMO_SPACE, 0, 1 ) - final_graph = optional_minus_graph + pynutil.insert("integer: \"") + self.graph + pynutil.insert("\"") + final_graph = ( + optional_minus_graph + pynutil.insert('integer: "') + self.graph + pynutil.insert('"') + ) final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/es/taggers/date.py b/fun_text_processing/inverse_text_normalization/es/taggers/date.py index 940fe411e..7204f5545 100644 --- a/fun_text_processing/inverse_text_normalization/es/taggers/date.py +++ b/fun_text_processing/inverse_text_normalization/es/taggers/date.py @@ -1,13 +1,16 @@ - import pynini from fun_text_processing.inverse_text_normalization.es.utils import get_abs_path -from fun_text_processing.text_normalization.en.graph_utils import GraphFst, delete_extra_space, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + GraphFst, + delete_extra_space, + delete_space, +) from pynini.lib import pynutil class DateFst(GraphFst): """ - Finite state transducer for classifying date, + Finite state transducer for classifying date, e.g. primero de enero -> date { day: "1" month: "enero" } e.g. uno de enero -> date { day: "1" month: "enero" } """ @@ -33,12 +36,14 @@ class DateFst(GraphFst): # can use "primero" for 1st day of the month graph_1_to_31 = pynini.union(graph_1_to_31, pynini.cross("primero", "1")) - day_graph = pynutil.insert("day: \"") + graph_1_to_31 + pynutil.insert("\"") + day_graph = pynutil.insert('day: "') + graph_1_to_31 + pynutil.insert('"') month_graph = pynini.string_file(get_abs_path("data/months.tsv")) - month_graph = pynutil.insert("month: \"") + month_graph + pynutil.insert("\"") + month_graph = pynutil.insert('month: "') + month_graph + pynutil.insert('"') - graph_dm = day_graph + delete_space + pynutil.delete("de") + delete_extra_space + month_graph + graph_dm = ( + day_graph + delete_space + pynutil.delete("de") + delete_extra_space + month_graph + ) final_graph = graph_dm final_graph += pynutil.insert(" preserve_order: true") diff --git a/fun_text_processing/inverse_text_normalization/es/taggers/decimal.py b/fun_text_processing/inverse_text_normalization/es/taggers/decimal.py index a2c4cc026..2a3dd5043 100644 --- a/fun_text_processing/inverse_text_normalization/es/taggers/decimal.py +++ b/fun_text_processing/inverse_text_normalization/es/taggers/decimal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.es.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import ( @@ -10,18 +9,22 @@ from fun_text_processing.text_normalization.en.graph_utils import ( from pynini.lib import pynutil -def get_quantity(decimal: 'pynini.FstLike', cardinal_up_to_million: 'pynini.FstLike') -> 'pynini.FstLike': +def get_quantity( + decimal: "pynini.FstLike", cardinal_up_to_million: "pynini.FstLike" +) -> "pynini.FstLike": """ Returns FST that transforms either a cardinal or decimal followed by a quantity into a numeral, e.g. one million -> integer_part: "1" quantity: "million" e.g. one point five million -> integer_part: "1" fractional_part: "5" quantity: "million" - Args: + Args: decimal: decimal FST cardinal_up_to_million: cardinal FST """ numbers = cardinal_up_to_million @ ( - pynutil.delete(pynini.closure("0")) + pynini.difference(DAMO_DIGIT, "0") + pynini.closure(DAMO_DIGIT) + pynutil.delete(pynini.closure("0")) + + pynini.difference(DAMO_DIGIT, "0") + + pynini.closure(DAMO_DIGIT) ) suffix = pynini.union( @@ -37,15 +40,17 @@ def get_quantity(decimal: 'pynini.FstLike', cardinal_up_to_million: 'pynini.FstL "cuatrillones", ) res = ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + numbers - + pynutil.insert("\"") + + pynutil.insert('"') + delete_extra_space - + pynutil.insert("quantity: \"") + + pynutil.insert('quantity: "') + suffix - + pynutil.insert("\"") + + pynutil.insert('"') + ) + res |= ( + decimal + delete_extra_space + pynutil.insert('quantity: "') + suffix + pynutil.insert('"') ) - res |= decimal + delete_extra_space + pynutil.insert("quantity: \"") + suffix + pynutil.insert("\"") return res @@ -59,7 +64,7 @@ class DecimalFst(GraphFst): This decimal rule assumes that decimals can be pronounced as: (a cardinal) + ('coma' or 'punto') plus (any sequence of cardinals <1000, including 'zero') - Also writes large numbers in shortened form, e.g. + Also writes large numbers in shortened form, e.g. e.g. uno coma dos seis millón -> decimal { negative: "false" integer_part: "1" morphosyntactic_features: "," fractional_part: "26" quantity: "millón" } e.g. dos millones -> decimal { negative: "false" integer_part: "2" quantity: "millones" } e.g. mil ochocientos veinticuatro millones -> decimal { negative: "false" integer_part: "1824" quantity: "millones" } @@ -77,17 +82,23 @@ class DecimalFst(GraphFst): self.graph = graph_decimal # decimal point can be denoted by 'coma' or 'punto' - decimal_point = pynini.cross("coma", "morphosyntactic_features: \",\"") - decimal_point |= pynini.cross("punto", "morphosyntactic_features: \".\"") + decimal_point = pynini.cross("coma", 'morphosyntactic_features: ","') + decimal_point |= pynini.cross("punto", 'morphosyntactic_features: "."') optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("menos", "\"true\"") + delete_extra_space, 0, 1 + pynutil.insert("negative: ") + pynini.cross("menos", '"true"') + delete_extra_space, + 0, + 1, ) - graph_fractional = pynutil.insert("fractional_part: \"") + graph_decimal + pynutil.insert("\"") + graph_fractional = ( + pynutil.insert('fractional_part: "') + graph_decimal + pynutil.insert('"') + ) - cardinal_graph = cardinal.graph_no_exception | pynini.string_file(get_abs_path("data/numbers/zero.tsv")) - graph_integer = pynutil.insert("integer_part: \"") + cardinal_graph + pynutil.insert("\"") + cardinal_graph = cardinal.graph_no_exception | pynini.string_file( + get_abs_path("data/numbers/zero.tsv") + ) + graph_integer = pynutil.insert('integer_part: "') + cardinal_graph + pynutil.insert('"') final_graph_wo_sign = ( pynini.closure(graph_integer + delete_extra_space, 0, 1) + decimal_point @@ -99,6 +110,8 @@ class DecimalFst(GraphFst): self.final_graph_wo_negative = final_graph_wo_sign | get_quantity( final_graph_wo_sign, cardinal.numbers_up_to_million ) - final_graph |= optional_graph_negative + get_quantity(final_graph_wo_sign, cardinal.numbers_up_to_million) + final_graph |= optional_graph_negative + get_quantity( + final_graph_wo_sign, cardinal.numbers_up_to_million + ) final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/es/taggers/electronic.py b/fun_text_processing/inverse_text_normalization/es/taggers/electronic.py index d4148f9c1..76c5523ac 100644 --- a/fun_text_processing/inverse_text_normalization/es/taggers/electronic.py +++ b/fun_text_processing/inverse_text_normalization/es/taggers/electronic.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.es.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import DAMO_ALPHA, GraphFst, insert_space @@ -29,33 +28,48 @@ class ElectronicFst(GraphFst): accepted_username = alpha_num | symbols process_dot = pynini.cross("punto", ".") username = ( - pynutil.insert("username: \"") + pynutil.insert('username: "') + alpha_num + delete_extra_space + pynini.closure(accepted_username + delete_extra_space) + alpha_num - + pynutil.insert("\"") + + pynutil.insert('"') ) single_alphanum = pynini.closure(alpha_num + delete_extra_space) + alpha_num - server = single_alphanum | pynini.string_file(get_abs_path("data/electronic/server_name.tsv")).invert() - domain = single_alphanum | pynini.string_file(get_abs_path("data/electronic/domain.tsv")).invert() + server = ( + single_alphanum + | pynini.string_file(get_abs_path("data/electronic/server_name.tsv")).invert() + ) + domain = ( + single_alphanum + | pynini.string_file(get_abs_path("data/electronic/domain.tsv")).invert() + ) domain_graph = ( - pynutil.insert("domain: \"") + pynutil.insert('domain: "') + server + delete_extra_space + process_dot + delete_extra_space + domain - + pynutil.insert("\"") + + pynutil.insert('"') ) graph = ( - username + delete_extra_space + pynutil.delete("arroba") + insert_space + delete_extra_space + domain_graph + username + + delete_extra_space + + pynutil.delete("arroba") + + insert_space + + delete_extra_space + + domain_graph ) ############# url ### - protocol_end = pynini.cross(pynini.union("www", "w w w", "doble ve doble ve doble ve"), "www") + protocol_end = pynini.cross( + pynini.union("www", "w w w", "doble ve doble ve doble ve"), "www" + ) protocol_start = pynini.cross(pynini.union("http", "h t t p", "hache te te pe"), "http") - protocol_start |= pynini.cross(pynini.union("https", "h t t p s", "hache te te pe ese"), "https") + protocol_start |= pynini.cross( + pynini.union("https", "h t t p s", "hache te te pe ese"), "https" + ) protocol_start += pynini.cross(" dos puntos barra barra ", "://") # e.g. .com, .es @@ -63,7 +77,13 @@ class ElectronicFst(GraphFst): delete_extra_space + symbols + delete_extra_space - + (domain | pynini.closure(accepted_username + delete_extra_space,) + accepted_username) + + ( + domain + | pynini.closure( + accepted_username + delete_extra_space, + ) + + accepted_username + ) ) protocol = ( @@ -75,7 +95,7 @@ class ElectronicFst(GraphFst): + (pynini.closure(delete_extra_space + accepted_username, 1) | server) + pynini.closure(ending, 1) ) - protocol = pynutil.insert("protocol: \"") + protocol + pynutil.insert("\"") + protocol = pynutil.insert('protocol: "') + protocol + pynutil.insert('"') graph |= protocol ######## diff --git a/fun_text_processing/inverse_text_normalization/es/taggers/measure.py b/fun_text_processing/inverse_text_normalization/es/taggers/measure.py index cf9af0c2c..0092a0ad9 100644 --- a/fun_text_processing/inverse_text_normalization/es/taggers/measure.py +++ b/fun_text_processing/inverse_text_normalization/es/taggers/measure.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.es.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import ( @@ -14,7 +13,7 @@ from pynini.lib import pynutil class MeasureFst(GraphFst): """ Finite state transducer for classifying measure - e.g. menos doce kilogramos -> measure { cardinal { negative: "true" integer: "12" } units: "kg" } + e.g. menos doce kilogramos -> measure { cardinal { negative: "true" integer: "12" } units: "kg" } Args: cardinal: CardinalFst @@ -32,22 +31,37 @@ class MeasureFst(GraphFst): graph_unit_plural = pynini.invert(graph_unit_plural) # plural -> abbr optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("menos", "\"true\"") + delete_extra_space, 0, 1 + pynutil.insert("negative: ") + pynini.cross("menos", '"true"') + delete_extra_space, + 0, + 1, ) unit_singular = convert_space(graph_unit_singular) unit_plural = convert_space(graph_unit_plural) - unit_misc = pynutil.insert("/") + pynutil.delete("por") + delete_space + convert_space(graph_unit_singular) + unit_misc = ( + pynutil.insert("/") + + pynutil.delete("por") + + delete_space + + convert_space(graph_unit_singular) + ) unit_singular = ( - pynutil.insert("units: \"") - + (unit_singular | unit_misc | pynutil.add_weight(unit_singular + delete_space + unit_misc, 0.01)) - + pynutil.insert("\"") + pynutil.insert('units: "') + + ( + unit_singular + | unit_misc + | pynutil.add_weight(unit_singular + delete_space + unit_misc, 0.01) + ) + + pynutil.insert('"') ) unit_plural = ( - pynutil.insert("units: \"") - + (unit_plural | unit_misc | pynutil.add_weight(unit_plural + delete_space + unit_misc, 0.01)) - + pynutil.insert("\"") + pynutil.insert('units: "') + + ( + unit_plural + | unit_misc + | pynutil.add_weight(unit_plural + delete_space + unit_misc, 0.01) + ) + + pynutil.insert('"') ) subgraph_decimal = ( @@ -61,9 +75,9 @@ class MeasureFst(GraphFst): subgraph_cardinal = ( pynutil.insert("cardinal { ") + optional_graph_negative - + pynutil.insert("integer: \"") + + pynutil.insert('integer: "') + ((DAMO_SIGMA - "un" - "una" - "uno") @ cardinal_graph) - + pynutil.insert("\"") + + pynutil.insert('"') + pynutil.insert(" }") + delete_extra_space + unit_plural @@ -71,9 +85,9 @@ class MeasureFst(GraphFst): subgraph_cardinal |= ( pynutil.insert("cardinal { ") + optional_graph_negative - + pynutil.insert("integer: \"") + + pynutil.insert('integer: "') + (pynini.cross("un", "1") | pynini.cross("una", "1") | pynini.cross("uno", "1")) - + pynutil.insert("\"") + + pynutil.insert('"') + pynutil.insert(" }") + delete_extra_space + unit_singular diff --git a/fun_text_processing/inverse_text_normalization/es/taggers/money.py b/fun_text_processing/inverse_text_normalization/es/taggers/money.py index 3b32959a1..742484b2e 100644 --- a/fun_text_processing/inverse_text_normalization/es/taggers/money.py +++ b/fun_text_processing/inverse_text_normalization/es/taggers/money.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.es.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import ( @@ -35,27 +34,36 @@ class MoneyFst(GraphFst): unit_plural = pynini.string_file(get_abs_path("data/currency_plural.tsv")) unit_plural = pynini.invert(unit_plural) - graph_unit_singular = pynutil.insert("currency: \"") + convert_space(unit_singular) + pynutil.insert("\"") - graph_unit_plural = pynutil.insert("currency: \"") + convert_space(unit_plural) + pynutil.insert("\"") + graph_unit_singular = ( + pynutil.insert('currency: "') + convert_space(unit_singular) + pynutil.insert('"') + ) + graph_unit_plural = ( + pynutil.insert('currency: "') + convert_space(unit_plural) + pynutil.insert('"') + ) - add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | (pynutil.insert("0") + DAMO_DIGIT) + add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | ( + pynutil.insert("0") + DAMO_DIGIT + ) # twelve dollars (and) fifty cents, zero cents cents_standalone = ( - pynutil.insert("morphosyntactic_features: \",\"") # always use a comma in the decimal + pynutil.insert('morphosyntactic_features: ","') # always use a comma in the decimal + insert_space - + pynutil.insert("fractional_part: \"") + + pynutil.insert('fractional_part: "') + pynini.union( - pynutil.add_weight(((DAMO_SIGMA - "un") @ cardinal_graph), -0.7) @ add_leading_zero_to_double_digit + pynutil.add_weight(((DAMO_SIGMA - "un") @ cardinal_graph), -0.7) + @ add_leading_zero_to_double_digit + delete_space + pynutil.delete(pynini.union("centavos", "céntimos")), - pynini.cross("un", "01") + delete_space + pynutil.delete(pynini.union("centavo", "céntimo")), + pynini.cross("un", "01") + + delete_space + + pynutil.delete(pynini.union("centavo", "céntimo")), ) - + pynutil.insert("\"") + + pynutil.insert('"') ) optional_cents_standalone = pynini.closure( delete_space - + pynini.closure((pynutil.delete("con") | pynutil.delete('y')) + delete_space, 0, 1) + + pynini.closure((pynutil.delete("con") | pynutil.delete("y")) + delete_space, 0, 1) + insert_space + cents_standalone, 0, @@ -65,34 +73,34 @@ class MoneyFst(GraphFst): # setenta y cinco dólares con sesenta y tres~$75,63 optional_cents_suffix = pynini.closure( delete_extra_space - + pynutil.insert("morphosyntactic_features: \",\"") # always use a comma in the decimal + + pynutil.insert('morphosyntactic_features: ","') # always use a comma in the decimal + insert_space - + pynutil.insert("fractional_part: \"") - + pynini.closure((pynutil.delete("con") | pynutil.delete('y')) + delete_space, 0, 1) + + pynutil.insert('fractional_part: "') + + pynini.closure((pynutil.delete("con") | pynutil.delete("y")) + delete_space, 0, 1) + pynutil.add_weight(cardinal_graph @ add_leading_zero_to_double_digit, -0.7) - + pynutil.insert("\""), + + pynutil.insert('"'), 0, 1, ) graph_integer = ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + ((DAMO_SIGMA - "un" - "una") @ cardinal_graph) - + pynutil.insert("\"") + + pynutil.insert('"') + delete_extra_space + graph_unit_plural + (optional_cents_standalone | optional_cents_suffix) ) graph_integer |= ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + (pynini.cross("un", "1") | pynini.cross("una", "1")) - + pynutil.insert("\"") + + pynutil.insert('"') + delete_extra_space + graph_unit_singular + (optional_cents_standalone | optional_cents_suffix) ) graph_decimal = graph_decimal_final + delete_extra_space + graph_unit_plural - graph_decimal |= pynutil.insert("currency: \"$\" integer_part: \"0\" ") + cents_standalone + graph_decimal |= pynutil.insert('currency: "$" integer_part: "0" ') + cents_standalone final_graph = graph_integer | graph_decimal final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/es/taggers/ordinal.py b/fun_text_processing/inverse_text_normalization/es/taggers/ordinal.py index eb1a0edef..014e7ee08 100644 --- a/fun_text_processing/inverse_text_normalization/es/taggers/ordinal.py +++ b/fun_text_processing/inverse_text_normalization/es/taggers/ordinal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.es.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import DAMO_SIGMA, GraphFst, delete_space @@ -11,7 +10,7 @@ class OrdinalFst(GraphFst): vigésimo primero -> ordinal { integer: "21" morphosyntactic_features: "o" } This class converts ordinal up to "millesímo" (one thousandth) exclusive. - Cardinals below ten are not converted (in order to avoid + Cardinals below ten are not converted (in order to avoid e.g. "primero hice ..." -> "1.º hice...", "segunda guerra mundial" -> "2.ª guerra mundial" and any other odd conversions.) @@ -32,7 +31,13 @@ class OrdinalFst(GraphFst): graph_ties = pynini.string_file(get_abs_path("data/ordinals/ties.tsv")) graph_hundreds = pynini.string_file(get_abs_path("data/ordinals/hundreds.tsv")) - ordinal_graph_union = pynini.union(graph_digit, graph_teens, graph_twenties, graph_ties, graph_hundreds,) + ordinal_graph_union = pynini.union( + graph_digit, + graph_teens, + graph_twenties, + graph_ties, + graph_hundreds, + ) accept_o_endings = DAMO_SIGMA + pynini.accep("o") accept_a_endings = DAMO_SIGMA + pynini.accep("a") @@ -44,34 +49,42 @@ class OrdinalFst(GraphFst): # 'optional_numbers_in_front' have negative weight so we always # include them if they're there - optional_numbers_in_front = (pynutil.add_weight(ordinal_graph_union, -0.1) + delete_space.closure()).closure() + optional_numbers_in_front = ( + pynutil.add_weight(ordinal_graph_union, -0.1) + delete_space.closure() + ).closure() graph_o_suffix = (optional_numbers_in_front + ordinal_graph_o) @ cardinal_graph graph_a_suffix = (optional_numbers_in_front + ordinal_graph_a) @ cardinal_graph graph_er_suffix = (optional_numbers_in_front + ordinal_graph_er) @ cardinal_graph # don't convert ordinals from one to nine inclusive - graph_exception = pynini.project(pynini.union(graph_digit), 'input') - graph_o_suffix = (pynini.project(graph_o_suffix, "input") - graph_exception.arcsort()) @ graph_o_suffix - graph_a_suffix = (pynini.project(graph_a_suffix, "input") - graph_exception.arcsort()) @ graph_a_suffix - graph_er_suffix = (pynini.project(graph_er_suffix, "input") - graph_exception.arcsort()) @ graph_er_suffix + graph_exception = pynini.project(pynini.union(graph_digit), "input") + graph_o_suffix = ( + pynini.project(graph_o_suffix, "input") - graph_exception.arcsort() + ) @ graph_o_suffix + graph_a_suffix = ( + pynini.project(graph_a_suffix, "input") - graph_exception.arcsort() + ) @ graph_a_suffix + graph_er_suffix = ( + pynini.project(graph_er_suffix, "input") - graph_exception.arcsort() + ) @ graph_er_suffix graph = ( - pynutil.insert("integer: \"") + pynutil.insert('integer: "') + graph_o_suffix - + pynutil.insert("\"") - + pynutil.insert(" morphosyntactic_features: \"o\"") + + pynutil.insert('"') + + pynutil.insert(' morphosyntactic_features: "o"') ) graph |= ( - pynutil.insert("integer: \"") + pynutil.insert('integer: "') + graph_a_suffix - + pynutil.insert("\"") - + pynutil.insert(" morphosyntactic_features: \"a\"") + + pynutil.insert('"') + + pynutil.insert(' morphosyntactic_features: "a"') ) graph |= ( - pynutil.insert("integer: \"") + pynutil.insert('integer: "') + graph_er_suffix - + pynutil.insert("\"") - + pynutil.insert(" morphosyntactic_features: \"er\"") + + pynutil.insert('"') + + pynutil.insert(' morphosyntactic_features: "er"') ) final_graph = self.add_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/es/taggers/punctuation.py b/fun_text_processing/inverse_text_normalization/es/taggers/punctuation.py index 73eb283b0..dfe16facf 100644 --- a/fun_text_processing/inverse_text_normalization/es/taggers/punctuation.py +++ b/fun_text_processing/inverse_text_normalization/es/taggers/punctuation.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import GraphFst from pynini.lib import pynutil @@ -13,9 +12,9 @@ class PunctuationFst(GraphFst): def __init__(self): super().__init__(name="punctuation", kind="classify") - s = "!#$%&\'()*+,-./:;<=>?@^_`{|}~" + s = "!#$%&'()*+,-./:;<=>?@^_`{|}~" punct = pynini.union(*s) - graph = pynutil.insert("name: \"") + punct + pynutil.insert("\"") + graph = pynutil.insert('name: "') + punct + pynutil.insert('"') self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/es/taggers/telephone.py b/fun_text_processing/inverse_text_normalization/es/taggers/telephone.py index fa6833654..d8c90a134 100644 --- a/fun_text_processing/inverse_text_normalization/es/taggers/telephone.py +++ b/fun_text_processing/inverse_text_normalization/es/taggers/telephone.py @@ -1,13 +1,16 @@ - import pynini from fun_text_processing.inverse_text_normalization.es.utils import get_abs_path -from fun_text_processing.text_normalization.en.graph_utils import GraphFst, delete_space, insert_space +from fun_text_processing.text_normalization.en.graph_utils import ( + GraphFst, + delete_space, + insert_space, +) from pynini.lib import pynutil class TelephoneFst(GraphFst): """ - Finite state transducer for classifying telephone numbers, e.g. + Finite state transducer for classifying telephone numbers, e.g. uno dos tres uno dos tres cinco seis siete ocho -> { number_part: "123-123-5678" }. If 10 digits are spoken, they are grouped as 3+3+4 (eg. 123-456-7890). If 9 digits are spoken, they are grouped as 3+3+3 (eg. 123-456-789). @@ -116,9 +119,15 @@ class TelephoneFst(GraphFst): + double_digits ) - number_part = pynini.union(ten_digit_graph, nine_digit_graph, eight_digit_graph,) + number_part = pynini.union( + ten_digit_graph, + nine_digit_graph, + eight_digit_graph, + ) - number_part = pynutil.insert("number_part: \"") + pynini.invert(number_part) + pynutil.insert("\"") + number_part = ( + pynutil.insert('number_part: "') + pynini.invert(number_part) + pynutil.insert('"') + ) graph = number_part final_graph = self.add_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/es/taggers/time.py b/fun_text_processing/inverse_text_normalization/es/taggers/time.py index 99f8d3d9c..9fa5cf220 100644 --- a/fun_text_processing/inverse_text_normalization/es/taggers/time.py +++ b/fun_text_processing/inverse_text_normalization/es/taggers/time.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.inverse_text_normalization.es.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import ( @@ -28,14 +26,14 @@ class TimeFst(GraphFst): e.g. cuarto para las dos -> time { minutes: "45" hours: "la 1" } Note that times on the hour (e.g. "las dos" i.e. "two o'clock") do not get - converted into a time format. This is to avoid converting phrases that are + converted into a time format. This is to avoid converting phrases that are not part of a time phrase (e.g. "las dos personas" i.e. "the two people") e.g. las dos -> tokens { name: "las" } tokens { name: "dos" } - However, if a time on the hour is followed by a suffix (indicating 'a.m.' + However, if a time on the hour is followed by a suffix (indicating 'a.m.' or 'p.m.'), it will be converted. e.g. las dos pe eme -> time { hours: "las 2" minutes: "00" suffix: "p.m." } - Note that although the TimeFst verbalizer can accept 'zone' (timezone) fields, + Note that although the TimeFst verbalizer can accept 'zone' (timezone) fields, so far the rules have not been added to the TimeFst tagger to process timezones (to keep the rules simple, and because timezones are not very often specified in Spanish.) @@ -71,21 +69,29 @@ class TimeFst(GraphFst): graph_minute = graph_1_to_100 @ pynini.union(*digits_1_to_59) graph_minute_verbose = pynini.cross("media", "30") | pynini.cross("cuarto", "15") - final_graph_hour = pynutil.insert("hours: \"") + (graph_1oclock | graph_hour) + pynutil.insert("\"") - - final_graph_minute = ( - pynutil.insert("minutes: \"") - + pynini.closure((pynutil.delete("y") | pynutil.delete("con")) + delete_space, 0, 1) - + (graph_minute | graph_minute_verbose) - + pynutil.insert("\"") + final_graph_hour = ( + pynutil.insert('hours: "') + (graph_1oclock | graph_hour) + pynutil.insert('"') ) - final_suffix = pynutil.insert("suffix: \"") + convert_space(suffix_graph) + pynutil.insert("\"") + final_graph_minute = ( + pynutil.insert('minutes: "') + + pynini.closure((pynutil.delete("y") | pynutil.delete("con")) + delete_space, 0, 1) + + (graph_minute | graph_minute_verbose) + + pynutil.insert('"') + ) + + final_suffix = ( + pynutil.insert('suffix: "') + convert_space(suffix_graph) + pynutil.insert('"') + ) final_suffix_optional = pynini.closure(delete_space + insert_space + final_suffix, 0, 1) # las nueve a eme (only convert on-the-hour times if they are followed by a suffix) graph_hsuffix = ( - final_graph_hour + delete_extra_space + pynutil.insert("minutes: \"00\"") + insert_space + final_suffix + final_graph_hour + + delete_extra_space + + pynutil.insert('minutes: "00"') + + insert_space + + final_suffix ) # las nueve y veinticinco @@ -93,22 +99,25 @@ class TimeFst(GraphFst): # un cuarto para las cinco graph_mh = ( - pynutil.insert("minutes: \"") - + pynini.union(pynini.cross("un cuarto para", "45"), pynini.cross("cuarto para", "45"),) - + pynutil.insert("\"") + pynutil.insert('minutes: "') + + pynini.union( + pynini.cross("un cuarto para", "45"), + pynini.cross("cuarto para", "45"), + ) + + pynutil.insert('"') + delete_extra_space - + pynutil.insert("hours: \"") + + pynutil.insert('hours: "') + time_to_graph - + pynutil.insert("\"") + + pynutil.insert('"') ) # las diez menos diez graph_time_to = ( - pynutil.insert("hours: \"") + pynutil.insert('hours: "') + time_to_graph - + pynutil.insert("\"") + + pynutil.insert('"') + delete_extra_space - + pynutil.insert("minutes: \"") + + pynutil.insert('minutes: "') + delete_space + pynutil.delete("menos") + delete_space @@ -119,7 +128,7 @@ class TimeFst(GraphFst): pynini.cross("veinte", "40"), pynini.cross("veinticinco", "30"), ) - + pynutil.insert("\"") + + pynutil.insert('"') ) final_graph = pynini.union( (graph_hm | graph_mh | graph_time_to) + final_suffix_optional, graph_hsuffix diff --git a/fun_text_processing/inverse_text_normalization/es/taggers/tokenize_and_classify.py b/fun_text_processing/inverse_text_normalization/es/taggers/tokenize_and_classify.py index 068c6f57d..3851de2f4 100644 --- a/fun_text_processing/inverse_text_normalization/es/taggers/tokenize_and_classify.py +++ b/fun_text_processing/inverse_text_normalization/es/taggers/tokenize_and_classify.py @@ -1,4 +1,3 @@ - import os import pynini @@ -28,7 +27,7 @@ import logging class ClassifyFst(GraphFst): """ Final class that composes all other classification grammars. This class can process an entire sentence, that is lower cased. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. Args: @@ -82,10 +81,16 @@ class ClassifyFst(GraphFst): | pynutil.add_weight(word_graph, 100) ) - punct = pynutil.insert("tokens { ") + pynutil.add_weight(punct_graph, weight=1.1) + pynutil.insert(" }") + punct = ( + pynutil.insert("tokens { ") + + pynutil.add_weight(punct_graph, weight=1.1) + + pynutil.insert(" }") + ) token = pynutil.insert("tokens { ") + classify + pynutil.insert(" }") token_plus_punct = ( - pynini.closure(punct + pynutil.insert(" ")) + token + pynini.closure(pynutil.insert(" ") + punct) + pynini.closure(punct + pynutil.insert(" ")) + + token + + pynini.closure(pynutil.insert(" ") + punct) ) graph = token_plus_punct + pynini.closure(delete_extra_space + token_plus_punct) diff --git a/fun_text_processing/inverse_text_normalization/es/taggers/whitelist.py b/fun_text_processing/inverse_text_normalization/es/taggers/whitelist.py index 947ecd67b..09c328601 100644 --- a/fun_text_processing/inverse_text_normalization/es/taggers/whitelist.py +++ b/fun_text_processing/inverse_text_normalization/es/taggers/whitelist.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.es.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import GraphFst, convert_space @@ -16,5 +15,5 @@ class WhiteListFst(GraphFst): super().__init__(name="whitelist", kind="classify") whitelist = pynini.string_file(get_abs_path("data/whitelist.tsv")).invert() - graph = pynutil.insert("name: \"") + convert_space(whitelist) + pynutil.insert("\"") + graph = pynutil.insert('name: "') + convert_space(whitelist) + pynutil.insert('"') self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/es/taggers/word.py b/fun_text_processing/inverse_text_normalization/es/taggers/word.py index 37775e4bb..9758f6885 100644 --- a/fun_text_processing/inverse_text_normalization/es/taggers/word.py +++ b/fun_text_processing/inverse_text_normalization/es/taggers/word.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_SPACE, GraphFst from pynini.lib import pynutil @@ -12,5 +11,5 @@ class WordFst(GraphFst): def __init__(self): super().__init__(name="word", kind="classify") - word = pynutil.insert("name: \"") + pynini.closure(DAMO_NOT_SPACE, 1) + pynutil.insert("\"") + word = pynutil.insert('name: "') + pynini.closure(DAMO_NOT_SPACE, 1) + pynutil.insert('"') self.fst = word.optimize() diff --git a/fun_text_processing/inverse_text_normalization/es/utils.py b/fun_text_processing/inverse_text_normalization/es/utils.py index 9cbd175cb..c6fdb7435 100644 --- a/fun_text_processing/inverse_text_normalization/es/utils.py +++ b/fun_text_processing/inverse_text_normalization/es/utils.py @@ -1,4 +1,3 @@ - import os @@ -8,7 +7,7 @@ def get_abs_path(rel_path): Args: rel_path: relative path to this file - + Returns absolute path """ - return os.path.dirname(os.path.abspath(__file__)) + '/' + rel_path + return os.path.dirname(os.path.abspath(__file__)) + "/" + rel_path diff --git a/fun_text_processing/inverse_text_normalization/es/verbalizers/__init__.py b/fun_text_processing/inverse_text_normalization/es/verbalizers/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/es/verbalizers/__init__.py +++ b/fun_text_processing/inverse_text_normalization/es/verbalizers/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/inverse_text_normalization/es/verbalizers/cardinal.py b/fun_text_processing/inverse_text_normalization/es/verbalizers/cardinal.py index 143651acc..005fae0bc 100644 --- a/fun_text_processing/inverse_text_normalization/es/verbalizers/cardinal.py +++ b/fun_text_processing/inverse_text_normalization/es/verbalizers/cardinal.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -15,9 +18,9 @@ class CardinalFst(GraphFst): optional_sign = pynini.closure( pynutil.delete("negative:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + DAMO_NOT_QUOTE - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space, 0, 1, @@ -25,9 +28,9 @@ class CardinalFst(GraphFst): graph = ( pynutil.delete("integer:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) self.numbers = graph graph = optional_sign + graph diff --git a/fun_text_processing/inverse_text_normalization/es/verbalizers/date.py b/fun_text_processing/inverse_text_normalization/es/verbalizers/date.py index 02148d8d3..b708adeb1 100644 --- a/fun_text_processing/inverse_text_normalization/es/verbalizers/date.py +++ b/fun_text_processing/inverse_text_normalization/es/verbalizers/date.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_NOT_QUOTE, @@ -21,16 +20,16 @@ class DateFst(GraphFst): month = ( pynutil.delete("month:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) day = ( pynutil.delete("day:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) # day month @@ -40,9 +39,9 @@ class DateFst(GraphFst): pynutil.delete("preserve_order:") + delete_space + pynutil.delete("true") + delete_space | pynutil.delete("field_order:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + DAMO_NOT_QUOTE - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space ) diff --git a/fun_text_processing/inverse_text_normalization/es/verbalizers/decimal.py b/fun_text_processing/inverse_text_normalization/es/verbalizers/decimal.py index 721275268..346e15b4a 100644 --- a/fun_text_processing/inverse_text_normalization/es/verbalizers/decimal.py +++ b/fun_text_processing/inverse_text_normalization/es/verbalizers/decimal.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -15,35 +18,35 @@ class DecimalFst(GraphFst): def __init__(self): super().__init__(name="decimal", kind="verbalize") - optionl_sign = pynini.closure(pynini.cross("negative: \"true\"", "-") + delete_space, 0, 1) + optionl_sign = pynini.closure(pynini.cross('negative: "true"', "-") + delete_space, 0, 1) integer = ( pynutil.delete("integer_part:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_integer = pynini.closure(integer + delete_space, 0, 1) - decimal_point = pynini.cross("morphosyntactic_features: \",\"", ",") - decimal_point |= pynini.cross("morphosyntactic_features: \".\"", ".") + decimal_point = pynini.cross('morphosyntactic_features: ","', ",") + decimal_point |= pynini.cross('morphosyntactic_features: "."', ".") fractional = ( decimal_point + delete_space + pynutil.delete("fractional_part:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_fractional = pynini.closure(fractional + delete_space, 0, 1) quantity = ( pynutil.delete("quantity:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_quantity = pynini.closure(pynutil.insert(" ") + quantity + delete_space, 0, 1) graph = optional_integer + optional_fractional + optional_quantity diff --git a/fun_text_processing/inverse_text_normalization/es/verbalizers/electronic.py b/fun_text_processing/inverse_text_normalization/es/verbalizers/electronic.py index 53fd9d946..bcdf689a4 100644 --- a/fun_text_processing/inverse_text_normalization/es/verbalizers/electronic.py +++ b/fun_text_processing/inverse_text_normalization/es/verbalizers/electronic.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -16,24 +19,24 @@ class ElectronicFst(GraphFst): user_name = ( pynutil.delete("username:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) domain = ( pynutil.delete("domain:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) protocol = ( pynutil.delete("protocol:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) graph = user_name + delete_space + pynutil.insert("@") + domain diff --git a/fun_text_processing/inverse_text_normalization/es/verbalizers/measure.py b/fun_text_processing/inverse_text_normalization/es/verbalizers/measure.py index f8512b09a..b2005c305 100644 --- a/fun_text_processing/inverse_text_normalization/es/verbalizers/measure.py +++ b/fun_text_processing/inverse_text_normalization/es/verbalizers/measure.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_CHAR, GraphFst, delete_space from pynini.lib import pynutil @@ -16,13 +15,13 @@ class MeasureFst(GraphFst): def __init__(self, decimal: GraphFst, cardinal: GraphFst): super().__init__(name="measure", kind="verbalize") - optional_sign = pynini.closure(pynini.cross("negative: \"true\"", "-"), 0, 1) + optional_sign = pynini.closure(pynini.cross('negative: "true"', "-"), 0, 1) unit = ( pynutil.delete("units:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space ) graph_decimal = ( diff --git a/fun_text_processing/inverse_text_normalization/es/verbalizers/money.py b/fun_text_processing/inverse_text_normalization/es/verbalizers/money.py index e35a2fb99..566a3854d 100644 --- a/fun_text_processing/inverse_text_normalization/es/verbalizers/money.py +++ b/fun_text_processing/inverse_text_normalization/es/verbalizers/money.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_CHAR, GraphFst, delete_space from pynini.lib import pynutil @@ -18,9 +17,9 @@ class MoneyFst(GraphFst): unit = ( pynutil.delete("currency:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) graph = unit + delete_space + decimal.numbers delete_tokens = self.delete_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/es/verbalizers/ordinal.py b/fun_text_processing/inverse_text_normalization/es/verbalizers/ordinal.py index 3e0b6b1f6..017cddf74 100644 --- a/fun_text_processing/inverse_text_normalization/es/verbalizers/ordinal.py +++ b/fun_text_processing/inverse_text_normalization/es/verbalizers/ordinal.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -15,15 +18,15 @@ class OrdinalFst(GraphFst): graph = ( pynutil.delete("integer:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) replace_suffix = pynini.union( - pynini.cross(" morphosyntactic_features: \"o\"", ".º"), - pynini.cross(" morphosyntactic_features: \"a\"", ".ª"), - pynini.cross(" morphosyntactic_features: \"er\"", ".ᵉʳ"), + pynini.cross(' morphosyntactic_features: "o"', ".º"), + pynini.cross(' morphosyntactic_features: "a"', ".ª"), + pynini.cross(' morphosyntactic_features: "er"', ".ᵉʳ"), ) graph = graph + replace_suffix diff --git a/fun_text_processing/inverse_text_normalization/es/verbalizers/telephone.py b/fun_text_processing/inverse_text_normalization/es/verbalizers/telephone.py index 09bc2e6d5..8c30c0843 100644 --- a/fun_text_processing/inverse_text_normalization/es/verbalizers/telephone.py +++ b/fun_text_processing/inverse_text_normalization/es/verbalizers/telephone.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst from pynini.lib import pynutil @@ -8,12 +7,16 @@ class TelephoneFst(GraphFst): """ Finite state transducer for verbalizing telephone, e.g. telephone { number_part: "123-123-5678" } - -> 123-123-5678 + -> 123-123-5678 """ def __init__(self): super().__init__(name="telephone", kind="verbalize") - number_part = pynutil.delete("number_part: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + number_part = ( + pynutil.delete('number_part: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + ) delete_tokens = self.delete_tokens(number_part) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/inverse_text_normalization/es/verbalizers/time.py b/fun_text_processing/inverse_text_normalization/es/verbalizers/time.py index 4b8b6fcde..321467848 100644 --- a/fun_text_processing/inverse_text_normalization/es/verbalizers/time.py +++ b/fun_text_processing/inverse_text_normalization/es/verbalizers/time.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_CHAR, @@ -19,32 +18,34 @@ class TimeFst(GraphFst): def __init__(self): super().__init__(name="time", kind="verbalize") - add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | (pynutil.insert("0") + DAMO_DIGIT) + add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | ( + pynutil.insert("0") + DAMO_DIGIT + ) # hour includes preposition ("la" or "las") hour = ( pynutil.delete("hours:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.union("la ", "las ") + pynini.closure(DAMO_DIGIT, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) minute = ( pynutil.delete("minutes:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_DIGIT, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) suffix = ( delete_space + insert_space + pynutil.delete("suffix:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_suffix = pynini.closure(suffix, 0, 1) zone = ( @@ -52,9 +53,9 @@ class TimeFst(GraphFst): + insert_space + pynutil.delete("zone:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_zone = pynini.closure(zone, 0, 1) graph = ( diff --git a/fun_text_processing/inverse_text_normalization/es/verbalizers/verbalize.py b/fun_text_processing/inverse_text_normalization/es/verbalizers/verbalize.py index b1f048731..8a25015b0 100644 --- a/fun_text_processing/inverse_text_normalization/es/verbalizers/verbalize.py +++ b/fun_text_processing/inverse_text_normalization/es/verbalizers/verbalize.py @@ -1,4 +1,3 @@ - from fun_text_processing.inverse_text_normalization.es.verbalizers.cardinal import CardinalFst from fun_text_processing.inverse_text_normalization.es.verbalizers.date import DateFst from fun_text_processing.inverse_text_normalization.es.verbalizers.decimal import DecimalFst @@ -15,7 +14,7 @@ from fun_text_processing.text_normalization.en.graph_utils import GraphFst class VerbalizeFst(GraphFst): """ Composes other verbalizer grammars. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. """ diff --git a/fun_text_processing/inverse_text_normalization/es/verbalizers/verbalize_final.py b/fun_text_processing/inverse_text_normalization/es/verbalizers/verbalize_final.py index 65182b967..19e6e4946 100644 --- a/fun_text_processing/inverse_text_normalization/es/verbalizers/verbalize_final.py +++ b/fun_text_processing/inverse_text_normalization/es/verbalizers/verbalize_final.py @@ -1,14 +1,17 @@ - import pynini from fun_text_processing.inverse_text_normalization.es.verbalizers.verbalize import VerbalizeFst from fun_text_processing.inverse_text_normalization.es.verbalizers.word import WordFst -from fun_text_processing.text_normalization.en.graph_utils import GraphFst, delete_extra_space, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + GraphFst, + delete_extra_space, + delete_space, +) from pynini.lib import pynutil class VerbalizeFinalFst(GraphFst): """ - Finite state transducer that verbalizes an entire sentence, e.g. + Finite state transducer that verbalizes an entire sentence, e.g. tokens { name: "its" } tokens { time { hours: "12" minutes: "30" } } tokens { name: "now" } -> its 12:30 now """ diff --git a/fun_text_processing/inverse_text_normalization/es/verbalizers/whitelist.py b/fun_text_processing/inverse_text_normalization/es/verbalizers/whitelist.py index 56724908a..da019c5cf 100644 --- a/fun_text_processing/inverse_text_normalization/es/verbalizers/whitelist.py +++ b/fun_text_processing/inverse_text_normalization/es/verbalizers/whitelist.py @@ -1,7 +1,10 @@ - - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_CHAR, DAMO_SIGMA, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_CHAR, + DAMO_SIGMA, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -16,9 +19,9 @@ class WhiteListFst(GraphFst): graph = ( pynutil.delete("name:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) - graph = graph @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + graph = graph @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/es/verbalizers/word.py b/fun_text_processing/inverse_text_normalization/es/verbalizers/word.py index 99861dc31..1d33cd156 100644 --- a/fun_text_processing/inverse_text_normalization/es/verbalizers/word.py +++ b/fun_text_processing/inverse_text_normalization/es/verbalizers/word.py @@ -1,6 +1,10 @@ - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_CHAR, DAMO_SIGMA, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_CHAR, + DAMO_SIGMA, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -13,7 +17,13 @@ class WordFst(GraphFst): def __init__(self): super().__init__(name="word", kind="verbalize") chars = pynini.closure(DAMO_CHAR - " ", 1) - char = pynutil.delete("name:") + delete_space + pynutil.delete("\"") + chars + pynutil.delete("\"") - graph = char @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + char = ( + pynutil.delete("name:") + + delete_space + + pynutil.delete('"') + + chars + + pynutil.delete('"') + ) + graph = char @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/export_models.py b/fun_text_processing/inverse_text_normalization/export_models.py index 9460ad9c9..0597f7f79 100644 --- a/fun_text_processing/inverse_text_normalization/export_models.py +++ b/fun_text_processing/inverse_text_normalization/export_models.py @@ -3,11 +3,16 @@ from time import perf_counter from argparse import ArgumentParser from fun_text_processing.text_normalization.en.graph_utils import generator_main + def parse_args(): parser = ArgumentParser() - + parser.add_argument( - "--language", help="language", choices=['de', 'en', 'es', 'fr', 'id', 'ja', 'ko', 'pt', 'ru', 'vi', 'zh'], default="en", type=str + "--language", + help="language", + choices=["de", "en", "es", "fr", "id", "ja", "ko", "pt", "ru", "vi", "zh"], + default="en", + type=str, ) parser.add_argument( "--export_dir", @@ -17,46 +22,95 @@ def parse_args(): ) return parser.parse_args() -def get_grammars(lang: str='en'): - if lang=='de': - from fun_text_processing.inverse_text_normalization.de.taggers.tokenize_and_classify import ClassifyFst - from fun_text_processing.inverse_text_normalization.de.verbalizers.verbalize_final import VerbalizeFinalFst - elif lang=='en': - from fun_text_processing.inverse_text_normalization.en.taggers.tokenize_and_classify import ClassifyFst - from fun_text_processing.inverse_text_normalization.en.verbalizers.verbalize_final import VerbalizeFinalFst - elif lang=='es': - from fun_text_processing.inverse_text_normalization.es.taggers.tokenize_and_classify import ClassifyFst - from fun_text_processing.inverse_text_normalization.es.verbalizers.verbalize_final import VerbalizeFinalFst - elif lang=='fr': - from fun_text_processing.inverse_text_normalization.fr.taggers.tokenize_and_classify import ClassifyFst - from fun_text_processing.inverse_text_normalization.fr.verbalizers.verbalize_final import VerbalizeFinalFst - elif lang=='id': - from fun_text_processing.inverse_text_normalization.id.taggers.tokenize_and_classify import ClassifyFst - from fun_text_processing.inverse_text_normalization.id.verbalizers.verbalize_final import VerbalizeFinalFst - elif lang=='ja': - from fun_text_processing.inverse_text_normalization.ja.taggers.tokenize_and_classify import ClassifyFst - from fun_text_processing.inverse_text_normalization.ja.verbalizers.verbalize_final import VerbalizeFinalFst - elif lang=='ko': - from fun_text_processing.inverse_text_normalization.ko.taggers.tokenize_and_classify import ClassifyFst - from fun_text_processing.inverse_text_normalization.ko.verbalizers.verbalize_final import VerbalizeFinalFst - elif lang=='pt': - from fun_text_processing.inverse_text_normalization.pt.taggers.tokenize_and_classify import ClassifyFst - from fun_text_processing.inverse_text_normalization.pt.verbalizers.verbalize_final import VerbalizeFinalFst - elif lang=='ru': - from fun_text_processing.inverse_text_normalization.ru.taggers.tokenize_and_classify import ClassifyFst - from fun_text_processing.inverse_text_normalization.ru.verbalizers.verbalize_final import VerbalizeFinalFst - elif lang=='vi': - from fun_text_processing.inverse_text_normalization.vi.taggers.tokenize_and_classify import ClassifyFst - from fun_text_processing.inverse_text_normalization.vi.verbalizers.verbalize_final import VerbalizeFinalFst - elif lang=='zh': - from fun_text_processing.inverse_text_normalization.zh.taggers.tokenize_and_classify import ClassifyFst - from fun_text_processing.inverse_text_normalization.zh.verbalizers.verbalize_final import VerbalizeFinalFst - else: - from fun_text_processing.inverse_text_normalization.en.taggers.tokenize_and_classify import ClassifyFst - from fun_text_processing.inverse_text_normalization.en.verbalizers.verbalize_final import VerbalizeFinalFst +def get_grammars(lang: str = "en"): + if lang == "de": + from fun_text_processing.inverse_text_normalization.de.taggers.tokenize_and_classify import ( + ClassifyFst, + ) + from fun_text_processing.inverse_text_normalization.de.verbalizers.verbalize_final import ( + VerbalizeFinalFst, + ) + elif lang == "en": + from fun_text_processing.inverse_text_normalization.en.taggers.tokenize_and_classify import ( + ClassifyFst, + ) + from fun_text_processing.inverse_text_normalization.en.verbalizers.verbalize_final import ( + VerbalizeFinalFst, + ) + elif lang == "es": + from fun_text_processing.inverse_text_normalization.es.taggers.tokenize_and_classify import ( + ClassifyFst, + ) + from fun_text_processing.inverse_text_normalization.es.verbalizers.verbalize_final import ( + VerbalizeFinalFst, + ) + elif lang == "fr": + from fun_text_processing.inverse_text_normalization.fr.taggers.tokenize_and_classify import ( + ClassifyFst, + ) + from fun_text_processing.inverse_text_normalization.fr.verbalizers.verbalize_final import ( + VerbalizeFinalFst, + ) + elif lang == "id": + from fun_text_processing.inverse_text_normalization.id.taggers.tokenize_and_classify import ( + ClassifyFst, + ) + from fun_text_processing.inverse_text_normalization.id.verbalizers.verbalize_final import ( + VerbalizeFinalFst, + ) + elif lang == "ja": + from fun_text_processing.inverse_text_normalization.ja.taggers.tokenize_and_classify import ( + ClassifyFst, + ) + from fun_text_processing.inverse_text_normalization.ja.verbalizers.verbalize_final import ( + VerbalizeFinalFst, + ) + elif lang == "ko": + from fun_text_processing.inverse_text_normalization.ko.taggers.tokenize_and_classify import ( + ClassifyFst, + ) + from fun_text_processing.inverse_text_normalization.ko.verbalizers.verbalize_final import ( + VerbalizeFinalFst, + ) + elif lang == "pt": + from fun_text_processing.inverse_text_normalization.pt.taggers.tokenize_and_classify import ( + ClassifyFst, + ) + from fun_text_processing.inverse_text_normalization.pt.verbalizers.verbalize_final import ( + VerbalizeFinalFst, + ) + elif lang == "ru": + from fun_text_processing.inverse_text_normalization.ru.taggers.tokenize_and_classify import ( + ClassifyFst, + ) + from fun_text_processing.inverse_text_normalization.ru.verbalizers.verbalize_final import ( + VerbalizeFinalFst, + ) + elif lang == "vi": + from fun_text_processing.inverse_text_normalization.vi.taggers.tokenize_and_classify import ( + ClassifyFst, + ) + from fun_text_processing.inverse_text_normalization.vi.verbalizers.verbalize_final import ( + VerbalizeFinalFst, + ) + elif lang == "zh": + from fun_text_processing.inverse_text_normalization.zh.taggers.tokenize_and_classify import ( + ClassifyFst, + ) + from fun_text_processing.inverse_text_normalization.zh.verbalizers.verbalize_final import ( + VerbalizeFinalFst, + ) + else: + from fun_text_processing.inverse_text_normalization.en.taggers.tokenize_and_classify import ( + ClassifyFst, + ) + from fun_text_processing.inverse_text_normalization.en.verbalizers.verbalize_final import ( + VerbalizeFinalFst, + ) + + return ClassifyFst().fst, VerbalizeFinalFst().fst - return ClassifyFst().fst, VerbalizeFinalFst().fst if __name__ == "__main__": args = parse_args() @@ -70,5 +124,4 @@ if __name__ == "__main__": tagger_fst, verbalizer_fst = get_grammars(args.language) generator_main(tagger_far_file, {"tokenize_and_classify": tagger_fst}) generator_main(verbalizer_far_file, {"verbalize": verbalizer_fst}) - print(f'Time to generate graph: {round(perf_counter() - start_time, 2)} sec') - + print(f"Time to generate graph: {round(perf_counter() - start_time, 2)} sec") diff --git a/fun_text_processing/inverse_text_normalization/fr/__init__.py b/fun_text_processing/inverse_text_normalization/fr/__init__.py index d30c1f8b2..a5fc721f1 100644 --- a/fun_text_processing/inverse_text_normalization/fr/__init__.py +++ b/fun_text_processing/inverse_text_normalization/fr/__init__.py @@ -1,4 +1,7 @@ - -from fun_text_processing.inverse_text_normalization.fr.taggers.tokenize_and_classify import ClassifyFst +from fun_text_processing.inverse_text_normalization.fr.taggers.tokenize_and_classify import ( + ClassifyFst, +) from fun_text_processing.inverse_text_normalization.fr.verbalizers.verbalize import VerbalizeFst -from fun_text_processing.inverse_text_normalization.fr.verbalizers.verbalize_final import VerbalizeFinalFst +from fun_text_processing.inverse_text_normalization.fr.verbalizers.verbalize_final import ( + VerbalizeFinalFst, +) diff --git a/fun_text_processing/inverse_text_normalization/fr/data/__init__.py b/fun_text_processing/inverse_text_normalization/fr/data/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/fr/data/__init__.py +++ b/fun_text_processing/inverse_text_normalization/fr/data/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/inverse_text_normalization/fr/data/electronic/__init__.py b/fun_text_processing/inverse_text_normalization/fr/data/electronic/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/fr/data/electronic/__init__.py +++ b/fun_text_processing/inverse_text_normalization/fr/data/electronic/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/inverse_text_normalization/fr/data/measurements/__init__.py b/fun_text_processing/inverse_text_normalization/fr/data/measurements/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/fr/data/measurements/__init__.py +++ b/fun_text_processing/inverse_text_normalization/fr/data/measurements/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/inverse_text_normalization/fr/data/money/__init__.py b/fun_text_processing/inverse_text_normalization/fr/data/money/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/fr/data/money/__init__.py +++ b/fun_text_processing/inverse_text_normalization/fr/data/money/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/inverse_text_normalization/fr/data/numbers/__init__.py b/fun_text_processing/inverse_text_normalization/fr/data/numbers/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/fr/data/numbers/__init__.py +++ b/fun_text_processing/inverse_text_normalization/fr/data/numbers/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/inverse_text_normalization/fr/data/ordinals/__init__.py b/fun_text_processing/inverse_text_normalization/fr/data/ordinals/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/fr/data/ordinals/__init__.py +++ b/fun_text_processing/inverse_text_normalization/fr/data/ordinals/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/inverse_text_normalization/fr/data/time/__init__.py b/fun_text_processing/inverse_text_normalization/fr/data/time/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/fr/data/time/__init__.py +++ b/fun_text_processing/inverse_text_normalization/fr/data/time/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/inverse_text_normalization/fr/graph_utils.py b/fun_text_processing/inverse_text_normalization/fr/graph_utils.py index f8b50864d..4b2cc420e 100644 --- a/fun_text_processing/inverse_text_normalization/fr/graph_utils.py +++ b/fun_text_processing/inverse_text_normalization/fr/graph_utils.py @@ -1,4 +1,3 @@ - import os import string from pathlib import Path @@ -19,9 +18,9 @@ DAMO_UPPER = pynini.union(*string.ascii_uppercase).optimize() DAMO_ALPHA = pynini.union(DAMO_LOWER, DAMO_UPPER).optimize() DAMO_ALNUM = pynini.union(DAMO_DIGIT, DAMO_ALPHA).optimize() DAMO_HEX = pynini.union(*string.hexdigits).optimize() -DAMO_NON_BREAKING_SPACE = u"\u00A0" +DAMO_NON_BREAKING_SPACE = "\u00A0" DAMO_SPACE = " " -DAMO_WHITE_SPACE = pynini.union(" ", "\t", "\n", "\r", u"\u00A0").optimize() +DAMO_WHITE_SPACE = pynini.union(" ", "\t", "\n", "\r", "\u00A0").optimize() DAMO_NOT_SPACE = pynini.difference(DAMO_CHAR, DAMO_WHITE_SPACE).optimize() DAMO_NOT_QUOTE = pynini.difference(DAMO_CHAR, r'"').optimize() @@ -49,7 +48,9 @@ graph_plural = plurals._priority_union( SINGULAR_TO_PLURAL = graph_plural PLURAL_TO_SINGULAR = pynini.invert(graph_plural) -TO_LOWER = pynini.union(*[pynini.cross(x, y) for x, y in zip(string.ascii_uppercase, string.ascii_lowercase)]) +TO_LOWER = pynini.union( + *[pynini.cross(x, y) for x, y in zip(string.ascii_uppercase, string.ascii_lowercase)] +) TO_UPPER = pynini.invert(TO_LOWER) @@ -65,7 +66,7 @@ def generator_main(file_name: str, graphs: Dict[str, pynini.FstLike]): for rule, graph in graphs.items(): exporter[rule] = graph.optimize() exporter.close() - print(f'Created {file_name}') + print(f"Created {file_name}") def get_plurals(fst): @@ -92,18 +93,20 @@ def get_singulars(fst): return PLURAL_TO_SINGULAR @ fst -def convert_space(fst) -> 'pynini.FstLike': +def convert_space(fst) -> "pynini.FstLike": """ Converts space to nonbreaking space. Used only in tagger grammars for transducing token values within quotes, e.g. name: "hello kitty" - This is making transducer significantly slower, so only use when there could be potential spaces within quotes, otherwise leave it. + This is making transducer significantly slower, so only use when there could be potential spaces within quotes, otherwise leave it. Args: fst: input fst Returns output fst where breaking spaces are converted to non breaking spaces """ - return fst @ pynini.cdrewrite(pynini.cross(DAMO_SPACE, DAMO_NON_BREAKING_SPACE), "", "", DAMO_SIGMA) + return fst @ pynini.cdrewrite( + pynini.cross(DAMO_SPACE, DAMO_NON_BREAKING_SPACE), "", "", DAMO_SIGMA + ) class GraphFst: @@ -123,9 +126,11 @@ class GraphFst: self._fst = None self.deterministic = deterministic - self.far_path = Path(os.path.dirname(__file__) + '/grammars/' + kind + '/' + name + '.far') + self.far_path = Path(os.path.dirname(__file__) + "/grammars/" + kind + "/" + name + ".far") if self.far_exist(): - self._fst = Far(self.far_path, mode="r", arc_type="standard", far_type="default").get_fst() + self._fst = Far( + self.far_path, mode="r", arc_type="standard", far_type="default" + ).get_fst() def far_exist(self) -> bool: """ @@ -134,26 +139,26 @@ class GraphFst: return self.far_path.exists() @property - def fst(self) -> 'pynini.FstLike': + def fst(self) -> "pynini.FstLike": return self._fst @fst.setter def fst(self, fst): self._fst = fst - def add_tokens(self, fst) -> 'pynini.FstLike': + def add_tokens(self, fst) -> "pynini.FstLike": """ Wraps class name around to given fst - Args: + Args: fst: input fst - + Returns: Fst: fst """ return pynutil.insert(f"{self.name} {{ ") + fst + pynutil.insert(" }") - def delete_tokens(self, fst) -> 'pynini.FstLike': + def delete_tokens(self, fst) -> "pynini.FstLike": """ Deletes class name wrap around output of given fst @@ -172,4 +177,4 @@ class GraphFst: + delete_space + pynutil.delete("}") ) - return res @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + return res @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) diff --git a/fun_text_processing/inverse_text_normalization/fr/taggers/__init__.py b/fun_text_processing/inverse_text_normalization/fr/taggers/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/fr/taggers/__init__.py +++ b/fun_text_processing/inverse_text_normalization/fr/taggers/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/inverse_text_normalization/fr/taggers/cardinal.py b/fun_text_processing/inverse_text_normalization/fr/taggers/cardinal.py index e78f86743..39901a954 100644 --- a/fun_text_processing/inverse_text_normalization/fr/taggers/cardinal.py +++ b/fun_text_processing/inverse_text_normalization/fr/taggers/cardinal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.fr.graph_utils import ( DAMO_CHAR, @@ -13,7 +12,7 @@ from fun_text_processing.inverse_text_normalization.fr.utils import get_abs_path from pynini.lib import pynutil -def rewrite(cardinal: 'pynini.FstLike') -> 'pynini.FstLike': +def rewrite(cardinal: "pynini.FstLike") -> "pynini.FstLike": """ Function to rewrite cardinals written in traditional orthograph (no '-' for numbers >100) to current orthography ('-' between all words in number string) @@ -21,9 +20,9 @@ def rewrite(cardinal: 'pynini.FstLike') -> 'pynini.FstLike': In cases where original orthography is current, or string is mixture of two orthographies, will render invalid form that will not pass through CardinalFst e.g. deux-mille cent-vingt-trois -> "deux##vingt-trois" ('#' is not accepted in cardinal FST and will fail to convert.) - e.g. deux + e.g. deux - Args: + Args: cardinal: cardinal FST """ @@ -52,9 +51,13 @@ def rewrite(cardinal: 'pynini.FstLike') -> 'pynini.FstLike': # Old orthography will not have these strings. Replacing with character to mark. targets_for_filtering = ("-" + targets) | ("-" + targets + "-") | (targets + "-") - filter = pynini.cdrewrite(pynini.cross(targets_for_filtering, "#"), "", "", DAMO_SIGMA) # Invalid for cardinal + filter = pynini.cdrewrite( + pynini.cross(targets_for_filtering, "#"), "", "", DAMO_SIGMA + ) # Invalid for cardinal - old_orthography_sigma = pynini.difference(DAMO_CHAR, "#") # Marked character removed from sigma_star. + old_orthography_sigma = pynini.difference( + DAMO_CHAR, "#" + ) # Marked character removed from sigma_star. old_orthography_sigma.closure() # Only accept strings that occur in old orthography. (This avoids tying two non-related numbers together.) @@ -76,13 +79,13 @@ def rewrite(cardinal: 'pynini.FstLike') -> 'pynini.FstLike': class CardinalFst(GraphFst): """ Finite state transducer for classifying cardinals - e.g. mois vingt-trois -> cardinal { negative: "-" integer: "23"} + e.g. mois vingt-trois -> cardinal { negative: "-" integer: "23"} This class converts cardinals up to (but not including) "un-quatrillion", i.e up to "one septillion" in English (10^{24}). - Cardinals below nine are not converted (in order to avoid + Cardinals below nine are not converted (in order to avoid "j'ai un pomme." --> "j'ai 1 pomme" and any other odd conversions.) This transducer accomodates both traditional hyphenation of numbers ('-' for most numbers <100) - and current hyphenation (all elements of number are hyphenated), prioritizing the latter. + and current hyphenation (all elements of number are hyphenated), prioritizing the latter. e.g cent cinquante et un -> cardinal { integer: "151"} cent-cinquante-et-un -> cardinal { integer: "151"} This is done through a context dependent rewrite that attempts to map old spelling to new. @@ -102,7 +105,8 @@ class CardinalFst(GraphFst): graph_tens_component = pynini.union(graph_tens_component, graph_teens, graph_ties_unique) graph_tens_component_with_leading_zeros = pynini.union( - graph_tens_component, (pynutil.insert("0") + (graph_digit | pynutil.insert("0", weight=0.01))) + graph_tens_component, + (pynutil.insert("0") + (graph_digit | pynutil.insert("0", weight=0.01))), ) # Hundreds components @@ -111,14 +115,18 @@ class CardinalFst(GraphFst): "cents", "00" ) # Only used as terminus of hundred sequence. deux cents -> 200, deux cent un -> 201 - graph_digit_no_one = pynini.project(pynini.union("un", "une"), 'input') - graph_digit_no_one = (pynini.project(graph_digit, "input") - graph_digit_no_one.arcsort()) @ graph_digit + graph_digit_no_one = pynini.project(pynini.union("un", "une"), "input") + graph_digit_no_one = ( + pynini.project(graph_digit, "input") - graph_digit_no_one.arcsort() + ) @ graph_digit graph_hundreds_component_singular = ( graph_digit_no_one + delete_hyphen + graph_cent_singular ) # Regular way: [1-9] * 100 - graph_hundreds_component_singular = pynini.union(graph_hundreds_component_singular, pynini.cross("cent", "1")) + graph_hundreds_component_singular = pynini.union( + graph_hundreds_component_singular, pynini.cross("cent", "1") + ) graph_hundreds_component_singular += delete_hyphen graph_hundreds_component_singular += graph_tens_component_with_leading_zeros @@ -141,14 +149,20 @@ class CardinalFst(GraphFst): graph_tens_of_hundreds_component_singular = ( graph_tens_component + delete_hyphen + graph_cent_singular ) # Tens of hundreds. e.g. 1900 = nineteen hundred/ 'dix neuf cents" - graph_tens_of_hundreds_component_singular += delete_hyphen + graph_tens_component_with_leading_zeros - graph_tens_of_hundreds_component_plural = graph_tens_component + delete_hyphen + graph_cent_plural + graph_tens_of_hundreds_component_singular += ( + delete_hyphen + graph_tens_component_with_leading_zeros + ) + graph_tens_of_hundreds_component_plural = ( + graph_tens_component + delete_hyphen + graph_cent_plural + ) graph_tens_of_hundred_component = ( graph_tens_of_hundreds_component_plural | graph_tens_of_hundreds_component_singular ) graph_thousands = pynini.union( - graph_hundreds_component_at_least_one_none_zero_digit + delete_hyphen + pynutil.delete("mille"), + graph_hundreds_component_at_least_one_none_zero_digit + + delete_hyphen + + pynutil.delete("mille"), pynutil.insert("001") + pynutil.delete("mille"), # because 'mille', not 'un mille' pynutil.insert("000", weight=0.1), ) @@ -176,8 +190,11 @@ class CardinalFst(GraphFst): ) graph_mille_billion = pynini.union( - graph_hundreds_component_at_least_one_none_zero_digit + delete_hyphen + pynutil.delete("mille"), - pynutil.insert("001") + pynutil.delete("mille"), # because we say 'mille', not 'un mille' + graph_hundreds_component_at_least_one_none_zero_digit + + delete_hyphen + + pynutil.delete("mille"), + pynutil.insert("001") + + pynutil.delete("mille"), # because we say 'mille', not 'un mille' ) graph_mille_billion += delete_hyphen + ( graph_millions | pynutil.insert("000") + pynutil.delete("billions") @@ -226,7 +243,10 @@ class CardinalFst(GraphFst): ) graph = graph @ pynini.union( - pynutil.delete(pynini.closure("0")) + pynini.difference(DAMO_DIGIT, "0") + pynini.closure(DAMO_DIGIT), "0" + pynutil.delete(pynini.closure("0")) + + pynini.difference(DAMO_DIGIT, "0") + + pynini.closure(DAMO_DIGIT), + "0", ) graph = rewrite(graph) @@ -234,32 +254,34 @@ class CardinalFst(GraphFst): self.graph_no_exception = graph.optimize() # save self.numbers_up_to_thousand for use in DecimalFst - digits_up_to_thousand = DAMO_DIGIT | (DAMO_DIGIT ** 2) | (DAMO_DIGIT ** 3) + digits_up_to_thousand = DAMO_DIGIT | (DAMO_DIGIT**2) | (DAMO_DIGIT**3) numbers_up_to_thousand = pynini.compose(graph, digits_up_to_thousand).optimize() self.numbers_up_to_thousand = numbers_up_to_thousand # save self.numbers_up_to_million for use in DecimalFst digits_up_to_million = ( DAMO_DIGIT - | (DAMO_DIGIT ** 2) - | (DAMO_DIGIT ** 3) - | (DAMO_DIGIT ** 4) - | (DAMO_DIGIT ** 5) - | (DAMO_DIGIT ** 6) + | (DAMO_DIGIT**2) + | (DAMO_DIGIT**3) + | (DAMO_DIGIT**4) + | (DAMO_DIGIT**5) + | (DAMO_DIGIT**6) ) numbers_up_to_million = pynini.compose(graph, digits_up_to_million).optimize() self.numbers_up_to_million = numbers_up_to_million # don't convert cardinals from zero to nine inclusive - graph_exception = pynini.project(pynini.union(graph_digit, graph_zero), 'input') + graph_exception = pynini.project(pynini.union(graph_digit, graph_zero), "input") self.graph = (pynini.project(graph, "input") - graph_exception.arcsort()) @ graph optional_minus_graph = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("moins", "\"-\"") + DAMO_SPACE, 0, 1 + pynutil.insert("negative: ") + pynini.cross("moins", '"-"') + DAMO_SPACE, 0, 1 ) - final_graph = optional_minus_graph + pynutil.insert("integer: \"") + self.graph + pynutil.insert("\"") + final_graph = ( + optional_minus_graph + pynutil.insert('integer: "') + self.graph + pynutil.insert('"') + ) final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/fr/taggers/date.py b/fun_text_processing/inverse_text_normalization/fr/taggers/date.py index 9e513c011..99972062e 100644 --- a/fun_text_processing/inverse_text_normalization/fr/taggers/date.py +++ b/fun_text_processing/inverse_text_normalization/fr/taggers/date.py @@ -1,6 +1,8 @@ - import pynini -from fun_text_processing.inverse_text_normalization.fr.graph_utils import GraphFst, delete_extra_space +from fun_text_processing.inverse_text_normalization.fr.graph_utils import ( + GraphFst, + delete_extra_space, +) from fun_text_processing.inverse_text_normalization.fr.utils import get_abs_path from pynini.lib import pynutil @@ -27,12 +29,16 @@ class DateFst(GraphFst): year_graph = self.cardinal month_graph = pynini.string_file(get_abs_path("data/months.tsv")) - month_graph = pynutil.insert("month: \"") + month_graph + pynutil.insert("\"") + month_graph = pynutil.insert('month: "') + month_graph + pynutil.insert('"') - day_graph = self.cardinal | pynini.cross("premier", "1") # Premier is only ordinal used for dates - day_graph = pynutil.insert("day: \"") + day_graph + pynutil.insert("\"") + day_graph = self.cardinal | pynini.cross( + "premier", "1" + ) # Premier is only ordinal used for dates + day_graph = pynutil.insert('day: "') + day_graph + pynutil.insert('"') optional_graph_year = pynini.closure( - delete_extra_space + pynutil.insert("year: \"") + year_graph + pynutil.insert("\""), 0, 1, + delete_extra_space + pynutil.insert('year: "') + year_graph + pynutil.insert('"'), + 0, + 1, ) graph_dmy = day_graph + delete_extra_space + month_graph + optional_graph_year diff --git a/fun_text_processing/inverse_text_normalization/fr/taggers/decimal.py b/fun_text_processing/inverse_text_normalization/fr/taggers/decimal.py index 2199bc710..9a00a6a47 100644 --- a/fun_text_processing/inverse_text_normalization/fr/taggers/decimal.py +++ b/fun_text_processing/inverse_text_normalization/fr/taggers/decimal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.fr.graph_utils import ( DAMO_DIGIT, @@ -11,20 +10,24 @@ from fun_text_processing.inverse_text_normalization.fr.utils import get_abs_path from pynini.lib import pynutil -def get_quantity(decimal: 'pynini.FstLike', cardinal_up_to_thousand: 'pynini.FstLike') -> 'pynini.FstLike': +def get_quantity( + decimal: "pynini.FstLike", cardinal_up_to_thousand: "pynini.FstLike" +) -> "pynini.FstLike": """ Returns FST that transforms either a cardinal or decimal followed by a quantity into a numeral, e.g. one million -> integer_part: "1" quantity: "million" e.g. one point five million -> integer_part: "1" fractional_part: "5" quantity: "million" - Will tag cases up to denominations of tens of hundreds of thousand. 'douze cent mille millions' -> 1 200 000 millions + Will tag cases up to denominations of tens of hundreds of thousand. 'douze cent mille millions' -> 1 200 000 millions - Args: + Args: decimal: decimal FST cardinal_up_to_million: cardinal FST """ numbers = cardinal_up_to_thousand @ ( - pynutil.delete(pynini.closure("0")) + pynini.difference(DAMO_DIGIT, "0") + pynini.closure(DAMO_DIGIT) + pynutil.delete(pynini.closure("0")) + + pynini.difference(DAMO_DIGIT, "0") + + pynini.closure(DAMO_DIGIT) ) suffix = pynini.union( @@ -42,17 +45,19 @@ def get_quantity(decimal: 'pynini.FstLike', cardinal_up_to_thousand: 'pynini.Fst "trilliards", ) res = ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + numbers - + pynutil.insert("\"") + + pynutil.insert('"') + ( pynini.union(delete_hyphen, delete_extra_space) ) # Can be written either as 'deux-millions' or 'deux millions' depending on whether it registers as a noun or part of cardinal. - + pynutil.insert(" quantity: \"") + + pynutil.insert(' quantity: "') + suffix - + pynutil.insert("\"") + + pynutil.insert('"') + ) + res |= ( + decimal + delete_extra_space + pynutil.insert(' quantity: "') + suffix + pynutil.insert('"') ) - res |= decimal + delete_extra_space + pynutil.insert(" quantity: \"") + suffix + pynutil.insert("\"") return res @@ -65,7 +70,7 @@ class DecimalFst(GraphFst): This decimal rule assumes that decimals can be pronounced as: (a cardinal) + ('virgule') plus (any sequence of cardinals <1 million, including 'zero') - Also writes large numbers in shortened form, e.g. + Also writes large numbers in shortened form, e.g. e.g. un virgule deux-six-million -> decimal { negative: "false" integer_part: "1" fractional_part: "26" quantity: "million" } e.g. deux-million -> decimal { negative: "false" integer_part: "2" quantity: "millions" } e.g. moins cent-vingt-quatre-millions -> decimal { negative: "true" integer_part: "124" quantity: "millions" } @@ -86,15 +91,21 @@ class DecimalFst(GraphFst): graph_fractional_separator = pynutil.delete("virgule") # Possible negatives - optional_graph_negative = pynutil.insert("negative: ") + pynini.cross("moins", "\"true\"") + delete_extra_space + optional_graph_negative = ( + pynutil.insert("negative: ") + pynini.cross("moins", '"true"') + delete_extra_space + ) optional_graph_negative = optional_graph_negative.ques # Fractional portion - graph_fractional = pynutil.insert("fractional_part: \"") + graph_decimal + pynutil.insert("\"") + graph_fractional = ( + pynutil.insert('fractional_part: "') + graph_decimal + pynutil.insert('"') + ) # Integers - cardinal_graph = cardinal.graph_no_exception | pynini.string_file(get_abs_path("data/numbers/zero.tsv")) - graph_integer = pynutil.insert("integer_part: \"") + cardinal_graph + pynutil.insert("\"") + cardinal_graph = cardinal.graph_no_exception | pynini.string_file( + get_abs_path("data/numbers/zero.tsv") + ) + graph_integer = pynutil.insert('integer_part: "') + cardinal_graph + pynutil.insert('"') # Final graphs final_graph_wo_sign = ( diff --git a/fun_text_processing/inverse_text_normalization/fr/taggers/electronic.py b/fun_text_processing/inverse_text_normalization/fr/taggers/electronic.py index 04dfcbd17..8c4d217ae 100644 --- a/fun_text_processing/inverse_text_normalization/fr/taggers/electronic.py +++ b/fun_text_processing/inverse_text_normalization/fr/taggers/electronic.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.fr.graph_utils import DAMO_ALPHA, GraphFst, insert_space +from fun_text_processing.inverse_text_normalization.fr.graph_utils import ( + DAMO_ALPHA, + GraphFst, + insert_space, +) from fun_text_processing.inverse_text_normalization.fr.utils import get_abs_path from pynini.lib import pynutil @@ -30,24 +33,26 @@ class ElectronicFst(GraphFst): accepted_username = alpha_num | symbols process_dot = pynini.cross("point", ".") username = ( - pynutil.insert("username: \"") + pynutil.insert('username: "') + alpha_num + delete_extra_space + pynini.closure(accepted_username + delete_extra_space) + alpha_num - + pynutil.insert("\"") + + pynutil.insert('"') ) single_alphanum = pynini.closure(alpha_num + delete_extra_space) + alpha_num - server = single_alphanum | pynini.string_file(get_abs_path("data/electronic/server_name.tsv")) + server = single_alphanum | pynini.string_file( + get_abs_path("data/electronic/server_name.tsv") + ) domain = single_alphanum | pynini.string_file(get_abs_path("data/electronic/domain.tsv")) domain_graph = ( - pynutil.insert("domain: \"") + pynutil.insert('domain: "') + server + delete_extra_space + process_dot + delete_extra_space + domain - + pynutil.insert("\"") + + pynutil.insert('"') ) graph = ( username @@ -59,9 +64,13 @@ class ElectronicFst(GraphFst): ) ############# url ### - protocol_end = pynini.cross(pynini.union("www", "w w w", "double vé double vé double vé"), "www") + protocol_end = pynini.cross( + pynini.union("www", "w w w", "double vé double vé double vé"), "www" + ) protocol_start = pynini.cross(pynini.union("http", "h t t p", "ache té té pé"), "http") - protocol_start |= pynini.cross(pynini.union("https", "h t t p s", "ache té té pé esse"), "https") + protocol_start |= pynini.cross( + pynini.union("https", "h t t p s", "ache té té pé esse"), "https" + ) protocol_start += pynini.cross( pynini.union( " deux-points barre oblique barre oblique ", @@ -89,7 +98,7 @@ class ElectronicFst(GraphFst): + (pynini.closure(delete_extra_space + accepted_username, 1) | server) + pynini.closure(ending, 1) ) - protocol = pynutil.insert("protocol: \"") + protocol + pynutil.insert("\"") + protocol = pynutil.insert('protocol: "') + protocol + pynutil.insert('"') graph |= protocol ######## diff --git a/fun_text_processing/inverse_text_normalization/fr/taggers/fraction.py b/fun_text_processing/inverse_text_normalization/fr/taggers/fraction.py index 324f77ee9..b6a7b0fd9 100644 --- a/fun_text_processing/inverse_text_normalization/fr/taggers/fraction.py +++ b/fun_text_processing/inverse_text_normalization/fr/taggers/fraction.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.fr.graph_utils import ( DAMO_CHAR, @@ -16,7 +15,7 @@ class FractionFst(GraphFst): e.g. demi -> tokens { fraction { numerator: "1" denominator: "2" } } e.g. un et demi -> tokens { fraction { integer_part: "1" numerator: "1" denominator: "2" } } e.g. trois et deux centième -> tokens { fraction { integer_part: "3" numerator: "2" denominator: "100" } } - + Args: cardinal: OrdinalFst """ @@ -26,7 +25,9 @@ class FractionFst(GraphFst): # integer_part # numerator # denominator graph_cardinal = cardinal.graph_no_exception - graph_strip_undo_root_change = pynini.string_file(get_abs_path("data/fractions.tsv")) # add in absolute path + graph_strip_undo_root_change = pynini.string_file( + get_abs_path("data/fractions.tsv") + ) # add in absolute path graph_strip_no_root_change = pynutil.delete("ième") # For no change to root graph_strip_no_root_change += pynutil.delete("s").ques # for plurals @@ -35,16 +36,18 @@ class FractionFst(GraphFst): self.fractional = ((pynini.closure(DAMO_CHAR) + graph_strip) @ graph_cardinal).optimize() - integer = pynutil.insert("integer_part: \"") + graph_cardinal + pynutil.insert("\" ") + integer = pynutil.insert('integer_part: "') + graph_cardinal + pynutil.insert('" ') integer += delete_space integer += pynutil.delete("et") # used to demarcate integer and fractional parts - numerator = pynutil.insert("numerator: \"") + graph_cardinal + pynutil.insert("\"") - denominator = pynutil.insert(" denominator: \"") + self.fractional + pynutil.insert("\"") + numerator = pynutil.insert('numerator: "') + graph_cardinal + pynutil.insert('"') + denominator = pynutil.insert(' denominator: "') + self.fractional + pynutil.insert('"') # Demi (half) can occur alone without explicit numerator. - graph_demi_component = pynutil.delete("demi") + pynutil.delete("e").ques + pynutil.delete("s").ques - graph_demi_component += pynutil.insert("numerator: \"1\" denominator: \"2\"") + graph_demi_component = ( + pynutil.delete("demi") + pynutil.delete("e").ques + pynutil.delete("s").ques + ) + graph_demi_component += pynutil.insert('numerator: "1" denominator: "2"') graph_fraction_component = numerator + delete_space + denominator graph_fraction_component |= graph_demi_component @@ -55,7 +58,9 @@ class FractionFst(GraphFst): self.final_graph_wo_negative = graph optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("moins", "\"true\"") + delete_extra_space, 0, 1 + pynutil.insert("negative: ") + pynini.cross("moins", '"true"') + delete_extra_space, + 0, + 1, ) graph = optional_graph_negative + graph final_graph = self.add_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/fr/taggers/measure.py b/fun_text_processing/inverse_text_normalization/fr/taggers/measure.py index 1cface9c2..07854a838 100644 --- a/fun_text_processing/inverse_text_normalization/fr/taggers/measure.py +++ b/fun_text_processing/inverse_text_normalization/fr/taggers/measure.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.fr.graph_utils import ( GraphFst, @@ -35,15 +34,22 @@ class MeasureFst(GraphFst): unit = graph_prefix.ques + unit optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("moins", "\"true\"") + delete_extra_space, 0, 1 + pynutil.insert("negative: ") + pynini.cross("moins", '"true"') + delete_extra_space, + 0, + 1, ) - unit_misc = pynutil.insert("/") + (pynutil.delete("par") | pynutil.delete("à")) + delete_space + unit + unit_misc = ( + pynutil.insert("/") + + (pynutil.delete("par") | pynutil.delete("à")) + + delete_space + + unit + ) unit = ( - pynutil.insert("units: \"") + pynutil.insert('units: "') + (unit | unit_misc | pynutil.add_weight(unit + delete_space + unit_misc, 0.01)) - + pynutil.insert("\"") + + pynutil.insert('"') ) subgraph_decimal = ( @@ -67,9 +73,9 @@ class MeasureFst(GraphFst): subgraph_cardinal = ( pynutil.insert("cardinal { ") + optional_graph_negative - + pynutil.insert("integer: \"") + + pynutil.insert('integer: "') + cardinal_graph - + pynutil.insert("\"") + + pynutil.insert('"') + pynutil.insert(" }") + delete_extra_space + unit diff --git a/fun_text_processing/inverse_text_normalization/fr/taggers/money.py b/fun_text_processing/inverse_text_normalization/fr/taggers/money.py index b2e8b3a03..1236f53fa 100644 --- a/fun_text_processing/inverse_text_normalization/fr/taggers/money.py +++ b/fun_text_processing/inverse_text_normalization/fr/taggers/money.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.fr.graph_utils import ( DAMO_DIGIT, @@ -52,10 +51,14 @@ class MoneyFst(GraphFst): ) # recognizes all currencies # Graphs for large round amounts ('deux billiards d'euros', 'un milliard de dollars') - graph_de = pynini.union("de ", "des ", "d'") # the use of de/d'only occurs with round amounts + graph_de = pynini.union( + "de ", "des ", "d'" + ) # the use of de/d'only occurs with round amounts graph_currency_component_large_round_amounts = graph_de + accept_all_currency graph_currency_component_large_round_amounts = ( - pynutil.insert(" currency: \"") + graph_currency_component_large_round_amounts + pynutil.insert("\"") + pynutil.insert(' currency: "') + + graph_currency_component_large_round_amounts + + pynutil.insert('"') ) graph_money_large_round_amounts = ( @@ -64,21 +67,25 @@ class MoneyFst(GraphFst): graph_money_large_round_amounts += graph_currency_component_large_round_amounts # For standard currency - add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | (pynutil.insert("0") + DAMO_DIGIT) + add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | ( + pynutil.insert("0") + DAMO_DIGIT + ) # Graphs integer denomination for large denominations (e.g. $) - graph_integer_component_major = pynutil.insert("integer_part: \"") + cardinal_graph + pynutil.insert("\"") + graph_integer_component_major = ( + pynutil.insert('integer_part: "') + cardinal_graph + pynutil.insert('"') + ) graph_integer_component_major += delete_space graph_currency_component_major = ( - pynutil.insert(" currency: \"") + convert_currency_major + pynutil.insert("\"") + pynutil.insert(' currency: "') + convert_currency_major + pynutil.insert('"') ) graph_decimal_component_major = ( delete_space - + pynutil.insert(" fractional_part: \"") + + pynutil.insert(' fractional_part: "') + (cardinal_graph @ add_leading_zero_to_double_digit) - + pynutil.insert("\"") + + pynutil.insert('"') ) # Rare cases where 'et' will separate major and minor denominations. @@ -96,21 +103,23 @@ class MoneyFst(GraphFst): ) # For cases when only small denominations are used. - graph_integer_component_minor = pynutil.insert("integer_part: \"0\"") + graph_integer_component_minor = pynutil.insert('integer_part: "0"') graph_decimal_component_minor = ( - pynutil.insert(" fractional_part: \"") + pynutil.insert(' fractional_part: "') + (cardinal_graph @ add_leading_zero_to_double_digit) - + pynutil.insert("\"") + + pynutil.insert('"') ) graph_decimal_component_minor += delete_extra_space graph_currency_component_minor = ( - pynutil.insert(" currency: \"") + convert_currency_minor + pynutil.insert("\"") + pynutil.insert(' currency: "') + convert_currency_minor + pynutil.insert('"') ) graph_money_minor = ( - graph_integer_component_minor + graph_decimal_component_minor + graph_currency_component_minor + graph_integer_component_minor + + graph_decimal_component_minor + + graph_currency_component_minor ) graph_money_standard_amounts = graph_money_major | graph_money_minor diff --git a/fun_text_processing/inverse_text_normalization/fr/taggers/ordinal.py b/fun_text_processing/inverse_text_normalization/fr/taggers/ordinal.py index 633f849f0..faeb62fba 100644 --- a/fun_text_processing/inverse_text_normalization/fr/taggers/ordinal.py +++ b/fun_text_processing/inverse_text_normalization/fr/taggers/ordinal.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.fr.graph_utils import DAMO_SIGMA, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.fr.graph_utils import ( + DAMO_SIGMA, + GraphFst, + delete_space, +) from fun_text_processing.inverse_text_normalization.fr.utils import get_abs_path from pynini.lib import pynutil @@ -10,7 +13,7 @@ class OrdinalFst(GraphFst): Finite state transducer for classifying ordinal vingt-deuxième -> ordinal { integer: "22" morphosyntactic_features: "e" } - Also notes specific nouns that have unique normalization conventions. + Also notes specific nouns that have unique normalization conventions. For instance, 'siècles' are rendered in roman numerals when given an ordinal adjective. e.g. dix-neuvième siècle -> XIXe @@ -41,24 +44,24 @@ class OrdinalFst(GraphFst): graph_morpheme_component += pynini.accep("s").ques # In case of plurals. # Concatenate with cardinal graph. - graph_ordinal = pynutil.insert("integer: \"") + graph_integer_component + pynutil.insert("\"") + graph_ordinal = pynutil.insert('integer: "') + graph_integer_component + pynutil.insert('"') graph_ordinal += ( - pynutil.insert(" morphosyntactic_features: \"") + graph_morpheme_component + pynutil.insert(' morphosyntactic_features: "') + graph_morpheme_component ) # Leave open in case further morphems occur # Primer has a different subscript depending on gender, need to take note if # 'premier' or 'première' - graph_firsts = pynutil.insert("integer: \"1\" morphosyntactic_features: \"") + graph_firsts + graph_firsts = pynutil.insert('integer: "1" morphosyntactic_features: "') + graph_firsts # Second used 'd' as a superscript. - graph_second = pynutil.insert("integer: \"2\" morphosyntactic_features: \"") + graph_second + graph_second = pynutil.insert('integer: "2" morphosyntactic_features: "') + graph_second graph = graph_firsts | graph_second | graph_ordinal # For roman numerals. Carries over designation to verbalizer graph_special_ordinals = pynutil.insert("/") + delete_space + graph_special_ordinals - graph += graph_special_ordinals.ques + pynutil.insert("\"") + graph += graph_special_ordinals.ques + pynutil.insert('"') final_graph = self.add_tokens(graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/fr/taggers/punctuation.py b/fun_text_processing/inverse_text_normalization/fr/taggers/punctuation.py index 5e7f706a7..ae5dc1e2f 100644 --- a/fun_text_processing/inverse_text_normalization/fr/taggers/punctuation.py +++ b/fun_text_processing/inverse_text_normalization/fr/taggers/punctuation.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.fr.graph_utils import GraphFst from pynini.lib import pynutil @@ -13,11 +12,11 @@ class PunctuationFst(GraphFst): def __init__(self): super().__init__(name="punctuation", kind="classify") - s = "!#$%&\'()*+,-./:;<=>?@^_`{|}~" + s = "!#$%&'()*+,-./:;<=>?@^_`{|}~" guillemets = "\u00AB" + "\u00BB" # quotation marks in French. s += guillemets punct = pynini.union(*s) - graph = pynutil.insert("name: \"") + punct + pynutil.insert("\"") + graph = pynutil.insert('name: "') + punct + pynutil.insert('"') self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/fr/taggers/telephone.py b/fun_text_processing/inverse_text_normalization/fr/taggers/telephone.py index 6705ade07..efa41415b 100644 --- a/fun_text_processing/inverse_text_normalization/fr/taggers/telephone.py +++ b/fun_text_processing/inverse_text_normalization/fr/taggers/telephone.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.fr.graph_utils import ( GraphFst, @@ -13,7 +12,7 @@ from pynini.lib import pynutil class TelephoneFst(GraphFst): """ Finite state transducer for classifying telephone numbers. Assumes conventional grouping for Metropolitan France (and overseas departments) - (two number sequences are grouped as individual cardinals) or digit by digit (chiffre-par-chiffre) e.g. + (two number sequences are grouped as individual cardinals) or digit by digit (chiffre-par-chiffre) e.g. "zero un quatre-vingt-deux zero deux vingt-deux cinquante" -> { number_part: "01 42 02 22 50" } "zero un quatre deux zero deux deux deux cinq zero" -> { number_part: "01 42 02 22 50" } @@ -57,12 +56,16 @@ class TelephoneFst(GraphFst): # Paired digits graph_pair_digits_and_ties = double_digits | graph_pair_all_digits - graph_digits_and_ties = pynini.closure(graph_pair_digits_and_ties + delete_space + insert_space, 3, 3) - graph_digits_and_ties = graph_first_pair + graph_digits_and_ties + graph_pair_digits_and_ties + graph_digits_and_ties = pynini.closure( + graph_pair_digits_and_ties + delete_space + insert_space, 3, 3 + ) + graph_digits_and_ties = ( + graph_first_pair + graph_digits_and_ties + graph_pair_digits_and_ties + ) number_part = pynini.union(graph_all_digits, graph_digits_and_ties) - number_part = pynutil.insert("number_part: \"") + number_part + pynutil.insert("\"") + number_part = pynutil.insert('number_part: "') + number_part + pynutil.insert('"') graph = number_part final_graph = self.add_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/fr/taggers/time.py b/fun_text_processing/inverse_text_normalization/fr/taggers/time.py index ff534f741..f62efccf8 100644 --- a/fun_text_processing/inverse_text_normalization/fr/taggers/time.py +++ b/fun_text_processing/inverse_text_normalization/fr/taggers/time.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.fr.graph_utils import GraphFst, delete_space from fun_text_processing.inverse_text_normalization.fr.utils import get_abs_path @@ -52,9 +51,9 @@ class TimeFst(GraphFst): graph_demi += pynini.accep("e").ques # people vary on feminine or masculine form graph_demi = pynini.cross(graph_demi, "30") - graph_quart = pynini.accep('quart') + graph_quart = pynini.accep("quart") graph_quart = pynini.accep("le ").ques + graph_quart # sometimes used - graph_quart = pynini.cross(graph_quart, '15') + graph_quart = pynini.cross(graph_quart, "15") graph_trois_quart = pynini.cross("trois quarts", "45") graph_fractions = pynini.union(graph_demi, graph_quart, graph_trois_quart) @@ -63,12 +62,16 @@ class TimeFst(GraphFst): # Hours component is usually just a cardinal + 'heures' (ignored in case of 'midi/minuit'). graph_hours_component = pynini.union(graph_hours, graph_midi, graph_minuit) - graph_hours_component = pynutil.insert("hours: \"") + graph_hours_component + pynutil.insert("\"") + graph_hours_component = ( + pynutil.insert('hours: "') + graph_hours_component + pynutil.insert('"') + ) graph_hours_component += delete_space # Minutes component graph_minutes_component = ( - pynutil.insert(" minutes: \"") + pynini.union(graph_minutes, graph_et_fractions) + pynutil.insert("\"") + pynutil.insert(' minutes: "') + + pynini.union(graph_minutes, graph_et_fractions) + + pynutil.insert('"') ) # Hour and minutes together. For 'demi' and 'qart', 'et' is used as a conjunction. @@ -80,18 +83,22 @@ class TimeFst(GraphFst): graph_hours_to_component = graph_hours | graph_midi | graph_minuit graph_hours_to_component @= graph_hours_to - graph_hours_to_component = pynutil.insert("hours: \"") + graph_hours_to_component + pynutil.insert("\"") + graph_hours_to_component = ( + pynutil.insert('hours: "') + graph_hours_to_component + pynutil.insert('"') + ) graph_hours_to_component += delete_space graph_minutes_to_component = pynini.union(graph_minutes, graph_fractions) graph_minutes_to_component @= graph_minutes_to - graph_minutes_to_component = pynutil.insert(" minutes: \"") + graph_minutes_to_component + pynutil.insert("\"") + graph_minutes_to_component = ( + pynutil.insert(' minutes: "') + graph_minutes_to_component + pynutil.insert('"') + ) graph_time_to = graph_hours_to_component + graph_moins + graph_minutes_to_component graph_time_no_suffix = graph_time_standard | graph_time_to - graph_suffix_component = pynutil.insert(" suffix: \"") + graph_suffix + pynutil.insert("\"") + graph_suffix_component = pynutil.insert(' suffix: "') + graph_suffix + pynutil.insert('"') graph_suffix_component = delete_space + graph_suffix_component graph_suffix_component = graph_suffix_component.ques diff --git a/fun_text_processing/inverse_text_normalization/fr/taggers/tokenize_and_classify.py b/fun_text_processing/inverse_text_normalization/fr/taggers/tokenize_and_classify.py index 3f722f658..69d64e6d2 100644 --- a/fun_text_processing/inverse_text_normalization/fr/taggers/tokenize_and_classify.py +++ b/fun_text_processing/inverse_text_normalization/fr/taggers/tokenize_and_classify.py @@ -1,4 +1,3 @@ - import os import pynini @@ -29,7 +28,7 @@ import logging class ClassifyFst(GraphFst): """ Final class that composes all other classification grammars. This class can process an entire sentence, that is lower cased. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. Args: @@ -87,10 +86,16 @@ class ClassifyFst(GraphFst): | pynutil.add_weight(word_graph, 100) ) - punct = pynutil.insert("tokens { ") + pynutil.add_weight(punct_graph, weight=1.1) + pynutil.insert(" }") + punct = ( + pynutil.insert("tokens { ") + + pynutil.add_weight(punct_graph, weight=1.1) + + pynutil.insert(" }") + ) token = pynutil.insert("tokens { ") + classify + pynutil.insert(" }") token_plus_punct = ( - pynini.closure(punct + pynutil.insert(" ")) + token + pynini.closure(pynutil.insert(" ") + punct) + pynini.closure(punct + pynutil.insert(" ")) + + token + + pynini.closure(pynutil.insert(" ") + punct) ) graph = token_plus_punct + pynini.closure(delete_extra_space + token_plus_punct) diff --git a/fun_text_processing/inverse_text_normalization/fr/taggers/whitelist.py b/fun_text_processing/inverse_text_normalization/fr/taggers/whitelist.py index 9db273871..f333e91d8 100644 --- a/fun_text_processing/inverse_text_normalization/fr/taggers/whitelist.py +++ b/fun_text_processing/inverse_text_normalization/fr/taggers/whitelist.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.inverse_text_normalization.fr.graph_utils import GraphFst, convert_space from fun_text_processing.inverse_text_normalization.fr.utils import get_abs_path @@ -17,5 +15,5 @@ class WhiteListFst(GraphFst): super().__init__(name="whitelist", kind="classify") whitelist = pynini.string_file(get_abs_path("data/whitelist.tsv")) - graph = pynutil.insert("name: \"") + convert_space(whitelist) + pynutil.insert("\"") + graph = pynutil.insert('name: "') + convert_space(whitelist) + pynutil.insert('"') self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/fr/taggers/word.py b/fun_text_processing/inverse_text_normalization/fr/taggers/word.py index b17a7be13..a2db671ce 100644 --- a/fun_text_processing/inverse_text_normalization/fr/taggers/word.py +++ b/fun_text_processing/inverse_text_normalization/fr/taggers/word.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.fr.graph_utils import DAMO_NOT_SPACE, GraphFst from pynini.lib import pynutil @@ -12,5 +11,5 @@ class WordFst(GraphFst): def __init__(self): super().__init__(name="word", kind="classify") - word = pynutil.insert("name: \"") + pynini.closure(DAMO_NOT_SPACE, 1) + pynutil.insert("\"") + word = pynutil.insert('name: "') + pynini.closure(DAMO_NOT_SPACE, 1) + pynutil.insert('"') self.fst = word.optimize() diff --git a/fun_text_processing/inverse_text_normalization/fr/utils.py b/fun_text_processing/inverse_text_normalization/fr/utils.py index 9cbd175cb..c6fdb7435 100644 --- a/fun_text_processing/inverse_text_normalization/fr/utils.py +++ b/fun_text_processing/inverse_text_normalization/fr/utils.py @@ -1,4 +1,3 @@ - import os @@ -8,7 +7,7 @@ def get_abs_path(rel_path): Args: rel_path: relative path to this file - + Returns absolute path """ - return os.path.dirname(os.path.abspath(__file__)) + '/' + rel_path + return os.path.dirname(os.path.abspath(__file__)) + "/" + rel_path diff --git a/fun_text_processing/inverse_text_normalization/fr/verbalizers/__init__.py b/fun_text_processing/inverse_text_normalization/fr/verbalizers/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/fr/verbalizers/__init__.py +++ b/fun_text_processing/inverse_text_normalization/fr/verbalizers/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/inverse_text_normalization/fr/verbalizers/cardinal.py b/fun_text_processing/inverse_text_normalization/fr/verbalizers/cardinal.py index c901e3efb..28087c93f 100644 --- a/fun_text_processing/inverse_text_normalization/fr/verbalizers/cardinal.py +++ b/fun_text_processing/inverse_text_normalization/fr/verbalizers/cardinal.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.fr.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.fr.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -15,9 +18,9 @@ class CardinalFst(GraphFst): optional_sign = pynini.closure( pynutil.delete("negative:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + DAMO_NOT_QUOTE - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space, 0, 1, @@ -25,9 +28,9 @@ class CardinalFst(GraphFst): graph = ( pynutil.delete("integer:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) self.numbers = graph graph = optional_sign + graph diff --git a/fun_text_processing/inverse_text_normalization/fr/verbalizers/date.py b/fun_text_processing/inverse_text_normalization/fr/verbalizers/date.py index da2071b9a..85e971c36 100644 --- a/fun_text_processing/inverse_text_normalization/fr/verbalizers/date.py +++ b/fun_text_processing/inverse_text_normalization/fr/verbalizers/date.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.fr.graph_utils import ( DAMO_NOT_QUOTE, @@ -18,29 +17,29 @@ class DateFst(GraphFst): def __init__(self): super().__init__(name="date", kind="verbalize") - convert_primer = pynini.cross('1', '1ᵉʳ') + convert_primer = pynini.cross("1", "1ᵉʳ") day = ( pynutil.delete("day:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + ( pynini.closure(DAMO_NOT_QUOTE, 1) | pynutil.add_weight(convert_primer, -1) ) # first of the month is ordinal - + pynutil.delete("\"") + + pynutil.delete('"') ) month = ( pynutil.delete("month:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) year = ( pynutil.delete("year:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) # day month @@ -51,9 +50,9 @@ class DateFst(GraphFst): pynutil.delete("preserve_order:") + delete_space + pynutil.delete("true") + delete_space | pynutil.delete("field_order:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + DAMO_NOT_QUOTE - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space ) diff --git a/fun_text_processing/inverse_text_normalization/fr/verbalizers/decimal.py b/fun_text_processing/inverse_text_normalization/fr/verbalizers/decimal.py index daa532e90..fadb7dac5 100644 --- a/fun_text_processing/inverse_text_normalization/fr/verbalizers/decimal.py +++ b/fun_text_processing/inverse_text_normalization/fr/verbalizers/decimal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.fr.graph_utils import ( DAMO_DIGIT, @@ -12,8 +11,8 @@ from pynini.lib import pynutil class NumberParser(GraphFst): """ - Finite state transducer for parsing strings of digis. Breaks up digit strings into groups of three for - strings of digits of four or more (inclusive). Groupings are separated by non-breaking space. + Finite state transducer for parsing strings of digis. Breaks up digit strings into groups of three for + strings of digits of four or more (inclusive). Groupings are separated by non-breaking space. e.g. '1000' -> '1 000' e.g. '1000,33333' -> '1 000,333 33 """ @@ -32,11 +31,12 @@ class DecimalFst(GraphFst): super().__init__(name="decimal", kind="verbalize") # Need parser to group digits by threes - exactly_three_digits = DAMO_DIGIT ** 3 + exactly_three_digits = DAMO_DIGIT**3 at_most_three_digits = pynini.closure(DAMO_DIGIT, 1, 3) space_every_three_integer = ( - at_most_three_digits + (pynutil.insert(DAMO_NON_BREAKING_SPACE) + exactly_three_digits).closure() + at_most_three_digits + + (pynutil.insert(DAMO_NON_BREAKING_SPACE) + exactly_three_digits).closure() ) space_every_three_decimal = ( pynini.accep(",") @@ -46,13 +46,13 @@ class DecimalFst(GraphFst): group_by_threes = space_every_three_integer | space_every_three_decimal self.group_by_threes = group_by_threes - optional_sign = pynini.closure(pynini.cross("negative: \"true\"", "-") + delete_space, 0, 1) + optional_sign = pynini.closure(pynini.cross('negative: "true"', "-") + delete_space, 0, 1) integer = ( pynutil.delete("integer_part:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) integer = integer @ group_by_threes optional_integer = pynini.closure(integer + delete_space, 0, 1) @@ -60,18 +60,18 @@ class DecimalFst(GraphFst): pynutil.insert(",") + pynutil.delete("fractional_part:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) fractional = fractional @ group_by_threes optional_fractional = pynini.closure(fractional + delete_space, 0, 1) quantity = ( pynutil.delete("quantity:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_quantity = pynini.closure(pynutil.insert(" ") + quantity + delete_space, 0, 1) graph = (optional_integer + optional_fractional + optional_quantity).optimize() diff --git a/fun_text_processing/inverse_text_normalization/fr/verbalizers/electronic.py b/fun_text_processing/inverse_text_normalization/fr/verbalizers/electronic.py index 06139e0c8..866af3fd8 100644 --- a/fun_text_processing/inverse_text_normalization/fr/verbalizers/electronic.py +++ b/fun_text_processing/inverse_text_normalization/fr/verbalizers/electronic.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.fr.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.fr.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -15,16 +18,16 @@ class ElectronicFst(GraphFst): user_name = ( pynutil.delete("username:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) domain = ( pynutil.delete("domain:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) graph = user_name + delete_space + pynutil.insert("@") + domain diff --git a/fun_text_processing/inverse_text_normalization/fr/verbalizers/fraction.py b/fun_text_processing/inverse_text_normalization/fr/verbalizers/fraction.py index cecd3e2ee..d73e41cb2 100644 --- a/fun_text_processing/inverse_text_normalization/fr/verbalizers/fraction.py +++ b/fun_text_processing/inverse_text_normalization/fr/verbalizers/fraction.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.fr.graph_utils import ( DAMO_NOT_QUOTE, @@ -18,23 +17,27 @@ class FractionFst(GraphFst): def __init__(self): super().__init__(name="fraction", kind="verbalize") - optional_sign = pynini.closure(pynini.cross("negative: \"true\"", "-") + delete_space, 0, 1) + optional_sign = pynini.closure(pynini.cross('negative: "true"', "-") + delete_space, 0, 1) integer = ( - pynutil.delete("integer_part: \"") + pynutil.delete('integer_part: "') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') + insert_space ) - numerator = pynutil.delete("numerator: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") - - denominator = ( - pynutil.insert('/') - + pynutil.delete("denominator: \"") - + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + numerator = ( + pynutil.delete('numerator: "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') ) - graph = (pynini.closure(integer + delete_space, 0, 1) + numerator + delete_space + denominator).optimize() + denominator = ( + pynutil.insert("/") + + pynutil.delete('denominator: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + ) + + graph = ( + pynini.closure(integer + delete_space, 0, 1) + numerator + delete_space + denominator + ).optimize() self.numbers = graph delete_tokens = self.delete_tokens(optional_sign + graph) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/inverse_text_normalization/fr/verbalizers/measure.py b/fun_text_processing/inverse_text_normalization/fr/verbalizers/measure.py index 6f0eb0cea..f6deac89a 100644 --- a/fun_text_processing/inverse_text_normalization/fr/verbalizers/measure.py +++ b/fun_text_processing/inverse_text_normalization/fr/verbalizers/measure.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.fr.graph_utils import DAMO_CHAR, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.fr.graph_utils import ( + DAMO_CHAR, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -17,13 +20,13 @@ class MeasureFst(GraphFst): def __init__(self, decimal: GraphFst, cardinal: GraphFst, fraction: GraphFst): super().__init__(name="measure", kind="verbalize") - optional_sign = pynini.closure(pynini.cross("negative: \"true\"", "-"), 0, 1) + optional_sign = pynini.closure(pynini.cross('negative: "true"', "-"), 0, 1) unit = ( pynutil.delete("units:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space ) graph_decimal = ( @@ -40,7 +43,8 @@ class MeasureFst(GraphFst): + delete_space + optional_sign + delete_space - + cardinal.numbers @ decimal.group_by_threes # measurements most obey three by three spacing + + cardinal.numbers + @ decimal.group_by_threes # measurements most obey three by three spacing + delete_space + pynutil.delete("}") ) @@ -54,6 +58,11 @@ class MeasureFst(GraphFst): + pynutil.delete("}") ) - graph = (graph_cardinal | graph_decimal | graph_fraction) + delete_space + pynutil.insert(" ") + unit + graph = ( + (graph_cardinal | graph_decimal | graph_fraction) + + delete_space + + pynutil.insert(" ") + + unit + ) delete_tokens = self.delete_tokens(graph) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/inverse_text_normalization/fr/verbalizers/money.py b/fun_text_processing/inverse_text_normalization/fr/verbalizers/money.py index fdeb35ec1..8c8d0a0b1 100644 --- a/fun_text_processing/inverse_text_normalization/fr/verbalizers/money.py +++ b/fun_text_processing/inverse_text_normalization/fr/verbalizers/money.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.fr.graph_utils import ( DAMO_NOT_QUOTE, @@ -23,9 +22,9 @@ class MoneyFst(GraphFst): unit = ( pynutil.delete("currency:") + delete_extra_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) graph = decimal.numbers + delete_space + unit delete_tokens = self.delete_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/fr/verbalizers/ordinal.py b/fun_text_processing/inverse_text_normalization/fr/verbalizers/ordinal.py index 92d715b61..aa4623743 100644 --- a/fun_text_processing/inverse_text_normalization/fr/verbalizers/ordinal.py +++ b/fun_text_processing/inverse_text_normalization/fr/verbalizers/ordinal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.fr.graph_utils import ( DAMO_DIGIT, @@ -26,9 +25,9 @@ class OrdinalFst(GraphFst): graph_integer = ( pynutil.delete("integer:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) replace_suffix = pynini.union( @@ -37,22 +36,26 @@ class OrdinalFst(GraphFst): pynini.cross("r", "ʳ"), pynini.cross("s", "ˢ"), ) - replace_suffix = pynutil.delete(" morphosyntactic_features: \"") + replace_suffix.plus + replace_suffix = pynutil.delete(' morphosyntactic_features: "') + replace_suffix.plus graph_arabic = graph_integer + replace_suffix.plus # For roman. - graph_roman_digits = pynini.string_file(get_abs_path("data/roman/digits_large.tsv")).invert() + graph_roman_digits = pynini.string_file( + get_abs_path("data/roman/digits_large.tsv") + ).invert() graph_roman_ties = pynini.string_file(get_abs_path("data/roman/ties_large.tsv")).invert() - graph_roman_hundreds = pynini.string_file(get_abs_path("data/roman/hundreds_large.tsv")).invert() + graph_roman_hundreds = pynini.string_file( + get_abs_path("data/roman/hundreds_large.tsv") + ).invert() graph_roman_zero_digit = pynutil.delete("0") - graph_roman_hundreds = DAMO_DIGIT ** 3 @ ( + graph_roman_hundreds = DAMO_DIGIT**3 @ ( graph_roman_hundreds + pynini.union(graph_roman_ties, graph_roman_zero_digit) + pynini.union(graph_roman_digits, graph_roman_zero_digit) ) - graph_roman_ties = DAMO_DIGIT ** 2 @ ( + graph_roman_ties = DAMO_DIGIT**2 @ ( graph_roman_ties + pynini.union(graph_roman_digits, graph_roman_zero_digit) ) graph_roman_digits = DAMO_DIGIT @ graph_roman_digits @@ -62,7 +65,7 @@ class OrdinalFst(GraphFst): graph_roman = (graph_integer @ graph_roman_integers) + replace_suffix graph_roman += pynini.cross("/", " ") + "siècle" - graph = (graph_roman | graph_arabic) + pynutil.delete("\"") + graph = (graph_roman | graph_arabic) + pynutil.delete('"') delete_tokens = self.delete_tokens(graph) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/inverse_text_normalization/fr/verbalizers/telephone.py b/fun_text_processing/inverse_text_normalization/fr/verbalizers/telephone.py index 2985f5083..c35d7e5d8 100644 --- a/fun_text_processing/inverse_text_normalization/fr/verbalizers/telephone.py +++ b/fun_text_processing/inverse_text_normalization/fr/verbalizers/telephone.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.fr.graph_utils import DAMO_NOT_QUOTE, GraphFst from pynini.lib import pynutil @@ -8,12 +7,16 @@ class TelephoneFst(GraphFst): """ Finite state transducer for verbalizing telephone, e.g. telephone { number_part: "02 33 43 53 22" } - -> 02 33 43 53 22 + -> 02 33 43 53 22 """ def __init__(self): super().__init__(name="telephone", kind="verbalize") - number_part = pynutil.delete("number_part: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + number_part = ( + pynutil.delete('number_part: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + ) delete_tokens = self.delete_tokens(number_part) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/inverse_text_normalization/fr/verbalizers/time.py b/fun_text_processing/inverse_text_normalization/fr/verbalizers/time.py index c75c69c29..0de583fbf 100644 --- a/fun_text_processing/inverse_text_normalization/fr/verbalizers/time.py +++ b/fun_text_processing/inverse_text_normalization/fr/verbalizers/time.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.fr.graph_utils import ( DAMO_DIGIT, @@ -15,7 +14,7 @@ class TimeFst(GraphFst): Finite state transducer for verbalizing time, e.g. time { hours: "8" minutes: "30" suffix: "du matin"} -> 8 h 30 time { hours: "8" minutes: "30" } -> 8 h 30 - time { hours: "8" minutes: "30" suffix: "du soir"} -> 20 h 30 + time { hours: "8" minutes: "30" suffix: "du soir"} -> 20 h 30 """ def __init__(self): @@ -23,25 +22,32 @@ class TimeFst(GraphFst): hour_to_night = pynini.string_file(get_abs_path("data/time/hour_to_night.tsv")) - day_suffixes = pynutil.delete("suffix: \"am\"") - night_suffixes = pynutil.delete("suffix: \"pm\"") + day_suffixes = pynutil.delete('suffix: "am"') + night_suffixes = pynutil.delete('suffix: "pm"') hour = ( pynutil.delete("hours:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_DIGIT, 1, 2) - + pynutil.delete("\"") + + pynutil.delete('"') ) minute = ( pynutil.delete("minutes:") + delete_extra_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_DIGIT, 1, 2) - + pynutil.delete("\"") + + pynutil.delete('"') ) - graph = hour + delete_extra_space + pynutil.insert("h") + minute.ques + delete_space + day_suffixes.ques + graph = ( + hour + + delete_extra_space + + pynutil.insert("h") + + minute.ques + + delete_space + + day_suffixes.ques + ) graph |= ( hour @ hour_to_night diff --git a/fun_text_processing/inverse_text_normalization/fr/verbalizers/verbalize.py b/fun_text_processing/inverse_text_normalization/fr/verbalizers/verbalize.py index 40271aaa0..55d5eb2ac 100644 --- a/fun_text_processing/inverse_text_normalization/fr/verbalizers/verbalize.py +++ b/fun_text_processing/inverse_text_normalization/fr/verbalizers/verbalize.py @@ -1,6 +1,3 @@ - - - from fun_text_processing.inverse_text_normalization.fr.graph_utils import GraphFst from fun_text_processing.inverse_text_normalization.fr.verbalizers.cardinal import CardinalFst from fun_text_processing.inverse_text_normalization.fr.verbalizers.date import DateFst @@ -18,7 +15,7 @@ from fun_text_processing.inverse_text_normalization.fr.verbalizers.whitelist imp class VerbalizeFst(GraphFst): """ Composes other verbalizer grammars. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. """ diff --git a/fun_text_processing/inverse_text_normalization/fr/verbalizers/verbalize_final.py b/fun_text_processing/inverse_text_normalization/fr/verbalizers/verbalize_final.py index f284aac49..c384df566 100644 --- a/fun_text_processing/inverse_text_normalization/fr/verbalizers/verbalize_final.py +++ b/fun_text_processing/inverse_text_normalization/fr/verbalizers/verbalize_final.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.fr.graph_utils import GraphFst, delete_extra_space, delete_space +from fun_text_processing.inverse_text_normalization.fr.graph_utils import ( + GraphFst, + delete_extra_space, + delete_space, +) from fun_text_processing.inverse_text_normalization.fr.verbalizers.verbalize import VerbalizeFst from fun_text_processing.inverse_text_normalization.fr.verbalizers.word import WordFst from pynini.lib import pynutil @@ -8,7 +11,7 @@ from pynini.lib import pynutil class VerbalizeFinalFst(GraphFst): """ - Finite state transducer that verbalizes an entire sentence, e.g. + Finite state transducer that verbalizes an entire sentence, e.g. tokens { name: "its" } tokens { time { hours: "12" minutes: "30" } } tokens { name: "now" } -> its 12:30 now """ diff --git a/fun_text_processing/inverse_text_normalization/fr/verbalizers/whitelist.py b/fun_text_processing/inverse_text_normalization/fr/verbalizers/whitelist.py index c94eebfea..889dc8878 100644 --- a/fun_text_processing/inverse_text_normalization/fr/verbalizers/whitelist.py +++ b/fun_text_processing/inverse_text_normalization/fr/verbalizers/whitelist.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.inverse_text_normalization.fr.graph_utils import ( DAMO_CHAR, @@ -21,9 +19,9 @@ class WhiteListFst(GraphFst): graph = ( pynutil.delete("name:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) - graph = graph @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + graph = graph @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/fr/verbalizers/word.py b/fun_text_processing/inverse_text_normalization/fr/verbalizers/word.py index 1339dac41..5a9fc1c90 100644 --- a/fun_text_processing/inverse_text_normalization/fr/verbalizers/word.py +++ b/fun_text_processing/inverse_text_normalization/fr/verbalizers/word.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.fr.graph_utils import ( DAMO_CHAR, @@ -18,7 +17,13 @@ class WordFst(GraphFst): def __init__(self): super().__init__(name="word", kind="verbalize") chars = pynini.closure(DAMO_CHAR - " ", 1) - char = pynutil.delete("name:") + delete_space + pynutil.delete("\"") + chars + pynutil.delete("\"") - graph = char @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + char = ( + pynutil.delete("name:") + + delete_space + + pynutil.delete('"') + + chars + + pynutil.delete('"') + ) + graph = char @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/id/__init__.py b/fun_text_processing/inverse_text_normalization/id/__init__.py index dd073ae25..096e8f133 100644 --- a/fun_text_processing/inverse_text_normalization/id/__init__.py +++ b/fun_text_processing/inverse_text_normalization/id/__init__.py @@ -1,3 +1,7 @@ -from fun_text_processing.inverse_text_normalization.id.taggers.tokenize_and_classify import ClassifyFst +from fun_text_processing.inverse_text_normalization.id.taggers.tokenize_and_classify import ( + ClassifyFst, +) from fun_text_processing.inverse_text_normalization.id.verbalizers.verbalize import VerbalizeFst -from fun_text_processing.inverse_text_normalization.id.verbalizers.verbalize_final import VerbalizeFinalFst +from fun_text_processing.inverse_text_normalization.id.verbalizers.verbalize_final import ( + VerbalizeFinalFst, +) diff --git a/fun_text_processing/inverse_text_normalization/id/clean_eval_data.py b/fun_text_processing/inverse_text_normalization/id/clean_eval_data.py index 97abe13a5..fc0a31576 100644 --- a/fun_text_processing/inverse_text_normalization/id/clean_eval_data.py +++ b/fun_text_processing/inverse_text_normalization/id/clean_eval_data.py @@ -53,7 +53,7 @@ class Filter: Args: processes given instance with process function - + Returns: processed instance if instance belongs to expected class type or original instance """ if instance.token_type != self.class_type: @@ -71,7 +71,9 @@ def process_cardinal_1(instance: Instance) -> Instance: normalized = instance.normalized un_normalized = re.sub(r"[^0-9]", "", un_normalized) normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_ordinal_1(instance: Instance) -> bool: @@ -84,7 +86,9 @@ def process_ordinal_1(instance: Instance) -> Instance: normalized = instance.normalized un_normalized = re.sub(r"[,\s]", "", un_normalized) normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_decimal_1(instance: Instance) -> bool: @@ -97,7 +101,9 @@ def process_decimal_1(instance: Instance) -> Instance: un_normalized = re.sub(r",", "", un_normalized) normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_measure_1(instance: Instance) -> bool: @@ -114,7 +120,9 @@ def process_measure_1(instance: Instance) -> Instance: normalized = re.sub(r"[^a-z\s]", "", normalized) normalized = re.sub(r"per ([a-z\s]*)s$", r"per \1", normalized) normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_money_1(instance: Instance) -> bool: @@ -131,7 +139,9 @@ def process_money_1(instance: Instance) -> Instance: un_normalized = re.sub(r"(\d)m\s*$", r"\1 juta", un_normalized) un_normalized = re.sub(r"(\d)bn?\s*$", r"\1 milyar", un_normalized) normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_time_1(instance: Instance) -> bool: @@ -146,7 +156,9 @@ def process_time_1(instance: Instance) -> Instance: un_normalized = re.sub(r"(\d)\s?p\s?m\s?", r"\1 p.m.", un_normalized) normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_plain_1(instance: Instance) -> bool: @@ -157,7 +169,9 @@ def filter_plain_1(instance: Instance) -> bool: def process_plain_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_punct_1(instance: Instance) -> bool: @@ -168,7 +182,9 @@ def filter_punct_1(instance: Instance) -> bool: def process_punct_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_date_1(instance: Instance) -> bool: @@ -181,7 +197,9 @@ def process_date_1(instance: Instance) -> Instance: un_normalized = re.sub(r",", "", un_normalized) normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_letters_1(instance: Instance) -> bool: @@ -193,7 +211,9 @@ def process_letters_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_verbatim_1(instance: Instance) -> bool: @@ -204,7 +224,9 @@ def filter_verbatim_1(instance: Instance) -> bool: def process_verbatim_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_digit_1(instance: Instance) -> bool: @@ -216,7 +238,9 @@ def process_digit_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_telephone_1(instance: Instance) -> bool: @@ -228,7 +252,9 @@ def process_telephone_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_electronic_1(instance: Instance) -> bool: @@ -240,7 +266,9 @@ def process_electronic_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_fraction_1(instance: Instance) -> bool: @@ -252,7 +280,9 @@ def process_fraction_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_address_1(instance: Instance) -> bool: @@ -264,27 +294,51 @@ def process_address_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) filters = [] -filters.append(Filter(class_type="CARDINAL", process_func=process_cardinal_1, filter_func=filter_cardinal_1)) -filters.append(Filter(class_type="ORDINAL", process_func=process_ordinal_1, filter_func=filter_ordinal_1)) -filters.append(Filter(class_type="DECIMAL", process_func=process_decimal_1, filter_func=filter_decimal_1)) -filters.append(Filter(class_type="MEASURE", process_func=process_measure_1, filter_func=filter_measure_1)) +filters.append( + Filter(class_type="CARDINAL", process_func=process_cardinal_1, filter_func=filter_cardinal_1) +) +filters.append( + Filter(class_type="ORDINAL", process_func=process_ordinal_1, filter_func=filter_ordinal_1) +) +filters.append( + Filter(class_type="DECIMAL", process_func=process_decimal_1, filter_func=filter_decimal_1) +) +filters.append( + Filter(class_type="MEASURE", process_func=process_measure_1, filter_func=filter_measure_1) +) filters.append(Filter(class_type="MONEY", process_func=process_money_1, filter_func=filter_money_1)) filters.append(Filter(class_type="TIME", process_func=process_time_1, filter_func=filter_time_1)) filters.append(Filter(class_type="DATE", process_func=process_date_1, filter_func=filter_date_1)) filters.append(Filter(class_type="PLAIN", process_func=process_plain_1, filter_func=filter_plain_1)) filters.append(Filter(class_type="PUNCT", process_func=process_punct_1, filter_func=filter_punct_1)) -filters.append(Filter(class_type="LETTERS", process_func=process_letters_1, filter_func=filter_letters_1)) -filters.append(Filter(class_type="VERBATIM", process_func=process_verbatim_1, filter_func=filter_verbatim_1)) +filters.append( + Filter(class_type="LETTERS", process_func=process_letters_1, filter_func=filter_letters_1) +) +filters.append( + Filter(class_type="VERBATIM", process_func=process_verbatim_1, filter_func=filter_verbatim_1) +) filters.append(Filter(class_type="DIGIT", process_func=process_digit_1, filter_func=filter_digit_1)) -filters.append(Filter(class_type="TELEPHONE", process_func=process_telephone_1, filter_func=filter_telephone_1)) -filters.append(Filter(class_type="ELECTRONIC", process_func=process_electronic_1, filter_func=filter_electronic_1)) -filters.append(Filter(class_type="FRACTION", process_func=process_fraction_1, filter_func=filter_fraction_1)) -filters.append(Filter(class_type="ADDRESS", process_func=process_address_1, filter_func=filter_address_1)) +filters.append( + Filter(class_type="TELEPHONE", process_func=process_telephone_1, filter_func=filter_telephone_1) +) +filters.append( + Filter( + class_type="ELECTRONIC", process_func=process_electronic_1, filter_func=filter_electronic_1 + ) +) +filters.append( + Filter(class_type="FRACTION", process_func=process_fraction_1, filter_func=filter_fraction_1) +) +filters.append( + Filter(class_type="ADDRESS", process_func=process_address_1, filter_func=filter_address_1) +) filters.append(Filter(class_type=EOS_TYPE, process_func=lambda x: x, filter_func=lambda x: True)) @@ -313,8 +367,10 @@ def filter_loaded_data(data: List[Instance], verbose: bool = False) -> List[Inst def parse_args(): parser = ArgumentParser() - parser.add_argument("--input", help="input file path", type=str, default='./id_with_types/output-00001-of-00100') - parser.add_argument("--verbose", help="print filtered instances", action='store_true') + parser.add_argument( + "--input", help="input file path", type=str, default="./id_with_types/output-00001-of-00100" + ) + parser.add_argument("--verbose", help="print filtered instances", action="store_true") return parser.parse_args() diff --git a/fun_text_processing/inverse_text_normalization/id/graph_utils.py b/fun_text_processing/inverse_text_normalization/id/graph_utils.py index ccc9fa90a..81a4da93d 100644 --- a/fun_text_processing/inverse_text_normalization/id/graph_utils.py +++ b/fun_text_processing/inverse_text_normalization/id/graph_utils.py @@ -18,9 +18,9 @@ DAMO_UPPER = pynini.union(*string.ascii_uppercase).optimize() DAMO_ALPHA = pynini.union(DAMO_LOWER, DAMO_UPPER).optimize() DAMO_ALNUM = pynini.union(DAMO_DIGIT, DAMO_ALPHA).optimize() DAMO_HEX = pynini.union(*string.hexdigits).optimize() -DAMO_NON_BREAKING_SPACE = u"\u00A0" +DAMO_NON_BREAKING_SPACE = "\u00A0" DAMO_SPACE = " " -DAMO_WHITE_SPACE = pynini.union(" ", "\t", "\n", "\r", u"\u00A0").optimize() +DAMO_WHITE_SPACE = pynini.union(" ", "\t", "\n", "\r", "\u00A0").optimize() DAMO_NOT_SPACE = pynini.difference(DAMO_CHAR, DAMO_WHITE_SPACE).optimize() DAMO_NOT_QUOTE = pynini.difference(DAMO_CHAR, r'"').optimize() @@ -35,31 +35,55 @@ insert_space = pynutil.insert(" ") delete_extra_space = pynini.cross(pynini.closure(DAMO_WHITE_SPACE, 1), " ") delete_preserve_order = pynini.closure( pynutil.delete(" preserve_order: true") - | (pynutil.delete(" field_order: \"") + DAMO_NOT_QUOTE + pynutil.delete("\"")) + | (pynutil.delete(' field_order: "') + DAMO_NOT_QUOTE + pynutil.delete('"')) ) suppletive = pynini.string_file(get_abs_path("data/suppletive.tsv")) # _v = pynini.union("a", "e", "i", "o", "u") _c = pynini.union( - "b", "c", "d", "f", "g", "h", "j", "k", "l", "m", "n", "p", "q", "r", "s", "t", "v", "w", "x", "y", "z" + "b", + "c", + "d", + "f", + "g", + "h", + "j", + "k", + "l", + "m", + "n", + "p", + "q", + "r", + "s", + "t", + "v", + "w", + "x", + "y", + "z", ) _ies = DAMO_SIGMA + _c + pynini.cross("y", "ies") _es = DAMO_SIGMA + pynini.union("s", "sh", "ch", "x", "z") + pynutil.insert("es") _s = DAMO_SIGMA + pynutil.insert("s") graph_plural = plurals._priority_union( - suppletive, plurals._priority_union(_ies, plurals._priority_union(_es, _s, DAMO_SIGMA), DAMO_SIGMA), DAMO_SIGMA + suppletive, + plurals._priority_union(_ies, plurals._priority_union(_es, _s, DAMO_SIGMA), DAMO_SIGMA), + DAMO_SIGMA, ).optimize() SINGULAR_TO_PLURAL = graph_plural PLURAL_TO_SINGULAR = pynini.invert(graph_plural) -TO_LOWER = pynini.union(*[pynini.cross(x, y) for x, y in zip(string.ascii_uppercase, string.ascii_lowercase)]) +TO_LOWER = pynini.union( + *[pynini.cross(x, y) for x, y in zip(string.ascii_uppercase, string.ascii_lowercase)] +) TO_UPPER = pynini.invert(TO_LOWER) MIN_NEG_WEIGHT = -0.0001 MIN_POS_WEIGHT = 0.0001 -def generator_main(file_name: str, graphs: Dict[str, 'pynini.FstLike']): +def generator_main(file_name: str, graphs: Dict[str, "pynini.FstLike"]): """ Exports graph as OpenFst finite state archive (FAR) file with given file name and rule name. @@ -71,7 +95,7 @@ def generator_main(file_name: str, graphs: Dict[str, 'pynini.FstLike']): for rule, graph in graphs.items(): exporter[rule] = graph.optimize() exporter.close() - print(f'Created {file_name}') + print(f"Created {file_name}") def get_plurals(fst): @@ -98,7 +122,7 @@ def get_singulars(fst): return PLURAL_TO_SINGULAR @ fst -def convert_space(fst) -> 'pynini.FstLike': +def convert_space(fst) -> "pynini.FstLike": """ Converts space to nonbreaking space. Used only in tagger grammars for transducing token values within quotes, e.g. name: "hello kitty" @@ -109,7 +133,9 @@ def convert_space(fst) -> 'pynini.FstLike': Returns output fst where breaking spaces are converted to non breaking spaces """ - return fst @ pynini.cdrewrite(pynini.cross(DAMO_SPACE, DAMO_NON_BREAKING_SPACE), "", "", DAMO_SIGMA) + return fst @ pynini.cdrewrite( + pynini.cross(DAMO_SPACE, DAMO_NON_BREAKING_SPACE), "", "", DAMO_SIGMA + ) class GraphFst: @@ -129,9 +155,11 @@ class GraphFst: self._fst = None self.deterministic = deterministic - self.far_path = Path(os.path.dirname(__file__) + '/grammars/' + kind + '/' + name + '.far') + self.far_path = Path(os.path.dirname(__file__) + "/grammars/" + kind + "/" + name + ".far") if self.far_exist(): - self._fst = Far(self.far_path, mode="r", arc_type="standard", far_type="default").get_fst() + self._fst = Far( + self.far_path, mode="r", arc_type="standard", far_type="default" + ).get_fst() def far_exist(self) -> bool: """ @@ -140,14 +168,14 @@ class GraphFst: return self.far_path.exists() @property - def fst(self) -> 'pynini.FstLike': + def fst(self) -> "pynini.FstLike": return self._fst @fst.setter def fst(self, fst): self._fst = fst - def add_tokens(self, fst) -> 'pynini.FstLike': + def add_tokens(self, fst) -> "pynini.FstLike": """ Wraps class name around to given fst @@ -159,7 +187,7 @@ class GraphFst: """ return pynutil.insert(f"{self.name} {{ ") + fst + pynutil.insert(" }") - def delete_tokens(self, fst) -> 'pynini.FstLike': + def delete_tokens(self, fst) -> "pynini.FstLike": """ Deletes class name wrap around output of given fst @@ -178,4 +206,4 @@ class GraphFst: + delete_space + pynutil.delete("}") ) - return res @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + return res @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) diff --git a/fun_text_processing/inverse_text_normalization/id/taggers/cardinal.py b/fun_text_processing/inverse_text_normalization/id/taggers/cardinal.py index d2f1a770e..f9e826ec2 100644 --- a/fun_text_processing/inverse_text_normalization/id/taggers/cardinal.py +++ b/fun_text_processing/inverse_text_normalization/id/taggers/cardinal.py @@ -15,7 +15,7 @@ class CardinalFst(GraphFst): """ Finite state transducer for classifying cardinals e.g. minus twenty three -> cardinal { integer: "23" negative: "-" } } - Numbers below thirteen are not converted. + Numbers below thirteen are not converted. """ def __init__(self, deterministic: bool = True): @@ -26,25 +26,28 @@ class CardinalFst(GraphFst): graph_teen = pynini.string_file(get_abs_path("data/numbers/teen.tsv")) graph_hundreds = pynini.string_file(get_abs_path("data/numbers/hundreds.tsv")) graph_thousand = pynini.string_file(get_abs_path("data/numbers/thousand.tsv")) - + graph_hundred = pynini.cross("ratus", "") | pynini.cross("seratus", "") - graph_hundred_component = pynini.union(graph_digit + delete_space + graph_hundred, pynutil.insert("0")) + graph_hundred_component = pynini.union( + graph_digit + delete_space + graph_hundred, pynutil.insert("0") + ) graph_hundred_component += delete_space graph_hundred_component += pynini.union( graph_teen | pynutil.insert("00"), (graph_ties | pynutil.insert("0")) + delete_space + (graph_digit | pynutil.insert("0")), ) - graph_one_hundred_component = pynini.union(pynini.cross("ratus", "1") | pynini.cross("seratus", "1")) + graph_one_hundred_component = pynini.union( + pynini.cross("ratus", "1") | pynini.cross("seratus", "1") + ) graph_one_hundred_component += delete_space graph_one_hundred_component += pynini.union( - graph_teen | pynutil.insert("00"), - (graph_ties | pynutil.insert("0")) + delete_space + ( - graph_digit | pynutil.insert("0")), + graph_teen | pynutil.insert("00"), + (graph_ties | pynutil.insert("0")) + delete_space + (graph_digit | pynutil.insert("0")), ) graph_hundred_component = graph_hundred_component | graph_one_hundred_component - + graph_hundred_component_at_least_one_none_zero_digit = graph_hundred_component @ ( pynini.closure(DAMO_DIGIT) + (DAMO_DIGIT - "0") + pynini.closure(DAMO_DIGIT) ) @@ -52,36 +55,57 @@ class CardinalFst(GraphFst): graph_hundred_component_at_least_one_none_zero_digit ) graph_thousand = pynini.cross("ribu", "") | pynini.cross("seribu", "") - graph_one_thousand_component = pynini.union(pynini.cross("ribu", "1") | pynini.cross("seribu", "1")) - + graph_one_thousand_component = pynini.union( + pynini.cross("ribu", "1") | pynini.cross("seribu", "1") + ) + graph_thousands = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + (pynutil.delete("ribu") | pynutil.delete("seribu")), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + (pynutil.delete("ribu") | pynutil.delete("seribu")), pynutil.insert("000", weight=0.1), ) graph_thousands = graph_thousands | (pynutil.insert("00") + graph_one_thousand_component) graph_million = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + (pynutil.delete("juta") | pynutil.delete("sejuta")), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + (pynutil.delete("juta") | pynutil.delete("sejuta")), pynutil.insert("000", weight=0.1), ) graph_billion = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + (pynutil.delete("miliar") | pynutil.delete("semiliar") | pynutil.delete("milyar") | pynutil.delete("semilyar")), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + ( + pynutil.delete("miliar") + | pynutil.delete("semiliar") + | pynutil.delete("milyar") + | pynutil.delete("semilyar") + ), pynutil.insert("000", weight=0.1), ) graph_trillion = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + (pynutil.delete("triliun") | pynutil.delete("setriliun")), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + (pynutil.delete("triliun") | pynutil.delete("setriliun")), pynutil.insert("000", weight=0.1), ) graph_quadrillion = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("milion lipat empat"), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + pynutil.delete("milion lipat empat"), pynutil.insert("000", weight=0.1), ) graph_quintillion = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("triliun"), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + pynutil.delete("triliun"), pynutil.insert("000", weight=0.1), ) graph_sextillion = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("sextillion"), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + pynutil.delete("sextillion"), pynutil.insert("000", weight=0.1), ) @@ -106,17 +130,19 @@ class CardinalFst(GraphFst): ) graph = graph @ pynini.union( - pynutil.delete(pynini.closure("0")) + pynini.difference(DAMO_DIGIT, "0") + pynini.closure(DAMO_DIGIT), "0" + pynutil.delete(pynini.closure("0")) + + pynini.difference(DAMO_DIGIT, "0") + + pynini.closure(DAMO_DIGIT), + "0", ) labels_exception = ["nol"] graph_exception = pynini.union(*labels_exception) graph = ( - pynini.cdrewrite(pynutil.delete("dan"), DAMO_SPACE, DAMO_SPACE, - DAMO_SIGMA) - @ (DAMO_ALPHA + DAMO_SIGMA) - @ graph + pynini.cdrewrite(pynutil.delete("dan"), DAMO_SPACE, DAMO_SPACE, DAMO_SIGMA) + @ (DAMO_ALPHA + DAMO_SIGMA) + @ graph ) self.graph_no_exception = graph @@ -124,12 +150,12 @@ class CardinalFst(GraphFst): self.graph = (pynini.project(graph, "input") - graph_exception.arcsort()) @ graph optional_minus_graph = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("kurang", "\"-\"") + DAMO_SPACE, 0, 1 + pynutil.insert("negative: ") + pynini.cross("kurang", '"-"') + DAMO_SPACE, 0, 1 ) - final_graph = optional_minus_graph + pynutil.insert("integer: \"") + self.graph + pynutil.insert("\"") + final_graph = ( + optional_minus_graph + pynutil.insert('integer: "') + self.graph + pynutil.insert('"') + ) final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() - - diff --git a/fun_text_processing/inverse_text_normalization/id/taggers/date.py b/fun_text_processing/inverse_text_normalization/id/taggers/date.py index 7ae0a14a1..82147fe43 100644 --- a/fun_text_processing/inverse_text_normalization/id/taggers/date.py +++ b/fun_text_processing/inverse_text_normalization/id/taggers/date.py @@ -62,7 +62,9 @@ def _get_year_graph(): def _get_thousands_graph(): graph_ties = _get_ties_graph() - graph_hundred_component = (graph_digit + delete_space + pynutil.delete("ratus")) | pynutil.insert("0") + graph_hundred_component = ( + graph_digit + delete_space + pynutil.delete("ratus") + ) | pynutil.insert("0") graph = ( graph_digit + delete_space @@ -89,7 +91,7 @@ def _get_year_graph(): class DateFst(GraphFst): """ - Finite state transducer for classifying date, + Finite state transducer for classifying date, e.g. january fifth twenty twelve -> date { month: "january" day: "5" year: "2012" preserve_order: true } e.g. the fifth of january twenty twelve -> date { day: "5" month: "january" year: "2012" preserve_order: true } e.g. twenty twenty -> date { year: "2012" preserve_order: true } @@ -107,18 +109,26 @@ class DateFst(GraphFst): year_graph = pynutil.add_weight(year_graph, YEAR_WEIGHT) month_graph = _get_month_graph() - month_graph = pynutil.insert("month: \"") + month_graph + pynutil.insert("\"") + month_graph = pynutil.insert('month: "') + month_graph + pynutil.insert('"') - day_graph = pynutil.insert("day: \"") + pynutil.add_weight(ordinal_graph, -0.7) + pynutil.insert("\"") + day_graph = ( + pynutil.insert('day: "') + pynutil.add_weight(ordinal_graph, -0.7) + pynutil.insert('"') + ) graph_year = ( delete_extra_space - + pynutil.insert("year: \"") + + pynutil.insert('year: "') + pynutil.add_weight(year_graph, -YEAR_WEIGHT) - + pynutil.insert("\"") + + pynutil.insert('"') + ) + optional_graph_year = pynini.closure( + graph_year, + 0, + 1, ) - optional_graph_year = pynini.closure(graph_year, 0, 1,) graph_mdy = month_graph + ( - (delete_extra_space + day_graph) | graph_year | (delete_extra_space + day_graph + graph_year) + (delete_extra_space + day_graph) + | graph_year + | (delete_extra_space + day_graph + graph_year) ) graph_dmy = ( pynutil.delete("the") @@ -130,7 +140,9 @@ class DateFst(GraphFst): + month_graph + optional_graph_year ) - graph_year = pynutil.insert("year: \"") + (year_graph | _get_range_graph()) + pynutil.insert("\"") + graph_year = ( + pynutil.insert('year: "') + (year_graph | _get_range_graph()) + pynutil.insert('"') + ) final_graph = graph_mdy | graph_dmy | graph_year final_graph += pynutil.insert(" preserve_order: true") diff --git a/fun_text_processing/inverse_text_normalization/id/taggers/decimal.py b/fun_text_processing/inverse_text_normalization/id/taggers/decimal.py index c51774855..98b42059d 100644 --- a/fun_text_processing/inverse_text_normalization/id/taggers/decimal.py +++ b/fun_text_processing/inverse_text_normalization/id/taggers/decimal.py @@ -9,30 +9,41 @@ from fun_text_processing.inverse_text_normalization.id.graph_utils import ( from pynini.lib import pynutil import pdb -def get_quantity(decimal: 'pynini.FstLike', cardinal_up_to_hundred: 'pynini.FstLike') -> 'pynini.FstLike': + +def get_quantity( + decimal: "pynini.FstLike", cardinal_up_to_hundred: "pynini.FstLike" +) -> "pynini.FstLike": """ Returns FST that transforms either a cardinal or decimal followed by a quantity into a numeral, e.g. one million -> integer_part: "1" quantity: "million" e.g. one point five million -> integer_part: "1" fractional_part: "5" quantity: "million" - Args: + Args: decimal: decimal FST cardinal_up_to_hundred: cardinal FST """ numbers = cardinal_up_to_hundred @ ( - pynutil.delete(pynini.closure("0")) + pynini.difference(DAMO_DIGIT, "0") + pynini.closure(DAMO_DIGIT) + pynutil.delete(pynini.closure("0")) + + pynini.difference(DAMO_DIGIT, "0") + + pynini.closure(DAMO_DIGIT) ) suffix = pynini.union("juta", "miliar", "triliun", "kuadriliun", "triliun", "sextillion") res = ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + numbers - + pynutil.insert("\"") + + pynutil.insert('"') + delete_extra_space - + pynutil.insert("quantity: \"") + + pynutil.insert('quantity: "') + suffix - + pynutil.insert("\"") + + pynutil.insert('"') + ) + res |= ( + decimal + + delete_extra_space + + pynutil.insert('quantity: "') + + (suffix | "ribu") + + pynutil.insert('"') ) - res |= decimal + delete_extra_space + pynutil.insert("quantity: \"") + (suffix | "ribu") + pynutil.insert("\"") return res @@ -52,21 +63,30 @@ class DecimalFst(GraphFst): graph_decimal = pynini.string_file(get_abs_path("data/numbers/digit.tsv")) - graph_decimal |= pynini.string_file(get_abs_path("data/numbers/zero.tsv")) | pynini.cross("o", "0") + graph_decimal |= pynini.string_file(get_abs_path("data/numbers/zero.tsv")) | pynini.cross( + "o", "0" + ) graph_decimal = pynini.closure(graph_decimal + delete_space) + graph_decimal self.graph = graph_decimal - point = pynutil.delete("point") #titik + point = pynutil.delete("point") # titik optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("kurang", "\"true\"") + delete_extra_space, 0, 1 + pynutil.insert("negative: ") + pynini.cross("kurang", '"true"') + delete_extra_space, + 0, + 1, ) - graph_fractional = pynutil.insert("fractional_part: \"") + graph_decimal + pynutil.insert("\"") - graph_integer = pynutil.insert("integer_part: \"") + cardinal_graph + pynutil.insert("\"") + graph_fractional = ( + pynutil.insert('fractional_part: "') + graph_decimal + pynutil.insert('"') + ) + graph_integer = pynutil.insert('integer_part: "') + cardinal_graph + pynutil.insert('"') final_graph_wo_sign = ( - pynini.closure(graph_integer + delete_extra_space, 0, 1) + point + delete_extra_space + graph_fractional + pynini.closure(graph_integer + delete_extra_space, 0, 1) + + point + + delete_extra_space + + graph_fractional ) final_graph = optional_graph_negative + final_graph_wo_sign diff --git a/fun_text_processing/inverse_text_normalization/id/taggers/electronic.py b/fun_text_processing/inverse_text_normalization/id/taggers/electronic.py index b5aba4665..3d401f07d 100644 --- a/fun_text_processing/inverse_text_normalization/id/taggers/electronic.py +++ b/fun_text_processing/inverse_text_normalization/id/taggers/electronic.py @@ -1,6 +1,10 @@ import pynini from fun_text_processing.inverse_text_normalization.id.utils import get_abs_path -from fun_text_processing.inverse_text_normalization.id.graph_utils import DAMO_ALPHA, GraphFst, insert_space +from fun_text_processing.inverse_text_normalization.id.graph_utils import ( + DAMO_ALPHA, + GraphFst, + insert_space, +) from pynini.lib import pynutil @@ -24,35 +28,50 @@ class ElectronicFst(GraphFst): accepted_username = alpha_num | symbols process_dot = pynini.cross("dot", ".") - username = (alpha_num + pynini.closure(delete_extra_space + accepted_username)) | pynutil.add_weight( - pynini.closure(DAMO_ALPHA, 1), weight=0.0001 - ) - username = pynutil.insert("username: \"") + username + pynutil.insert("\"") + username = ( + alpha_num + pynini.closure(delete_extra_space + accepted_username) + ) | pynutil.add_weight(pynini.closure(DAMO_ALPHA, 1), weight=0.0001) + username = pynutil.insert('username: "') + username + pynutil.insert('"') single_alphanum = pynini.closure(alpha_num + delete_extra_space) + alpha_num - server = single_alphanum | pynini.string_file(get_abs_path("data/electronic/server_name.tsv")) + server = single_alphanum | pynini.string_file( + get_abs_path("data/electronic/server_name.tsv") + ) domain = single_alphanum | pynini.string_file(get_abs_path("data/electronic/domain.tsv")) domain_graph = ( - pynutil.insert("domain: \"") + pynutil.insert('domain: "') + server + delete_extra_space + process_dot + delete_extra_space + domain - + pynutil.insert("\"") + + pynutil.insert('"') + ) + graph = ( + username + + delete_extra_space + + pynutil.delete("at") + + insert_space + + delete_extra_space + + domain_graph ) - graph = username + delete_extra_space + pynutil.delete("at") + insert_space + delete_extra_space + domain_graph ############# url ### protocol_end = pynini.cross(pynini.union("w w w", "www"), "www") - protocol_start = (pynini.cross("h t t p", "http") | pynini.cross("h t t p s", "https")) + pynini.cross( - " colon slash slash ", "://" - ) + protocol_start = ( + pynini.cross("h t t p", "http") | pynini.cross("h t t p s", "https") + ) + pynini.cross(" colon slash slash ", "://") # .com, ending = ( delete_extra_space + symbols + delete_extra_space - + (domain | pynini.closure(accepted_username + delete_extra_space,) + accepted_username) + + ( + domain + | pynini.closure( + accepted_username + delete_extra_space, + ) + + accepted_username + ) ) protocol_default = ( @@ -63,12 +82,18 @@ class ElectronicFst(GraphFst): + pynini.closure(ending, 1) ).optimize() protocol = ( - pynini.closure(protocol_start, 0, 1) + protocol_end + delete_extra_space + process_dot + protocol_default + pynini.closure(protocol_start, 0, 1) + + protocol_end + + delete_extra_space + + process_dot + + protocol_default ).optimize() - protocol |= pynini.closure(protocol_end + delete_extra_space + process_dot, 0, 1) + protocol_default + protocol |= ( + pynini.closure(protocol_end + delete_extra_space + process_dot, 0, 1) + protocol_default + ) - protocol = pynutil.insert("protocol: \"") + protocol.optimize() + pynutil.insert("\"") + protocol = pynutil.insert('protocol: "') + protocol.optimize() + pynutil.insert('"') graph |= protocol final_graph = self.add_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/id/taggers/measure.py b/fun_text_processing/inverse_text_normalization/id/taggers/measure.py index 458469b97..519cc083c 100644 --- a/fun_text_processing/inverse_text_normalization/id/taggers/measure.py +++ b/fun_text_processing/inverse_text_normalization/id/taggers/measure.py @@ -31,22 +31,37 @@ class MeasureFst(GraphFst): graph_unit_plural = get_singulars(graph_unit_singular) # plural -> abbr optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("kurang", "\"true\"") + delete_extra_space, 0, 1 + pynutil.insert("negative: ") + pynini.cross("kurang", '"true"') + delete_extra_space, + 0, + 1, ) unit_singular = convert_space(graph_unit_singular) unit_plural = convert_space(graph_unit_plural) - unit_misc = pynutil.insert("/") + pynutil.delete("per") + delete_space + convert_space(graph_unit_singular) + unit_misc = ( + pynutil.insert("/") + + pynutil.delete("per") + + delete_space + + convert_space(graph_unit_singular) + ) unit_singular = ( - pynutil.insert("units: \"") - + (unit_singular | unit_misc | pynutil.add_weight(unit_singular + delete_space + unit_misc, 0.01)) - + pynutil.insert("\"") + pynutil.insert('units: "') + + ( + unit_singular + | unit_misc + | pynutil.add_weight(unit_singular + delete_space + unit_misc, 0.01) + ) + + pynutil.insert('"') ) unit_plural = ( - pynutil.insert("units: \"") - + (unit_plural | unit_misc | pynutil.add_weight(unit_plural + delete_space + unit_misc, 0.01)) - + pynutil.insert("\"") + pynutil.insert('units: "') + + ( + unit_plural + | unit_misc + | pynutil.add_weight(unit_plural + delete_space + unit_misc, 0.01) + ) + + pynutil.insert('"') ) subgraph_decimal = ( @@ -60,9 +75,9 @@ class MeasureFst(GraphFst): subgraph_cardinal = ( pynutil.insert("cardinal { ") + optional_graph_negative - + pynutil.insert("integer: \"") + + pynutil.insert('integer: "') + ((DAMO_SIGMA - "satu") @ cardinal_graph) - + pynutil.insert("\"") + + pynutil.insert('"') + pynutil.insert(" }") + delete_extra_space + unit_plural @@ -70,9 +85,9 @@ class MeasureFst(GraphFst): subgraph_cardinal |= ( pynutil.insert("cardinal { ") + optional_graph_negative - + pynutil.insert("integer: \"") + + pynutil.insert('integer: "') + pynini.cross("satu", "1") - + pynutil.insert("\"") + + pynutil.insert('"') + pynutil.insert(" }") + delete_extra_space + unit_singular diff --git a/fun_text_processing/inverse_text_normalization/id/taggers/money.py b/fun_text_processing/inverse_text_normalization/id/taggers/money.py index 61d20d522..fbe7af4d2 100644 --- a/fun_text_processing/inverse_text_normalization/id/taggers/money.py +++ b/fun_text_processing/inverse_text_normalization/id/taggers/money.py @@ -32,8 +32,11 @@ class MoneyFst(GraphFst): # add support for missing hundred (only for 3 digit numbers) # "one fifty" -> "one hundred fifty" with_hundred = pynini.compose( - pynini.closure(DAMO_NOT_SPACE) + pynini.accep(" ") + pynutil.insert("ratus ") + DAMO_SIGMA, - pynini.compose(cardinal_graph, DAMO_DIGIT ** 3), + pynini.closure(DAMO_NOT_SPACE) + + pynini.accep(" ") + + pynutil.insert("ratus ") + + DAMO_SIGMA, + pynini.compose(cardinal_graph, DAMO_DIGIT**3), ) cardinal_graph |= with_hundred graph_decimal_final = decimal.final_graph_wo_negative @@ -42,20 +45,27 @@ class MoneyFst(GraphFst): unit_singular = pynini.invert(unit) unit_plural = get_singulars(unit_singular) - graph_unit_singular = pynutil.insert("currency: \"") + convert_space(unit_singular) + pynutil.insert("\"") - graph_unit_plural = pynutil.insert("currency: \"") + convert_space(unit_plural) + pynutil.insert("\"") + graph_unit_singular = ( + pynutil.insert('currency: "') + convert_space(unit_singular) + pynutil.insert('"') + ) + graph_unit_plural = ( + pynutil.insert('currency: "') + convert_space(unit_plural) + pynutil.insert('"') + ) - add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | (pynutil.insert("0") + DAMO_DIGIT) + add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | ( + pynutil.insert("0") + DAMO_DIGIT + ) # twelve dollars (and) fifty cents, zero cents cents_standalone = ( - pynutil.insert("fractional_part: \"") + pynutil.insert('fractional_part: "') + pynini.union( - pynutil.add_weight(((DAMO_SIGMA - "satu") @ cardinal_graph), -0.7) @ add_leading_zero_to_double_digit + pynutil.add_weight(((DAMO_SIGMA - "satu") @ cardinal_graph), -0.7) + @ add_leading_zero_to_double_digit + delete_space + pynutil.delete("sen"), pynini.cross("satu", "01") + delete_space + pynutil.delete("sen"), ) - + pynutil.insert("\"") + + pynutil.insert('"') ) optional_cents_standalone = pynini.closure( @@ -69,31 +79,31 @@ class MoneyFst(GraphFst): # twelve dollars fifty, only after integer optional_cents_suffix = pynini.closure( delete_extra_space - + pynutil.insert("fractional_part: \"") + + pynutil.insert('fractional_part: "') + pynutil.add_weight(cardinal_graph @ add_leading_zero_to_double_digit, -0.7) - + pynutil.insert("\""), + + pynutil.insert('"'), 0, 1, ) graph_integer = ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + ((DAMO_SIGMA - "satu") @ cardinal_graph) - + pynutil.insert("\"") + + pynutil.insert('"') + delete_extra_space + graph_unit_plural + (optional_cents_standalone | optional_cents_suffix) ) graph_integer |= ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + pynini.cross("satu", "1") - + pynutil.insert("\"") + + pynutil.insert('"') + delete_extra_space + graph_unit_singular + (optional_cents_standalone | optional_cents_suffix) ) graph_decimal = graph_decimal_final + delete_extra_space + graph_unit_plural - graph_decimal |= pynutil.insert("currency: \"$\" integer_part: \"0\" ") + cents_standalone + graph_decimal |= pynutil.insert('currency: "$" integer_part: "0" ') + cents_standalone final_graph = graph_integer | graph_decimal final_graph = self.add_tokens(final_graph) diff --git a/fun_text_processing/inverse_text_normalization/id/taggers/ordinal.py b/fun_text_processing/inverse_text_normalization/id/taggers/ordinal.py index 3d68d9659..46c036313 100644 --- a/fun_text_processing/inverse_text_normalization/id/taggers/ordinal.py +++ b/fun_text_processing/inverse_text_normalization/id/taggers/ordinal.py @@ -21,9 +21,9 @@ class OrdinalFst(GraphFst): graph_teens = pynini.string_file(get_abs_path("data/ordinals/teen.tsv")) graph = pynini.closure(DAMO_CHAR) + pynini.union( graph_digit, graph_teens, pynini.cross("tieth", "ty"), pynini.cross("th", "") - ) # TODO + ) # TODO self.graph = graph @ cardinal_graph - final_graph = pynutil.insert("integer: \"") + self.graph + pynutil.insert("\"") + final_graph = pynutil.insert('integer: "') + self.graph + pynutil.insert('"') final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/id/taggers/punctuation.py b/fun_text_processing/inverse_text_normalization/id/taggers/punctuation.py index 611fdfaca..fee28ae1e 100644 --- a/fun_text_processing/inverse_text_normalization/id/taggers/punctuation.py +++ b/fun_text_processing/inverse_text_normalization/id/taggers/punctuation.py @@ -12,9 +12,9 @@ class PunctuationFst(GraphFst): def __init__(self): super().__init__(name="punctuation", kind="classify") - s = "!#$%&\'()*+,-./:;<=>?@^_`{|}~" + s = "!#$%&'()*+,-./:;<=>?@^_`{|}~" punct = pynini.union(*s) - graph = pynutil.insert("name: \"") + punct + pynutil.insert("\"") + graph = pynutil.insert('name: "') + punct + pynutil.insert('"') self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/id/taggers/telephone.py b/fun_text_processing/inverse_text_normalization/id/taggers/telephone.py index 8875d2906..6256fed3a 100644 --- a/fun_text_processing/inverse_text_normalization/id/taggers/telephone.py +++ b/fun_text_processing/inverse_text_normalization/id/taggers/telephone.py @@ -23,7 +23,7 @@ def get_serial_number(cardinal): class TelephoneFst(GraphFst): """ - Finite state transducer for classifying telephone numbers, e.g. + Finite state transducer for classifying telephone numbers, e.g. one two three one two three five six seven eight -> { number_part: "123-123-5678" } This class also support card number and IP format. @@ -60,61 +60,89 @@ class TelephoneFst(GraphFst): double_digit.invert() # to handle cases like "one twenty three" - two_digit_cardinal = pynini.compose(cardinal.graph_no_exception, DAMO_DIGIT ** 2) + two_digit_cardinal = pynini.compose(cardinal.graph_no_exception, DAMO_DIGIT**2) double_digit_to_digit = ( - pynini.compose(double_digit, str_to_digit + pynutil.delete(" ") + str_to_digit) | two_digit_cardinal + pynini.compose(double_digit, str_to_digit + pynutil.delete(" ") + str_to_digit) + | two_digit_cardinal ) - single_or_double_digit = (pynutil.add_weight(double_digit_to_digit, -0.0001) | str_to_digit).optimize() + single_or_double_digit = ( + pynutil.add_weight(double_digit_to_digit, -0.0001) | str_to_digit + ).optimize() single_or_double_digit |= ( single_or_double_digit - + pynini.closure(pynutil.add_weight(pynutil.delete(" ") + single_or_double_digit, 0.0001)) + + pynini.closure( + pynutil.add_weight(pynutil.delete(" ") + single_or_double_digit, 0.0001) + ) ).optimize() number_part = pynini.compose( single_or_double_digit, - DAMO_DIGIT ** 2 + pynutil.insert(" ") + DAMO_DIGIT ** 4 + pynutil.insert("-") + DAMO_DIGIT ** 4, + DAMO_DIGIT**2 + + pynutil.insert(" ") + + DAMO_DIGIT**4 + + pynutil.insert("-") + + DAMO_DIGIT**4, ).optimize() - number_part = pynutil.insert("number_part: \"") + number_part.optimize() + pynutil.insert("\"") + number_part = ( + pynutil.insert('number_part: "') + number_part.optimize() + pynutil.insert('"') + ) cardinal_option = pynini.compose(single_or_double_digit, DAMO_DIGIT ** (2, 3)) country_code = ( - pynutil.insert("country_code: \"") + pynutil.insert('country_code: "') + pynini.closure(pynini.cross("ditambah ", "+"), 0, 1) - + ((pynini.closure(str_to_digit + pynutil.delete(" "), 0, 2) + str_to_digit) | cardinal_option) - + pynutil.insert("\"") + + ( + (pynini.closure(str_to_digit + pynutil.delete(" "), 0, 2) + str_to_digit) + | cardinal_option + ) + + pynutil.insert('"') ) - optional_country_code = pynini.closure(country_code + pynutil.delete(" ") + insert_space, 0, 1).optimize() + optional_country_code = pynini.closure( + country_code + pynutil.delete(" ") + insert_space, 0, 1 + ).optimize() graph = optional_country_code + number_part # credit card number - space_four_digits = insert_space + DAMO_DIGIT ** 4 - credit_card_graph = pynini.compose(single_or_double_digit, DAMO_DIGIT ** 4 + space_four_digits ** 3).optimize() - graph |= pynutil.insert("number_part: \"") + credit_card_graph.optimize() + pynutil.insert("\"") + space_four_digits = insert_space + DAMO_DIGIT**4 + credit_card_graph = pynini.compose( + single_or_double_digit, DAMO_DIGIT**4 + space_four_digits**3 + ).optimize() + graph |= ( + pynutil.insert('number_part: "') + credit_card_graph.optimize() + pynutil.insert('"') + ) # SSN ssn_graph = pynini.compose( single_or_double_digit, - DAMO_DIGIT ** 3 + pynutil.insert("-") + DAMO_DIGIT ** 2 + pynutil.insert("-") + DAMO_DIGIT ** 4, + DAMO_DIGIT**3 + + pynutil.insert("-") + + DAMO_DIGIT**2 + + pynutil.insert("-") + + DAMO_DIGIT**4, ).optimize() - graph |= pynutil.insert("number_part: \"") + ssn_graph.optimize() + pynutil.insert("\"") + graph |= pynutil.insert('number_part: "') + ssn_graph.optimize() + pynutil.insert('"') # ip - digit_or_double = pynini.closure(str_to_digit + pynutil.delete(" "), 0, 1) + double_digit_to_digit - digit_or_double |= double_digit_to_digit + pynini.closure(pynutil.delete(" ") + str_to_digit, 0, 1) + digit_or_double = ( + pynini.closure(str_to_digit + pynutil.delete(" "), 0, 1) + double_digit_to_digit + ) + digit_or_double |= double_digit_to_digit + pynini.closure( + pynutil.delete(" ") + str_to_digit, 0, 1 + ) digit_or_double |= str_to_digit + (pynutil.delete(" ") + str_to_digit) ** (0, 2) digit_or_double |= cardinal_option digit_or_double = digit_or_double.optimize() ip_graph = digit_or_double + (pynini.cross(" dot ", ".") + digit_or_double) ** 3 - graph |= pynutil.insert("number_part: \"") + ip_graph.optimize() + pynutil.insert("\"") + graph |= pynutil.insert('number_part: "') + ip_graph.optimize() + pynutil.insert('"') graph |= ( - pynutil.insert("number_part: \"") + pynutil.insert('number_part: "') + pynutil.add_weight(get_serial_number(cardinal=cardinal), weight=0.0001) - + pynutil.insert("\"") + + pynutil.insert('"') ) final_graph = self.add_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/id/taggers/time.py b/fun_text_processing/inverse_text_normalization/id/taggers/time.py index 3a866f508..3071c9f84 100644 --- a/fun_text_processing/inverse_text_normalization/id/taggers/time.py +++ b/fun_text_processing/inverse_text_normalization/id/taggers/time.py @@ -33,7 +33,22 @@ class TimeFst(GraphFst): cardinal = pynutil.add_weight(CardinalFst().graph_no_exception, weight=-0.7) - labels_hour = ["nol", "satu", "dua", "tiga", "empat", "lima", "enam", "tujuh", "delapan", "sembilan", "sepuluh", "sebelas", "duabelas", "tigabelas"] + labels_hour = [ + "nol", + "satu", + "dua", + "tiga", + "empat", + "lima", + "enam", + "tujuh", + "delapan", + "sembilan", + "sepuluh", + "sebelas", + "duabelas", + "tigabelas", + ] labels_minute_single = [num_to_word(x) for x in range(1, 10)] labels_minute_double = [num_to_word(x) for x in range(10, 60)] @@ -44,21 +59,23 @@ class TimeFst(GraphFst): graph_minute_verbose = pynini.cross("setengah", "30") | pynini.cross("seperempat", "15") oclock = pynini.cross(pynini.union("jam", "pukul"), "") - final_graph_hour = pynutil.insert("hours: \"") + graph_hour + pynutil.insert("\"") + final_graph_hour = pynutil.insert('hours: "') + graph_hour + pynutil.insert('"') graph_minute = ( oclock + pynutil.insert("00") | pynutil.delete("o") + delete_space + graph_minute_single | graph_minute_double ) - final_suffix = pynutil.insert("suffix: \"") + convert_space(suffix_graph) + pynutil.insert("\"") + final_suffix = ( + pynutil.insert('suffix: "') + convert_space(suffix_graph) + pynutil.insert('"') + ) final_suffix = delete_space + insert_space + final_suffix final_suffix_optional = pynini.closure(final_suffix, 0, 1) final_time_zone_optional = pynini.closure( delete_space + insert_space - + pynutil.insert("zone: \"") + + pynutil.insert('zone: "') + convert_space(time_zone_graph) - + pynutil.insert("\""), + + pynutil.insert('"'), 0, 1, ) @@ -67,13 +84,17 @@ class TimeFst(GraphFst): # two o eight, two thirty five (am/pm) # two pm/am graph_hm = ( - final_graph_hour + delete_extra_space + pynutil.insert("minutes: \"") + graph_minute + pynutil.insert("\"") + final_graph_hour + + delete_extra_space + + pynutil.insert('minutes: "') + + graph_minute + + pynutil.insert('"') ) # 10 past four, quarter past four, half past four graph_m_past_h = ( - pynutil.insert("minutes: \"") + pynutil.insert('minutes: "') + pynini.union(graph_minute_single, graph_minute_double, graph_minute_verbose) - + pynutil.insert("\"") + + pynutil.insert('"') + delete_space + pynutil.delete("setengah") + delete_extra_space @@ -81,42 +102,46 @@ class TimeFst(GraphFst): ) graph_quarter_time = ( - pynutil.insert("minutes: \"") + pynutil.insert('minutes: "') + pynini.cross("seperempat", "45") - + pynutil.insert("\"") + + pynutil.insert('"') + delete_space + pynutil.delete(pynini.union("to", "till")) + delete_extra_space - + pynutil.insert("hours: \"") + + pynutil.insert('hours: "') + to_hour_graph - + pynutil.insert("\"") + + pynutil.insert('"') ) graph_m_to_h_suffix_time = ( - pynutil.insert("minutes: \"") + pynutil.insert('minutes: "') + ((graph_minute_single | graph_minute_double).optimize() @ minute_to_graph) - + pynutil.insert("\"") - + pynini.closure(delete_space + pynutil.delete(pynini.union("min", "mins", "menit")), 0, 1) + + pynutil.insert('"') + + pynini.closure( + delete_space + pynutil.delete(pynini.union("min", "mins", "menit")), 0, 1 + ) + delete_space + pynutil.delete(pynini.union("sampai")) + delete_extra_space - + pynutil.insert("hours: \"") + + pynutil.insert('hours: "') + to_hour_graph - + pynutil.insert("\"") + + pynutil.insert('"') + final_suffix ) graph_h = ( final_graph_hour + delete_extra_space - + pynutil.insert("minutes: \"") + + pynutil.insert('minutes: "') + (pynutil.insert("00") | graph_minute) - + pynutil.insert("\"") + + pynutil.insert('"') + final_suffix + final_time_zone_optional ) final_graph = ( - (graph_hm | graph_m_past_h | graph_quarter_time) + final_suffix_optional + final_time_zone_optional + (graph_hm | graph_m_past_h | graph_quarter_time) + + final_suffix_optional + + final_time_zone_optional ) final_graph |= graph_h final_graph |= graph_m_to_h_suffix_time diff --git a/fun_text_processing/inverse_text_normalization/id/taggers/tokenize_and_classify.py b/fun_text_processing/inverse_text_normalization/id/taggers/tokenize_and_classify.py index 2460b34b5..8037c250c 100644 --- a/fun_text_processing/inverse_text_normalization/id/taggers/tokenize_and_classify.py +++ b/fun_text_processing/inverse_text_normalization/id/taggers/tokenize_and_classify.py @@ -27,7 +27,7 @@ import logging class ClassifyFst(GraphFst): """ Final class that composes all other classification grammars. This class can process an entire sentence, that is lower cased. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. Args: @@ -80,10 +80,16 @@ class ClassifyFst(GraphFst): | pynutil.add_weight(word_graph, 100) ) - punct = pynutil.insert("tokens { ") + pynutil.add_weight(punct_graph, weight=1.1) + pynutil.insert(" }") + punct = ( + pynutil.insert("tokens { ") + + pynutil.add_weight(punct_graph, weight=1.1) + + pynutil.insert(" }") + ) token = pynutil.insert("tokens { ") + classify + pynutil.insert(" }") token_plus_punct = ( - pynini.closure(punct + pynutil.insert(" ")) + token + pynini.closure(pynutil.insert(" ") + punct) + pynini.closure(punct + pynutil.insert(" ")) + + token + + pynini.closure(pynutil.insert(" ") + punct) ) graph = token_plus_punct + pynini.closure(delete_extra_space + token_plus_punct) diff --git a/fun_text_processing/inverse_text_normalization/id/taggers/whitelist.py b/fun_text_processing/inverse_text_normalization/id/taggers/whitelist.py index 87a54c153..c238f2caf 100644 --- a/fun_text_processing/inverse_text_normalization/id/taggers/whitelist.py +++ b/fun_text_processing/inverse_text_normalization/id/taggers/whitelist.py @@ -15,5 +15,5 @@ class WhiteListFst(GraphFst): super().__init__(name="whitelist", kind="classify") whitelist = pynini.string_file(get_abs_path("data/whitelist.tsv")).invert() - graph = pynutil.insert("name: \"") + convert_space(whitelist) + pynutil.insert("\"") + graph = pynutil.insert('name: "') + convert_space(whitelist) + pynutil.insert('"') self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/id/taggers/word.py b/fun_text_processing/inverse_text_normalization/id/taggers/word.py index 76098c57e..8af6ed9f8 100644 --- a/fun_text_processing/inverse_text_normalization/id/taggers/word.py +++ b/fun_text_processing/inverse_text_normalization/id/taggers/word.py @@ -11,5 +11,5 @@ class WordFst(GraphFst): def __init__(self): super().__init__(name="word", kind="classify") - word = pynutil.insert("name: \"") + pynini.closure(DAMO_NOT_SPACE, 1) + pynutil.insert("\"") + word = pynutil.insert('name: "') + pynini.closure(DAMO_NOT_SPACE, 1) + pynutil.insert('"') self.fst = word.optimize() diff --git a/fun_text_processing/inverse_text_normalization/id/utils.py b/fun_text_processing/inverse_text_normalization/id/utils.py index 9b9786777..38a95bac5 100644 --- a/fun_text_processing/inverse_text_normalization/id/utils.py +++ b/fun_text_processing/inverse_text_normalization/id/utils.py @@ -13,7 +13,7 @@ def num_to_word(x: Union[str, int]): Args x: integer - Returns: spoken representation + Returns: spoken representation """ if isinstance(x, int): x = str(x) @@ -27,7 +27,7 @@ def get_abs_path(rel_path): Args: rel_path: relative path to this file - + Returns absolute path """ - return os.path.dirname(os.path.abspath(__file__)) + '/' + rel_path + return os.path.dirname(os.path.abspath(__file__)) + "/" + rel_path diff --git a/fun_text_processing/inverse_text_normalization/id/verbalizers/cardinal.py b/fun_text_processing/inverse_text_normalization/id/verbalizers/cardinal.py index 741a62b17..e5ed7eac6 100644 --- a/fun_text_processing/inverse_text_normalization/id/verbalizers/cardinal.py +++ b/fun_text_processing/inverse_text_normalization/id/verbalizers/cardinal.py @@ -1,5 +1,9 @@ import pynini -from fun_text_processing.inverse_text_normalization.id.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.id.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -14,9 +18,9 @@ class CardinalFst(GraphFst): optional_sign = pynini.closure( pynutil.delete("negative:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + DAMO_NOT_QUOTE - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space, 0, 1, @@ -24,9 +28,9 @@ class CardinalFst(GraphFst): graph = ( pynutil.delete("integer:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) self.numbers = graph graph = optional_sign + graph diff --git a/fun_text_processing/inverse_text_normalization/id/verbalizers/date.py b/fun_text_processing/inverse_text_normalization/id/verbalizers/date.py index 82ad1fcfb..995836cbf 100644 --- a/fun_text_processing/inverse_text_normalization/id/verbalizers/date.py +++ b/fun_text_processing/inverse_text_normalization/id/verbalizers/date.py @@ -20,43 +20,47 @@ class DateFst(GraphFst): month = ( pynutil.delete("month:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) day = ( pynutil.delete("day:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) year = ( pynutil.delete("year:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') ) # month (day) year graph_mdy = ( - month + pynini.closure(delete_extra_space + day, 0, 1) + pynini.closure(delete_extra_space + year, 0, 1) + month + + pynini.closure(delete_extra_space + day, 0, 1) + + pynini.closure(delete_extra_space + year, 0, 1) ) # (day) month year graph_dmy = ( - pynini.closure(day + delete_extra_space, 0, 1) + month + pynini.closure(delete_extra_space + year, 0, 1) + pynini.closure(day + delete_extra_space, 0, 1) + + month + + pynini.closure(delete_extra_space + year, 0, 1) ) optional_preserve_order = pynini.closure( pynutil.delete("preserve_order:") + delete_space + pynutil.delete("true") + delete_space | pynutil.delete("field_order:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + DAMO_NOT_QUOTE - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space ) diff --git a/fun_text_processing/inverse_text_normalization/id/verbalizers/decimal.py b/fun_text_processing/inverse_text_normalization/id/verbalizers/decimal.py index 044780dc5..99cf63469 100644 --- a/fun_text_processing/inverse_text_normalization/id/verbalizers/decimal.py +++ b/fun_text_processing/inverse_text_normalization/id/verbalizers/decimal.py @@ -1,5 +1,9 @@ import pynini -from fun_text_processing.inverse_text_normalization.id.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.id.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -11,30 +15,30 @@ class DecimalFst(GraphFst): def __init__(self): super().__init__(name="decimal", kind="verbalize") - optionl_sign = pynini.closure(pynini.cross("negative: \"true\"", "-") + delete_space, 0, 1) + optionl_sign = pynini.closure(pynini.cross('negative: "true"', "-") + delete_space, 0, 1) integer = ( pynutil.delete("integer_part:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_integer = pynini.closure(integer + delete_space, 0, 1) fractional = ( pynutil.insert(".") + pynutil.delete("fractional_part:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_fractional = pynini.closure(fractional + delete_space, 0, 1) quantity = ( pynutil.delete("quantity:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_quantity = pynini.closure(pynutil.insert(" ") + quantity + delete_space, 0, 1) graph = optional_integer + optional_fractional + optional_quantity diff --git a/fun_text_processing/inverse_text_normalization/id/verbalizers/electronic.py b/fun_text_processing/inverse_text_normalization/id/verbalizers/electronic.py index d88424bd3..4994cb90c 100644 --- a/fun_text_processing/inverse_text_normalization/id/verbalizers/electronic.py +++ b/fun_text_processing/inverse_text_normalization/id/verbalizers/electronic.py @@ -1,5 +1,9 @@ import pynini -from fun_text_processing.inverse_text_normalization.id.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.id.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -14,24 +18,24 @@ class ElectronicFst(GraphFst): user_name = ( pynutil.delete("username:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) domain = ( pynutil.delete("domain:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) protocol = ( pynutil.delete("protocol:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) graph = user_name + delete_space + pynutil.insert("@") + domain diff --git a/fun_text_processing/inverse_text_normalization/id/verbalizers/fraction.py b/fun_text_processing/inverse_text_normalization/id/verbalizers/fraction.py index 15557da4f..a424b258d 100644 --- a/fun_text_processing/inverse_text_normalization/id/verbalizers/fraction.py +++ b/fun_text_processing/inverse_text_normalization/id/verbalizers/fraction.py @@ -3,7 +3,7 @@ from fun_text_processing.inverse_text_normalization.id.graph_utils import GraphF class FractionFst(GraphFst): """ - Finite state transducer for verbalizing fraction, + Finite state transducer for verbalizing fraction, """ def __init__(self): diff --git a/fun_text_processing/inverse_text_normalization/id/verbalizers/measure.py b/fun_text_processing/inverse_text_normalization/id/verbalizers/measure.py index 8695cc9b9..470169d76 100644 --- a/fun_text_processing/inverse_text_normalization/id/verbalizers/measure.py +++ b/fun_text_processing/inverse_text_normalization/id/verbalizers/measure.py @@ -1,5 +1,9 @@ import pynini -from fun_text_processing.inverse_text_normalization.id.graph_utils import DAMO_CHAR, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.id.graph_utils import ( + DAMO_CHAR, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -15,13 +19,13 @@ class MeasureFst(GraphFst): def __init__(self, decimal: GraphFst, cardinal: GraphFst): super().__init__(name="measure", kind="verbalize") - optional_sign = pynini.closure(pynini.cross("negative: \"true\"", "-"), 0, 1) + optional_sign = pynini.closure(pynini.cross('negative: "true"', "-"), 0, 1) unit = ( pynutil.delete("units:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space ) graph_decimal = ( diff --git a/fun_text_processing/inverse_text_normalization/id/verbalizers/money.py b/fun_text_processing/inverse_text_normalization/id/verbalizers/money.py index d02faaa47..11bf87544 100644 --- a/fun_text_processing/inverse_text_normalization/id/verbalizers/money.py +++ b/fun_text_processing/inverse_text_normalization/id/verbalizers/money.py @@ -1,5 +1,9 @@ import pynini -from fun_text_processing.inverse_text_normalization.id.graph_utils import DAMO_CHAR, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.id.graph_utils import ( + DAMO_CHAR, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -17,9 +21,9 @@ class MoneyFst(GraphFst): unit = ( pynutil.delete("currency:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) graph = unit + delete_space + decimal.numbers delete_tokens = self.delete_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/id/verbalizers/ordinal.py b/fun_text_processing/inverse_text_normalization/id/verbalizers/ordinal.py index 0571f7861..2a6801f6c 100644 --- a/fun_text_processing/inverse_text_normalization/id/verbalizers/ordinal.py +++ b/fun_text_processing/inverse_text_normalization/id/verbalizers/ordinal.py @@ -1,5 +1,10 @@ import pynini -from fun_text_processing.inverse_text_normalization.id.graph_utils import DAMO_NOT_QUOTE, DAMO_SIGMA, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.id.graph_utils import ( + DAMO_NOT_QUOTE, + DAMO_SIGMA, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -14,9 +19,9 @@ class OrdinalFst(GraphFst): graph = ( pynutil.delete("integer:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) convert_eleven = pynini.cross("11", "11th") convert_twelve = pynini.cross("12", "12th") diff --git a/fun_text_processing/inverse_text_normalization/id/verbalizers/telephone.py b/fun_text_processing/inverse_text_normalization/id/verbalizers/telephone.py index fb41d58c6..23320571b 100644 --- a/fun_text_processing/inverse_text_normalization/id/verbalizers/telephone.py +++ b/fun_text_processing/inverse_text_normalization/id/verbalizers/telephone.py @@ -7,17 +7,21 @@ class TelephoneFst(GraphFst): """ Finite state transducer for verbalizing telephone, e.g. telephone { number_part: "123-123-5678" } - -> 123-123-5678 + -> 123-123-5678 """ def __init__(self): super().__init__(name="telephone", kind="verbalize") - number_part = pynutil.delete("number_part: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") - optional_country_code = pynini.closure( - pynutil.delete("country_code: \"") + number_part = ( + pynutil.delete('number_part: "') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') + ) + optional_country_code = pynini.closure( + pynutil.delete('country_code: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + pynini.accep(" "), 0, 1, diff --git a/fun_text_processing/inverse_text_normalization/id/verbalizers/time.py b/fun_text_processing/inverse_text_normalization/id/verbalizers/time.py index da02ae2e2..b98fc9912 100644 --- a/fun_text_processing/inverse_text_normalization/id/verbalizers/time.py +++ b/fun_text_processing/inverse_text_normalization/id/verbalizers/time.py @@ -19,29 +19,31 @@ class TimeFst(GraphFst): def __init__(self): super().__init__(name="time", kind="verbalize") - add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | (pynutil.insert("0") + DAMO_DIGIT) + add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | ( + pynutil.insert("0") + DAMO_DIGIT + ) hour = ( pynutil.delete("hours:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_DIGIT, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) minute = ( pynutil.delete("minutes:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_DIGIT, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) suffix = ( delete_space + insert_space + pynutil.delete("suffix:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_suffix = pynini.closure(suffix, 0, 1) zone = ( @@ -49,9 +51,9 @@ class TimeFst(GraphFst): + insert_space + pynutil.delete("zone:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_zone = pynini.closure(zone, 0, 1) graph = ( diff --git a/fun_text_processing/inverse_text_normalization/id/verbalizers/verbalize.py b/fun_text_processing/inverse_text_normalization/id/verbalizers/verbalize.py index 8271474b3..7aaef2aef 100644 --- a/fun_text_processing/inverse_text_normalization/id/verbalizers/verbalize.py +++ b/fun_text_processing/inverse_text_normalization/id/verbalizers/verbalize.py @@ -14,7 +14,7 @@ from fun_text_processing.inverse_text_normalization.id.graph_utils import GraphF class VerbalizeFst(GraphFst): """ Composes other verbalizer grammars. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. """ diff --git a/fun_text_processing/inverse_text_normalization/id/verbalizers/verbalize_final.py b/fun_text_processing/inverse_text_normalization/id/verbalizers/verbalize_final.py index d5fcf589a..014dac1b0 100644 --- a/fun_text_processing/inverse_text_normalization/id/verbalizers/verbalize_final.py +++ b/fun_text_processing/inverse_text_normalization/id/verbalizers/verbalize_final.py @@ -1,13 +1,17 @@ import pynini from fun_text_processing.inverse_text_normalization.id.verbalizers.verbalize import VerbalizeFst from fun_text_processing.inverse_text_normalization.id.verbalizers.word import WordFst -from fun_text_processing.inverse_text_normalization.id.graph_utils import GraphFst, delete_extra_space, delete_space +from fun_text_processing.inverse_text_normalization.id.graph_utils import ( + GraphFst, + delete_extra_space, + delete_space, +) from pynini.lib import pynutil class VerbalizeFinalFst(GraphFst): """ - Finite state transducer that verbalizes an entire sentence, e.g. + Finite state transducer that verbalizes an entire sentence, e.g. tokens { name: "its" } tokens { time { hours: "12" minutes: "30" } } tokens { name: "now" } -> its 12:30 now """ diff --git a/fun_text_processing/inverse_text_normalization/id/verbalizers/whitelist.py b/fun_text_processing/inverse_text_normalization/id/verbalizers/whitelist.py index f1ced4776..445020254 100644 --- a/fun_text_processing/inverse_text_normalization/id/verbalizers/whitelist.py +++ b/fun_text_processing/inverse_text_normalization/id/verbalizers/whitelist.py @@ -1,5 +1,10 @@ import pynini -from fun_text_processing.inverse_text_normalization.id.graph_utils import DAMO_CHAR, DAMO_SIGMA, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.id.graph_utils import ( + DAMO_CHAR, + DAMO_SIGMA, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -14,9 +19,9 @@ class WhiteListFst(GraphFst): graph = ( pynutil.delete("name:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) - graph = graph @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + graph = graph @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/id/verbalizers/word.py b/fun_text_processing/inverse_text_normalization/id/verbalizers/word.py index 9a638f2ac..079c6136c 100644 --- a/fun_text_processing/inverse_text_normalization/id/verbalizers/word.py +++ b/fun_text_processing/inverse_text_normalization/id/verbalizers/word.py @@ -1,5 +1,10 @@ import pynini -from fun_text_processing.inverse_text_normalization.id.graph_utils import DAMO_CHAR, DAMO_SIGMA, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.id.graph_utils import ( + DAMO_CHAR, + DAMO_SIGMA, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -12,7 +17,13 @@ class WordFst(GraphFst): def __init__(self): super().__init__(name="word", kind="verbalize") chars = pynini.closure(DAMO_CHAR - " ", 1) - char = pynutil.delete("name:") + delete_space + pynutil.delete("\"") + chars + pynutil.delete("\"") - graph = char @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + char = ( + pynutil.delete("name:") + + delete_space + + pynutil.delete('"') + + chars + + pynutil.delete('"') + ) + graph = char @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/inverse_normalize.py b/fun_text_processing/inverse_text_normalization/inverse_normalize.py index fe6aa3198..2dffd5d98 100644 --- a/fun_text_processing/inverse_text_normalization/inverse_normalize.py +++ b/fun_text_processing/inverse_text_normalization/inverse_normalize.py @@ -21,73 +21,102 @@ class InverseNormalizer(Normalizer): overwrite_cache: set to True to overwrite .far files """ - def __init__(self, lang: str = 'en', cache_dir: str = None, overwrite_cache: bool = False, - enable_standalone_number: bool = True, - enable_0_to_9: bool = True): + def __init__( + self, + lang: str = "en", + cache_dir: str = None, + overwrite_cache: bool = False, + enable_standalone_number: bool = True, + enable_0_to_9: bool = True, + ): - if lang == 'en': - from fun_text_processing.inverse_text_normalization.en.taggers.tokenize_and_classify import ClassifyFst + if lang == "en": + from fun_text_processing.inverse_text_normalization.en.taggers.tokenize_and_classify import ( + ClassifyFst, + ) from fun_text_processing.inverse_text_normalization.en.verbalizers.verbalize_final import ( VerbalizeFinalFst, ) - elif lang == 'id': - from fun_text_processing.inverse_text_normalization.id.taggers.tokenize_and_classify import ClassifyFst + elif lang == "id": + from fun_text_processing.inverse_text_normalization.id.taggers.tokenize_and_classify import ( + ClassifyFst, + ) from fun_text_processing.inverse_text_normalization.id.verbalizers.verbalize_final import ( VerbalizeFinalFst, ) - elif lang == 'ja': - from fun_text_processing.inverse_text_normalization.ja.taggers.tokenize_and_classify import ClassifyFst + elif lang == "ja": + from fun_text_processing.inverse_text_normalization.ja.taggers.tokenize_and_classify import ( + ClassifyFst, + ) from fun_text_processing.inverse_text_normalization.ja.verbalizers.verbalize_final import ( VerbalizeFinalFst, ) - elif lang == 'es': - from fun_text_processing.inverse_text_normalization.es.taggers.tokenize_and_classify import ClassifyFst + elif lang == "es": + from fun_text_processing.inverse_text_normalization.es.taggers.tokenize_and_classify import ( + ClassifyFst, + ) from fun_text_processing.inverse_text_normalization.es.verbalizers.verbalize_final import ( VerbalizeFinalFst, ) - elif lang == 'pt': - from fun_text_processing.inverse_text_normalization.pt.taggers.tokenize_and_classify import ClassifyFst + elif lang == "pt": + from fun_text_processing.inverse_text_normalization.pt.taggers.tokenize_and_classify import ( + ClassifyFst, + ) from fun_text_processing.inverse_text_normalization.pt.verbalizers.verbalize_final import ( VerbalizeFinalFst, ) - elif lang == 'ru': - from fun_text_processing.inverse_text_normalization.ru.taggers.tokenize_and_classify import ClassifyFst + elif lang == "ru": + from fun_text_processing.inverse_text_normalization.ru.taggers.tokenize_and_classify import ( + ClassifyFst, + ) from fun_text_processing.inverse_text_normalization.ru.verbalizers.verbalize_final import ( VerbalizeFinalFst, ) - elif lang == 'de': - from fun_text_processing.inverse_text_normalization.de.taggers.tokenize_and_classify import ClassifyFst + elif lang == "de": + from fun_text_processing.inverse_text_normalization.de.taggers.tokenize_and_classify import ( + ClassifyFst, + ) from fun_text_processing.inverse_text_normalization.de.verbalizers.verbalize_final import ( VerbalizeFinalFst, ) - elif lang == 'fr': - from fun_text_processing.inverse_text_normalization.fr.taggers.tokenize_and_classify import ClassifyFst + elif lang == "fr": + from fun_text_processing.inverse_text_normalization.fr.taggers.tokenize_and_classify import ( + ClassifyFst, + ) from fun_text_processing.inverse_text_normalization.fr.verbalizers.verbalize_final import ( VerbalizeFinalFst, ) - elif lang == 'vi': - from fun_text_processing.inverse_text_normalization.vi.taggers.tokenize_and_classify import ClassifyFst + elif lang == "vi": + from fun_text_processing.inverse_text_normalization.vi.taggers.tokenize_and_classify import ( + ClassifyFst, + ) from fun_text_processing.inverse_text_normalization.vi.verbalizers.verbalize_final import ( VerbalizeFinalFst, ) - elif lang == 'ko': - from fun_text_processing.inverse_text_normalization.ko.taggers.tokenize_and_classify import ClassifyFst + elif lang == "ko": + from fun_text_processing.inverse_text_normalization.ko.taggers.tokenize_and_classify import ( + ClassifyFst, + ) from fun_text_processing.inverse_text_normalization.ko.verbalizers.verbalize_final import ( VerbalizeFinalFst, ) - elif lang == 'zh': - from fun_text_processing.inverse_text_normalization.zh.taggers.tokenize_and_classify import ClassifyFst + elif lang == "zh": + from fun_text_processing.inverse_text_normalization.zh.taggers.tokenize_and_classify import ( + ClassifyFst, + ) from fun_text_processing.inverse_text_normalization.zh.verbalizers.verbalize_final import ( VerbalizeFinalFst, ) - elif lang == 'tl': - from fun_text_processing.inverse_text_normalization.tl.taggers.tokenize_and_classify import ClassifyFst + elif lang == "tl": + from fun_text_processing.inverse_text_normalization.tl.taggers.tokenize_and_classify import ( + ClassifyFst, + ) from fun_text_processing.inverse_text_normalization.tl.verbalizers.verbalize_final import ( VerbalizeFinalFst, ) @@ -129,51 +158,62 @@ class InverseNormalizer(Normalizer): def str2bool(s, default=False): s = s.lower() - if s == 'true': + if s == "true": return True - elif s == 'false': + elif s == "false": return False else: return default + def parse_args(): parser = ArgumentParser() input = parser.add_mutually_exclusive_group() input.add_argument("--text", dest="input_string", help="input string", type=str) input.add_argument("--input_file", dest="input_file", help="input file path", type=str) - parser.add_argument('--output_file', dest="output_file", help="output file path", type=str) + parser.add_argument("--output_file", dest="output_file", help="output file path", type=str) parser.add_argument( - "--language", help="language", choices=['en', 'id', 'ja', 'de', 'es', 'pt', 'ru', 'fr', 'vi', 'ko', 'zh', 'tl'], default="en", type=str + "--language", + help="language", + choices=["en", "id", "ja", "de", "es", "pt", "ru", "fr", "vi", "ko", "zh", "tl"], + default="en", + type=str, + ) + parser.add_argument("--verbose", help="print info for debugging", action="store_true") + parser.add_argument( + "--overwrite_cache", help="set to True to re-create .far grammar files", action="store_true" ) - parser.add_argument("--verbose", help="print info for debugging", action='store_true') - parser.add_argument("--overwrite_cache", help="set to True to re-create .far grammar files", action="store_true") parser.add_argument( "--cache_dir", help="path to a dir with .far grammar file. Set to None to avoid using cache", default=None, type=str, ) - parser.add_argument('--enable_standalone_number', type=str, - default='True', - help='enable standalone number') - parser.add_argument('--enable_0_to_9', type=str, - default='True', - help='enable convert number 0 to 9') + parser.add_argument( + "--enable_standalone_number", type=str, default="True", help="enable standalone number" + ) + parser.add_argument( + "--enable_0_to_9", type=str, default="True", help="enable convert number 0 to 9" + ) return parser.parse_args() if __name__ == "__main__": args = parse_args() start_time = perf_counter() - if args.language == 'ja': - inverse_normalizer = InverseNormalizer(lang=args.language, cache_dir=args.cache_dir, overwrite_cache=args.overwrite_cache, - enable_standalone_number=str2bool(args.enable_standalone_number), - enable_0_to_9=str2bool(args.enable_0_to_9)) + if args.language == "ja": + inverse_normalizer = InverseNormalizer( + lang=args.language, + cache_dir=args.cache_dir, + overwrite_cache=args.overwrite_cache, + enable_standalone_number=str2bool(args.enable_standalone_number), + enable_0_to_9=str2bool(args.enable_0_to_9), + ) else: inverse_normalizer = InverseNormalizer( lang=args.language, cache_dir=args.cache_dir, overwrite_cache=args.overwrite_cache ) - print(f'Time to generate graph: {round(perf_counter() - start_time, 2)} sec') + print(f"Time to generate graph: {round(perf_counter() - start_time, 2)} sec") if args.input_string: print(inverse_normalizer.inverse_normalize(args.input_string, verbose=args.verbose)) diff --git a/fun_text_processing/inverse_text_normalization/ja/__init__.py b/fun_text_processing/inverse_text_normalization/ja/__init__.py index a3740c25e..4f73ff2c3 100644 --- a/fun_text_processing/inverse_text_normalization/ja/__init__.py +++ b/fun_text_processing/inverse_text_normalization/ja/__init__.py @@ -1,3 +1,7 @@ -from fun_text_processing.inverse_text_normalization.en.taggers.tokenize_and_classify import ClassifyFst +from fun_text_processing.inverse_text_normalization.en.taggers.tokenize_and_classify import ( + ClassifyFst, +) from fun_text_processing.inverse_text_normalization.en.verbalizers.verbalize import VerbalizeFst -from fun_text_processing.inverse_text_normalization.en.verbalizers.verbalize_final import VerbalizeFinalFst +from fun_text_processing.inverse_text_normalization.en.verbalizers.verbalize_final import ( + VerbalizeFinalFst, +) diff --git a/fun_text_processing/inverse_text_normalization/ja/clean_eval_data.py b/fun_text_processing/inverse_text_normalization/ja/clean_eval_data.py index c1e1bcfaa..e975bf5ec 100644 --- a/fun_text_processing/inverse_text_normalization/ja/clean_eval_data.py +++ b/fun_text_processing/inverse_text_normalization/ja/clean_eval_data.py @@ -53,7 +53,7 @@ class Filter: Args: processes given instance with process function - + Returns: processed instance if instance belongs to expected class type or original instance """ if instance.token_type != self.class_type: @@ -71,7 +71,9 @@ def process_cardinal_1(instance: Instance) -> Instance: normalized = instance.normalized un_normalized = re.sub(r"[^0-9]", "", un_normalized) normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_ordinal_1(instance: Instance) -> bool: @@ -84,7 +86,9 @@ def process_ordinal_1(instance: Instance) -> Instance: normalized = instance.normalized un_normalized = re.sub(r"[,\s]", "", un_normalized) normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_decimal_1(instance: Instance) -> bool: @@ -97,7 +101,9 @@ def process_decimal_1(instance: Instance) -> Instance: un_normalized = re.sub(r",", "", un_normalized) normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_measure_1(instance: Instance) -> bool: @@ -114,7 +120,9 @@ def process_measure_1(instance: Instance) -> Instance: normalized = re.sub(r"[^a-z\s]", "", normalized) normalized = re.sub(r"per ([a-z\s]*)s$", r"per \1", normalized) normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_money_1(instance: Instance) -> bool: @@ -131,7 +139,9 @@ def process_money_1(instance: Instance) -> Instance: un_normalized = re.sub(r"(\d)m\s*$", r"\1 million", un_normalized) un_normalized = re.sub(r"(\d)bn?\s*$", r"\1 billion", un_normalized) normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_time_1(instance: Instance) -> bool: @@ -146,7 +156,9 @@ def process_time_1(instance: Instance) -> Instance: un_normalized = re.sub(r"(\d)\s?p\s?m\s?", r"\1 p.m.", un_normalized) normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_plain_1(instance: Instance) -> bool: @@ -157,7 +169,9 @@ def filter_plain_1(instance: Instance) -> bool: def process_plain_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_punct_1(instance: Instance) -> bool: @@ -168,7 +182,9 @@ def filter_punct_1(instance: Instance) -> bool: def process_punct_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_date_1(instance: Instance) -> bool: @@ -181,7 +197,9 @@ def process_date_1(instance: Instance) -> Instance: un_normalized = re.sub(r",", "", un_normalized) normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_letters_1(instance: Instance) -> bool: @@ -193,7 +211,9 @@ def process_letters_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_verbatim_1(instance: Instance) -> bool: @@ -204,7 +224,9 @@ def filter_verbatim_1(instance: Instance) -> bool: def process_verbatim_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_digit_1(instance: Instance) -> bool: @@ -216,7 +238,9 @@ def process_digit_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_telephone_1(instance: Instance) -> bool: @@ -228,7 +252,9 @@ def process_telephone_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_electronic_1(instance: Instance) -> bool: @@ -240,7 +266,9 @@ def process_electronic_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_fraction_1(instance: Instance) -> bool: @@ -252,7 +280,9 @@ def process_fraction_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_address_1(instance: Instance) -> bool: @@ -264,27 +294,51 @@ def process_address_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) filters = [] -filters.append(Filter(class_type="CARDINAL", process_func=process_cardinal_1, filter_func=filter_cardinal_1)) -filters.append(Filter(class_type="ORDINAL", process_func=process_ordinal_1, filter_func=filter_ordinal_1)) -filters.append(Filter(class_type="DECIMAL", process_func=process_decimal_1, filter_func=filter_decimal_1)) -filters.append(Filter(class_type="MEASURE", process_func=process_measure_1, filter_func=filter_measure_1)) +filters.append( + Filter(class_type="CARDINAL", process_func=process_cardinal_1, filter_func=filter_cardinal_1) +) +filters.append( + Filter(class_type="ORDINAL", process_func=process_ordinal_1, filter_func=filter_ordinal_1) +) +filters.append( + Filter(class_type="DECIMAL", process_func=process_decimal_1, filter_func=filter_decimal_1) +) +filters.append( + Filter(class_type="MEASURE", process_func=process_measure_1, filter_func=filter_measure_1) +) filters.append(Filter(class_type="MONEY", process_func=process_money_1, filter_func=filter_money_1)) filters.append(Filter(class_type="TIME", process_func=process_time_1, filter_func=filter_time_1)) filters.append(Filter(class_type="DATE", process_func=process_date_1, filter_func=filter_date_1)) filters.append(Filter(class_type="PLAIN", process_func=process_plain_1, filter_func=filter_plain_1)) filters.append(Filter(class_type="PUNCT", process_func=process_punct_1, filter_func=filter_punct_1)) -filters.append(Filter(class_type="LETTERS", process_func=process_letters_1, filter_func=filter_letters_1)) -filters.append(Filter(class_type="VERBATIM", process_func=process_verbatim_1, filter_func=filter_verbatim_1)) +filters.append( + Filter(class_type="LETTERS", process_func=process_letters_1, filter_func=filter_letters_1) +) +filters.append( + Filter(class_type="VERBATIM", process_func=process_verbatim_1, filter_func=filter_verbatim_1) +) filters.append(Filter(class_type="DIGIT", process_func=process_digit_1, filter_func=filter_digit_1)) -filters.append(Filter(class_type="TELEPHONE", process_func=process_telephone_1, filter_func=filter_telephone_1)) -filters.append(Filter(class_type="ELECTRONIC", process_func=process_electronic_1, filter_func=filter_electronic_1)) -filters.append(Filter(class_type="FRACTION", process_func=process_fraction_1, filter_func=filter_fraction_1)) -filters.append(Filter(class_type="ADDRESS", process_func=process_address_1, filter_func=filter_address_1)) +filters.append( + Filter(class_type="TELEPHONE", process_func=process_telephone_1, filter_func=filter_telephone_1) +) +filters.append( + Filter( + class_type="ELECTRONIC", process_func=process_electronic_1, filter_func=filter_electronic_1 + ) +) +filters.append( + Filter(class_type="FRACTION", process_func=process_fraction_1, filter_func=filter_fraction_1) +) +filters.append( + Filter(class_type="ADDRESS", process_func=process_address_1, filter_func=filter_address_1) +) filters.append(Filter(class_type=EOS_TYPE, process_func=lambda x: x, filter_func=lambda x: True)) @@ -313,8 +367,10 @@ def filter_loaded_data(data: List[Instance], verbose: bool = False) -> List[Inst def parse_args(): parser = ArgumentParser() - parser.add_argument("--input", help="input file path", type=str, default='./en_with_types/output-00001-of-00100') - parser.add_argument("--verbose", help="print filtered instances", action='store_true') + parser.add_argument( + "--input", help="input file path", type=str, default="./en_with_types/output-00001-of-00100" + ) + parser.add_argument("--verbose", help="print filtered instances", action="store_true") return parser.parse_args() diff --git a/fun_text_processing/inverse_text_normalization/ja/graph_utils.py b/fun_text_processing/inverse_text_normalization/ja/graph_utils.py index d645118f8..644a0aeec 100644 --- a/fun_text_processing/inverse_text_normalization/ja/graph_utils.py +++ b/fun_text_processing/inverse_text_normalization/ja/graph_utils.py @@ -35,7 +35,9 @@ delete_extra_space = pynini.cross(pynini.closure(DAMO_WHITE_SPACE, 1), " ") delete_hyphen = pynutil.delete(pynini.closure("-", 0, 1)) insert_hyphen = pynutil.insert("-") -TO_LOWER = pynini.union(*[pynini.cross(x, y) for x, y in zip(string.ascii_uppercase, string.ascii_lowercase)]) +TO_LOWER = pynini.union( + *[pynini.cross(x, y) for x, y in zip(string.ascii_uppercase, string.ascii_lowercase)] +) TO_UPPER = pynini.invert(TO_LOWER) @@ -65,7 +67,9 @@ def convert_space(fst) -> "pynini.FstLike": Returns output fst where breaking spaces are converted to non breaking spaces """ - return fst @ pynini.cdrewrite(pynini.cross(DAMO_SPACE, DAMO_NON_BREAKING_SPACE), "", "", DAMO_SIGMA) + return fst @ pynini.cdrewrite( + pynini.cross(DAMO_SPACE, DAMO_NON_BREAKING_SPACE), "", "", DAMO_SIGMA + ) class GraphFst: @@ -87,7 +91,9 @@ class GraphFst: self.far_path = Path(os.path.dirname(__file__) + "/grammars/" + kind + "/" + name + ".far") if self.far_exist(): - self._fst = Far(self.far_path, mode="r", arc_type="standard", far_type="default").get_fst() + self._fst = Far( + self.far_path, mode="r", arc_type="standard", far_type="default" + ).get_fst() def far_exist(self) -> bool: """ diff --git a/fun_text_processing/inverse_text_normalization/ja/taggers/cardinal.py b/fun_text_processing/inverse_text_normalization/ja/taggers/cardinal.py index 2b2af28f2..520309348 100644 --- a/fun_text_processing/inverse_text_normalization/ja/taggers/cardinal.py +++ b/fun_text_processing/inverse_text_normalization/ja/taggers/cardinal.py @@ -22,109 +22,153 @@ class CardinalFst(GraphFst): """ Finite state transducer for classifying cardinals e.g. minus twenty three -> cardinal { integer: "23" negative: "-" } } - Numbers below thirteen are not converted. + Numbers below thirteen are not converted. """ - def __init__(self, enable_standalone_number: bool = True, - enable_0_to_9: bool = True): + def __init__(self, enable_standalone_number: bool = True, enable_0_to_9: bool = True): super().__init__(name="cardinal", kind="classify") self.enable_standalone_number = enable_standalone_number self.enable_0_to_9 = enable_0_to_9 zero = string_file(get_abs_path("data/numbers/zero.tsv")) digit = string_file(get_abs_path("data/numbers/digit.tsv")) hundred_digit = string_file(get_abs_path("data/numbers/hundred_digit.tsv")) - sign = string_file(get_abs_path('data/numbers/sign.tsv')) - dot = string_file(get_abs_path('data/numbers/dot.tsv')) - ties = string_file(get_abs_path('data/numbers/ties.tsv')) + sign = string_file(get_abs_path("data/numbers/sign.tsv")) + dot = string_file(get_abs_path("data/numbers/dot.tsv")) + ties = string_file(get_abs_path("data/numbers/ties.tsv")) graph_teen = string_file(get_abs_path("data/numbers/teen.tsv")) - addzero = insert('0') - digits = zero | digit # 0 ~ 9 + addzero = insert("0") + digits = zero | digit # 0 ~ 9 teen = graph_teen - teen |= cross('十', '1') + (digit | addzero) + teen |= cross("十", "1") + (digit | addzero) tens = ties + addzero | (ties + (digit | addzero)) - hundred = (digit + delete('百') + (tens - | teen - | add_weight(zero + digit, 0.1) - | add_weight(digit + addzero, 0.5) - | add_weight(addzero**2, 1.0))) - hundred |= (cross("百", "1") + (tens - | teen - | add_weight(zero + digit, 0.1) - | add_weight(digit + addzero, 0.5) - | add_weight(addzero**2, 1.0))) + hundred = ( + digit + + delete("百") + + ( + tens + | teen + | add_weight(zero + digit, 0.1) + | add_weight(digit + addzero, 0.5) + | add_weight(addzero**2, 1.0) + ) + ) + hundred |= cross("百", "1") + ( + tens + | teen + | add_weight(zero + digit, 0.1) + | add_weight(digit + addzero, 0.5) + | add_weight(addzero**2, 1.0) + ) hundred |= hundred_digit - thousand = ((hundred | teen | tens | digits) + delete('千') + ( - hundred - | add_weight(zero + tens, 0.1) - | add_weight(addzero + zero + digit, 0.5) - | add_weight(digit + addzero**2, 0.8) - | add_weight(addzero**3, 1.0))) - ten_thousand = ((thousand | hundred | teen | tens | digits) - + delete('万') - + (thousand - | add_weight(zero + hundred, 0.1) - | add_weight(addzero + zero + tens, 0.5) - | add_weight(addzero + addzero + zero + digit, 0.5) - | add_weight(digit + addzero**3, 0.8) - | add_weight(addzero**4, 1.0))) + thousand = ( + (hundred | teen | tens | digits) + + delete("千") + + ( + hundred + | add_weight(zero + tens, 0.1) + | add_weight(addzero + zero + digit, 0.5) + | add_weight(digit + addzero**2, 0.8) + | add_weight(addzero**3, 1.0) + ) + ) + ten_thousand = ( + (thousand | hundred | teen | tens | digits) + + delete("万") + + ( + thousand + | add_weight(zero + hundred, 0.1) + | add_weight(addzero + zero + tens, 0.5) + | add_weight(addzero + addzero + zero + digit, 0.5) + | add_weight(digit + addzero**3, 0.8) + | add_weight(addzero**4, 1.0) + ) + ) - hundred_thousand = ((ten_thousand | thousand | hundred | teen | tens | digits) - + delete('十万') - + (ten_thousand - | add_weight(zero + thousand, 0.1) - | add_weight(addzero + zero + hundred, 0.5) - | add_weight(addzero + addzero + zero + tens, 0.5) - | add_weight(addzero**3 + zero + digit, 0.5) - | add_weight(digit + addzero**4, 0.8) - | add_weight(addzero**5, 1.0))) + hundred_thousand = ( + (ten_thousand | thousand | hundred | teen | tens | digits) + + delete("十万") + + ( + ten_thousand + | add_weight(zero + thousand, 0.1) + | add_weight(addzero + zero + hundred, 0.5) + | add_weight(addzero + addzero + zero + tens, 0.5) + | add_weight(addzero**3 + zero + digit, 0.5) + | add_weight(digit + addzero**4, 0.8) + | add_weight(addzero**5, 1.0) + ) + ) - million = ((hundred_thousand | ten_thousand | thousand | hundred | teen | tens | digits) - + delete('百万') - + (hundred_thousand - | add_weight(zero + ten_thousand, 0.1) - | add_weight(addzero + zero + thousand, 0.5) - | add_weight(addzero + addzero + zero + hundred, 0.5) - | add_weight(addzero ** 3 + zero + tens, 0.5) - | add_weight(addzero ** 4 + zero + digit, - 0.5) - | add_weight(digit + addzero ** 5, 0.8) - | add_weight(addzero ** 6, 1.0))) + million = ( + (hundred_thousand | ten_thousand | thousand | hundred | teen | tens | digits) + + delete("百万") + + ( + hundred_thousand + | add_weight(zero + ten_thousand, 0.1) + | add_weight(addzero + zero + thousand, 0.5) + | add_weight(addzero + addzero + zero + hundred, 0.5) + | add_weight(addzero**3 + zero + tens, 0.5) + | add_weight(addzero**4 + zero + digit, 0.5) + | add_weight(digit + addzero**5, 0.8) + | add_weight(addzero**6, 1.0) + ) + ) # 1亿 - hundred_million = ((million | hundred_thousand | ten_thousand | thousand | hundred | teen | tens | digits) - + delete('億') - + (add_weight(zero + million, 0.1) - | add_weight(addzero + zero + hundred_thousand, 0.5) - | add_weight(addzero**2 + zero + ten_thousand, 0.5) - | add_weight(addzero**3 + zero + thousand, 0.5) - | add_weight(addzero**4 + hundred, 0.5) - | add_weight(addzero**5 + tens, 0.5) - | add_weight(addzero**6 + digit, 0.5) - | add_weight(digit + addzero**7, 0.8) - | add_weight(addzero**8, 1.0))) + hundred_million = ( + (million | hundred_thousand | ten_thousand | thousand | hundred | teen | tens | digits) + + delete("億") + + ( + add_weight(zero + million, 0.1) + | add_weight(addzero + zero + hundred_thousand, 0.5) + | add_weight(addzero**2 + zero + ten_thousand, 0.5) + | add_weight(addzero**3 + zero + thousand, 0.5) + | add_weight(addzero**4 + hundred, 0.5) + | add_weight(addzero**5 + tens, 0.5) + | add_weight(addzero**6 + digit, 0.5) + | add_weight(digit + addzero**7, 0.8) + | add_weight(addzero**8, 1.0) + ) + ) # 1兆 - hundred_billion = ((hundred_million | million | hundred_thousand | ten_thousand | thousand | hundred | teen | tens | digits) - + delete('兆') - + (add_weight(addzero**3 + zero + hundred_million, 0.1) - | add_weight(addzero**4 + zero + million, 0.5) - | add_weight(addzero**5 + zero + hundred_thousand, 0.5) - | add_weight(addzero**6 + zero + ten_thousand, 0.5) - | add_weight(addzero**7 + zero + thousand, 0.5) - | add_weight(addzero**8 + hundred, 0.5) - | add_weight(addzero**9 + tens, 0.5) - | add_weight(addzero**10 + digit, 0.5) - | add_weight(digit + addzero**11, 0.8) - | add_weight(addzero**12, 1.0))) + hundred_billion = ( + ( + hundred_million + | million + | hundred_thousand + | ten_thousand + | thousand + | hundred + | teen + | tens + | digits + ) + + delete("兆") + + ( + add_weight(addzero**3 + zero + hundred_million, 0.1) + | add_weight(addzero**4 + zero + million, 0.5) + | add_weight(addzero**5 + zero + hundred_thousand, 0.5) + | add_weight(addzero**6 + zero + ten_thousand, 0.5) + | add_weight(addzero**7 + zero + thousand, 0.5) + | add_weight(addzero**8 + hundred, 0.5) + | add_weight(addzero**9 + tens, 0.5) + | add_weight(addzero**10 + digit, 0.5) + | add_weight(digit + addzero**11, 0.8) + | add_weight(addzero**12, 1.0) + ) + ) # 1.11, 1.01 - number = digits | teen | tens | hundred | thousand | ten_thousand | hundred_thousand | million + number = ( + digits | teen | tens | hundred | thousand | ten_thousand | hundred_thousand | million + ) # number = digits | teen | tens | hundred | thousand | ten_thousand | hundred_thousand | million | hundred_million | hundred_billion # 兆/亿 - number = ((number + accep('兆') + delete('零').ques).ques + - (number + accep('億') + delete('零').ques).ques + number - | (number + accep('兆') + delete('〇').ques).ques + - (number + accep('億') + delete('〇').ques).ques + number) + number = (number + accep("兆") + delete("零").ques).ques + ( + number + accep("億") + delete("零").ques + ).ques + number | (number + accep("兆") + delete("〇").ques).ques + ( + number + accep("億") + delete("〇").ques + ).ques + number number = sign.ques + number + (dot + digits.plus).ques self.number = number.optimize() @@ -133,27 +177,40 @@ class CardinalFst(GraphFst): # cardinal string like 127.0.0.1, used in ID, IP, etc. cardinal = digit.plus + (dot + digits.plus).plus # float number like 1.11 - cardinal |= (number + dot + digits.plus) + cardinal |= number + dot + digits.plus # cardinal string like 110 or 12306 or 13125617878, used in phone - cardinal |= (digits**3 | digits**5 | digits**10 | digits**11 | digits**12) + cardinal |= digits**3 | digits**5 | digits**10 | digits**11 | digits**12 # cardinal string like 23 if self.enable_standalone_number: if self.enable_0_to_9: cardinal |= number else: - number_two_plus = (digits + digits.plus) | teen | tens | hundred | thousand | ten_thousand | hundred_thousand | million | hundred_million | hundred_billion + number_two_plus = ( + (digits + digits.plus) + | teen + | tens + | hundred + | thousand + | ten_thousand + | hundred_thousand + | million + | hundred_million + | hundred_billion + ) cardinal |= number_two_plus - labels_exception = [''] + labels_exception = [""] graph_exception = pynini.union(*labels_exception) self.graph_no_exception = cardinal self.graph = (pynini.project(cardinal, "input") - graph_exception.arcsort()) @ cardinal optional_minus_graph = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("マイナス", "\"-\"") + DAMO_SPACE, 0, 1 + pynutil.insert("negative: ") + pynini.cross("マイナス", '"-"') + DAMO_SPACE, 0, 1 ) - final_graph = optional_minus_graph + pynutil.insert("integer: \"") + self.graph + pynutil.insert("\"") + final_graph = ( + optional_minus_graph + pynutil.insert('integer: "') + self.graph + pynutil.insert('"') + ) final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() @@ -170,11 +227,16 @@ class CardinalFst(GraphFst): graph_teen | pynutil.insert("00"), (ties | pynutil.insert("0")) + (digit | pynutil.insert("0")), ) - graph_hundred_component = graph_hundred_component | graph_a_hundred_digit_component | graph_one_hundred_component | graph_hundred_ties_component + graph_hundred_component = ( + graph_hundred_component + | graph_a_hundred_digit_component + | graph_one_hundred_component + | graph_hundred_ties_component + ) # graph_hundred_component_at_least_one_none_zero_digit = graph_hundred_component @ ( pynini.closure(DAMO_DIGIT) + (DAMO_DIGIT - "0") + pynini.closure(DAMO_DIGIT) ) self.graph_hundred_component_at_least_one_none_zero_digit = ( graph_hundred_component_at_least_one_none_zero_digit - ) \ No newline at end of file + ) diff --git a/fun_text_processing/inverse_text_normalization/ja/taggers/date.py b/fun_text_processing/inverse_text_normalization/ja/taggers/date.py index 2c2db4029..37689be63 100644 --- a/fun_text_processing/inverse_text_normalization/ja/taggers/date.py +++ b/fun_text_processing/inverse_text_normalization/ja/taggers/date.py @@ -62,7 +62,9 @@ def _get_year_graph(): def _get_thousands_graph(): graph_ties = _get_ties_graph() - graph_hundred_component = (graph_digit + delete_space + pynutil.delete("百")) | pynutil.insert("0") + graph_hundred_component = ( + graph_digit + delete_space + pynutil.delete("百") + ) | pynutil.insert("0") graph_hundred_component |= (pynini.cross("百", "1")) | pynutil.insert("0") graph = ( graph_digit @@ -98,7 +100,7 @@ def _get_year_graph(): class DateFst(GraphFst): """ - Finite state transducer for classifying date, + Finite state transducer for classifying date, e.g. january fifth twenty twelve -> date { month: "january" day: "5" year: "2012" preserve_order: true } e.g. the fifth of january twenty twelve -> date { day: "5" month: "january" year: "2012" preserve_order: true } e.g. twenty twenty -> date { year: "2012" preserve_order: true } @@ -116,18 +118,26 @@ class DateFst(GraphFst): year_graph = pynutil.add_weight(year_graph, YEAR_WEIGHT) month_graph = _get_month_graph() - month_graph = pynutil.insert("month: \"") + month_graph + pynutil.insert("\"") + month_graph = pynutil.insert('month: "') + month_graph + pynutil.insert('"') - day_graph = pynutil.insert("day: \"") + pynutil.add_weight(ordinal_graph, -0.7) + pynutil.insert("\"") + day_graph = ( + pynutil.insert('day: "') + pynutil.add_weight(ordinal_graph, -0.7) + pynutil.insert('"') + ) graph_year = ( delete_extra_space - + pynutil.insert("year: \"") + + pynutil.insert('year: "') + pynutil.add_weight(year_graph, -YEAR_WEIGHT) - + pynutil.insert("\"") + + pynutil.insert('"') + ) + optional_graph_year = pynini.closure( + graph_year, + 0, + 1, ) - optional_graph_year = pynini.closure(graph_year, 0, 1,) graph_mdy = month_graph + ( - (delete_extra_space + day_graph) | graph_year | (delete_extra_space + day_graph + graph_year) + (delete_extra_space + day_graph) + | graph_year + | (delete_extra_space + day_graph + graph_year) ) graph_dmy = ( pynutil.delete("the") @@ -139,7 +149,9 @@ class DateFst(GraphFst): + month_graph + optional_graph_year ) - graph_year = pynutil.insert("year: \"") + (year_graph | _get_range_graph()) + pynutil.insert("\"") + graph_year = ( + pynutil.insert('year: "') + (year_graph | _get_range_graph()) + pynutil.insert('"') + ) final_graph = graph_mdy | graph_dmy | graph_year final_graph += pynutil.insert(" preserve_order: true") diff --git a/fun_text_processing/inverse_text_normalization/ja/taggers/decimal.py b/fun_text_processing/inverse_text_normalization/ja/taggers/decimal.py index 5ae9e8bc7..1e7919651 100644 --- a/fun_text_processing/inverse_text_normalization/ja/taggers/decimal.py +++ b/fun_text_processing/inverse_text_normalization/ja/taggers/decimal.py @@ -9,30 +9,42 @@ from fun_text_processing.inverse_text_normalization.ja.graph_utils import ( from pynini.lib import pynutil -def get_quantity(decimal: 'pynini.FstLike', cardinal_up_to_hundred: 'pynini.FstLike') -> 'pynini.FstLike': +def get_quantity( + decimal: "pynini.FstLike", cardinal_up_to_hundred: "pynini.FstLike" +) -> "pynini.FstLike": """ Returns FST that transforms either a cardinal or decimal followed by a quantity into a numeral, e.g. one million -> integer_part: "1" quantity: "million" e.g. one point five million -> integer_part: "1" fractional_part: "5" quantity: "million" - Args: + Args: decimal: decimal FST cardinal_up_to_hundred: cardinal FST """ numbers = cardinal_up_to_hundred @ ( - pynutil.delete(pynini.closure("0")) + pynini.difference(DAMO_DIGIT, "0") + pynini.closure(DAMO_DIGIT) + pynutil.delete(pynini.closure("0")) + + pynini.difference(DAMO_DIGIT, "0") + + pynini.closure(DAMO_DIGIT) + ) + suffix = pynini.union( + "万", "百万", "千万", "億" "十億", "trillion", "quadrillion", "quintillion", "sextillion" ) - suffix = pynini.union("万", "百万", "千万", "億" "十億", "trillion", "quadrillion", "quintillion", "sextillion") res = ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + numbers - + pynutil.insert("\"") + + pynutil.insert('"') + delete_extra_space - + pynutil.insert("quantity: \"") + + pynutil.insert('quantity: "') + suffix - + pynutil.insert("\"") + + pynutil.insert('"') + ) + res |= ( + decimal + + delete_extra_space + + pynutil.insert('quantity: "') + + (suffix | "千") + + pynutil.insert('"') ) - res |= decimal + delete_extra_space + pynutil.insert("quantity: \"") + (suffix | "千") + pynutil.insert("\"") return res @@ -51,7 +63,11 @@ class DecimalFst(GraphFst): cardinal_graph = cardinal.graph_no_exception graph_decimal = pynini.string_file(get_abs_path("data/numbers/digit.tsv")) - graph_decimal |= pynini.string_file(get_abs_path("data/numbers/zero.tsv")) | pynini.cross("零", "０") | pynini.cross("〇", "０") + graph_decimal |= ( + pynini.string_file(get_abs_path("data/numbers/zero.tsv")) + | pynini.cross("零", "０") + | pynini.cross("〇", "０") + ) graph_decimal = pynini.closure(graph_decimal + delete_space) + graph_decimal self.graph = graph_decimal @@ -59,13 +75,20 @@ class DecimalFst(GraphFst): point = pynutil.delete("点") optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("マイナス", "\"true\"") + delete_extra_space, 0, 1 + pynutil.insert("negative: ") + pynini.cross("マイナス", '"true"') + delete_extra_space, + 0, + 1, ) - graph_fractional = pynutil.insert("fractional_part: \"") + graph_decimal + pynutil.insert("\"") - graph_integer = pynutil.insert("integer_part: \"") + cardinal_graph + pynutil.insert("\"") + graph_fractional = ( + pynutil.insert('fractional_part: "') + graph_decimal + pynutil.insert('"') + ) + graph_integer = pynutil.insert('integer_part: "') + cardinal_graph + pynutil.insert('"') final_graph_wo_sign = ( - pynini.closure(graph_integer + delete_extra_space, 0, 1) + point + delete_extra_space + graph_fractional + pynini.closure(graph_integer + delete_extra_space, 0, 1) + + point + + delete_extra_space + + graph_fractional ) final_graph = optional_graph_negative + final_graph_wo_sign diff --git a/fun_text_processing/inverse_text_normalization/ja/taggers/electronic.py b/fun_text_processing/inverse_text_normalization/ja/taggers/electronic.py index 88c7d9c33..dfd13a691 100644 --- a/fun_text_processing/inverse_text_normalization/ja/taggers/electronic.py +++ b/fun_text_processing/inverse_text_normalization/ja/taggers/electronic.py @@ -1,6 +1,10 @@ import pynini from fun_text_processing.inverse_text_normalization.ja.utils import get_abs_path -from fun_text_processing.inverse_text_normalization.ja.graph_utils import DAMO_ALPHA, GraphFst, insert_space +from fun_text_processing.inverse_text_normalization.ja.graph_utils import ( + DAMO_ALPHA, + GraphFst, + insert_space, +) from pynini.lib import pynutil @@ -24,35 +28,52 @@ class ElectronicFst(GraphFst): accepted_username = alpha_num | symbols process_dot = pynini.cross("ドット", ".") - username = (alpha_num + pynini.closure(delete_extra_space + accepted_username)) | pynutil.add_weight( - pynini.closure(DAMO_ALPHA, 1), weight=0.0001 - ) - username = pynutil.insert("username: \"") + username + pynutil.insert("\"") + username = ( + alpha_num + pynini.closure(delete_extra_space + accepted_username) + ) | pynutil.add_weight(pynini.closure(DAMO_ALPHA, 1), weight=0.0001) + username = pynutil.insert('username: "') + username + pynutil.insert('"') single_alphanum = pynini.closure(alpha_num + delete_extra_space) + alpha_num - server = single_alphanum | pynini.string_file(get_abs_path("data/electronic/server_name.tsv")) + server = single_alphanum | pynini.string_file( + get_abs_path("data/electronic/server_name.tsv") + ) domain = single_alphanum | pynini.string_file(get_abs_path("data/electronic/domain.tsv")) domain_graph = ( - pynutil.insert("domain: \"") + pynutil.insert('domain: "') + server + delete_extra_space + process_dot + delete_extra_space + domain - + pynutil.insert("\"") + + pynutil.insert('"') + ) + graph = ( + username + + delete_extra_space + + pynutil.delete("at") + + insert_space + + delete_extra_space + + domain_graph ) - graph = username + delete_extra_space + pynutil.delete("at") + insert_space + delete_extra_space + domain_graph ############# url ### protocol_end = pynini.cross(pynini.union("w w w", "www"), "www") - protocol_start = (pynini.cross("h t t p", "http") | pynini.cross("h t t p s", "https")) + pynini.cross( - "コロンスラッシュスラッシュ", "://" # colon slash slash + protocol_start = ( + pynini.cross("h t t p", "http") | pynini.cross("h t t p s", "https") + ) + pynini.cross( + "コロンスラッシュスラッシュ", "://" # colon slash slash ) # .com, ending = ( delete_extra_space + symbols + delete_extra_space - + (domain | pynini.closure(accepted_username + delete_extra_space,) + accepted_username) + + ( + domain + | pynini.closure( + accepted_username + delete_extra_space, + ) + + accepted_username + ) ) protocol_default = ( @@ -63,12 +84,18 @@ class ElectronicFst(GraphFst): + pynini.closure(ending, 1) ).optimize() protocol = ( - pynini.closure(protocol_start, 0, 1) + protocol_end + delete_extra_space + process_dot + protocol_default + pynini.closure(protocol_start, 0, 1) + + protocol_end + + delete_extra_space + + process_dot + + protocol_default ).optimize() - protocol |= pynini.closure(protocol_end + delete_extra_space + process_dot, 0, 1) + protocol_default + protocol |= ( + pynini.closure(protocol_end + delete_extra_space + process_dot, 0, 1) + protocol_default + ) - protocol = pynutil.insert("protocol: \"") + protocol.optimize() + pynutil.insert("\"") + protocol = pynutil.insert('protocol: "') + protocol.optimize() + pynutil.insert('"') graph |= protocol ######## diff --git a/fun_text_processing/inverse_text_normalization/ja/taggers/fraction.py b/fun_text_processing/inverse_text_normalization/ja/taggers/fraction.py index 2e66a96f2..5ae4a00fe 100644 --- a/fun_text_processing/inverse_text_normalization/ja/taggers/fraction.py +++ b/fun_text_processing/inverse_text_normalization/ja/taggers/fraction.py @@ -1,5 +1,11 @@ import pynini -from fun_text_processing.inverse_text_normalization.ja.graph_utils import GraphFst, delete_extra_space, delete_space, insert_space, DAMO_CHAR +from fun_text_processing.inverse_text_normalization.ja.graph_utils import ( + GraphFst, + delete_extra_space, + delete_space, + insert_space, + DAMO_CHAR, +) from pynini.lib import pynutil @@ -18,9 +24,11 @@ class FractionFst(GraphFst): super().__init__(name="fraction", kind="classify") graph_cardinal = cardinal.graph_no_exception - graph_four = pynini.cross("クォーター", "4") # quarter + graph_four = pynini.cross("クォーター", "4") # quarter - denominator = pynutil.insert('denominator: "') + (graph_cardinal | graph_four) + pynutil.insert('"') + denominator = ( + pynutil.insert('denominator: "') + (graph_cardinal | graph_four) + pynutil.insert('"') + ) fraction_component = pynutil.delete(pynini.union("分の", "割る")) numerator = pynutil.insert('numerator: "') + graph_cardinal + pynutil.insert('"') @@ -40,4 +48,4 @@ class FractionFst(GraphFst): graph = optional_graph_negative + graph final_graph = self.add_tokens(graph) - self.fst = final_graph.optimize() \ No newline at end of file + self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ja/taggers/measure.py b/fun_text_processing/inverse_text_normalization/ja/taggers/measure.py index a81204a30..8b0be1f27 100644 --- a/fun_text_processing/inverse_text_normalization/ja/taggers/measure.py +++ b/fun_text_processing/inverse_text_normalization/ja/taggers/measure.py @@ -31,7 +31,9 @@ class MeasureFst(GraphFst): graph_unit_plural = graph_unit_singular # plural -> abbr optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("マイナス", "\"true\"") + delete_extra_space, 0, 1 + pynutil.insert("negative: ") + pynini.cross("マイナス", '"true"') + delete_extra_space, + 0, + 1, ) unit_singular = graph_unit_singular @@ -39,15 +41,23 @@ class MeasureFst(GraphFst): unit_misc = pynutil.insert("/") + pynutil.delete("每") + delete_space + graph_unit_singular unit_singular = ( - pynutil.insert("units: \"") - + (unit_singular | unit_misc | pynutil.add_weight(unit_singular + delete_space + unit_misc, 0.01)) - + pynutil.insert("\"") + pynutil.insert('units: "') + + ( + unit_singular + | unit_misc + | pynutil.add_weight(unit_singular + delete_space + unit_misc, 0.01) + ) + + pynutil.insert('"') ) unit_plural = ( - pynutil.insert("units: \"") - + (unit_plural | unit_misc | pynutil.add_weight(unit_plural + delete_space + unit_misc, 0.01)) - + pynutil.insert("\"") + pynutil.insert('units: "') + + ( + unit_plural + | unit_misc + | pynutil.add_weight(unit_plural + delete_space + unit_misc, 0.01) + ) + + pynutil.insert('"') ) subgraph_decimal = ( @@ -69,9 +79,9 @@ class MeasureFst(GraphFst): subgraph_cardinal = ( pynutil.insert("cardinal { ") + optional_graph_negative - + pynutil.insert("integer: \"") + + pynutil.insert('integer: "') + ((DAMO_SIGMA - "一") @ cardinal_graph) - + pynutil.insert("\"") + + pynutil.insert('"') + pynutil.insert(" }") + delete_extra_space + unit_plural @@ -79,21 +89,21 @@ class MeasureFst(GraphFst): subgraph_cardinal |= ( pynutil.insert("cardinal { ") + optional_graph_negative - + pynutil.insert("integer: \"") + + pynutil.insert('integer: "') + pynini.cross("一", "1") - + pynutil.insert("\"") + + pynutil.insert('"') + pynutil.insert(" }") + delete_extra_space + unit_singular ) subgraph_cardinal |= ( - pynutil.insert("cardinal { ") - + optional_graph_negative - + pynutil.insert("integer: \"") - + ((DAMO_SIGMA - "一") @ cardinal_graph) - + pynutil.insert("\"") - + pynutil.insert(" }") - + unit_singular + pynutil.insert("cardinal { ") + + optional_graph_negative + + pynutil.insert('integer: "') + + ((DAMO_SIGMA - "一") @ cardinal_graph) + + pynutil.insert('"') + + pynutil.insert(" }") + + unit_singular ) final_graph = subgraph_decimal | subgraph_cardinal final_graph = self.add_tokens(final_graph) diff --git a/fun_text_processing/inverse_text_normalization/ja/taggers/money.py b/fun_text_processing/inverse_text_normalization/ja/taggers/money.py index c2351538f..63e2c1458 100644 --- a/fun_text_processing/inverse_text_normalization/ja/taggers/money.py +++ b/fun_text_processing/inverse_text_normalization/ja/taggers/money.py @@ -33,7 +33,7 @@ class MoneyFst(GraphFst): # "one fifty" -> "one hundred fifty" with_hundred = pynini.compose( pynini.closure(DAMO_NOT_SPACE) + pynini.accep(" ") + pynutil.insert("百") + DAMO_SIGMA, - pynini.compose(cardinal_graph, DAMO_DIGIT ** 3), + pynini.compose(cardinal_graph, DAMO_DIGIT**3), ) cardinal_graph |= with_hundred graph_decimal_final = decimal.final_graph_wo_negative @@ -43,25 +43,32 @@ class MoneyFst(GraphFst): unit_plural = unit_singular # unit_plural = get_singulars(unit_singular) - graph_unit_singular = pynutil.insert("currency: \"") + convert_space(unit_singular) + pynutil.insert("\"") - graph_unit_plural = pynutil.insert("currency: \"") + convert_space(unit_plural) + pynutil.insert("\"") + graph_unit_singular = ( + pynutil.insert('currency: "') + convert_space(unit_singular) + pynutil.insert('"') + ) + graph_unit_plural = ( + pynutil.insert('currency: "') + convert_space(unit_plural) + pynutil.insert('"') + ) - add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | (pynutil.insert("0") + DAMO_DIGIT) + add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | ( + pynutil.insert("0") + DAMO_DIGIT + ) # twelve dollars (and) fifty cents, zero cents cents_standalone = ( - pynutil.insert("fractional_part: \"") + pynutil.insert('fractional_part: "') + pynini.union( - pynutil.add_weight(((DAMO_SIGMA - "一") @ cardinal_graph), -0.7) @ add_leading_zero_to_double_digit + pynutil.add_weight(((DAMO_SIGMA - "一") @ cardinal_graph), -0.7) + @ add_leading_zero_to_double_digit + delete_space - + pynutil.delete("セント"), #cent - pynini.cross("一", "01") + delete_space + pynutil.delete("セント"), #cent + + pynutil.delete("セント"), # cent + pynini.cross("一", "01") + delete_space + pynutil.delete("セント"), # cent ) - + pynutil.insert("\"") + + pynutil.insert('"') ) optional_cents_standalone = pynini.closure( delete_space - + pynini.closure(pynutil.delete("と") + delete_space, 0, 1) # and + + pynini.closure(pynutil.delete("と") + delete_space, 0, 1) # and + insert_space + cents_standalone, 0, @@ -70,31 +77,31 @@ class MoneyFst(GraphFst): # twelve dollars fifty, only after integer optional_cents_suffix = pynini.closure( delete_extra_space - + pynutil.insert("fractional_part: \"") + + pynutil.insert('fractional_part: "') + pynutil.add_weight(cardinal_graph @ add_leading_zero_to_double_digit, -0.7) - + pynutil.insert("\""), + + pynutil.insert('"'), 0, 1, ) graph_integer = ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + ((DAMO_SIGMA - "一") @ cardinal_graph) - + pynutil.insert("\"") + + pynutil.insert('"') + delete_extra_space + graph_unit_plural + (optional_cents_standalone | optional_cents_suffix) ) graph_integer |= ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + pynini.cross("一", "1") - + pynutil.insert("\"") + + pynutil.insert('"') + delete_extra_space + graph_unit_singular + (optional_cents_standalone | optional_cents_suffix) ) graph_decimal = graph_decimal_final + delete_extra_space + graph_unit_plural - graph_decimal |= pynutil.insert("currency: \"$\" integer_part: \"0\" ") + cents_standalone + graph_decimal |= pynutil.insert('currency: "$" integer_part: "0" ') + cents_standalone final_graph = graph_integer | graph_decimal final_graph = self.add_tokens(final_graph) diff --git a/fun_text_processing/inverse_text_normalization/ja/taggers/ordinal.py b/fun_text_processing/inverse_text_normalization/ja/taggers/ordinal.py index 3dd6fcb9b..6803811ce 100644 --- a/fun_text_processing/inverse_text_normalization/ja/taggers/ordinal.py +++ b/fun_text_processing/inverse_text_normalization/ja/taggers/ordinal.py @@ -24,24 +24,32 @@ class OrdinalFst(GraphFst): teen = pynini.string_file(get_abs_path("data/ordinals/teen.tsv")) zero = pynini.string_file(get_abs_path("data/numbers/zero.tsv")) hundred_digit = pynini.string_file(get_abs_path("data/numbers/hundred_digit.tsv")) - addzero = insert('0') - tens = ties + addzero | (digit + delete('十') + (digit | addzero)) - hundred = (digit + delete('百') + (tens - | teen - | add_weight(zero + digit, 0.1) - | add_weight(digit + addzero, 0.5) - | add_weight(addzero**2, 1.0))) - hundred |= (cross("百", "1") + (tens - | teen - | add_weight(zero + digit, 0.1) - | add_weight(digit + addzero, 0.5) - | add_weight(addzero**2, 1.0))) + addzero = insert("0") + tens = ties + addzero | (digit + delete("十") + (digit | addzero)) + hundred = ( + digit + + delete("百") + + ( + tens + | teen + | add_weight(zero + digit, 0.1) + | add_weight(digit + addzero, 0.5) + | add_weight(addzero**2, 1.0) + ) + ) + hundred |= cross("百", "1") + ( + tens + | teen + | add_weight(zero + digit, 0.1) + | add_weight(digit + addzero, 0.5) + | add_weight(addzero**2, 1.0) + ) hundred |= hundred_digit ordinal = digit | teen | tens | hundred graph = pynini.closure(DAMO_CHAR, 1) + ordinal self.graph = graph @ cardinal_graph - final_graph = pynutil.insert("integer: \"") + self.graph + pynutil.insert("\"") + final_graph = pynutil.insert('integer: "') + self.graph + pynutil.insert('"') final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ja/taggers/preprocessor.py b/fun_text_processing/inverse_text_normalization/ja/taggers/preprocessor.py index b1015001f..2c66ed3bc 100644 --- a/fun_text_processing/inverse_text_normalization/ja/taggers/preprocessor.py +++ b/fun_text_processing/inverse_text_normalization/ja/taggers/preprocessor.py @@ -6,14 +6,17 @@ from pynini.lib import pynutil class PreProcessor(GraphFst): def __init__( - self, halfwidth_to_fullwidth: bool = True, + self, + halfwidth_to_fullwidth: bool = True, ): super().__init__(name="PreProcessor", kind="processor") - graph = pynini.cdrewrite('', '', '', DAMO_SIGMA) + graph = pynini.cdrewrite("", "", "", DAMO_SIGMA) if halfwidth_to_fullwidth: - halfwidth_to_fullwidth_graph = pynini.string_file(get_abs_path('data/char/halfwidth_to_fullwidth.tsv')) - graph @= pynini.cdrewrite(halfwidth_to_fullwidth_graph, '', '', DAMO_SIGMA) + halfwidth_to_fullwidth_graph = pynini.string_file( + get_abs_path("data/char/halfwidth_to_fullwidth.tsv") + ) + graph @= pynini.cdrewrite(halfwidth_to_fullwidth_graph, "", "", DAMO_SIGMA) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ja/taggers/punctuation.py b/fun_text_processing/inverse_text_normalization/ja/taggers/punctuation.py index 4a9a03d5a..7c79ac877 100644 --- a/fun_text_processing/inverse_text_normalization/ja/taggers/punctuation.py +++ b/fun_text_processing/inverse_text_normalization/ja/taggers/punctuation.py @@ -12,9 +12,9 @@ class PunctuationFst(GraphFst): def __init__(self): super().__init__(name="punctuation", kind="classify") - s = "!#$%&\'()*+,-./:;<=>?@^_`{|}~、。，！【】「」《》￥（）——・" + s = "!#$%&'()*+,-./:;<=>?@^_`{|}~、。，！【】「」《》￥（）——・" punct = pynini.union(*s) - graph = pynutil.insert("name: \"") + punct + pynutil.insert("\"") + graph = pynutil.insert('name: "') + punct + pynutil.insert('"') self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ja/taggers/telephone.py b/fun_text_processing/inverse_text_normalization/ja/taggers/telephone.py index 51174bcd5..192ff18cb 100644 --- a/fun_text_processing/inverse_text_normalization/ja/taggers/telephone.py +++ b/fun_text_processing/inverse_text_normalization/ja/taggers/telephone.py @@ -23,7 +23,7 @@ def get_serial_number(cardinal): class TelephoneFst(GraphFst): """ - Finite state transducer for classifying telephone numbers, e.g. + Finite state transducer for classifying telephone numbers, e.g. one two three one two three five six seven eight -> { number_part: "123-123-5678" } This class also support card number and IP format. @@ -61,61 +61,89 @@ class TelephoneFst(GraphFst): double_digit.invert() # to handle cases like "one twenty three" - two_digit_cardinal = pynini.compose(cardinal.graph_no_exception, DAMO_DIGIT ** 2) + two_digit_cardinal = pynini.compose(cardinal.graph_no_exception, DAMO_DIGIT**2) double_digit_to_digit = ( - pynini.compose(double_digit, str_to_digit + pynutil.delete(" ") + str_to_digit) | two_digit_cardinal + pynini.compose(double_digit, str_to_digit + pynutil.delete(" ") + str_to_digit) + | two_digit_cardinal ) - single_or_double_digit = (pynutil.add_weight(double_digit_to_digit, -0.0001) | str_to_digit).optimize() + single_or_double_digit = ( + pynutil.add_weight(double_digit_to_digit, -0.0001) | str_to_digit + ).optimize() single_or_double_digit |= ( single_or_double_digit - + pynini.closure(pynutil.add_weight(pynutil.delete(" ") + single_or_double_digit, 0.0001)) + + pynini.closure( + pynutil.add_weight(pynutil.delete(" ") + single_or_double_digit, 0.0001) + ) ).optimize() number_part = pynini.compose( single_or_double_digit, - DAMO_DIGIT ** 3 + pynutil.insert("-") + DAMO_DIGIT ** 3 + pynutil.insert("-") + DAMO_DIGIT ** 4, + DAMO_DIGIT**3 + + pynutil.insert("-") + + DAMO_DIGIT**3 + + pynutil.insert("-") + + DAMO_DIGIT**4, ).optimize() - number_part = pynutil.insert("number_part: \"") + number_part.optimize() + pynutil.insert("\"") + number_part = ( + pynutil.insert('number_part: "') + number_part.optimize() + pynutil.insert('"') + ) cardinal_option = pynini.compose(single_or_double_digit, DAMO_DIGIT ** (2, 3)) country_code = ( - pynutil.insert("country_code: \"") - + pynini.closure(pynini.cross("プラス ", "+"), 0, 1) #plus - + ((pynini.closure(str_to_digit + pynutil.delete(" "), 0, 2) + str_to_digit) | cardinal_option) - + pynutil.insert("\"") + pynutil.insert('country_code: "') + + pynini.closure(pynini.cross("プラス ", "+"), 0, 1) # plus + + ( + (pynini.closure(str_to_digit + pynutil.delete(" "), 0, 2) + str_to_digit) + | cardinal_option + ) + + pynutil.insert('"') ) - optional_country_code = pynini.closure(country_code + pynutil.delete(" ") + insert_space, 0, 1).optimize() + optional_country_code = pynini.closure( + country_code + pynutil.delete(" ") + insert_space, 0, 1 + ).optimize() graph = optional_country_code + number_part # credit card number - space_four_digits = insert_space + DAMO_DIGIT ** 4 - credit_card_graph = pynini.compose(single_or_double_digit, DAMO_DIGIT ** 4 + space_four_digits ** 3).optimize() - graph |= pynutil.insert("number_part: \"") + credit_card_graph.optimize() + pynutil.insert("\"") + space_four_digits = insert_space + DAMO_DIGIT**4 + credit_card_graph = pynini.compose( + single_or_double_digit, DAMO_DIGIT**4 + space_four_digits**3 + ).optimize() + graph |= ( + pynutil.insert('number_part: "') + credit_card_graph.optimize() + pynutil.insert('"') + ) # SSN ssn_graph = pynini.compose( single_or_double_digit, - DAMO_DIGIT ** 3 + pynutil.insert("-") + DAMO_DIGIT ** 2 + pynutil.insert("-") + DAMO_DIGIT ** 4, + DAMO_DIGIT**3 + + pynutil.insert("-") + + DAMO_DIGIT**2 + + pynutil.insert("-") + + DAMO_DIGIT**4, ).optimize() - graph |= pynutil.insert("number_part: \"") + ssn_graph.optimize() + pynutil.insert("\"") + graph |= pynutil.insert('number_part: "') + ssn_graph.optimize() + pynutil.insert('"') # ip - digit_or_double = pynini.closure(str_to_digit + pynutil.delete(" "), 0, 1) + double_digit_to_digit - digit_or_double |= double_digit_to_digit + pynini.closure(pynutil.delete(" ") + str_to_digit, 0, 1) + digit_or_double = ( + pynini.closure(str_to_digit + pynutil.delete(" "), 0, 1) + double_digit_to_digit + ) + digit_or_double |= double_digit_to_digit + pynini.closure( + pynutil.delete(" ") + str_to_digit, 0, 1 + ) digit_or_double |= str_to_digit + (pynutil.delete(" ") + str_to_digit) ** (0, 2) digit_or_double |= cardinal_option digit_or_double = digit_or_double.optimize() - ip_graph = digit_or_double + (pynini.cross(" 点 ", ".") + digit_or_double) ** 3 #dot + ip_graph = digit_or_double + (pynini.cross(" 点 ", ".") + digit_or_double) ** 3 # dot - graph |= pynutil.insert("number_part: \"") + ip_graph.optimize() + pynutil.insert("\"") + graph |= pynutil.insert('number_part: "') + ip_graph.optimize() + pynutil.insert('"') graph |= ( - pynutil.insert("number_part: \"") + pynutil.insert('number_part: "') + pynutil.add_weight(get_serial_number(cardinal=cardinal), weight=0.0001) - + pynutil.insert("\"") + + pynutil.insert('"') ) final_graph = self.add_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/ja/taggers/time.py b/fun_text_processing/inverse_text_normalization/ja/taggers/time.py index 883a39bdf..ee92f8e20 100644 --- a/fun_text_processing/inverse_text_normalization/ja/taggers/time.py +++ b/fun_text_processing/inverse_text_normalization/ja/taggers/time.py @@ -42,24 +42,28 @@ class TimeFst(GraphFst): graph_minute_single = pynini.union(*labels_minute_single) @ cardinal graph_minute_double = pynini.union(*labels_minute_double) @ cardinal - graph_minute_verbose = pynini.cross("半", "30") | pynini.cross("クォーター", "15") # half, quarter + graph_minute_verbose = pynini.cross("半", "30") | pynini.cross( + "クォーター", "15" + ) # half, quarter oclock = pynini.cross(pynini.union("時", "o' clock", "o clock", "o'clock", "oclock"), "") - final_graph_hour = pynutil.insert("hours: \"") + graph_hour + pynutil.insert("\"") + final_graph_hour = pynutil.insert('hours: "') + graph_hour + pynutil.insert('"') graph_minute = ( oclock + pynutil.insert("00") | pynutil.delete(pynini.union("〇", "零")) + delete_space + graph_minute_single | graph_minute_double ) - final_suffix = pynutil.insert("suffix: \"") + convert_space(suffix_graph) + pynutil.insert("\"") + final_suffix = ( + pynutil.insert('suffix: "') + convert_space(suffix_graph) + pynutil.insert('"') + ) final_suffix = delete_space + insert_space + final_suffix final_suffix_optional = pynini.closure(final_suffix, 0, 1) final_time_zone_optional = pynini.closure( delete_space + insert_space - + pynutil.insert("zone: \"") + + pynutil.insert('zone: "') + convert_space(time_zone_graph) - + pynutil.insert("\""), + + pynutil.insert('"'), 0, 1, ) @@ -68,54 +72,65 @@ class TimeFst(GraphFst): # two o eight, two thirty five (am/pm) # two pm/am graph_hm = ( - final_graph_hour + delete_extra_space + pynutil.insert("minutes: \"") + graph_minute + pynutil.insert("\"") + final_graph_hour + + delete_extra_space + + pynutil.insert('minutes: "') + + graph_minute + + pynutil.insert('"') ) # 10 past four, quarter past four, half past four graph_m_past_h = ( - pynutil.insert("minutes: \"") + pynutil.insert('minutes: "') + pynini.union(graph_minute_single, graph_minute_double, graph_minute_verbose) - + pynutil.insert("\"") + + pynutil.insert('"') + delete_extra_space + final_graph_hour ) graph_quarter_time = ( - pynutil.insert("minutes: \"") - + pynini.cross("クォーター", "45") # quarter - + pynutil.insert("\"") + pynutil.insert('minutes: "') + + pynini.cross("クォーター", "45") # quarter + + pynutil.insert('"') + delete_space - + pynutil.delete(pynini.union("から", "to", "till")) # to, till + + pynutil.delete(pynini.union("から", "to", "till")) # to, till + delete_extra_space - + pynutil.insert("hours: \"") + + pynutil.insert('hours: "') + to_hour_graph - + pynutil.insert("\"") + + pynutil.insert('"') ) graph_m_to_h_suffix_time = ( - pynutil.insert("minutes: \"") + pynutil.insert('minutes: "') + ((graph_minute_single | graph_minute_double).optimize() @ minute_to_graph) - + pynutil.insert("\"") - + pynini.closure(delete_space + pynutil.delete(pynini.union("分", "min", "mins", "minute", "minutes")), 0, 1) + + pynutil.insert('"') + + pynini.closure( + delete_space + + pynutil.delete(pynini.union("分", "min", "mins", "minute", "minutes")), + 0, + 1, + ) + delete_space - + pynutil.delete(pynini.union("から", "to", "till")) # to, till + + pynutil.delete(pynini.union("から", "to", "till")) # to, till + delete_extra_space - + pynutil.insert("hours: \"") + + pynutil.insert('hours: "') + to_hour_graph - + pynutil.insert("\"") + + pynutil.insert('"') + final_suffix ) graph_h = ( final_graph_hour + delete_extra_space - + pynutil.insert("minutes: \"") + + pynutil.insert('minutes: "') + (pynutil.insert("00") | graph_minute) - + pynutil.insert("\"") + + pynutil.insert('"') + final_suffix + final_time_zone_optional ) final_graph = ( - (graph_hm | graph_m_past_h | graph_quarter_time) + final_suffix_optional + final_time_zone_optional + (graph_hm | graph_m_past_h | graph_quarter_time) + + final_suffix_optional + + final_time_zone_optional ) final_graph |= graph_h final_graph |= graph_m_to_h_suffix_time diff --git a/fun_text_processing/inverse_text_normalization/ja/taggers/tokenize_and_classify.py b/fun_text_processing/inverse_text_normalization/ja/taggers/tokenize_and_classify.py index 68ab2114e..67bb0cda7 100644 --- a/fun_text_processing/inverse_text_normalization/ja/taggers/tokenize_and_classify.py +++ b/fun_text_processing/inverse_text_normalization/ja/taggers/tokenize_and_classify.py @@ -21,7 +21,7 @@ from fun_text_processing.inverse_text_normalization.ja.graph_utils import ( delete_extra_space, delete_space, generator_main, - insert_space + insert_space, ) from pynini.lib import pynutil @@ -31,7 +31,7 @@ import logging class ClassifyFst(GraphFst): """ Final class that composes all other classification grammars. This class can process an entire sentence, that is lower cased. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. Args: @@ -39,9 +39,13 @@ class ClassifyFst(GraphFst): overwrite_cache: set to True to overwrite .far files """ - def __init__(self, cache_dir: str = None, overwrite_cache: bool = False, + def __init__( + self, + cache_dir: str = None, + overwrite_cache: bool = False, enable_standalone_number: bool = True, - enable_0_to_9: bool = True): + enable_0_to_9: bool = True, + ): super().__init__(name="tokenize_and_classify", kind="classify") self.convert_number = enable_standalone_number self.enable_0_to_9 = enable_0_to_9 @@ -93,12 +97,20 @@ class ClassifyFst(GraphFst): | pynutil.add_weight(word_graph, 100) ) - punct = pynutil.insert("tokens { ") + pynutil.add_weight(punct_graph, weight=1.1) + pynutil.insert(" }") + punct = ( + pynutil.insert("tokens { ") + + pynutil.add_weight(punct_graph, weight=1.1) + + pynutil.insert(" }") + ) token = pynutil.insert("tokens { ") + classify + pynutil.insert(" }") token_plus_punct = ( - pynini.closure(punct + pynutil.insert(" ")) + token + pynini.closure(pynutil.insert(" ") + punct) + pynini.closure(punct + pynutil.insert(" ")) + + token + + pynini.closure(pynutil.insert(" ") + punct) + ) + graph = token_plus_punct + pynini.closure( + pynini.union(insert_space, delete_extra_space) + token_plus_punct ) - graph = token_plus_punct + pynini.closure(pynini.union(insert_space, delete_extra_space) + token_plus_punct) graph = delete_space + graph + delete_space diff --git a/fun_text_processing/inverse_text_normalization/ja/taggers/whitelist.py b/fun_text_processing/inverse_text_normalization/ja/taggers/whitelist.py index 9ef40193f..5bc040831 100644 --- a/fun_text_processing/inverse_text_normalization/ja/taggers/whitelist.py +++ b/fun_text_processing/inverse_text_normalization/ja/taggers/whitelist.py @@ -15,5 +15,5 @@ class WhiteListFst(GraphFst): super().__init__(name="whitelist", kind="classify") whitelist = pynini.string_file(get_abs_path("data/whitelist.tsv")).invert() - graph = pynutil.insert("name: \"") + convert_space(whitelist) + pynutil.insert("\"") + graph = pynutil.insert('name: "') + convert_space(whitelist) + pynutil.insert('"') self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ja/taggers/word.py b/fun_text_processing/inverse_text_normalization/ja/taggers/word.py index b5879d7ac..3fc491577 100644 --- a/fun_text_processing/inverse_text_normalization/ja/taggers/word.py +++ b/fun_text_processing/inverse_text_normalization/ja/taggers/word.py @@ -1,5 +1,9 @@ import pynini -from fun_text_processing.inverse_text_normalization.ja.graph_utils import DAMO_NOT_SPACE, GraphFst, DAMO_CHAR +from fun_text_processing.inverse_text_normalization.ja.graph_utils import ( + DAMO_NOT_SPACE, + GraphFst, + DAMO_CHAR, +) from pynini.lib import pynutil @@ -11,6 +15,6 @@ class WordFst(GraphFst): def __init__(self): super().__init__(name="word", kind="classify") - word = pynutil.insert("name: \"") + DAMO_NOT_SPACE + pynutil.insert("\"") + word = pynutil.insert('name: "') + DAMO_NOT_SPACE + pynutil.insert('"') self.fst = word.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ja/utils.py b/fun_text_processing/inverse_text_normalization/ja/utils.py index 9b9786777..38a95bac5 100644 --- a/fun_text_processing/inverse_text_normalization/ja/utils.py +++ b/fun_text_processing/inverse_text_normalization/ja/utils.py @@ -13,7 +13,7 @@ def num_to_word(x: Union[str, int]): Args x: integer - Returns: spoken representation + Returns: spoken representation """ if isinstance(x, int): x = str(x) @@ -27,7 +27,7 @@ def get_abs_path(rel_path): Args: rel_path: relative path to this file - + Returns absolute path """ - return os.path.dirname(os.path.abspath(__file__)) + '/' + rel_path + return os.path.dirname(os.path.abspath(__file__)) + "/" + rel_path diff --git a/fun_text_processing/inverse_text_normalization/ja/verbalizers/cardinal.py b/fun_text_processing/inverse_text_normalization/ja/verbalizers/cardinal.py index 03f6915fd..522692a7a 100644 --- a/fun_text_processing/inverse_text_normalization/ja/verbalizers/cardinal.py +++ b/fun_text_processing/inverse_text_normalization/ja/verbalizers/cardinal.py @@ -1,8 +1,12 @@ - # -*- coding: utf-8 -*- import pynini -from fun_text_processing.inverse_text_normalization.ja.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space, DAMO_CHAR +from fun_text_processing.inverse_text_normalization.ja.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, + DAMO_CHAR, +) from pynini.lib import pynutil @@ -21,9 +25,9 @@ class CardinalFst(GraphFst): optional_sign = pynini.closure( pynutil.delete("negative:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + DAMO_NOT_QUOTE - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space, 0, 1, @@ -31,9 +35,9 @@ class CardinalFst(GraphFst): graph = ( pynutil.delete("integer:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) # if self.enable_standalone_number: diff --git a/fun_text_processing/inverse_text_normalization/ja/verbalizers/date.py b/fun_text_processing/inverse_text_normalization/ja/verbalizers/date.py index d017ae482..87fcc1380 100644 --- a/fun_text_processing/inverse_text_normalization/ja/verbalizers/date.py +++ b/fun_text_processing/inverse_text_normalization/ja/verbalizers/date.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.inverse_text_normalization.ja.graph_utils import ( DAMO_NOT_QUOTE, @@ -22,43 +20,47 @@ class DateFst(GraphFst): month = ( pynutil.delete("month:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) day = ( pynutil.delete("day:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) year = ( pynutil.delete("year:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') ) # month (day) year graph_mdy = ( - month + pynini.closure(delete_extra_space + day, 0, 1) + pynini.closure(delete_extra_space + year, 0, 1) + month + + pynini.closure(delete_extra_space + day, 0, 1) + + pynini.closure(delete_extra_space + year, 0, 1) ) # (day) month year graph_dmy = ( - pynini.closure(day + delete_extra_space, 0, 1) + month + pynini.closure(delete_extra_space + year, 0, 1) + pynini.closure(day + delete_extra_space, 0, 1) + + month + + pynini.closure(delete_extra_space + year, 0, 1) ) optional_preserve_order = pynini.closure( pynutil.delete("preserve_order:") + delete_space + pynutil.delete("true") + delete_space | pynutil.delete("field_order:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + DAMO_NOT_QUOTE - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space ) diff --git a/fun_text_processing/inverse_text_normalization/ja/verbalizers/decimal.py b/fun_text_processing/inverse_text_normalization/ja/verbalizers/decimal.py index 81ed2b48b..cc69c0f70 100644 --- a/fun_text_processing/inverse_text_normalization/ja/verbalizers/decimal.py +++ b/fun_text_processing/inverse_text_normalization/ja/verbalizers/decimal.py @@ -1,7 +1,9 @@ - - import pynini -from fun_text_processing.inverse_text_normalization.ja.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.ja.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -13,30 +15,30 @@ class DecimalFst(GraphFst): def __init__(self): super().__init__(name="decimal", kind="verbalize") - optionl_sign = pynini.closure(pynini.cross("negative: \"true\"", "-") + delete_space, 0, 1) + optionl_sign = pynini.closure(pynini.cross('negative: "true"', "-") + delete_space, 0, 1) integer = ( pynutil.delete("integer_part:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_integer = pynini.closure(integer + delete_space, 0, 1) fractional = ( pynutil.insert(".") + pynutil.delete("fractional_part:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_fractional = pynini.closure(fractional + delete_space, 0, 1) quantity = ( pynutil.delete("quantity:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_quantity = pynini.closure(pynutil.insert(" ") + quantity + delete_space, 0, 1) graph = optional_integer + optional_fractional + optional_quantity diff --git a/fun_text_processing/inverse_text_normalization/ja/verbalizers/electronic.py b/fun_text_processing/inverse_text_normalization/ja/verbalizers/electronic.py index faa2ec8eb..7e674c54b 100644 --- a/fun_text_processing/inverse_text_normalization/ja/verbalizers/electronic.py +++ b/fun_text_processing/inverse_text_normalization/ja/verbalizers/electronic.py @@ -1,7 +1,9 @@ - - import pynini -from fun_text_processing.inverse_text_normalization.ja.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.ja.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -16,24 +18,24 @@ class ElectronicFst(GraphFst): user_name = ( pynutil.delete("username:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) domain = ( pynutil.delete("domain:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) protocol = ( pynutil.delete("protocol:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) graph = user_name + delete_space + pynutil.insert("@") + domain diff --git a/fun_text_processing/inverse_text_normalization/ja/verbalizers/fraction.py b/fun_text_processing/inverse_text_normalization/ja/verbalizers/fraction.py index 503f64f6a..c1bb028f9 100644 --- a/fun_text_processing/inverse_text_normalization/ja/verbalizers/fraction.py +++ b/fun_text_processing/inverse_text_normalization/ja/verbalizers/fraction.py @@ -1,6 +1,11 @@ - import pynini -from fun_text_processing.inverse_text_normalization.ja.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space, DAMO_SIGMA, delete_extra_space +from fun_text_processing.inverse_text_normalization.ja.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, + DAMO_SIGMA, + delete_extra_space, +) from pynini.lib import pynutil @@ -14,7 +19,9 @@ class FractionFst(GraphFst): def __init__(self): super().__init__(name="fraction", kind="verbalize") optional_sign = pynini.closure(pynini.cross('negative: "true"', "-") + delete_space, 0, 1) - numerator = pynutil.delete('numerator: "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') + numerator = ( + pynutil.delete('numerator: "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') + ) denominator = ( pynutil.insert("/") diff --git a/fun_text_processing/inverse_text_normalization/ja/verbalizers/measure.py b/fun_text_processing/inverse_text_normalization/ja/verbalizers/measure.py index bb0dee2fc..a9f6448e4 100644 --- a/fun_text_processing/inverse_text_normalization/ja/verbalizers/measure.py +++ b/fun_text_processing/inverse_text_normalization/ja/verbalizers/measure.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.ja.graph_utils import DAMO_CHAR, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.ja.graph_utils import ( + DAMO_CHAR, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -16,13 +19,13 @@ class MeasureFst(GraphFst): def __init__(self, decimal: GraphFst, cardinal: GraphFst): super().__init__(name="measure", kind="verbalize") - optional_sign = pynini.closure(pynini.cross("negative: \"true\"", "-"), 0, 1) + optional_sign = pynini.closure(pynini.cross('negative: "true"', "-"), 0, 1) unit = ( pynutil.delete("units:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space ) graph_decimal = ( diff --git a/fun_text_processing/inverse_text_normalization/ja/verbalizers/money.py b/fun_text_processing/inverse_text_normalization/ja/verbalizers/money.py index 510e8d65a..142008b7c 100644 --- a/fun_text_processing/inverse_text_normalization/ja/verbalizers/money.py +++ b/fun_text_processing/inverse_text_normalization/ja/verbalizers/money.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.ja.graph_utils import DAMO_CHAR, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.ja.graph_utils import ( + DAMO_CHAR, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -18,9 +21,9 @@ class MoneyFst(GraphFst): unit = ( pynutil.delete("currency:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) graph = unit + delete_space + decimal.numbers delete_tokens = self.delete_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/ja/verbalizers/ordinal.py b/fun_text_processing/inverse_text_normalization/ja/verbalizers/ordinal.py index 6d665c7b3..0458859c6 100644 --- a/fun_text_processing/inverse_text_normalization/ja/verbalizers/ordinal.py +++ b/fun_text_processing/inverse_text_normalization/ja/verbalizers/ordinal.py @@ -1,6 +1,10 @@ - import pynini -from fun_text_processing.inverse_text_normalization.ja.graph_utils import DAMO_NOT_QUOTE, DAMO_SIGMA, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.ja.graph_utils import ( + DAMO_NOT_QUOTE, + DAMO_SIGMA, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -16,11 +20,11 @@ class OrdinalFst(GraphFst): graph = ( pynutil.delete("integer:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynutil.insert("第", weight=0.01) # + DAMO_NOT_QUOTE + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) # convert_hundred = pynini.cross("第百", "第100") # convert_eleven = pynini.cross("11", "十一") @@ -30,7 +34,6 @@ class OrdinalFst(GraphFst): # convert_two = pynini.cross("2", "第二") # convert_three = pynini.cross("3", "第三") - # suffix = pynini.cdrewrite( # convert_hundred # # convert_eleven diff --git a/fun_text_processing/inverse_text_normalization/ja/verbalizers/telephone.py b/fun_text_processing/inverse_text_normalization/ja/verbalizers/telephone.py index 0c727bbaf..ec28d5dea 100644 --- a/fun_text_processing/inverse_text_normalization/ja/verbalizers/telephone.py +++ b/fun_text_processing/inverse_text_normalization/ja/verbalizers/telephone.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.ja.graph_utils import DAMO_NOT_QUOTE, GraphFst from pynini.lib import pynutil @@ -8,17 +7,21 @@ class TelephoneFst(GraphFst): """ Finite state transducer for verbalizing telephone, e.g. telephone { number_part: "123-123-5678" } - -> 123-123-5678 + -> 123-123-5678 """ def __init__(self): super().__init__(name="telephone", kind="verbalize") - number_part = pynutil.delete("number_part: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") - optional_country_code = pynini.closure( - pynutil.delete("country_code: \"") + number_part = ( + pynutil.delete('number_part: "') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') + ) + optional_country_code = pynini.closure( + pynutil.delete('country_code: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + pynini.accep(" "), 0, 1, diff --git a/fun_text_processing/inverse_text_normalization/ja/verbalizers/time.py b/fun_text_processing/inverse_text_normalization/ja/verbalizers/time.py index 5fd56aac3..1b69664ce 100644 --- a/fun_text_processing/inverse_text_normalization/ja/verbalizers/time.py +++ b/fun_text_processing/inverse_text_normalization/ja/verbalizers/time.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.ja.graph_utils import ( DAMO_CHAR, @@ -20,29 +19,31 @@ class TimeFst(GraphFst): def __init__(self): super().__init__(name="time", kind="verbalize") - add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | (pynutil.insert("0") + DAMO_DIGIT) + add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | ( + pynutil.insert("0") + DAMO_DIGIT + ) hour = ( pynutil.delete("hours:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_DIGIT, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) minute = ( pynutil.delete("minutes:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_DIGIT, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) suffix = ( delete_space + insert_space + pynutil.delete("suffix:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_suffix = pynini.closure(suffix, 0, 1) zone = ( @@ -50,9 +51,9 @@ class TimeFst(GraphFst): + insert_space + pynutil.delete("zone:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_zone = pynini.closure(zone, 0, 1) graph = ( diff --git a/fun_text_processing/inverse_text_normalization/ja/verbalizers/verbalize.py b/fun_text_processing/inverse_text_normalization/ja/verbalizers/verbalize.py index 1f2d07116..3a373e4eb 100644 --- a/fun_text_processing/inverse_text_normalization/ja/verbalizers/verbalize.py +++ b/fun_text_processing/inverse_text_normalization/ja/verbalizers/verbalize.py @@ -1,4 +1,3 @@ - from fun_text_processing.inverse_text_normalization.ja.verbalizers.cardinal import CardinalFst from fun_text_processing.inverse_text_normalization.ja.verbalizers.date import DateFst from fun_text_processing.inverse_text_normalization.ja.verbalizers.decimal import DecimalFst @@ -16,7 +15,7 @@ from fun_text_processing.inverse_text_normalization.ja.graph_utils import GraphF class VerbalizeFst(GraphFst): """ Composes other verbalizer grammars. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. """ diff --git a/fun_text_processing/inverse_text_normalization/ja/verbalizers/verbalize_final.py b/fun_text_processing/inverse_text_normalization/ja/verbalizers/verbalize_final.py index c18c6b04f..d16450fac 100644 --- a/fun_text_processing/inverse_text_normalization/ja/verbalizers/verbalize_final.py +++ b/fun_text_processing/inverse_text_normalization/ja/verbalizers/verbalize_final.py @@ -1,14 +1,17 @@ - import pynini from fun_text_processing.inverse_text_normalization.ja.verbalizers.verbalize import VerbalizeFst from fun_text_processing.inverse_text_normalization.ja.verbalizers.word import WordFst -from fun_text_processing.inverse_text_normalization.ja.graph_utils import GraphFst, delete_extra_space, delete_space +from fun_text_processing.inverse_text_normalization.ja.graph_utils import ( + GraphFst, + delete_extra_space, + delete_space, +) from pynini.lib import pynutil class VerbalizeFinalFst(GraphFst): """ - Finite state transducer that verbalizes an entire sentence, e.g. + Finite state transducer that verbalizes an entire sentence, e.g. tokens { name: "its" } tokens { time { hours: "12" minutes: "30" } } tokens { name: "now" } -> its 12:30 now """ diff --git a/fun_text_processing/inverse_text_normalization/ja/verbalizers/whitelist.py b/fun_text_processing/inverse_text_normalization/ja/verbalizers/whitelist.py index abe1c3ae5..f0ba56441 100644 --- a/fun_text_processing/inverse_text_normalization/ja/verbalizers/whitelist.py +++ b/fun_text_processing/inverse_text_normalization/ja/verbalizers/whitelist.py @@ -1,7 +1,10 @@ - - import pynini -from fun_text_processing.inverse_text_normalization.ja.graph_utils import DAMO_CHAR, DAMO_SIGMA, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.ja.graph_utils import ( + DAMO_CHAR, + DAMO_SIGMA, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -16,9 +19,9 @@ class WhiteListFst(GraphFst): graph = ( pynutil.delete("name:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) - graph = graph @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + graph = graph @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ja/verbalizers/word.py b/fun_text_processing/inverse_text_normalization/ja/verbalizers/word.py index 47c278fc5..66f2bf10f 100644 --- a/fun_text_processing/inverse_text_normalization/ja/verbalizers/word.py +++ b/fun_text_processing/inverse_text_normalization/ja/verbalizers/word.py @@ -1,6 +1,10 @@ - import pynini -from fun_text_processing.inverse_text_normalization.ja.graph_utils import DAMO_CHAR, DAMO_SIGMA, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.ja.graph_utils import ( + DAMO_CHAR, + DAMO_SIGMA, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -13,7 +17,13 @@ class WordFst(GraphFst): def __init__(self): super().__init__(name="word", kind="verbalize") chars = pynini.closure(DAMO_CHAR - " ", 1) - char = pynutil.delete("name:") + delete_space + pynutil.delete("\"") + chars + pynutil.delete("\"") - graph = char @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + char = ( + pynutil.delete("name:") + + delete_space + + pynutil.delete('"') + + chars + + pynutil.delete('"') + ) + graph = char @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ko/__init__.py b/fun_text_processing/inverse_text_normalization/ko/__init__.py index 888da72a3..e78024567 100755 --- a/fun_text_processing/inverse_text_normalization/ko/__init__.py +++ b/fun_text_processing/inverse_text_normalization/ko/__init__.py @@ -1,4 +1,7 @@ - -from fun_text_processing.inverse_text_normalization.ko.taggers.tokenize_and_classify import ClassifyFst +from fun_text_processing.inverse_text_normalization.ko.taggers.tokenize_and_classify import ( + ClassifyFst, +) from fun_text_processing.inverse_text_normalization.ko.verbalizers.verbalize import VerbalizeFst -from fun_text_processing.inverse_text_normalization.ko.verbalizers.verbalize_final import VerbalizeFinalFst +from fun_text_processing.inverse_text_normalization.ko.verbalizers.verbalize_final import ( + VerbalizeFinalFst, +) diff --git a/fun_text_processing/inverse_text_normalization/ko/clean_eval_data.py b/fun_text_processing/inverse_text_normalization/ko/clean_eval_data.py index f59f8d014..e975bf5ec 100755 --- a/fun_text_processing/inverse_text_normalization/ko/clean_eval_data.py +++ b/fun_text_processing/inverse_text_normalization/ko/clean_eval_data.py @@ -1,4 +1,3 @@ - from argparse import ArgumentParser from typing import List @@ -54,7 +53,7 @@ class Filter: Args: processes given instance with process function - + Returns: processed instance if instance belongs to expected class type or original instance """ if instance.token_type != self.class_type: @@ -72,7 +71,9 @@ def process_cardinal_1(instance: Instance) -> Instance: normalized = instance.normalized un_normalized = re.sub(r"[^0-9]", "", un_normalized) normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_ordinal_1(instance: Instance) -> bool: @@ -85,7 +86,9 @@ def process_ordinal_1(instance: Instance) -> Instance: normalized = instance.normalized un_normalized = re.sub(r"[,\s]", "", un_normalized) normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_decimal_1(instance: Instance) -> bool: @@ -98,7 +101,9 @@ def process_decimal_1(instance: Instance) -> Instance: un_normalized = re.sub(r",", "", un_normalized) normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_measure_1(instance: Instance) -> bool: @@ -115,7 +120,9 @@ def process_measure_1(instance: Instance) -> Instance: normalized = re.sub(r"[^a-z\s]", "", normalized) normalized = re.sub(r"per ([a-z\s]*)s$", r"per \1", normalized) normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_money_1(instance: Instance) -> bool: @@ -132,7 +139,9 @@ def process_money_1(instance: Instance) -> Instance: un_normalized = re.sub(r"(\d)m\s*$", r"\1 million", un_normalized) un_normalized = re.sub(r"(\d)bn?\s*$", r"\1 billion", un_normalized) normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_time_1(instance: Instance) -> bool: @@ -147,7 +156,9 @@ def process_time_1(instance: Instance) -> Instance: un_normalized = re.sub(r"(\d)\s?p\s?m\s?", r"\1 p.m.", un_normalized) normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_plain_1(instance: Instance) -> bool: @@ -158,7 +169,9 @@ def filter_plain_1(instance: Instance) -> bool: def process_plain_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_punct_1(instance: Instance) -> bool: @@ -169,7 +182,9 @@ def filter_punct_1(instance: Instance) -> bool: def process_punct_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_date_1(instance: Instance) -> bool: @@ -182,7 +197,9 @@ def process_date_1(instance: Instance) -> Instance: un_normalized = re.sub(r",", "", un_normalized) normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_letters_1(instance: Instance) -> bool: @@ -194,7 +211,9 @@ def process_letters_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_verbatim_1(instance: Instance) -> bool: @@ -205,7 +224,9 @@ def filter_verbatim_1(instance: Instance) -> bool: def process_verbatim_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_digit_1(instance: Instance) -> bool: @@ -217,7 +238,9 @@ def process_digit_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_telephone_1(instance: Instance) -> bool: @@ -229,7 +252,9 @@ def process_telephone_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_electronic_1(instance: Instance) -> bool: @@ -241,7 +266,9 @@ def process_electronic_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_fraction_1(instance: Instance) -> bool: @@ -253,7 +280,9 @@ def process_fraction_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_address_1(instance: Instance) -> bool: @@ -265,27 +294,51 @@ def process_address_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) filters = [] -filters.append(Filter(class_type="CARDINAL", process_func=process_cardinal_1, filter_func=filter_cardinal_1)) -filters.append(Filter(class_type="ORDINAL", process_func=process_ordinal_1, filter_func=filter_ordinal_1)) -filters.append(Filter(class_type="DECIMAL", process_func=process_decimal_1, filter_func=filter_decimal_1)) -filters.append(Filter(class_type="MEASURE", process_func=process_measure_1, filter_func=filter_measure_1)) +filters.append( + Filter(class_type="CARDINAL", process_func=process_cardinal_1, filter_func=filter_cardinal_1) +) +filters.append( + Filter(class_type="ORDINAL", process_func=process_ordinal_1, filter_func=filter_ordinal_1) +) +filters.append( + Filter(class_type="DECIMAL", process_func=process_decimal_1, filter_func=filter_decimal_1) +) +filters.append( + Filter(class_type="MEASURE", process_func=process_measure_1, filter_func=filter_measure_1) +) filters.append(Filter(class_type="MONEY", process_func=process_money_1, filter_func=filter_money_1)) filters.append(Filter(class_type="TIME", process_func=process_time_1, filter_func=filter_time_1)) filters.append(Filter(class_type="DATE", process_func=process_date_1, filter_func=filter_date_1)) filters.append(Filter(class_type="PLAIN", process_func=process_plain_1, filter_func=filter_plain_1)) filters.append(Filter(class_type="PUNCT", process_func=process_punct_1, filter_func=filter_punct_1)) -filters.append(Filter(class_type="LETTERS", process_func=process_letters_1, filter_func=filter_letters_1)) -filters.append(Filter(class_type="VERBATIM", process_func=process_verbatim_1, filter_func=filter_verbatim_1)) +filters.append( + Filter(class_type="LETTERS", process_func=process_letters_1, filter_func=filter_letters_1) +) +filters.append( + Filter(class_type="VERBATIM", process_func=process_verbatim_1, filter_func=filter_verbatim_1) +) filters.append(Filter(class_type="DIGIT", process_func=process_digit_1, filter_func=filter_digit_1)) -filters.append(Filter(class_type="TELEPHONE", process_func=process_telephone_1, filter_func=filter_telephone_1)) -filters.append(Filter(class_type="ELECTRONIC", process_func=process_electronic_1, filter_func=filter_electronic_1)) -filters.append(Filter(class_type="FRACTION", process_func=process_fraction_1, filter_func=filter_fraction_1)) -filters.append(Filter(class_type="ADDRESS", process_func=process_address_1, filter_func=filter_address_1)) +filters.append( + Filter(class_type="TELEPHONE", process_func=process_telephone_1, filter_func=filter_telephone_1) +) +filters.append( + Filter( + class_type="ELECTRONIC", process_func=process_electronic_1, filter_func=filter_electronic_1 + ) +) +filters.append( + Filter(class_type="FRACTION", process_func=process_fraction_1, filter_func=filter_fraction_1) +) +filters.append( + Filter(class_type="ADDRESS", process_func=process_address_1, filter_func=filter_address_1) +) filters.append(Filter(class_type=EOS_TYPE, process_func=lambda x: x, filter_func=lambda x: True)) @@ -314,8 +367,10 @@ def filter_loaded_data(data: List[Instance], verbose: bool = False) -> List[Inst def parse_args(): parser = ArgumentParser() - parser.add_argument("--input", help="input file path", type=str, default='./en_with_types/output-00001-of-00100') - parser.add_argument("--verbose", help="print filtered instances", action='store_true') + parser.add_argument( + "--input", help="input file path", type=str, default="./en_with_types/output-00001-of-00100" + ) + parser.add_argument("--verbose", help="print filtered instances", action="store_true") return parser.parse_args() diff --git a/fun_text_processing/inverse_text_normalization/ko/graph_utils.py b/fun_text_processing/inverse_text_normalization/ko/graph_utils.py index c9f737d75..3091083fe 100755 --- a/fun_text_processing/inverse_text_normalization/ko/graph_utils.py +++ b/fun_text_processing/inverse_text_normalization/ko/graph_utils.py @@ -1,4 +1,3 @@ - import os import string from pathlib import Path @@ -19,9 +18,9 @@ DAMO_UPPER = pynini.union(*string.ascii_uppercase).optimize() DAMO_ALPHA = pynini.union(DAMO_LOWER, DAMO_UPPER).optimize() DAMO_ALNUM = pynini.union(DAMO_DIGIT, DAMO_ALPHA).optimize() DAMO_HEX = pynini.union(*string.hexdigits).optimize() -DAMO_NON_BREAKING_SPACE = u"\u00A0" +DAMO_NON_BREAKING_SPACE = "\u00A0" DAMO_SPACE = " " -DAMO_WHITE_SPACE = pynini.union(" ", "\t", "\n", "\r", u"\u00A0").optimize() +DAMO_WHITE_SPACE = pynini.union(" ", "\t", "\n", "\r", "\u00A0").optimize() DAMO_NOT_SPACE = pynini.difference(DAMO_CHAR, DAMO_WHITE_SPACE).optimize() DAMO_NOT_QUOTE = pynini.difference(DAMO_CHAR, r'"').optimize() @@ -36,31 +35,55 @@ insert_space = pynutil.insert(" ") delete_extra_space = pynini.cross(pynini.closure(DAMO_WHITE_SPACE, 1), " ") delete_preserve_order = pynini.closure( pynutil.delete(" preserve_order: true") - | (pynutil.delete(" field_order: \"") + DAMO_NOT_QUOTE + pynutil.delete("\"")) + | (pynutil.delete(' field_order: "') + DAMO_NOT_QUOTE + pynutil.delete('"')) ) suppletive = pynini.string_file(get_abs_path("data/suppletive.tsv")) # _v = pynini.union("a", "e", "i", "o", "u") _c = pynini.union( - "b", "c", "d", "f", "g", "h", "j", "k", "l", "m", "n", "p", "q", "r", "s", "t", "v", "w", "x", "y", "z" + "b", + "c", + "d", + "f", + "g", + "h", + "j", + "k", + "l", + "m", + "n", + "p", + "q", + "r", + "s", + "t", + "v", + "w", + "x", + "y", + "z", ) _ies = DAMO_SIGMA + _c + pynini.cross("y", "ies") _es = DAMO_SIGMA + pynini.union("s", "sh", "ch", "x", "z") + pynutil.insert("es") _s = DAMO_SIGMA + pynutil.insert("s") graph_plural = plurals._priority_union( - suppletive, plurals._priority_union(_ies, plurals._priority_union(_es, _s, DAMO_SIGMA), DAMO_SIGMA), DAMO_SIGMA + suppletive, + plurals._priority_union(_ies, plurals._priority_union(_es, _s, DAMO_SIGMA), DAMO_SIGMA), + DAMO_SIGMA, ).optimize() SINGULAR_TO_PLURAL = graph_plural PLURAL_TO_SINGULAR = pynini.invert(graph_plural) -TO_LOWER = pynini.union(*[pynini.cross(x, y) for x, y in zip(string.ascii_uppercase, string.ascii_lowercase)]) +TO_LOWER = pynini.union( + *[pynini.cross(x, y) for x, y in zip(string.ascii_uppercase, string.ascii_lowercase)] +) TO_UPPER = pynini.invert(TO_LOWER) MIN_NEG_WEIGHT = -0.0001 MIN_POS_WEIGHT = 0.0001 -def generator_main(file_name: str, graphs: Dict[str, 'pynini.FstLike']): +def generator_main(file_name: str, graphs: Dict[str, "pynini.FstLike"]): """ Exports graph as OpenFst finite state archive (FAR) file with given file name and rule name. @@ -72,7 +95,7 @@ def generator_main(file_name: str, graphs: Dict[str, 'pynini.FstLike']): for rule, graph in graphs.items(): exporter[rule] = graph.optimize() exporter.close() - print(f'Created {file_name}') + print(f"Created {file_name}") def get_plurals(fst): @@ -99,7 +122,7 @@ def get_singulars(fst): return PLURAL_TO_SINGULAR @ fst -def convert_space(fst) -> 'pynini.FstLike': +def convert_space(fst) -> "pynini.FstLike": """ Converts space to nonbreaking space. Used only in tagger grammars for transducing token values within quotes, e.g. name: "hello kitty" @@ -110,7 +133,9 @@ def convert_space(fst) -> 'pynini.FstLike': Returns output fst where breaking spaces are converted to non breaking spaces """ - return fst @ pynini.cdrewrite(pynini.cross(DAMO_SPACE, DAMO_NON_BREAKING_SPACE), "", "", DAMO_SIGMA) + return fst @ pynini.cdrewrite( + pynini.cross(DAMO_SPACE, DAMO_NON_BREAKING_SPACE), "", "", DAMO_SIGMA + ) class GraphFst: @@ -130,9 +155,11 @@ class GraphFst: self._fst = None self.deterministic = deterministic - self.far_path = Path(os.path.dirname(__file__) + '/grammars/' + kind + '/' + name + '.far') + self.far_path = Path(os.path.dirname(__file__) + "/grammars/" + kind + "/" + name + ".far") if self.far_exist(): - self._fst = Far(self.far_path, mode="r", arc_type="standard", far_type="default").get_fst() + self._fst = Far( + self.far_path, mode="r", arc_type="standard", far_type="default" + ).get_fst() def far_exist(self) -> bool: """ @@ -141,14 +168,14 @@ class GraphFst: return self.far_path.exists() @property - def fst(self) -> 'pynini.FstLike': + def fst(self) -> "pynini.FstLike": return self._fst @fst.setter def fst(self, fst): self._fst = fst - def add_tokens(self, fst) -> 'pynini.FstLike': + def add_tokens(self, fst) -> "pynini.FstLike": """ Wraps class name around to given fst @@ -160,7 +187,7 @@ class GraphFst: """ return pynutil.insert(f"{self.name} {{ ") + fst + pynutil.insert(" }") - def delete_tokens(self, fst) -> 'pynini.FstLike': + def delete_tokens(self, fst) -> "pynini.FstLike": """ Deletes class name wrap around output of given fst @@ -179,4 +206,4 @@ class GraphFst: + delete_space + pynutil.delete("}") ) - return res @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + return res @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) diff --git a/fun_text_processing/inverse_text_normalization/ko/taggers/cardinal.py b/fun_text_processing/inverse_text_normalization/ko/taggers/cardinal.py index 0a5490537..6f21d1565 100755 --- a/fun_text_processing/inverse_text_normalization/ko/taggers/cardinal.py +++ b/fun_text_processing/inverse_text_normalization/ko/taggers/cardinal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.ko.utils import get_abs_path from fun_text_processing.inverse_text_normalization.ko.graph_utils import ( @@ -9,10 +8,10 @@ from fun_text_processing.inverse_text_normalization.ko.graph_utils import ( DAMO_CHAR, GraphFst, delete_space, - ) from pynini.lib import pynutil + class CardinalFst(GraphFst): """ Finite state transducer for classifying cardinals @@ -25,87 +24,119 @@ class CardinalFst(GraphFst): graph_zero = pynini.string_file(get_abs_path("data/numbers/zero.tsv")) graph_digit = pynini.string_file(get_abs_path("data/numbers/digit.tsv")) - graph_teens_without_zero = pynini.string_file(get_abs_path("data/numbers/digit_teens_without_zero.tsv")) - graph_teens = pynini.string_file(get_abs_path("data/numbers/digit_teens.tsv")) + graph_teens_without_zero = pynini.string_file( + get_abs_path("data/numbers/digit_teens_without_zero.tsv") + ) + graph_teens = pynini.string_file(get_abs_path("data/numbers/digit_teens.tsv")) graph_inh_digit = pynini.string_file(get_abs_path("data/numbers/digit_inherent_digit.tsv")) - graph_inh_teen_without_zero = pynini.string_file(get_abs_path("data/numbers/digit_inherent_teens_without_zero.tsv")) + graph_inh_teen_without_zero = pynini.string_file( + get_abs_path("data/numbers/digit_inherent_teens_without_zero.tsv") + ) graph_inh_teen = pynini.string_file(get_abs_path("data/numbers/digit_inherent_teens.tsv")) - graph_inh_teen_others = pynini.string_file(get_abs_path("data/numbers/digit_inherent_others.tsv")) - + graph_inh_teen_others = pynini.string_file( + get_abs_path("data/numbers/digit_inherent_others.tsv") + ) + graph_less_hundred_num_inh_p1 = graph_inh_teen_without_zero + graph_inh_digit - graph_less_hundred_num_inh = pynini.union(graph_inh_teen, graph_less_hundred_num_inh_p1, graph_inh_teen_others) + graph_less_hundred_num_inh = pynini.union( + graph_inh_teen, graph_less_hundred_num_inh_p1, graph_inh_teen_others + ) graph_less_hundred_num_p1 = graph_teens_without_zero + graph_digit graph_less_hundred_num = pynini.union(graph_less_hundred_num_p1, graph_teens) - #digits + # digits addzero = pynutil.insert("0") zero = graph_zero digits_combine = graph_digit | graph_inh_digit | zero digits = graph_digit | zero digit = graph_digit - #teens + # teens teens_combine = graph_less_hundred_num | graph_less_hundred_num_inh - #teens = graph_less_hundred_num + # teens = graph_less_hundred_num teens = teens_combine - - #hundred, #백 单位 百 - hundred = ( digit + pynutil.delete("백") + (teens - | pynutil.add_weight(zero + digit, 0.1) - | pynutil.add_weight(digit + addzero, 0.5) - | pynutil.add_weight(addzero**2, 1.0))) + + # hundred, #백 单位 百 + hundred = ( + digit + + pynutil.delete("백") + + ( + teens + | pynutil.add_weight(zero + digit, 0.1) + | pynutil.add_weight(digit + addzero, 0.5) + | pynutil.add_weight(addzero**2, 1.0) + ) + ) graph_hundred_component_at_least_one_none_zero_digit = hundred @ ( - pynini.closure(DAMO_DIGIT) + (DAMO_DIGIT - "0") + pynini.closure(DAMO_DIGIT) + pynini.closure(DAMO_DIGIT) + (DAMO_DIGIT - "0") + pynini.closure(DAMO_DIGIT) ) self.graph_hundred_component_at_least_one_none_zero_digit = ( graph_hundred_component_at_least_one_none_zero_digit ) - ##thousand 천 千单位 - thousand = ((hundred | teens | digits) + pynutil.delete("천") + ( - hundred - | pynutil.add_weight(zero + teens, 0.1) - | pynutil.add_weight(addzero + zero + digit, 0.5) - | pynutil.add_weight(digit + addzero**2, 0.8) - | pynutil.add_weight(addzero**3, 1.0))) + ##thousand 천 千单位 + thousand = ( + (hundred | teens | digits) + + pynutil.delete("천") + + ( + hundred + | pynutil.add_weight(zero + teens, 0.1) + | pynutil.add_weight(addzero + zero + digit, 0.5) + | pynutil.add_weight(digit + addzero**2, 0.8) + | pynutil.add_weight(addzero**3, 1.0) + ) + ) ##만 单位万 - ten_thousand = ((thousand | hundred | teens | digits) - + pynutil.delete("만") - + pynini.cross(" ", "").ques - + (thousand - | pynutil.add_weight(zero + hundred, 0.1) - | pynutil.add_weight(addzero + zero + teens, 0.5) - | pynutil.add_weight(addzero + addzero + zero + digit, 0.5) - | pynutil.add_weight(digit + addzero**3, 0.8) - | pynutil.add_weight(addzero**4, 1.0))) + ten_thousand = ( + (thousand | hundred | teens | digits) + + pynutil.delete("만") + + pynini.cross(" ", "").ques + + ( + thousand + | pynutil.add_weight(zero + hundred, 0.1) + | pynutil.add_weight(addzero + zero + teens, 0.5) + | pynutil.add_weight(addzero + addzero + zero + digit, 0.5) + | pynutil.add_weight(digit + addzero**3, 0.8) + | pynutil.add_weight(addzero**4, 1.0) + ) + ) ##조, 单位兆， 억, 单位亿 number = digits | teens | hundred | thousand | ten_thousand - - ## ques is equal to pynini.closure(, 0, 1) - number = ((number + pynini.accep("조").ques + pynini.cross(" ", "").ques).ques + - (number + pynini.accep("억").ques + pynini.cross(" ", "").ques).ques + number) - graph = number | graph_less_hundred_num_inh | graph_inh_digit | graph_inh_teen | graph_inh_teen_others - #labels_exception = [num_to_word(x) for x in range(0, 13)] - labels_exception = [ 'zzzzzzzzz' ] + ## ques is equal to pynini.closure(, 0, 1) + number = ( + (number + pynini.accep("조").ques + pynini.cross(" ", "").ques).ques + + (number + pynini.accep("억").ques + pynini.cross(" ", "").ques).ques + + number + ) + + graph = ( + number + | graph_less_hundred_num_inh + | graph_inh_digit + | graph_inh_teen + | graph_inh_teen_others + ) + # labels_exception = [num_to_word(x) for x in range(0, 13)] + labels_exception = ["zzzzzzzzz"] graph_exception = pynini.union(*labels_exception) - + self.graph_no_exception = graph self.graph = (pynini.project(graph, "input") - graph_exception.arcsort()) @ graph optional_minus_graph = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("마이너스", "\"-\"") + DAMO_SPACE, 0, 1 + pynutil.insert("negative: ") + pynini.cross("마이너스", '"-"') + DAMO_SPACE, 0, 1 ) - final_graph = optional_minus_graph + pynutil.insert("integer: \"") + self.graph + pynutil.insert("\"") + final_graph = ( + optional_minus_graph + pynutil.insert('integer: "') + self.graph + pynutil.insert('"') + ) final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() - - diff --git a/fun_text_processing/inverse_text_normalization/ko/taggers/date.py b/fun_text_processing/inverse_text_normalization/ko/taggers/date.py index 8d31fc77d..870bfaa09 100755 --- a/fun_text_processing/inverse_text_normalization/ko/taggers/date.py +++ b/fun_text_processing/inverse_text_normalization/ko/taggers/date.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.ko.utils import get_abs_path from fun_text_processing.inverse_text_normalization.ko.graph_utils import ( @@ -12,16 +11,20 @@ from pynini.lib import pynutil graph_zero = pynini.string_file(get_abs_path("data/numbers/zero.tsv")).optimize() graph_digit = pynini.string_file(get_abs_path("data/numbers/digit.tsv")).optimize() -graph_digit_inh = pynini.string_file(get_abs_path("data/numbers/digit_inherent_digit.tsv")).optimize() +graph_digit_inh = pynini.string_file( + get_abs_path("data/numbers/digit_inherent_digit.tsv") +).optimize() + def _get_month_graph(): """ Transducer for month, e.g. march -> march """ month_graph = pynini.string_file(get_abs_path("data/months.tsv")) - #print(month_graph) + # print(month_graph) return month_graph + def _get_day_graph(): """ Transducer for month, e.g. march -> march @@ -29,21 +32,23 @@ def _get_day_graph(): day_graph_num = pynini.string_file(get_abs_path("data/day.tsv")) day_graph_inh = pynini.string_file(get_abs_path("data/day_inherent.tsv")) day_graph = pynini.union(day_graph_num, day_graph_inh) - #print(day_graph) + # print(day_graph) return day_graph + def _get_year_graph(): """ Transducer for year, e.g. twenty twenty -> 2020 """ - digit = (graph_digit | graph_digit_inh) + digit = graph_digit | graph_digit_inh zero = graph_zero - year_graph_4num = digit + (digit | zero)**3 + year_graph_4num = digit + (digit | zero) ** 3 year_graph_2num = digit**2 year_graph = pynini.union(year_graph_4num, year_graph_2num) return year_graph + class DateFst(GraphFst): """ Finite state transducer for classifying date, @@ -59,26 +64,30 @@ class DateFst(GraphFst): year_graph = _get_year_graph() + pynini.accep("년") YEAR_WEIGHT = 0.001 - year_graph = pynutil.insert("year: \"") + pynutil.add_weight(year_graph, YEAR_WEIGHT) + pynutil.insert("\"") - #year_graph_space = pynutil.insert("year: \"") + pynutil.add_weight(year_graph, YEAR_WEIGHT) + pynutil.insert("\"") + pynutil.insert(" ") - #year_graph = pynutil.insert("year: \"") + year_graph + pynutil.insert("\"") + year_graph = ( + pynutil.insert('year: "') + + pynutil.add_weight(year_graph, YEAR_WEIGHT) + + pynutil.insert('"') + ) + # year_graph_space = pynutil.insert("year: \"") + pynutil.add_weight(year_graph, YEAR_WEIGHT) + pynutil.insert("\"") + pynutil.insert(" ") + # year_graph = pynutil.insert("year: \"") + year_graph + pynutil.insert("\"") MONTH_WEIGHT = -0.001 month_graph = _get_month_graph() + pynini.cross("", "월") - #month_graph = pynutil.insert("month: \"") + pynutil.add_weight(month_graph, MONTH_WEIGHT) + pynutil.insert("\"") - month_graph = pynutil.insert("month: \"") + month_graph + pynutil.insert("\"") - #month_graph_space = pynutil.insert("month: \"") + month_graph + pynutil.insert("\"") + pynutil.insert(" ") - + # month_graph = pynutil.insert("month: \"") + pynutil.add_weight(month_graph, MONTH_WEIGHT) + pynutil.insert("\"") + month_graph = pynutil.insert('month: "') + month_graph + pynutil.insert('"') + # month_graph_space = pynutil.insert("month: \"") + month_graph + pynutil.insert("\"") + pynutil.insert(" ") + day_graph = _get_day_graph() + pynini.cross("", "일") DAY_WEIGHT = -0.7 - #day_graph = pynutil.insert("day: \"") + pynutil.add_weight(day_graph, DAY_WEIGHT) + pynutil.insert("\"") - day_graph = pynutil.insert("day: \"") + day_graph + pynutil.insert("\"") - #day_graph_space = pynutil.insert("day: \"") + day_graph + pynutil.insert("\"") + pynutil.insert(" ") + # day_graph = pynutil.insert("day: \"") + pynutil.add_weight(day_graph, DAY_WEIGHT) + pynutil.insert("\"") + day_graph = pynutil.insert('day: "') + day_graph + pynutil.insert('"') + # day_graph_space = pynutil.insert("day: \"") + day_graph + pynutil.insert("\"") + pynutil.insert(" ") graph_ymd = year_graph + delete_space + month_graph + delete_space + day_graph graph_md = month_graph + delete_space + day_graph - graph_ym = year_graph + delete_space + month_graph - + graph_ym = year_graph + delete_space + month_graph + final_graph = graph_ymd | graph_md | graph_ym | year_graph | month_graph | day_graph final_graph += pynutil.insert(" preserve_order: true") diff --git a/fun_text_processing/inverse_text_normalization/ko/taggers/decimal.py b/fun_text_processing/inverse_text_normalization/ko/taggers/decimal.py index 3c3071087..9cbecd1c8 100755 --- a/fun_text_processing/inverse_text_normalization/ko/taggers/decimal.py +++ b/fun_text_processing/inverse_text_normalization/ko/taggers/decimal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.ko.utils import get_abs_path from fun_text_processing.inverse_text_normalization.ko.graph_utils import ( @@ -10,32 +9,42 @@ from fun_text_processing.inverse_text_normalization.ko.graph_utils import ( from pynini.lib import pynutil -def get_quantity(decimal: 'pynini.FstLike', cardinal_up_to_hundred: 'pynini.FstLike') -> 'pynini.FstLike': +def get_quantity( + decimal: "pynini.FstLike", cardinal_up_to_hundred: "pynini.FstLike" +) -> "pynini.FstLike": """ Returns FST that transforms either a cardinal or decimal followed by a quantity into a numeral, e.g. one million -> integer_part: "1" quantity: "million" e.g. one point five million -> integer_part: "1" fractional_part: "5" quantity: "million" - Args: + Args: decimal: decimal FST cardinal_up_to_hundred: cardinal FST """ numbers = cardinal_up_to_hundred @ ( - pynutil.delete(pynini.closure("0")) + pynini.difference(DAMO_DIGIT, "0") + pynini.closure(DAMO_DIGIT) + pynutil.delete(pynini.closure("0")) + + pynini.difference(DAMO_DIGIT, "0") + + pynini.closure(DAMO_DIGIT) ) - #"만", "백만", "천만", "억", "조", 万、百万、千万、亿、兆 - #천 千 + # "만", "백만", "천만", "억", "조", 万、百万、千万、亿、兆 + # 천 千 suffix = pynini.union("만", "백만", "천만", "억", "조") res = ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + numbers - + pynutil.insert("\"") + + pynutil.insert('"') + delete_extra_space - + pynutil.insert("quantity: \"") + + pynutil.insert('quantity: "') + suffix - + pynutil.insert("\"") + + pynutil.insert('"') + ) + res |= ( + decimal + + delete_extra_space + + pynutil.insert('quantity: "') + + (suffix | "천") + + pynutil.insert('"') ) - res |= decimal + delete_extra_space + pynutil.insert("quantity: \"") + (suffix | "천") + pynutil.insert("\"") return res @@ -61,18 +70,25 @@ class DecimalFst(GraphFst): ##마이너스 负 optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("마이너스", "\"true\"") + delete_extra_space, 0, 1 + pynutil.insert("negative: ") + pynini.cross("마이너스", '"true"') + delete_extra_space, + 0, + 1, ) - graph_fractional = pynutil.insert("fractional_part: \"") + graph_decimal + pynutil.insert("\"") - - #점 点 - graph_integer = pynutil.insert("integer_part: \"") + cardinal_graph + pynutil.delete("점") + pynutil.insert("\"") - - final_graph_wo_sign = ( - graph_integer + pynini.cross(" ", " ") + graph_fractional + graph_fractional = ( + pynutil.insert('fractional_part: "') + graph_decimal + pynutil.insert('"') ) + # 점 点 + graph_integer = ( + pynutil.insert('integer_part: "') + + cardinal_graph + + pynutil.delete("점") + + pynutil.insert('"') + ) + + final_graph_wo_sign = graph_integer + pynini.cross(" ", " ") + graph_fractional + final_graph = optional_graph_negative + delete_space + final_graph_wo_sign self.final_graph_wo_negative = final_graph_wo_sign | get_quantity( diff --git a/fun_text_processing/inverse_text_normalization/ko/taggers/electronic.py b/fun_text_processing/inverse_text_normalization/ko/taggers/electronic.py index 8ab0f787f..f65e71795 100755 --- a/fun_text_processing/inverse_text_normalization/ko/taggers/electronic.py +++ b/fun_text_processing/inverse_text_normalization/ko/taggers/electronic.py @@ -1,7 +1,10 @@ - import pynini from fun_text_processing.inverse_text_normalization.ko.utils import get_abs_path -from fun_text_processing.inverse_text_normalization.ko.graph_utils import DAMO_ALPHA, GraphFst, insert_space +from fun_text_processing.inverse_text_normalization.ko.graph_utils import ( + DAMO_ALPHA, + GraphFst, + insert_space, +) from pynini.lib import pynutil @@ -25,35 +28,50 @@ class ElectronicFst(GraphFst): accepted_username = alpha_num | symbols process_dot = pynini.cross("점", ".") - username = (alpha_num + pynini.closure(delete_extra_space + accepted_username)) | pynutil.add_weight( - pynini.closure(DAMO_ALPHA, 1), weight=0.0001 - ) - username = pynutil.insert("username: \"") + username + pynutil.insert("\"") + username = ( + alpha_num + pynini.closure(delete_extra_space + accepted_username) + ) | pynutil.add_weight(pynini.closure(DAMO_ALPHA, 1), weight=0.0001) + username = pynutil.insert('username: "') + username + pynutil.insert('"') single_alphanum = pynini.closure(alpha_num + delete_extra_space) + alpha_num - server = single_alphanum | pynini.string_file(get_abs_path("data/electronic/server_name.tsv")) + server = single_alphanum | pynini.string_file( + get_abs_path("data/electronic/server_name.tsv") + ) domain = single_alphanum | pynini.string_file(get_abs_path("data/electronic/domain.tsv")) domain_graph = ( - pynutil.insert("domain: \"") + pynutil.insert('domain: "') + server + delete_extra_space + process_dot + delete_extra_space + domain - + pynutil.insert("\"") + + pynutil.insert('"') + ) + graph = ( + username + + delete_extra_space + + pynutil.delete("에서") + + insert_space + + delete_extra_space + + domain_graph ) - graph = username + delete_extra_space + pynutil.delete("에서") + insert_space + delete_extra_space + domain_graph ############# url ### protocol_end = pynini.cross(pynini.union("w w w", "www"), "www") - protocol_start = (pynini.cross("h t t p", "http") | pynini.cross("h t t p s", "https")) + pynini.cross( - " 콜론 슬래시 슬래시 ", "://" - ) + protocol_start = ( + pynini.cross("h t t p", "http") | pynini.cross("h t t p s", "https") + ) + pynini.cross(" 콜론 슬래시 슬래시 ", "://") # .com, ending = ( delete_extra_space + symbols + delete_extra_space - + (domain | pynini.closure(accepted_username + delete_extra_space,) + accepted_username) + + ( + domain + | pynini.closure( + accepted_username + delete_extra_space, + ) + + accepted_username + ) ) protocol_default = ( @@ -64,12 +82,18 @@ class ElectronicFst(GraphFst): + pynini.closure(ending, 1) ).optimize() protocol = ( - pynini.closure(protocol_start, 0, 1) + protocol_end + delete_extra_space + process_dot + protocol_default + pynini.closure(protocol_start, 0, 1) + + protocol_end + + delete_extra_space + + process_dot + + protocol_default ).optimize() - protocol |= pynini.closure(protocol_end + delete_extra_space + process_dot, 0, 1) + protocol_default + protocol |= ( + pynini.closure(protocol_end + delete_extra_space + process_dot, 0, 1) + protocol_default + ) - protocol = pynutil.insert("protocol: \"") + protocol.optimize() + pynutil.insert("\"") + protocol = pynutil.insert('protocol: "') + protocol.optimize() + pynutil.insert('"') graph |= protocol final_graph = self.add_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/ko/taggers/fraction.py b/fun_text_processing/inverse_text_normalization/ko/taggers/fraction.py index e841f8854..b57aed505 100755 --- a/fun_text_processing/inverse_text_normalization/ko/taggers/fraction.py +++ b/fun_text_processing/inverse_text_normalization/ko/taggers/fraction.py @@ -1,4 +1,3 @@ - from fun_text_processing.inverse_text_normalization.ko.graph_utils import ( DAMO_NOT_QUOTE, GraphFst, @@ -12,6 +11,7 @@ from fun_text_processing.inverse_text_normalization.ko.graph_utils import ( import pynini from pynini.lib import pynutil + class FractionFst(GraphFst): """ Finite state transducer for classifying fraction @@ -20,14 +20,19 @@ class FractionFst(GraphFst): def __init__(self, cardinal: GraphFst): super().__init__(name="fraction", kind="classify") # integer_part # numerator # denominator - + graph_cardinal = cardinal.graph_no_exception - - #without the integerate part - #分子 + + # without the integerate part + # 分子 numerator = pynutil.insert('numerator: "') + graph_cardinal + pynutil.insert('"') - #分母 - denominator = pynutil.insert('denominator: "') + graph_cardinal + pynutil.delete("분의") + pynutil.insert('"') + # 分母 + denominator = ( + pynutil.insert('denominator: "') + + graph_cardinal + + pynutil.delete("분의") + + pynutil.insert('"') + ) ## graph_fraction_component = denominator + pynini.cross(" ", " ") + numerator @@ -40,9 +45,9 @@ class FractionFst(GraphFst): ##负 optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("마이너스", "\"true\"") + DAMO_SPACE, 0, 1 + pynutil.insert("negative: ") + pynini.cross("마이너스", '"true"') + DAMO_SPACE, 0, 1 ) graph = optional_graph_negative + graph final_graph = self.add_tokens(graph) - self.fst = final_graph.optimize() + self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ko/taggers/measure.py b/fun_text_processing/inverse_text_normalization/ko/taggers/measure.py index 71645d3b3..9e9ace5e1 100755 --- a/fun_text_processing/inverse_text_normalization/ko/taggers/measure.py +++ b/fun_text_processing/inverse_text_normalization/ko/taggers/measure.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.ko.utils import get_abs_path from fun_text_processing.inverse_text_normalization.ko.graph_utils import ( @@ -29,19 +28,17 @@ class MeasureFst(GraphFst): decimal_graph = decimal.final_graph_wo_negative unit_graph = pynini.string_file(get_abs_path("data/measurements.tsv")) - + graph_unit = pynini.invert(unit_graph) # singular -> abbr ## 마이너 负 optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("마이너", "\"true\"") + delete_extra_space, 0, 1 + pynutil.insert("negative: ") + pynini.cross("마이너", '"true"') + delete_extra_space, + 0, + 1, ) - graph_units = ( - pynutil.insert("units: \"") - + graph_unit - + pynutil.insert("\"") - ) + graph_units = pynutil.insert('units: "') + graph_unit + pynutil.insert('"') subgraph_decimal = ( pynutil.insert("decimal { ") @@ -54,14 +51,14 @@ class MeasureFst(GraphFst): subgraph_cardinal = ( pynutil.insert("cardinal { ") + optional_graph_negative - + pynutil.insert("integer: \"") + + pynutil.insert('integer: "') + cardinal_graph - + pynutil.insert("\"") + + pynutil.insert('"') + pynutil.insert(" }") + delete_extra_space + graph_units ) - + final_graph = subgraph_decimal | subgraph_cardinal final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ko/taggers/money.py b/fun_text_processing/inverse_text_normalization/ko/taggers/money.py index 4e41b94a7..9f1bad967 100755 --- a/fun_text_processing/inverse_text_normalization/ko/taggers/money.py +++ b/fun_text_processing/inverse_text_normalization/ko/taggers/money.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.ko.utils import get_abs_path from fun_text_processing.inverse_text_normalization.ko.graph_utils import ( @@ -34,12 +33,12 @@ class MoneyFst(GraphFst): unit = pynini.string_file(get_abs_path("data/currency.tsv")).invert() - graph_unit = pynutil.insert("currency: \"") + unit + pynutil.insert("\"") + graph_unit = pynutil.insert('currency: "') + unit + pynutil.insert('"') graph_integer = ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + cardinal_graph - + pynutil.insert("\"") + + pynutil.insert('"') + delete_extra_space + graph_unit ) diff --git a/fun_text_processing/inverse_text_normalization/ko/taggers/punctuation.py b/fun_text_processing/inverse_text_normalization/ko/taggers/punctuation.py index 0aac488d8..1edb9aca9 100755 --- a/fun_text_processing/inverse_text_normalization/ko/taggers/punctuation.py +++ b/fun_text_processing/inverse_text_normalization/ko/taggers/punctuation.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.ko.graph_utils import GraphFst from pynini.lib import pynutil @@ -13,9 +12,9 @@ class PunctuationFst(GraphFst): def __init__(self): super().__init__(name="punctuation", kind="classify") - s = ",.?" #here, we only support three type of punctuation + s = ",.?" # here, we only support three type of punctuation punct = pynini.union(*s) - graph = pynutil.insert("name: \"") + punct + pynutil.insert("\"") + graph = pynutil.insert('name: "') + punct + pynutil.insert('"') self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ko/taggers/telephone.py b/fun_text_processing/inverse_text_normalization/ko/taggers/telephone.py index 17473a7af..4a8343c65 100755 --- a/fun_text_processing/inverse_text_normalization/ko/taggers/telephone.py +++ b/fun_text_processing/inverse_text_normalization/ko/taggers/telephone.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.ko.utils import get_abs_path from fun_text_processing.inverse_text_normalization.ko.graph_utils import ( @@ -24,7 +23,7 @@ def get_serial_number(cardinal): class TelephoneFst(GraphFst): """ - Finite state transducer for classifying telephone numbers, e.g. + Finite state transducer for classifying telephone numbers, e.g. one two three one two three five six seven eight -> { number_part: "123-123-5678" } This class also support card number and IP format. @@ -44,36 +43,36 @@ class TelephoneFst(GraphFst): graph_digit = pynini.string_file(get_abs_path("data/numbers/digit.tsv")) graph_zero = pynini.string_file(get_abs_path("data/numbers/zero.tsv")) graph_dot = pynini.string_file(get_abs_path("data/numbers/dot.tsv")) - + graph_digits = graph_digit | graph_zero phone_number_graph = graph_digits**9 | graph_digits**10 | graph_digits**11 - + country_code = ( - pynutil.insert("country_code: \"") + pynutil.insert('country_code: "') + pynini.closure(pynini.cross("더한", "+"), 0, 1) + (pynini.closure(graph_digits, 0, 2) + graph_digits) - + pynutil.insert("\"") + + pynutil.insert('"') ) - optional_country_code = pynini.closure(country_code + pynutil.delete(" ") + insert_space, 0, 1).optimize() + optional_country_code = pynini.closure( + country_code + pynutil.delete(" ") + insert_space, 0, 1 + ).optimize() grpah_phone_number = ( - pynutil.insert("number_part: \"") - + phone_number_graph - + pynutil.insert("\"") - ) + pynutil.insert('number_part: "') + phone_number_graph + pynutil.insert('"') + ) graph = optional_country_code + grpah_phone_number # ip ip_graph = graph_digit.plus + (graph_dot + graph_digits.plus).plus - graph |= pynutil.insert("number_part: \"") + ip_graph.optimize() + pynutil.insert("\"") + graph |= pynutil.insert('number_part: "') + ip_graph.optimize() + pynutil.insert('"') graph |= ( - pynutil.insert("number_part: \"") + pynutil.insert('number_part: "') + pynutil.add_weight(get_serial_number(cardinal=cardinal), weight=0.0001) - + pynutil.insert("\"") + + pynutil.insert('"') ) final_graph = self.add_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/ko/taggers/time.py b/fun_text_processing/inverse_text_normalization/ko/taggers/time.py index 24c82bc66..f814ca57d 100755 --- a/fun_text_processing/inverse_text_normalization/ko/taggers/time.py +++ b/fun_text_processing/inverse_text_normalization/ko/taggers/time.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.ko.taggers.cardinal import CardinalFst from fun_text_processing.inverse_text_normalization.ko.utils import get_abs_path, num_to_word @@ -29,54 +28,66 @@ class TimeFst(GraphFst): suffix_graph = pynini.string_file(get_abs_path("data/time/time_suffix.tsv")) time_zone_graph = pynini.invert(pynini.string_file(get_abs_path("data/time/time_zone.tsv"))) - + hour_graph = pynini.string_file(get_abs_path("data/time/hours.tsv")) minute_graph = pynini.string_file(get_abs_path("data/time/minutes.tsv")) second_graph = pynini.string_file(get_abs_path("data/time/seconds.tsv")) # only used for < 1000 thousand -> 0 weight - #cardinal = pynutil.add_weight(CardinalFst().graph_no_exception, weight=-0.7) + # cardinal = pynutil.add_weight(CardinalFst().graph_no_exception, weight=-0.7) graph_hour = hour_graph graph_minute = minute_graph graph_second = second_graph - - final_graph_hour = pynutil.insert("hours: \"") + graph_hour + pynutil.insert("\"") - - final_suffix = pynutil.insert("suffix: \"") + convert_space(suffix_graph) + pynutil.insert("\"") + + final_graph_hour = pynutil.insert('hours: "') + graph_hour + pynutil.insert('"') + + final_suffix = ( + pynutil.insert('suffix: "') + convert_space(suffix_graph) + pynutil.insert('"') + ) final_suffix = delete_space + insert_space + final_suffix final_suffix_optional = pynini.closure(final_suffix, 0, 1) final_time_zone_optional = pynini.closure( delete_space + insert_space - + pynutil.insert("zone: \"") + + pynutil.insert('zone: "') + convert_space(time_zone_graph) - + pynutil.insert("\""), + + pynutil.insert('"'), 0, 1, ) graph_hm = ( - final_graph_hour + delete_extra_space + pynutil.insert("minutes: \"") + graph_minute + pynutil.insert("\"") + final_graph_hour + + delete_extra_space + + pynutil.insert('minutes: "') + + graph_minute + + pynutil.insert('"') ) graph_hms = ( - final_graph_hour + delete_extra_space + pynutil.insert("minutes: \"") + graph_minute + pynutil.insert("\"") + delete_extra_space + pynutil.insert("seconds: \"") + graph_second + pynutil.insert("\"") + final_graph_hour + + delete_extra_space + + pynutil.insert('minutes: "') + + graph_minute + + pynutil.insert('"') + + delete_extra_space + + pynutil.insert('seconds: "') + + graph_second + + pynutil.insert('"') ) graph_h = ( final_graph_hour + delete_extra_space - + pynutil.insert("minutes: \"") + + pynutil.insert('minutes: "') + (pynutil.insert("00") | graph_minute) - + pynutil.insert("\"") + + pynutil.insert('"') + final_suffix + final_time_zone_optional ) - final_graph = ( - (graph_hm | graph_hms) + final_suffix_optional + final_time_zone_optional - ) + final_graph = (graph_hm | graph_hms) + final_suffix_optional + final_time_zone_optional final_graph |= graph_h diff --git a/fun_text_processing/inverse_text_normalization/ko/taggers/tokenize_and_classify.py b/fun_text_processing/inverse_text_normalization/ko/taggers/tokenize_and_classify.py index b9a250278..015321cae 100755 --- a/fun_text_processing/inverse_text_normalization/ko/taggers/tokenize_and_classify.py +++ b/fun_text_processing/inverse_text_normalization/ko/taggers/tokenize_and_classify.py @@ -1,4 +1,3 @@ - import os import pynini @@ -28,7 +27,7 @@ import logging class ClassifyFst(GraphFst): """ Final class that composes all other classification grammars. This class can process an entire sentence, that is lower cased. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. Args: @@ -81,10 +80,16 @@ class ClassifyFst(GraphFst): | pynutil.add_weight(word_graph, 100) ) - punct = pynutil.insert("tokens { ") + pynutil.add_weight(punct_graph, weight=1.1) + pynutil.insert(" }") + punct = ( + pynutil.insert("tokens { ") + + pynutil.add_weight(punct_graph, weight=1.1) + + pynutil.insert(" }") + ) token = pynutil.insert("tokens { ") + classify + pynutil.insert(" }") token_plus_punct = ( - pynini.closure(punct + pynutil.insert(" ")) + token + pynini.closure(pynutil.insert(" ") + punct) + pynini.closure(punct + pynutil.insert(" ")) + + token + + pynini.closure(pynutil.insert(" ") + punct) ) graph = token_plus_punct + pynini.closure(delete_extra_space + token_plus_punct) diff --git a/fun_text_processing/inverse_text_normalization/ko/taggers/whitelist.py b/fun_text_processing/inverse_text_normalization/ko/taggers/whitelist.py index 928a7f597..97733ae75 100755 --- a/fun_text_processing/inverse_text_normalization/ko/taggers/whitelist.py +++ b/fun_text_processing/inverse_text_normalization/ko/taggers/whitelist.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.ko.utils import get_abs_path from fun_text_processing.inverse_text_normalization.ko.graph_utils import GraphFst, convert_space @@ -16,5 +15,5 @@ class WhiteListFst(GraphFst): super().__init__(name="whitelist", kind="classify") whitelist = pynini.string_file(get_abs_path("data/whitelist.tsv")).invert() - graph = pynutil.insert("name: \"") + convert_space(whitelist) + pynutil.insert("\"") + graph = pynutil.insert('name: "') + convert_space(whitelist) + pynutil.insert('"') self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ko/taggers/word.py b/fun_text_processing/inverse_text_normalization/ko/taggers/word.py index 04aa2bc9e..5fa7a9e93 100755 --- a/fun_text_processing/inverse_text_normalization/ko/taggers/word.py +++ b/fun_text_processing/inverse_text_normalization/ko/taggers/word.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.ko.graph_utils import DAMO_NOT_SPACE, GraphFst from pynini.lib import pynutil @@ -12,5 +11,5 @@ class WordFst(GraphFst): def __init__(self): super().__init__(name="word", kind="classify") - word = pynutil.insert("name: \"") + pynini.closure(DAMO_NOT_SPACE, 1) + pynutil.insert("\"") + word = pynutil.insert('name: "') + pynini.closure(DAMO_NOT_SPACE, 1) + pynutil.insert('"') self.fst = word.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ko/utils.py b/fun_text_processing/inverse_text_normalization/ko/utils.py index 75d66e62f..c073d630f 100755 --- a/fun_text_processing/inverse_text_normalization/ko/utils.py +++ b/fun_text_processing/inverse_text_normalization/ko/utils.py @@ -1,4 +1,3 @@ - import csv import os from typing import Union @@ -7,28 +6,30 @@ import inflect _inflect = inflect.engine() + def num_to_word(x: Union[str, int]): """ converts integer to spoken representation Args x: integer - Returns: spoken representation + Returns: spoken representation """ if isinstance(x, int): x = str(x) x = _inflect.number_to_words(str(x)).replace("-", " ").replace(",", "") return x + def get_abs_path(rel_path): """ Get absolute path Args: rel_path: relative path to this file - + Returns absolute path """ - return os.path.dirname(os.path.abspath(__file__)) + '/' + rel_path + return os.path.dirname(os.path.abspath(__file__)) + "/" + rel_path def load_labels(abs_path): @@ -47,7 +48,7 @@ def load_labels(abs_path): def augment_labels_with_punct_at_end(labels): """ - augments labels: if key ends on a punctuation that value does not have, add a new label + augments labels: if key ends on a punctuation that value does not have, add a new label where the value maintains the punctuation Args: diff --git a/fun_text_processing/inverse_text_normalization/ko/verbalizers/cardinal.py b/fun_text_processing/inverse_text_normalization/ko/verbalizers/cardinal.py index 8ffcf43db..4acadc660 100755 --- a/fun_text_processing/inverse_text_normalization/ko/verbalizers/cardinal.py +++ b/fun_text_processing/inverse_text_normalization/ko/verbalizers/cardinal.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.ko.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.ko.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -15,9 +18,9 @@ class CardinalFst(GraphFst): optional_sign = pynini.closure( pynutil.delete("negative:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + DAMO_NOT_QUOTE - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space, 0, 1, @@ -25,9 +28,9 @@ class CardinalFst(GraphFst): graph = ( pynutil.delete("integer:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) self.numbers = graph graph = optional_sign + graph diff --git a/fun_text_processing/inverse_text_normalization/ko/verbalizers/date.py b/fun_text_processing/inverse_text_normalization/ko/verbalizers/date.py index 56358b8d5..cd9ab1072 100755 --- a/fun_text_processing/inverse_text_normalization/ko/verbalizers/date.py +++ b/fun_text_processing/inverse_text_normalization/ko/verbalizers/date.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.ko.graph_utils import ( DAMO_NOT_QUOTE, @@ -21,60 +20,67 @@ class DateFst(GraphFst): month = ( pynutil.delete("month:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') + pynutil.insert(" ") ) day = ( pynutil.delete("day:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') + pynutil.insert(" ") ) year = ( pynutil.delete("year:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') + pynutil.insert(" ") ) # month (day) year graph_mdy = ( - month + pynini.closure(delete_extra_space + day, 0, 1) + pynini.closure(delete_extra_space + year, 0, 1) + month + + pynini.closure(delete_extra_space + day, 0, 1) + + pynini.closure(delete_extra_space + year, 0, 1) ) # (day) month year graph_dmy = ( - pynini.closure(day + delete_extra_space, 0, 1) + month + pynini.closure(delete_extra_space + year, 0, 1) + pynini.closure(day + delete_extra_space, 0, 1) + + month + + pynini.closure(delete_extra_space + year, 0, 1) ) optional_preserve_order = pynini.closure( pynutil.delete("preserve_order:") + delete_space + pynutil.delete("true") + delete_space - | pynutil.delete("field_order:") + delete_space + pynutil.delete("\"") + DAMO_NOT_QUOTE + pynutil.delete("\"") + delete_space + | pynutil.delete("field_order:") + + delete_space + + pynutil.delete('"') + + DAMO_NOT_QUOTE + + pynutil.delete('"') + + delete_space ) # year month day - graph_ymd = ( - year + month + day - ) + graph_ymd = year + month + day # month day - graph_md = ( - month + day - ) + graph_md = month + day # year month - graph_ym = ( - year + month - ) - + graph_ym = year + month + # add some grammars - final_graph = (graph_mdy | year | graph_dmy | graph_ymd | graph_md | graph_ym | month | day) + delete_space + optional_preserve_order + final_graph = ( + (graph_mdy | year | graph_dmy | graph_ymd | graph_md | graph_ym | month | day) + + delete_space + + optional_preserve_order + ) delete_tokens = self.delete_tokens(final_graph) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ko/verbalizers/decimal.py b/fun_text_processing/inverse_text_normalization/ko/verbalizers/decimal.py index 6e089b37d..e9ee60339 100755 --- a/fun_text_processing/inverse_text_normalization/ko/verbalizers/decimal.py +++ b/fun_text_processing/inverse_text_normalization/ko/verbalizers/decimal.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.ko.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.ko.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -12,30 +15,30 @@ class DecimalFst(GraphFst): def __init__(self): super().__init__(name="decimal", kind="verbalize") - optionl_sign = pynini.closure(pynini.cross("negative: \"true\"", "-") + delete_space, 0, 1) + optionl_sign = pynini.closure(pynini.cross('negative: "true"', "-") + delete_space, 0, 1) integer = ( pynutil.delete("integer_part:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_integer = pynini.closure(integer + delete_space, 0, 1) fractional = ( pynutil.insert(".") + pynutil.delete("fractional_part:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_fractional = pynini.closure(fractional + delete_space, 0, 1) quantity = ( pynutil.delete("quantity:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_quantity = pynini.closure(pynutil.insert(" ") + quantity + delete_space, 0, 1) graph = optional_integer + optional_fractional + optional_quantity diff --git a/fun_text_processing/inverse_text_normalization/ko/verbalizers/electronic.py b/fun_text_processing/inverse_text_normalization/ko/verbalizers/electronic.py index e21782723..4acfe3f90 100755 --- a/fun_text_processing/inverse_text_normalization/ko/verbalizers/electronic.py +++ b/fun_text_processing/inverse_text_normalization/ko/verbalizers/electronic.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.ko.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.ko.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -15,24 +18,24 @@ class ElectronicFst(GraphFst): user_name = ( pynutil.delete("username:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) domain = ( pynutil.delete("domain:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) protocol = ( pynutil.delete("protocol:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) graph = user_name + delete_space + pynutil.insert("@") + domain diff --git a/fun_text_processing/inverse_text_normalization/ko/verbalizers/fraction.py b/fun_text_processing/inverse_text_normalization/ko/verbalizers/fraction.py index 4acf7ef38..5dd7cdc3f 100755 --- a/fun_text_processing/inverse_text_normalization/ko/verbalizers/fraction.py +++ b/fun_text_processing/inverse_text_normalization/ko/verbalizers/fraction.py @@ -1,4 +1,3 @@ - from fun_text_processing.inverse_text_normalization.ko.graph_utils import ( DAMO_NOT_QUOTE, GraphFst, @@ -9,31 +8,34 @@ from fun_text_processing.inverse_text_normalization.ko.graph_utils import ( import pynini from pynini.lib import pynutil + class FractionFst(GraphFst): """ - Finite state transducer for verbalizing fraction, + Finite state transducer for verbalizing fraction, """ def __init__(self): super().__init__(name="fraction", kind="verbalize") optional_sign = pynini.closure(pynini.cross('negative: "true"', "-") + delete_space, 0, 1) numerator = ( - pynutil.delete('numerator:') + pynutil.delete("numerator:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) denominator = ( - pynutil.delete('denominator:') + pynutil.delete("denominator:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) - graph = (optional_sign + numerator + delete_space + pynutil.insert("/") + denominator).optimize() + graph = ( + optional_sign + numerator + delete_space + pynutil.insert("/") + denominator + ).optimize() self.numbers = graph delete_tokens = self.delete_tokens(optional_sign + graph) - self.fst = delete_tokens.optimize() + self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ko/verbalizers/measure.py b/fun_text_processing/inverse_text_normalization/ko/verbalizers/measure.py index 56067bdc2..2dd25f773 100755 --- a/fun_text_processing/inverse_text_normalization/ko/verbalizers/measure.py +++ b/fun_text_processing/inverse_text_normalization/ko/verbalizers/measure.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.ko.graph_utils import DAMO_CHAR, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.ko.graph_utils import ( + DAMO_CHAR, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -16,13 +19,13 @@ class MeasureFst(GraphFst): def __init__(self, decimal: GraphFst, cardinal: GraphFst): super().__init__(name="measure", kind="verbalize") - optional_sign = pynini.closure(pynini.cross("negative: \"true\"", "-"), 0, 1) + optional_sign = pynini.closure(pynini.cross('negative: "true"', "-"), 0, 1) unit = ( pynutil.delete("units:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space ) graph_decimal = ( diff --git a/fun_text_processing/inverse_text_normalization/ko/verbalizers/money.py b/fun_text_processing/inverse_text_normalization/ko/verbalizers/money.py index 5988283c1..572cc11d8 100755 --- a/fun_text_processing/inverse_text_normalization/ko/verbalizers/money.py +++ b/fun_text_processing/inverse_text_normalization/ko/verbalizers/money.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.ko.graph_utils import DAMO_CHAR, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.ko.graph_utils import ( + DAMO_CHAR, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -18,9 +21,9 @@ class MoneyFst(GraphFst): unit = ( pynutil.delete("currency:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) graph = unit + delete_space + decimal.numbers delete_tokens = self.delete_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/ko/verbalizers/ordinal.py b/fun_text_processing/inverse_text_normalization/ko/verbalizers/ordinal.py index 98b0e11c7..f6c9de5a8 100755 --- a/fun_text_processing/inverse_text_normalization/ko/verbalizers/ordinal.py +++ b/fun_text_processing/inverse_text_normalization/ko/verbalizers/ordinal.py @@ -1,6 +1,10 @@ - import pynini -from fun_text_processing.inverse_text_normalization.ko.graph_utils import DAMO_NOT_QUOTE, DAMO_SIGMA, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.ko.graph_utils import ( + DAMO_NOT_QUOTE, + DAMO_SIGMA, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -15,9 +19,9 @@ class OrdinalFst(GraphFst): graph = ( pynutil.delete("integer:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) convert_eleven = pynini.cross("11", "11th") convert_twelve = pynini.cross("12", "12th") diff --git a/fun_text_processing/inverse_text_normalization/ko/verbalizers/telephone.py b/fun_text_processing/inverse_text_normalization/ko/verbalizers/telephone.py index ac90a7fe6..e0c475a48 100755 --- a/fun_text_processing/inverse_text_normalization/ko/verbalizers/telephone.py +++ b/fun_text_processing/inverse_text_normalization/ko/verbalizers/telephone.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.ko.graph_utils import DAMO_NOT_QUOTE, GraphFst from pynini.lib import pynutil @@ -8,17 +7,21 @@ class TelephoneFst(GraphFst): """ Finite state transducer for verbalizing telephone, e.g. telephone { number_part: "123-123-5678" } - -> 123-123-5678 + -> 123-123-5678 """ def __init__(self): super().__init__(name="telephone", kind="verbalize") - number_part = pynutil.delete("number_part: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") - optional_country_code = pynini.closure( - pynutil.delete("country_code: \"") + number_part = ( + pynutil.delete('number_part: "') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') + ) + optional_country_code = pynini.closure( + pynutil.delete('country_code: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + pynini.accep(" "), 0, 1, diff --git a/fun_text_processing/inverse_text_normalization/ko/verbalizers/time.py b/fun_text_processing/inverse_text_normalization/ko/verbalizers/time.py index 9897c7e37..26bbc2645 100755 --- a/fun_text_processing/inverse_text_normalization/ko/verbalizers/time.py +++ b/fun_text_processing/inverse_text_normalization/ko/verbalizers/time.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.ko.graph_utils import ( DAMO_CHAR, @@ -20,40 +19,42 @@ class TimeFst(GraphFst): def __init__(self): super().__init__(name="time", kind="verbalize") - add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | (pynutil.insert("0") + DAMO_DIGIT) - #hour + add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | ( + pynutil.insert("0") + DAMO_DIGIT + ) + # hour hour = ( pynutil.delete("hours:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_DIGIT, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) - #minute + # minute minute = ( pynutil.delete("minutes:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_DIGIT, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) - #seconds + # seconds second = ( pynutil.delete("seconds:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_DIGIT, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) - + suffix = ( delete_space + insert_space + pynutil.delete("suffix:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_suffix = pynini.closure(suffix, 0, 1) zone = ( @@ -61,18 +62,31 @@ class TimeFst(GraphFst): + insert_space + pynutil.delete("zone:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_zone = pynini.closure(zone, 0, 1) - #hms - graph_hms = (hour @ add_leading_zero_to_double_digit) + delete_space + pynutil.insert(":") + (minute @ add_leading_zero_to_double_digit) + delete_space + pynutil.insert(":") + second - #hm - graph_hm = (hour @ add_leading_zero_to_double_digit) + delete_space + pynutil.insert(":") + (minute @ add_leading_zero_to_double_digit) - - graph = ( (graph_hms | graph_hm) + optional_suffix + optional_zone ) + # hms + graph_hms = ( + (hour @ add_leading_zero_to_double_digit) + + delete_space + + pynutil.insert(":") + + (minute @ add_leading_zero_to_double_digit) + + delete_space + + pynutil.insert(":") + + second + ) + # hm + graph_hm = ( + (hour @ add_leading_zero_to_double_digit) + + delete_space + + pynutil.insert(":") + + (minute @ add_leading_zero_to_double_digit) + ) + + graph = (graph_hms | graph_hm) + optional_suffix + optional_zone delete_tokens = self.delete_tokens(graph) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ko/verbalizers/verbalize.py b/fun_text_processing/inverse_text_normalization/ko/verbalizers/verbalize.py index 5085a81fb..91960d8ab 100755 --- a/fun_text_processing/inverse_text_normalization/ko/verbalizers/verbalize.py +++ b/fun_text_processing/inverse_text_normalization/ko/verbalizers/verbalize.py @@ -1,4 +1,3 @@ - from fun_text_processing.inverse_text_normalization.ko.verbalizers.cardinal import CardinalFst from fun_text_processing.inverse_text_normalization.ko.verbalizers.date import DateFst from fun_text_processing.inverse_text_normalization.ko.verbalizers.decimal import DecimalFst @@ -16,7 +15,7 @@ from fun_text_processing.inverse_text_normalization.ko.graph_utils import GraphF class VerbalizeFst(GraphFst): """ Composes other verbalizer grammars. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. """ diff --git a/fun_text_processing/inverse_text_normalization/ko/verbalizers/verbalize_final.py b/fun_text_processing/inverse_text_normalization/ko/verbalizers/verbalize_final.py index dfd0b3909..cc67060d6 100755 --- a/fun_text_processing/inverse_text_normalization/ko/verbalizers/verbalize_final.py +++ b/fun_text_processing/inverse_text_normalization/ko/verbalizers/verbalize_final.py @@ -1,14 +1,17 @@ - import pynini from fun_text_processing.inverse_text_normalization.ko.verbalizers.verbalize import VerbalizeFst from fun_text_processing.inverse_text_normalization.ko.verbalizers.word import WordFst -from fun_text_processing.inverse_text_normalization.ko.graph_utils import GraphFst, delete_extra_space, delete_space +from fun_text_processing.inverse_text_normalization.ko.graph_utils import ( + GraphFst, + delete_extra_space, + delete_space, +) from pynini.lib import pynutil class VerbalizeFinalFst(GraphFst): """ - Finite state transducer that verbalizes an entire sentence, e.g. + Finite state transducer that verbalizes an entire sentence, e.g. tokens { name: "its" } tokens { time { hours: "12" minutes: "30" } } tokens { name: "now" } -> its 12:30 now """ diff --git a/fun_text_processing/inverse_text_normalization/ko/verbalizers/whitelist.py b/fun_text_processing/inverse_text_normalization/ko/verbalizers/whitelist.py index 60b755af8..654ef2e94 100755 --- a/fun_text_processing/inverse_text_normalization/ko/verbalizers/whitelist.py +++ b/fun_text_processing/inverse_text_normalization/ko/verbalizers/whitelist.py @@ -1,6 +1,10 @@ - import pynini -from fun_text_processing.inverse_text_normalization.ko.graph_utils import DAMO_CHAR, DAMO_SIGMA, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.ko.graph_utils import ( + DAMO_CHAR, + DAMO_SIGMA, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -15,9 +19,9 @@ class WhiteListFst(GraphFst): graph = ( pynutil.delete("name:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) - graph = graph @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + graph = graph @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ko/verbalizers/word.py b/fun_text_processing/inverse_text_normalization/ko/verbalizers/word.py index 1d01e8431..3243520ce 100755 --- a/fun_text_processing/inverse_text_normalization/ko/verbalizers/word.py +++ b/fun_text_processing/inverse_text_normalization/ko/verbalizers/word.py @@ -1,6 +1,10 @@ - import pynini -from fun_text_processing.inverse_text_normalization.ko.graph_utils import DAMO_CHAR, DAMO_SIGMA, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.ko.graph_utils import ( + DAMO_CHAR, + DAMO_SIGMA, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -13,7 +17,13 @@ class WordFst(GraphFst): def __init__(self): super().__init__(name="word", kind="verbalize") chars = pynini.closure(DAMO_CHAR - " ", 1) - char = pynutil.delete("name:") + delete_space + pynutil.delete("\"") + chars + pynutil.delete("\"") - graph = char @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + char = ( + pynutil.delete("name:") + + delete_space + + pynutil.delete('"') + + chars + + pynutil.delete('"') + ) + graph = char @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/pt/__init__.py b/fun_text_processing/inverse_text_normalization/pt/__init__.py index 498c4e26e..427ceec46 100644 --- a/fun_text_processing/inverse_text_normalization/pt/__init__.py +++ b/fun_text_processing/inverse_text_normalization/pt/__init__.py @@ -1,6 +1,7 @@ - - - -from fun_text_processing.inverse_text_normalization.pt.taggers.tokenize_and_classify import ClassifyFst +from fun_text_processing.inverse_text_normalization.pt.taggers.tokenize_and_classify import ( + ClassifyFst, +) from fun_text_processing.inverse_text_normalization.pt.verbalizers.verbalize import VerbalizeFst -from fun_text_processing.inverse_text_normalization.pt.verbalizers.verbalize_final import VerbalizeFinalFst +from fun_text_processing.inverse_text_normalization.pt.verbalizers.verbalize_final import ( + VerbalizeFinalFst, +) diff --git a/fun_text_processing/inverse_text_normalization/pt/data/__init__.py b/fun_text_processing/inverse_text_normalization/pt/data/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/pt/data/__init__.py +++ b/fun_text_processing/inverse_text_normalization/pt/data/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/inverse_text_normalization/pt/data/electronic/__init__.py b/fun_text_processing/inverse_text_normalization/pt/data/electronic/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/pt/data/electronic/__init__.py +++ b/fun_text_processing/inverse_text_normalization/pt/data/electronic/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/inverse_text_normalization/pt/data/numbers/__init__.py b/fun_text_processing/inverse_text_normalization/pt/data/numbers/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/pt/data/numbers/__init__.py +++ b/fun_text_processing/inverse_text_normalization/pt/data/numbers/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/inverse_text_normalization/pt/data/ordinals/__init__.py b/fun_text_processing/inverse_text_normalization/pt/data/ordinals/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/pt/data/ordinals/__init__.py +++ b/fun_text_processing/inverse_text_normalization/pt/data/ordinals/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/inverse_text_normalization/pt/data/time/__init__.py b/fun_text_processing/inverse_text_normalization/pt/data/time/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/pt/data/time/__init__.py +++ b/fun_text_processing/inverse_text_normalization/pt/data/time/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/inverse_text_normalization/pt/taggers/__init__.py b/fun_text_processing/inverse_text_normalization/pt/taggers/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/pt/taggers/__init__.py +++ b/fun_text_processing/inverse_text_normalization/pt/taggers/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/inverse_text_normalization/pt/taggers/cardinal.py b/fun_text_processing/inverse_text_normalization/pt/taggers/cardinal.py index 465cc44c4..8521e2a62 100644 --- a/fun_text_processing/inverse_text_normalization/pt/taggers/cardinal.py +++ b/fun_text_processing/inverse_text_normalization/pt/taggers/cardinal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.pt.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import ( @@ -50,7 +49,9 @@ class CardinalFst(GraphFst): graph_hundred_component += delete_space graph_hundred_component += pynini.union( graph_twenties | graph_teen | pynutil.insert("00"), - (graph_ties | pynutil.insert("0")) + delete_space + (graph_digit | pynutil.insert("0")), + (graph_ties | pynutil.insert("0")) + + delete_space + + (graph_digit | pynutil.insert("0")), ) graph_hundred_component = pynini.union(graph_hundred_component, graph_one_hundred) @@ -59,8 +60,11 @@ class CardinalFst(GraphFst): ) graph_thousands = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("mil"), - pynutil.insert("001") + pynutil.delete("mil"), # because we say 'mil', not 'hum mil' + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + pynutil.delete("mil"), + pynutil.insert("001") + + pynutil.delete("mil"), # because we say 'mil', not 'hum mil' pynutil.insert("000", weight=0.01), ) @@ -126,7 +130,9 @@ class CardinalFst(GraphFst): ) graph = graph @ pynini.union( - pynutil.delete(pynini.closure("0")) + pynini.difference(DAMO_DIGIT, "0") + pynini.closure(DAMO_DIGIT), + pynutil.delete(pynini.closure("0")) + + pynini.difference(DAMO_DIGIT, "0") + + pynini.closure(DAMO_DIGIT), "0", ) @@ -138,7 +144,9 @@ class CardinalFst(GraphFst): else: graph_e = ( - pynutil.delete(DAMO_WHITE_SPACE.plus) + pynutil.delete("e") + pynutil.delete(DAMO_WHITE_SPACE.plus) + pynutil.delete(DAMO_WHITE_SPACE.plus) + + pynutil.delete("e") + + pynutil.delete(DAMO_WHITE_SPACE.plus) ) graph_ties_component = pynini.union( @@ -147,7 +155,9 @@ class CardinalFst(GraphFst): pynutil.add_weight(pynutil.insert("0") + graph_digit, 0.1), ) @ (pynini.closure(DAMO_DIGIT) + (DAMO_DIGIT - "0") + pynini.closure(DAMO_DIGIT)) - graph_hundreds_except_hundred = (pynini.project(graph_hundreds, "input") - "cento") @ graph_hundreds + graph_hundreds_except_hundred = ( + pynini.project(graph_hundreds, "input") - "cento" + ) @ graph_hundreds graph_hundred_component_prefix_e = pynini.union( graph_one_hundred, @@ -156,9 +166,9 @@ class CardinalFst(GraphFst): ) @ (pynini.closure(DAMO_DIGIT) + (DAMO_DIGIT - "0") + pynini.closure(DAMO_DIGIT)) graph_hundred_component_prefix_e = graph_hundred_component_prefix_e.optimize() - graph_hundred_component_no_prefix = pynini.union(graph_hundreds + graph_e + graph_ties_component,) @ ( - pynini.closure(DAMO_DIGIT) + (DAMO_DIGIT - "0") + pynini.closure(DAMO_DIGIT) - ) + graph_hundred_component_no_prefix = pynini.union( + graph_hundreds + graph_e + graph_ties_component, + ) @ (pynini.closure(DAMO_DIGIT) + (DAMO_DIGIT - "0") + pynini.closure(DAMO_DIGIT)) graph_hundred_component_no_prefix = graph_hundred_component_no_prefix.optimize() graph_mil_prefix_e = pynini.union( @@ -276,7 +286,10 @@ class CardinalFst(GraphFst): + delete_space + (pynutil.delete("quintilhão") | pynutil.delete("quintilhões")) ) - + ((graph_e + graph_quatrilhao_prefix_e) | (delete_space + graph_quatrilhao_no_prefix)) + + ( + (graph_e + graph_quatrilhao_prefix_e) + | (delete_space + graph_quatrilhao_no_prefix) + ) ) graph_quintilhao_no_prefix = pynini.union( @@ -288,7 +301,10 @@ class CardinalFst(GraphFst): ) | pynutil.insert("000", weight=0.1) ) - + ((graph_e + graph_quatrilhao_prefix_e) | (delete_space + graph_quatrilhao_no_prefix)) + + ( + (graph_e + graph_quatrilhao_prefix_e) + | (delete_space + graph_quatrilhao_no_prefix) + ) ) graph_sextilhao_prefix_e = pynini.union( @@ -297,7 +313,10 @@ class CardinalFst(GraphFst): + delete_space + (pynutil.delete("sextilhão") | pynutil.delete("sextilhões")) ) - + ((graph_e + graph_quintilhao_prefix_e) | (delete_space + graph_quintilhao_no_prefix)) + + ( + (graph_e + graph_quintilhao_prefix_e) + | (delete_space + graph_quintilhao_no_prefix) + ) ) graph_sextilhao_no_prefix = pynini.union( @@ -309,7 +328,10 @@ class CardinalFst(GraphFst): ) | pynutil.insert("000", weight=0.1) ) - + ((graph_e + graph_quintilhao_prefix_e) | (delete_space + graph_quintilhao_no_prefix)) + + ( + (graph_e + graph_quintilhao_prefix_e) + | (delete_space + graph_quintilhao_no_prefix) + ) ) graph = pynini.union( @@ -327,7 +349,9 @@ class CardinalFst(GraphFst): ).optimize() graph = graph @ pynini.union( - pynutil.delete(pynini.closure("0")) + pynini.difference(DAMO_DIGIT, "0") + pynini.closure(DAMO_DIGIT), + pynutil.delete(pynini.closure("0")) + + pynini.difference(DAMO_DIGIT, "0") + + pynini.closure(DAMO_DIGIT), "0", ) @@ -335,18 +359,18 @@ class CardinalFst(GraphFst): self.graph_no_exception = graph # save self.numbers_up_to_thousand for use in DecimalFst - digits_up_to_thousand = DAMO_DIGIT | (DAMO_DIGIT ** 2) | (DAMO_DIGIT ** 3) + digits_up_to_thousand = DAMO_DIGIT | (DAMO_DIGIT**2) | (DAMO_DIGIT**3) numbers_up_to_thousand = pynini.compose(graph, digits_up_to_thousand).optimize() self.numbers_up_to_thousand = numbers_up_to_thousand # save self.numbers_up_to_million for use in DecimalFst digits_up_to_million = ( DAMO_DIGIT - | (DAMO_DIGIT ** 2) - | (DAMO_DIGIT ** 3) - | (DAMO_DIGIT ** 4) - | (DAMO_DIGIT ** 5) - | (DAMO_DIGIT ** 6) + | (DAMO_DIGIT**2) + | (DAMO_DIGIT**3) + | (DAMO_DIGIT**4) + | (DAMO_DIGIT**5) + | (DAMO_DIGIT**6) ) numbers_up_to_million = pynini.compose(graph, digits_up_to_million).optimize() self.numbers_up_to_million = numbers_up_to_million @@ -357,15 +381,17 @@ class CardinalFst(GraphFst): self.digits_from_year = digits_from_year # don't convert cardinals from zero to nine inclusive - graph_exception = pynini.project(pynini.union(graph_digit, graph_zero), 'input') + graph_exception = pynini.project(pynini.union(graph_digit, graph_zero), "input") self.graph = (pynini.project(graph, "input") - graph_exception.arcsort()) @ graph optional_minus_graph = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("menos", "\"-\"") + DAMO_SPACE, 0, 1 + pynutil.insert("negative: ") + pynini.cross("menos", '"-"') + DAMO_SPACE, 0, 1 ) - final_graph = optional_minus_graph + pynutil.insert("integer: \"") + self.graph + pynutil.insert("\"") + final_graph = ( + optional_minus_graph + pynutil.insert('integer: "') + self.graph + pynutil.insert('"') + ) final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/pt/taggers/date.py b/fun_text_processing/inverse_text_normalization/pt/taggers/date.py index 4c9eb0148..ed247a0e7 100644 --- a/fun_text_processing/inverse_text_normalization/pt/taggers/date.py +++ b/fun_text_processing/inverse_text_normalization/pt/taggers/date.py @@ -1,9 +1,10 @@ - - - import pynini from fun_text_processing.inverse_text_normalization.pt.utils import get_abs_path -from fun_text_processing.text_normalization.en.graph_utils import GraphFst, delete_extra_space, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + GraphFst, + delete_extra_space, + delete_space, +) from pynini.lib import pynutil @@ -37,17 +38,19 @@ class DateFst(GraphFst): # can use "primeiro" for 1st day of the month graph_1_to_31 = pynini.union(graph_1_to_31, pynini.cross("primeiro", "01")) - day_graph = pynutil.insert("day: \"") + graph_1_to_31 + pynutil.insert("\"") + day_graph = pynutil.insert('day: "') + graph_1_to_31 + pynutil.insert('"') month_name_graph = pynini.string_file(get_abs_path("data/months.tsv")) - month_name_graph = pynutil.insert("month: \"") + month_name_graph + pynutil.insert("\"") + month_name_graph = pynutil.insert('month: "') + month_name_graph + pynutil.insert('"') # vinte do oito -> 20/08 digits_1_to_12 = [str("{:0>2d}").format(digits) for digits in range(1, 13)] graph_1_to_12 = graph_1_to_100 @ pynini.union(*digits_1_to_12) - month_number_graph = pynutil.insert("month: \"") + graph_1_to_12 + pynutil.insert("\"") + month_number_graph = pynutil.insert('month: "') + graph_1_to_12 + pynutil.insert('"') - graph_dm = day_graph + delete_space + pynutil.delete("de") + delete_extra_space + month_name_graph + graph_dm = ( + day_graph + delete_space + pynutil.delete("de") + delete_extra_space + month_name_graph + ) graph_dm |= ( day_graph @@ -55,16 +58,16 @@ class DateFst(GraphFst): + pynutil.delete("do") + delete_extra_space + month_number_graph - + pynutil.insert(" morphosyntactic_features: \"/\"") + + pynutil.insert(' morphosyntactic_features: "/"') ) graph_year = ( delete_space + pynutil.delete("de") + delete_extra_space - + pynutil.insert("year: \"") + + pynutil.insert('year: "') + digits_from_year - + pynutil.insert("\"") + + pynutil.insert('"') ) graph_dmy = graph_dm + graph_year.ques diff --git a/fun_text_processing/inverse_text_normalization/pt/taggers/decimal.py b/fun_text_processing/inverse_text_normalization/pt/taggers/decimal.py index a3c2f5d54..8cec15efb 100644 --- a/fun_text_processing/inverse_text_normalization/pt/taggers/decimal.py +++ b/fun_text_processing/inverse_text_normalization/pt/taggers/decimal.py @@ -1,6 +1,3 @@ - - - import pynini from fun_text_processing.inverse_text_normalization.pt.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import ( @@ -12,7 +9,9 @@ from fun_text_processing.text_normalization.en.graph_utils import ( from pynini.lib import pynutil -def get_quantity(decimal: 'pynini.FstLike', cardinal_up_to_million: 'pynini.FstLike') -> 'pynini.FstLike': +def get_quantity( + decimal: "pynini.FstLike", cardinal_up_to_million: "pynini.FstLike" +) -> "pynini.FstLike": """ Returns FST that transforms either a cardinal or decimal followed by a quantity into a numeral, e.g. one million -> integer_part: "1" quantity: "million" @@ -23,7 +22,9 @@ def get_quantity(decimal: 'pynini.FstLike', cardinal_up_to_million: 'pynini.FstL cardinal_up_to_million: cardinal FST """ numbers = cardinal_up_to_million @ ( - pynutil.delete(pynini.closure("0")) + pynini.difference(DAMO_DIGIT, "0") + pynini.closure(DAMO_DIGIT) + pynutil.delete(pynini.closure("0")) + + pynini.difference(DAMO_DIGIT, "0") + + pynini.closure(DAMO_DIGIT) ) suffix = pynini.union( @@ -41,15 +42,17 @@ def get_quantity(decimal: 'pynini.FstLike', cardinal_up_to_million: 'pynini.FstL "sextilhões", ) res = ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + numbers - + pynutil.insert("\"") + + pynutil.insert('"') + delete_extra_space - + pynutil.insert("quantity: \"") + + pynutil.insert('quantity: "') + suffix - + pynutil.insert("\"") + + pynutil.insert('"') + ) + res |= ( + decimal + delete_extra_space + pynutil.insert('quantity: "') + suffix + pynutil.insert('"') ) - res |= decimal + delete_extra_space + pynutil.insert("quantity: \"") + suffix + pynutil.insert("\"") return res @@ -81,17 +84,23 @@ class DecimalFst(GraphFst): self.graph = graph_decimal # decimal point can be denoted by 'vírgula' or 'ponto' - decimal_point = pynini.cross("vírgula", "morphosyntactic_features: \",\"") - decimal_point |= pynini.cross("ponto", "morphosyntactic_features: \".\"") + decimal_point = pynini.cross("vírgula", 'morphosyntactic_features: ","') + decimal_point |= pynini.cross("ponto", 'morphosyntactic_features: "."') optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("menos", "\"true\"") + delete_extra_space, 0, 1 + pynutil.insert("negative: ") + pynini.cross("menos", '"true"') + delete_extra_space, + 0, + 1, ) - graph_fractional = pynutil.insert("fractional_part: \"") + graph_decimal + pynutil.insert("\"") + graph_fractional = ( + pynutil.insert('fractional_part: "') + graph_decimal + pynutil.insert('"') + ) - cardinal_graph = cardinal.graph_no_exception | pynini.string_file(get_abs_path("data/numbers/zero.tsv")) - graph_integer = pynutil.insert("integer_part: \"") + cardinal_graph + pynutil.insert("\"") + cardinal_graph = cardinal.graph_no_exception | pynini.string_file( + get_abs_path("data/numbers/zero.tsv") + ) + graph_integer = pynutil.insert('integer_part: "') + cardinal_graph + pynutil.insert('"') final_graph_wo_sign = ( pynini.closure(graph_integer + delete_extra_space, 0, 1) + decimal_point @@ -103,6 +112,8 @@ class DecimalFst(GraphFst): self.final_graph_wo_negative = final_graph_wo_sign | get_quantity( final_graph_wo_sign, cardinal.numbers_up_to_million ) - final_graph |= optional_graph_negative + get_quantity(final_graph_wo_sign, cardinal.numbers_up_to_million) + final_graph |= optional_graph_negative + get_quantity( + final_graph_wo_sign, cardinal.numbers_up_to_million + ) final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/pt/taggers/electronic.py b/fun_text_processing/inverse_text_normalization/pt/taggers/electronic.py index 1aee1be83..abe8c373a 100644 --- a/fun_text_processing/inverse_text_normalization/pt/taggers/electronic.py +++ b/fun_text_processing/inverse_text_normalization/pt/taggers/electronic.py @@ -1,6 +1,3 @@ - - - import pynini from fun_text_processing.inverse_text_normalization.pt.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import DAMO_ALPHA, GraphFst, insert_space @@ -31,33 +28,46 @@ class ElectronicFst(GraphFst): accepted_username = alpha_num | symbols process_dot = pynini.cross("ponto", ".") username = ( - pynutil.insert("username: \"") + pynutil.insert('username: "') + alpha_num + delete_extra_space + pynini.closure(accepted_username + delete_extra_space) + alpha_num - + pynutil.insert("\"") + + pynutil.insert('"') ) single_alphanum = pynini.closure(alpha_num + delete_extra_space) + alpha_num - server = single_alphanum | pynini.string_file(get_abs_path("data/electronic/server_name.tsv")).invert() - domain = single_alphanum | pynini.string_file(get_abs_path("data/electronic/domain.tsv")).invert() + server = ( + single_alphanum + | pynini.string_file(get_abs_path("data/electronic/server_name.tsv")).invert() + ) + domain = ( + single_alphanum + | pynini.string_file(get_abs_path("data/electronic/domain.tsv")).invert() + ) domain_graph = ( - pynutil.insert("domain: \"") + pynutil.insert('domain: "') + server + delete_extra_space + process_dot + delete_extra_space + domain - + pynutil.insert("\"") + + pynutil.insert('"') ) graph = ( - username + delete_extra_space + pynutil.delete("arroba") + insert_space + delete_extra_space + domain_graph + username + + delete_extra_space + + pynutil.delete("arroba") + + insert_space + + delete_extra_space + + domain_graph ) ############# url ### protocol_end = pynini.cross(pynini.union("www", "w w w", "dáblio dáblio dáblio"), "www") protocol_start = pynini.cross(pynini.union("http", "h t t p", "agá tê tê pê"), "http") - protocol_start |= pynini.cross(pynini.union("https", "h t t p s", "agá tê tê pê ésse"), "https") + protocol_start |= pynini.cross( + pynini.union("https", "h t t p s", "agá tê tê pê ésse"), "https" + ) protocol_start += pynini.cross(" dois pontos barra barra ", "://") # e.g. .com, .es @@ -77,7 +87,7 @@ class ElectronicFst(GraphFst): + (pynini.closure(delete_extra_space + accepted_username, 1) | server) + pynini.closure(ending, 1) ) - protocol = pynutil.insert("protocol: \"") + protocol + pynutil.insert("\"") + protocol = pynutil.insert('protocol: "') + protocol + pynutil.insert('"') graph |= protocol ######## diff --git a/fun_text_processing/inverse_text_normalization/pt/taggers/measure.py b/fun_text_processing/inverse_text_normalization/pt/taggers/measure.py index 1859cd92e..9e1397348 100644 --- a/fun_text_processing/inverse_text_normalization/pt/taggers/measure.py +++ b/fun_text_processing/inverse_text_normalization/pt/taggers/measure.py @@ -1,6 +1,3 @@ - - - import pynini from fun_text_processing.inverse_text_normalization.pt.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import ( @@ -28,26 +25,45 @@ class MeasureFst(GraphFst): cardinal_graph = cardinal.graph_no_exception - graph_unit_singular = pynini.string_file(get_abs_path("data/measurements_singular.tsv")).invert() - graph_unit_plural = pynini.string_file(get_abs_path("data/measurements_plural.tsv")).invert() + graph_unit_singular = pynini.string_file( + get_abs_path("data/measurements_singular.tsv") + ).invert() + graph_unit_plural = pynini.string_file( + get_abs_path("data/measurements_plural.tsv") + ).invert() optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("menos", "\"true\"") + delete_extra_space, 0, 1 + pynutil.insert("negative: ") + pynini.cross("menos", '"true"') + delete_extra_space, + 0, + 1, ) unit_singular = convert_space(graph_unit_singular) unit_plural = convert_space(graph_unit_plural) - unit_misc = pynutil.insert("/") + pynutil.delete("por") + delete_space + convert_space(graph_unit_singular) + unit_misc = ( + pynutil.insert("/") + + pynutil.delete("por") + + delete_space + + convert_space(graph_unit_singular) + ) unit_singular = ( - pynutil.insert("units: \"") - + (unit_singular | unit_misc | pynutil.add_weight(unit_singular + delete_space + unit_misc, 0.01)) - + pynutil.insert("\"") + pynutil.insert('units: "') + + ( + unit_singular + | unit_misc + | pynutil.add_weight(unit_singular + delete_space + unit_misc, 0.01) + ) + + pynutil.insert('"') ) unit_plural = ( - pynutil.insert("units: \"") - + (unit_plural | unit_misc | pynutil.add_weight(unit_plural + delete_space + unit_misc, 0.01)) - + pynutil.insert("\"") + pynutil.insert('units: "') + + ( + unit_plural + | unit_misc + | pynutil.add_weight(unit_plural + delete_space + unit_misc, 0.01) + ) + + pynutil.insert('"') ) subgraph_decimal = ( @@ -61,9 +77,9 @@ class MeasureFst(GraphFst): subgraph_cardinal = ( pynutil.insert("cardinal { ") + optional_graph_negative - + pynutil.insert("integer: \"") + + pynutil.insert('integer: "') + ((DAMO_SIGMA - "um" - "uma") @ cardinal_graph) - + pynutil.insert("\"") + + pynutil.insert('"') + pynutil.insert(" }") + delete_extra_space + unit_plural @@ -71,9 +87,9 @@ class MeasureFst(GraphFst): subgraph_cardinal |= ( pynutil.insert("cardinal { ") + optional_graph_negative - + pynutil.insert("integer: \"") + + pynutil.insert('integer: "') + (pynini.cross("um", "1") | pynini.cross("uma", "1")) - + pynutil.insert("\"") + + pynutil.insert('"') + pynutil.insert(" }") + delete_extra_space + unit_singular diff --git a/fun_text_processing/inverse_text_normalization/pt/taggers/money.py b/fun_text_processing/inverse_text_normalization/pt/taggers/money.py index a163e8358..5e5347bd0 100644 --- a/fun_text_processing/inverse_text_normalization/pt/taggers/money.py +++ b/fun_text_processing/inverse_text_normalization/pt/taggers/money.py @@ -1,6 +1,3 @@ - - - import pynini from fun_text_processing.inverse_text_normalization.pt.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import ( @@ -35,15 +32,21 @@ class MoneyFst(GraphFst): unit_singular = pynini.string_file(get_abs_path("data/currency_singular.tsv")).invert() unit_plural = pynini.string_file(get_abs_path("data/currency_plural.tsv")).invert() - graph_unit_singular = pynutil.insert("currency: \"") + convert_space(unit_singular) + pynutil.insert("\"") - graph_unit_plural = pynutil.insert("currency: \"") + convert_space(unit_plural) + pynutil.insert("\"") + graph_unit_singular = ( + pynutil.insert('currency: "') + convert_space(unit_singular) + pynutil.insert('"') + ) + graph_unit_plural = ( + pynutil.insert('currency: "') + convert_space(unit_plural) + pynutil.insert('"') + ) - add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | (pynutil.insert("0") + DAMO_DIGIT) + add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | ( + pynutil.insert("0") + DAMO_DIGIT + ) # twelve dollars (and) fifty cents, zero cents cents_standalone = ( - pynutil.insert("morphosyntactic_features: \",\"") # always use a comma in the decimal + pynutil.insert('morphosyntactic_features: ","') # always use a comma in the decimal + insert_space - + pynutil.insert("fractional_part: \"") + + pynutil.insert('fractional_part: "') + pynini.union( pynutil.add_weight(((DAMO_SIGMA - "um" - "uma") @ cardinal_graph), -0.7) @ add_leading_zero_to_double_digit @@ -51,12 +54,12 @@ class MoneyFst(GraphFst): + pynutil.delete(pynini.union("centavos")), pynini.cross("um", "01") + delete_space + pynutil.delete(pynini.union("centavo")), ) - + pynutil.insert("\"") + + pynutil.insert('"') ) optional_cents_standalone = pynini.closure( delete_space - + pynini.closure((pynutil.delete("com") | pynutil.delete('e')) + delete_space, 0, 1) + + pynini.closure((pynutil.delete("com") | pynutil.delete("e")) + delete_space, 0, 1) + insert_space + cents_standalone, 0, @@ -67,37 +70,37 @@ class MoneyFst(GraphFst): # setenta e cinco dólares com sessenta e três ~ $75,63 optional_cents_suffix = pynini.closure( delete_extra_space - + pynutil.insert("morphosyntactic_features: \",\"") # always use a comma in the decimal + + pynutil.insert('morphosyntactic_features: ","') # always use a comma in the decimal + insert_space - + pynutil.insert("fractional_part: \"") - + pynini.closure((pynutil.delete("com") | pynutil.delete('e')) + delete_space, 0, 1) + + pynutil.insert('fractional_part: "') + + pynini.closure((pynutil.delete("com") | pynutil.delete("e")) + delete_space, 0, 1) + pynutil.add_weight(cardinal_graph @ add_leading_zero_to_double_digit, -0.7) - + pynutil.insert("\""), + + pynutil.insert('"'), 0, 1, ) graph_integer = ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + ((DAMO_SIGMA - "um" - "uma") @ cardinal_graph) - + pynutil.insert("\"") + + pynutil.insert('"') + delete_extra_space + graph_unit_plural + (optional_cents_standalone | optional_cents_suffix) ) graph_integer |= ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + (pynini.cross("um", "1") | pynini.cross("uma", "1")) - + pynutil.insert("\"") + + pynutil.insert('"') + delete_extra_space + graph_unit_singular + (optional_cents_standalone | optional_cents_suffix) ) graph_cents_standalone = pynini.union( - pynutil.insert("currency: \"R$\" integer_part: \"0\" ") + cents_standalone, + pynutil.insert('currency: "R$" integer_part: "0" ') + cents_standalone, pynutil.add_weight( - pynutil.insert("integer_part: \"0\" ") + pynutil.insert('integer_part: "0" ') + cents_standalone + delete_extra_space + pynutil.delete("de") @@ -108,7 +111,10 @@ class MoneyFst(GraphFst): ) graph_decimal = ( - graph_decimal_final + delete_extra_space + (pynutil.delete("de") + delete_space).ques + graph_unit_plural + graph_decimal_final + + delete_extra_space + + (pynutil.delete("de") + delete_space).ques + + graph_unit_plural ) graph_decimal |= graph_cents_standalone final_graph = graph_integer | graph_decimal diff --git a/fun_text_processing/inverse_text_normalization/pt/taggers/ordinal.py b/fun_text_processing/inverse_text_normalization/pt/taggers/ordinal.py index a2f30b519..2670d062b 100644 --- a/fun_text_processing/inverse_text_normalization/pt/taggers/ordinal.py +++ b/fun_text_processing/inverse_text_normalization/pt/taggers/ordinal.py @@ -1,6 +1,3 @@ - - - import pynini from fun_text_processing.inverse_text_normalization.pt.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import DAMO_SIGMA, GraphFst, delete_space @@ -33,7 +30,9 @@ class OrdinalFst(GraphFst): ordinal_graph_union = pynini.union( pynutil.add_weight(graph_digit, 0.4), - pynutil.add_weight(graph_ties + ((delete_space + graph_digit) | pynutil.insert("0")), 0.2), + pynutil.add_weight( + graph_ties + ((delete_space + graph_digit) | pynutil.insert("0")), 0.2 + ), graph_hundreds + ((delete_space + graph_ties) | pynutil.insert("0")) + ((delete_space + graph_digit) | pynutil.insert("0")), @@ -47,26 +46,32 @@ class OrdinalFst(GraphFst): # 'optional_numbers_in_front' have negative weight so we always # include them if they're there - optional_in_front = (pynutil.add_weight(ordinal_graph_union, -0.1) + delete_space.closure()).closure() + optional_in_front = ( + pynutil.add_weight(ordinal_graph_union, -0.1) + delete_space.closure() + ).closure() graph_o_suffix = optional_in_front + ordinal_graph_o graph_a_suffix = optional_in_front + ordinal_graph_a # don't convert ordinals from one to nine inclusive - graph_exception = pynini.project(pynini.union(graph_digit), 'input') - graph_o_suffix = (pynini.project(graph_o_suffix, "input") - graph_exception.arcsort()) @ graph_o_suffix - graph_a_suffix = (pynini.project(graph_a_suffix, "input") - graph_exception.arcsort()) @ graph_a_suffix + graph_exception = pynini.project(pynini.union(graph_digit), "input") + graph_o_suffix = ( + pynini.project(graph_o_suffix, "input") - graph_exception.arcsort() + ) @ graph_o_suffix + graph_a_suffix = ( + pynini.project(graph_a_suffix, "input") - graph_exception.arcsort() + ) @ graph_a_suffix graph = ( - pynutil.insert("integer: \"") + pynutil.insert('integer: "') + graph_o_suffix - + pynutil.insert("\"") - + pynutil.insert(" morphosyntactic_features: \"o\"") + + pynutil.insert('"') + + pynutil.insert(' morphosyntactic_features: "o"') ) graph |= ( - pynutil.insert("integer: \"") + pynutil.insert('integer: "') + graph_a_suffix - + pynutil.insert("\"") - + pynutil.insert(" morphosyntactic_features: \"a\"") + + pynutil.insert('"') + + pynutil.insert(' morphosyntactic_features: "a"') ) final_graph = self.add_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/pt/taggers/punctuation.py b/fun_text_processing/inverse_text_normalization/pt/taggers/punctuation.py index ef63904c2..dfe16facf 100644 --- a/fun_text_processing/inverse_text_normalization/pt/taggers/punctuation.py +++ b/fun_text_processing/inverse_text_normalization/pt/taggers/punctuation.py @@ -1,6 +1,3 @@ - - - import pynini from fun_text_processing.text_normalization.en.graph_utils import GraphFst from pynini.lib import pynutil @@ -15,9 +12,9 @@ class PunctuationFst(GraphFst): def __init__(self): super().__init__(name="punctuation", kind="classify") - s = "!#$%&\'()*+,-./:;<=>?@^_`{|}~" + s = "!#$%&'()*+,-./:;<=>?@^_`{|}~" punct = pynini.union(*s) - graph = pynutil.insert("name: \"") + punct + pynutil.insert("\"") + graph = pynutil.insert('name: "') + punct + pynutil.insert('"') self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/pt/taggers/telephone.py b/fun_text_processing/inverse_text_normalization/pt/taggers/telephone.py index 934363e82..3edda648d 100755 --- a/fun_text_processing/inverse_text_normalization/pt/taggers/telephone.py +++ b/fun_text_processing/inverse_text_normalization/pt/taggers/telephone.py @@ -1,9 +1,10 @@ - - - import pynini from fun_text_processing.inverse_text_normalization.pt.utils import get_abs_path -from fun_text_processing.text_normalization.en.graph_utils import GraphFst, delete_space, insert_space +from fun_text_processing.text_normalization.en.graph_utils import ( + GraphFst, + delete_space, + insert_space, +) from pynini.lib import pynutil @@ -111,9 +112,13 @@ class TelephoneFst(GraphFst): + double_digits ) - number_part = pynini.union(eleven_digit_graph, ten_digit_graph, nine_digit_graph, eight_digit_graph) + number_part = pynini.union( + eleven_digit_graph, ten_digit_graph, nine_digit_graph, eight_digit_graph + ) - number_part = pynutil.insert("number_part: \"") + pynini.invert(number_part) + pynutil.insert("\"") + number_part = ( + pynutil.insert('number_part: "') + pynini.invert(number_part) + pynutil.insert('"') + ) graph = number_part final_graph = self.add_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/pt/taggers/time.py b/fun_text_processing/inverse_text_normalization/pt/taggers/time.py index 826c819aa..4b59e76e6 100755 --- a/fun_text_processing/inverse_text_normalization/pt/taggers/time.py +++ b/fun_text_processing/inverse_text_normalization/pt/taggers/time.py @@ -1,9 +1,10 @@ - - - import pynini from fun_text_processing.inverse_text_normalization.pt.utils import get_abs_path -from fun_text_processing.text_normalization.en.graph_utils import GraphFst, delete_space, insert_space +from fun_text_processing.text_normalization.en.graph_utils import ( + GraphFst, + delete_space, + insert_space, +) from pynini.lib import pynutil @@ -74,37 +75,43 @@ class TimeFst(GraphFst): # à uma e meia -> 1:30 # às três e meia -> 3:30 graph_hours_at_prefix_singular = ( - pynutil.insert("morphosyntactic_features: \"") + pynutil.insert('morphosyntactic_features: "') + (pynini.cross("à", "à") | pynini.cross("a", "à")) - + pynutil.insert("\" ") + + pynutil.insert('" ') + delete_space ) graph_hours_at_singular = ( graph_hours_at_prefix_singular - + pynutil.insert("hours: \"") + + pynutil.insert('hours: "') + graph_uma - + pynutil.insert("\"") + + pynutil.insert('"') + (delete_space + graph_hour).ques ) graph_hours_at_prefix_plural = ( - pynutil.insert("morphosyntactic_features: \"") + pynutil.insert('morphosyntactic_features: "') + (pynini.cross("às", "às") | pynini.cross("as", "às")) - + pynutil.insert("\" ") + + pynutil.insert('" ') + delete_space ) graph_hours_at_plural = ( graph_hours_at_prefix_plural - + pynutil.insert("hours: \"") + + pynutil.insert('hours: "') + graph_2_to_23 - + pynutil.insert("\"") + + pynutil.insert('"') + (delete_space + graph_hour).ques ) final_graph_hour_at = graph_hours_at_singular | graph_hours_at_plural - graph_minutes_component_without_zero = graph_e + graph_1_to_59 + (delete_space + graph_minute).ques - graph_minutes_component_without_zero |= graph_e_meia + pynutil.delete(delete_space + pynini.accep("hora")).ques + graph_minutes_component_without_zero = ( + graph_e + graph_1_to_59 + (delete_space + graph_minute).ques + ) + graph_minutes_component_without_zero |= ( + graph_e_meia + pynutil.delete(delete_space + pynini.accep("hora")).ques + ) final_graph_minute = ( - pynutil.insert(" minutes: \"") + graph_minutes_component_without_zero + pynutil.insert("\"") + pynutil.insert(' minutes: "') + + graph_minutes_component_without_zero + + pynutil.insert('"') ) graph_hm = final_graph_hour_at + final_graph_minute @@ -112,43 +119,45 @@ class TimeFst(GraphFst): # à uma hora -> 1:00 graph_hours_at_singular_with_hour = ( graph_hours_at_prefix_singular - + pynutil.insert("hours: \"") + + pynutil.insert('hours: "') + graph_uma - + pynutil.insert("\"") + + pynutil.insert('"') + delete_space + graph_hour ) graph_hours_at_plural_with_hour = ( graph_hours_at_prefix_plural - + pynutil.insert("hours: \"") + + pynutil.insert('hours: "') + graph_2_to_23 - + pynutil.insert("\"") + + pynutil.insert('"') + delete_space + graph_hour ) - graph_hm |= (graph_hours_at_singular_with_hour | graph_hours_at_plural_with_hour) + pynutil.insert( - " minutes: \"00\"", weight=0.2 - ) + graph_hm |= ( + graph_hours_at_singular_with_hour | graph_hours_at_plural_with_hour + ) + pynutil.insert(' minutes: "00"', weight=0.2) # meio dia e meia -> 12:30 # meia noite e meia -> 0:30 graph_minutes_without_zero = ( - pynutil.insert(" minutes: \"") + graph_minutes_component_without_zero + pynutil.insert("\"") + pynutil.insert(' minutes: "') + + graph_minutes_component_without_zero + + pynutil.insert('"') ) graph_meio_min = ( - pynutil.insert("hours: \"") + pynutil.insert('hours: "') + (graph_meio_dia | graph_meia_noite) - + pynutil.insert("\"") + + pynutil.insert('"') + graph_minutes_without_zero ) graph_meio_min |= ( - pynutil.insert("hours: \"") + pynutil.insert('hours: "') + graph_meio_dia - + pynutil.insert("\" minutes: \"") + + pynutil.insert('" minutes: "') + graph_e_meio - + pynutil.insert("\"") + + pynutil.insert('"') ) graph_hm |= graph_meio_min @@ -165,35 +174,46 @@ class TimeFst(GraphFst): graph_para_o = pynutil.delete("para") | pynutil.delete("para o") | pynutil.delete("pro") graph_pra_min = ( - pynutil.insert("morphosyntactic_features: \"") - + (pynini.cross("à", "à") | pynini.cross("às", "às") | pynini.cross("a", "à") | pynini.cross("as", "às")) - + pynutil.insert("\" ") + pynutil.insert('morphosyntactic_features: "') + + ( + pynini.cross("à", "à") + | pynini.cross("às", "às") + | pynini.cross("a", "à") + | pynini.cross("as", "às") + ) + + pynutil.insert('" ') + delete_space ) graph_pra_min += ( - pynutil.insert("minutes: \"") + pynutil.insert('minutes: "') + (graph_1_to_59 @ graph_minutes_to) - + pynutil.insert("\" ") + + pynutil.insert('" ') + (delete_space + graph_minute).ques ) graph_pra_hour = ( - pynutil.insert("hours: \"") + pynutil.insert('hours: "') + (graph_2_to_23 @ graph_hours_to) - + pynutil.insert("\"") + + pynutil.insert('"') + (delete_space + graph_hour).ques ) - graph_pra_hour |= pynutil.insert("hours: \"") + (graph_meia_noite @ graph_hours_to) + pynutil.insert("\"") + graph_pra_hour |= ( + pynutil.insert('hours: "') + (graph_meia_noite @ graph_hours_to) + pynutil.insert('"') + ) graph_pra = graph_pra_min + delete_space + graph_para_a + delete_space + graph_pra_hour # às quinze pro meio dia -> às 11:45 graph_pro = graph_pra_min + delete_space + graph_para_o + delete_space - graph_pro += pynutil.insert(" hours: \"") + (graph_meio_dia @ graph_hours_to) + pynutil.insert("\"") + graph_pro += ( + pynutil.insert(' hours: "') + (graph_meio_dia @ graph_hours_to) + pynutil.insert('"') + ) graph_mh = graph_pra | graph_pro # optional suffix - final_suffix = pynutil.insert("suffix: \"") + (graph_suffix_am | graph_suffix_pm) + pynutil.insert("\"") + final_suffix = ( + pynutil.insert('suffix: "') + (graph_suffix_am | graph_suffix_pm) + pynutil.insert('"') + ) final_suffix_optional = pynini.closure(delete_space + insert_space + final_suffix, 0, 1) final_graph = pynini.union((graph_hm | graph_mh) + final_suffix_optional).optimize() diff --git a/fun_text_processing/inverse_text_normalization/pt/taggers/tokenize_and_classify.py b/fun_text_processing/inverse_text_normalization/pt/taggers/tokenize_and_classify.py index 7ad13224b..ac84c0c34 100644 --- a/fun_text_processing/inverse_text_normalization/pt/taggers/tokenize_and_classify.py +++ b/fun_text_processing/inverse_text_normalization/pt/taggers/tokenize_and_classify.py @@ -1,6 +1,3 @@ - - - import os import pynini @@ -83,10 +80,16 @@ class ClassifyFst(GraphFst): | pynutil.add_weight(word_graph, 100) ) - punct = pynutil.insert("tokens { ") + pynutil.add_weight(punct_graph, weight=1.1) + pynutil.insert(" }") + punct = ( + pynutil.insert("tokens { ") + + pynutil.add_weight(punct_graph, weight=1.1) + + pynutil.insert(" }") + ) token = pynutil.insert("tokens { ") + classify + pynutil.insert(" }") token_plus_punct = ( - pynini.closure(punct + pynutil.insert(" ")) + token + pynini.closure(pynutil.insert(" ") + punct) + pynini.closure(punct + pynutil.insert(" ")) + + token + + pynini.closure(pynutil.insert(" ") + punct) ) graph = token_plus_punct + pynini.closure(delete_extra_space + token_plus_punct) diff --git a/fun_text_processing/inverse_text_normalization/pt/taggers/whitelist.py b/fun_text_processing/inverse_text_normalization/pt/taggers/whitelist.py index 9aa2ffc0b..2ffb4387f 100644 --- a/fun_text_processing/inverse_text_normalization/pt/taggers/whitelist.py +++ b/fun_text_processing/inverse_text_normalization/pt/taggers/whitelist.py @@ -1,6 +1,3 @@ - - - import pynini from fun_text_processing.inverse_text_normalization.pt.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import GraphFst, convert_space @@ -18,5 +15,5 @@ class WhiteListFst(GraphFst): super().__init__(name="whitelist", kind="classify") whitelist = pynini.string_file(get_abs_path("data/whitelist.tsv")).invert() - graph = pynutil.insert("name: \"") + convert_space(whitelist) + pynutil.insert("\"") + graph = pynutil.insert('name: "') + convert_space(whitelist) + pynutil.insert('"') self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/pt/taggers/word.py b/fun_text_processing/inverse_text_normalization/pt/taggers/word.py index c6bae2503..9758f6885 100644 --- a/fun_text_processing/inverse_text_normalization/pt/taggers/word.py +++ b/fun_text_processing/inverse_text_normalization/pt/taggers/word.py @@ -1,6 +1,3 @@ - - - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_SPACE, GraphFst from pynini.lib import pynutil @@ -14,5 +11,5 @@ class WordFst(GraphFst): def __init__(self): super().__init__(name="word", kind="classify") - word = pynutil.insert("name: \"") + pynini.closure(DAMO_NOT_SPACE, 1) + pynutil.insert("\"") + word = pynutil.insert('name: "') + pynini.closure(DAMO_NOT_SPACE, 1) + pynutil.insert('"') self.fst = word.optimize() diff --git a/fun_text_processing/inverse_text_normalization/pt/utils.py b/fun_text_processing/inverse_text_normalization/pt/utils.py index ed78496f8..c6fdb7435 100644 --- a/fun_text_processing/inverse_text_normalization/pt/utils.py +++ b/fun_text_processing/inverse_text_normalization/pt/utils.py @@ -1,6 +1,3 @@ - - - import os @@ -13,4 +10,4 @@ def get_abs_path(rel_path): Returns absolute path """ - return os.path.dirname(os.path.abspath(__file__)) + '/' + rel_path + return os.path.dirname(os.path.abspath(__file__)) + "/" + rel_path diff --git a/fun_text_processing/inverse_text_normalization/pt/verbalizers/__init__.py b/fun_text_processing/inverse_text_normalization/pt/verbalizers/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/pt/verbalizers/__init__.py +++ b/fun_text_processing/inverse_text_normalization/pt/verbalizers/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/inverse_text_normalization/pt/verbalizers/cardinal.py b/fun_text_processing/inverse_text_normalization/pt/verbalizers/cardinal.py index 1befda1b5..005fae0bc 100644 --- a/fun_text_processing/inverse_text_normalization/pt/verbalizers/cardinal.py +++ b/fun_text_processing/inverse_text_normalization/pt/verbalizers/cardinal.py @@ -1,8 +1,9 @@ - - - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -17,9 +18,9 @@ class CardinalFst(GraphFst): optional_sign = pynini.closure( pynutil.delete("negative:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + DAMO_NOT_QUOTE - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space, 0, 1, @@ -27,9 +28,9 @@ class CardinalFst(GraphFst): graph = ( pynutil.delete("integer:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) self.numbers = graph graph = optional_sign + graph diff --git a/fun_text_processing/inverse_text_normalization/pt/verbalizers/date.py b/fun_text_processing/inverse_text_normalization/pt/verbalizers/date.py index 759c43487..ef18b1b45 100644 --- a/fun_text_processing/inverse_text_normalization/pt/verbalizers/date.py +++ b/fun_text_processing/inverse_text_normalization/pt/verbalizers/date.py @@ -1,6 +1,3 @@ - - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_NOT_QUOTE, @@ -23,23 +20,23 @@ class DateFst(GraphFst): month = ( pynutil.delete("month:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) day = ( pynutil.delete("day:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) year = ( pynutil.delete("year:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) # day month @@ -56,7 +53,7 @@ class DateFst(GraphFst): + delete_space + pynutil.insert("/") + month - + pynutil.delete(" morphosyntactic_features: \"/\"") + + pynutil.delete(' morphosyntactic_features: "/"') + (delete_space + pynutil.insert("/") + year).ques ) @@ -64,9 +61,9 @@ class DateFst(GraphFst): pynutil.delete("preserve_order:") + delete_space + pynutil.delete("true") + delete_space | pynutil.delete("field_order:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + DAMO_NOT_QUOTE - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space ) diff --git a/fun_text_processing/inverse_text_normalization/pt/verbalizers/decimal.py b/fun_text_processing/inverse_text_normalization/pt/verbalizers/decimal.py index f8366f1f9..346e15b4a 100644 --- a/fun_text_processing/inverse_text_normalization/pt/verbalizers/decimal.py +++ b/fun_text_processing/inverse_text_normalization/pt/verbalizers/decimal.py @@ -1,8 +1,9 @@ - - - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -17,35 +18,35 @@ class DecimalFst(GraphFst): def __init__(self): super().__init__(name="decimal", kind="verbalize") - optionl_sign = pynini.closure(pynini.cross("negative: \"true\"", "-") + delete_space, 0, 1) + optionl_sign = pynini.closure(pynini.cross('negative: "true"', "-") + delete_space, 0, 1) integer = ( pynutil.delete("integer_part:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_integer = pynini.closure(integer + delete_space, 0, 1) - decimal_point = pynini.cross("morphosyntactic_features: \",\"", ",") - decimal_point |= pynini.cross("morphosyntactic_features: \".\"", ".") + decimal_point = pynini.cross('morphosyntactic_features: ","', ",") + decimal_point |= pynini.cross('morphosyntactic_features: "."', ".") fractional = ( decimal_point + delete_space + pynutil.delete("fractional_part:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_fractional = pynini.closure(fractional + delete_space, 0, 1) quantity = ( pynutil.delete("quantity:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_quantity = pynini.closure(pynutil.insert(" ") + quantity + delete_space, 0, 1) graph = optional_integer + optional_fractional + optional_quantity diff --git a/fun_text_processing/inverse_text_normalization/pt/verbalizers/electronic.py b/fun_text_processing/inverse_text_normalization/pt/verbalizers/electronic.py index e4cb1a718..bcdf689a4 100644 --- a/fun_text_processing/inverse_text_normalization/pt/verbalizers/electronic.py +++ b/fun_text_processing/inverse_text_normalization/pt/verbalizers/electronic.py @@ -1,8 +1,9 @@ - - - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -18,24 +19,24 @@ class ElectronicFst(GraphFst): user_name = ( pynutil.delete("username:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) domain = ( pynutil.delete("domain:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) protocol = ( pynutil.delete("protocol:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) graph = user_name + delete_space + pynutil.insert("@") + domain diff --git a/fun_text_processing/inverse_text_normalization/pt/verbalizers/measure.py b/fun_text_processing/inverse_text_normalization/pt/verbalizers/measure.py index c0b83dddd..b2005c305 100644 --- a/fun_text_processing/inverse_text_normalization/pt/verbalizers/measure.py +++ b/fun_text_processing/inverse_text_normalization/pt/verbalizers/measure.py @@ -1,6 +1,3 @@ - - - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_CHAR, GraphFst, delete_space from pynini.lib import pynutil @@ -18,13 +15,13 @@ class MeasureFst(GraphFst): def __init__(self, decimal: GraphFst, cardinal: GraphFst): super().__init__(name="measure", kind="verbalize") - optional_sign = pynini.closure(pynini.cross("negative: \"true\"", "-"), 0, 1) + optional_sign = pynini.closure(pynini.cross('negative: "true"', "-"), 0, 1) unit = ( pynutil.delete("units:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space ) graph_decimal = ( diff --git a/fun_text_processing/inverse_text_normalization/pt/verbalizers/money.py b/fun_text_processing/inverse_text_normalization/pt/verbalizers/money.py index 9cde5f47e..ae8baecbb 100644 --- a/fun_text_processing/inverse_text_normalization/pt/verbalizers/money.py +++ b/fun_text_processing/inverse_text_normalization/pt/verbalizers/money.py @@ -1,8 +1,10 @@ - - - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_CHAR, GraphFst, delete_space, insert_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_CHAR, + GraphFst, + delete_space, + insert_space, +) from pynini.lib import pynutil @@ -20,9 +22,9 @@ class MoneyFst(GraphFst): unit = ( pynutil.delete("currency:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) graph = unit + delete_space + insert_space + decimal.numbers delete_tokens = self.delete_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/pt/verbalizers/ordinal.py b/fun_text_processing/inverse_text_normalization/pt/verbalizers/ordinal.py index 0db15410d..1dc3971c8 100644 --- a/fun_text_processing/inverse_text_normalization/pt/verbalizers/ordinal.py +++ b/fun_text_processing/inverse_text_normalization/pt/verbalizers/ordinal.py @@ -1,8 +1,9 @@ - - - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -17,14 +18,14 @@ class OrdinalFst(GraphFst): graph = ( pynutil.delete("integer:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) replace_suffix = pynini.union( - pynini.cross(" morphosyntactic_features: \"o\"", "º"), - pynini.cross(" morphosyntactic_features: \"a\"", "ª"), + pynini.cross(' morphosyntactic_features: "o"', "º"), + pynini.cross(' morphosyntactic_features: "a"', "ª"), ) graph = graph + replace_suffix diff --git a/fun_text_processing/inverse_text_normalization/pt/verbalizers/telephone.py b/fun_text_processing/inverse_text_normalization/pt/verbalizers/telephone.py index ef2ba51b7..8c30c0843 100644 --- a/fun_text_processing/inverse_text_normalization/pt/verbalizers/telephone.py +++ b/fun_text_processing/inverse_text_normalization/pt/verbalizers/telephone.py @@ -1,6 +1,3 @@ - - - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst from pynini.lib import pynutil @@ -16,6 +13,10 @@ class TelephoneFst(GraphFst): def __init__(self): super().__init__(name="telephone", kind="verbalize") - number_part = pynutil.delete("number_part: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + number_part = ( + pynutil.delete('number_part: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + ) delete_tokens = self.delete_tokens(number_part) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/inverse_text_normalization/pt/verbalizers/time.py b/fun_text_processing/inverse_text_normalization/pt/verbalizers/time.py index 9e3597ef3..0a408cde9 100755 --- a/fun_text_processing/inverse_text_normalization/pt/verbalizers/time.py +++ b/fun_text_processing/inverse_text_normalization/pt/verbalizers/time.py @@ -1,6 +1,3 @@ - - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_DIGIT, @@ -21,14 +18,16 @@ class TimeFst(GraphFst): def __init__(self): super().__init__(name="time", kind="verbalize") - add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | (pynutil.insert("0") + DAMO_DIGIT) + add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | ( + pynutil.insert("0") + DAMO_DIGIT + ) prefix = ( pynutil.delete("morphosyntactic_features:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space + insert_space ) @@ -37,25 +36,25 @@ class TimeFst(GraphFst): hour = ( pynutil.delete("hours:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_DIGIT, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) minute = ( pynutil.delete("minutes:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_DIGIT, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) suffix = ( delete_space + insert_space + pynutil.delete("suffix:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_suffix = pynini.closure(suffix, 0, 1) diff --git a/fun_text_processing/inverse_text_normalization/pt/verbalizers/verbalize.py b/fun_text_processing/inverse_text_normalization/pt/verbalizers/verbalize.py index fbc0873c6..f8d42636d 100644 --- a/fun_text_processing/inverse_text_normalization/pt/verbalizers/verbalize.py +++ b/fun_text_processing/inverse_text_normalization/pt/verbalizers/verbalize.py @@ -1,6 +1,3 @@ - - - from fun_text_processing.inverse_text_normalization.pt.verbalizers.cardinal import CardinalFst from fun_text_processing.inverse_text_normalization.pt.verbalizers.date import DateFst from fun_text_processing.inverse_text_normalization.pt.verbalizers.decimal import DecimalFst diff --git a/fun_text_processing/inverse_text_normalization/pt/verbalizers/verbalize_final.py b/fun_text_processing/inverse_text_normalization/pt/verbalizers/verbalize_final.py index 34f0d694c..c0fc8b2cb 100644 --- a/fun_text_processing/inverse_text_normalization/pt/verbalizers/verbalize_final.py +++ b/fun_text_processing/inverse_text_normalization/pt/verbalizers/verbalize_final.py @@ -1,10 +1,11 @@ - - - import pynini from fun_text_processing.inverse_text_normalization.pt.verbalizers.verbalize import VerbalizeFst from fun_text_processing.inverse_text_normalization.pt.verbalizers.word import WordFst -from fun_text_processing.text_normalization.en.graph_utils import GraphFst, delete_extra_space, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + GraphFst, + delete_extra_space, + delete_space, +) from pynini.lib import pynutil diff --git a/fun_text_processing/inverse_text_normalization/pt/verbalizers/whitelist.py b/fun_text_processing/inverse_text_normalization/pt/verbalizers/whitelist.py index 4a25f67d7..70240ef37 100644 --- a/fun_text_processing/inverse_text_normalization/pt/verbalizers/whitelist.py +++ b/fun_text_processing/inverse_text_normalization/pt/verbalizers/whitelist.py @@ -1,6 +1,10 @@ - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_CHAR, DAMO_SIGMA, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_CHAR, + DAMO_SIGMA, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -15,9 +19,9 @@ class WhiteListFst(GraphFst): graph = ( pynutil.delete("name:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) - graph = graph @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + graph = graph @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/pt/verbalizers/word.py b/fun_text_processing/inverse_text_normalization/pt/verbalizers/word.py index 590dc65bb..1d33cd156 100644 --- a/fun_text_processing/inverse_text_normalization/pt/verbalizers/word.py +++ b/fun_text_processing/inverse_text_normalization/pt/verbalizers/word.py @@ -1,8 +1,10 @@ - - - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_CHAR, DAMO_SIGMA, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_CHAR, + DAMO_SIGMA, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -15,7 +17,13 @@ class WordFst(GraphFst): def __init__(self): super().__init__(name="word", kind="verbalize") chars = pynini.closure(DAMO_CHAR - " ", 1) - char = pynutil.delete("name:") + delete_space + pynutil.delete("\"") + chars + pynutil.delete("\"") - graph = char @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + char = ( + pynutil.delete("name:") + + delete_space + + pynutil.delete('"') + + chars + + pynutil.delete('"') + ) + graph = char @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ru/__init__.py b/fun_text_processing/inverse_text_normalization/ru/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/ru/__init__.py +++ b/fun_text_processing/inverse_text_normalization/ru/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/inverse_text_normalization/ru/taggers/cardinal.py b/fun_text_processing/inverse_text_normalization/ru/taggers/cardinal.py index 029997bfd..086ab4e9e 100644 --- a/fun_text_processing/inverse_text_normalization/ru/taggers/cardinal.py +++ b/fun_text_processing/inverse_text_normalization/ru/taggers/cardinal.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_DIGIT, GraphFst, insert_space from pynini.lib import pynutil @@ -7,7 +5,7 @@ from pynini.lib import pynutil class CardinalFst(GraphFst): """ - Finite state transducer for classifying cardinals, e.g. + Finite state transducer for classifying cardinals, e.g. "тысяча один" -> cardinal { integer: "1 001" } Args: @@ -23,11 +21,11 @@ class CardinalFst(GraphFst): self.graph = graph.invert().optimize() optional_sign = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("минус ", "\"-\"") + insert_space, 0, 1 + pynutil.insert("negative: ") + pynini.cross("минус ", '"-"') + insert_space, 0, 1 ) # do not invert numbers less than 10 graph = pynini.compose(graph, DAMO_DIGIT ** (2, ...)) - graph = optional_sign + pynutil.insert("integer: \"") + graph + pynutil.insert("\"") + graph = optional_sign + pynutil.insert('integer: "') + graph + pynutil.insert('"') graph = self.add_tokens(graph) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ru/taggers/date.py b/fun_text_processing/inverse_text_normalization/ru/taggers/date.py index 57b2d170c..1de12243d 100644 --- a/fun_text_processing/inverse_text_normalization/ru/taggers/date.py +++ b/fun_text_processing/inverse_text_normalization/ru/taggers/date.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import GraphFst from pynini.lib import pynutil @@ -20,5 +18,5 @@ class DateFst(GraphFst): super().__init__(name="date", kind="classify", deterministic=deterministic) graph = pynini.invert(tn_date.final_graph).optimize() - graph = self.add_tokens(pynutil.insert("day: \"") + graph + pynutil.insert("\"")) + graph = self.add_tokens(pynutil.insert('day: "') + graph + pynutil.insert('"')) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ru/taggers/decimals.py b/fun_text_processing/inverse_text_normalization/ru/taggers/decimals.py index 22a2b89da..71a39e678 100644 --- a/fun_text_processing/inverse_text_normalization/ru/taggers/decimals.py +++ b/fun_text_processing/inverse_text_normalization/ru/taggers/decimals.py @@ -1,7 +1,9 @@ - - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_SPACE, GraphFst, delete_extra_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_SPACE, + GraphFst, + delete_extra_space, +) from pynini.lib import pynutil @@ -20,17 +22,25 @@ class DecimalFst(GraphFst): super().__init__(name="decimal", kind="classify", deterministic=deterministic) optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("минус", "\"true\"") + delete_extra_space, 0, 1 + pynutil.insert("negative: ") + pynini.cross("минус", '"true"') + delete_extra_space, + 0, + 1, ) graph_fractional_part = pynini.invert(tn_decimal.graph_fractional).optimize() graph_integer_part = pynini.invert(tn_decimal.integer_part).optimize() optional_graph_quantity = pynini.invert(tn_decimal.optional_quantity).optimize() - graph_fractional = pynutil.insert("fractional_part: \"") + graph_fractional_part + pynutil.insert("\"") - graph_integer = pynutil.insert("integer_part: \"") + graph_integer_part + pynutil.insert("\"") - optional_graph_quantity = pynutil.insert("quantity: \"") + optional_graph_quantity + pynutil.insert("\"") - optional_graph_quantity = pynini.closure(pynini.accep(DAMO_SPACE) + optional_graph_quantity, 0, 1) + graph_fractional = ( + pynutil.insert('fractional_part: "') + graph_fractional_part + pynutil.insert('"') + ) + graph_integer = pynutil.insert('integer_part: "') + graph_integer_part + pynutil.insert('"') + optional_graph_quantity = ( + pynutil.insert('quantity: "') + optional_graph_quantity + pynutil.insert('"') + ) + optional_graph_quantity = pynini.closure( + pynini.accep(DAMO_SPACE) + optional_graph_quantity, 0, 1 + ) self.final_graph_wo_sign = ( graph_integer + pynini.accep(DAMO_SPACE) + graph_fractional + optional_graph_quantity diff --git a/fun_text_processing/inverse_text_normalization/ru/taggers/electronic.py b/fun_text_processing/inverse_text_normalization/ru/taggers/electronic.py index c59f225ae..f2277421d 100644 --- a/fun_text_processing/inverse_text_normalization/ru/taggers/electronic.py +++ b/fun_text_processing/inverse_text_normalization/ru/taggers/electronic.py @@ -1,5 +1,3 @@ - - from fun_text_processing.text_normalization.en.graph_utils import GraphFst from pynini.lib import pynutil @@ -20,6 +18,6 @@ class ElectronicFst(GraphFst): graph = tn_electronic.final_graph graph = graph.invert().optimize() - graph = pynutil.insert("username: \"") + graph + pynutil.insert("\"") + graph = pynutil.insert('username: "') + graph + pynutil.insert('"') graph = self.add_tokens(graph) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ru/taggers/measure.py b/fun_text_processing/inverse_text_normalization/ru/taggers/measure.py index 920316b4e..b70fc1922 100644 --- a/fun_text_processing/inverse_text_normalization/ru/taggers/measure.py +++ b/fun_text_processing/inverse_text_normalization/ru/taggers/measure.py @@ -1,5 +1,3 @@ - - from fun_text_processing.text_normalization.en.graph_utils import GraphFst from pynini.lib import pynutil @@ -20,6 +18,6 @@ class MeasureFst(GraphFst): tn_measure = tn_measure.tagger_graph_default @ tn_measure.verbalizer_graph graph = tn_measure.invert().optimize() - graph = pynutil.insert("cardinal { integer: \"") + graph + pynutil.insert("\" }") + graph = pynutil.insert('cardinal { integer: "') + graph + pynutil.insert('" }') graph = self.add_tokens(graph) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ru/taggers/money.py b/fun_text_processing/inverse_text_normalization/ru/taggers/money.py index 89a0b9630..3edd6ffb1 100644 --- a/fun_text_processing/inverse_text_normalization/ru/taggers/money.py +++ b/fun_text_processing/inverse_text_normalization/ru/taggers/money.py @@ -1,5 +1,3 @@ - - from fun_text_processing.text_normalization.en.graph_utils import GraphFst from pynini.lib import pynutil @@ -20,6 +18,6 @@ class MoneyFst(GraphFst): graph = tn_money.final_graph graph = graph.invert().optimize() - graph = pynutil.insert("integer_part: \"") + graph + pynutil.insert("\"") + graph = pynutil.insert('integer_part: "') + graph + pynutil.insert('"') graph = self.add_tokens(graph) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ru/taggers/ordinal.py b/fun_text_processing/inverse_text_normalization/ru/taggers/ordinal.py index b3b1a7814..8695dcb3b 100644 --- a/fun_text_processing/inverse_text_normalization/ru/taggers/ordinal.py +++ b/fun_text_processing/inverse_text_normalization/ru/taggers/ordinal.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_DIGIT, GraphFst from pynini.lib import pynutil @@ -26,6 +24,6 @@ class OrdinalFst(GraphFst): # do not invert numbers less than 10 graph = pynini.compose(graph, DAMO_DIGIT ** (2, ...)) - graph = pynutil.insert("integer: \"") + graph + pynutil.insert("\"") + graph = pynutil.insert('integer: "') + graph + pynutil.insert('"') graph = self.add_tokens(graph) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ru/taggers/telephone.py b/fun_text_processing/inverse_text_normalization/ru/taggers/telephone.py index aa7370a65..be09c9956 100644 --- a/fun_text_processing/inverse_text_normalization/ru/taggers/telephone.py +++ b/fun_text_processing/inverse_text_normalization/ru/taggers/telephone.py @@ -1,5 +1,3 @@ - - from fun_text_processing.text_normalization.en.graph_utils import GraphFst from pynini.lib import pynutil @@ -20,6 +18,6 @@ class TelephoneFst(GraphFst): tn_telephone = tn_telephone.final_graph graph = tn_telephone.invert().optimize() - graph = pynutil.insert("number_part: \"") + graph + pynutil.insert("\"") + graph = pynutil.insert('number_part: "') + graph + pynutil.insert('"') graph = self.add_tokens(graph) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ru/taggers/time.py b/fun_text_processing/inverse_text_normalization/ru/taggers/time.py index b725a6e82..5a5f2880e 100644 --- a/fun_text_processing/inverse_text_normalization/ru/taggers/time.py +++ b/fun_text_processing/inverse_text_normalization/ru/taggers/time.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_SPACE, GraphFst from fun_text_processing.text_normalization.ru.verbalizers.time import TimeFst as TNTimeVerbalizer @@ -24,30 +22,32 @@ class TimeFst(GraphFst): tn_time_verbalizer = TNTimeVerbalizer().graph tn_time_graph_preserve_order = pynini.compose(tn_time_tagger, tn_time_verbalizer).optimize() graph_preserve_order = pynini.invert(tn_time_graph_preserve_order).optimize() - graph_preserve_order = pynutil.insert("hours: \"") + graph_preserve_order + pynutil.insert("\"") + graph_preserve_order = ( + pynutil.insert('hours: "') + graph_preserve_order + pynutil.insert('"') + ) # "пятнадцать минут шестого" -> 17:15 # Requires permutations for the correct verbalization m_next_h = ( - pynutil.insert("minutes: \"") + pynutil.insert('minutes: "') + pynini.invert(tn_time.minutes).optimize() - + pynutil.insert("\"") + + pynutil.insert('"') + pynini.accep(DAMO_SPACE) - + pynutil.insert("hours: \"") + + pynutil.insert('hours: "') + pynini.invert(tn_time.increment_hour_ordinal).optimize() - + pynutil.insert("\"") + + pynutil.insert('"') ).optimize() # "без пятнадцати минут шесть" -> 17:45 # Requires permutation for the correct verbalization m_to_h = ( - pynini.cross("без ", "minutes: \"") + pynini.cross("без ", 'minutes: "') + pynini.invert(tn_time.mins_to_h) - + pynutil.insert("\"") + + pynutil.insert('"') + pynini.accep(DAMO_SPACE) - + pynutil.insert("hours: \"") + + pynutil.insert('hours: "') + pynini.invert(tn_time.increment_hour_cardinal).optimize() - + pynutil.insert("\"") + + pynutil.insert('"') ) graph_reserve_order = m_next_h | m_to_h diff --git a/fun_text_processing/inverse_text_normalization/ru/taggers/tokenize_and_classify.py b/fun_text_processing/inverse_text_normalization/ru/taggers/tokenize_and_classify.py index ee5128c84..6f7f231f3 100644 --- a/fun_text_processing/inverse_text_normalization/ru/taggers/tokenize_and_classify.py +++ b/fun_text_processing/inverse_text_normalization/ru/taggers/tokenize_and_classify.py @@ -1,4 +1,3 @@ - import os import pynini @@ -20,7 +19,9 @@ from fun_text_processing.text_normalization.en.graph_utils import ( delete_space, generator_main, ) -from fun_text_processing.text_normalization.ru.taggers.tokenize_and_classify import ClassifyFst as TNClassifyFst +from fun_text_processing.text_normalization.ru.taggers.tokenize_and_classify import ( + ClassifyFst as TNClassifyFst, +) from pynini.lib import pynutil import logging @@ -29,7 +30,7 @@ import logging class ClassifyFst(GraphFst): """ Final class that composes all other classification grammars. This class can process an entire sentence, that is lower cased. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. Args: @@ -50,7 +51,7 @@ class ClassifyFst(GraphFst): else: logging.info(f"Creating ClassifyFst grammars. This might take some time...") tn_classify = TNClassifyFst( - input_case='cased', deterministic=False, cache_dir=cache_dir, overwrite_cache=True + input_case="cased", deterministic=False, cache_dir=cache_dir, overwrite_cache=True ) cardinal = CardinalFst(tn_cardinal=tn_classify.cardinal) cardinal_graph = cardinal.fst @@ -85,13 +86,21 @@ class ClassifyFst(GraphFst): | pynutil.add_weight(word_graph, 100) ) - punct = pynutil.insert("tokens { ") + pynutil.add_weight(punct_graph, weight=1.1) + pynutil.insert(" }") + punct = ( + pynutil.insert("tokens { ") + + pynutil.add_weight(punct_graph, weight=1.1) + + pynutil.insert(" }") + ) token = pynutil.insert("tokens { ") + classify + pynutil.insert(" }") token_plus_punct = ( - pynini.closure(punct + pynutil.insert(" ")) + token + pynini.closure(pynutil.insert(" ") + punct) + pynini.closure(punct + pynutil.insert(" ")) + + token + + pynini.closure(pynutil.insert(" ") + punct) ) - graph = token_plus_punct + pynini.closure(pynutil.add_weight(delete_extra_space, 1.1) + token_plus_punct) + graph = token_plus_punct + pynini.closure( + pynutil.add_weight(delete_extra_space, 1.1) + token_plus_punct + ) graph = delete_space + graph + delete_space self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ru/taggers/whitelist.py b/fun_text_processing/inverse_text_normalization/ru/taggers/whitelist.py index 601e92923..8564ff029 100644 --- a/fun_text_processing/inverse_text_normalization/ru/taggers/whitelist.py +++ b/fun_text_processing/inverse_text_normalization/ru/taggers/whitelist.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import GraphFst, convert_space from fun_text_processing.text_normalization.ru.utils import get_abs_path @@ -19,5 +17,5 @@ class WhiteListFst(GraphFst): def __init__(self, deterministic: bool = True): super().__init__(name="whitelist", kind="classify", deterministic=deterministic) whitelist = pynini.string_file(get_abs_path("data/whitelist.tsv")).invert() - graph = pynutil.insert("name: \"") + convert_space(whitelist) + pynutil.insert("\"") + graph = pynutil.insert('name: "') + convert_space(whitelist) + pynutil.insert('"') self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ru/verbalizers/cardinal.py b/fun_text_processing/inverse_text_normalization/ru/verbalizers/cardinal.py index ab5afee81..d83043e25 100644 --- a/fun_text_processing/inverse_text_normalization/ru/verbalizers/cardinal.py +++ b/fun_text_processing/inverse_text_normalization/ru/verbalizers/cardinal.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -20,16 +23,19 @@ class CardinalFst(GraphFst): optional_sign = pynini.closure( pynutil.delete("negative:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + DAMO_NOT_QUOTE - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space, 0, 1, ) graph = ( - optional_sign + pynutil.delete("integer: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + optional_sign + + pynutil.delete('integer: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') ) delete_tokens = self.delete_tokens(graph) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ru/verbalizers/date.py b/fun_text_processing/inverse_text_normalization/ru/verbalizers/date.py index aa74b771f..e1f530704 100644 --- a/fun_text_processing/inverse_text_normalization/ru/verbalizers/date.py +++ b/fun_text_processing/inverse_text_normalization/ru/verbalizers/date.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst from pynini.lib import pynutil @@ -12,6 +11,6 @@ class DateFst(GraphFst): def __init__(self): super().__init__(name="date", kind="verbalize") - graph = pynutil.delete("day: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + graph = pynutil.delete('day: "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') delete_tokens = self.delete_tokens(graph.optimize()) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ru/verbalizers/decimal.py b/fun_text_processing/inverse_text_normalization/ru/verbalizers/decimal.py index da1928d55..db730c748 100644 --- a/fun_text_processing/inverse_text_normalization/ru/verbalizers/decimal.py +++ b/fun_text_processing/inverse_text_normalization/ru/verbalizers/decimal.py @@ -1,6 +1,10 @@ - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, DAMO_SPACE, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + DAMO_SPACE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -17,12 +21,14 @@ class DecimalFst(GraphFst): def __init__(self, deterministic: bool = True): super().__init__(name="decimal", kind="verbalize", deterministic=deterministic) - optional_sign = pynini.closure(pynini.cross("negative: \"true\" ", "-"), 0, 1) + optional_sign = pynini.closure(pynini.cross('negative: "true" ', "-"), 0, 1) - integer = pynutil.delete(" \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + integer = pynutil.delete(' "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') integer_part = pynutil.delete("integer_part:") + integer fractional_part = pynutil.delete("fractional_part:") + integer - optional_quantity = pynini.closure(pynini.accep(DAMO_SPACE) + pynutil.delete("quantity:") + integer, 0, 1) + optional_quantity = pynini.closure( + pynini.accep(DAMO_SPACE) + pynutil.delete("quantity:") + integer, 0, 1 + ) graph = optional_sign + integer_part + delete_space + fractional_part + optional_quantity delete_tokens = self.delete_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/ru/verbalizers/electronic.py b/fun_text_processing/inverse_text_normalization/ru/verbalizers/electronic.py index 378212b0a..585c86561 100644 --- a/fun_text_processing/inverse_text_normalization/ru/verbalizers/electronic.py +++ b/fun_text_processing/inverse_text_normalization/ru/verbalizers/electronic.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst from pynini.lib import pynutil @@ -13,6 +12,8 @@ class ElectronicFst(GraphFst): def __init__(self): super().__init__(name="electronic", kind="verbalize") - graph = pynutil.delete("username: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + graph = ( + pynutil.delete('username: "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') + ) delete_tokens = self.delete_tokens(graph) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ru/verbalizers/measure.py b/fun_text_processing/inverse_text_normalization/ru/verbalizers/measure.py index 0a119bb16..55b2d3038 100644 --- a/fun_text_processing/inverse_text_normalization/ru/verbalizers/measure.py +++ b/fun_text_processing/inverse_text_normalization/ru/verbalizers/measure.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -14,9 +17,9 @@ class MeasureFst(GraphFst): super().__init__(name="measure", kind="verbalize") graph = ( - pynutil.delete(" cardinal { integer: \"") + pynutil.delete(' cardinal { integer: "') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space + pynutil.delete("}") ) diff --git a/fun_text_processing/inverse_text_normalization/ru/verbalizers/money.py b/fun_text_processing/inverse_text_normalization/ru/verbalizers/money.py index 531dac75a..f1f376c52 100644 --- a/fun_text_processing/inverse_text_normalization/ru/verbalizers/money.py +++ b/fun_text_processing/inverse_text_normalization/ru/verbalizers/money.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst from pynini.lib import pynutil @@ -13,6 +12,10 @@ class MoneyFst(GraphFst): def __init__(self): super().__init__(name="money", kind="verbalize") - graph = pynutil.delete("integer_part: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + graph = ( + pynutil.delete('integer_part: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + ) delete_tokens = self.delete_tokens(graph) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ru/verbalizers/ordinal.py b/fun_text_processing/inverse_text_normalization/ru/verbalizers/ordinal.py index 7e175d374..348b3a907 100644 --- a/fun_text_processing/inverse_text_normalization/ru/verbalizers/ordinal.py +++ b/fun_text_processing/inverse_text_normalization/ru/verbalizers/ordinal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst from pynini.lib import pynutil @@ -18,6 +17,6 @@ class OrdinalFst(GraphFst): super().__init__(name="ordinal", kind="verbalize", deterministic=deterministic) value = pynini.closure(DAMO_NOT_QUOTE) - graph = pynutil.delete("integer: \"") + value + pynutil.delete("\"") + graph = pynutil.delete('integer: "') + value + pynutil.delete('"') delete_tokens = self.delete_tokens(graph) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ru/verbalizers/telephone.py b/fun_text_processing/inverse_text_normalization/ru/verbalizers/telephone.py index 5a0f3c745..9fb3fec48 100644 --- a/fun_text_processing/inverse_text_normalization/ru/verbalizers/telephone.py +++ b/fun_text_processing/inverse_text_normalization/ru/verbalizers/telephone.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst from pynini.lib import pynutil @@ -13,6 +12,10 @@ class TelephoneFst(GraphFst): def __init__(self): super().__init__(name="telephone", kind="verbalize") - graph = pynutil.delete("number_part: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + graph = ( + pynutil.delete('number_part: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + ) delete_tokens = self.delete_tokens(graph) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/inverse_text_normalization/ru/verbalizers/time.py b/fun_text_processing/inverse_text_normalization/ru/verbalizers/time.py index abc6ccc3e..92889ddf5 100644 --- a/fun_text_processing/inverse_text_normalization/ru/verbalizers/time.py +++ b/fun_text_processing/inverse_text_normalization/ru/verbalizers/time.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -13,16 +16,21 @@ class TimeFst(GraphFst): def __init__(self): super().__init__(name="time", kind="verbalize") hour = ( - pynutil.delete("hours: ") + pynutil.delete("\"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + pynutil.delete("hours: ") + + pynutil.delete('"') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') ) minutes = ( pynutil.delete("minutes: ") - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) - graph_preserve_order = pynutil.delete("hours: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + graph_preserve_order = ( + pynutil.delete('hours: "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') + ) # for cases that require permutations for the correct verbalization graph_reverse_order = hour + delete_space + pynutil.insert(":") + minutes + delete_space diff --git a/fun_text_processing/inverse_text_normalization/ru/verbalizers/verbalize.py b/fun_text_processing/inverse_text_normalization/ru/verbalizers/verbalize.py index 6683fbfc3..a9e4aedfc 100644 --- a/fun_text_processing/inverse_text_normalization/ru/verbalizers/verbalize.py +++ b/fun_text_processing/inverse_text_normalization/ru/verbalizers/verbalize.py @@ -1,4 +1,3 @@ - from fun_text_processing.inverse_text_normalization.en.verbalizers.whitelist import WhiteListFst from fun_text_processing.inverse_text_normalization.ru.verbalizers.cardinal import CardinalFst from fun_text_processing.inverse_text_normalization.ru.verbalizers.date import DateFst @@ -15,7 +14,7 @@ from fun_text_processing.text_normalization.en.graph_utils import GraphFst class VerbalizeFst(GraphFst): """ Composes other verbalizer grammars. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. """ diff --git a/fun_text_processing/inverse_text_normalization/ru/verbalizers/verbalize_final.py b/fun_text_processing/inverse_text_normalization/ru/verbalizers/verbalize_final.py index 5ea1ae8b8..05d50ca29 100644 --- a/fun_text_processing/inverse_text_normalization/ru/verbalizers/verbalize_final.py +++ b/fun_text_processing/inverse_text_normalization/ru/verbalizers/verbalize_final.py @@ -1,14 +1,17 @@ - import pynini from fun_text_processing.inverse_text_normalization.en.verbalizers.word import WordFst from fun_text_processing.inverse_text_normalization.ru.verbalizers.verbalize import VerbalizeFst -from fun_text_processing.text_normalization.en.graph_utils import GraphFst, delete_extra_space, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + GraphFst, + delete_extra_space, + delete_space, +) from pynini.lib import pynutil class VerbalizeFinalFst(GraphFst): """ - Finite state transducer that verbalizes an entire sentence, e.g. + Finite state transducer that verbalizes an entire sentence, e.g. tokens { name: "its" } tokens { time { hours: "12" minutes: "30" } } tokens { name: "now" } -> its 12:30 now """ diff --git a/fun_text_processing/inverse_text_normalization/run_evaluate.py b/fun_text_processing/inverse_text_normalization/run_evaluate.py index 7fb4a554d..76e6e3ce2 100644 --- a/fun_text_processing/inverse_text_normalization/run_evaluate.py +++ b/fun_text_processing/inverse_text_normalization/run_evaluate.py @@ -10,17 +10,21 @@ from fun_text_processing.text_normalization.data_loader_utils import ( ) -''' +""" Runs Evaluation on data in the format of : \t\t<`self` if trivial class or normalized text> like the Google text normalization data https://www.kaggle.com/richardwilliamsproat/text-normalization-for-english-russian-and-polish -''' +""" def parse_args(): parser = ArgumentParser() parser.add_argument("--input", help="input file path", type=str) parser.add_argument( - "--lang", help="language", choices=['en', 'id', 'ja', 'de', 'es', 'pt', 'ru', 'fr', 'vi', 'ko', 'zh', 'fil'], default="en", type=str + "--lang", + help="language", + choices=["en", "id", "ja", "de", "es", "pt", "ru", "fr", "vi", "ko", "zh", "fil"], + default="en", + type=str, ) parser.add_argument( "--cat", @@ -30,7 +34,9 @@ def parse_args(): default=None, choices=known_types, ) - parser.add_argument("--filter", action='store_true', help="clean data for inverse normalization purposes") + parser.add_argument( + "--filter", action="store_true", help="clean data for inverse normalization purposes" + ) return parser.parse_args() @@ -38,8 +44,10 @@ if __name__ == "__main__": # Example usage: # python run_evaluate.py --input= --cat= --filter args = parse_args() - if args.lang == 'en': - from fun_text_processing.inverse_text_normalization.en.clean_eval_data import filter_loaded_data + if args.lang == "en": + from fun_text_processing.inverse_text_normalization.en.clean_eval_data import ( + filter_loaded_data, + ) file_path = args.input inverse_normalizer = InverseNormalizer() @@ -69,31 +77,38 @@ if __name__ == "__main__": print(" - Data: " + str(len(tokens_normalized)) + " tokens") tokens_prediction = inverse_normalizer.inverse_normalize_list(tokens_normalized) print(" - Denormalized. Evaluating...") - token_accuracy[token_type] = evaluate(tokens_prediction, tokens_un_normalized, input=tokens_normalized) + token_accuracy[token_type] = evaluate( + tokens_prediction, tokens_un_normalized, input=tokens_normalized + ) print(" - Accuracy: " + str(token_accuracy[token_type])) - token_count_per_type = {token_type: len(tokens_per_type[token_type][0]) for token_type in tokens_per_type} + token_count_per_type = { + token_type: len(tokens_per_type[token_type][0]) for token_type in tokens_per_type + } token_weighted_accuracy = [ - token_count_per_type[token_type] * accuracy for token_type, accuracy in token_accuracy.items() + token_count_per_type[token_type] * accuracy + for token_type, accuracy in token_accuracy.items() ] print("- Accuracy: " + str(sum(token_weighted_accuracy) / sum(token_count_per_type.values()))) - print(" - Total: " + str(sum(token_count_per_type.values())), '\n') + print(" - Total: " + str(sum(token_count_per_type.values())), "\n") for token_type in token_accuracy: if token_type not in known_types: raise ValueError("Unexpected token type: " + token_type) if args.category is None: - c1 = ['Class', 'sent level'] + known_types - c2 = ['Num Tokens', len(sentences_normalized)] + [ - token_count_per_type[known_type] if known_type in tokens_per_type else '0' for known_type in known_types + c1 = ["Class", "sent level"] + known_types + c2 = ["Num Tokens", len(sentences_normalized)] + [ + token_count_per_type[known_type] if known_type in tokens_per_type else "0" + for known_type in known_types ] c3 = ["Denormalization", sentences_accuracy] + [ - token_accuracy[known_type] if known_type in token_accuracy else '0' for known_type in known_types + token_accuracy[known_type] if known_type in token_accuracy else "0" + for known_type in known_types ] for i in range(len(c1)): - print(f'{str(c1[i]):10s} | {str(c2[i]):10s} | {str(c3[i]):5s}') + print(f"{str(c1[i]):10s} | {str(c2[i]):10s} | {str(c3[i]):5s}") else: - print(f'numbers\t{token_count_per_type[args.category]}') - print(f'Denormalization\t{token_accuracy[args.category]}') + print(f"numbers\t{token_count_per_type[args.category]}") + print(f"Denormalization\t{token_accuracy[args.category]}") diff --git a/fun_text_processing/inverse_text_normalization/tl/__init__.py b/fun_text_processing/inverse_text_normalization/tl/__init__.py index 3622fb249..9cd5f7b2f 100755 --- a/fun_text_processing/inverse_text_normalization/tl/__init__.py +++ b/fun_text_processing/inverse_text_normalization/tl/__init__.py @@ -1,4 +1,7 @@ - -from fun_text_processing.inverse_text_normalization.tl.taggers.tokenize_and_classify import ClassifyFst +from fun_text_processing.inverse_text_normalization.tl.taggers.tokenize_and_classify import ( + ClassifyFst, +) from fun_text_processing.inverse_text_normalization.tl.verbalizers.verbalize import VerbalizeFst -from fun_text_processing.inverse_text_normalization.tl.verbalizers.verbalize_final import VerbalizeFinalFst +from fun_text_processing.inverse_text_normalization.tl.verbalizers.verbalize_final import ( + VerbalizeFinalFst, +) diff --git a/fun_text_processing/inverse_text_normalization/tl/clean_eval_data.py b/fun_text_processing/inverse_text_normalization/tl/clean_eval_data.py index f59f8d014..e975bf5ec 100755 --- a/fun_text_processing/inverse_text_normalization/tl/clean_eval_data.py +++ b/fun_text_processing/inverse_text_normalization/tl/clean_eval_data.py @@ -1,4 +1,3 @@ - from argparse import ArgumentParser from typing import List @@ -54,7 +53,7 @@ class Filter: Args: processes given instance with process function - + Returns: processed instance if instance belongs to expected class type or original instance """ if instance.token_type != self.class_type: @@ -72,7 +71,9 @@ def process_cardinal_1(instance: Instance) -> Instance: normalized = instance.normalized un_normalized = re.sub(r"[^0-9]", "", un_normalized) normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_ordinal_1(instance: Instance) -> bool: @@ -85,7 +86,9 @@ def process_ordinal_1(instance: Instance) -> Instance: normalized = instance.normalized un_normalized = re.sub(r"[,\s]", "", un_normalized) normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_decimal_1(instance: Instance) -> bool: @@ -98,7 +101,9 @@ def process_decimal_1(instance: Instance) -> Instance: un_normalized = re.sub(r",", "", un_normalized) normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_measure_1(instance: Instance) -> bool: @@ -115,7 +120,9 @@ def process_measure_1(instance: Instance) -> Instance: normalized = re.sub(r"[^a-z\s]", "", normalized) normalized = re.sub(r"per ([a-z\s]*)s$", r"per \1", normalized) normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_money_1(instance: Instance) -> bool: @@ -132,7 +139,9 @@ def process_money_1(instance: Instance) -> Instance: un_normalized = re.sub(r"(\d)m\s*$", r"\1 million", un_normalized) un_normalized = re.sub(r"(\d)bn?\s*$", r"\1 billion", un_normalized) normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_time_1(instance: Instance) -> bool: @@ -147,7 +156,9 @@ def process_time_1(instance: Instance) -> Instance: un_normalized = re.sub(r"(\d)\s?p\s?m\s?", r"\1 p.m.", un_normalized) normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_plain_1(instance: Instance) -> bool: @@ -158,7 +169,9 @@ def filter_plain_1(instance: Instance) -> bool: def process_plain_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_punct_1(instance: Instance) -> bool: @@ -169,7 +182,9 @@ def filter_punct_1(instance: Instance) -> bool: def process_punct_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_date_1(instance: Instance) -> bool: @@ -182,7 +197,9 @@ def process_date_1(instance: Instance) -> Instance: un_normalized = re.sub(r",", "", un_normalized) normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_letters_1(instance: Instance) -> bool: @@ -194,7 +211,9 @@ def process_letters_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_verbatim_1(instance: Instance) -> bool: @@ -205,7 +224,9 @@ def filter_verbatim_1(instance: Instance) -> bool: def process_verbatim_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_digit_1(instance: Instance) -> bool: @@ -217,7 +238,9 @@ def process_digit_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_telephone_1(instance: Instance) -> bool: @@ -229,7 +252,9 @@ def process_telephone_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_electronic_1(instance: Instance) -> bool: @@ -241,7 +266,9 @@ def process_electronic_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_fraction_1(instance: Instance) -> bool: @@ -253,7 +280,9 @@ def process_fraction_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_address_1(instance: Instance) -> bool: @@ -265,27 +294,51 @@ def process_address_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) filters = [] -filters.append(Filter(class_type="CARDINAL", process_func=process_cardinal_1, filter_func=filter_cardinal_1)) -filters.append(Filter(class_type="ORDINAL", process_func=process_ordinal_1, filter_func=filter_ordinal_1)) -filters.append(Filter(class_type="DECIMAL", process_func=process_decimal_1, filter_func=filter_decimal_1)) -filters.append(Filter(class_type="MEASURE", process_func=process_measure_1, filter_func=filter_measure_1)) +filters.append( + Filter(class_type="CARDINAL", process_func=process_cardinal_1, filter_func=filter_cardinal_1) +) +filters.append( + Filter(class_type="ORDINAL", process_func=process_ordinal_1, filter_func=filter_ordinal_1) +) +filters.append( + Filter(class_type="DECIMAL", process_func=process_decimal_1, filter_func=filter_decimal_1) +) +filters.append( + Filter(class_type="MEASURE", process_func=process_measure_1, filter_func=filter_measure_1) +) filters.append(Filter(class_type="MONEY", process_func=process_money_1, filter_func=filter_money_1)) filters.append(Filter(class_type="TIME", process_func=process_time_1, filter_func=filter_time_1)) filters.append(Filter(class_type="DATE", process_func=process_date_1, filter_func=filter_date_1)) filters.append(Filter(class_type="PLAIN", process_func=process_plain_1, filter_func=filter_plain_1)) filters.append(Filter(class_type="PUNCT", process_func=process_punct_1, filter_func=filter_punct_1)) -filters.append(Filter(class_type="LETTERS", process_func=process_letters_1, filter_func=filter_letters_1)) -filters.append(Filter(class_type="VERBATIM", process_func=process_verbatim_1, filter_func=filter_verbatim_1)) +filters.append( + Filter(class_type="LETTERS", process_func=process_letters_1, filter_func=filter_letters_1) +) +filters.append( + Filter(class_type="VERBATIM", process_func=process_verbatim_1, filter_func=filter_verbatim_1) +) filters.append(Filter(class_type="DIGIT", process_func=process_digit_1, filter_func=filter_digit_1)) -filters.append(Filter(class_type="TELEPHONE", process_func=process_telephone_1, filter_func=filter_telephone_1)) -filters.append(Filter(class_type="ELECTRONIC", process_func=process_electronic_1, filter_func=filter_electronic_1)) -filters.append(Filter(class_type="FRACTION", process_func=process_fraction_1, filter_func=filter_fraction_1)) -filters.append(Filter(class_type="ADDRESS", process_func=process_address_1, filter_func=filter_address_1)) +filters.append( + Filter(class_type="TELEPHONE", process_func=process_telephone_1, filter_func=filter_telephone_1) +) +filters.append( + Filter( + class_type="ELECTRONIC", process_func=process_electronic_1, filter_func=filter_electronic_1 + ) +) +filters.append( + Filter(class_type="FRACTION", process_func=process_fraction_1, filter_func=filter_fraction_1) +) +filters.append( + Filter(class_type="ADDRESS", process_func=process_address_1, filter_func=filter_address_1) +) filters.append(Filter(class_type=EOS_TYPE, process_func=lambda x: x, filter_func=lambda x: True)) @@ -314,8 +367,10 @@ def filter_loaded_data(data: List[Instance], verbose: bool = False) -> List[Inst def parse_args(): parser = ArgumentParser() - parser.add_argument("--input", help="input file path", type=str, default='./en_with_types/output-00001-of-00100') - parser.add_argument("--verbose", help="print filtered instances", action='store_true') + parser.add_argument( + "--input", help="input file path", type=str, default="./en_with_types/output-00001-of-00100" + ) + parser.add_argument("--verbose", help="print filtered instances", action="store_true") return parser.parse_args() diff --git a/fun_text_processing/inverse_text_normalization/tl/graph_utils.py b/fun_text_processing/inverse_text_normalization/tl/graph_utils.py index 833776d36..1b6300955 100755 --- a/fun_text_processing/inverse_text_normalization/tl/graph_utils.py +++ b/fun_text_processing/inverse_text_normalization/tl/graph_utils.py @@ -1,4 +1,3 @@ - import os import string from pathlib import Path @@ -19,9 +18,9 @@ DAMO_UPPER = pynini.union(*string.ascii_uppercase).optimize() DAMO_ALPHA = pynini.union(DAMO_LOWER, DAMO_UPPER).optimize() DAMO_ALNUM = pynini.union(DAMO_DIGIT, DAMO_ALPHA).optimize() DAMO_HEX = pynini.union(*string.hexdigits).optimize() -DAMO_NON_BREAKING_SPACE = u"\u00A0" +DAMO_NON_BREAKING_SPACE = "\u00A0" DAMO_SPACE = " " -DAMO_WHITE_SPACE = pynini.union(" ", "\t", "\n", "\r", u"\u00A0").optimize() +DAMO_WHITE_SPACE = pynini.union(" ", "\t", "\n", "\r", "\u00A0").optimize() DAMO_NOT_SPACE = pynini.difference(DAMO_CHAR, DAMO_WHITE_SPACE).optimize() DAMO_NOT_QUOTE = pynini.difference(DAMO_CHAR, r'"').optimize() @@ -36,31 +35,55 @@ insert_space = pynutil.insert(" ") delete_extra_space = pynini.cross(pynini.closure(DAMO_WHITE_SPACE, 1), " ") delete_preserve_order = pynini.closure( pynutil.delete(" preserve_order: true") - | (pynutil.delete(" field_order: \"") + DAMO_NOT_QUOTE + pynutil.delete("\"")) + | (pynutil.delete(' field_order: "') + DAMO_NOT_QUOTE + pynutil.delete('"')) ) suppletive = pynini.string_file(get_abs_path("data/suppletive.tsv")) # _v = pynini.union("a", "e", "i", "o", "u") _c = pynini.union( - "b", "c", "d", "f", "g", "h", "j", "k", "l", "m", "n", "p", "q", "r", "s", "t", "v", "w", "x", "y", "z" + "b", + "c", + "d", + "f", + "g", + "h", + "j", + "k", + "l", + "m", + "n", + "p", + "q", + "r", + "s", + "t", + "v", + "w", + "x", + "y", + "z", ) _ies = DAMO_SIGMA + _c + pynini.cross("y", "ies") _es = DAMO_SIGMA + pynini.union("s", "sh", "ch", "x", "z") + pynutil.insert("es") _s = DAMO_SIGMA + pynutil.insert("s") graph_plural = plurals._priority_union( - suppletive, plurals._priority_union(_ies, plurals._priority_union(_es, _s, DAMO_SIGMA), DAMO_SIGMA), DAMO_SIGMA + suppletive, + plurals._priority_union(_ies, plurals._priority_union(_es, _s, DAMO_SIGMA), DAMO_SIGMA), + DAMO_SIGMA, ).optimize() SINGULAR_TO_PLURAL = graph_plural PLURAL_TO_SINGULAR = pynini.invert(graph_plural) -TO_LOWER = pynini.union(*[pynini.cross(x, y) for x, y in zip(string.ascii_uppercase, string.ascii_lowercase)]) +TO_LOWER = pynini.union( + *[pynini.cross(x, y) for x, y in zip(string.ascii_uppercase, string.ascii_lowercase)] +) TO_UPPER = pynini.invert(TO_LOWER) MIN_NEG_WEIGHT = -0.0001 MIN_POS_WEIGHT = 0.0001 -def generator_main(file_name: str, graphs: Dict[str, 'pynini.FstLike']): +def generator_main(file_name: str, graphs: Dict[str, "pynini.FstLike"]): """ Exports graph as OpenFst finite state archive (FAR) file with given file name and rule name. @@ -72,7 +95,7 @@ def generator_main(file_name: str, graphs: Dict[str, 'pynini.FstLike']): for rule, graph in graphs.items(): exporter[rule] = graph.optimize() exporter.close() - print(f'Created {file_name}') + print(f"Created {file_name}") def get_plurals(fst): @@ -99,7 +122,7 @@ def get_singulars(fst): return PLURAL_TO_SINGULAR @ fst -def convert_space(fst) -> 'pynini.FstLike': +def convert_space(fst) -> "pynini.FstLike": """ Converts space to nonbreaking space. Used only in tagger grammars for transducing token values within quotes, e.g. name: "hello kitty" @@ -110,7 +133,9 @@ def convert_space(fst) -> 'pynini.FstLike': Returns output fst where breaking spaces are converted to non breaking spaces """ - return fst @ pynini.cdrewrite(pynini.cross(DAMO_SPACE, DAMO_NON_BREAKING_SPACE), "", "", DAMO_SIGMA) + return fst @ pynini.cdrewrite( + pynini.cross(DAMO_SPACE, DAMO_NON_BREAKING_SPACE), "", "", DAMO_SIGMA + ) class GraphFst: @@ -130,9 +155,11 @@ class GraphFst: self._fst = None self.deterministic = deterministic - self.far_path = Path(os.path.dirname(__file__) + '/grammars/' + kind + '/' + name + '.far') + self.far_path = Path(os.path.dirname(__file__) + "/grammars/" + kind + "/" + name + ".far") if self.far_exist(): - self._fst = Far(self.far_path, mode="r", arc_type="standard", far_type="default").get_fst() + self._fst = Far( + self.far_path, mode="r", arc_type="standard", far_type="default" + ).get_fst() def far_exist(self) -> bool: """ @@ -141,14 +168,14 @@ class GraphFst: return self.far_path.exists() @property - def fst(self) -> 'pynini.FstLike': + def fst(self) -> "pynini.FstLike": return self._fst @fst.setter def fst(self, fst): self._fst = fst - def add_tokens(self, fst) -> 'pynini.FstLike': + def add_tokens(self, fst) -> "pynini.FstLike": """ Wraps class name around to given fst @@ -160,7 +187,7 @@ class GraphFst: """ return pynutil.insert(f"{self.name} {{ ") + fst + pynutil.insert(" }") - def delete_tokens(self, fst) -> 'pynini.FstLike': + def delete_tokens(self, fst) -> "pynini.FstLike": """ Deletes class name wrap around output of given fst @@ -179,4 +206,4 @@ class GraphFst: + delete_space + pynutil.delete("}") ) - return res @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + return res @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) diff --git a/fun_text_processing/inverse_text_normalization/tl/taggers/cardinal.py b/fun_text_processing/inverse_text_normalization/tl/taggers/cardinal.py index 1bbeb9f8f..6beb3391b 100755 --- a/fun_text_processing/inverse_text_normalization/tl/taggers/cardinal.py +++ b/fun_text_processing/inverse_text_normalization/tl/taggers/cardinal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.tl.utils import get_abs_path, num_to_word from fun_text_processing.inverse_text_normalization.tl.graph_utils import ( @@ -16,7 +15,7 @@ class CardinalFst(GraphFst): """ Finite state transducer for classifying cardinals e.g. minus twenty three -> cardinal { integer: "23" negative: "-" } } - Numbers below thirteen are not converted. + Numbers below thirteen are not converted. """ def __init__(self): @@ -31,7 +30,7 @@ class CardinalFst(GraphFst): graph_ties = pynini.string_file(get_abs_path("data/numbers/ties.tsv")) graph_teen = pynini.string_file(get_abs_path("data/numbers/teen.tsv")) - + graph_teens_using_ties = graph_ties + pynini.cross(" ", "") + graph_digit addzero = pynutil.insert("0") @@ -46,7 +45,6 @@ class CardinalFst(GraphFst): ##两位数 graph_teens = graph_teen | graph_teens_using_ties teens = graph_teens - ##三位数, daan 百，raan 百（只有4和9的时候） graph_hundred1 = pynutil.delete("daan") @@ -57,12 +55,20 @@ class CardinalFst(GraphFst): graph_hundred_component1 = graph_digit_ties1 + delete_space + graph_hundred1 graph_hundred_component2 = graph_digit_ties2 + delete_space + graph_hundred2 - + graph_hundred_component = graph_hundred_component1 | graph_hundred_component2 - - hundred = ( (graph_hundred_component + pynutil.insert("00")) | (graph_hundred_component + delete_space + delete_at + delete_space + (teens - | pynutil.add_weight(addzero + digit, 0.1) - | pynutil.add_weight(digit + addzero, 0.5)))) + + hundred = (graph_hundred_component + pynutil.insert("00")) | ( + graph_hundred_component + + delete_space + + delete_at + + delete_space + + ( + teens + | pynutil.add_weight(addzero + digit, 0.1) + | pynutil.add_weight(digit + addzero, 0.5) + ) + ) graph_hundred_component_at_least_one_none_zero_digit = hundred @ ( pynini.closure(DAMO_DIGIT) + (DAMO_DIGIT - "0") + pynini.closure(DAMO_DIGIT) ) @@ -71,15 +77,22 @@ class CardinalFst(GraphFst): ) ##千， libo 表示千 - thousand = ((hundred | teens | digit_ties) + delete_space + pynutil.delete("libo") + delete_space + ( - hundred - | (delete_at + delete_space).ques + pynutil.add_weight(addzero + teens, 0.1) - | (delete_at + delete_space).ques + pynutil.add_weight(addzero**2 + digit, 0.5) - | pynutil.add_weight(digit + addzero**2, 0.8) - | pynutil.add_weight(addzero**3, 1.0))) - + thousand = ( + (hundred | teens | digit_ties) + + delete_space + + pynutil.delete("libo") + + delete_space + + ( + hundred + | (delete_at + delete_space).ques + pynutil.add_weight(addzero + teens, 0.1) + | (delete_at + delete_space).ques + pynutil.add_weight(addzero**2 + digit, 0.5) + | pynutil.add_weight(digit + addzero**2, 0.8) + | pynutil.add_weight(addzero**3, 1.0) + ) + ) + ##百万，milyon表示百万 - #million = (((hundred | teens | digit_ties) + delete_space + pynutil.delete("milyon") | pynutil.insert("000", weight=0.1))+ delete_space + ( + # million = (((hundred | teens | digit_ties) + delete_space + pynutil.delete("milyon") | pynutil.insert("000", weight=0.1))+ delete_space + ( # thousand # | pynutil.add_weight(addzero + hundred, 0.1) # | (delete_at + delete_space).ques + pynutil.add_weight(addzero**2 + teens, 0.5) @@ -88,36 +101,48 @@ class CardinalFst(GraphFst): # | pynutil.add_weight(addzero**4, 1.0))) graph_million = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("milyon"), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + pynutil.delete("milyon"), pynutil.insert("000", weight=0.1), ) ##bilyon bilyon表示十亿 graph_billion = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("bilyon"), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + pynutil.delete("bilyon"), pynutil.insert("000", weight=0.1), ) - + ##trilyon trilyon表示兆 graph_trillion = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("trilyon"), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + pynutil.delete("trilyon"), pynutil.insert("000", weight=0.1), ) ## graph_quadrillion = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("quadrilyon"), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + pynutil.delete("quadrilyon"), pynutil.insert("000", weight=0.1), ) graph_quintillion = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("quintilyon"), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + pynutil.delete("quintilyon"), pynutil.insert("000", weight=0.1), ) graph_sextillion = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("sextilyon"), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + pynutil.delete("sextilyon"), pynutil.insert("000", weight=0.1), ) - + # graph = pynini.union( graph_sextillion @@ -142,30 +167,35 @@ class CardinalFst(GraphFst): graph_zero, ) - #graph = zero | digits | teens | hundred | thousand | million + # graph = zero | digits | teens | hundred | thousand | million graph = graph @ pynini.union( - pynutil.delete(pynini.closure("0")) + pynini.difference(DAMO_DIGIT, "0") + pynini.closure(DAMO_DIGIT), "0" + pynutil.delete(pynini.closure("0")) + + pynini.difference(DAMO_DIGIT, "0") + + pynini.closure(DAMO_DIGIT), + "0", ) labels_exception = [num_to_word(x) for x in range(0, 13)] graph_exception = pynini.union(*labels_exception) - #graph = ( + # graph = ( # pynini.cdrewrite(pynutil.delete("and"), DAMO_SPACE, DAMO_SPACE, DAMO_SIGMA) # @ (DAMO_ALPHA + DAMO_SIGMA) # @ graph - #) + # ) self.graph_no_exception = graph self.graph = (pynini.project(graph, "input") - graph_exception.arcsort()) @ graph optional_minus_graph = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("minus", "\"-\"") + DAMO_SPACE, 0, 1 + pynutil.insert("negative: ") + pynini.cross("minus", '"-"') + DAMO_SPACE, 0, 1 ) - final_graph = optional_minus_graph + pynutil.insert("integer: \"") + self.graph + pynutil.insert("\"") + final_graph = ( + optional_minus_graph + pynutil.insert('integer: "') + self.graph + pynutil.insert('"') + ) final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/tl/taggers/date.py b/fun_text_processing/inverse_text_normalization/tl/taggers/date.py index 4f01bc3cb..fac6f0e39 100755 --- a/fun_text_processing/inverse_text_normalization/tl/taggers/date.py +++ b/fun_text_processing/inverse_text_normalization/tl/taggers/date.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.tl.utils import get_abs_path from fun_text_processing.inverse_text_normalization.tl.graph_utils import ( @@ -63,7 +62,9 @@ def _get_year_graph(): def _get_thousands_graph(): graph_ties = _get_ties_graph() - graph_hundred_component = (graph_digit + delete_space + pynutil.delete("hundred")) | pynutil.insert("0") + graph_hundred_component = ( + graph_digit + delete_space + pynutil.delete("hundred") + ) | pynutil.insert("0") graph = ( graph_digit + delete_space @@ -90,7 +91,7 @@ def _get_year_graph(): class DateFst(GraphFst): """ - Finite state transducer for classifying date, + Finite state transducer for classifying date, e.g. january fifth twenty twelve -> date { month: "january" day: "5" year: "2012" preserve_order: true } e.g. the fifth of january twenty twelve -> date { day: "5" month: "january" year: "2012" preserve_order: true } e.g. twenty twenty -> date { year: "2012" preserve_order: true } @@ -108,18 +109,26 @@ class DateFst(GraphFst): year_graph = pynutil.add_weight(year_graph, YEAR_WEIGHT) month_graph = _get_month_graph() - month_graph = pynutil.insert("month: \"") + month_graph + pynutil.insert("\"") + month_graph = pynutil.insert('month: "') + month_graph + pynutil.insert('"') - day_graph = pynutil.insert("day: \"") + pynutil.add_weight(ordinal_graph, -0.7) + pynutil.insert("\"") + day_graph = ( + pynutil.insert('day: "') + pynutil.add_weight(ordinal_graph, -0.7) + pynutil.insert('"') + ) graph_year = ( delete_extra_space - + pynutil.insert("year: \"") + + pynutil.insert('year: "') + pynutil.add_weight(year_graph, -YEAR_WEIGHT) - + pynutil.insert("\"") + + pynutil.insert('"') + ) + optional_graph_year = pynini.closure( + graph_year, + 0, + 1, ) - optional_graph_year = pynini.closure(graph_year, 0, 1,) graph_mdy = month_graph + ( - (delete_extra_space + day_graph) | graph_year | (delete_extra_space + day_graph + graph_year) + (delete_extra_space + day_graph) + | graph_year + | (delete_extra_space + day_graph + graph_year) ) graph_dmy = ( pynutil.delete("the") @@ -131,7 +140,9 @@ class DateFst(GraphFst): + month_graph + optional_graph_year ) - graph_year = pynutil.insert("year: \"") + (year_graph | _get_range_graph()) + pynutil.insert("\"") + graph_year = ( + pynutil.insert('year: "') + (year_graph | _get_range_graph()) + pynutil.insert('"') + ) final_graph = graph_mdy | graph_dmy | graph_year final_graph += pynutil.insert(" preserve_order: true") diff --git a/fun_text_processing/inverse_text_normalization/tl/taggers/decimal.py b/fun_text_processing/inverse_text_normalization/tl/taggers/decimal.py index a0771fa3b..24516f796 100755 --- a/fun_text_processing/inverse_text_normalization/tl/taggers/decimal.py +++ b/fun_text_processing/inverse_text_normalization/tl/taggers/decimal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.tl.utils import get_abs_path from fun_text_processing.inverse_text_normalization.tl.graph_utils import ( @@ -10,30 +9,40 @@ from fun_text_processing.inverse_text_normalization.tl.graph_utils import ( from pynini.lib import pynutil -def get_quantity(decimal: 'pynini.FstLike', cardinal_up_to_hundred: 'pynini.FstLike') -> 'pynini.FstLike': +def get_quantity( + decimal: "pynini.FstLike", cardinal_up_to_hundred: "pynini.FstLike" +) -> "pynini.FstLike": """ Returns FST that transforms either a cardinal or decimal followed by a quantity into a numeral, e.g. one million -> integer_part: "1" quantity: "million" e.g. one point five million -> integer_part: "1" fractional_part: "5" quantity: "million" - Args: + Args: decimal: decimal FST cardinal_up_to_hundred: cardinal FST """ numbers = cardinal_up_to_hundred @ ( - pynutil.delete(pynini.closure("0")) + pynini.difference(DAMO_DIGIT, "0") + pynini.closure(DAMO_DIGIT) + pynutil.delete(pynini.closure("0")) + + pynini.difference(DAMO_DIGIT, "0") + + pynini.closure(DAMO_DIGIT) ) suffix = pynini.union("milyon", "bilyon", "trilyon", "quadrilyon", "quintilyon", "sextilyon") res = ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + numbers - + pynutil.insert("\"") + + pynutil.insert('"') + delete_extra_space - + pynutil.insert("quantity: \"") + + pynutil.insert('quantity: "') + suffix - + pynutil.insert("\"") + + pynutil.insert('"') + ) + res |= ( + decimal + + delete_extra_space + + pynutil.insert('quantity: "') + + (suffix | "libo") + + pynutil.insert('"') ) - res |= decimal + delete_extra_space + pynutil.insert("quantity: \"") + (suffix | "libo") + pynutil.insert("\"") return res @@ -52,22 +61,31 @@ class DecimalFst(GraphFst): cardinal_graph = cardinal.graph_no_exception graph_decimal = pynini.string_file(get_abs_path("data/numbers/digit.tsv")) - graph_decimal |= pynini.string_file(get_abs_path("data/numbers/zero.tsv")) | pynini.cross("o", "0") + graph_decimal |= pynini.string_file(get_abs_path("data/numbers/zero.tsv")) | pynini.cross( + "o", "0" + ) graph_decimal = pynini.closure(graph_decimal + delete_space) + graph_decimal - + self.graph = graph_decimal point = pynini.cross("punto", "") optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("minus", "\"true\"") + delete_extra_space, 0, 1 + pynutil.insert("negative: ") + pynini.cross("minus", '"true"') + delete_extra_space, + 0, + 1, ) - graph_fractional = pynutil.insert("fractional_part: \"") + graph_decimal + pynutil.insert("\"") - graph_integer = pynutil.insert("integer_part: \"") + cardinal_graph + pynutil.insert("\"") + graph_fractional = ( + pynutil.insert('fractional_part: "') + graph_decimal + pynutil.insert('"') + ) + graph_integer = pynutil.insert('integer_part: "') + cardinal_graph + pynutil.insert('"') final_graph_wo_sign = ( - pynini.closure(graph_integer + delete_extra_space, 0, 1) + point + delete_extra_space + graph_fractional + pynini.closure(graph_integer + delete_extra_space, 0, 1) + + point + + delete_extra_space + + graph_fractional ) final_graph = optional_graph_negative + final_graph_wo_sign diff --git a/fun_text_processing/inverse_text_normalization/tl/taggers/electronic.py b/fun_text_processing/inverse_text_normalization/tl/taggers/electronic.py index 5c9d804d1..fcf7a892c 100755 --- a/fun_text_processing/inverse_text_normalization/tl/taggers/electronic.py +++ b/fun_text_processing/inverse_text_normalization/tl/taggers/electronic.py @@ -1,7 +1,10 @@ - import pynini from fun_text_processing.inverse_text_normalization.tl.utils import get_abs_path -from fun_text_processing.inverse_text_normalization.tl.graph_utils import DAMO_ALPHA, GraphFst, insert_space +from fun_text_processing.inverse_text_normalization.tl.graph_utils import ( + DAMO_ALPHA, + GraphFst, + insert_space, +) from pynini.lib import pynutil @@ -25,35 +28,50 @@ class ElectronicFst(GraphFst): accepted_username = alpha_num | symbols process_dot = pynini.cross("dot", ".") - username = (alpha_num + pynini.closure(delete_extra_space + accepted_username)) | pynutil.add_weight( - pynini.closure(DAMO_ALPHA, 1), weight=0.0001 - ) - username = pynutil.insert("username: \"") + username + pynutil.insert("\"") + username = ( + alpha_num + pynini.closure(delete_extra_space + accepted_username) + ) | pynutil.add_weight(pynini.closure(DAMO_ALPHA, 1), weight=0.0001) + username = pynutil.insert('username: "') + username + pynutil.insert('"') single_alphanum = pynini.closure(alpha_num + delete_extra_space) + alpha_num - server = single_alphanum | pynini.string_file(get_abs_path("data/electronic/server_name.tsv")) + server = single_alphanum | pynini.string_file( + get_abs_path("data/electronic/server_name.tsv") + ) domain = single_alphanum | pynini.string_file(get_abs_path("data/electronic/domain.tsv")) domain_graph = ( - pynutil.insert("domain: \"") + pynutil.insert('domain: "') + server + delete_extra_space + process_dot + delete_extra_space + domain - + pynutil.insert("\"") + + pynutil.insert('"') + ) + graph = ( + username + + delete_extra_space + + pynutil.delete("at") + + insert_space + + delete_extra_space + + domain_graph ) - graph = username + delete_extra_space + pynutil.delete("at") + insert_space + delete_extra_space + domain_graph ############# url ### protocol_end = pynini.cross(pynini.union("w w w", "www"), "www") - protocol_start = (pynini.cross("h t t p", "http") | pynini.cross("h t t p s", "https")) + pynini.cross( - " colon slash slash ", "://" - ) + protocol_start = ( + pynini.cross("h t t p", "http") | pynini.cross("h t t p s", "https") + ) + pynini.cross(" colon slash slash ", "://") # .com, ending = ( delete_extra_space + symbols + delete_extra_space - + (domain | pynini.closure(accepted_username + delete_extra_space,) + accepted_username) + + ( + domain + | pynini.closure( + accepted_username + delete_extra_space, + ) + + accepted_username + ) ) protocol_default = ( @@ -64,12 +82,18 @@ class ElectronicFst(GraphFst): + pynini.closure(ending, 1) ).optimize() protocol = ( - pynini.closure(protocol_start, 0, 1) + protocol_end + delete_extra_space + process_dot + protocol_default + pynini.closure(protocol_start, 0, 1) + + protocol_end + + delete_extra_space + + process_dot + + protocol_default ).optimize() - protocol |= pynini.closure(protocol_end + delete_extra_space + process_dot, 0, 1) + protocol_default + protocol |= ( + pynini.closure(protocol_end + delete_extra_space + process_dot, 0, 1) + protocol_default + ) - protocol = pynutil.insert("protocol: \"") + protocol.optimize() + pynutil.insert("\"") + protocol = pynutil.insert('protocol: "') + protocol.optimize() + pynutil.insert('"') graph |= protocol ######## diff --git a/fun_text_processing/inverse_text_normalization/tl/taggers/fraction.py b/fun_text_processing/inverse_text_normalization/tl/taggers/fraction.py index 3330c876e..f1de69f53 100755 --- a/fun_text_processing/inverse_text_normalization/tl/taggers/fraction.py +++ b/fun_text_processing/inverse_text_normalization/tl/taggers/fraction.py @@ -1,4 +1,3 @@ - from fun_text_processing.inverse_text_normalization.tl.graph_utils import GraphFst diff --git a/fun_text_processing/inverse_text_normalization/tl/taggers/measure.py b/fun_text_processing/inverse_text_normalization/tl/taggers/measure.py index ad5572b00..6b4e145d1 100755 --- a/fun_text_processing/inverse_text_normalization/tl/taggers/measure.py +++ b/fun_text_processing/inverse_text_normalization/tl/taggers/measure.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.tl.utils import get_abs_path from fun_text_processing.inverse_text_normalization.tl.graph_utils import ( @@ -32,22 +31,37 @@ class MeasureFst(GraphFst): graph_unit_plural = get_singulars(graph_unit_singular) # plural -> abbr optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("minus", "\"true\"") + delete_extra_space, 0, 1 + pynutil.insert("negative: ") + pynini.cross("minus", '"true"') + delete_extra_space, + 0, + 1, ) unit_singular = convert_space(graph_unit_singular) unit_plural = convert_space(graph_unit_plural) - unit_misc = pynutil.insert("/") + pynutil.delete("per") + delete_space + convert_space(graph_unit_singular) + unit_misc = ( + pynutil.insert("/") + + pynutil.delete("per") + + delete_space + + convert_space(graph_unit_singular) + ) unit_singular = ( - pynutil.insert("units: \"") - + (unit_singular | unit_misc | pynutil.add_weight(unit_singular + delete_space + unit_misc, 0.01)) - + pynutil.insert("\"") + pynutil.insert('units: "') + + ( + unit_singular + | unit_misc + | pynutil.add_weight(unit_singular + delete_space + unit_misc, 0.01) + ) + + pynutil.insert('"') ) unit_plural = ( - pynutil.insert("units: \"") - + (unit_plural | unit_misc | pynutil.add_weight(unit_plural + delete_space + unit_misc, 0.01)) - + pynutil.insert("\"") + pynutil.insert('units: "') + + ( + unit_plural + | unit_misc + | pynutil.add_weight(unit_plural + delete_space + unit_misc, 0.01) + ) + + pynutil.insert('"') ) subgraph_decimal = ( @@ -61,9 +75,9 @@ class MeasureFst(GraphFst): subgraph_cardinal = ( pynutil.insert("cardinal { ") + optional_graph_negative - + pynutil.insert("integer: \"") + + pynutil.insert('integer: "') + ((DAMO_SIGMA - "one") @ cardinal_graph) - + pynutil.insert("\"") + + pynutil.insert('"') + pynutil.insert(" }") + delete_extra_space + unit_plural @@ -71,9 +85,9 @@ class MeasureFst(GraphFst): subgraph_cardinal |= ( pynutil.insert("cardinal { ") + optional_graph_negative - + pynutil.insert("integer: \"") + + pynutil.insert('integer: "') + pynini.cross("one", "1") - + pynutil.insert("\"") + + pynutil.insert('"') + pynutil.insert(" }") + delete_extra_space + unit_singular diff --git a/fun_text_processing/inverse_text_normalization/tl/taggers/money.py b/fun_text_processing/inverse_text_normalization/tl/taggers/money.py index aca058c45..822ef3882 100755 --- a/fun_text_processing/inverse_text_normalization/tl/taggers/money.py +++ b/fun_text_processing/inverse_text_normalization/tl/taggers/money.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.tl.utils import get_abs_path from fun_text_processing.inverse_text_normalization.tl.graph_utils import ( @@ -33,8 +32,11 @@ class MoneyFst(GraphFst): # add support for missing hundred (only for 3 digit numbers) # "one fifty" -> "one hundred fifty" with_hundred = pynini.compose( - pynini.closure(DAMO_NOT_SPACE) + pynini.accep(" ") + pynutil.insert("hundred ") + DAMO_SIGMA, - pynini.compose(cardinal_graph, DAMO_DIGIT ** 3), + pynini.closure(DAMO_NOT_SPACE) + + pynini.accep(" ") + + pynutil.insert("hundred ") + + DAMO_SIGMA, + pynini.compose(cardinal_graph, DAMO_DIGIT**3), ) cardinal_graph |= with_hundred graph_decimal_final = decimal.final_graph_wo_negative @@ -43,20 +45,27 @@ class MoneyFst(GraphFst): unit_singular = pynini.invert(unit) unit_plural = get_singulars(unit_singular) - graph_unit_singular = pynutil.insert("currency: \"") + convert_space(unit_singular) + pynutil.insert("\"") - graph_unit_plural = pynutil.insert("currency: \"") + convert_space(unit_plural) + pynutil.insert("\"") + graph_unit_singular = ( + pynutil.insert('currency: "') + convert_space(unit_singular) + pynutil.insert('"') + ) + graph_unit_plural = ( + pynutil.insert('currency: "') + convert_space(unit_plural) + pynutil.insert('"') + ) - add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | (pynutil.insert("0") + DAMO_DIGIT) + add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | ( + pynutil.insert("0") + DAMO_DIGIT + ) # twelve dollars (and) fifty cents, zero cents cents_standalone = ( - pynutil.insert("fractional_part: \"") + pynutil.insert('fractional_part: "') + pynini.union( - pynutil.add_weight(((DAMO_SIGMA - "one") @ cardinal_graph), -0.7) @ add_leading_zero_to_double_digit + pynutil.add_weight(((DAMO_SIGMA - "one") @ cardinal_graph), -0.7) + @ add_leading_zero_to_double_digit + delete_space + pynutil.delete("cents"), pynini.cross("one", "01") + delete_space + pynutil.delete("cent"), ) - + pynutil.insert("\"") + + pynutil.insert('"') ) optional_cents_standalone = pynini.closure( @@ -70,31 +79,31 @@ class MoneyFst(GraphFst): # twelve dollars fifty, only after integer optional_cents_suffix = pynini.closure( delete_extra_space - + pynutil.insert("fractional_part: \"") + + pynutil.insert('fractional_part: "') + pynutil.add_weight(cardinal_graph @ add_leading_zero_to_double_digit, -0.7) - + pynutil.insert("\""), + + pynutil.insert('"'), 0, 1, ) graph_integer = ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + ((DAMO_SIGMA - "one") @ cardinal_graph) - + pynutil.insert("\"") + + pynutil.insert('"') + delete_extra_space + graph_unit_plural + (optional_cents_standalone | optional_cents_suffix) ) graph_integer |= ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + pynini.cross("one", "1") - + pynutil.insert("\"") + + pynutil.insert('"') + delete_extra_space + graph_unit_singular + (optional_cents_standalone | optional_cents_suffix) ) graph_decimal = graph_decimal_final + delete_extra_space + graph_unit_plural - graph_decimal |= pynutil.insert("currency: \"$\" integer_part: \"0\" ") + cents_standalone + graph_decimal |= pynutil.insert('currency: "$" integer_part: "0" ') + cents_standalone final_graph = graph_integer | graph_decimal final_graph = self.add_tokens(final_graph) diff --git a/fun_text_processing/inverse_text_normalization/tl/taggers/ordinal.py b/fun_text_processing/inverse_text_normalization/tl/taggers/ordinal.py index 599a02307..9fd1d5db1 100755 --- a/fun_text_processing/inverse_text_normalization/tl/taggers/ordinal.py +++ b/fun_text_processing/inverse_text_normalization/tl/taggers/ordinal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.tl.utils import get_abs_path from fun_text_processing.inverse_text_normalization.tl.graph_utils import DAMO_CHAR, GraphFst @@ -25,6 +24,6 @@ class OrdinalFst(GraphFst): ) self.graph = graph @ cardinal_graph - final_graph = pynutil.insert("integer: \"") + self.graph + pynutil.insert("\"") + final_graph = pynutil.insert('integer: "') + self.graph + pynutil.insert('"') final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/tl/taggers/punctuation.py b/fun_text_processing/inverse_text_normalization/tl/taggers/punctuation.py index fc48a6665..59cd89157 100755 --- a/fun_text_processing/inverse_text_normalization/tl/taggers/punctuation.py +++ b/fun_text_processing/inverse_text_normalization/tl/taggers/punctuation.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.tl.graph_utils import GraphFst from pynini.lib import pynutil @@ -13,9 +12,9 @@ class PunctuationFst(GraphFst): def __init__(self): super().__init__(name="punctuation", kind="classify") - s = "!#$%&\'()*+,-./:;<=>?@^_`{|}~" + s = "!#$%&'()*+,-./:;<=>?@^_`{|}~" punct = pynini.union(*s) - graph = pynutil.insert("name: \"") + punct + pynutil.insert("\"") + graph = pynutil.insert('name: "') + punct + pynutil.insert('"') self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/tl/taggers/telephone.py b/fun_text_processing/inverse_text_normalization/tl/taggers/telephone.py index 790e3967a..9c960f6e6 100755 --- a/fun_text_processing/inverse_text_normalization/tl/taggers/telephone.py +++ b/fun_text_processing/inverse_text_normalization/tl/taggers/telephone.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.tl.utils import get_abs_path from fun_text_processing.inverse_text_normalization.tl.graph_utils import ( @@ -24,7 +23,7 @@ def get_serial_number(cardinal): class TelephoneFst(GraphFst): """ - Finite state transducer for classifying telephone numbers, e.g. + Finite state transducer for classifying telephone numbers, e.g. one two three one two three five six seven eight -> { number_part: "123-123-5678" } This class also support card number and IP format. @@ -61,61 +60,89 @@ class TelephoneFst(GraphFst): double_digit.invert() # to handle cases like "one twenty three" - two_digit_cardinal = pynini.compose(cardinal.graph_no_exception, DAMO_DIGIT ** 2) + two_digit_cardinal = pynini.compose(cardinal.graph_no_exception, DAMO_DIGIT**2) double_digit_to_digit = ( - pynini.compose(double_digit, str_to_digit + pynutil.delete(" ") + str_to_digit) | two_digit_cardinal + pynini.compose(double_digit, str_to_digit + pynutil.delete(" ") + str_to_digit) + | two_digit_cardinal ) - single_or_double_digit = (pynutil.add_weight(double_digit_to_digit, -0.0001) | str_to_digit).optimize() + single_or_double_digit = ( + pynutil.add_weight(double_digit_to_digit, -0.0001) | str_to_digit + ).optimize() single_or_double_digit |= ( single_or_double_digit - + pynini.closure(pynutil.add_weight(pynutil.delete(" ") + single_or_double_digit, 0.0001)) + + pynini.closure( + pynutil.add_weight(pynutil.delete(" ") + single_or_double_digit, 0.0001) + ) ).optimize() number_part = pynini.compose( single_or_double_digit, - DAMO_DIGIT ** 3 + pynutil.insert("-") + DAMO_DIGIT ** 3 + pynutil.insert("-") + DAMO_DIGIT ** 4, + DAMO_DIGIT**3 + + pynutil.insert("-") + + DAMO_DIGIT**3 + + pynutil.insert("-") + + DAMO_DIGIT**4, ).optimize() - number_part = pynutil.insert("number_part: \"") + number_part.optimize() + pynutil.insert("\"") + number_part = ( + pynutil.insert('number_part: "') + number_part.optimize() + pynutil.insert('"') + ) cardinal_option = pynini.compose(single_or_double_digit, DAMO_DIGIT ** (2, 3)) country_code = ( - pynutil.insert("country_code: \"") + pynutil.insert('country_code: "') + pynini.closure(pynini.cross("plus ", "+"), 0, 1) - + ((pynini.closure(str_to_digit + pynutil.delete(" "), 0, 2) + str_to_digit) | cardinal_option) - + pynutil.insert("\"") + + ( + (pynini.closure(str_to_digit + pynutil.delete(" "), 0, 2) + str_to_digit) + | cardinal_option + ) + + pynutil.insert('"') ) - optional_country_code = pynini.closure(country_code + pynutil.delete(" ") + insert_space, 0, 1).optimize() + optional_country_code = pynini.closure( + country_code + pynutil.delete(" ") + insert_space, 0, 1 + ).optimize() graph = optional_country_code + number_part # credit card number - space_four_digits = insert_space + DAMO_DIGIT ** 4 - credit_card_graph = pynini.compose(single_or_double_digit, DAMO_DIGIT ** 4 + space_four_digits ** 3).optimize() - graph |= pynutil.insert("number_part: \"") + credit_card_graph.optimize() + pynutil.insert("\"") + space_four_digits = insert_space + DAMO_DIGIT**4 + credit_card_graph = pynini.compose( + single_or_double_digit, DAMO_DIGIT**4 + space_four_digits**3 + ).optimize() + graph |= ( + pynutil.insert('number_part: "') + credit_card_graph.optimize() + pynutil.insert('"') + ) # SSN ssn_graph = pynini.compose( single_or_double_digit, - DAMO_DIGIT ** 3 + pynutil.insert("-") + DAMO_DIGIT ** 2 + pynutil.insert("-") + DAMO_DIGIT ** 4, + DAMO_DIGIT**3 + + pynutil.insert("-") + + DAMO_DIGIT**2 + + pynutil.insert("-") + + DAMO_DIGIT**4, ).optimize() - graph |= pynutil.insert("number_part: \"") + ssn_graph.optimize() + pynutil.insert("\"") + graph |= pynutil.insert('number_part: "') + ssn_graph.optimize() + pynutil.insert('"') # ip - digit_or_double = pynini.closure(str_to_digit + pynutil.delete(" "), 0, 1) + double_digit_to_digit - digit_or_double |= double_digit_to_digit + pynini.closure(pynutil.delete(" ") + str_to_digit, 0, 1) + digit_or_double = ( + pynini.closure(str_to_digit + pynutil.delete(" "), 0, 1) + double_digit_to_digit + ) + digit_or_double |= double_digit_to_digit + pynini.closure( + pynutil.delete(" ") + str_to_digit, 0, 1 + ) digit_or_double |= str_to_digit + (pynutil.delete(" ") + str_to_digit) ** (0, 2) digit_or_double |= cardinal_option digit_or_double = digit_or_double.optimize() ip_graph = digit_or_double + (pynini.cross(" dot ", ".") + digit_or_double) ** 3 - graph |= pynutil.insert("number_part: \"") + ip_graph.optimize() + pynutil.insert("\"") + graph |= pynutil.insert('number_part: "') + ip_graph.optimize() + pynutil.insert('"') graph |= ( - pynutil.insert("number_part: \"") + pynutil.insert('number_part: "') + pynutil.add_weight(get_serial_number(cardinal=cardinal), weight=0.0001) - + pynutil.insert("\"") + + pynutil.insert('"') ) final_graph = self.add_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/tl/taggers/time.py b/fun_text_processing/inverse_text_normalization/tl/taggers/time.py index 232c13826..2ee72f367 100755 --- a/fun_text_processing/inverse_text_normalization/tl/taggers/time.py +++ b/fun_text_processing/inverse_text_normalization/tl/taggers/time.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.tl.taggers.cardinal import CardinalFst from fun_text_processing.inverse_text_normalization.tl.utils import get_abs_path, num_to_word @@ -46,21 +45,23 @@ class TimeFst(GraphFst): graph_minute_verbose = pynini.cross("half", "30") | pynini.cross("quarter", "15") oclock = pynini.cross(pynini.union("o' clock", "o clock", "o'clock", "oclock"), "") - final_graph_hour = pynutil.insert("hours: \"") + graph_hour + pynutil.insert("\"") + final_graph_hour = pynutil.insert('hours: "') + graph_hour + pynutil.insert('"') graph_minute = ( oclock + pynutil.insert("00") | pynutil.delete("o") + delete_space + graph_minute_single | graph_minute_double ) - final_suffix = pynutil.insert("suffix: \"") + convert_space(suffix_graph) + pynutil.insert("\"") + final_suffix = ( + pynutil.insert('suffix: "') + convert_space(suffix_graph) + pynutil.insert('"') + ) final_suffix = delete_space + insert_space + final_suffix final_suffix_optional = pynini.closure(final_suffix, 0, 1) final_time_zone_optional = pynini.closure( delete_space + insert_space - + pynutil.insert("zone: \"") + + pynutil.insert('zone: "') + convert_space(time_zone_graph) - + pynutil.insert("\""), + + pynutil.insert('"'), 0, 1, ) @@ -69,13 +70,17 @@ class TimeFst(GraphFst): # two o eight, two thirty five (am/pm) # two pm/am graph_hm = ( - final_graph_hour + delete_extra_space + pynutil.insert("minutes: \"") + graph_minute + pynutil.insert("\"") + final_graph_hour + + delete_extra_space + + pynutil.insert('minutes: "') + + graph_minute + + pynutil.insert('"') ) # 10 past four, quarter past four, half past four graph_m_past_h = ( - pynutil.insert("minutes: \"") + pynutil.insert('minutes: "') + pynini.union(graph_minute_single, graph_minute_double, graph_minute_verbose) - + pynutil.insert("\"") + + pynutil.insert('"') + delete_space + pynutil.delete("past") + delete_extra_space @@ -83,42 +88,48 @@ class TimeFst(GraphFst): ) graph_quarter_time = ( - pynutil.insert("minutes: \"") + pynutil.insert('minutes: "') + pynini.cross("quarter", "45") - + pynutil.insert("\"") + + pynutil.insert('"') + delete_space + pynutil.delete(pynini.union("to", "till")) + delete_extra_space - + pynutil.insert("hours: \"") + + pynutil.insert('hours: "') + to_hour_graph - + pynutil.insert("\"") + + pynutil.insert('"') ) graph_m_to_h_suffix_time = ( - pynutil.insert("minutes: \"") + pynutil.insert('minutes: "') + ((graph_minute_single | graph_minute_double).optimize() @ minute_to_graph) - + pynutil.insert("\"") - + pynini.closure(delete_space + pynutil.delete(pynini.union("min", "mins", "minute", "minutes")), 0, 1) + + pynutil.insert('"') + + pynini.closure( + delete_space + pynutil.delete(pynini.union("min", "mins", "minute", "minutes")), + 0, + 1, + ) + delete_space + pynutil.delete(pynini.union("to", "till")) + delete_extra_space - + pynutil.insert("hours: \"") + + pynutil.insert('hours: "') + to_hour_graph - + pynutil.insert("\"") + + pynutil.insert('"') + final_suffix ) graph_h = ( final_graph_hour + delete_extra_space - + pynutil.insert("minutes: \"") + + pynutil.insert('minutes: "') + (pynutil.insert("00") | graph_minute) - + pynutil.insert("\"") + + pynutil.insert('"') + final_suffix + final_time_zone_optional ) final_graph = ( - (graph_hm | graph_m_past_h | graph_quarter_time) + final_suffix_optional + final_time_zone_optional + (graph_hm | graph_m_past_h | graph_quarter_time) + + final_suffix_optional + + final_time_zone_optional ) final_graph |= graph_h final_graph |= graph_m_to_h_suffix_time diff --git a/fun_text_processing/inverse_text_normalization/tl/taggers/tokenize_and_classify.py b/fun_text_processing/inverse_text_normalization/tl/taggers/tokenize_and_classify.py index a5c67ddd5..9f3073a5d 100755 --- a/fun_text_processing/inverse_text_normalization/tl/taggers/tokenize_and_classify.py +++ b/fun_text_processing/inverse_text_normalization/tl/taggers/tokenize_and_classify.py @@ -1,4 +1,3 @@ - import os import pynini @@ -28,7 +27,7 @@ import logging class ClassifyFst(GraphFst): """ Final class that composes all other classification grammars. This class can process an entire sentence, that is lower cased. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. Args: @@ -81,10 +80,16 @@ class ClassifyFst(GraphFst): | pynutil.add_weight(word_graph, 100) ) - punct = pynutil.insert("tokens { ") + pynutil.add_weight(punct_graph, weight=1.1) + pynutil.insert(" }") + punct = ( + pynutil.insert("tokens { ") + + pynutil.add_weight(punct_graph, weight=1.1) + + pynutil.insert(" }") + ) token = pynutil.insert("tokens { ") + classify + pynutil.insert(" }") token_plus_punct = ( - pynini.closure(punct + pynutil.insert(" ")) + token + pynini.closure(pynutil.insert(" ") + punct) + pynini.closure(punct + pynutil.insert(" ")) + + token + + pynini.closure(pynutil.insert(" ") + punct) ) graph = token_plus_punct + pynini.closure(delete_extra_space + token_plus_punct) diff --git a/fun_text_processing/inverse_text_normalization/tl/taggers/whitelist.py b/fun_text_processing/inverse_text_normalization/tl/taggers/whitelist.py index 0c725d693..d555a6c53 100755 --- a/fun_text_processing/inverse_text_normalization/tl/taggers/whitelist.py +++ b/fun_text_processing/inverse_text_normalization/tl/taggers/whitelist.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.tl.utils import get_abs_path from fun_text_processing.inverse_text_normalization.tl.graph_utils import GraphFst, convert_space @@ -16,5 +15,5 @@ class WhiteListFst(GraphFst): super().__init__(name="whitelist", kind="classify") whitelist = pynini.string_file(get_abs_path("data/whitelist.tsv")).invert() - graph = pynutil.insert("name: \"") + convert_space(whitelist) + pynutil.insert("\"") + graph = pynutil.insert('name: "') + convert_space(whitelist) + pynutil.insert('"') self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/tl/taggers/word.py b/fun_text_processing/inverse_text_normalization/tl/taggers/word.py index 12d2cbaac..fe2226afd 100755 --- a/fun_text_processing/inverse_text_normalization/tl/taggers/word.py +++ b/fun_text_processing/inverse_text_normalization/tl/taggers/word.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.tl.graph_utils import DAMO_NOT_SPACE, GraphFst from pynini.lib import pynutil @@ -12,5 +11,5 @@ class WordFst(GraphFst): def __init__(self): super().__init__(name="word", kind="classify") - word = pynutil.insert("name: \"") + pynini.closure(DAMO_NOT_SPACE, 1) + pynutil.insert("\"") + word = pynutil.insert('name: "') + pynini.closure(DAMO_NOT_SPACE, 1) + pynutil.insert('"') self.fst = word.optimize() diff --git a/fun_text_processing/inverse_text_normalization/tl/utils.py b/fun_text_processing/inverse_text_normalization/tl/utils.py index 042d6e0eb..38a95bac5 100755 --- a/fun_text_processing/inverse_text_normalization/tl/utils.py +++ b/fun_text_processing/inverse_text_normalization/tl/utils.py @@ -1,4 +1,3 @@ - import os from typing import Union @@ -14,7 +13,7 @@ def num_to_word(x: Union[str, int]): Args x: integer - Returns: spoken representation + Returns: spoken representation """ if isinstance(x, int): x = str(x) @@ -28,7 +27,7 @@ def get_abs_path(rel_path): Args: rel_path: relative path to this file - + Returns absolute path """ - return os.path.dirname(os.path.abspath(__file__)) + '/' + rel_path + return os.path.dirname(os.path.abspath(__file__)) + "/" + rel_path diff --git a/fun_text_processing/inverse_text_normalization/tl/verbalizers/cardinal.py b/fun_text_processing/inverse_text_normalization/tl/verbalizers/cardinal.py index 2643007f3..0c58f8c1f 100755 --- a/fun_text_processing/inverse_text_normalization/tl/verbalizers/cardinal.py +++ b/fun_text_processing/inverse_text_normalization/tl/verbalizers/cardinal.py @@ -1,7 +1,9 @@ - - import pynini -from fun_text_processing.inverse_text_normalization.tl.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.tl.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -16,9 +18,9 @@ class CardinalFst(GraphFst): optional_sign = pynini.closure( pynutil.delete("negative:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + DAMO_NOT_QUOTE - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space, 0, 1, @@ -26,9 +28,9 @@ class CardinalFst(GraphFst): graph = ( pynutil.delete("integer:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) self.numbers = graph graph = optional_sign + graph diff --git a/fun_text_processing/inverse_text_normalization/tl/verbalizers/date.py b/fun_text_processing/inverse_text_normalization/tl/verbalizers/date.py index 486963734..91fefeebb 100755 --- a/fun_text_processing/inverse_text_normalization/tl/verbalizers/date.py +++ b/fun_text_processing/inverse_text_normalization/tl/verbalizers/date.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.tl.graph_utils import ( DAMO_NOT_QUOTE, @@ -21,43 +20,47 @@ class DateFst(GraphFst): month = ( pynutil.delete("month:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) day = ( pynutil.delete("day:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) year = ( pynutil.delete("year:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') ) # month (day) year graph_mdy = ( - month + pynini.closure(delete_extra_space + day, 0, 1) + pynini.closure(delete_extra_space + year, 0, 1) + month + + pynini.closure(delete_extra_space + day, 0, 1) + + pynini.closure(delete_extra_space + year, 0, 1) ) # (day) month year graph_dmy = ( - pynini.closure(day + delete_extra_space, 0, 1) + month + pynini.closure(delete_extra_space + year, 0, 1) + pynini.closure(day + delete_extra_space, 0, 1) + + month + + pynini.closure(delete_extra_space + year, 0, 1) ) optional_preserve_order = pynini.closure( pynutil.delete("preserve_order:") + delete_space + pynutil.delete("true") + delete_space | pynutil.delete("field_order:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + DAMO_NOT_QUOTE - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space ) diff --git a/fun_text_processing/inverse_text_normalization/tl/verbalizers/decimal.py b/fun_text_processing/inverse_text_normalization/tl/verbalizers/decimal.py index 7d0ffa755..929db408c 100755 --- a/fun_text_processing/inverse_text_normalization/tl/verbalizers/decimal.py +++ b/fun_text_processing/inverse_text_normalization/tl/verbalizers/decimal.py @@ -1,7 +1,9 @@ - - import pynini -from fun_text_processing.inverse_text_normalization.tl.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.tl.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -13,30 +15,30 @@ class DecimalFst(GraphFst): def __init__(self): super().__init__(name="decimal", kind="verbalize") - optionl_sign = pynini.closure(pynini.cross("negative: \"true\"", "-") + delete_space, 0, 1) + optionl_sign = pynini.closure(pynini.cross('negative: "true"', "-") + delete_space, 0, 1) integer = ( pynutil.delete("integer_part:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_integer = pynini.closure(integer + delete_space, 0, 1) fractional = ( pynutil.insert(".") + pynutil.delete("fractional_part:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_fractional = pynini.closure(fractional + delete_space, 0, 1) quantity = ( pynutil.delete("quantity:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_quantity = pynini.closure(pynutil.insert(" ") + quantity + delete_space, 0, 1) graph = optional_integer + optional_fractional + optional_quantity diff --git a/fun_text_processing/inverse_text_normalization/tl/verbalizers/electronic.py b/fun_text_processing/inverse_text_normalization/tl/verbalizers/electronic.py index 63bd0e1de..97cc9d38a 100755 --- a/fun_text_processing/inverse_text_normalization/tl/verbalizers/electronic.py +++ b/fun_text_processing/inverse_text_normalization/tl/verbalizers/electronic.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.tl.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.tl.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -15,24 +18,24 @@ class ElectronicFst(GraphFst): user_name = ( pynutil.delete("username:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) domain = ( pynutil.delete("domain:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) protocol = ( pynutil.delete("protocol:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) graph = user_name + delete_space + pynutil.insert("@") + domain diff --git a/fun_text_processing/inverse_text_normalization/tl/verbalizers/fraction.py b/fun_text_processing/inverse_text_normalization/tl/verbalizers/fraction.py index 613de0f95..6b24cafe3 100755 --- a/fun_text_processing/inverse_text_normalization/tl/verbalizers/fraction.py +++ b/fun_text_processing/inverse_text_normalization/tl/verbalizers/fraction.py @@ -1,10 +1,9 @@ - from fun_text_processing.inverse_text_normalization.tl.graph_utils import GraphFst class FractionFst(GraphFst): """ - Finite state transducer for verbalizing fraction, + Finite state transducer for verbalizing fraction, """ def __init__(self): diff --git a/fun_text_processing/inverse_text_normalization/tl/verbalizers/measure.py b/fun_text_processing/inverse_text_normalization/tl/verbalizers/measure.py index 10e688c63..c1f9da26a 100755 --- a/fun_text_processing/inverse_text_normalization/tl/verbalizers/measure.py +++ b/fun_text_processing/inverse_text_normalization/tl/verbalizers/measure.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.tl.graph_utils import DAMO_CHAR, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.tl.graph_utils import ( + DAMO_CHAR, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -16,13 +19,13 @@ class MeasureFst(GraphFst): def __init__(self, decimal: GraphFst, cardinal: GraphFst): super().__init__(name="measure", kind="verbalize") - optional_sign = pynini.closure(pynini.cross("negative: \"true\"", "-"), 0, 1) + optional_sign = pynini.closure(pynini.cross('negative: "true"', "-"), 0, 1) unit = ( pynutil.delete("units:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space ) graph_decimal = ( diff --git a/fun_text_processing/inverse_text_normalization/tl/verbalizers/money.py b/fun_text_processing/inverse_text_normalization/tl/verbalizers/money.py index cc76b837b..816c87940 100755 --- a/fun_text_processing/inverse_text_normalization/tl/verbalizers/money.py +++ b/fun_text_processing/inverse_text_normalization/tl/verbalizers/money.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.tl.graph_utils import DAMO_CHAR, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.tl.graph_utils import ( + DAMO_CHAR, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -18,9 +21,9 @@ class MoneyFst(GraphFst): unit = ( pynutil.delete("currency:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) graph = unit + delete_space + decimal.numbers delete_tokens = self.delete_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/tl/verbalizers/ordinal.py b/fun_text_processing/inverse_text_normalization/tl/verbalizers/ordinal.py index 5649d0507..c6810fb7a 100755 --- a/fun_text_processing/inverse_text_normalization/tl/verbalizers/ordinal.py +++ b/fun_text_processing/inverse_text_normalization/tl/verbalizers/ordinal.py @@ -1,6 +1,10 @@ - import pynini -from fun_text_processing.inverse_text_normalization.tl.graph_utils import DAMO_NOT_QUOTE, DAMO_SIGMA, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.tl.graph_utils import ( + DAMO_NOT_QUOTE, + DAMO_SIGMA, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -15,9 +19,9 @@ class OrdinalFst(GraphFst): graph = ( pynutil.delete("integer:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) convert_eleven = pynini.cross("11", "11th") convert_twelve = pynini.cross("12", "12th") diff --git a/fun_text_processing/inverse_text_normalization/tl/verbalizers/telephone.py b/fun_text_processing/inverse_text_normalization/tl/verbalizers/telephone.py index c59f91651..0c245e5c4 100755 --- a/fun_text_processing/inverse_text_normalization/tl/verbalizers/telephone.py +++ b/fun_text_processing/inverse_text_normalization/tl/verbalizers/telephone.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.tl.graph_utils import DAMO_NOT_QUOTE, GraphFst from pynini.lib import pynutil @@ -8,17 +7,21 @@ class TelephoneFst(GraphFst): """ Finite state transducer for verbalizing telephone, e.g. telephone { number_part: "123-123-5678" } - -> 123-123-5678 + -> 123-123-5678 """ def __init__(self): super().__init__(name="telephone", kind="verbalize") - number_part = pynutil.delete("number_part: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") - optional_country_code = pynini.closure( - pynutil.delete("country_code: \"") + number_part = ( + pynutil.delete('number_part: "') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') + ) + optional_country_code = pynini.closure( + pynutil.delete('country_code: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + pynini.accep(" "), 0, 1, diff --git a/fun_text_processing/inverse_text_normalization/tl/verbalizers/time.py b/fun_text_processing/inverse_text_normalization/tl/verbalizers/time.py index 10907defe..7c0834779 100755 --- a/fun_text_processing/inverse_text_normalization/tl/verbalizers/time.py +++ b/fun_text_processing/inverse_text_normalization/tl/verbalizers/time.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.tl.graph_utils import ( DAMO_CHAR, @@ -20,29 +19,31 @@ class TimeFst(GraphFst): def __init__(self): super().__init__(name="time", kind="verbalize") - add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | (pynutil.insert("0") + DAMO_DIGIT) + add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | ( + pynutil.insert("0") + DAMO_DIGIT + ) hour = ( pynutil.delete("hours:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_DIGIT, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) minute = ( pynutil.delete("minutes:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_DIGIT, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) suffix = ( delete_space + insert_space + pynutil.delete("suffix:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_suffix = pynini.closure(suffix, 0, 1) zone = ( @@ -50,9 +51,9 @@ class TimeFst(GraphFst): + insert_space + pynutil.delete("zone:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_zone = pynini.closure(zone, 0, 1) graph = ( diff --git a/fun_text_processing/inverse_text_normalization/tl/verbalizers/verbalize.py b/fun_text_processing/inverse_text_normalization/tl/verbalizers/verbalize.py index 17ff479e3..c7d47d2e9 100755 --- a/fun_text_processing/inverse_text_normalization/tl/verbalizers/verbalize.py +++ b/fun_text_processing/inverse_text_normalization/tl/verbalizers/verbalize.py @@ -1,4 +1,3 @@ - from fun_text_processing.inverse_text_normalization.tl.verbalizers.cardinal import CardinalFst from fun_text_processing.inverse_text_normalization.tl.verbalizers.date import DateFst from fun_text_processing.inverse_text_normalization.tl.verbalizers.decimal import DecimalFst @@ -15,7 +14,7 @@ from fun_text_processing.inverse_text_normalization.tl.graph_utils import GraphF class VerbalizeFst(GraphFst): """ Composes other verbalizer grammars. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. """ diff --git a/fun_text_processing/inverse_text_normalization/tl/verbalizers/verbalize_final.py b/fun_text_processing/inverse_text_normalization/tl/verbalizers/verbalize_final.py index bb3eaf9af..d7fca6e6c 100755 --- a/fun_text_processing/inverse_text_normalization/tl/verbalizers/verbalize_final.py +++ b/fun_text_processing/inverse_text_normalization/tl/verbalizers/verbalize_final.py @@ -1,14 +1,17 @@ - import pynini from fun_text_processing.inverse_text_normalization.tl.verbalizers.verbalize import VerbalizeFst from fun_text_processing.inverse_text_normalization.tl.verbalizers.word import WordFst -from fun_text_processing.inverse_text_normalization.tl.graph_utils import GraphFst, delete_extra_space, delete_space +from fun_text_processing.inverse_text_normalization.tl.graph_utils import ( + GraphFst, + delete_extra_space, + delete_space, +) from pynini.lib import pynutil class VerbalizeFinalFst(GraphFst): """ - Finite state transducer that verbalizes an entire sentence, e.g. + Finite state transducer that verbalizes an entire sentence, e.g. tokens { name: "its" } tokens { time { hours: "12" minutes: "30" } } tokens { name: "now" } -> its 12:30 now """ diff --git a/fun_text_processing/inverse_text_normalization/tl/verbalizers/whitelist.py b/fun_text_processing/inverse_text_normalization/tl/verbalizers/whitelist.py index 137529aa0..94ed9d4f1 100755 --- a/fun_text_processing/inverse_text_normalization/tl/verbalizers/whitelist.py +++ b/fun_text_processing/inverse_text_normalization/tl/verbalizers/whitelist.py @@ -1,6 +1,10 @@ - import pynini -from fun_text_processing.inverse_text_normalization.tl.graph_utils import DAMO_CHAR, DAMO_SIGMA, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.tl.graph_utils import ( + DAMO_CHAR, + DAMO_SIGMA, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -15,9 +19,9 @@ class WhiteListFst(GraphFst): graph = ( pynutil.delete("name:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) - graph = graph @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + graph = graph @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/tl/verbalizers/word.py b/fun_text_processing/inverse_text_normalization/tl/verbalizers/word.py index 2f43e3891..485a4589a 100755 --- a/fun_text_processing/inverse_text_normalization/tl/verbalizers/word.py +++ b/fun_text_processing/inverse_text_normalization/tl/verbalizers/word.py @@ -1,6 +1,10 @@ - import pynini -from fun_text_processing.inverse_text_normalization.tl.graph_utils import DAMO_CHAR, DAMO_SIGMA, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.tl.graph_utils import ( + DAMO_CHAR, + DAMO_SIGMA, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -13,7 +17,13 @@ class WordFst(GraphFst): def __init__(self): super().__init__(name="word", kind="verbalize") chars = pynini.closure(DAMO_CHAR - " ", 1) - char = pynutil.delete("name:") + delete_space + pynutil.delete("\"") + chars + pynutil.delete("\"") - graph = char @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + char = ( + pynutil.delete("name:") + + delete_space + + pynutil.delete('"') + + chars + + pynutil.delete('"') + ) + graph = char @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/vi/__init__.py b/fun_text_processing/inverse_text_normalization/vi/__init__.py index f84f840b1..90d3cbc71 100644 --- a/fun_text_processing/inverse_text_normalization/vi/__init__.py +++ b/fun_text_processing/inverse_text_normalization/vi/__init__.py @@ -1,5 +1,7 @@ - - -from fun_text_processing.inverse_text_normalization.vi.taggers.tokenize_and_classify import ClassifyFst +from fun_text_processing.inverse_text_normalization.vi.taggers.tokenize_and_classify import ( + ClassifyFst, +) from fun_text_processing.inverse_text_normalization.vi.verbalizers.verbalize import VerbalizeFst -from fun_text_processing.inverse_text_normalization.vi.verbalizers.verbalize_final import VerbalizeFinalFst +from fun_text_processing.inverse_text_normalization.vi.verbalizers.verbalize_final import ( + VerbalizeFinalFst, +) diff --git a/fun_text_processing/inverse_text_normalization/vi/data/math/__init__.py b/fun_text_processing/inverse_text_normalization/vi/data/math/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/inverse_text_normalization/vi/data/math/__init__.py +++ b/fun_text_processing/inverse_text_normalization/vi/data/math/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/inverse_text_normalization/vi/graph_utils.py b/fun_text_processing/inverse_text_normalization/vi/graph_utils.py index aac828933..644a0aeec 100644 --- a/fun_text_processing/inverse_text_normalization/vi/graph_utils.py +++ b/fun_text_processing/inverse_text_normalization/vi/graph_utils.py @@ -1,4 +1,3 @@ - import os import string from pathlib import Path @@ -36,7 +35,9 @@ delete_extra_space = pynini.cross(pynini.closure(DAMO_WHITE_SPACE, 1), " ") delete_hyphen = pynutil.delete(pynini.closure("-", 0, 1)) insert_hyphen = pynutil.insert("-") -TO_LOWER = pynini.union(*[pynini.cross(x, y) for x, y in zip(string.ascii_uppercase, string.ascii_lowercase)]) +TO_LOWER = pynini.union( + *[pynini.cross(x, y) for x, y in zip(string.ascii_uppercase, string.ascii_lowercase)] +) TO_UPPER = pynini.invert(TO_LOWER) @@ -66,7 +67,9 @@ def convert_space(fst) -> "pynini.FstLike": Returns output fst where breaking spaces are converted to non breaking spaces """ - return fst @ pynini.cdrewrite(pynini.cross(DAMO_SPACE, DAMO_NON_BREAKING_SPACE), "", "", DAMO_SIGMA) + return fst @ pynini.cdrewrite( + pynini.cross(DAMO_SPACE, DAMO_NON_BREAKING_SPACE), "", "", DAMO_SIGMA + ) class GraphFst: @@ -88,7 +91,9 @@ class GraphFst: self.far_path = Path(os.path.dirname(__file__) + "/grammars/" + kind + "/" + name + ".far") if self.far_exist(): - self._fst = Far(self.far_path, mode="r", arc_type="standard", far_type="default").get_fst() + self._fst = Far( + self.far_path, mode="r", arc_type="standard", far_type="default" + ).get_fst() def far_exist(self) -> bool: """ diff --git a/fun_text_processing/inverse_text_normalization/vi/taggers/cardinal.py b/fun_text_processing/inverse_text_normalization/vi/taggers/cardinal.py index e168f3fc5..fd5deb35d 100644 --- a/fun_text_processing/inverse_text_normalization/vi/taggers/cardinal.py +++ b/fun_text_processing/inverse_text_normalization/vi/taggers/cardinal.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.inverse_text_normalization.vi.graph_utils import ( DAMO_DIGIT, @@ -62,7 +60,9 @@ class CardinalFst(GraphFst): zero + delete_space + (graph_digit | graph_four), ) ) - graph_hundred_component = graph_hundred_ties_component | (pynutil.insert("00") + delete_space + graph_digit) + graph_hundred_component = graph_hundred_ties_component | ( + pynutil.insert("00") + delete_space + graph_digit + ) graph_hundred_component_at_least_one_none_zero_digit = graph_hundred_component @ ( pynini.closure(DAMO_DIGIT) + (DAMO_DIGIT - "0") + pynini.closure(DAMO_DIGIT) @@ -80,7 +80,9 @@ class CardinalFst(GraphFst): ) graph_ten_thousand = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("vạn"), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + pynutil.delete("vạn"), pynutil.insert("0000", weight=0.1), ) @@ -90,7 +92,9 @@ class CardinalFst(GraphFst): ) graph_million = pynini.union( - graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("triệu"), + graph_hundred_component_at_least_one_none_zero_digit + + delete_space + + pynutil.delete("triệu"), pynutil.insert("000", weight=0.1), ) graph_billion = pynini.union( @@ -108,7 +112,11 @@ class CardinalFst(GraphFst): + graph_thousands + delete_space + graph_hundred_ties_zero, - graph_ten_thousand + delete_space + graph_ten_thousand_suffix + delete_space + graph_hundred_ties_zero, + graph_ten_thousand + + delete_space + + graph_ten_thousand_suffix + + delete_space + + graph_hundred_ties_zero, graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete(pynini.union("nghìn", "ngàn")) @@ -119,7 +127,10 @@ class CardinalFst(GraphFst): ) graph = graph @ pynini.union( - pynutil.delete(pynini.closure("0")) + pynini.difference(DAMO_DIGIT, "0") + pynini.closure(DAMO_DIGIT), "0", + pynutil.delete(pynini.closure("0")) + + pynini.difference(DAMO_DIGIT, "0") + + pynini.closure(DAMO_DIGIT), + "0", ) # don't convert cardinals from zero to nine inclusive @@ -130,10 +141,16 @@ class CardinalFst(GraphFst): self.graph = (pynini.project(graph, "input") - graph_exception.arcsort()) @ graph optional_minus_graph = pynini.closure( - pynutil.insert("negative: ") + pynini.cross(pynini.union("âm", "trừ"), '"-"') + DAMO_SPACE, 0, 1, + pynutil.insert("negative: ") + + pynini.cross(pynini.union("âm", "trừ"), '"-"') + + DAMO_SPACE, + 0, + 1, ) - final_graph = optional_minus_graph + pynutil.insert('integer: "') + self.graph + pynutil.insert('"') + final_graph = ( + optional_minus_graph + pynutil.insert('integer: "') + self.graph + pynutil.insert('"') + ) final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/vi/taggers/date.py b/fun_text_processing/inverse_text_normalization/vi/taggers/date.py index e0783c894..5110d259f 100644 --- a/fun_text_processing/inverse_text_normalization/vi/taggers/date.py +++ b/fun_text_processing/inverse_text_normalization/vi/taggers/date.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.vi.graph_utils import GraphFst, delete_extra_space, delete_space +from fun_text_processing.inverse_text_normalization.vi.graph_utils import ( + GraphFst, + delete_extra_space, + delete_space, +) from fun_text_processing.inverse_text_normalization.vi.utils import get_abs_path from pynini.lib import pynutil @@ -30,7 +33,10 @@ def _get_ties_graph(): optional_ten = pynini.closure(delete_space + graph_ten, 0, 1) graph = pynini.union( - ties_graph + optional_ten + delete_space + (graph_digit | graph_one | graph_four | graph_five), + ties_graph + + optional_ten + + delete_space + + (graph_digit | graph_one | graph_four | graph_five), ties_graph + delete_space + graph_ten + pynutil.insert("0"), ) return graph @@ -44,7 +50,10 @@ def _get_year_graph(): def _get_digits_graph(): zero = pynini.cross((pynini.union("linh", "lẻ")), "0") four = pynini.cross("tư", "4") - graph = pynini.union(zero + delete_space + (graph_digit | four), graph_zero + delete_space + graph_digit,) + graph = pynini.union( + zero + delete_space + (graph_digit | four), + graph_zero + delete_space + graph_digit, + ) graph.optimize() return graph @@ -86,7 +95,12 @@ def _get_year_graph(): + (graph_teen | graph_ties | graph_digits) | graph_thousands | graph_hundreds - | (graph_digit + pynutil.insert("0") + delete_space + (graph_ties | graph_digits | graph_teen)) + | ( + graph_digit + + pynutil.insert("0") + + delete_space + + (graph_ties | graph_digits | graph_teen) + ) ) year_graph.optimize() return year_graph @@ -114,7 +128,9 @@ class DateFst(GraphFst): month_graph = pynutil.insert('month: "') + month_graph + pynutil.insert('"') month_exception = pynini.project(pynini.cross("năm", "5"), "input") - month_graph_exception = (pynini.project(month_graph, "input") - month_exception.arcsort()) @ month_graph + month_graph_exception = ( + pynini.project(month_graph, "input") - month_exception.arcsort() + ) @ month_graph day_graph = pynutil.insert('day: "') + cardinal_graph + pynutil.insert('"') # day_suffix = pynini.union("ngày", "mùng") @@ -132,12 +148,23 @@ class DateFst(GraphFst): optional_graph_year = pynini.closure(graph_year, 0, 1) graph_my = pynutil.delete("tháng") + delete_space + month_graph + graph_year graph_dmy = ( - day_graph + delete_space + pynutil.delete("tháng") + delete_extra_space + month_graph + optional_graph_year + day_graph + + delete_space + + pynutil.delete("tháng") + + delete_extra_space + + month_graph + + optional_graph_year ) graph_year = ( - pynutil.delete("năm") + delete_extra_space + pynutil.insert('year: "') + year_graph + pynutil.insert('"') + pynutil.delete("năm") + + delete_extra_space + + pynutil.insert('year: "') + + year_graph + + pynutil.insert('"') ) - final_graph = (graph_dmy | graph_my | graph_month | graph_year) + pynutil.insert(" preserve_order: true") + final_graph = (graph_dmy | graph_my | graph_month | graph_year) + pynutil.insert( + " preserve_order: true" + ) final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/vi/taggers/decimal.py b/fun_text_processing/inverse_text_normalization/vi/taggers/decimal.py index 1b2098dd3..ed3065627 100644 --- a/fun_text_processing/inverse_text_normalization/vi/taggers/decimal.py +++ b/fun_text_processing/inverse_text_normalization/vi/taggers/decimal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.vi.graph_utils import ( DAMO_DIGIT, @@ -12,7 +11,9 @@ from pynini.lib import pynutil graph_digit = pynini.string_file(get_abs_path("data/numbers/digit.tsv")) -def get_quantity(decimal: "pynini.FstLike", cardinal_up_to_hundred: "pynini.FstLike") -> "pynini.FstLike": +def get_quantity( + decimal: "pynini.FstLike", cardinal_up_to_hundred: "pynini.FstLike" +) -> "pynini.FstLike": """ Returns FST that transforms either a cardinal or decimal followed by a quantity into a numeral, e.g. một triệu -> integer_part: "1" quantity: "triệu" @@ -23,7 +24,9 @@ def get_quantity(decimal: "pynini.FstLike", cardinal_up_to_hundred: "pynini.FstL cardinal_up_to_hundred: cardinal FST """ numbers = cardinal_up_to_hundred @ ( - pynutil.delete(pynini.closure("0")) + pynini.difference(DAMO_DIGIT, "0") + pynini.closure(DAMO_DIGIT) + pynutil.delete(pynini.closure("0")) + + pynini.difference(DAMO_DIGIT, "0") + + pynini.closure(DAMO_DIGIT) ) suffix = pynini.union("triệu", "tỉ", "tỷ", "vạn") graph_four = pynini.cross("tư", "4") @@ -88,30 +91,40 @@ class DecimalFst(GraphFst): graph_decimal = pynini.union( graph_decimal, graph_four, - pynini.closure(graph_decimal + delete_space, 1) + (graph_decimal | graph_four | graph_five | graph_one), + pynini.closure(graph_decimal + delete_space, 1) + + (graph_decimal | graph_four | graph_five | graph_one), ) self.graph = graph_decimal point = pynutil.delete("chấm") | pynutil.delete("phẩy") optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + pynini.cross(pynini.union("âm", "trừ"), '"true"') + delete_extra_space, + pynutil.insert("negative: ") + + pynini.cross(pynini.union("âm", "trừ"), '"true"') + + delete_extra_space, 0, 1, ) - graph_fractional = pynutil.insert('fractional_part: "') + graph_decimal + pynutil.insert('"') + graph_fractional = ( + pynutil.insert('fractional_part: "') + graph_decimal + pynutil.insert('"') + ) graph_integer = pynutil.insert('integer_part: "') + cardinal_graph + pynutil.insert('"') final_graph_wo_sign = ( - pynini.closure(graph_integer + delete_extra_space, 0, 1) + point + delete_extra_space + graph_fractional + pynini.closure(graph_integer + delete_extra_space, 0, 1) + + point + + delete_extra_space + + graph_fractional ) final_graph = optional_graph_negative + final_graph_wo_sign self.final_graph_wo_negative = final_graph_wo_sign | get_quantity( - final_graph_wo_sign, cardinal.graph_hundred_component_at_least_one_none_zero_digit, + final_graph_wo_sign, + cardinal.graph_hundred_component_at_least_one_none_zero_digit, ) final_graph |= optional_graph_negative + get_quantity( - final_graph_wo_sign, cardinal.graph_hundred_component_at_least_one_none_zero_digit, + final_graph_wo_sign, + cardinal.graph_hundred_component_at_least_one_none_zero_digit, ) final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/vi/taggers/electronic.py b/fun_text_processing/inverse_text_normalization/vi/taggers/electronic.py index 00f0e18b6..b70afe3a4 100644 --- a/fun_text_processing/inverse_text_normalization/vi/taggers/electronic.py +++ b/fun_text_processing/inverse_text_normalization/vi/taggers/electronic.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.vi.graph_utils import DAMO_ALPHA, GraphFst, insert_space +from fun_text_processing.inverse_text_normalization.vi.graph_utils import ( + DAMO_ALPHA, + GraphFst, + insert_space, +) from fun_text_processing.inverse_text_normalization.vi.utils import get_abs_path from pynini.lib import pynutil @@ -32,7 +35,9 @@ class ElectronicFst(GraphFst): + pynutil.insert('"') ) single_alphanum = pynini.closure(alpha_num + delete_extra_space) + alpha_num - server = single_alphanum | pynini.string_file(get_abs_path("data/electronic/server_name.tsv")) + server = single_alphanum | pynini.string_file( + get_abs_path("data/electronic/server_name.tsv") + ) domain = single_alphanum | pynini.string_file(get_abs_path("data/electronic/domain.tsv")) multi_domain = ( pynini.closure(process_dot + delete_extra_space + domain + delete_extra_space) @@ -40,7 +45,13 @@ class ElectronicFst(GraphFst): + delete_extra_space + domain ) - domain_graph = pynutil.insert('domain: "') + server + delete_extra_space + multi_domain + pynutil.insert('"') + domain_graph = ( + pynutil.insert('domain: "') + + server + + delete_extra_space + + multi_domain + + pynutil.insert('"') + ) graph = ( username + delete_extra_space @@ -52,9 +63,9 @@ class ElectronicFst(GraphFst): ############# url ### protocol_end = pynini.cross(pynini.union("w w w", "www"), "www") - protocol_start = (pynini.cross("h t t p", "http") | pynini.cross("h t t p s", "https")) + pynini.cross( - " hai chấm sẹc sẹc ", "://" - ) + protocol_start = ( + pynini.cross("h t t p", "http") | pynini.cross("h t t p s", "https") + ) + pynini.cross(" hai chấm sẹc sẹc ", "://") # .com, ending = ( delete_extra_space diff --git a/fun_text_processing/inverse_text_normalization/vi/taggers/fraction.py b/fun_text_processing/inverse_text_normalization/vi/taggers/fraction.py index 60c4737cb..da15c5e2b 100644 --- a/fun_text_processing/inverse_text_normalization/vi/taggers/fraction.py +++ b/fun_text_processing/inverse_text_normalization/vi/taggers/fraction.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.vi.graph_utils import GraphFst, delete_extra_space, delete_space +from fun_text_processing.inverse_text_normalization.vi.graph_utils import ( + GraphFst, + delete_extra_space, + delete_space, +) from pynini.lib import pynutil @@ -24,9 +27,13 @@ class FractionFst(GraphFst): numerator = pynutil.insert('numerator: "') + graph_cardinal + pynutil.insert('"') fraction_component = pynutil.delete(pynini.union("phần", "trên", "chia")) - denominator = pynutil.insert('denominator: "') + (graph_cardinal | graph_four) + pynutil.insert('"') + denominator = ( + pynutil.insert('denominator: "') + (graph_cardinal | graph_four) + pynutil.insert('"') + ) - graph_fraction_component = numerator + delete_space + fraction_component + delete_extra_space + denominator + graph_fraction_component = ( + numerator + delete_space + fraction_component + delete_extra_space + denominator + ) self.graph_fraction_component = graph_fraction_component graph = graph_fraction_component @@ -34,7 +41,9 @@ class FractionFst(GraphFst): self.final_graph_wo_negative = graph optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + pynini.cross(pynini.union("âm", "trừ"), '"true"') + delete_extra_space, + pynutil.insert("negative: ") + + pynini.cross(pynini.union("âm", "trừ"), '"true"') + + delete_extra_space, 0, 1, ) diff --git a/fun_text_processing/inverse_text_normalization/vi/taggers/measure.py b/fun_text_processing/inverse_text_normalization/vi/taggers/measure.py index b711b5f69..9b2843077 100644 --- a/fun_text_processing/inverse_text_normalization/vi/taggers/measure.py +++ b/fun_text_processing/inverse_text_normalization/vi/taggers/measure.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.vi.graph_utils import ( GraphFst, @@ -34,17 +33,28 @@ class MeasureFst(GraphFst): graph_unit_singular = pynini.invert(graph_unit) # singular -> abbr optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + pynini.cross(pynini.union("âm", "trừ"), '"true"') + delete_extra_space, + pynutil.insert("negative: ") + + pynini.cross(pynini.union("âm", "trừ"), '"true"') + + delete_extra_space, 0, 1, ) unit_singular = convert_space(graph_unit_singular) - unit_misc = pynutil.insert("/") + pynutil.delete("trên") + delete_space + convert_space(graph_unit_singular) + unit_misc = ( + pynutil.insert("/") + + pynutil.delete("trên") + + delete_space + + convert_space(graph_unit_singular) + ) unit_singular = ( pynutil.insert('units: "') - + (unit_singular | unit_misc | pynutil.add_weight(unit_singular + delete_space + unit_misc, 0.01)) + + ( + unit_singular + | unit_misc + | pynutil.add_weight(unit_singular + delete_space + unit_misc, 0.01) + ) + pynutil.insert('"') ) diff --git a/fun_text_processing/inverse_text_normalization/vi/taggers/money.py b/fun_text_processing/inverse_text_normalization/vi/taggers/money.py index 5f25f54ad..da441b111 100644 --- a/fun_text_processing/inverse_text_normalization/vi/taggers/money.py +++ b/fun_text_processing/inverse_text_normalization/vi/taggers/money.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.vi.graph_utils import ( DAMO_DIGIT, @@ -32,15 +31,22 @@ class MoneyFst(GraphFst): unit = pynini.string_file(get_abs_path("data/currency.tsv")) unit_singular = pynini.invert(unit) - graph_unit_singular = pynutil.insert('currency: "') + convert_space(unit_singular) + pynutil.insert('"') + graph_unit_singular = ( + pynutil.insert('currency: "') + convert_space(unit_singular) + pynutil.insert('"') + ) - add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | (pynutil.insert("0") + DAMO_DIGIT) + add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | ( + pynutil.insert("0") + DAMO_DIGIT + ) # twelve dollars fifty, only after integer optional_cents_suffix = pynini.closure( delete_extra_space + pynutil.insert('fractional_part: "') - + (pynutil.add_weight(cardinal_graph @ add_leading_zero_to_double_digit, -0.7) | graph_half) + + ( + pynutil.add_weight(cardinal_graph @ add_leading_zero_to_double_digit, -0.7) + | graph_half + ) + pynutil.insert('"'), 0, 1, @@ -55,7 +61,9 @@ class MoneyFst(GraphFst): + optional_cents_suffix ) - graph_decimal = graph_decimal_final + delete_extra_space + graph_unit_singular + optional_cents_suffix + graph_decimal = ( + graph_decimal_final + delete_extra_space + graph_unit_singular + optional_cents_suffix + ) final_graph = graph_integer | graph_decimal final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/vi/taggers/ordinal.py b/fun_text_processing/inverse_text_normalization/vi/taggers/ordinal.py index e1e27ecab..7ab827380 100644 --- a/fun_text_processing/inverse_text_normalization/vi/taggers/ordinal.py +++ b/fun_text_processing/inverse_text_normalization/vi/taggers/ordinal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.vi.graph_utils import GraphFst, delete_space from fun_text_processing.inverse_text_normalization.vi.utils import get_abs_path @@ -19,6 +18,12 @@ class OrdinalFst(GraphFst): graph = graph_digit self.graph = graph - final_graph = pynutil.insert('integer: "') + graph_ordinal + delete_space + self.graph + pynutil.insert('"') + final_graph = ( + pynutil.insert('integer: "') + + graph_ordinal + + delete_space + + self.graph + + pynutil.insert('"') + ) final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/vi/taggers/punctuation.py b/fun_text_processing/inverse_text_normalization/vi/taggers/punctuation.py index b05da586f..4cdfaa7ce 100644 --- a/fun_text_processing/inverse_text_normalization/vi/taggers/punctuation.py +++ b/fun_text_processing/inverse_text_normalization/vi/taggers/punctuation.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.vi.graph_utils import GraphFst from pynini.lib import pynutil diff --git a/fun_text_processing/inverse_text_normalization/vi/taggers/telephone.py b/fun_text_processing/inverse_text_normalization/vi/taggers/telephone.py index 7c738a4db..1353c007a 100644 --- a/fun_text_processing/inverse_text_normalization/vi/taggers/telephone.py +++ b/fun_text_processing/inverse_text_normalization/vi/taggers/telephone.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.vi.graph_utils import GraphFst, delete_space from fun_text_processing.inverse_text_normalization.vi.utils import get_abs_path @@ -16,7 +15,9 @@ class TelephoneFst(GraphFst): graph_zero = pynini.string_file(get_abs_path("data/numbers/zero.tsv")) graph_digit = pynini.string_file(get_abs_path("data/numbers/digit.tsv")) digit = graph_digit | graph_zero - last_digit = digit | pynini.cross("mốt", "1") | pynini.cross("tư", "4") | pynini.cross("lăm", "5") + last_digit = ( + digit | pynini.cross("mốt", "1") | pynini.cross("tư", "4") | pynini.cross("lăm", "5") + ) graph_number_part = pynini.closure(digit + delete_space, 2) + last_digit number_part = pynutil.insert('number_part: "') + graph_number_part + pynutil.insert('"') diff --git a/fun_text_processing/inverse_text_normalization/vi/taggers/time.py b/fun_text_processing/inverse_text_normalization/vi/taggers/time.py index 6bf757928..d5d2c3703 100644 --- a/fun_text_processing/inverse_text_normalization/vi/taggers/time.py +++ b/fun_text_processing/inverse_text_normalization/vi/taggers/time.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.inverse_text_normalization.vi.graph_utils import ( GraphFst, @@ -37,7 +35,9 @@ class TimeFst(GraphFst): optional_minute = pynini.closure(delete_space + minute, 0, 1) second = pynini.cross("giây", "") - final_graph_hour = pynutil.insert('hours: "') + graph_hours + pynutil.insert('"') + delete_space + oclock + final_graph_hour = ( + pynutil.insert('hours: "') + graph_hours + pynutil.insert('"') + delete_space + oclock + ) graph_minute = graph_minutes + optional_minute graph_second = graph_minutes + delete_space + second final_time_zone_optional = pynini.closure( diff --git a/fun_text_processing/inverse_text_normalization/vi/taggers/tokenize_and_classify.py b/fun_text_processing/inverse_text_normalization/vi/taggers/tokenize_and_classify.py index 82beedd67..2543568f6 100644 --- a/fun_text_processing/inverse_text_normalization/vi/taggers/tokenize_and_classify.py +++ b/fun_text_processing/inverse_text_normalization/vi/taggers/tokenize_and_classify.py @@ -1,4 +1,3 @@ - import os import pynini @@ -86,10 +85,16 @@ class ClassifyFst(GraphFst): | pynutil.add_weight(word_graph, 100) ) - punct = pynutil.insert("tokens { ") + pynutil.add_weight(punct_graph, weight=1.1) + pynutil.insert(" }") + punct = ( + pynutil.insert("tokens { ") + + pynutil.add_weight(punct_graph, weight=1.1) + + pynutil.insert(" }") + ) token = pynutil.insert("tokens { ") + classify + pynutil.insert(" }") token_plus_punct = ( - pynini.closure(punct + pynutil.insert(" ")) + token + pynini.closure(pynutil.insert(" ") + punct) + pynini.closure(punct + pynutil.insert(" ")) + + token + + pynini.closure(pynutil.insert(" ") + punct) ) graph = token_plus_punct + pynini.closure(delete_extra_space + token_plus_punct) diff --git a/fun_text_processing/inverse_text_normalization/vi/taggers/whitelist.py b/fun_text_processing/inverse_text_normalization/vi/taggers/whitelist.py index 2e3961220..dec1f6b27 100644 --- a/fun_text_processing/inverse_text_normalization/vi/taggers/whitelist.py +++ b/fun_text_processing/inverse_text_normalization/vi/taggers/whitelist.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.vi.graph_utils import GraphFst, convert_space from fun_text_processing.inverse_text_normalization.vi.utils import get_abs_path diff --git a/fun_text_processing/inverse_text_normalization/vi/taggers/word.py b/fun_text_processing/inverse_text_normalization/vi/taggers/word.py index 5762cfca4..77b28ce0c 100644 --- a/fun_text_processing/inverse_text_normalization/vi/taggers/word.py +++ b/fun_text_processing/inverse_text_normalization/vi/taggers/word.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.vi.graph_utils import DAMO_NOT_SPACE, GraphFst from pynini.lib import pynutil diff --git a/fun_text_processing/inverse_text_normalization/vi/utils.py b/fun_text_processing/inverse_text_normalization/vi/utils.py index 64bf5c698..c6fdb7435 100644 --- a/fun_text_processing/inverse_text_normalization/vi/utils.py +++ b/fun_text_processing/inverse_text_normalization/vi/utils.py @@ -1,5 +1,3 @@ - - import os diff --git a/fun_text_processing/inverse_text_normalization/vi/verbalizers/cardinal.py b/fun_text_processing/inverse_text_normalization/vi/verbalizers/cardinal.py index f34694496..a54b4ea80 100644 --- a/fun_text_processing/inverse_text_normalization/vi/verbalizers/cardinal.py +++ b/fun_text_processing/inverse_text_normalization/vi/verbalizers/cardinal.py @@ -1,7 +1,9 @@ - - import pynini -from fun_text_processing.inverse_text_normalization.vi.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.vi.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil diff --git a/fun_text_processing/inverse_text_normalization/vi/verbalizers/date.py b/fun_text_processing/inverse_text_normalization/vi/verbalizers/date.py index 767df49d5..d87beab36 100644 --- a/fun_text_processing/inverse_text_normalization/vi/verbalizers/date.py +++ b/fun_text_processing/inverse_text_normalization/vi/verbalizers/date.py @@ -1,7 +1,9 @@ - - import pynini -from fun_text_processing.inverse_text_normalization.vi.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.vi.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -55,7 +57,11 @@ class DateFst(GraphFst): + delete_space ) - final_graph = (graph_y | graph_m | graph_dm | graph_dmy | graph_my) + delete_space + optional_preserve_order + final_graph = ( + (graph_y | graph_m | graph_dm | graph_dmy | graph_my) + + delete_space + + optional_preserve_order + ) delete_tokens = self.delete_tokens(final_graph) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/inverse_text_normalization/vi/verbalizers/decimal.py b/fun_text_processing/inverse_text_normalization/vi/verbalizers/decimal.py index 330e6d8c2..595735d5a 100644 --- a/fun_text_processing/inverse_text_normalization/vi/verbalizers/decimal.py +++ b/fun_text_processing/inverse_text_normalization/vi/verbalizers/decimal.py @@ -1,7 +1,9 @@ - - import pynini -from fun_text_processing.inverse_text_normalization.vi.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.vi.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil diff --git a/fun_text_processing/inverse_text_normalization/vi/verbalizers/electronic.py b/fun_text_processing/inverse_text_normalization/vi/verbalizers/electronic.py index 38a82a329..32cad1fb0 100644 --- a/fun_text_processing/inverse_text_normalization/vi/verbalizers/electronic.py +++ b/fun_text_processing/inverse_text_normalization/vi/verbalizers/electronic.py @@ -1,7 +1,9 @@ - - import pynini -from fun_text_processing.inverse_text_normalization.vi.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.vi.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil diff --git a/fun_text_processing/inverse_text_normalization/vi/verbalizers/fraction.py b/fun_text_processing/inverse_text_normalization/vi/verbalizers/fraction.py index 6bd2d75b9..29edff551 100644 --- a/fun_text_processing/inverse_text_normalization/vi/verbalizers/fraction.py +++ b/fun_text_processing/inverse_text_normalization/vi/verbalizers/fraction.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.vi.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.vi.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -14,7 +17,9 @@ class FractionFst(GraphFst): def __init__(self): super().__init__(name="fraction", kind="verbalize") optional_sign = pynini.closure(pynini.cross('negative: "true"', "-") + delete_space, 0, 1) - numerator = pynutil.delete('numerator: "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') + numerator = ( + pynutil.delete('numerator: "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') + ) denominator = ( pynutil.insert("/") diff --git a/fun_text_processing/inverse_text_normalization/vi/verbalizers/measure.py b/fun_text_processing/inverse_text_normalization/vi/verbalizers/measure.py index 535a14376..f4088b754 100644 --- a/fun_text_processing/inverse_text_normalization/vi/verbalizers/measure.py +++ b/fun_text_processing/inverse_text_normalization/vi/verbalizers/measure.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.vi.graph_utils import ( DAMO_CHAR, diff --git a/fun_text_processing/inverse_text_normalization/vi/verbalizers/money.py b/fun_text_processing/inverse_text_normalization/vi/verbalizers/money.py index 4e3175d4c..71de0f377 100644 --- a/fun_text_processing/inverse_text_normalization/vi/verbalizers/money.py +++ b/fun_text_processing/inverse_text_normalization/vi/verbalizers/money.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.vi.graph_utils import DAMO_CHAR, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.vi.graph_utils import ( + DAMO_CHAR, + GraphFst, + delete_space, +) from pynini.lib import pynutil diff --git a/fun_text_processing/inverse_text_normalization/vi/verbalizers/ordinal.py b/fun_text_processing/inverse_text_normalization/vi/verbalizers/ordinal.py index 79a4325ff..0d152d1f3 100644 --- a/fun_text_processing/inverse_text_normalization/vi/verbalizers/ordinal.py +++ b/fun_text_processing/inverse_text_normalization/vi/verbalizers/ordinal.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.vi.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.vi.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil diff --git a/fun_text_processing/inverse_text_normalization/vi/verbalizers/telephone.py b/fun_text_processing/inverse_text_normalization/vi/verbalizers/telephone.py index 43fc3a809..bb603d620 100644 --- a/fun_text_processing/inverse_text_normalization/vi/verbalizers/telephone.py +++ b/fun_text_processing/inverse_text_normalization/vi/verbalizers/telephone.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.vi.graph_utils import DAMO_NOT_QUOTE, GraphFst from pynini.lib import pynutil @@ -14,6 +13,10 @@ class TelephoneFst(GraphFst): def __init__(self): super().__init__(name="telephone", kind="verbalize") - number_part = pynutil.delete('number_part: "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') + number_part = ( + pynutil.delete('number_part: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + ) delete_tokens = self.delete_tokens(number_part) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/inverse_text_normalization/vi/verbalizers/time.py b/fun_text_processing/inverse_text_normalization/vi/verbalizers/time.py index 0f918739b..5e3b3dc76 100644 --- a/fun_text_processing/inverse_text_normalization/vi/verbalizers/time.py +++ b/fun_text_processing/inverse_text_normalization/vi/verbalizers/time.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.vi.graph_utils import ( DAMO_CHAR, @@ -22,7 +21,9 @@ class TimeFst(GraphFst): def __init__(self): super().__init__(name="time", kind="verbalize") - add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | (pynutil.insert("0") + DAMO_DIGIT) + add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | ( + pynutil.insert("0") + DAMO_DIGIT + ) hour = ( pynutil.delete("hours:") + delete_space @@ -55,12 +56,18 @@ class TimeFst(GraphFst): ) optional_zone = pynini.closure(zone, 0, 1) optional_second = pynini.closure( - delete_space + pynutil.insert(":") + (second @ add_leading_zero_to_double_digit), 0, 1, + delete_space + pynutil.insert(":") + (second @ add_leading_zero_to_double_digit), + 0, + 1, ) graph_h = hour + pynutil.insert("h") graph_hms = ( - hour + delete_space + pynutil.insert(":") + (minute @ add_leading_zero_to_double_digit) + optional_second + hour + + delete_space + + pynutil.insert(":") + + (minute @ add_leading_zero_to_double_digit) + + optional_second ) graph_ms = ( minute diff --git a/fun_text_processing/inverse_text_normalization/vi/verbalizers/verbalize.py b/fun_text_processing/inverse_text_normalization/vi/verbalizers/verbalize.py index 9b592acab..f99c31fe9 100644 --- a/fun_text_processing/inverse_text_normalization/vi/verbalizers/verbalize.py +++ b/fun_text_processing/inverse_text_normalization/vi/verbalizers/verbalize.py @@ -1,4 +1,3 @@ - from fun_text_processing.inverse_text_normalization.vi.graph_utils import GraphFst from fun_text_processing.inverse_text_normalization.vi.verbalizers.cardinal import CardinalFst from fun_text_processing.inverse_text_normalization.vi.verbalizers.date import DateFst diff --git a/fun_text_processing/inverse_text_normalization/vi/verbalizers/verbalize_final.py b/fun_text_processing/inverse_text_normalization/vi/verbalizers/verbalize_final.py index 00928cf39..fe6d1263c 100644 --- a/fun_text_processing/inverse_text_normalization/vi/verbalizers/verbalize_final.py +++ b/fun_text_processing/inverse_text_normalization/vi/verbalizers/verbalize_final.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.vi.graph_utils import GraphFst, delete_extra_space, delete_space +from fun_text_processing.inverse_text_normalization.vi.graph_utils import ( + GraphFst, + delete_extra_space, + delete_space, +) from fun_text_processing.inverse_text_normalization.vi.verbalizers.verbalize import VerbalizeFst from fun_text_processing.inverse_text_normalization.vi.verbalizers.word import WordFst from pynini.lib import pynutil diff --git a/fun_text_processing/inverse_text_normalization/vi/verbalizers/whitelist.py b/fun_text_processing/inverse_text_normalization/vi/verbalizers/whitelist.py index 0a08a7b14..45741219a 100644 --- a/fun_text_processing/inverse_text_normalization/vi/verbalizers/whitelist.py +++ b/fun_text_processing/inverse_text_normalization/vi/verbalizers/whitelist.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.inverse_text_normalization.vi.graph_utils import ( DAMO_CHAR, @@ -25,5 +23,5 @@ class WhiteListFst(GraphFst): + pynini.closure(DAMO_CHAR - " ", 1) + pynutil.delete('"') ) - graph = graph @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + graph = graph @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/vi/verbalizers/word.py b/fun_text_processing/inverse_text_normalization/vi/verbalizers/word.py index 095d4af56..07cf18986 100644 --- a/fun_text_processing/inverse_text_normalization/vi/verbalizers/word.py +++ b/fun_text_processing/inverse_text_normalization/vi/verbalizers/word.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.vi.graph_utils import ( DAMO_CHAR, @@ -18,7 +17,13 @@ class WordFst(GraphFst): def __init__(self): super().__init__(name="word", kind="verbalize") chars = pynini.closure(DAMO_CHAR - " ", 1) - char = pynutil.delete("name:") + delete_space + pynutil.delete('"') + chars + pynutil.delete('"') - graph = char @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + char = ( + pynutil.delete("name:") + + delete_space + + pynutil.delete('"') + + chars + + pynutil.delete('"') + ) + graph = char @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/zh/__init__.py b/fun_text_processing/inverse_text_normalization/zh/__init__.py index 6b2a5553d..b21a9317a 100755 --- a/fun_text_processing/inverse_text_normalization/zh/__init__.py +++ b/fun_text_processing/inverse_text_normalization/zh/__init__.py @@ -1,5 +1,7 @@ - - -from fun_text_processing.inverse_text_normalization.zh.taggers.tokenize_and_classify import ClassifyFst +from fun_text_processing.inverse_text_normalization.zh.taggers.tokenize_and_classify import ( + ClassifyFst, +) from fun_text_processing.inverse_text_normalization.zh.verbalizers.verbalize import VerbalizeFst -from fun_text_processing.inverse_text_normalization.zh.verbalizers.verbalize_final import VerbalizeFinalFst +from fun_text_processing.inverse_text_normalization.zh.verbalizers.verbalize_final import ( + VerbalizeFinalFst, +) diff --git a/fun_text_processing/inverse_text_normalization/zh/graph_utils.py b/fun_text_processing/inverse_text_normalization/zh/graph_utils.py index 5349863a7..9e213b83e 100755 --- a/fun_text_processing/inverse_text_normalization/zh/graph_utils.py +++ b/fun_text_processing/inverse_text_normalization/zh/graph_utils.py @@ -1,4 +1,3 @@ - import os import string from pathlib import Path @@ -20,10 +19,10 @@ DAMO_ALPHA = pynini.union(DAMO_LOWER, DAMO_UPPER).optimize() DAMO_ALNUM = pynini.union(DAMO_DIGIT, DAMO_ALPHA).optimize() DAMO_HEX = pynini.union(*string.hexdigits).optimize() -DAMO_NON_BREAKING_SPACE = u"\u00A0" +DAMO_NON_BREAKING_SPACE = "\u00A0" DAMO_SPACE = " " -DAMO_WHITE_SPACE = pynini.union(" ", "\t", "\n", "\r", u"\u00A0").optimize() +DAMO_WHITE_SPACE = pynini.union(" ", "\t", "\n", "\r", "\u00A0").optimize() DAMO_NOT_SPACE = pynini.difference(DAMO_CHAR, DAMO_WHITE_SPACE).optimize() DAMO_NOT_QUOTE = pynini.difference(DAMO_CHAR, r'"').optimize() @@ -39,30 +38,55 @@ delete_extra_space = pynini.cross(pynini.closure(DAMO_WHITE_SPACE, 1), " ") delete_preserve_order = pynini.closure( pynutil.delete(" preserve_order: true") - | (pynutil.delete(" field_order: \"") + DAMO_NOT_QUOTE + pynutil.delete("\"")) + | (pynutil.delete(' field_order: "') + DAMO_NOT_QUOTE + pynutil.delete('"')) ) suppletive = pynini.string_file(get_abs_path("data/suppletive.tsv")) # _v = pynini.union("a", "e", "i", "o", "u") _c = pynini.union( - "b", "c", "d", "f", "g", "h", "j", "k", "l", "m", "n", "p", "q", "r", "s", "t", "v", "w", "x", "y", "z" + "b", + "c", + "d", + "f", + "g", + "h", + "j", + "k", + "l", + "m", + "n", + "p", + "q", + "r", + "s", + "t", + "v", + "w", + "x", + "y", + "z", ) _ies = DAMO_SIGMA + _c + pynini.cross("y", "ies") _es = DAMO_SIGMA + pynini.union("s", "sh", "ch", "x", "z") + pynutil.insert("es") _s = DAMO_SIGMA + pynutil.insert("s") graph_plural = plurals._priority_union( - suppletive, plurals._priority_union(_ies, plurals._priority_union(_es, _s, DAMO_SIGMA), DAMO_SIGMA), DAMO_SIGMA + suppletive, + plurals._priority_union(_ies, plurals._priority_union(_es, _s, DAMO_SIGMA), DAMO_SIGMA), + DAMO_SIGMA, ).optimize() SINGULAR_TO_PLURAL = graph_plural PLURAL_TO_SINGULAR = pynini.invert(graph_plural) -TO_LOWER = pynini.union(*[pynini.cross(x, y) for x, y in zip(string.ascii_uppercase, string.ascii_lowercase)]) +TO_LOWER = pynini.union( + *[pynini.cross(x, y) for x, y in zip(string.ascii_uppercase, string.ascii_lowercase)] +) TO_UPPER = pynini.invert(TO_LOWER) MIN_NEG_WEIGHT = -0.0001 MIN_POS_WEIGHT = 0.0001 -def generator_main(file_name: str, graphs: Dict[str, 'pynini.FstLike']): + +def generator_main(file_name: str, graphs: Dict[str, "pynini.FstLike"]): """ Exports graph as OpenFst finite state archive (FAR) file with given file name and rule name. @@ -74,7 +98,8 @@ def generator_main(file_name: str, graphs: Dict[str, 'pynini.FstLike']): for rule, graph in graphs.items(): exporter[rule] = graph.optimize() exporter.close() - print(f'Created {file_name}') + print(f"Created {file_name}") + def get_plurals(fst): """ @@ -96,7 +121,7 @@ def get_singulars(fst): return PLURAL_TO_SINGULAR @ fst -def convert_space(fst) -> 'pynini.FstLike': +def convert_space(fst) -> "pynini.FstLike": """ Converts space to nonbreaking space. Used only in tagger grammars for transducing token values within quotes, e.g. name: "hello kitty" @@ -105,7 +130,10 @@ def convert_space(fst) -> 'pynini.FstLike': fst: input fst Returns output fst where breaking spaces are converted to non breaking spaces """ - return fst @ pynini.cdrewrite(pynini.cross(DAMO_SPACE, DAMO_NON_BREAKING_SPACE), "", "", DAMO_SIGMA) + return fst @ pynini.cdrewrite( + pynini.cross(DAMO_SPACE, DAMO_NON_BREAKING_SPACE), "", "", DAMO_SIGMA + ) + class GraphFst: """ @@ -124,9 +152,11 @@ class GraphFst: self._fst = None self.deterministic = deterministic - self.far_path = Path(os.path.dirname(__file__) + '/grammars/' + kind + '/' + name + '.far') + self.far_path = Path(os.path.dirname(__file__) + "/grammars/" + kind + "/" + name + ".far") if self.far_exist(): - self._fst = Far(self.far_path, mode="r", arc_type="standard", far_type="default").get_fst() + self._fst = Far( + self.far_path, mode="r", arc_type="standard", far_type="default" + ).get_fst() def far_exist(self) -> bool: """ @@ -135,14 +165,14 @@ class GraphFst: return self.far_path.exists() @property - def fst(self) -> 'pynini.FstLike': + def fst(self) -> "pynini.FstLike": return self._fst @fst.setter def fst(self, fst): self._fst = fst - def add_tokens(self, fst) -> 'pynini.FstLike': + def add_tokens(self, fst) -> "pynini.FstLike": """ Wraps class name around to given fst @@ -154,7 +184,7 @@ class GraphFst: """ return pynutil.insert(f"{self.name} {{ ") + fst + pynutil.insert(" }") - def delete_tokens(self, fst) -> 'pynini.FstLike': + def delete_tokens(self, fst) -> "pynini.FstLike": """ Deletes class name wrap around output of given fst @@ -173,4 +203,4 @@ class GraphFst: + delete_space + pynutil.delete("}") ) - return res @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + return res @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) diff --git a/fun_text_processing/inverse_text_normalization/zh/taggers/cardinal.py b/fun_text_processing/inverse_text_normalization/zh/taggers/cardinal.py index 61db01b6f..6037effc6 100755 --- a/fun_text_processing/inverse_text_normalization/zh/taggers/cardinal.py +++ b/fun_text_processing/inverse_text_normalization/zh/taggers/cardinal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.zh.utils import get_abs_path from fun_text_processing.inverse_text_normalization.zh.graph_utils import ( @@ -16,7 +15,7 @@ class CardinalFst(GraphFst): """ Finite state transducer for classifying cardinals e.g. minus twenty three -> cardinal { integer: "23" negative: "-" } } - Numbers below thirteen are not converted. + Numbers below thirteen are not converted. """ def __init__(self): @@ -28,22 +27,27 @@ class CardinalFst(GraphFst): addzero = pynutil.insert("0") graph_digits = graph_digit | graph_zero - - #单位， 十 - #十一 + + # 单位， 十 + # 十一 graph_teen = pynini.cross("十", "1") + (graph_digit | addzero) - + ##二十一 graph_tens = graph_digit + pynutil.delete("十") + (graph_digit | addzero) - - #单位， 百 - #一百一十 一百零六 一百二十 两百 - graph_hundred = graph_digit + pynutil.delete("百") + ( graph_tens - | graph_teen - | pynutil.add_weight(graph_zero + graph_digit, 0.1) - | pynutil.add_weight(graph_digit + addzero, 0.5) - | pynutil.add_weight(addzero**2, 0.1)) + # 单位， 百 + # 一百一十 一百零六 一百二十 两百 + graph_hundred = ( + graph_digit + + pynutil.delete("百") + + ( + graph_tens + | graph_teen + | pynutil.add_weight(graph_zero + graph_digit, 0.1) + | pynutil.add_weight(graph_digit + addzero, 0.5) + | pynutil.add_weight(addzero**2, 0.1) + ) + ) graph_hundred_component_at_least_one_none_zero_digit = graph_hundred @ ( pynini.closure(DAMO_DIGIT) + (DAMO_DIGIT - "0") + pynini.closure(DAMO_DIGIT) @@ -53,38 +57,58 @@ class CardinalFst(GraphFst): graph_hundred_component_at_least_one_none_zero_digit ) - #单位， 千 - #一千两百三十四 一千零三十四， 一千零四 一千二百 - #五千 - graph_thousand = ( (graph_hundred | graph_teen | graph_tens | graph_digits) + pynutil.delete("千") + - ( graph_hundred - | pynutil.add_weight(graph_zero + graph_tens, 0.1 ) - | pynutil.add_weight(addzero + graph_zero + graph_digit, 0.5) - | pynutil.add_weight(graph_digit + addzero**2, 0.8) - | pynutil.add_weight(addzero**3, 1.0))) + # 单位， 千 + # 一千两百三十四 一千零三十四， 一千零四 一千二百 + # 五千 + graph_thousand = ( + (graph_hundred | graph_teen | graph_tens | graph_digits) + + pynutil.delete("千") + + ( + graph_hundred + | pynutil.add_weight(graph_zero + graph_tens, 0.1) + | pynutil.add_weight(addzero + graph_zero + graph_digit, 0.5) + | pynutil.add_weight(graph_digit + addzero**2, 0.8) + | pynutil.add_weight(addzero**3, 1.0) + ) + ) - #单位，万 - graph_ten_thousand = ( (graph_thousand | graph_hundred | graph_teen | graph_tens | graph_digit) + pynutil.delete("万") + - ( graph_thousand - | pynutil.add_weight(graph_zero + graph_hundred, 0.1) - | pynutil.add_weight(addzero + graph_zero + graph_tens, 0.5) - | pynutil.add_weight(addzero + addzero + graph_zero + graph_digit, 0.5) - | pynutil.add_weight(graph_digit + addzero**3, 0.8) - | pynutil.add_weight(addzero**4, 1.0))) - - #单位， 亿/兆 - graph_number = graph_digits | graph_teen | graph_tens | graph_hundred | graph_thousand | graph_ten_thousand + # 单位，万 + graph_ten_thousand = ( + (graph_thousand | graph_hundred | graph_teen | graph_tens | graph_digit) + + pynutil.delete("万") + + ( + graph_thousand + | pynutil.add_weight(graph_zero + graph_hundred, 0.1) + | pynutil.add_weight(addzero + graph_zero + graph_tens, 0.5) + | pynutil.add_weight(addzero + addzero + graph_zero + graph_digit, 0.5) + | pynutil.add_weight(graph_digit + addzero**3, 0.8) + | pynutil.add_weight(addzero**4, 1.0) + ) + ) - graph_number = ((graph_number + pynini.accep('兆') + pynutil.delete('零').ques).ques + - (graph_number + pynini.accep('亿') + pynutil.delete('零').ques).ques + graph_number) + # 单位， 亿/兆 + graph_number = ( + graph_digits + | graph_teen + | graph_tens + | graph_hundred + | graph_thousand + | graph_ten_thousand + ) + + graph_number = ( + (graph_number + pynini.accep("兆") + pynutil.delete("零").ques).ques + + (graph_number + pynini.accep("亿") + pynutil.delete("零").ques).ques + + graph_number + ) ##符号 graph = graph_number - #IP - graph |= (graph_digit.plus + ( graph_dot + graph_digit.plus).plus) + # IP + graph |= graph_digit.plus + (graph_dot + graph_digit.plus).plus - #just fill - labels_exception = [ 'ttttt' ] + # just fill + labels_exception = ["ttttt"] graph_exception = pynini.union(*labels_exception) self.graph_no_exception = graph @@ -92,10 +116,12 @@ class CardinalFst(GraphFst): self.graph = (pynini.project(graph, "input") - graph_exception.arcsort()) @ graph optional_minus_graph = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("负", "\"-\""), 0, 1 + pynutil.insert("negative: ") + pynini.cross("负", '"-"'), 0, 1 ) - final_graph = optional_minus_graph + pynutil.insert("integer: \"") + self.graph + pynutil.insert("\"") + final_graph = ( + optional_minus_graph + pynutil.insert('integer: "') + self.graph + pynutil.insert('"') + ) final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/zh/taggers/date.py b/fun_text_processing/inverse_text_normalization/zh/taggers/date.py index 6aec08450..0089957c6 100755 --- a/fun_text_processing/inverse_text_normalization/zh/taggers/date.py +++ b/fun_text_processing/inverse_text_normalization/zh/taggers/date.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.zh.utils import get_abs_path from fun_text_processing.inverse_text_normalization.zh.graph_utils import ( @@ -13,6 +12,7 @@ from pynini.lib import pynutil graph_digit = pynini.string_file(get_abs_path("data/numbers/digit.tsv")).optimize() graph_zero = pynini.string_file(get_abs_path("data/numbers/zero.tsv")).optimize() + def _get_month_graph(): """ Transducer for month, e.g. march -> march @@ -20,6 +20,7 @@ def _get_month_graph(): month_graph = pynini.string_file(get_abs_path("data/months.tsv")) return month_graph + def _get_day_graph(): """ Tranducer for day, @@ -27,19 +28,20 @@ def _get_day_graph(): day_graph = pynini.string_file(get_abs_path("data/days.tsv")) return day_graph + def _get_year_graph(): """ Transducer for year, """ - year_graph = graph_digit + (graph_digit|graph_zero)**3 + year_graph = graph_digit + (graph_digit | graph_zero) ** 3 year_graph |= graph_digit**2 - + return year_graph class DateFst(GraphFst): """ - Finite state transducer for classifying date, + Finite state transducer for classifying date, e.g. january fifth twenty twelve -> date { month: "january" day: "5" year: "2012" preserve_order: true } e.g. the fifth of january twenty twelve -> date { day: "5" month: "january" year: "2012" preserve_order: true } e.g. twenty twenty -> date { year: "2012" preserve_order: true } @@ -54,25 +56,22 @@ class DateFst(GraphFst): year_graph = _get_year_graph() year_graph = ( - pynutil.insert("year: \"") - + year_graph - + pynutil.delete("年") - + pynutil.insert("\"") + pynutil.insert('year: "') + year_graph + pynutil.delete("年") + pynutil.insert('"') ) month_graph = _get_month_graph() - month_graph = pynutil.insert("month: \"") + month_graph + pynutil.insert("\"") + month_graph = pynutil.insert('month: "') + month_graph + pynutil.insert('"') day_graph = _get_day_graph() - day_graph = pynutil.insert("day: \"") + day_graph + pynutil.insert("\"") - - graph_md = month_graph + pynutil.insert(" ") + day_graph + day_graph = pynutil.insert('day: "') + day_graph + pynutil.insert('"') + + graph_md = month_graph + pynutil.insert(" ") + day_graph graph_ym = year_graph + pynutil.insert(" ") + month_graph graph_ymd = year_graph + pynutil.insert(" ") + month_graph + pynutil.insert(" ") + day_graph - + graph_year = year_graph - final_graph = (graph_ymd | graph_ym | graph_md | graph_year) + final_graph = graph_ymd | graph_ym | graph_md | graph_year final_graph += pynutil.insert(" preserve_order: true") final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/zh/taggers/decimal.py b/fun_text_processing/inverse_text_normalization/zh/taggers/decimal.py index 1016035f1..a8e1a47e8 100755 --- a/fun_text_processing/inverse_text_normalization/zh/taggers/decimal.py +++ b/fun_text_processing/inverse_text_normalization/zh/taggers/decimal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.zh.utils import get_abs_path from fun_text_processing.inverse_text_normalization.zh.graph_utils import ( @@ -26,7 +25,7 @@ class DecimalFst(GraphFst): graph_digit = pynini.string_file(get_abs_path("data/numbers/digit.tsv")) graph_zero = pynini.string_file(get_abs_path("data/numbers/zero.tsv")) - graph_dot =pynini.string_file(get_abs_path("data/numbers/dot.tsv")) + graph_dot = pynini.string_file(get_abs_path("data/numbers/dot.tsv")) graph_digits = graph_digit | graph_zero @@ -37,11 +36,17 @@ class DecimalFst(GraphFst): point = pynini.cross("点", "") optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + (pynini.cross("负", "\"true\"") | pynini.cross("负的", "\"true\"")) + delete_extra_space, 0, 1 + pynutil.insert("negative: ") + + (pynini.cross("负", '"true"') | pynini.cross("负的", '"true"')) + + delete_extra_space, + 0, + 1, ) - graph_fractional = pynutil.insert("fractional_part: \"") + graph_decimal + pynutil.insert("\"") - graph_integer = pynutil.insert("integer_part: \"") + cardinal_graph + pynutil.insert("\"") + graph_fractional = ( + pynutil.insert('fractional_part: "') + graph_decimal + pynutil.insert('"') + ) + graph_integer = pynutil.insert('integer_part: "') + cardinal_graph + pynutil.insert('"') final_graph_wo_sign = ( pynini.closure(graph_integer, 0, 1) + point + pynutil.insert(" ") + graph_fractional ) diff --git a/fun_text_processing/inverse_text_normalization/zh/taggers/electronic.py b/fun_text_processing/inverse_text_normalization/zh/taggers/electronic.py index 573194fa6..02da3c594 100755 --- a/fun_text_processing/inverse_text_normalization/zh/taggers/electronic.py +++ b/fun_text_processing/inverse_text_normalization/zh/taggers/electronic.py @@ -1,7 +1,10 @@ - import pynini from fun_text_processing.inverse_text_normalization.zh.utils import get_abs_path -from fun_text_processing.inverse_text_normalization.zh.graph_utils import DAMO_ALPHA, GraphFst, insert_space +from fun_text_processing.inverse_text_normalization.zh.graph_utils import ( + DAMO_ALPHA, + GraphFst, + insert_space, +) from pynini.lib import pynutil @@ -25,35 +28,50 @@ class ElectronicFst(GraphFst): accepted_username = alpha_num | symbols process_dot = pynini.cross("点", ".") - username = (alpha_num + pynini.closure(delete_extra_space + accepted_username)) | pynutil.add_weight( - pynini.closure(DAMO_ALPHA, 1), weight=0.0001 - ) - username = pynutil.insert("username: \"") + username + pynutil.insert("\"") + username = ( + alpha_num + pynini.closure(delete_extra_space + accepted_username) + ) | pynutil.add_weight(pynini.closure(DAMO_ALPHA, 1), weight=0.0001) + username = pynutil.insert('username: "') + username + pynutil.insert('"') single_alphanum = pynini.closure(alpha_num + delete_extra_space) + alpha_num - server = single_alphanum | pynini.string_file(get_abs_path("data/electronic/server_name.tsv")) + server = single_alphanum | pynini.string_file( + get_abs_path("data/electronic/server_name.tsv") + ) domain = single_alphanum | pynini.string_file(get_abs_path("data/electronic/domain.tsv")) domain_graph = ( - pynutil.insert("domain: \"") + pynutil.insert('domain: "') + server + delete_extra_space + process_dot + delete_extra_space + domain - + pynutil.insert("\"") + + pynutil.insert('"') + ) + graph = ( + username + + delete_extra_space + + (pynutil.delete("at") | pynutil.delete("艾特")) + + insert_space + + delete_extra_space + + domain_graph ) - graph = username + delete_extra_space + (pynutil.delete("at")|pynutil.delete("艾特")) + insert_space + delete_extra_space + domain_graph ############# url ### protocol_end = pynini.cross(pynini.union("w w w", "www"), "www") - protocol_start = (pynini.cross("h t t p", "http") | pynini.cross("h t t p s", "https")) + pynini.cross( - "冒号斜杆斜杠", "://" - ) + protocol_start = ( + pynini.cross("h t t p", "http") | pynini.cross("h t t p s", "https") + ) + pynini.cross("冒号斜杆斜杠", "://") # .com, ending = ( delete_extra_space + symbols + delete_extra_space - + (domain | pynini.closure(accepted_username + delete_extra_space,) + accepted_username) + + ( + domain + | pynini.closure( + accepted_username + delete_extra_space, + ) + + accepted_username + ) ) protocol_default = ( @@ -64,12 +82,19 @@ class ElectronicFst(GraphFst): + pynini.closure(ending, 1) ).optimize() protocol = ( - pynini.closure(protocol_start, 0, 1) + protocol_end + pynini.closure(pynini.cross(" ", ""), 0, 1) + process_dot + pynini.closure(pynini.cross(" ", ""), 0, 1) + protocol_default + pynini.closure(protocol_start, 0, 1) + + protocol_end + + pynini.closure(pynini.cross(" ", ""), 0, 1) + + process_dot + + pynini.closure(pynini.cross(" ", ""), 0, 1) + + protocol_default ).optimize() - protocol |= pynini.closure(protocol_end + delete_extra_space + process_dot, 0, 1) + protocol_default + protocol |= ( + pynini.closure(protocol_end + delete_extra_space + process_dot, 0, 1) + protocol_default + ) - protocol = pynutil.insert("protocol: \"") + protocol.optimize() + pynutil.insert("\"") + protocol = pynutil.insert('protocol: "') + protocol.optimize() + pynutil.insert('"') graph |= protocol ######## diff --git a/fun_text_processing/inverse_text_normalization/zh/taggers/fraction.py b/fun_text_processing/inverse_text_normalization/zh/taggers/fraction.py index e64f4f4e9..9b2386eec 100755 --- a/fun_text_processing/inverse_text_normalization/zh/taggers/fraction.py +++ b/fun_text_processing/inverse_text_normalization/zh/taggers/fraction.py @@ -1,4 +1,3 @@ - from fun_text_processing.inverse_text_normalization.zh.graph_utils import ( DAMO_NOT_QUOTE, GraphFst, @@ -13,6 +12,7 @@ from fun_text_processing.inverse_text_normalization.zh.graph_utils import ( import pynini from pynini.lib import pynutil + class FractionFst(GraphFst): """ Finite state transducer for classifying fraction @@ -23,11 +23,16 @@ class FractionFst(GraphFst): # integer_part # numerator # denominator graph_cardinal = cardinal.graph_no_exception - #without the integerate part - #分子 + # without the integerate part + # 分子 numerator = pynutil.insert('numerator: "') + graph_cardinal + pynutil.insert('"') - #分母 - denominator = pynutil.insert('denominator: "') + graph_cardinal + pynutil.delete("分之") + pynutil.insert('"') + # 分母 + denominator = ( + pynutil.insert('denominator: "') + + graph_cardinal + + pynutil.delete("分之") + + pynutil.insert('"') + ) ## graph_fraction_component = denominator + pynutil.insert(" ") + numerator @@ -40,7 +45,11 @@ class FractionFst(GraphFst): ##负 optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + (pynini.cross("负", "\"true\"")|pynini.cross("负的", "\"true\"")) + DAMO_SPACE, 0, 1 + pynutil.insert("negative: ") + + (pynini.cross("负", '"true"') | pynini.cross("负的", '"true"')) + + DAMO_SPACE, + 0, + 1, ) graph = optional_graph_negative + graph diff --git a/fun_text_processing/inverse_text_normalization/zh/taggers/measure.py b/fun_text_processing/inverse_text_normalization/zh/taggers/measure.py index 311a5d246..f6594a5c2 100755 --- a/fun_text_processing/inverse_text_normalization/zh/taggers/measure.py +++ b/fun_text_processing/inverse_text_normalization/zh/taggers/measure.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.zh.utils import get_abs_path from fun_text_processing.inverse_text_normalization.zh.graph_utils import ( @@ -30,24 +29,28 @@ class MeasureFst(GraphFst): graph_unit_en = pynini.string_file(get_abs_path("data/measurements_en.tsv")) graph_unit_zh = pynini.string_file(get_abs_path("data/measurements_zh.tsv")) graph_sign = pynini.string_file(get_abs_path("data/numbers/sign.tsv")) - - graph_units = graph_unit_en | ((pynini.accep("亿") | pynini.accep("兆") | pynini.accep("万")).ques + graph_unit_zh) + graph_units = graph_unit_en | ( + (pynini.accep("亿") | pynini.accep("兆") | pynini.accep("万")).ques + graph_unit_zh + ) optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + (pynini.cross("负", "\"true\"") | pynini.cross("负的", "\"true\"")) + delete_extra_space, 0, 1 + pynutil.insert("negative: ") + + (pynini.cross("负", '"true"') | pynini.cross("负的", '"true"')) + + delete_extra_space, + 0, + 1, ) - percent = (pynutil.delete('百分') + - pynutil.delete('之').ques + (cardinal_graph | pynini.cross('百', '100')) - + pynutil.insert('%')) - - unit_all = ( - pynutil.insert("units: \"") - + graph_units - + pynutil.insert("\"") + percent = ( + pynutil.delete("百分") + + pynutil.delete("之").ques + + (cardinal_graph | pynini.cross("百", "100")) + + pynutil.insert("%") ) + unit_all = pynutil.insert('units: "') + graph_units + pynutil.insert('"') + subgraph_decimal = ( pynutil.insert("decimal { ") + optional_graph_negative @@ -59,9 +62,9 @@ class MeasureFst(GraphFst): subgraph_cardinal = ( pynutil.insert("cardinal { ") + optional_graph_negative - + pynutil.insert("integer: \"") + + pynutil.insert('integer: "') + cardinal_graph - + pynutil.insert("\"") + + pynutil.insert('"') + pynutil.insert(" }") + pynutil.insert(" ") + unit_all diff --git a/fun_text_processing/inverse_text_normalization/zh/taggers/money.py b/fun_text_processing/inverse_text_normalization/zh/taggers/money.py index 53dc5e230..548e3fcad 100755 --- a/fun_text_processing/inverse_text_normalization/zh/taggers/money.py +++ b/fun_text_processing/inverse_text_normalization/zh/taggers/money.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.zh.utils import get_abs_path from fun_text_processing.inverse_text_normalization.zh.graph_utils import ( @@ -33,17 +32,19 @@ class MoneyFst(GraphFst): decimal_graph = decimal.final_graph_wo_negative graph_code = pynini.string_file(get_abs_path("data/money/code.tsv")) graph_symbol = pynini.string_file(get_abs_path("data/money/symbol.tsv")) - - graph_unit = pynutil.insert("currency: \"") + (graph_code | graph_symbol) + pynutil.insert("\"") + + graph_unit = ( + pynutil.insert('currency: "') + (graph_code | graph_symbol) + pynutil.insert('"') + ) graph_integer = ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + cardinal_graph - + pynutil.insert("\"") + + pynutil.insert('"') + pynutil.insert(" ") + graph_unit ) - + graph_decimal = decimal_graph + pynutil.insert(" ") + graph_unit final_graph = graph_integer | graph_decimal diff --git a/fun_text_processing/inverse_text_normalization/zh/taggers/punctuation.py b/fun_text_processing/inverse_text_normalization/zh/taggers/punctuation.py index 3df5e38a2..47eae03ee 100755 --- a/fun_text_processing/inverse_text_normalization/zh/taggers/punctuation.py +++ b/fun_text_processing/inverse_text_normalization/zh/taggers/punctuation.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.zh.graph_utils import GraphFst from pynini.lib import pynutil @@ -13,9 +12,9 @@ class PunctuationFst(GraphFst): def __init__(self): super().__init__(name="punctuation", kind="classify") - s = "!#$%&\'()*+,-./:;<=>?@^_`{|}~" + s = "!#$%&'()*+,-./:;<=>?@^_`{|}~" punct = pynini.union(*s) - graph = pynutil.insert("name: \"") + punct + pynutil.insert("\"") + graph = pynutil.insert('name: "') + punct + pynutil.insert('"') self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/zh/taggers/telephone.py b/fun_text_processing/inverse_text_normalization/zh/taggers/telephone.py index 8f15ebbdc..24efdef1a 100755 --- a/fun_text_processing/inverse_text_normalization/zh/taggers/telephone.py +++ b/fun_text_processing/inverse_text_normalization/zh/taggers/telephone.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.zh.utils import get_abs_path from fun_text_processing.inverse_text_normalization.zh.graph_utils import ( @@ -24,7 +23,7 @@ def get_serial_number(cardinal): class TelephoneFst(GraphFst): """ - Finite state transducer for classifying telephone numbers, e.g. + Finite state transducer for classifying telephone numbers, e.g. one two three one two three five six seven eight -> { number_part: "123-123-5678" } This class also support card number and IP format. @@ -45,52 +44,58 @@ class TelephoneFst(GraphFst): graph_zero = pynini.string_file(get_abs_path("data/numbers/zero.tsv")) graph_dot = pynini.string_file(get_abs_path("data/numbers/dot.tsv")) - graph_digits = graph_digit | graph_zero - + graph_digits = graph_digit | graph_zero + # to handle cases like "one twenty three" phone_number_graph_w_country_code = graph_digits**11 - phone_number_graph_wo_country_code = graph_digits**3 | graph_digits**5 | graph_digits**11 | graph_digits**8 | graph_digits**12 + phone_number_graph_wo_country_code = ( + graph_digits**3 + | graph_digits**5 + | graph_digits**11 + | graph_digits**8 + | graph_digits**12 + ) country_code = ( - pynutil.insert("country_code: \"") + pynutil.insert('country_code: "') + pynini.closure(pynini.cross("加", "+"), 0, 1) + ((pynini.closure(graph_digits, 0, 2) + graph_digits)) - + pynutil.insert("\"") + + pynutil.insert('"') ) optional_country_code = pynini.closure(country_code + insert_space, 0, 1).optimize() - + grpah_phone_number_with_country_code = ( - pynutil.insert("number_part: \"") + pynutil.insert('number_part: "') + pynutil.add_weight(phone_number_graph_w_country_code.optimize(), 0.1) - + pynutil.insert("\"") + + pynutil.insert('"') ) - #grpah_phone_number_with_country_code = ( + # grpah_phone_number_with_country_code = ( # pynutil.insert("number_part: \"") # + phone_number_graph_w_country_code.optimize() # + pynutil.insert("\"") - #) + # ) graph_phone_with_country_code = optional_country_code + grpah_phone_number_with_country_code - + grpah_phone_number_wo_country_code = ( - pynutil.insert("number_part: \"") + pynutil.insert('number_part: "') + phone_number_graph_wo_country_code - + pynutil.insert("\"") - ) + + pynutil.insert('"') + ) # phone number graph = graph_phone_with_country_code | grpah_phone_number_wo_country_code - + # ip ip_graph = graph_digit.plus + (graph_dot + graph_digits.plus).plus - - graph |= pynutil.insert("number_part: \"") + ip_graph.optimize() + pynutil.insert("\"") + + graph |= pynutil.insert('number_part: "') + ip_graph.optimize() + pynutil.insert('"') graph |= ( - pynutil.insert("number_part: \"") + pynutil.insert('number_part: "') + pynutil.add_weight(get_serial_number(cardinal=cardinal), weight=0.0001) - + pynutil.insert("\"") + + pynutil.insert('"') ) final_graph = self.add_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/zh/taggers/time.py b/fun_text_processing/inverse_text_normalization/zh/taggers/time.py index fc2a37c62..9d9acb058 100755 --- a/fun_text_processing/inverse_text_normalization/zh/taggers/time.py +++ b/fun_text_processing/inverse_text_normalization/zh/taggers/time.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.zh.taggers.cardinal import CardinalFst from fun_text_processing.inverse_text_normalization.zh.utils import get_abs_path, num_to_word @@ -33,11 +32,17 @@ class TimeFst(GraphFst): second_graph = pynini.string_file(get_abs_path("data/time/second.tsv")) final_graph = ( - (pynutil.insert('noon: "') + noon_graph + pynutil.insert('" ')).ques + - pynutil.insert('hour: "') + hour_graph + pynutil.insert('"') + - pynutil.insert(' minute: "') + minute_graph + pynutil.delete('分').ques + pynutil.insert('"') + - (pynutil.insert(' second: "') + second_graph + pynutil.insert('"')).ques) - + (pynutil.insert('noon: "') + noon_graph + pynutil.insert('" ')).ques + + pynutil.insert('hour: "') + + hour_graph + + pynutil.insert('"') + + pynutil.insert(' minute: "') + + minute_graph + + pynutil.delete("分").ques + + pynutil.insert('"') + + (pynutil.insert(' second: "') + second_graph + pynutil.insert('"')).ques + ) + final_graph = self.add_tokens(final_graph.optimize()) self.fst = final_graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/zh/taggers/tokenize_and_classify.py b/fun_text_processing/inverse_text_normalization/zh/taggers/tokenize_and_classify.py index b4c86296d..60139d867 100755 --- a/fun_text_processing/inverse_text_normalization/zh/taggers/tokenize_and_classify.py +++ b/fun_text_processing/inverse_text_normalization/zh/taggers/tokenize_and_classify.py @@ -26,7 +26,7 @@ import logging class ClassifyFst(GraphFst): """ Final class that composes all other classification grammars. This class can process an entire sentence, that is lower cased. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. Args: @@ -79,10 +79,16 @@ class ClassifyFst(GraphFst): | pynutil.add_weight(word_graph, 100) ) - punct = pynutil.insert("tokens { ") + pynutil.add_weight(punct_graph, weight=1.1) + pynutil.insert(" }") + punct = ( + pynutil.insert("tokens { ") + + pynutil.add_weight(punct_graph, weight=1.1) + + pynutil.insert(" }") + ) token = pynutil.insert("tokens { ") + classify + pynutil.insert(" }") token_plus_punct = ( - pynini.closure(punct + pynutil.insert(" ")) + token + pynini.closure(pynutil.insert(" ") + punct) + pynini.closure(punct + pynutil.insert(" ")) + + token + + pynini.closure(pynutil.insert(" ") + punct) ) graph = token_plus_punct + pynini.closure(delete_extra_space + token_plus_punct) diff --git a/fun_text_processing/inverse_text_normalization/zh/taggers/whitelist.py b/fun_text_processing/inverse_text_normalization/zh/taggers/whitelist.py index 328736059..3dc9db3f2 100755 --- a/fun_text_processing/inverse_text_normalization/zh/taggers/whitelist.py +++ b/fun_text_processing/inverse_text_normalization/zh/taggers/whitelist.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.zh.utils import get_abs_path from fun_text_processing.inverse_text_normalization.zh.graph_utils import GraphFst, convert_space @@ -16,5 +15,5 @@ class WhiteListFst(GraphFst): super().__init__(name="whitelist", kind="classify") whitelist = pynini.string_file(get_abs_path("data/whitelist.tsv")).invert() - graph = pynutil.insert("name: \"") + convert_space(whitelist) + pynutil.insert("\"") + graph = pynutil.insert('name: "') + convert_space(whitelist) + pynutil.insert('"') self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/zh/taggers/word.py b/fun_text_processing/inverse_text_normalization/zh/taggers/word.py index 59d77bb8f..ac1e89a57 100755 --- a/fun_text_processing/inverse_text_normalization/zh/taggers/word.py +++ b/fun_text_processing/inverse_text_normalization/zh/taggers/word.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.zh.graph_utils import DAMO_NOT_SPACE, GraphFst from pynini.lib import pynutil @@ -12,5 +11,5 @@ class WordFst(GraphFst): def __init__(self): super().__init__(name="word", kind="classify") - word = pynutil.insert("name: \"") + pynini.closure(DAMO_NOT_SPACE, 1) + pynutil.insert("\"") + word = pynutil.insert('name: "') + pynini.closure(DAMO_NOT_SPACE, 1) + pynutil.insert('"') self.fst = word.optimize() diff --git a/fun_text_processing/inverse_text_normalization/zh/utils.py b/fun_text_processing/inverse_text_normalization/zh/utils.py index 281e632e4..6e9d84280 100755 --- a/fun_text_processing/inverse_text_normalization/zh/utils.py +++ b/fun_text_processing/inverse_text_normalization/zh/utils.py @@ -1,14 +1,12 @@ - - import csv import os from typing import Union import inflect -UNIT_1e01 = '十' -UNIT_1e02 = '百' -UNIT_1e03 = '千' -UNIT_1e04 = '万' +UNIT_1e01 = "十" +UNIT_1e02 = "百" +UNIT_1e03 = "千" +UNIT_1e04 = "万" _inflect = inflect.engine() @@ -18,23 +16,24 @@ def num_to_word(x: Union[str, int]): converts integer to spoken representation Args x: integer - Returns: spoken representation + Returns: spoken representation """ if isinstance(x, int): x = str(x) x = _inflect.number_to_words(str(x)).replace("-", " ").replace(",", "") return x + def get_abs_path(rel_path): """ Get absolute path Args: rel_path: relative path to this file - + Returns absolute path """ - return os.path.dirname(os.path.abspath(__file__)) + '/' + rel_path + return os.path.dirname(os.path.abspath(__file__)) + "/" + rel_path def load_labels(abs_path): @@ -53,7 +52,7 @@ def load_labels(abs_path): def augment_labels_with_punct_at_end(labels): """ - augments labels: if key ends on a punctuation that value does not have, add a new label + augments labels: if key ends on a punctuation that value does not have, add a new label where the value maintains the punctuation Args: diff --git a/fun_text_processing/inverse_text_normalization/zh/verbalizers/cardinal.py b/fun_text_processing/inverse_text_normalization/zh/verbalizers/cardinal.py index 19a8c5637..8e667c670 100755 --- a/fun_text_processing/inverse_text_normalization/zh/verbalizers/cardinal.py +++ b/fun_text_processing/inverse_text_normalization/zh/verbalizers/cardinal.py @@ -1,7 +1,10 @@ - - import pynini -from fun_text_processing.inverse_text_normalization.zh.graph_utils import DAMO_SPACE, DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.zh.graph_utils import ( + DAMO_SPACE, + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -16,18 +19,18 @@ class CardinalFst(GraphFst): optional_sign = pynini.closure( pynutil.delete("negative:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + DAMO_NOT_QUOTE - + pynutil.delete("\""), + + pynutil.delete('"'), 0, 1, ) graph = ( pynutil.delete("integer:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) self.numbers = graph graph = optional_sign + graph diff --git a/fun_text_processing/inverse_text_normalization/zh/verbalizers/date.py b/fun_text_processing/inverse_text_normalization/zh/verbalizers/date.py index f95f920ab..11ccc3877 100755 --- a/fun_text_processing/inverse_text_normalization/zh/verbalizers/date.py +++ b/fun_text_processing/inverse_text_normalization/zh/verbalizers/date.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.zh.graph_utils import ( DAMO_NOT_QUOTE, @@ -21,52 +20,53 @@ class DateFst(GraphFst): month = ( pynutil.delete("month:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) day = ( pynutil.delete("day:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) year = ( pynutil.delete("year:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') ) graph_ymd = ( - year + pynini.cross(" ", "年") + month + pynini.cross(" ", "月") + day + pynutil.insert("日") + year + + pynini.cross(" ", "年") + + month + + pynini.cross(" ", "月") + + day + + pynutil.insert("日") ) - graph_md = ( - month + pynini.cross(" ", "月") + day + pynutil.insert("日") - ) + graph_md = month + pynini.cross(" ", "月") + day + pynutil.insert("日") - graph_ym = ( - year + pynini.cross(" ", "年") + month + pynutil.insert("月") - ) + graph_ym = year + pynini.cross(" ", "年") + month + pynutil.insert("月") - graph_year = ( - year + pynutil.insert("年") - ) + graph_year = year + pynutil.insert("年") optional_preserve_order = pynini.closure( pynutil.delete("preserve_order:") + delete_space + pynutil.delete("true") + delete_space | pynutil.delete("field_order:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + DAMO_NOT_QUOTE - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space ) - final_graph = (graph_ymd | graph_ym | graph_md | graph_year) + delete_space + optional_preserve_order + final_graph = ( + (graph_ymd | graph_ym | graph_md | graph_year) + delete_space + optional_preserve_order + ) delete_tokens = self.delete_tokens(final_graph) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/inverse_text_normalization/zh/verbalizers/decimal.py b/fun_text_processing/inverse_text_normalization/zh/verbalizers/decimal.py index 7aa063e39..a9884c806 100755 --- a/fun_text_processing/inverse_text_normalization/zh/verbalizers/decimal.py +++ b/fun_text_processing/inverse_text_normalization/zh/verbalizers/decimal.py @@ -1,7 +1,9 @@ - - import pynini -from fun_text_processing.inverse_text_normalization.zh.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.zh.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -13,30 +15,30 @@ class DecimalFst(GraphFst): def __init__(self): super().__init__(name="decimal", kind="verbalize") - optionl_sign = pynini.closure(pynini.cross("negative: \"true\"", "-") + delete_space, 0, 1) + optionl_sign = pynini.closure(pynini.cross('negative: "true"', "-") + delete_space, 0, 1) integer = ( pynutil.delete("integer_part:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_integer = pynini.closure(integer + delete_space, 0, 1) fractional = ( pynutil.insert(".") + pynutil.delete("fractional_part:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_fractional = pynini.closure(fractional + delete_space, 0, 1) quantity = ( pynutil.delete("quantity:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_quantity = pynini.closure(pynutil.insert(" ") + quantity + delete_space, 0, 1) graph = optional_integer + optional_fractional + optional_quantity diff --git a/fun_text_processing/inverse_text_normalization/zh/verbalizers/electronic.py b/fun_text_processing/inverse_text_normalization/zh/verbalizers/electronic.py index 56881e703..661fe97ef 100755 --- a/fun_text_processing/inverse_text_normalization/zh/verbalizers/electronic.py +++ b/fun_text_processing/inverse_text_normalization/zh/verbalizers/electronic.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.zh.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.zh.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -15,24 +18,24 @@ class ElectronicFst(GraphFst): user_name = ( pynutil.delete("username:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) domain = ( pynutil.delete("domain:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) protocol = ( pynutil.delete("protocol:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) graph = user_name + delete_space + pynutil.insert("@") + domain diff --git a/fun_text_processing/inverse_text_normalization/zh/verbalizers/fraction.py b/fun_text_processing/inverse_text_normalization/zh/verbalizers/fraction.py index 003bb0a62..2b7fe4493 100755 --- a/fun_text_processing/inverse_text_normalization/zh/verbalizers/fraction.py +++ b/fun_text_processing/inverse_text_normalization/zh/verbalizers/fraction.py @@ -1,4 +1,3 @@ - from fun_text_processing.inverse_text_normalization.zh.graph_utils import ( DAMO_NOT_QUOTE, GraphFst, @@ -9,31 +8,34 @@ from fun_text_processing.inverse_text_normalization.zh.graph_utils import ( import pynini from pynini.lib import pynutil + class FractionFst(GraphFst): """ - Finite state transducer for verbalizing fraction, + Finite state transducer for verbalizing fraction, """ def __init__(self): super().__init__(name="fraction", kind="verbalize") optional_sign = pynini.closure(pynini.cross('negative: "true"', "-") + delete_space, 0, 1) numerator = ( - pynutil.delete('numerator:') + pynutil.delete("numerator:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) denominator = ( - pynutil.delete('denominator:') + pynutil.delete("denominator:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) - graph = (optional_sign + numerator + delete_space + pynutil.insert("/") + denominator).optimize() + graph = ( + optional_sign + numerator + delete_space + pynutil.insert("/") + denominator + ).optimize() self.numbers = graph delete_tokens = self.delete_tokens(optional_sign + graph) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/inverse_text_normalization/zh/verbalizers/measure.py b/fun_text_processing/inverse_text_normalization/zh/verbalizers/measure.py index ab4188dc1..78fd864f2 100755 --- a/fun_text_processing/inverse_text_normalization/zh/verbalizers/measure.py +++ b/fun_text_processing/inverse_text_normalization/zh/verbalizers/measure.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.zh.graph_utils import DAMO_CHAR, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.zh.graph_utils import ( + DAMO_CHAR, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -16,13 +19,13 @@ class MeasureFst(GraphFst): def __init__(self, decimal: GraphFst, cardinal: GraphFst): super().__init__(name="measure", kind="verbalize") - optional_sign = pynini.closure(pynini.cross("negative: \"true\"", "-"), 0, 1) + optional_sign = pynini.closure(pynini.cross('negative: "true"', "-"), 0, 1) unit = ( pynutil.delete("units:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space ) graph_decimal = ( diff --git a/fun_text_processing/inverse_text_normalization/zh/verbalizers/money.py b/fun_text_processing/inverse_text_normalization/zh/verbalizers/money.py index 5aaac24fc..74fedde8b 100755 --- a/fun_text_processing/inverse_text_normalization/zh/verbalizers/money.py +++ b/fun_text_processing/inverse_text_normalization/zh/verbalizers/money.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.inverse_text_normalization.zh.graph_utils import DAMO_CHAR, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.zh.graph_utils import ( + DAMO_CHAR, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -18,9 +21,9 @@ class MoneyFst(GraphFst): unit = ( pynutil.delete("currency:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) graph = unit + delete_space + decimal.numbers delete_tokens = self.delete_tokens(graph) diff --git a/fun_text_processing/inverse_text_normalization/zh/verbalizers/telephone.py b/fun_text_processing/inverse_text_normalization/zh/verbalizers/telephone.py index a4e408e14..475542028 100755 --- a/fun_text_processing/inverse_text_normalization/zh/verbalizers/telephone.py +++ b/fun_text_processing/inverse_text_normalization/zh/verbalizers/telephone.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.zh.graph_utils import DAMO_NOT_QUOTE, GraphFst from pynini.lib import pynutil @@ -8,17 +7,21 @@ class TelephoneFst(GraphFst): """ Finite state transducer for verbalizing telephone, e.g. telephone { number_part: "123-123-5678" } - -> 123-123-5678 + -> 123-123-5678 """ def __init__(self): super().__init__(name="telephone", kind="verbalize") - number_part = pynutil.delete("number_part: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") - optional_country_code = pynini.closure( - pynutil.delete("country_code: \"") + number_part = ( + pynutil.delete('number_part: "') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') + ) + optional_country_code = pynini.closure( + pynutil.delete('country_code: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + pynini.accep(" "), 0, 1, diff --git a/fun_text_processing/inverse_text_normalization/zh/verbalizers/time.py b/fun_text_processing/inverse_text_normalization/zh/verbalizers/time.py index 60a0d5e14..2fdc91970 100755 --- a/fun_text_processing/inverse_text_normalization/zh/verbalizers/time.py +++ b/fun_text_processing/inverse_text_normalization/zh/verbalizers/time.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.inverse_text_normalization.zh.graph_utils import ( DAMO_CHAR, @@ -20,29 +19,31 @@ class TimeFst(GraphFst): def __init__(self): super().__init__(name="time", kind="verbalize") - add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | (pynutil.insert("0") + DAMO_DIGIT) - + add_leading_zero_to_double_digit = (DAMO_DIGIT + DAMO_DIGIT) | ( + pynutil.insert("0") + DAMO_DIGIT + ) + hour = ( pynutil.delete("hour:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_DIGIT, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) minute = ( pynutil.delete("minute:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_DIGIT, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) second = ( pynutil.delete("second:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_DIGIT, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) suffix = ( @@ -50,9 +51,9 @@ class TimeFst(GraphFst): + pynutil.insert(" ") + pynutil.delete("noon:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_suffix = pynini.closure(suffix, 0, 1) zone = ( @@ -60,9 +61,9 @@ class TimeFst(GraphFst): + pynutil.insert(" ") + pynutil.delete("zone:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_zone = pynini.closure(zone, 0, 1) graph = ( diff --git a/fun_text_processing/inverse_text_normalization/zh/verbalizers/verbalize.py b/fun_text_processing/inverse_text_normalization/zh/verbalizers/verbalize.py index 63a57a3cf..0f66b2c38 100755 --- a/fun_text_processing/inverse_text_normalization/zh/verbalizers/verbalize.py +++ b/fun_text_processing/inverse_text_normalization/zh/verbalizers/verbalize.py @@ -1,4 +1,3 @@ - from fun_text_processing.inverse_text_normalization.zh.verbalizers.cardinal import CardinalFst from fun_text_processing.inverse_text_normalization.zh.verbalizers.date import DateFst from fun_text_processing.inverse_text_normalization.zh.verbalizers.decimal import DecimalFst @@ -15,7 +14,7 @@ from fun_text_processing.inverse_text_normalization.zh.graph_utils import GraphF class VerbalizeFst(GraphFst): """ Composes other verbalizer grammars. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. """ diff --git a/fun_text_processing/inverse_text_normalization/zh/verbalizers/verbalize_final.py b/fun_text_processing/inverse_text_normalization/zh/verbalizers/verbalize_final.py index 62badfc83..a6ecda16f 100755 --- a/fun_text_processing/inverse_text_normalization/zh/verbalizers/verbalize_final.py +++ b/fun_text_processing/inverse_text_normalization/zh/verbalizers/verbalize_final.py @@ -1,14 +1,17 @@ - import pynini from fun_text_processing.inverse_text_normalization.zh.verbalizers.verbalize import VerbalizeFst from fun_text_processing.inverse_text_normalization.zh.verbalizers.word import WordFst -from fun_text_processing.inverse_text_normalization.zh.graph_utils import GraphFst, delete_extra_space, delete_space +from fun_text_processing.inverse_text_normalization.zh.graph_utils import ( + GraphFst, + delete_extra_space, + delete_space, +) from pynini.lib import pynutil class VerbalizeFinalFst(GraphFst): """ - Finite state transducer that verbalizes an entire sentence, e.g. + Finite state transducer that verbalizes an entire sentence, e.g. tokens { name: "its" } tokens { time { hours: "12" minutes: "30" } } tokens { name: "now" } -> its 12:30 now """ diff --git a/fun_text_processing/inverse_text_normalization/zh/verbalizers/whitelist.py b/fun_text_processing/inverse_text_normalization/zh/verbalizers/whitelist.py index 42fca618e..902b373e2 100755 --- a/fun_text_processing/inverse_text_normalization/zh/verbalizers/whitelist.py +++ b/fun_text_processing/inverse_text_normalization/zh/verbalizers/whitelist.py @@ -1,7 +1,10 @@ - - import pynini -from fun_text_processing.inverse_text_normalization.zh.graph_utils import DAMO_CHAR, DAMO_SIGMA, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.zh.graph_utils import ( + DAMO_CHAR, + DAMO_SIGMA, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -16,9 +19,9 @@ class WhiteListFst(GraphFst): graph = ( pynutil.delete("name:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) - graph = graph @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + graph = graph @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) self.fst = graph.optimize() diff --git a/fun_text_processing/inverse_text_normalization/zh/verbalizers/word.py b/fun_text_processing/inverse_text_normalization/zh/verbalizers/word.py index 93c30f348..d93db03fa 100755 --- a/fun_text_processing/inverse_text_normalization/zh/verbalizers/word.py +++ b/fun_text_processing/inverse_text_normalization/zh/verbalizers/word.py @@ -1,6 +1,10 @@ - import pynini -from fun_text_processing.inverse_text_normalization.zh.graph_utils import DAMO_CHAR, DAMO_SIGMA, GraphFst, delete_space +from fun_text_processing.inverse_text_normalization.zh.graph_utils import ( + DAMO_CHAR, + DAMO_SIGMA, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -13,7 +17,13 @@ class WordFst(GraphFst): def __init__(self): super().__init__(name="word", kind="verbalize") chars = pynini.closure(DAMO_CHAR - " ", 1) - char = pynutil.delete("name:") + delete_space + pynutil.delete("\"") + chars + pynutil.delete("\"") - graph = char @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + char = ( + pynutil.delete("name:") + + delete_space + + pynutil.delete('"') + + chars + + pynutil.delete('"') + ) + graph = char @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) self.fst = graph.optimize() diff --git a/fun_text_processing/num2words/.coveragerc b/fun_text_processing/num2words/.coveragerc deleted file mode 100644 index 6d88b72e2..000000000 --- a/fun_text_processing/num2words/.coveragerc +++ /dev/null @@ -1,5 +0,0 @@ -[run] -branch = true -source = - num2words - tests diff --git a/fun_text_processing/num2words/.gitignore b/fun_text_processing/num2words/.gitignore deleted file mode 100644 index e5a9b78de..000000000 --- a/fun_text_processing/num2words/.gitignore +++ /dev/null @@ -1,7 +0,0 @@ -*.pyc -build -dist -.idea/ -*.egg-info -/.tox -.eggs/ diff --git a/fun_text_processing/num2words/.travis.yml b/fun_text_processing/num2words/.travis.yml deleted file mode 100644 index aa600d03c..000000000 --- a/fun_text_processing/num2words/.travis.yml +++ /dev/null @@ -1,20 +0,0 @@ -sudo: false -language: python -python: - - "2.7" - - "3.4" - - "3.5" - - "3.6" -matrix: - include: - - { python: 3.6, env: TOXENV=flake8 } - - { python: 3.6, env: TOXENV=isort } - # Py37 requires xenial distrubution and sudo - # See travis-ci/travis-ci#9069 - - { python: 3.7, dist: xenial, sudo: true } - -install: - - pip install tox-travis - - pip install coveralls -script: tox -after_success: if [ -e .coverage ]; then coveralls; fi diff --git a/fun_text_processing/num2words/MANIFEST.in b/fun_text_processing/num2words/MANIFEST.in deleted file mode 100644 index 2d05b8cd9..000000000 --- a/fun_text_processing/num2words/MANIFEST.in +++ /dev/null @@ -1 +0,0 @@ -include bin/num2words diff --git a/fun_text_processing/num2words/docker-compose.yml b/fun_text_processing/num2words/docker-compose.yml deleted file mode 100644 index d2b033ff5..000000000 --- a/fun_text_processing/num2words/docker-compose.yml +++ /dev/null @@ -1,7 +0,0 @@ -version: '3.0' -services: - web: - image: python:3-alpine - command: python3 -m http.server 8080 - volumes: - - .:/num2words diff --git a/fun_text_processing/num2words/num2words/__init__.py b/fun_text_processing/num2words/num2words/__init__.py deleted file mode 100644 index 3d0192bd6..000000000 --- a/fun_text_processing/num2words/num2words/__init__.py +++ /dev/null @@ -1,147 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import unicode_literals - -from . import (lang_AR, lang_AZ, lang_BN, lang_BG, lang_CA, lang_CZ, lang_DE, lang_DK, - lang_EN, lang_EL, lang_EN_IN, - lang_ES, lang_ES_CO, lang_ES_NI, lang_ES_VE, lang_FI, lang_FR, - lang_FR_BE, lang_FR_CH, lang_FR_DZ, lang_HE, lang_HU, lang_ID, - lang_IT, lang_IW, lang_JA, lang_KH, lang_KN, lang_KO, lang_KZ, lang_LT, - lang_LO, lang_LV, lang_MN, lang_MY, - lang_NE, lang_NL, lang_NO, lang_PL, lang_PT, lang_PT_BR, lang_RO, - lang_RU, lang_SI, lang_SL, lang_SR, lang_SV, lang_SW, lang_TE, lang_TH, lang_TR, - lang_UK, lang_UR, lang_VI) - -CONVERTER_CLASSES = { - 'ar': lang_AR.Num2Word_AR(), - 'az': lang_AZ.Num2Word_AZ(), - 'bn': lang_BN.Num2Word_BN(), - 'bg': lang_BG.Num2Word_BG(), - 'ca': lang_CA.Num2Word_CA(), - 'cz': lang_CZ.Num2Word_CZ(), - 'en': lang_EN.Num2Word_EN(), - 'el': lang_EL.Num2Word_EL(), - 'en_IN': lang_EN_IN.Num2Word_EN_IN(), - 'fr': lang_FR.Num2Word_FR(), - 'fr_CH': lang_FR_CH.Num2Word_FR_CH(), - 'fr_BE': lang_FR_BE.Num2Word_FR_BE(), - 'fr_DZ': lang_FR_DZ.Num2Word_FR_DZ(), - 'de': lang_DE.Num2Word_DE(), - 'fi': lang_FI.Num2Word_FI(), - 'es': lang_ES.Num2Word_ES(), - 'es_CO': lang_ES_CO.Num2Word_ES_CO(), - 'es_NI': lang_ES_NI.Num2Word_ES_NI(), - 'es_VE': lang_ES_VE.Num2Word_ES_VE(), - 'id': lang_ID.Num2Word_ID(), - 'iw': lang_IW.Num2Word_IW(), - 'ja': lang_JA.Num2Word_JA(), - 'kh': lang_KH.Num2Word_KH(), - 'kn': lang_KN.Num2Word_KN(), - 'ko': lang_KO.Num2Word_KO(), - 'kz': lang_KZ.Num2Word_KZ(), - 'lo': lang_LO.Num2Word_LO(), - 'lt': lang_LT.Num2Word_LT(), - 'lv': lang_LV.Num2Word_LV(), - 'mn': lang_MN.Num2Word_MN(), - 'my': lang_MY.Num2Word_MY(), - 'pl': lang_PL.Num2Word_PL(), - 'ro': lang_RO.Num2Word_RO(), - 'ru': lang_RU.Num2Word_RU(), - 'si': lang_SI.Num2Word_SI(), - 'sl': lang_SL.Num2Word_SL(), - 'sr': lang_SR.Num2Word_SR(), - 'sv': lang_SV.Num2Word_SV(), - 'sw': lang_SW.Num2Word_SW(), - 'no': lang_NO.Num2Word_NO(), - 'dk': lang_DK.Num2Word_DK(), - 'pt': lang_PT.Num2Word_PT(), - 'pt_BR': lang_PT_BR.Num2Word_PT_BR(), - 'he': lang_HE.Num2Word_HE(), - 'it': lang_IT.Num2Word_IT(), - 'vi': lang_VI.Num2Word_VI(), - 'th': lang_TH.Num2Word_TH(), - 'tr': lang_TR.Num2Word_TR(), - 'ne': lang_NE.Num2Word_NE(), - 'nl': lang_NL.Num2Word_NL(), - 'uk': lang_UK.Num2Word_UK(), - 'ur': lang_UR.Num2Word_UR(), - 'te': lang_TE.Num2Word_TE(), - 'hu': lang_HU.Num2Word_HU() -} - -CONVERTES_TYPES = ['cardinal', 'ordinal', 'ordinal_num', 'year', 'currency'] - -'ar': lang_AR.Num2Word_AR(), -'az': lang_AZ.Num2Word_AZ(), -'bn': lang_BN.Num2Word_BN(), -'bg': lang_BG.Num2Word_BG(), -'ca': lang_CA.Num2Word_CA(), -'cz': lang_CZ.Num2Word_CZ(), -'en': lang_EN.Num2Word_EN(), -'el': lang_EL.Num2Word_EL(), -'en_IN': lang_EN_IN.Num2Word_EN_IN(), -'fr': lang_FR.Num2Word_FR(), -'fr_CH': lang_FR_CH.Num2Word_FR_CH(), -'fr_BE': lang_FR_BE.Num2Word_FR_BE(), -'fr_DZ': lang_FR_DZ.Num2Word_FR_DZ(), -'de': lang_DE.Num2Word_DE(), -'fi': lang_FI.Num2Word_FI(), -'es': lang_ES.Num2Word_ES(), -'es_CO': lang_ES_CO.Num2Word_ES_CO(), -'es_NI': lang_ES_NI.Num2Word_ES_NI(), -'es_VE': lang_ES_VE.Num2Word_ES_VE(), -'id': lang_ID.Num2Word_ID(), -'iw': lang_IW.Num2Word_IW(), -'ja': lang_JA.Num2Word_JA(), -'kh': lang_KH.Num2Word_KH(), -'kn': lang_KN.Num2Word_KN(), -'ko': lang_KO.Num2Word_KO(), -'kz': lang_KZ.Num2Word_KZ(), -'lo': lang_LO.Num2Word_LO(), -'lt': lang_LT.Num2Word_LT(), -'lv': lang_LV.Num2Word_LV(), -'mn': lang_MN.Num2Word_MN(), -'my': lang_MY.Num2Word_MY(), -'pl': lang_PL.Num2Word_PL(), -'ro': lang_RO.Num2Word_RO(), -'ru': lang_RU.Num2Word_RU(), -'si': lang_SI.Num2Word_SI(), -'sl': lang_SL.Num2Word_SL(), -'sr': lang_SR.Num2Word_SR(), -'sv': lang_SV.Num2Word_SV(), -'sw': lang_SW.Num2Word_SW(), -'no': lang_NO.Num2Word_NO(), -'dk': lang_DK.Num2Word_DK(), -'pt': lang_PT.Num2Word_PT(), -'pt_BR': lang_PT_BR.Num2Word_PT_BR(), -'he': lang_HE.Num2Word_HE(), -'it': lang_IT.Num2Word_IT(), -'vi': lang_VI.Num2Word_VI(), -'th': lang_TH.Num2Word_TH(), -'tr': lang_TR.Num2Word_TR(), -'ne': lang_NE.Num2Word_NE(), -'nl': lang_NL.Num2Word_NL(), -'uk': lang_UK.Num2Word_UK(), -'ur': lang_UR.Num2Word_UR(), -'te': lang_TE.Num2Word_TE(), -'hu': lang_HU.Num2Word_HU() -def num2words(number, ordinal=False, lang='en', to='cardinal', **kwargs): - # We try the full language first - if lang not in CONVERTER_CLASSES: - # ... and then try only the first 2 letters - lang = lang[:2] - if lang not in CONVERTER_CLASSES: - raise NotImplementedError() - converter = CONVERTER_CLASSES[lang] - - if isinstance(number, str): - number = converter.str_to_number(number) - - # backwards compatible - if ordinal: - return converter.to_ordinal(number) - - if to not in CONVERTES_TYPES: - raise NotImplementedError() - - return getattr(converter, 'to_{}'.format(to))(number, **kwargs) diff --git a/fun_text_processing/num2words/num2words/base.py b/fun_text_processing/num2words/num2words/base.py deleted file mode 100644 index 738eef223..000000000 --- a/fun_text_processing/num2words/num2words/base.py +++ /dev/null @@ -1,300 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import unicode_literals - -import math -from collections import OrderedDict -from decimal import Decimal - -from .compat import to_s -from .currency import parse_currency_parts, prefix_currency - - -class Num2Word_Base(object): - CURRENCY_FORMS = {} - CURRENCY_ADJECTIVES = {} - - def __init__(self): - self.is_title = False - self.precision = 2 - self.exclude_title = [] - self.negword = "(-) " - self.pointword = "(.)" - self.errmsg_nonnum = "type(%s) not in [long, int, float]" - self.errmsg_floatord = "Cannot treat float %s as ordinal." - self.errmsg_negord = "Cannot treat negative num %s as ordinal." - self.errmsg_toobig = "abs(%s) must be less than %s." - - self.setup() - - # uses cards - if any(hasattr(self, field) for field in - ['high_numwords', 'mid_numwords', 'low_numwords']): - self.cards = OrderedDict() - self.set_numwords() - self.MAXVAL = 1000 * list(self.cards.keys())[0] - - def set_numwords(self): - self.set_high_numwords(self.high_numwords) - self.set_mid_numwords(self.mid_numwords) - self.set_low_numwords(self.low_numwords) - # self.set_labeled_numbers(self.labeled_numbers) - - def set_high_numwords(self, *args): - raise NotImplementedError - - # def set_labeled_numbers(self, labeled_numbers): - # for key, val in labeled_numbers.items(): - # self.cards[key] = val - - def set_mid_numwords(self, mid): - for key, val in mid: - self.cards[key] = val - - def set_low_numwords(self, numwords): - for word, n in zip(numwords, range(len(numwords) - 1, -1, -1)): - self.cards[n] = word - - def splitnum(self, value): - for elem in self.cards: - if elem > value: - continue - - out = [] - if value == 0: - div, mod = 1, 0 - else: - div, mod = divmod(value, elem) - - if div == 1: - out.append((self.cards[1], 1)) - else: - if div == value: # The system tallies, eg Roman Numerals - return [(div * self.cards[elem], div*elem)] - out.append(self.splitnum(div)) - - out.append((self.cards[elem], elem)) - - if mod: - out.append(self.splitnum(mod)) - - return out - - def parse_minus(self, num_str): - """Detach minus and return it as symbol with new num_str.""" - if num_str.startswith('-'): - # Extra spacing to compensate if there is no minus. - return '%s ' % self.negword, num_str[1:] - return '', num_str - - def str_to_number(self, value): - return Decimal(value) - - def to_cardinal(self, value): - try: - assert int(value) == value - except (ValueError, TypeError, AssertionError): - return self.to_cardinal_float(value) - - out = "" - if value < 0: - value = abs(value) - out = self.negword - - if value >= self.MAXVAL: - raise OverflowError(self.errmsg_toobig % (value, self.MAXVAL)) - - val = self.splitnum(value) - words, num = self.clean(val) - return self.title(out + words) - - def float2tuple(self, value): - pre = int(value) - - # Simple way of finding decimal places to update the precision - self.precision = abs(Decimal(str(value)).as_tuple().exponent) - - post = abs(value - pre) * 10**self.precision - if abs(round(post) - post) < 0.01: - # We generally floor all values beyond our precision (rather than - # rounding), but in cases where we have something like 1.239999999, - # which is probably due to python's handling of floats, we actually - # want to consider it as 1.24 instead of 1.23 - post = int(round(post)) - else: - post = int(math.floor(post)) - - return pre, post - - def to_cardinal_float(self, value): - try: - float(value) == value - except (ValueError, TypeError, AssertionError, AttributeError): - raise TypeError(self.errmsg_nonnum % value) - - pre, post = self.float2tuple(float(value)) - - post = str(post) - post = '0' * (self.precision - len(post)) + post - - out = [self.to_cardinal(pre)] - if self.precision: - out.append(self.title(self.pointword)) - - for i in range(self.precision): - curr = int(post[i]) - out.append(to_s(self.to_cardinal(curr))) - - return " ".join(out) - - def merge(self, curr, next): - raise NotImplementedError - - def clean(self, val): - out = val - while len(val) != 1: - out = [] - left, right = val[:2] - if isinstance(left, tuple) and isinstance(right, tuple): - out.append(self.merge(left, right)) - if val[2:]: - out.append(val[2:]) - else: - for elem in val: - if isinstance(elem, list): - if len(elem) == 1: - out.append(elem[0]) - else: - out.append(self.clean(elem)) - else: - out.append(elem) - val = out - return out[0] - - def title(self, value): - if self.is_title: - out = [] - value = value.split() - for word in value: - if word in self.exclude_title: - out.append(word) - else: - out.append(word[0].upper() + word[1:]) - value = " ".join(out) - return value - - def verify_ordinal(self, value): - if not value == int(value): - raise TypeError(self.errmsg_floatord % value) - if not abs(value) == value: - raise TypeError(self.errmsg_negord % value) - - def to_ordinal(self, value): - return self.to_cardinal(value) - - def to_ordinal_num(self, value): - return value - - # Trivial version - def inflect(self, value, text): - text = text.split("/") - if value == 1: - return text[0] - return "".join(text) - - # //CHECK: generalise? Any others like pounds/shillings/pence? - def to_splitnum(self, val, hightxt="", lowtxt="", jointxt="", - divisor=100, longval=True, cents=True): - out = [] - - if isinstance(val, float): - high, low = self.float2tuple(val) - else: - try: - high, low = val - except TypeError: - high, low = divmod(val, divisor) - - if high: - hightxt = self.title(self.inflect(high, hightxt)) - out.append(self.to_cardinal(high)) - if low: - if longval: - if hightxt: - out.append(hightxt) - if jointxt: - out.append(self.title(jointxt)) - elif hightxt: - out.append(hightxt) - - if low: - if cents: - out.append(self.to_cardinal(low)) - else: - out.append("%02d" % low) - if lowtxt and longval: - out.append(self.title(self.inflect(low, lowtxt))) - - return " ".join(out) - - def to_year(self, value, **kwargs): - return self.to_cardinal(value) - - def pluralize(self, n, forms): - """ - Should resolve gettext form: - http://docs.translatehouse.org/projects/localization-guide/en/latest/l10n/pluralforms.html - """ - raise NotImplementedError - - def _money_verbose(self, number, currency): - return self.to_cardinal(number) - - def _cents_verbose(self, number, currency): - return self.to_cardinal(number) - - def _cents_terse(self, number, currency): - return "%02d" % number - - def to_currency(self, val, currency='EUR', cents=True, separator=',', - adjective=False): - """ - Args: - val: Numeric value - currency (str): Currency code - cents (bool): Verbose cents - separator (str): Cent separator - adjective (bool): Prefix currency name with adjective - Returns: - str: Formatted string - - """ - left, right, is_negative = parse_currency_parts(val) - - try: - cr1, cr2 = self.CURRENCY_FORMS[currency] - - except KeyError: - raise NotImplementedError( - 'Currency code "%s" not implemented for "%s"' % - (currency, self.__class__.__name__)) - - if adjective and currency in self.CURRENCY_ADJECTIVES: - cr1 = prefix_currency(self.CURRENCY_ADJECTIVES[currency], cr1) - - minus_str = "%s " % self.negword if is_negative else "" - money_str = self._money_verbose(left, currency) - cents_str = self._cents_verbose(right, currency) \ - if cents else self._cents_terse(right, currency) - - return u'%s%s %s%s %s %s' % ( - minus_str, - money_str, - self.pluralize(left, cr1), - separator, - cents_str, - self.pluralize(right, cr2) - ) - - def setup(self): - pass diff --git a/fun_text_processing/num2words/num2words/compat.py b/fun_text_processing/num2words/num2words/compat.py deleted file mode 100644 index ef030eeef..000000000 --- a/fun_text_processing/num2words/num2words/compat.py +++ /dev/null @@ -1,14 +0,0 @@ -# -*- coding: utf-8 -*- - - -try: - strtype = basestring -except NameError: - strtype = str - - -def to_s(val): - try: - return unicode(val) - except NameError: - return str(val) diff --git a/fun_text_processing/num2words/num2words/currency.py b/fun_text_processing/num2words/num2words/currency.py deleted file mode 100644 index d7f999f4a..000000000 --- a/fun_text_processing/num2words/num2words/currency.py +++ /dev/null @@ -1,35 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import division - -from decimal import ROUND_HALF_UP, Decimal - - -def parse_currency_parts(value, is_int_with_cents=True): - if isinstance(value, int): - if is_int_with_cents: - # assume cents if value is integer - negative = value < 0 - value = abs(value) - integer, cents = divmod(value, 100) - else: - negative = value < 0 - integer, cents = abs(value), 0 - - else: - value = Decimal(value) - value = value.quantize( - Decimal('.01'), - rounding=ROUND_HALF_UP - ) - negative = value < 0 - value = abs(value) - integer, fraction = divmod(value, 1) - integer = int(integer) - cents = int(fraction * 100) - - return integer, cents, negative - - -def prefix_currency(prefix, base): - return tuple("%s %s" % (prefix, i) for i in base) diff --git a/fun_text_processing/num2words/num2words/lang_AR.py b/fun_text_processing/num2words/num2words/lang_AR.py deleted file mode 100644 index a0027a835..000000000 --- a/fun_text_processing/num2words/num2words/lang_AR.py +++ /dev/null @@ -1,352 +0,0 @@ -# -*- coding: utf-8 -*- -# Copyright (c) 2003, Taro Ogawa. All Rights Reserved. -# Copyright (c) 2013, Savoir-faire Linux inc. All Rights Reserved. -# Copyright (c) 2018, Abdullah Alhazmy, Alhazmy13. All Rights Reserved. - - -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, -# MA 02110-1301 USA - -import re -from decimal import Decimal -from math import floor - -CURRENCY_SR = [("ريال", "ريالان", "ريالات", "ريالاً"), - ("هللة", "هللتان", "هللات", "هللة")] -CURRENCY_EGP = [("جنيه", "جنيهان", "جنيهات", "جنيهاً"), - ("قرش", "قرشان", "قروش", "قرش")] -CURRENCY_KWD = [("دينار", "ديناران", "دينارات", "ديناراً"), - ("فلس", "فلسان", "فلس", "فلس")] - -ARABIC_ONES = [ - "", "واحد", "اثنان", "ثلاثة", "أربعة", "خمسة", "ستة", "سبعة", "ثمانية", - "تسعة", - "عشرة", "أحد عشر", "اثنا عشر", "ثلاثة عشر", "أربعة عشر", "خمسة عشر", - "ستة عشر", "سبعة عشر", "ثمانية عشر", - "تسعة عشر" -] - - -class Num2Word_AR(object): - errmsg_too_big = "Too large" - max_num = 10 ** 36 - - def __init__(self): - self.number = 0 - self.arabicPrefixText = "" - self.arabicSuffixText = "" - self.integer_value = 0 - self._decimalValue = "" - self.partPrecision = 2 - self.currency_unit = CURRENCY_SR[0] - self.currency_subunit = CURRENCY_SR[1] - self.isCurrencyPartNameFeminine = True - self.isCurrencyNameFeminine = False - self.separator = 'و' - - self.arabicOnes = ARABIC_ONES - self.arabicFeminineOnes = [ - "", "إحدى", "اثنتان", "ثلاث", "أربع", "خمس", "ست", "سبع", "ثمان", - "تسع", - "عشر", "إحدى عشرة", "اثنتا عشرة", "ثلاث عشرة", "أربع عشرة", - "خمس عشرة", "ست عشرة", "سبع عشرة", "ثماني عشرة", - "تسع عشرة" - ] - self.arabicOrdinal = [ - "", "اول", "ثاني", "ثالث", "رابع", "خامس", "سادس", "سابع", "ثامن", - "تاسع", "عاشر", "حادي عشر", "ثاني عشر", "ثالث عشر", "رابع عشر", - "خامس عشر", "سادس عشر", "سابع عشر", "ثامن عشر", "تاسع عشر" - ] - self.arabicTens = [ - "عشرون", "ثلاثون", "أربعون", "خمسون", "ستون", "سبعون", "ثمانون", - "تسعون" - ] - self.arabicHundreds = [ - "", "مائة", "مئتان", "ثلاثمائة", "أربعمائة", "خمسمائة", "ستمائة", - "سبعمائة", "ثمانمائة", "تسعمائة" - ] - self.arabicAppendedTwos = [ - "مئتا", "ألفا", "مليونا", "مليارا", "تريليونا", "كوادريليونا", - "كوينتليونا", "سكستيليونا" - ] - self.arabicTwos = [ - "مئتان", "ألفان", "مليونان", "ملياران", "تريليونان", - "كوادريليونان", "كوينتليونان", "سكستيليونان" - ] - self.arabicGroup = [ - "مائة", "ألف", "مليون", "مليار", "تريليون", "كوادريليون", - "كوينتليون", "سكستيليون" - ] - self.arabicAppendedGroup = [ - "", "ألفاً", "مليوناً", "ملياراً", "تريليوناً", "كوادريليوناً", - "كوينتليوناً", "سكستيليوناً" - ] - self.arabicPluralGroups = [ - "", "آلاف", "ملايين", "مليارات", "تريليونات", "كوادريليونات", - "كوينتليونات", "سكستيليونات" - ] - - def number_to_arabic(self, arabic_prefix_text, arabic_suffix_text): - self.arabicPrefixText = arabic_prefix_text - self.arabicSuffixText = arabic_suffix_text - self.extract_integer_and_decimal_parts() - - def extract_integer_and_decimal_parts(self): - re.split('\\.', str(self.number)) - splits = re.split('\\.', str(self.number)) - - self.integer_value = int(splits[0]) - if len(splits) > 1: - self._decimalValue = int(self.decimal_value(splits[1])) - else: - self._decimalValue = 0 - - def decimal_value(self, decimal_part): - - if self.partPrecision is not len(decimal_part): - decimal_part_length = len(decimal_part) - - decimal_part_builder = decimal_part - for i in range(0, self.partPrecision - decimal_part_length): - decimal_part_builder += '0' - decimal_part = decimal_part_builder - - if len(decimal_part) <= self.partPrecision: - dec = len(decimal_part) - else: - dec = self.partPrecision - result = decimal_part[0: dec] - else: - result = decimal_part - - for i in range(len(result), self.partPrecision): - result += '0' - return result - - def digit_feminine_status(self, digit, group_level): - if group_level == -1: - if self.isCurrencyPartNameFeminine: - return self.arabicFeminineOnes[int(digit)] - else: - return self.arabicOnes[int(digit)] - elif group_level == 0: - if self.isCurrencyNameFeminine: - return self.arabicFeminineOnes[int(digit)] - else: - return self.arabicOnes[int(digit)] - - else: - return self.arabicOnes[int(digit)] - - def process_arabic_group(self, group_number, group_level, - remaining_number): - tens = Decimal(group_number) % Decimal(100) - hundreds = Decimal(group_number) / Decimal(100) - ret_val = "" - - if int(hundreds) > 0: - if tens == 0 and int(hundreds) == 2: - ret_val = "{}".format(self.arabicAppendedTwos[0]) - else: - ret_val = "{}".format(self.arabicHundreds[int(hundreds)]) - - if tens > 0: - if tens < 20: - if tens == 2 and int(hundreds) == 0 and group_level > 0: - if self.integer_value in [2000, 2000000, 2000000000, - 2000000000000, 2000000000000000, - 2000000000000000000]: - ret_val = "{}".format( - self.arabicAppendedTwos[int(group_level)]) - else: - ret_val = "{}".format( - self.arabicTwos[int(group_level)]) - else: - if ret_val != "": - ret_val += " و " - - if tens == 1 and group_level > 0 and hundreds == 0: - ret_val += "" - elif (tens == 1 or tens == 2) and ( - group_level == 0 or group_level == -1) and \ - hundreds == 0 and remaining_number == 0: - ret_val += "" - else: - ret_val += self.digit_feminine_status(int(tens), - group_level) - else: - ones = tens % 10 - tens = (tens / 10) - 2 - if ones > 0: - if ret_val != "" and tens < 4: - ret_val += " و " - - ret_val += self.digit_feminine_status(ones, group_level) - if ret_val != "" and ones != 0: - ret_val += " و " - - ret_val += self.arabicTens[int(tens)] - - return ret_val - - def convert(self, value): - self.number = "{:.9f}".format(value) - self.number_to_arabic(self.arabicPrefixText, self.arabicSuffixText) - return self.convert_to_arabic() - - def convert_to_arabic(self): - temp_number = Decimal(self.number) - - if temp_number == Decimal(0): - return "صفر" - - decimal_string = self.process_arabic_group(self._decimalValue, - -1, - Decimal(0)) - ret_val = "" - group = 0 - - while temp_number > Decimal(0): - - number_to_process = int( - Decimal(str(temp_number)) % Decimal(str(1000))) - temp_number = int(Decimal(temp_number) / Decimal(1000)) - - group_description = \ - self.process_arabic_group(number_to_process, - group, - Decimal(floor(temp_number))) - if group_description != '': - if group > 0: - if ret_val != "": - ret_val = "{} و {}".format("", ret_val) - if number_to_process != 2: - if number_to_process % 100 != 1: - if 3 <= number_to_process <= 10: - ret_val = "{} {}".format( - self.arabicPluralGroups[group], ret_val) - else: - if ret_val != "": - ret_val = "{} {}".format( - self.arabicAppendedGroup[group], - ret_val) - else: - ret_val = "{} {}".format( - self.arabicGroup[group], ret_val) - - else: - ret_val = "{} {}".format(self.arabicGroup[group], - ret_val) - ret_val = "{} {}".format(group_description, ret_val) - group += 1 - formatted_number = "" - if self.arabicPrefixText != "": - formatted_number += "{} ".format(self.arabicPrefixText) - formatted_number += ret_val - if self.integer_value != 0: - remaining100 = int(self.integer_value % 100) - - if remaining100 == 0: - formatted_number += self.currency_unit[0] - elif remaining100 == 1: - formatted_number += self.currency_unit[0] - elif remaining100 == 2: - if self.integer_value == 2: - formatted_number += self.currency_unit[1] - else: - formatted_number += self.currency_unit[0] - elif 3 <= remaining100 <= 10: - formatted_number += self.currency_unit[2] - elif 11 <= remaining100 <= 99: - formatted_number += self.currency_unit[3] - if self._decimalValue != 0: - formatted_number += " {} ".format(self.separator) - formatted_number += decimal_string - - if self._decimalValue != 0: - formatted_number += " " - remaining100 = int(self._decimalValue % 100) - - if remaining100 == 0: - formatted_number += self.currency_subunit[0] - elif remaining100 == 1: - formatted_number += self.currency_subunit[0] - elif remaining100 == 2: - formatted_number += self.currency_subunit[1] - elif 3 <= remaining100 <= 10: - formatted_number += self.currency_subunit[2] - elif 11 <= remaining100 <= 99: - formatted_number += self.currency_subunit[3] - - if self.arabicSuffixText != "": - formatted_number += " {}".format(self.arabicSuffixText) - - return formatted_number - - def validate_number(self, number): - if number >= self.max_num: - raise OverflowError(self.errmsg_too_big) - return number - - def set_currency_prefer(self, currency): - if currency == 'EGP': - self.currency_unit = CURRENCY_EGP[0] - self.currency_subunit = CURRENCY_EGP[1] - elif currency == 'KWD': - self.currency_unit = CURRENCY_KWD[0] - self.currency_subunit = CURRENCY_KWD[1] - else: - self.currency_unit = CURRENCY_SR[0] - self.currency_subunit = CURRENCY_SR[1] - - def to_currency(self, value, currency='SR', prefix='', suffix=''): - self.set_currency_prefer(currency) - self.isCurrencyNameFeminine = False - self.separator = "و" - self.arabicOnes = ARABIC_ONES - self.arabicPrefixText = prefix - self.arabicSuffixText = suffix - return self.convert(value=value) - - def to_ordinal(self, number, prefix=''): - if number <= 19: - return "{}".format(self.arabicOrdinal[number]) - if number < 100: - self.isCurrencyNameFeminine = True - else: - self.isCurrencyNameFeminine = False - self.currency_subunit = ('', '', '', '') - self.currency_unit = ('', '', '', '') - self.arabicPrefixText = prefix - self.arabicSuffixText = "" - return "{}".format(self.convert(abs(number)).strip()) - - def to_year(self, value): - value = self.validate_number(value) - return self.to_cardinal(value) - - def to_ordinal_num(self, value): - return self.to_ordinal(value).strip() - - def to_cardinal(self, number): - number = self.validate_number(number) - minus = '' - if number < 0: - minus = 'سالب ' - self.separator = ',' - self.currency_subunit = ('', '', '', '') - self.currency_unit = ('', '', '', '') - self.arabicPrefixText = "" - self.arabicSuffixText = "" - self.arabicOnes = ARABIC_ONES - return minus + self.convert(value=abs(number)).strip() diff --git a/fun_text_processing/num2words/num2words/lang_AZ.py b/fun_text_processing/num2words/num2words/lang_AZ.py deleted file mode 100644 index a16d4bff9..000000000 --- a/fun_text_processing/num2words/num2words/lang_AZ.py +++ /dev/null @@ -1,89 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import division, print_function, unicode_literals - -from . import lang_EU - - -class Num2Word_AZ(lang_EU.Num2Word_EU): - def set_high_numwords(self, high): - max = 3 + 3 * len(high) - for word, n in zip(high, range(max, 3, -3)): - self.cards[10 ** n] = word + " milyon" - - def setup(self): - super(Num2Word_AZ, self).setup() - - self.negword = "mənfi " - self.pointword = "nöqtə" - self.exclude_title = ["və", "nöqtə", "mənfi"] - - self.mid_numwords = [(1000, "min"), (100, "yüz"), - (90, "doxsan"), (80, "səksən"), (70, "yetmiş"), - (60, "altmış"), (50, "əlli"), (40, "qırx"), - (30, "otuz")] - self.low_numwords = ['iyirmi', 'on doqquz', 'on səkkiz', 'on yeddi', 'on altı', 'on beş', 'on dörd', 'on üç', 'on iki', 'on bir', 'on', 'doqquz', 'səkkiz', 'yeddi', 'altı', 'beş', 'dörd', 'üç', 'iki', 'bir', 'sıfır'] - self.ords = {'bir': 'birinci', - 'iki': 'ikinci', - 'üç': 'üçüncü', - 'dörd': 'dördüncü', - 'beş': 'beşinci', - 'altı': 'altıncı', - 'yeddi': 'yeddinci', - 'səkkiz': 'səkkizinci', - 'doqquz': 'doqquzuncu', - 'on': 'onuncu'} - - def merge(self, lpair, rpair): - ltext, lnum = lpair - rtext, rnum = rpair - if lnum == 1 and rnum < 100: - return (rtext, rnum) - elif 100 > lnum > rnum: - return ("%s %s" % (ltext, rtext), lnum + rnum) - elif lnum >= 100 > rnum: - return ("%s və %s" % (ltext, rtext), lnum + rnum) - elif rnum > lnum: - return ("%s %s" % (ltext, rtext), lnum * rnum) - return ("%s %s" % (ltext, rtext), lnum + rnum) - - def to_ordinal(self, value): - self.verify_ordinal(value) - outwords = self.to_cardinal(value).split(" ") - lastwords = outwords[-1].split("-") - lastword = lastwords[-1].lower() - try: - lastword = self.ords[lastword] - except KeyError: - if lastword[-1] == "y": - lastword = lastword[:-1] + "ie" - lastword += "th" - lastwords[-1] = self.title(lastword) - outwords[-1] = "-".join(lastwords) - return " ".join(outwords) - - def to_ordinal_num(self, value): - self.verify_ordinal(value) - return "%s%s" % (value, self.to_ordinal(value)[-2:]) - - def to_year(self, val, suffix=None, longval=True): - if val < 0: - val = abs(val) - suffix = 'əsrdən əvvəl ' if not suffix else suffix - high, low = (val // 100, val % 100) - # If year is 00XX, X00X, or beyond 9999, go cardinal. - if (high == 0 - or (high % 10 == 0 and low < 10) - or high >= 100): - valtext = self.to_cardinal(val) - else: - hightext = self.to_cardinal(high) - if low == 0: - lowtext = "yüz" - elif low < 10: - lowtext = "oh-%s" % self.to_cardinal(low) - else: - lowtext = self.to_cardinal(low) - valtext = "%s %s" % (hightext, lowtext) - return (valtext if not suffix - else "%s %s" % (valtext, suffix)) diff --git a/fun_text_processing/num2words/num2words/lang_BG.py b/fun_text_processing/num2words/num2words/lang_BG.py deleted file mode 100644 index 804502b3b..000000000 --- a/fun_text_processing/num2words/num2words/lang_BG.py +++ /dev/null @@ -1,91 +0,0 @@ -# -*- coding: utf-8 -*- - -# Bulgarian num2words - -from __future__ import division, print_function, unicode_literals - -from . import lang_EU - - -class Num2Word_BG(lang_EU.Num2Word_EU): - def set_high_numwords(self, high): - max = 3 + 3 * len(high) - for word, n in zip(high, range(max, 3, -3)): - self.cards[10 ** n] = word + " милиона" - - def setup(self): - super(Num2Word_BG, self).setup() - - self.negword = "минус " - self.pointword = "точка" - self.exclude_title = ["и", "точка", "минус"] - - self.mid_numwords = [(1000, "хиляда"), (100, "сто"), - (90, "деветдесет"), (80, "осемдесет"), (70, "седемдесет"), - (60, "шестдесет"), (50, "петдесет"), (40, "четиридесет"), - (30, "тридесет")] - self.low_numwords = ['двадесет', 'деветнайсет', 'осемнадесет', 'седемнадесет', 'шестнадесет', 'петнадесет', 'четиринадесет', 'тринадесет', 'дванадесет', 'единадесет', 'десет', 'девет', 'осем', 'седем', 'шест', 'пет', 'четири', 'три', 'две', 'един', 'нула'] - self.ords = {'един': 'първи', -'две': 'второ', -'три': 'трети', -'четири': 'четвърто', -'пет': 'пети', -'шест': 'шесто', -'седем': 'седмо', -'осем': 'осмо', -'девет': 'девети', -'десет': 'десети'} - - def merge(self, lpair, rpair): - ltext, lnum = lpair - rtext, rnum = rpair - if lnum == 1 and rnum < 100: - return (rtext, rnum) - elif 100 > lnum > rnum: - return ("%s-%s" % (ltext, rtext), lnum + rnum) - elif lnum >= 100 > rnum: - return ("%s и %s" % (ltext, rtext), lnum + rnum) - elif rnum > lnum: - return ("%s %s" % (ltext, rtext), lnum * rnum) - return ("%s, %s" % (ltext, rtext), lnum + rnum) - - # def to_ordinal(self, value): - # self.verify_ordinal(value) - # outwords = self.to_cardinal(value).split(" ") - # lastwords = outwords[-1].split("-") - # lastword = lastwords[-1].lower() - # try: - # lastword = self.ords[lastword] - # except KeyError: - # if lastword[-1] == "y": - # lastword = lastword[:-1] + "ie" - # lastword += "th" - # lastwords[-1] = self.title(lastword) - # outwords[-1] = "-".join(lastwords) - # return " ".join(outwords) - - # def to_ordinal_num(self, value): - # self.verify_ordinal(value) - # return "%s%s" % (value, self.to_ordinal(value)[-2:]) - - def to_year(self, val, suffix=None, longval=True): - if val < 0: - val = abs(val) - suffix = 'г. пр.н.е' if not suffix else suffix - high, low = (val // 100, val % 100) - # If year is 00XX, X00X, or beyond 9999, go cardinal. - if (high == 0 - or (high % 10 == 0 and low < 10) - or high >= 100): - valtext = self.to_cardinal(val) - else: - hightext = self.to_cardinal(high) - if low == 0: - lowtext = "сто" - elif low < 10: - lowtext = "oh-%s" % self.to_cardinal(low) - else: - lowtext = self.to_cardinal(low) - valtext = "%s %s" % (hightext, lowtext) - return (valtext if not suffix - else "%s %s" % (valtext, suffix)) diff --git a/fun_text_processing/num2words/num2words/lang_BN.py b/fun_text_processing/num2words/num2words/lang_BN.py deleted file mode 100644 index a308e73fd..000000000 --- a/fun_text_processing/num2words/num2words/lang_BN.py +++ /dev/null @@ -1,108 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import division, print_function, unicode_literals - -from . import lang_EU - -# Bengali - -class Num2Word_BN(lang_EU.Num2Word_EU): - def set_high_numwords(self, high): - max = 3 + 3 * len(high) - for word, n in zip(high, range(max, 3, -3)): - self.cards[10 ** n] = word - - def setup(self): - super(Num2Word_BN, self).setup() - - self.negword = "বিয়োগ " - self.pointword = "পয়েন্ট" - self.exclude_title = ["এবং", "পয়েন্ট", "বিয়োগ"] - - self.high_numwords = [(100000000000, "ট্রিলিয়ন"), - (1000000000, "এক ট্রিলিয়ন"), - (100000000, "বিলিয়ন"), - (10000000, "এক শত মিলিয়ন"), - (1000000, "কোটি"), - (100000, "দশ লক্ষ"), - (10000, "দশ হাজার")] - - self.mid_numwords = [(1000, "হাজার"), - (100, "শত"), - (90, "নব্বই"), (80, "আশি"), (70, "সত্তর"), - (60, "ষাট"), (50, "পঞ্চাশ"), (40, "চল্লিশ"), - (30, "ত্রিশ")] - self.low_numwords = ["বিশ","উনিশ","আঠার","সতের","ষোল","পনের","চৌদ্দ","তেরো","বারো","এগারো","দশ","নয়টি","আট","সাত","ছয়","পাঁচ","চার","তিন","দুই","এক","শূন্য"] - self.ords = {"এক": "প্রথম", - "দুই": "দ্বিতীয়", - "তিন": "তৃতীয়", - "চার": "চতুর্থ", - "পাঁচ": "পঞ্চম", - "ছয়": "ষষ্ঠ", - "সাত": "সপ্তম", - "আট": "অষ্টম", - "নয়টি": "নবম", - "দশ": "দশম", - "এগারো": "একাদশ", - "বারো": "দ্বাদশ", - "তেরো": "ত্রয়োদশ", - "চৌদ্দ": "চতুর্দশ", - "পনের": "পঞ্চদশ", - "ষোল": "ষোড়শ", - "সতের": "সপ্তদশতম", - "আঠার": "অষ্টাদশ", - "উনিশ": "উনিশতম", - "বিশ": "বিংশতম"} - - - def merge(self, lpair, rpair): - ltext, lnum = lpair - rtext, rnum = rpair - if lnum == 1 and rnum < 100: - return (rtext, rnum) - elif 100 > lnum > rnum: - return ("%s%s" % (ltext, rtext), lnum + rnum) - elif lnum >= 100 > rnum: - return ("%s %s" % (ltext, rtext), lnum + rnum) - elif rnum > lnum: - return ("%s %s" % (ltext, rtext), lnum * rnum) - return ("%s %s" % (ltext, rtext), lnum + rnum) - - def to_ordinal(self, value): - self.verify_ordinal(value) - outwords = self.to_cardinal(value).split(" ") - lastwords = outwords[-1].split("-") - lastword = lastwords[-1].lower() - try: - lastword = self.ords[lastword] - except KeyError: - lastword += "তম" - lastwords[-1] = self.title(lastword) - outwords[-1] = " ".join(lastwords) - return " ".join(outwords) - - def to_ordinal_num(self, value): - self.verify_ordinal(value) - return "%s%s" % (value, self.to_ordinal(value)[-2:]) - - def to_year(self, val, suffix=None, longval=True): - if val < 0: - val = abs(val) - suffix = 'খ্রিস্টপূর্ব' if not suffix else suffix - high, low = (val // 100, val % 100) - # If year is 00XX, X00X, or beyond 9999, go cardinal. - if (high == 0 - or (high % 10 == 0 and low < 10) - or high >= 100): - valtext = self.to_cardinal(val) - else: - hightext = self.to_cardinal(high) - if low == 0: - lowtext = "শত" - elif low < 10: - lowtext = "%s" % self.to_cardinal(low) - else: - lowtext = self.to_cardinal(low) - valtext = "%s %s" % (hightext, lowtext) - return (valtext if not suffix - else "%s %s" % (valtext, suffix)) diff --git a/fun_text_processing/num2words/num2words/lang_CA.py b/fun_text_processing/num2words/num2words/lang_CA.py deleted file mode 100644 index 1910e5599..000000000 --- a/fun_text_processing/num2words/num2words/lang_CA.py +++ /dev/null @@ -1,89 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import division, print_function, unicode_literals - -from . import lang_EU - - -class Num2Word_CA(lang_EU.Num2Word_EU): - def set_high_numwords(self, high): - max = 3 + 3 * len(high) - for word, n in zip(high, range(max, 3, -3)): - self.cards[10 ** n] = word + "illions" - - def setup(self): - super(Num2Word_CA, self).setup() - - self.negword = "menys " - self.pointword = "punt" - self.exclude_title = ["i", "punt", "menys"] - - self.mid_numwords = [(1000, "mil"), (100, "cent"), - (90, "noranta"), (80, "vuitanta"), (70, "setanta"), - (60, "seixanta"), (50, "cinquanta"), (40, "quaranta"), - (30, "trenta")] - self.low_numwords = ['vint', 'dinou', 'divuit anys', 'disset', 'setze', 'quinze', 'catorze', 'tretze', 'dotze', 'onze', 'deu', 'nou', 'vuit', 'set', 'sis', 'cinc', 'quatre', 'tres', 'dos', 'un', 'zero'] - self.ords = {'un': 'primer', - 'dos': 'segon', - 'tres': 'tercer', - 'quatre': 'quart', - 'cinc': 'cinquè', - 'sis': 'sisè', - 'set': 'setè', - 'vuit': 'vuitè', - 'nou': 'novè', - 'deu': 'desè'} - - def merge(self, lpair, rpair): - ltext, lnum = lpair - rtext, rnum = rpair - if lnum == 1 and rnum < 100: - return (rtext, rnum) - elif 100 > lnum > rnum: - return ("%s-%s" % (ltext, rtext), lnum + rnum) - elif lnum >= 100 > rnum: - return ("%s i %s" % (ltext, rtext), lnum + rnum) - elif rnum > lnum: - return ("%s %s" % (ltext, rtext), lnum * rnum) - return ("%s, %s" % (ltext, rtext), lnum + rnum) - - def to_ordinal(self, value): - self.verify_ordinal(value) - outwords = self.to_cardinal(value).split(" ") - lastwords = outwords[-1].split("-") - lastword = lastwords[-1].lower() - try: - lastword = self.ords[lastword] - except KeyError: - if lastword[-1] == "y": - lastword = lastword[:-1] + "ie" - lastword += "th" - lastwords[-1] = self.title(lastword) - outwords[-1] = "-".join(lastwords) - return " ".join(outwords) - - def to_ordinal_num(self, value): - self.verify_ordinal(value) - return "%s%s" % (value, self.to_ordinal(value)[-2:]) - - def to_year(self, val, suffix=None, longval=True): - if val < 0: - val = abs(val) - suffix = 'BC' if not suffix else suffix - high, low = (val // 100, val % 100) - # If year is 00XX, X00X, or beyond 9999, go cardinal. - if (high == 0 - or (high % 10 == 0 and low < 10) - or high >= 100): - valtext = self.to_cardinal(val) - else: - hightext = self.to_cardinal(high) - if low == 0: - lowtext = "cent" - elif low < 10: - lowtext = "oh-%s" % self.to_cardinal(low) - else: - lowtext = self.to_cardinal(low) - valtext = "%s %s" % (hightext, lowtext) - return (valtext if not suffix - else "%s %s" % (valtext, suffix)) diff --git a/fun_text_processing/num2words/num2words/lang_CZ.py b/fun_text_processing/num2words/num2words/lang_CZ.py deleted file mode 100644 index 634f36138..000000000 --- a/fun_text_processing/num2words/num2words/lang_CZ.py +++ /dev/null @@ -1,142 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import unicode_literals - -from .base import Num2Word_Base -from .utils import get_digits, splitbyx - -ZERO = ('nula',) - -ONES = { - 1: ('jedna',), - 2: ('dva',), - 3: ('tři',), - 4: ('čtyři',), - 5: ('pět',), - 6: ('šest',), - 7: ('sedm',), - 8: ('osm',), - 9: ('devět',), -} - -TENS = { - 0: ('deset',), - 1: ('jedenáct',), - 2: ('dvanáct',), - 3: ('třináct',), - 4: ('čtrnáct',), - 5: ('patnáct',), - 6: ('šestnáct',), - 7: ('sedmnáct',), - 8: ('osmnáct',), - 9: ('devatenáct',), -} - -TWENTIES = { - 2: ('dvacet',), - 3: ('třicet',), - 4: ('čtyřicet',), - 5: ('padesát',), - 6: ('šedesát',), - 7: ('sedmdesát',), - 8: ('osmdesát',), - 9: ('devadesát',), -} - -HUNDREDS = { - 1: ('sto',), - 2: ('dvěstě',), - 3: ('třista',), - 4: ('čtyřista',), - 5: ('pětset',), - 6: ('šestset',), - 7: ('sedmset',), - 8: ('osmset',), - 9: ('devětset',), -} - -THOUSANDS = { - 1: ('tisíc', 'tisíce', 'tisíc'), # 10^3 - 2: ('milion', 'miliony', 'milionů'), # 10^6 - 3: ('miliarda', 'miliardy', 'miliard'), # 10^9 - 4: ('bilion', 'biliony', 'bilionů'), # 10^12 - 5: ('biliarda', 'biliardy', 'biliard'), # 10^15 - 6: ('trilion', 'triliony', 'trilionů'), # 10^18 - 7: ('triliarda', 'triliardy', 'triliard'), # 10^21 - 8: ('kvadrilion', 'kvadriliony', 'kvadrilionů'), # 10^24 - 9: ('kvadriliarda', 'kvadriliardy', 'kvadriliard'), # 10^27 - 10: ('quintillion', 'quintilliony', 'quintillionů'), # 10^30 -} - - -class Num2Word_CZ(Num2Word_Base): - CURRENCY_FORMS = { - 'CZK': ( - ('koruna', 'koruny', 'korun'), ('halíř', 'halíře', 'haléřů') - ), - 'EUR': ( - ('euro', 'euro', 'euro'), ('cent', 'centy', 'centů') - ), - } - - def setup(self): - self.negword = "mínus" - self.pointword = "celá" - - def to_cardinal(self, number): - n = str(number).replace(',', '.') - if '.' in n: - left, right = n.split('.') - leading_zero_count = len(right) - len(right.lstrip('0')) - decimal_part = ((ZERO[0] + ' ') * leading_zero_count + - self._int2word(int(right))) - return u'%s %s %s' % ( - self._int2word(int(left)), - self.pointword, - decimal_part - ) - else: - return self._int2word(int(n)) - - def pluralize(self, n, forms): - if n == 1: - form = 0 - elif 5 > n % 10 > 1 and (n % 100 < 10 or n % 100 > 20): - form = 1 - else: - form = 2 - return forms[form] - - def to_ordinal(self, number): - raise NotImplementedError() - - def _int2word(self, n): - if n == 0: - return ZERO[0] - - words = [] - chunks = list(splitbyx(str(n), 3)) - i = len(chunks) - for x in chunks: - i -= 1 - - if x == 0: - continue - - n1, n2, n3 = get_digits(x) - - if n3 > 0: - words.append(HUNDREDS[n3][0]) - - if n2 > 1: - words.append(TWENTIES[n2][0]) - - if n2 == 1: - words.append(TENS[n1][0]) - elif n1 > 0 and not (i > 0 and x == 1): - words.append(ONES[n1][0]) - - if i > 0: - words.append(self.pluralize(x, THOUSANDS[i])) - - return ' '.join(words) diff --git a/fun_text_processing/num2words/num2words/lang_DE.py b/fun_text_processing/num2words/num2words/lang_DE.py deleted file mode 100644 index 5e905ee3f..000000000 --- a/fun_text_processing/num2words/num2words/lang_DE.py +++ /dev/null @@ -1,142 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import print_function, unicode_literals - -import re - -from .lang_EU import Num2Word_EU - - -class Num2Word_DE(Num2Word_EU): - CURRENCY_FORMS = { - 'EUR': (('Euro', 'Euro'), ('Cent', 'Cent')), - 'GBP': (('Pfund', 'Pfund'), ('Penny', 'Pence')), - 'USD': (('Dollar', 'Dollar'), ('Cent', 'Cent')), - 'CNY': (('Yuan', 'Yuan'), ('Jiao', 'Fen')), - 'DEM': (('Mark', 'Mark'), ('Pfennig', 'Pfennig')), - } - - GIGA_SUFFIX = "illiarde" - MEGA_SUFFIX = "illion" - - def setup(self): - self.negword = "minus " - self.pointword = "Komma" - # "Cannot treat float %s as ordinal." - self.errmsg_floatord = ( - "Die Gleitkommazahl %s kann nicht in eine Ordnungszahl " + - "konvertiert werden." - ) - # "type(((type(%s)) ) not in [long, int, float]" - self.errmsg_nonnum = ( - "Nur Zahlen (type(%s)) können in Wörter konvertiert werden." - ) - # "Cannot treat negative num %s as ordinal." - self.errmsg_negord = ( - "Die negative Zahl %s kann nicht in eine Ordnungszahl " + - "konvertiert werden." - ) - # "abs(%s) must be less than %s." - self.errmsg_toobig = "Die Zahl %s muss kleiner als %s sein." - self.exclude_title = [] - - lows = ["Non", "Okt", "Sept", "Sext", "Quint", "Quadr", "Tr", "B", "M"] - units = ["", "un", "duo", "tre", "quattuor", "quin", "sex", "sept", - "okto", "novem"] - tens = ["dez", "vigint", "trigint", "quadragint", "quinquagint", - "sexagint", "septuagint", "oktogint", "nonagint"] - self.high_numwords = ( - ["zent"] + self.gen_high_numwords(units, tens, lows) - ) - self.mid_numwords = [(1000, "tausend"), (100, "hundert"), - (90, "neunzig"), (80, "achtzig"), (70, "siebzig"), - (60, "sechzig"), (50, "f\xFCnfzig"), - (40, "vierzig"), (30, "drei\xDFig")] - self.low_numwords = ["zwanzig", "neunzehn", "achtzehn", "siebzehn", - "sechzehn", "f\xFCnfzehn", "vierzehn", "dreizehn", - "zw\xF6lf", "elf", "zehn", "neun", "acht", - "sieben", "sechs", "f\xFCnf", "vier", "drei", - "zwei", "eins", "null"] - self.ords = {"eins": "ers", - "drei": "drit", - "acht": "ach", - "sieben": "sieb", - "ig": "igs", - "ert": "erts", - "end": "ends", - "ion": "ions", - "nen": "ns", - "rde": "rds", - "rden": "rds"} - - def merge(self, curr, next): - ctext, cnum, ntext, nnum = curr + next - - if cnum == 1: - if nnum == 100 or nnum == 1000: - return ("ein" + ntext, nnum) - elif nnum < 10 ** 6: - return next - ctext = "eine" - - if nnum > cnum: - if nnum >= 10 ** 6: - if cnum > 1: - if ntext.endswith("e"): - ntext += "n" - else: - ntext += "en" - ctext += " " - val = cnum * nnum - else: - if nnum < 10 < cnum < 100: - if nnum == 1: - ntext = "ein" - ntext, ctext = ctext, ntext + "und" - elif cnum >= 10 ** 6: - ctext += " " - val = cnum + nnum - - word = ctext + ntext - return (word, val) - - def to_ordinal(self, value): - self.verify_ordinal(value) - outword = self.to_cardinal(value).lower() - for key in self.ords: - if outword.endswith(key): - outword = outword[:len(outword) - len(key)] + self.ords[key] - break - - res = outword + "te" - - # Exception: "hundertste" is usually preferred over "einhundertste" - if res == "eintausendste" or res == "einhundertste": - res = res.replace("ein", "", 1) - # ... similarly for "millionste" etc. - res = re.sub(r'eine ([a-z]+(illion|illiard)ste)$', - lambda m: m.group(1), res) - # Ordinals involving "Million" etc. are written without a space. - # see https://de.wikipedia.org/wiki/Million#Sprachliches - res = re.sub(r' ([a-z]+(illion|illiard)ste)$', - lambda m: m.group(1), res) - - return res - - def to_ordinal_num(self, value): - self.verify_ordinal(value) - return str(value) + "." - - def to_currency(self, val, currency='EUR', cents=True, separator=' und', - adjective=False): - result = super(Num2Word_DE, self).to_currency( - val, currency=currency, cents=cents, separator=separator, - adjective=adjective) - # Handle exception, in german is "ein Euro" and not "eins Euro" - return result.replace("eins ", "ein ") - - def to_year(self, val, longval=True): - if not (val // 100) % 10: - return self.to_cardinal(val) - return self.to_splitnum(val, hightxt="hundert", longval=longval)\ - .replace(' ', '') diff --git a/fun_text_processing/num2words/num2words/lang_DK.py b/fun_text_processing/num2words/num2words/lang_DK.py deleted file mode 100644 index 6b527cf9d..000000000 --- a/fun_text_processing/num2words/num2words/lang_DK.py +++ /dev/null @@ -1,117 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import division, print_function, unicode_literals - -from . import lang_EU - - -class Num2Word_DK(lang_EU.Num2Word_EU): - GIGA_SUFFIX = "illarder" - MEGA_SUFFIX = "illioner" - - def setup(self): - super(Num2Word_DK, self).setup() - - self.negword = "minus " - self.pointword = "komma" - self.exclude_title = ["og", "komma", "minus"] - - self.mid_numwords = [(1000, "tusind"), (100, "hundrede"), - (90, "halvfems"), (80, "firs"), - (70, "halvfjerds"), (60, "treds"), - (50, "halvtreds"), (40, "fyrre"), (30, "tredive")] - self.low_numwords = ["tyve", "nitten", "atten", "sytten", - "seksten", "femten", "fjorten", "tretten", - "tolv", "elleve", "ti", "ni", "otte", - "syv", "seks", "fem", "fire", "tre", "to", - "et", "nul"] - self.ords = {"nul": "nul", - "et": "f\xf8rste", - "to": "anden", - "tre": "tredje", - "fire": "fjerde", - "fem": "femte", - "seks": "sjette", - "syv": "syvende", - "otte": "ottende", - "ni": "niende", - "ti": "tiende", - "elleve": "ellevte", - "tolv": "tolvte", - "tretten": "trett", - "fjorten": "fjort", - "femten": "femt", - "seksten": "sekst", - "sytten": "sytt", - "atten": "att", - "nitten": "nitt", - "tyve": "tyv"} - - def merge(self, curr, next): - ctext, cnum, ntext, nnum = curr + next - if next[1] == 100 or next[1] == 1000: - lst = list(next) - lst[0] = 'et' + lst[0] - next = tuple(lst) - - if cnum == 1: - if nnum < 10 ** 6 or self.ordflag: - return next - ctext = "en" - if nnum > cnum: - if nnum >= 10 ** 6: - ctext += " " - val = cnum * nnum - else: - if cnum >= 100 and cnum < 1000: - ctext += " og " - elif cnum >= 1000 and cnum <= 100000: - ctext += "e og " - if nnum < 10 < cnum < 100: - if nnum == 1: - ntext = "en" - ntext, ctext = ctext, ntext + "og" - elif cnum >= 10 ** 6: - ctext += " " - val = cnum + nnum - word = ctext + ntext - return (word, val) - - def to_ordinal(self, value): - self.verify_ordinal(value) - self.ordflag = True - outword = self.to_cardinal(value) - self.ordflag = False - for key in self.ords: - if outword.endswith(key): - outword = outword[:len(outword) - len(key)] + self.ords[key] - break - if value % 100 >= 30 and value % 100 <= 39 or value % 100 == 0: - outword += "te" - elif value % 100 > 12 or value % 100 == 0: - outword += "ende" - return outword - - def to_ordinal_num(self, value): - self.verify_ordinal(value) - vaerdte = (0, 1, 5, 6, 11, 12) - if value % 100 >= 30 and value % 100 <= 39 or value % 100 in vaerdte: - return str(value) + "te" - elif value % 100 == 2: - return str(value) + "en" - return str(value) + "ende" - - def to_currency(self, val, longval=True): - if val // 100 == 1 or val == 1: - ret = self.to_splitnum(val, hightxt="kr", lowtxt="\xf8re", - jointxt="og", longval=longval) - return "en " + ret[3:] - return self.to_splitnum(val, hightxt="kr", lowtxt="\xf8re", - jointxt="og", longval=longval) - - def to_year(self, val, longval=True): - if val == 1: - return 'en' - if not (val // 100) % 10: - return self.to_cardinal(val) - return self.to_splitnum(val, hightxt="hundrede", longval=longval) diff --git a/fun_text_processing/num2words/num2words/lang_EL.py b/fun_text_processing/num2words/num2words/lang_EL.py deleted file mode 100644 index c08cb1313..000000000 --- a/fun_text_processing/num2words/num2words/lang_EL.py +++ /dev/null @@ -1,115 +0,0 @@ -# -*- coding: utf-8 -*- - - -from __future__ import division, print_function, unicode_literals - -from . import lang_EU - - -class Num2Word_EL(lang_EU.Num2Word_EU): - def set_high_numwords(self, high): - max = 3 + 3 * len(high) - for word, n in zip(high, range(max, 3, -3)): - self.cards[10 ** n] = word # Todo: check "-illion" - - def setup(self): - super(Num2Word_EL, self).setup() - # lang_EU.Num2Word_EU.setup(self) - - self.negword = "μείον " - self.pointword = "σημείο" - self.exclude_title = ["και", "σημείο", "μείον"] - - self.high_numwords = [(1000000000, "ένα δισεκατομμύριο"), - (100000000, "εκατό εκατομμύρια"), - (10000000, "δέκα εκατομμύρια"), - (1000000, "ένα εκατομμύριο"), - (100000, "εκατό χιλιάδες"), - (10000, "δέκα χιλιάδες")] - - self.mid_numwords = [(1000, "χίλια"), (100, "εκατό"), - (90, "ενενήντα"), (80, "ογδόντα"), (70, "εβδομήντα"), - (60, "εξήντα"), (50, "πενήντα"), (40, "σαράντα"), - (30, "τριάντα")] - - self.low_numwords = ["είκοσι", - "δεκαεννέα", - "δεκαοχτώ", - "δεκαεπτά", - "δεκαέξι", - "δεκαπέντε", - "δεκατέσσερα", - "δεκατρία", - "δώδεκα", - "έντεκα", - "δέκα", - "εννιά", - "οχτώ", - "επτά", - "έξι", - "πέντε", - "τέσσερα", - "τρία", - "δύο", - "ένα", - "μηδέν"] - - self.ords = {"ένα": "πρώτος", - "δύο": "δεύτερος", - "τρία": "τρίτος", - "τέσσερα": "τέταρτος", - "πέντε": "πέμπτος", - "έξι": "έκτος", - "επτά": "έβδομος", - "οχτώ": "όγδοος", - "εννιά": "ένατος", - "δέκα": "δέκατος", - "έντεκα": "ενδέκατος", - "δώδεκα": "δωδέκατος"} - - def merge(self, lpair, rpair): - ltext, lnum = lpair - rtext, rnum = rpair - if lnum == 1 and rnum < 100: - return (rtext, rnum) - elif 100 > lnum > rnum: - return ("%s%s" % (ltext, rtext), lnum + rnum) - elif lnum >= 100 > rnum: - return ("%s%s" % (ltext, rtext), lnum + rnum) - elif rnum > lnum: - return ("%s%s" % (ltext, rtext), lnum * rnum) - return ("%s%s" % (ltext, rtext), lnum + rnum) - - def to_ordinal_num(self, value): - self.verify_ordinal(value) - return "%s%s" % (value, self.to_ordinal(value)) - - def to_ordinal(self, value): - self.verify_ordinal(value) - try: - ordinal_word = self.ords[value] - except KeyError: - ordinal_word = value + "τος" - return ordinal_word - - def to_year(self, val, suffix=None, longval=True): - if val < 0: - val = abs(val) - suffix = 'π.Χ' if not suffix else suffix - high, low = (val // 100, val % 100) - # If year is 00XX, X00X, or beyond 9999, go cardinal. - if (high == 0 - or (high % 10 == 0 and low < 10) - or high >= 100): - valtext = self.to_cardinal(val) - else: - hightext = self.to_cardinal(high) - if low == 0: - lowtext = "εκατό" - elif low < 10: - lowtext = "%s" % self.to_cardinal(low) - else: - lowtext = self.to_cardinal(low) - valtext = "%s %s" % (hightext, lowtext) - return (valtext if not suffix - else "%s %s" % (valtext, suffix)) diff --git a/fun_text_processing/num2words/num2words/lang_EN.py b/fun_text_processing/num2words/num2words/lang_EN.py deleted file mode 100644 index 408c8cf53..000000000 --- a/fun_text_processing/num2words/num2words/lang_EN.py +++ /dev/null @@ -1,95 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import division, print_function, unicode_literals - -from . import lang_EU - - -class Num2Word_EN(lang_EU.Num2Word_EU): - def set_high_numwords(self, high): - max = 3 + 3 * len(high) - for word, n in zip(high, range(max, 3, -3)): - self.cards[10 ** n] = word + "illion" - - def setup(self): - super(Num2Word_EN, self).setup() - - self.negword = "minus " - self.pointword = "point" - self.exclude_title = ["and", "point", "minus"] - - self.mid_numwords = [(1000, "thousand"), (100, "hundred"), - (90, "ninety"), (80, "eighty"), (70, "seventy"), - (60, "sixty"), (50, "fifty"), (40, "forty"), - (30, "thirty")] - self.low_numwords = ["twenty", "nineteen", "eighteen", "seventeen", - "sixteen", "fifteen", "fourteen", "thirteen", - "twelve", "eleven", "ten", "nine", "eight", - "seven", "six", "five", "four", "three", "two", - "one", "zero"] - self.ords = {"one": "first", - "two": "second", - "three": "third", - "four": "fourth", - "five": "fifth", - "six": "sixth", - "seven": "seventh", - "eight": "eighth", - "nine": "ninth", - "ten": "tenth", - "eleven": "eleventh", - "twelve": "twelfth"} - - def merge(self, lpair, rpair): - ltext, lnum = lpair - rtext, rnum = rpair - if lnum == 1 and rnum < 100: - return (rtext, rnum) - elif 100 > lnum > rnum: - return ("%s-%s" % (ltext, rtext), lnum + rnum) - elif lnum >= 100 > rnum: - return ("%s and %s" % (ltext, rtext), lnum + rnum) - elif rnum > lnum: - return ("%s %s" % (ltext, rtext), lnum * rnum) - return ("%s, %s" % (ltext, rtext), lnum + rnum) - - def to_ordinal(self, value): - self.verify_ordinal(value) - outwords = self.to_cardinal(value).split(" ") - lastwords = outwords[-1].split("-") - lastword = lastwords[-1].lower() - try: - lastword = self.ords[lastword] - except KeyError: - if lastword[-1] == "y": - lastword = lastword[:-1] + "ie" - lastword += "th" - lastwords[-1] = self.title(lastword) - outwords[-1] = "-".join(lastwords) - return " ".join(outwords) - - def to_ordinal_num(self, value): - self.verify_ordinal(value) - return "%s%s" % (value, self.to_ordinal(value)[-2:]) - - def to_year(self, val, suffix=None, longval=True): - if val < 0: - val = abs(val) - suffix = 'BC' if not suffix else suffix - high, low = (val // 100, val % 100) - # If year is 00XX, X00X, or beyond 9999, go cardinal. - if (high == 0 - or (high % 10 == 0 and low < 10) - or high >= 100): - valtext = self.to_cardinal(val) - else: - hightext = self.to_cardinal(high) - if low == 0: - lowtext = "hundred" - elif low < 10: - lowtext = "oh-%s" % self.to_cardinal(low) - else: - lowtext = self.to_cardinal(low) - valtext = "%s %s" % (hightext, lowtext) - return (valtext if not suffix - else "%s %s" % (valtext, suffix)) diff --git a/fun_text_processing/num2words/num2words/lang_EN_IN.py b/fun_text_processing/num2words/num2words/lang_EN_IN.py deleted file mode 100644 index 0736c27e8..000000000 --- a/fun_text_processing/num2words/num2words/lang_EN_IN.py +++ /dev/null @@ -1,11 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import unicode_literals - -from .lang_EN import Num2Word_EN - - -class Num2Word_EN_IN(Num2Word_EN): - def set_high_numwords(self, high): - self.cards[10 ** 7] = "crore" - self.cards[10 ** 5] = "lakh" diff --git a/fun_text_processing/num2words/num2words/lang_ES.py b/fun_text_processing/num2words/num2words/lang_ES.py deleted file mode 100644 index 0c2c985cb..000000000 --- a/fun_text_processing/num2words/num2words/lang_ES.py +++ /dev/null @@ -1,351 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import division, print_function, unicode_literals - -import math - -from .lang_EU import Num2Word_EU - -GENERIC_DOLLARS = ('dolar', 'dólares') -GENERIC_CENTS = ('centavo', 'centavos') -CURRENCIES_UNA = ('SLL', 'SEK', 'NOK', 'CZK', 'DKK', 'ISK', - 'SKK', 'GBP', 'CYP', 'EGP', 'FKP', 'GIP', - 'LBP', 'SDG', 'SHP', 'SSP', 'SYP', 'INR', - 'IDR', 'LKR', 'MUR', 'NPR', 'PKR', 'SCR', - 'ESP') - - -class Num2Word_ES(Num2Word_EU): - CURRENCY_FORMS = { - 'EUR': (('euro', 'euros'), ('céntimo', 'céntimos')), - 'ESP': (('peseta', 'pesetas'), ('céntimo', 'céntimos')), - 'USD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'PEN': (('sol', 'soles'), ('céntimo', 'céntimos')), - 'CRC': (('colón', 'colones'), GENERIC_CENTS), - 'AUD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'CAD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'GBP': (('libra', 'libras'), ('penny', 'pence')), - 'RUB': (('rublo', 'rublos'), ('kopeyka', 'kopeykas')), - 'SEK': (('corona', 'coronas'), ('öre', 'öre')), - 'NOK': (('corona', 'coronas'), ('øre', 'øre')), - 'PLN': (('zloty', 'zlotys'), ('grosz', 'groszy')), - 'MXN': (('peso', 'pesos'), GENERIC_CENTS), - 'RON': (('leu', 'leus'), ('ban', 'bani')), - 'INR': (('rupia', 'rupias'), ('paisa', 'paisas')), - 'HUF': (('florín', 'florines'), ('fillér', 'fillér')), - 'FRF': (('franco', 'francos'), ('céntimo', 'céntimos')), - 'CNY': (('yuan', 'yuanes'), ('fen', 'jiaos')), - 'CZK': (('corona', 'coronas'), ('haléř', 'haléř')), - 'NIO': (('córdoba', 'córdobas'), GENERIC_CENTS), - 'VES': (('bolívar', 'bolívares'), ('céntimo', 'céntimos')), - 'BRL': (('real', 'reales'), GENERIC_CENTS), - 'CHF': (('franco', 'francos'), ('céntimo', 'céntimos')), - 'JPY': (('yen', 'yenes'), ('sen', 'sen')), - 'KRW': (('won', 'wones'), ('jeon', 'jeon')), - 'KPW': (('won', 'wones'), ('chon', 'chon')), - 'TRY': (('lira', 'liras'), ('kuruş', 'kuruş')), - 'ZAR': (('rand', 'rands'), ('céntimo', 'céntimos')), - 'KZT': (('tenge', 'tenges'), ('tïın', 'tïın')), - 'UAH': (('hryvnia', 'hryvnias'), ('kopiyka', 'kopiykas')), - 'THB': (('baht', 'bahts'), ('satang', 'satang')), - 'AED': (('dirham', 'dirhams'), ('fils', 'fils')), - 'AFN': (('afghani', 'afghanis'), ('pul', 'puls')), - 'ALL': (('lek ', 'leke'), ('qindarkë', 'qindarka')), - 'AMD': (('dram', 'drams'), ('luma', 'lumas')), - 'ANG': (('florín', 'florines'), GENERIC_CENTS), - 'AOA': (('kwanza', 'kwanzas'), ('céntimo', 'céntimos')), - 'ARS': (('peso', 'pesos'), GENERIC_CENTS), - 'AWG': (('florín', 'florines'), GENERIC_CENTS), - 'AZN': (('manat', 'manat'), ('qəpik', 'qəpik')), - 'BBD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'BDT': (('taka', 'takas'), ('paisa', 'paisas')), - 'BGN': (('lev', 'leva'), ('stotinka', 'stotinki')), - 'BHD': (('dinar', 'dinares'), ('fils', 'fils')), - 'BIF': (('franco', 'francos'), ('céntimo', 'céntimos')), - 'BMD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'BND': (GENERIC_DOLLARS, GENERIC_CENTS), - 'BOB': (('boliviano', 'bolivianos'), GENERIC_CENTS), - 'BSD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'BTN': (('ngultrum', 'ngultrum'), ('chetrum', 'chetrum')), - 'BWP': (('pula', 'pulas'), ('thebe', 'thebes')), - 'BYN': (('rublo', 'rublos'), ('kópek', 'kópeks')), - 'BYR': (('rublo', 'rublos'), ('kópek', 'kópeks')), - 'BZD': (GENERIC_DOLLARS, ('céntimo', 'céntimos')), - 'CDF': (('franco', 'francos'), ('céntimo', 'céntimos')), - 'CLP': (('peso', 'pesos'), GENERIC_CENTS), - 'COP': (('peso', 'pesos'), GENERIC_CENTS), - 'CUP': (('peso', 'pesos'), GENERIC_CENTS), - 'CVE': (('escudo', 'escudos'), GENERIC_CENTS), - 'CYP': (('libra', 'libras'), ('céntimo', 'céntimos')), - 'DJF': (('franco', 'francos'), ('céntimo', 'céntimos')), - 'DKK': (('corona', 'coronas'), ('øre', 'øre')), - 'DOP': (('peso', 'pesos'), GENERIC_CENTS), - 'DZD': (('dinar', 'dinares'), ('céntimo', 'céntimos')), - 'ECS': (('sucre', 'sucres'), GENERIC_CENTS), - 'EGP': (('libra', 'libras'), ('piastra', 'piastras')), - 'ERN': (('nakfa', 'nakfas'), ('céntimo', 'céntimos')), - 'ETB': (('birr', 'birrs'), ('céntimo', 'céntimos')), - 'FJD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'FKP': (('libra', 'libras'), ('penny', 'peniques')), - 'GEL': (('lari', 'laris'), ('tetri', 'tetris')), - 'GHS': (('cedi', 'cedis'), ('pesewa', 'pesewas')), - 'GIP': (('libra', 'libras'), ('penique', 'peniques')), - 'GMD': (('dalasi', 'dalasis'), ('butut', 'bututs')), - 'GNF': (('franco', 'francos'), ('céntimo', 'céntimos')), - 'GTQ': (('quetzal', 'quetzales'), GENERIC_CENTS), - 'GYD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'HKD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'HNL': (('lempira', 'lempiras'), GENERIC_CENTS), - 'HRK': (('kuna', 'kunas'), ('lipa', 'lipas')), - 'HTG': (('gourde', 'gourdes'), ('céntimo', 'céntimos')), - 'IDR': (('rupia', 'rupias'), ('céntimo', 'céntimos')), - 'ILS': (('séquel', 'séqueles'), ('agora', 'agoras')), - 'IQD': (('dinar', 'dinares'), ('fils', 'fils')), - 'IRR': (('rial', 'riales'), ('dinar', 'dinares')), - 'ISK': (('corona', 'coronas'), ('eyrir', 'aurar')), - 'ITL': (('lira', 'liras'), ('céntimo', 'céntimos')), - 'JMD': (GENERIC_DOLLARS, ('céntimo', 'céntimos')), - 'JOD': (('dinar', 'dinares'), ('piastra', 'piastras')), - 'KES': (('chelín', 'chelines'), ('céntimo', 'céntimos')), - 'KGS': (('som', 'som'), ('tyiyn', 'tyiyn')), - 'KHR': (('riel', 'rieles'), ('céntimo', 'céntimos')), - 'KMF': (('franco', 'francos'), ('céntimo', 'céntimos')), - 'KWD': (('dinar', 'dinares'), ('fils', 'fils')), - 'KYD': (GENERIC_DOLLARS, ('céntimo', 'céntimos')), - 'LAK': (('kip', 'kips'), ('att', 'att')), - 'LBP': (('libra', 'libras'), ('piastra', 'piastras')), - 'LKR': (('rupia', 'rupias'), ('céntimo', 'céntimos')), - 'LRD': (GENERIC_DOLLARS, ('céntimo', 'céntimos')), - 'LSL': (('loti', 'lotis'), ('céntimo', 'céntimos')), - 'LTL': (('lita', 'litas'), ('céntimo', 'céntimos')), - 'LVL': (('lat', 'lats'), ('céntimo', 'céntimos')), - 'LYD': (('dinar', 'dinares'), ('dírham', 'dírhams')), - 'MAD': (('dírham', 'dirhams'), ('céntimo', 'céntimos')), - 'MDL': (('leu', 'lei'), ('ban', 'bani')), - 'MGA': (('ariary', 'ariaris'), ('iraimbilanja', 'iraimbilanja')), - 'MKD': (('denar', 'denares'), ('deni', 'denis')), - 'MMK': (('kiat', 'kiats'), ('pya', 'pyas')), - 'MNT': (('tugrik', 'tugriks'), ('möngö', 'möngö')), - 'MOP': (('pataca', 'patacas'), ('avo', 'avos')), - 'MRO': (('ouguiya', 'ouguiyas'), ('khoums', 'khoums')), - 'MRU': (('ouguiya', 'ouguiyas'), ('khoums', 'khoums')), - 'MUR': (('rupia', 'rupias'), ('céntimo', 'céntimos')), - 'MVR': (('rufiyaa', 'rufiyaas'), ('laari', 'laari')), - 'MWK': (('kuacha', 'kuachas'), ('tambala', 'tambalas')), - 'MYR': (('ringgit', 'ringgit'), ('céntimo', 'céntimos')), - 'MZN': (('metical', 'metical'), GENERIC_CENTS), - 'NAD': (GENERIC_DOLLARS, ('céntimo', 'céntimos')), - 'NGN': (('naira', 'nairas'), ('kobo', 'kobo')), - 'NPR': (('rupia', 'rupias'), ('paisa', 'paisas')), - 'NZD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'OMR': (('rial', 'riales'), ('baisa', 'baisa')), - 'PAB': (('balboa', 'balboas'), ('centésimo', 'centésimos')), - 'PGK': (('kina', 'kinas'), ('toea', 'toea')), - 'PHP': (('peso', 'pesos'), GENERIC_CENTS), - 'PKR': (('rupia', 'rupias'), ('paisa', 'paisas')), - 'PLZ': (('zloty', 'zlotys'), ('grosz', 'groszy')), - 'PYG': (('guaraní', 'guaranís'), ('céntimo', 'céntimos')), - 'QAR': (('rial', 'riales'), ('dírham', 'dírhams')), - 'QTQ': (('quetzal', 'quetzales'), GENERIC_CENTS), - 'RSD': (('dinar', 'dinares'), ('para', 'para')), - 'RUR': (('rublo', 'rublos'), ('kopek', 'kopeks')), - 'RWF': (('franco', 'francos'), ('céntimo', 'céntimos')), - 'SAR': (('riyal', 'riales'), ('halala', 'halalas')), - 'SBD': (GENERIC_DOLLARS, ('céntimo', 'céntimos')), - 'SCR': (('rupia', 'rupias'), ('céntimo', 'céntimos')), - 'SDG': (('libra', 'libras'), ('piastra', 'piastras')), - 'SGD': (GENERIC_DOLLARS, ('céntimo', 'céntimos')), - 'SHP': (('libra', 'libras'), ('penny', 'peniques')), - 'SKK': (('corona', 'coronas'), ('halier', 'haliers')), - 'SLL': (('leona', 'leonas'), ('céntimo', 'céntimos')), - 'SRD': (GENERIC_DOLLARS, ('céntimo', 'céntimos')), - 'SSP': (('libra', 'libras'), ('piastra', 'piastras')), - 'STD': (('dobra', 'dobras'), ('céntimo', 'céntimos')), - 'SVC': (('colón', 'colones'), GENERIC_CENTS), - 'SYP': (('libra', 'libras'), ('piastra', 'piastras')), - 'SZL': (('lilangeni', 'emalangeni'), ('céntimo', 'céntimos')), - 'TJS': (('somoni', 'somonis'), ('dirame', 'dirames')), - 'TMT': (('manat', 'manat'), ('tenge', 'tenge')), - 'TND': (('dinar', 'dinares'), ('milésimo', 'milésimos')), - 'TOP': (('paanga', 'paangas'), ('céntimo', 'céntimos')), - 'TTD': (GENERIC_DOLLARS, ('céntimo', 'céntimos')), - 'TWD': (('nuevo dólar', 'nuevos dolares'), ('céntimo', 'céntimos')), - 'TZS': (('chelín', 'chelines'), ('céntimo', 'céntimos')), - 'UAG': (('hryvnia', 'hryvnias'), ('kopiyka', 'kopiykas')), - 'UGX': (('chelín', 'chelines'), ('céntimo', 'céntimos')), - 'UYU': (('peso', 'pesos'), ('centésimo', 'centésimos')), - 'UZS': (('sum', 'sum'), ('tiyin', 'tiyin')), - 'VEF': (('bolívar fuerte', 'bolívares fuertes'), - ('céntimo', 'céntimos')), - 'VND': (('dong', 'dongs'), ('xu', 'xu')), - 'VUV': (('vatu', 'vatu'), ('nenhum', 'nenhum')), - 'WST': (('tala', 'tala'), GENERIC_CENTS), - 'XAF': (('franco CFA', 'francos CFA'), ('céntimo', 'céntimos')), - 'XCD': (GENERIC_DOLLARS, ('céntimo', 'céntimos')), - 'XOF': (('franco CFA', 'francos CFA'), ('céntimo', 'céntimos')), - 'XPF': (('franco CFP', 'francos CFP'), ('céntimo', 'céntimos')), - 'YER': (('rial', 'riales'), ('fils', 'fils')), - 'YUM': (('dinar', 'dinares'), ('para', 'para')), - 'ZMW': (('kwacha', 'kwachas'), ('ngwee', 'ngwee')), - 'ZRZ': (('zaire', 'zaires'), ('likuta', 'makuta')), - 'ZWL': (GENERIC_DOLLARS, ('céntimo', 'céntimos')), - 'ZWL': (GENERIC_DOLLARS, ('céntimo', 'céntimos')), - } - - # //CHECK: Is this sufficient?? - GIGA_SUFFIX = None - MEGA_SUFFIX = "illón" - - def setup(self): - lows = ["cuatr", "tr", "b", "m"] - self.high_numwords = self.gen_high_numwords([], [], lows) - self.negword = "menos " - self.pointword = "punto" - self.errmsg_nonnum = "Solo números pueden ser convertidos a palabras." - self.errmsg_toobig = ( - "Numero muy grande para ser convertido a palabras." - ) - self.gender_stem = "o" - self.exclude_title = ["y", "menos", "punto"] - self.mid_numwords = [(1000, "mil"), (100, "cien"), (90, "noventa"), - (80, "ochenta"), (70, "setenta"), (60, "sesenta"), - (50, "cincuenta"), (40, "cuarenta"), - (30, "treinta")] - self.low_numwords = ["veintinueve", "veintiocho", "veintisiete", - "veintiséis", "veinticinco", "veinticuatro", - "veintitrés", "veintidós", "veintiuno", - "veinte", "diecinueve", "dieciocho", "diecisiete", - "dieciseis", "quince", "catorce", "trece", "doce", - "once", "diez", "nueve", "ocho", "siete", "seis", - "cinco", "cuatro", "tres", "dos", "uno", "cero"] - self.ords = {1: "primer", - 2: "segund", - 3: "tercer", - 4: "cuart", - 5: "quint", - 6: "sext", - 7: "séptim", - 8: "octav", - 9: "noven", - 10: "décim", - 20: "vigésim", - 30: "trigésim", - 40: "quadragésim", - 50: "quincuagésim", - 60: "sexagésim", - 70: "septuagésim", - 80: "octogésim", - 90: "nonagésim", - 100: "centésim", - 200: "ducentésim", - 300: "tricentésim", - 400: "cuadrigentésim", - 500: "quingentésim", - 600: "sexcentésim", - 700: "septigentésim", - 800: "octigentésim", - 900: "noningentésim", - 1e3: "milésim", - 1e6: "millonésim", - 1e9: "billonésim", - 1e12: "trillonésim", - 1e15: "cuadrillonésim"} - - def merge(self, curr, next): - ctext, cnum, ntext, nnum = curr + next - - if cnum == 1: - if nnum < 1000000: - return next - ctext = "un" - elif cnum == 100 and not nnum % 1000 == 0: - ctext += "t" + self.gender_stem - - if nnum < cnum: - if cnum < 100: - return "%s y %s" % (ctext, ntext), cnum + nnum - return "%s %s" % (ctext, ntext), cnum + nnum - elif (not nnum % 1000000) and cnum > 1: - ntext = ntext[:-3] + "lones" - - if nnum == 100: - if cnum == 5: - ctext = "quinien" - ntext = "" - elif cnum == 7: - ctext = "sete" - elif cnum == 9: - ctext = "nove" - ntext += "t" + self.gender_stem + "s" - else: - ntext = " " + ntext - - return (ctext + ntext, cnum * nnum) - - def to_ordinal(self, value): - self.verify_ordinal(value) - if value == 0: - text = "" - elif value <= 10: - text = "%s%s" % (self.ords[value], self.gender_stem) - elif value <= 12: - text = ( - "%s%s%s" % (self.ords[10], self.gender_stem, - self.to_ordinal(value - 10)) - ) - elif value <= 100: - dec = (value // 10) * 10 - text = ( - "%s%s %s" % (self.ords[dec], self.gender_stem, - self.to_ordinal(value - dec)) - ) - elif value <= 1e3: - cen = (value // 100) * 100 - text = ( - "%s%s %s" % (self.ords[cen], self.gender_stem, - self.to_ordinal(value - cen)) - ) - elif value < 1e18: - # Round down to the nearest 1e(3n) - # dec contains the following: - # [ 1e3, 1e6): 1e3 - # [ 1e6, 1e9): 1e6 - # [ 1e9, 1e12): 1e9 - # [1e12, 1e15): 1e12 - # [1e15, 1e18): 1e15 - dec = 1000 ** int(math.log(int(value), 1000)) - - # Split the parts before and after the word for 'dec' - # eg (12, 345) = divmod(12_345, 1_000) - high_part, low_part = divmod(value, dec) - - cardinal = self.to_cardinal(high_part) if high_part != 1 else "" - text = ( - "%s%s%s %s" % (cardinal, self.ords[dec], self.gender_stem, - self.to_ordinal(low_part)) - ) - else: - text = self.to_cardinal(value) - return text.strip() - - def to_ordinal_num(self, value): - self.verify_ordinal(value) - return "%s%s" % (value, "º" if self.gender_stem == 'o' else "ª") - - def to_currency(self, val, currency='EUR', cents=True, separator=' con', - adjective=False): - result = super(Num2Word_ES, self).to_currency( - val, currency=currency, cents=cents, separator=separator, - adjective=adjective) - # Handle exception, in spanish is "un euro" and not "uno euro" - # except in this currencies: leona, corona, - # libra, rupia, lempira, peseta, is 'una' - # but only when it's first word, otherwise - # it's replaced in others words like 'veintiun' - if currency in CURRENCIES_UNA: - list_result = result.split(" ") - if list_result[0] == "uno": - list_result[0] = list_result[0].replace("uno", "una") - result = " ".join(list_result) - result = result.replace("uno", "un") - return result diff --git a/fun_text_processing/num2words/num2words/lang_ES_CO.py b/fun_text_processing/num2words/num2words/lang_ES_CO.py deleted file mode 100644 index f27b17b87..000000000 --- a/fun_text_processing/num2words/num2words/lang_ES_CO.py +++ /dev/null @@ -1,14 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import print_function, unicode_literals - -from .lang_ES import Num2Word_ES - - -class Num2Word_ES_CO(Num2Word_ES): - - def to_currency(self, val, longval=True, old=False): - result = self.to_splitnum(val, hightxt="peso/s", lowtxt="centavo/s", - divisor=1, jointxt="y", longval=longval) - # Handle exception, in spanish is "un euro" and not "uno euro" - return result.replace("uno", "un") diff --git a/fun_text_processing/num2words/num2words/lang_ES_NI.py b/fun_text_processing/num2words/num2words/lang_ES_NI.py deleted file mode 100644 index 5537eab72..000000000 --- a/fun_text_processing/num2words/num2words/lang_ES_NI.py +++ /dev/null @@ -1,18 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import print_function, unicode_literals - -from .lang_ES import Num2Word_ES - - -class Num2Word_ES_NI(Num2Word_ES): - CURRENCY_FORMS = { - 'NIO': (('córdoba', 'córdobas'), ('centavo', 'centavos')), - } - - def to_currency(self, val, currency='NIO', cents=True, separator=' con', - adjective=False): - result = super(Num2Word_ES, self).to_currency( - val, currency=currency, cents=cents, separator=separator, - adjective=adjective) - return result.replace("uno", "un") diff --git a/fun_text_processing/num2words/num2words/lang_ES_VE.py b/fun_text_processing/num2words/num2words/lang_ES_VE.py deleted file mode 100644 index d73acc206..000000000 --- a/fun_text_processing/num2words/num2words/lang_ES_VE.py +++ /dev/null @@ -1,17 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import print_function, unicode_literals - -from .lang_ES import Num2Word_ES - - -class Num2Word_ES_VE(Num2Word_ES): - - def to_currency(self, val, longval=True, old=False): - hightxt = "bolívar/es" if old else "bolívar/es fuerte/s" - result = self.to_splitnum( - val, hightxt=hightxt, lowtxt="centavo/s", - divisor=1, jointxt="y", longval=longval - ) - # Handle exception, in spanish is "un euro" and not "uno euro" - return result.replace("uno", "un") diff --git a/fun_text_processing/num2words/num2words/lang_EU.py b/fun_text_processing/num2words/num2words/lang_EU.py deleted file mode 100644 index 4e649210f..000000000 --- a/fun_text_processing/num2words/num2words/lang_EU.py +++ /dev/null @@ -1,76 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import unicode_literals - -from .base import Num2Word_Base - -GENERIC_DOLLARS = ('dollar', 'dollars') -GENERIC_CENTS = ('cent', 'cents') - - -class Num2Word_EU(Num2Word_Base): - CURRENCY_FORMS = { - 'AUD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'CAD': (GENERIC_DOLLARS, GENERIC_CENTS), - # repalced by EUR - 'EEK': (('kroon', 'kroons'), ('sent', 'senti')), - 'EUR': (('euro', 'euro'), GENERIC_CENTS), - 'GBP': (('pound sterling', 'pounds sterling'), ('penny', 'pence')), - # replaced by EUR - 'LTL': (('litas', 'litas'), GENERIC_CENTS), - # replaced by EUR - 'LVL': (('lat', 'lats'), ('santim', 'santims')), - 'USD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'RUB': (('rouble', 'roubles'), ('kopek', 'kopeks')), - 'SEK': (('krona', 'kronor'), ('öre', 'öre')), - 'NOK': (('krone', 'kroner'), ('øre', 'øre')), - 'PLN': (('zloty', 'zlotys', 'zlotu'), ('grosz', 'groszy')), - 'MXN': (('peso', 'pesos'), GENERIC_CENTS), - 'RON': (('leu', 'lei', 'de lei'), ('ban', 'bani', 'de bani')), - 'INR': (('rupee', 'rupees'), ('paisa', 'paise')), - 'HUF': (('forint', 'forint'), ('fillér', 'fillér')) - } - - CURRENCY_ADJECTIVES = { - 'AUD': 'Australian', - 'CAD': 'Canadian', - 'EEK': 'Estonian', - 'USD': 'US', - 'RUB': 'Russian', - 'NOK': 'Norwegian', - 'MXN': 'Mexican', - 'RON': 'Romanian', - 'INR': 'Indian', - 'HUF': 'Hungarian' - } - - GIGA_SUFFIX = "illiard" - MEGA_SUFFIX = "illion" - - def set_high_numwords(self, high): - cap = 3 + 6 * len(high) - - for word, n in zip(high, range(cap, 3, -6)): - if self.GIGA_SUFFIX: - self.cards[10 ** n] = word + self.GIGA_SUFFIX - - if self.MEGA_SUFFIX: - self.cards[10 ** (n - 3)] = word + self.MEGA_SUFFIX - - def gen_high_numwords(self, units, tens, lows): - out = [u + t for t in tens for u in units] - out.reverse() - return out + lows - - def pluralize(self, n, forms): - form = 0 if n == 1 else 1 - return forms[form] - - def setup(self): - lows = ["non", "oct", "sept", "sext", "quint", "quadr", "tr", "b", "m"] - units = ["", "un", "duo", "tre", "quattuor", "quin", "sex", "sept", - "octo", "novem"] - tens = ["dec", "vigint", "trigint", "quadragint", "quinquagint", - "sexagint", "septuagint", "octogint", "nonagint"] - self.high_numwords = ["cent"] + self.gen_high_numwords(units, tens, - lows) diff --git a/fun_text_processing/num2words/num2words/lang_FI.py b/fun_text_processing/num2words/num2words/lang_FI.py deleted file mode 100644 index 86f9e097b..000000000 --- a/fun_text_processing/num2words/num2words/lang_FI.py +++ /dev/null @@ -1,721 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import division, print_function, unicode_literals - -from collections import OrderedDict - -from . import lang_EU - -GENERIC_CENTS = ('sentti', 'senttiä') -GENERIC_CENTAVOS = ('centavo', 'centavoa') - -# grammatical cases -NOM = 10 # nominative: the dictionary form -GEN = 11 # genitive: ~of/'s -ACC = 12 # accusative: not used; either nominative or genitive -PTV = 13 # partitive: as an object -# locative cases (internal) -INE = 14 # inessive: in -ELA = 15 # elative: from/out of -ILL = 16 # illative: into -# locative cases (external) -ADE = 17 # adessive: at/on -ABL = 18 # ablative: from (after being at/on, not in) -ALL = 19 # allative: to -# essive -ESS = 20 # essive: as (in the role of) -TRANSL = 21 # translative: to (the role of; being sth) -# rare -INSTRUC = 22 # instructive: with (plural is the same as singular) -ABE = 23 # abessive: without -COM = 24 # comitative: together with (plural = singular) - -NAME_TO_CASE = { - 'nominative': NOM, - 'genitive': GEN, - 'accusative': ACC, - 'partitive': PTV, - 'inessive': INE, - 'elative': ELA, - 'illative': ILL, - 'adessive': ADE, - 'ablative': ABL, - 'allative': ALL, - 'essive': ESS, - 'translative': TRANSL, - 'instructive': INSTRUC, - 'abessive': ABE, - 'comitative': COM, -} - -# https://en.wikibooks.org/wiki/Finnish/Grammar-Vowel_harmony -BACK_TO_FRONT = { - 'a': 'ä', - 'o': 'ö', - 'u': 'y', -} - -# https://en.wiktionary.org/wiki/Appendix:Finnish_nominal_inflection -# CASE: (SINGULAR_SUFFIX+, PLURAL_SUFFIX+) -KOTUS_TYPE = { - - # Kotus type 5/risti, no gradation - 5: { - # grammatical - NOM: ('i', 'it'), - GEN: ('in', 'ien'), - PTV: ('ia', 'eja'), - # locative, internal - INE: ('issa', 'eissa'), - ELA: ('ista', 'eista'), - ILL: ('iin', 'eihin'), - # locative, external - ADE: ('illa', 'eilla'), - ABL: ('ilta', 'eilta'), - ALL: ('ille', 'eille'), - # essive - ESS: ('ina', 'eina'), - TRANSL: ('iksi', 'eiksi'), - # rare - INSTRUC: ('ein', 'ein'), - ABE: ('itta', 'eitta'), - COM: ('eine', 'eine'), # works better - }, - - # Kotus type 7/ovi, no gradation - 7: { - # grammatical - NOM: ('i', 'et'), - GEN: ('en', 'ien'), - PTV: ('ea', 'ia'), - # locative, internal - INE: ('essa', 'issa'), - ELA: ('esta', 'ista'), - ILL: ('een', 'iin'), - # locative, external - ADE: ('ella', 'illa'), - ABL: ('elta', 'ilta'), - ALL: ('elle', 'ille'), - # essive - ESS: ('ena', 'ina'), - TRANSL: ('eksi', 'iksi'), - # rare - INSTRUC: ('in', 'in'), - ABE: ('etta', 'itta'), - COM: ('ine', 'ine'), # works better - }, - - # Kotus type 8/nalle, no gradation - 8: { - # grammatical - NOM: ('e', 'et'), - GEN: ('en', ('ejen', 'ein')), - PTV: ('ea', 'eja'), - # locative, internal - INE: ('essa', 'eissa'), - ELA: ('esta', 'eista'), - ILL: ('een', 'eihin'), - # locative, external - ADE: ('ella', 'eilla'), - ABL: ('elta', 'eilta'), - ALL: ('elle', 'eille'), - # essive - ESS: ('ena', 'eina'), - TRANSL: ('eksi', 'eiksi'), - # rare - INSTRUC: ('ein', 'ein'), - ABE: ('etta', 'eitta'), - COM: ('eine', 'eine'), # works better - }, - - # Kotus type 9/kala, t-d gradation (sata) - 109: { - # grammatical - NOM: ('ta', 'dat'), - GEN: ('dan', ('tojen', 'tain')), - PTV: ('taa', 'toja'), - # locative, internal - INE: ('dassa', 'doissa'), - ELA: ('dasta', 'doista'), - ILL: ('taan', 'toihin'), - # locative, external - ADE: ('dalla', 'doilla'), - ABL: ('dalta', 'doilta'), - ALL: ('dalle', 'doille'), - # essive - ESS: ('tana', 'toina'), - TRANSL: ('daksi', 'doiksi'), - # rare - INSTRUC: ('doin', 'doin'), - ABE: ('datta', 'doitta'), - COM: ('toine', 'toine'), # works better - }, - - # Kotus type 10/koira, no gradation - 10: { - # grammatical - NOM: ('a', 'at'), - GEN: ('an', ('ien', 'ain')), - PTV: ('aa', 'ia'), - # locative, internal - INE: ('assa', 'issa'), - ELA: ('asta', 'ista'), - ILL: ('aan', 'iin'), - # locative, external - ADE: ('alla', 'illa'), - ABL: ('alta', 'ilta'), - ALL: ('alle', 'ille'), - # essive - ESS: ('ana', 'ina'), - TRANSL: ('aksi', 'iksi'), - # rare - INSTRUC: ('in', 'in'), - ABE: ('atta', 'itta'), - COM: ('ine', 'ine'), # works better - }, - - # Kotus type 27/käsi, t-d gradation - 27: { - # grammatical - NOM: ('si', 'det'), - GEN: ('den', ('sien', 'tten')), - PTV: ('tta', 'sia'), - # locative, internal - INE: ('dessa', 'sissa'), - ELA: ('desta', 'sista'), - ILL: ('teen', 'siin'), - # locative, external - ADE: ('della', 'silla'), - ABL: ('delta', 'silta'), - ALL: ('delle', 'sille'), - # essive - ESS: ('tena', 'sina'), - TRANSL: ('deksi', 'siksi'), - # rare - INSTRUC: ('sin', 'sin'), - ABE: ('detta', 'sitta'), - COM: ('sine', 'sine'), # works better - }, - - # Kotus type 31/kaksi, t-d gradation - 31: { - # grammatical - NOM: ('ksi', 'hdet'), - GEN: ('hden', 'ksien'), - PTV: ('hta', 'ksia'), - # locative, internal - INE: ('hdessa', 'ksissa'), - ELA: ('hdesta', 'ksista'), - ILL: ('hteen', 'ksiin'), - # locative, external - ADE: ('hdella', 'ksilla'), - ABL: ('hdelta', 'ksilta'), - ALL: ('hdelle', 'ksille'), - # essive - ESS: ('htena', 'ksina'), - TRANSL: ('hdeksi', 'ksiksi'), - # rare - INSTRUC: ('ksin', 'ksin'), - ABE: ('hdetta', 'ksitta'), - COM: ('ksine', 'ksine'), # works better - }, - - # Kotus type 32/sisar, no gradation - 32: { - # grammatical - NOM: ('', 'et'), - GEN: ('en', ('ien', 'ten')), - PTV: ('ta', 'ia'), - # locative, internal - INE: ('essa', 'issa'), - ELA: ('esta', 'ista'), - ILL: ('een', 'iin'), - # locative, external - ADE: ('ella', 'illa'), - ABL: ('elta', 'ilta'), - ALL: ('elle', 'ille'), - # essive - ESS: ('ena', 'ina'), - TRANSL: ('eksi', 'iksi'), - # rare - INSTRUC: ('in', 'in'), - ABE: ('etta', 'itta'), - COM: ('ine', 'ine'), # works better - }, - - # Kotus type 38/nainen, no gradation - 38: { - # grammatical - NOM: ('nen', 'set'), - GEN: ('sen', ('sten', 'sien')), - PTV: ('sta', 'sia'), - # locative, internal - INE: ('sessa', 'sissa'), - ELA: ('sesta', 'sista'), - ILL: ('seen', 'siin'), - # locative, external - ADE: ('sella', 'silla'), - ABL: ('selta', 'silta'), - ALL: ('selle', 'sille'), - # essive - ESS: ('sena', 'sina'), - TRANSL: ('seksi', 'siksi'), - # rare - INSTRUC: ('sin', 'sin'), - ABE: ('setta', 'sitta'), - COM: ('sine', 'sine'), # works better - }, - - # Kotus type 45/kahdeksas, nt-nn gradation - 45: { - # grammatical - NOM: ('s', 'nnet'), - GEN: ('nnen', 'nsien'), - PTV: ('tta', 'nsia'), - # locative, internal - INE: ('nnessa', 'nsissa'), - ELA: ('nnesta', 'nsista'), - ILL: ('nteen', 'nsiin'), - # locative, external - ADE: ('nnella', 'nsilla'), - ABL: ('nnelta', 'nsilta'), - ALL: ('nnelle', 'nsille'), - # essive - ESS: ('ntena', 'nsina'), - TRANSL: ('nneksi', 'nsiksi'), - # rare - INSTRUC: ('nsin', 'nsin'), - ABE: ('nnetta', 'nsitta'), - COM: ('nsine', 'nsine'), # works better - }, - - # Kotus type 46/tuhat, nt-nn gradation - 46: { - # grammatical - NOM: ('t', 'nnet'), - GEN: ('nnen', ('nsien', 'nten')), - PTV: ('tta', 'nsia'), - # locative, internal - INE: ('nnessa', 'nsissa'), - ELA: ('nnesta', 'nsista'), - ILL: ('nteen', 'nsiin'), - # locative, external - ADE: ('nnella', 'nsilla'), - ABL: ('nnelta', 'nsilta'), - ALL: ('nnelle', 'nsille'), - # essive - ESS: ('ntena', 'nsina'), - TRANSL: ('nneksi', 'nsiksi'), - # rare - INSTRUC: ('nsin', 'nsin'), - ABE: ('nnetta', 'nsitta'), - COM: ('nsine', 'nsine'), # works better - }, -} - -# kolme -KOTUS_TYPE[108] = { - c: (KOTUS_TYPE[8][c][0], KOTUS_TYPE[7][c][1]) - for c in KOTUS_TYPE[8] -} -KOTUS_TYPE[108][INSTRUC] = ('en', 'in') -KOTUS_TYPE[108][ABE] = ('etta', 'itta') -KOTUS_TYPE[108][COM] = ('ine', 'ine') - -# seitsemän, kahdeksan, yhdeksän -KOTUS_TYPE[110] = KOTUS_TYPE[10].copy() -KOTUS_TYPE[110][NOM] = ('an', 'at') - -# kymmenen -KOTUS_TYPE[132] = KOTUS_TYPE[32].copy() -KOTUS_TYPE[132][NOM] = ('en', 'et') - - -def inflect(parts, options): - if not isinstance(parts, list): - parts = [parts] - - out = '' - for part in parts: - # part is plain text, concat and continue - if not isinstance(part, tuple): - out += part - continue - # predefined case (kaksikymmentä, ...) - tmp_case = options.case - if len(part) == 3: - # override singular nominative only - if options.case == NOM and not options.plural: - tmp_case = part[2] - part = part[:2] - # stem and suffix - stem, kotus_type = part - suffix = KOTUS_TYPE[kotus_type][tmp_case][options.plural] - # many choices, choose preferred or first - if isinstance(suffix, tuple): - common = set(suffix) & set(options.prefer or set()) - if len(common) == 1: - suffix = common.pop() - else: - suffix = suffix[0] - # apply vowel harmony - if not set(BACK_TO_FRONT) & set(stem): - for back, front in BACK_TO_FRONT.items(): - suffix = suffix.replace(back, front) - # concat - out += stem + suffix - - return out - - -class Options(object): - def __init__(self, ordinal, case, plural, prefer): - self.ordinal = ordinal - self.case = case - self.plural = plural - self.prefer = prefer - - def variation(self, ordinal=None, case=None, plural=None, prefer=None): - return Options( - ordinal if ordinal is not None else self.ordinal, - case if case is not None else self.case, - plural if plural is not None else self.plural, - prefer if prefer is not None else self.prefer, - ) - - -class Num2Word_FI(lang_EU.Num2Word_EU): - CURRENCY_FORMS = { - 'BRL': (('real', 'realia'), GENERIC_CENTAVOS), - 'CHF': (('frangi', 'frangia'), ('rappen', 'rappenia')), - 'CNY': (('juan', 'juania'), ('fen', 'feniä')), - 'EUR': (('euro', 'euroa'), GENERIC_CENTS), - 'FIM': (('markka', 'markkaa'), ('penni', 'penniä')), # historical - 'INR': (('rupia', 'rupiaa'), ('paisa', 'paisaa')), - 'JPY': (('jeni', 'jeniä'), ('sen', 'seniä')), # rare subunit - 'KRW': (('won', 'wonia'), ('jeon', 'jeonia')), # rare subunit - 'KPW': (('won', 'wonia'), ('chon', 'chonia')), # rare subunit - 'MXN': (('peso', 'pesoa'), GENERIC_CENTAVOS), - 'RUB': (('rupla', 'ruplaa'), ('kopeekka', 'kopeekkaa')), - 'TRY': (('liira', 'liiraa'), ('kuruş', 'kuruşia')), - 'ZAR': (('randi', 'randia'), GENERIC_CENTS), - } - - # crowns - for curr_code in 'DKK', 'ISK', 'NOK', 'SEK': - CURRENCY_FORMS[curr_code] = (('kruunu', 'kruunua'), ('äyri', 'äyriä')) - - # dollars - for curr_code in 'AUD', 'CAD', 'HKD', 'NZD', 'SGD', 'USD': - CURRENCY_FORMS[curr_code] = ( - ('dollari', 'dollaria'), GENERIC_CENTS) - - # pounds - for curr_code in ('GBP',): - CURRENCY_FORMS[curr_code] = (('punta', 'puntaa'), ('penny', 'pennyä')) - - CURRENCY_ADJECTIVES = { - 'AUD': 'Australian', - 'BRL': 'Brasilian', - 'CAD': 'Kanadan', - 'CHF': 'Sveitsin', - 'DKK': 'Tanskan', - 'FIM': 'Suomen', # historical - 'GBP': 'Englannin', - 'HKD': 'Hongkongin', - 'INR': 'Intian', - 'ISK': 'Islannin', - 'KRW': 'Etelä-Korean', - 'KPW': 'Pohjois-Korean', - 'MXN': 'Meksikon', - 'NOK': 'Norjan', - 'NZD': 'Uuden-Seelannin', - 'RUB': 'Venäjän', - 'SEK': 'Ruotsin', - 'SGD': 'Singaporen', - 'TRY': 'Turkin', - 'USD': 'Yhdysvaltain', - 'ZAR': 'Etelä-Afrikan', - } - - def __init__(self): - self.ords = OrderedDict() - super(Num2Word_FI, self).__init__() - - def set_numwords(self): - self.set_high_numwords(self.high_numwords) - self.set_mid_numwords(self.mid_numwords, self.mid_ords) - self.set_low_numwords(self.low_numwords, self.low_ords) - - def set_high_numwords(self, high): - # references: - # https://fi.wikipedia.org/wiki/Suurten_lukujen_nimet - # https://en.wikipedia.org/wiki/Names_of_large_numbers#Standard_dictionary_numbers - - # translate to Finnish - replacements = [ - ("qu", "kv"), - ("x", "ks"), - ("c", "k"), - ("kent", "sent"), # applied after c -> k to cent - ] - translated = [] - for i, numword in enumerate(high): - # notes: - # - 1e6**9 can be either noviljoona or noniljoona - # - 1e6**38 and above are untested - - # 1e6**6 is sekstiljoona but 1e6**16 is sedekiljoona - if numword.startswith("sex") and numword != "sext": - numword = numword.replace("sex", "se") - # 1e6**7 is septiljoona but 1e6**17 is septendekiljoona - elif numword.startswith("sept") and numword != "sept": - numword = "septen" + numword[len("sept"):] - # 1e6**8 is oktiljoona but 1e6**18 is duodevigintiljoona - # (2 from 20) - elif numword.startswith("octo"): - numword = high[i + -10] - numword = "duode" + numword[len("octo"):] - # 1e6**9 is noniljoona but 1e6**19 is undevigintiljoona (1 from 20) - elif numword.startswith("nove"): - numword = high[i + -10] - numword = "unde" + numword[len("nove") + 1:] - - # apply general replacements to all numwords - for repl in replacements: - numword = numword.replace(repl[0], repl[1]) - translated.append(numword) - - max = 6 * len(translated) - for word, n in zip(translated, range(max, 0, -6)): - if n == 6: - # irregularity considering short scale and long scale - self.cards[10 ** 9] = ("miljard", 5) - self.ords[10 ** 9] = ("miljardi", 45) - self.cards[10 ** n] = (word + "iljoon", 10) - self.ords[10 ** n] = (word + "iljoona", 45) - - def set_mid_numwords(self, cards, ords): - for key, val in cards: - self.cards[key] = val - for key, val in ords: - self.ords[key] = val - - def set_low_numwords(self, cards, ords): - for key, val in cards: - self.cards[key] = val - for key, val in ords: - self.ords[key] = val - - def setup(self): - super(Num2Word_FI, self).setup() - - self.negword = "miinus " - self.pointword = "pilkku" - self.exclude_title = ["pilkku", "miinus"] - - self.mid_numwords = [ - (1000, ("tuha", 46)), - (100, ("sa", 109)), - (90, [("yhdeks", 110), ("kymmen", 132, PTV)]), - (80, [("kahdeks", 110), ("kymmen", 132, PTV)]), - (70, [("seitsem", 110), ("kymmen", 132, PTV)]), - (60, [("kuu", 27), ("kymmen", 132, PTV)]), - (50, [("vii", 27), ("kymmen", 132, PTV)]), - (40, [("nelj", 10), ("kymmen", 132, PTV)]), - (30, [("kolm", 108), ("kymmen", 132, PTV)]), - ] - - self.mid_ords = [ - (1000, ("tuhanne", 45)), - (100, ("sada", 45)), - (90, [("yhdeksä", 45), ("kymmene", 45)]), - (80, [("kahdeksa", 45), ("kymmene", 45)]), - (70, [("seitsemä", 45), ("kymmene", 45)]), - (60, [("kuude", 45), ("kymmene", 45)]), - (50, [("viide", 45), ("kymmene", 45)]), - (40, [("neljä", 45), ("kymmene", 45)]), - (30, [("kolma", 45), ("kymmene", 45)]), - ] - - self.low_numwords = [ - (20, [("ka", 31), ("kymmen", 132, PTV)]), - (19, [("yhdeks", 110), "toista"]), - (18, [("kahdeks", 110), "toista"]), - (17, [("seitsem", 110), "toista"]), - (16, [("kuu", 27), "toista"]), - (15, [("vii", 27), "toista"]), - (14, [("nelj", 10), "toista"]), - (13, [("kolm", 108), "toista"]), - (12, [("ka", 31), "toista"]), - (11, [("y", 31), "toista"]), - (10, ("kymmen", 132)), - (9, ("yhdeks", 110)), - (8, ("kahdeks", 110)), - (7, ("seitsem", 110)), - (6, ("kuu", 27)), - (5, ("vii", 27)), - (4, ("nelj", 10)), - (3, ("kolm", 108)), - (2, ("ka", 31)), - (1, ("y", 31)), - (0, ("noll", 10)), - ] - - self.low_ords = [ - (20, [("kahde", 45), ("kymmene", 45)]), - (19, [("yhdeksä", 45), "toista"]), - (18, [("kahdeksa", 45), "toista"]), - (17, [("seitsemä", 45), "toista"]), - (16, [("kuude", 45), "toista"]), - (15, [("viide", 45), "toista"]), - (14, [("neljä", 45), "toista"]), - (13, [("kolma", 45), "toista"]), - (12, [("kahde", 45), "toista"]), - (11, [("yhde", 45), "toista"]), - (10, ("kymmene", 45)), - (9, ("yhdeksä", 45)), - (8, ("kahdeksa", 45)), - (7, ("seitsemä", 45)), - (6, ("kuude", 45)), - (5, ("viide", 45)), - (4, ("neljä", 45)), - (3, ("kolma", 45)), - (2, ("toi", 38)), - (1, ("ensimmäi", 38)), - (0, ("nolla", 45)), - ] - - def merge(self, lpair, rpair, options): - ltext, lnum = lpair - rtext, rnum = rpair - - # http://www.kielitoimistonohjepankki.fi/ohje/49 - fmt = "%s%s" - # ignore lpair if lnum is 1 - if lnum == 1: - rtext = inflect(rtext, options) - return (rtext, rnum) - # rnum is added to lnum - elif lnum > rnum: - ltext = inflect(ltext, options) - rtext = inflect(rtext, options) - # separate groups with space - if lnum >= 1000: - fmt = "%s %s" - return (fmt % (ltext, rtext), lnum + rnum) - # rnum is multiplied by lnum - elif lnum < rnum: - if options.ordinal: - # kahdessadas, not toinensadas - if lnum == 2: - ltext = ("kahde", 45) - rtext = inflect(rtext, options) - else: - # kaksituhatta but kahdettuhannet - rcase = options.case - if options.case == NOM and not options.plural: - rcase = PTV - rtext = inflect(rtext, options.variation(case=rcase)) - ltext = inflect(ltext, options) - return (fmt % (ltext, rtext), lnum * rnum) - - def to_cardinal(self, value, case='nominative', plural=False, prefer=None): - case = NAME_TO_CASE[case] - options = Options(False, case, plural, prefer) - try: - assert int(value) == value - except (ValueError, TypeError, AssertionError): - if case != NOM: - raise NotImplementedError( - "Cases other than nominative are not implemented for " - "cardinal floating point numbers.") - return self.to_cardinal_float(value) - - out = "" - if value < 0: - value = abs(value) - out = self.negword - - if value >= self.MAXVAL: - raise OverflowError(self.errmsg_toobig % (value, self.MAXVAL)) - - val = self.splitnum(value, options) - words, num = self.clean(val, options) - return self.title(out + words) - - def to_ordinal(self, value, case='nominative', plural=False, prefer=None): - case = NAME_TO_CASE[case] - options = Options(True, case, plural, prefer) - - self.verify_ordinal(value) - if value >= self.MAXVAL: - raise OverflowError(self.errmsg_toobig % (value, self.MAXVAL)) - - val = self.splitnum(value, options) - words, num = self.clean(val, options) - return self.title(words) - - def to_ordinal_num(self, value, case='nominative', plural=False): - case = NAME_TO_CASE[case] - raise NotImplementedError - - def to_year(self, val, suffix=None, longval=True): - suffix = suffix or "" - if val < 0: - val = abs(val) - suffix = suffix or " ennen ajanlaskun alkua" - return self.to_cardinal(val).replace(" ", "") + suffix - - def to_currency(self, val, currency="EUR", cents=True, separator=" ja", - adjective=False): - return super(Num2Word_FI, self).to_currency( - val, currency=currency, cents=cents, separator=separator, - adjective=adjective) - - def splitnum(self, value, options): - elems = self.ords if options.ordinal else self.cards - for elem in elems: - if elem > value: - continue - - out = [] - if value == 0: - div, mod = 1, 0 - else: - div, mod = divmod(value, elem) - - if div == 1: - out.append((elems[1], 1)) - else: - if div == value: # The system tallies, eg Roman Numerals - return [(div * elems[elem], div*elem)] - out.append(self.splitnum(div, options)) - - out.append((elems[elem], elem)) - - if mod: - out.append(self.splitnum(mod, options)) - - return out - - def clean(self, val, options): - out = val - while len(val) != 1: - out = [] - left, right = val[:2] - if isinstance(left, tuple) and isinstance(right, tuple): - out.append(self.merge(left, right, options)) - if val[2:]: - out.append(val[2:]) - else: - for elem in val: - if isinstance(elem, list): - if len(elem) == 1: - out.append(elem[0]) - else: - out.append(self.clean(elem, options)) - else: - out.append(elem) - val = out - return out[0] diff --git a/fun_text_processing/num2words/num2words/lang_FR.py b/fun_text_processing/num2words/num2words/lang_FR.py deleted file mode 100644 index 133d8027d..000000000 --- a/fun_text_processing/num2words/num2words/lang_FR.py +++ /dev/null @@ -1,93 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import print_function, unicode_literals - -from .lang_EU import Num2Word_EU - - -class Num2Word_FR(Num2Word_EU): - CURRENCY_FORMS = { - 'EUR': (('euro', 'euros'), ('centime', 'centimes')), - 'USD': (('dollar', 'dollars'), ('cent', 'cents')), - 'FRF': (('franc', 'francs'), ('centime', 'centimes')), - 'GBP': (('livre', 'livres'), ('penny', 'pence')), - 'CNY': (('yuan', 'yuans'), ('fen', 'jiaos')), - } - - def setup(self): - Num2Word_EU.setup(self) - - self.negword = "moins " - self.pointword = "virgule" - self.errmsg_nonnum = ( - u"Seulement des nombres peuvent être convertis en mots." - ) - self.errmsg_toobig = u"Nombre trop grand pour être converti en mots." - self.exclude_title = ["et", "virgule", "moins"] - self.mid_numwords = [(1000, "mille"), (100, "cent"), - (80, "quatre-vingts"), (60, "soixante"), - (50, "cinquante"), (40, "quarante"), - (30, "trente")] - self.low_numwords = ["vingt", "dix-neuf", "dix-huit", "dix-sept", - "seize", "quinze", "quatorze", "treize", "douze", - "onze", "dix", "neuf", "huit", "sept", "six", - "cinq", "quatre", "trois", "deux", "un", "zéro"] - self.ords = { - "cinq": "cinquième", - "neuf": "neuvième", - } - - def merge(self, curr, next): - ctext, cnum, ntext, nnum = curr + next - - if cnum == 1: - if nnum < 1000000: - return next - else: - if (not (cnum - 80) % 100 - or (not cnum % 100 and cnum < 1000))\ - and nnum < 1000000 \ - and ctext[-1] == "s": - ctext = ctext[:-1] - if cnum < 1000 and nnum != 1000 and \ - ntext[-1] != "s" and not nnum % 100: - ntext += "s" - - if nnum < cnum < 100: - if nnum % 10 == 1 and cnum != 80: - return ("%s et %s" % (ctext, ntext), cnum + nnum) - return ("%s-%s" % (ctext, ntext), cnum + nnum) - if nnum > cnum: - return ("%s %s" % (ctext, ntext), cnum * nnum) - return ("%s %s" % (ctext, ntext), cnum + nnum) - - # Is this right for such things as 1001 - "mille unième" instead of - # "mille premier"?? "millième"?? - - def to_ordinal(self, value): - self.verify_ordinal(value) - if value == 1: - return "premier" - word = self.to_cardinal(value) - for src, repl in self.ords.items(): - if word.endswith(src): - word = word[:-len(src)] + repl - break - else: - if word[-1] == "e": - word = word[:-1] - word = word + "ième" - return word - - def to_ordinal_num(self, value): - self.verify_ordinal(value) - out = str(value) - out += "er" if value == 1 else "me" - return out - - def to_currency(self, val, currency='EUR', cents=True, separator=' et', - adjective=False): - result = super(Num2Word_FR, self).to_currency( - val, currency=currency, cents=cents, separator=separator, - adjective=adjective) - return result diff --git a/fun_text_processing/num2words/num2words/lang_FR_BE.py b/fun_text_processing/num2words/num2words/lang_FR_BE.py deleted file mode 100644 index feecd3603..000000000 --- a/fun_text_processing/num2words/num2words/lang_FR_BE.py +++ /dev/null @@ -1,34 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import print_function, unicode_literals - -from .lang_FR import Num2Word_FR - - -class Num2Word_FR_BE(Num2Word_FR): - def setup(self): - Num2Word_FR.setup(self) - - self.mid_numwords = [(1000, "mille"), (100, "cent"), (90, "nonante"), - (80, "quatre-vingt"), (70, "septante"), - (60, "soixante"), (50, "cinquante"), - (40, "quarante"), (30, "trente")] - - def merge(self, curr, next): - ctext, cnum, ntext, nnum = curr + next - - if cnum == 1: - if nnum < 1000000: - return next - - if cnum < 1000 and nnum != 1000 and\ - ntext[-1] != "s" and not nnum % 100: - ntext += "s" - - if nnum < cnum < 100: - if nnum % 10 == 1: - return ("%s et %s" % (ctext, ntext), cnum + nnum) - return ("%s-%s" % (ctext, ntext), cnum + nnum) - if nnum > cnum: - return ("%s %s" % (ctext, ntext), cnum * nnum) - return ("%s %s" % (ctext, ntext), cnum + nnum) diff --git a/fun_text_processing/num2words/num2words/lang_FR_CH.py b/fun_text_processing/num2words/num2words/lang_FR_CH.py deleted file mode 100644 index 72b98a1d6..000000000 --- a/fun_text_processing/num2words/num2words/lang_FR_CH.py +++ /dev/null @@ -1,33 +0,0 @@ -# -*- coding: utf-8 -*- -from __future__ import print_function, unicode_literals - -from .lang_FR import Num2Word_FR - - -class Num2Word_FR_CH(Num2Word_FR): - def setup(self): - Num2Word_FR.setup(self) - - self.mid_numwords = [(1000, "mille"), (100, "cent"), (90, "nonante"), - (80, "huitante"), (70, "septante"), - (60, "soixante"), (50, "cinquante"), - (40, "quarante"), (30, "trente")] - - def merge(self, curr, next): - ctext, cnum, ntext, nnum = curr + next - - if cnum == 1: - if nnum < 1000000: - return next - - if cnum < 1000 and nnum != 1000 and\ - ntext[-1] != "s" and not nnum % 100: - ntext += "s" - - if nnum < cnum < 100: - if nnum % 10 == 1: - return ("%s et %s" % (ctext, ntext), cnum + nnum) - return ("%s-%s" % (ctext, ntext), cnum + nnum) - if nnum > cnum: - return ("%s %s" % (ctext, ntext), cnum * nnum) - return ("%s %s" % (ctext, ntext), cnum + nnum) diff --git a/fun_text_processing/num2words/num2words/lang_FR_DZ.py b/fun_text_processing/num2words/num2words/lang_FR_DZ.py deleted file mode 100644 index 9920cb605..000000000 --- a/fun_text_processing/num2words/num2words/lang_FR_DZ.py +++ /dev/null @@ -1,18 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import unicode_literals - -from .lang_FR import Num2Word_FR - - -class Num2Word_FR_DZ(Num2Word_FR): - CURRENCY_FORMS = { - 'DIN': (('dinard', 'dinards'), ('centime', 'centimes')), - } - - def to_currency(self, val, currency='DIN', cents=True, separator=' et', - adjective=False): - result = super(Num2Word_FR, self).to_currency( - val, currency=currency, cents=cents, separator=separator, - adjective=adjective) - return result diff --git a/fun_text_processing/num2words/num2words/lang_HE.py b/fun_text_processing/num2words/num2words/lang_HE.py deleted file mode 100644 index 3117301a0..000000000 --- a/fun_text_processing/num2words/num2words/lang_HE.py +++ /dev/null @@ -1,144 +0,0 @@ -# -*- coding: utf-8 -*- - - -from __future__ import print_function, unicode_literals - -from .base import Num2Word_Base -from .utils import get_digits, splitbyx - -ZERO = (u'אפס',) - -ONES = { - 1: (u'אחת',), - 2: (u'שתים',), - 3: (u'שלש',), - 4: (u'ארבע',), - 5: (u'חמש',), - 6: (u'שש',), - 7: (u'שבע',), - 8: (u'שמונה',), - 9: (u'תשע',), -} - -TENS = { - 0: (u'עשר',), - 1: (u'אחת עשרה',), - 2: (u'שתים עשרה',), - 3: (u'שלש עשרה',), - 4: (u'ארבע עשרה',), - 5: (u'חמש עשרה',), - 6: (u'שש עשרה',), - 7: (u'שבע עשרה',), - 8: (u'שמונה עשרה',), - 9: (u'תשע עשרה',), -} - -TWENTIES = { - 2: (u'עשרים',), - 3: (u'שלשים',), - 4: (u'ארבעים',), - 5: (u'חמישים',), - 6: (u'ששים',), - 7: (u'שבעים',), - 8: (u'שמונים',), - 9: (u'תשעים',), -} - -HUNDRED = { - 1: (u'מאה',), - 2: (u'מאתיים',), - 3: (u'מאות',) -} - -THOUSANDS = { - 1: (u'אלף',), - 2: (u'אלפיים',), - 3: (u'שלשת אלפים',), - 4: (u'ארבעת אלפים',), - 5: (u'חמשת אלפים',), - 6: (u'ששת אלפים',), - 7: (u'שבעת אלפים',), - 8: (u'שמונת אלפים',), - 9: (u'תשעת אלפים',), -} - -AND = u'ו' - - -def pluralize(n, forms): - # gettext implementation: - # (n%10==1 && n%100!=11 ? 0 : n != 0 ? 1 : 2) - - form = 0 if (n % 10 == 1 and n % 100 != 11) else 1 if n != 0 else 2 - - return forms[form] - - -def int2word(n): - if n > 9999: # doesn't yet work for numbers this big - raise NotImplementedError() - - if n == 0: - return ZERO[0] - - words = [] - - chunks = list(splitbyx(str(n), 3)) - i = len(chunks) - for x in chunks: - i -= 1 - - if x == 0: - continue - - n1, n2, n3 = get_digits(x) - - if i > 0: - words.append(THOUSANDS[n1][0]) - continue - - if n3 > 0: - if n3 <= 2: - words.append(HUNDRED[n3][0]) - else: - words.append(ONES[n3][0] + ' ' + HUNDRED[3][0]) - - if n2 > 1: - words.append(TWENTIES[n2][0]) - - if n2 == 1: - words.append(TENS[n1][0]) - elif n1 > 0 and not (i > 0 and x == 1): - words.append(ONES[n1][0]) - - if i > 0: - words.append(THOUSANDS[i][0]) - - # source: https://hebrew-academy.org.il/2017/01/30/ו-החיבור-במספרים/ - if len(words) > 1: - words[-1] = AND + words[-1] - - return ' '.join(words) - - -def n2w(n): - return int2word(int(n)) - - -def to_currency(n, currency='EUR', cents=True, separator=','): - raise NotImplementedError() - - -class Num2Word_HE(Num2Word_Base): - def to_cardinal(self, number): - return n2w(number) - - def to_ordinal(self, number): - raise NotImplementedError() - - -if __name__ == '__main__': - yo = Num2Word_HE() - nums = [1, 11, 21, 24, 99, 100, 101, 200, 211, 345, 1000, 1011] - for num in nums: - print(num, yo.to_cardinal(num)) diff --git a/fun_text_processing/num2words/num2words/lang_HU.py b/fun_text_processing/num2words/num2words/lang_HU.py deleted file mode 100644 index 32591e9d0..000000000 --- a/fun_text_processing/num2words/num2words/lang_HU.py +++ /dev/null @@ -1,150 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import division, print_function, unicode_literals - -from . import lang_EU - -ZERO = 'nulla' - - -class Num2Word_HU(lang_EU.Num2Word_EU): - GIGA_SUFFIX = "illiárd" - MEGA_SUFFIX = "illió" - - def setup(self): - super(Num2Word_HU, self).setup() - - self.negword = "mínusz " - self.pointword = "egész" - - self.mid_numwords = [(1000, "ezer"), (100, "száz"), (90, "kilencven"), - (80, "nyolcvan"), (70, "hetven"), (60, "hatvan"), - (50, "ötven"), (40, "negyven"), (30, "harminc")] - - low_numwords = ["kilenc", "nyolc", "hét", "hat", "öt", "négy", "három", - "kettő", "egy"] - self.low_numwords = (['tizen' + w for w in low_numwords] - + ['tíz'] - + low_numwords) - self.low_numwords = (['huszon' + w for w in low_numwords] - + ['húsz'] - + self.low_numwords - + [ZERO]) - - self.partial_ords = { - 'nulla': 'nullad', - 'egy': 'egyed', - 'kettő': 'ketted', - 'három': 'harmad', - 'négy': 'negyed', - 'öt': 'ötöd', - 'hat': 'hatod', - 'hét': 'heted', - 'nyolc': 'nyolcad', - 'kilenc': 'kilenced', - 'tíz': 'tized', - 'húsz': 'huszad', - 'harminc': 'harmincad', - 'negyven': 'negyvened', - 'ötven': 'ötvened', - 'hatvan': 'hatvanad', - 'hetven': 'hetvened', - 'nyolcvan': 'nyolcvanad', - 'kilencven': 'kilencvened', - 'száz': 'század', - 'ezer': 'ezred', - 'illió': 'milliomod', - 'illiárd': 'milliárdod' - } - - def to_cardinal(self, value, zero=ZERO): - if int(value) != value: - return self.to_cardinal_float(value) - elif value < 0: - out = self.negword + self.to_cardinal(-value) - elif value == 0: - out = zero - elif zero == '' and value == 2: - out = 'két' - elif value < 30: - out = self.cards[value] - elif value < 100: - out = self.tens_to_cardinal(value) - elif value < 1000: - out = self.hundreds_to_cardinal(value) - elif value < 10**6: - out = self.thousands_to_cardinal(value) - else: - out = self.big_number_to_cardinal(value) - return out - - def tens_to_cardinal(self, value): - try: - return self.cards[value] - except KeyError: - return self.cards[value // 10 * 10] + self.to_cardinal(value % 10) - - def hundreds_to_cardinal(self, value): - hundreds = value // 100 - prefix = "száz" - if hundreds != 1: - prefix = self.to_cardinal(hundreds, zero="") + prefix - postfix = self.to_cardinal(value % 100, zero="") - return prefix + postfix - - def thousands_to_cardinal(self, value): - thousands = value // 1000 - prefix = "ezer" - if thousands != 1: - prefix = self.to_cardinal(thousands, zero="") + prefix - postfix = self.to_cardinal(value % 1000, zero="") - return prefix + ('' if value <= 2000 or not postfix else '-') + postfix - - def big_number_to_cardinal(self, value): - digits = len(str(value)) - digits = digits if digits % 3 != 0 else digits - 2 - exp = 10 ** (digits // 3 * 3) - rest = self.to_cardinal(value % exp, '') - return (self.to_cardinal(value // exp, '') + self.cards[exp] - + ('-' + rest if rest else '')) - - def to_ordinal(self, value): - if value < 0: - return self.negword + self.to_ordinal(-value) - if value == 1: - return 'első' - elif value == 2: - return 'második' - else: - out = self.to_cardinal(value) - for card_word, ord_word in self.partial_ords.items(): - if out[-len(card_word):] == card_word: - out = out[:-len(card_word)] + ord_word - break - return out + 'ik' - - def to_ordinal_num(self, value): - self.verify_ordinal(value) - return str(value) + '.' - - def to_year(self, val, suffix=None, longval=True): - # suffix is prefix here - prefix = '' - if val < 0 or suffix is not None: - val = abs(val) - prefix = (suffix + ' ' if suffix is not None else 'i. e. ') - return prefix + self.to_cardinal(val) - - def to_currency(self, val, currency='HUF', cents=True, separator=',', - adjective=False): - return super(Num2Word_HU, self).to_currency( - val, currency, cents, separator, adjective) - - def to_cardinal_float(self, value): - if abs(value) != value: - return self.negword + self.to_cardinal_float(-value) - left, right = str(value).split('.') - return (self.to_cardinal(int(left)) - + ' egész ' - + self.to_cardinal(int(right)) - + ' ' + self.partial_ords[self.cards[10 ** len(right)]]) diff --git a/fun_text_processing/num2words/num2words/lang_ID.py b/fun_text_processing/num2words/num2words/lang_ID.py deleted file mode 100644 index 2648918ef..000000000 --- a/fun_text_processing/num2words/num2words/lang_ID.py +++ /dev/null @@ -1,189 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import print_function, unicode_literals - - -class Num2Word_ID(): - BASE = {0: [], - 1: ["satu"], - 2: ["dua"], - 3: ["tiga"], - 4: ["empat"], - 5: ["lima"], - 6: ["enam"], - 7: ["tujuh"], - 8: ["delapan"], - 9: ["sembilan"]} - - TENS_TO = {3: "ribu", - 6: "juta", - 9: "miliar", - 12: "triliun", - 15: "kuadriliun", - 18: "kuantiliun", - 21: "sekstiliun", - 24: "septiliun", - 27: "oktiliun", - 30: "noniliun", - 33: "desiliun"} - - errmsg_floatord = "Cannot treat float number as ordinal" - errmsg_negord = "Cannot treat negative number as ordinal" - errmsg_toobig = "Too large" - max_num = 10 ** 36 - - def split_by_koma(self, number): - return str(number).split('.') - - def split_by_3(self, number): - """ - starting here, it groups the number by three from the tail - '1234567' -> (('1',),('234',),('567',)) - :param number:str - :rtype:tuple - """ - blocks = () - length = len(number) - - if length < 3: - blocks += ((number,),) - else: - len_of_first_block = length % 3 - - if len_of_first_block > 0: - first_block = number[0:len_of_first_block], - blocks += first_block, - - for i in range(len_of_first_block, length, 3): - next_block = (number[i:i + 3],), - blocks += next_block - - return blocks - - def spell(self, blocks): - """ - it adds the list of spelling to the blocks - ( - ('1',),('034',)) -> (('1',['satu']),('234',['tiga', 'puluh', 'empat']) - ) - :param blocks: tuple - :rtype: tuple - """ - word_blocks = () - first_block = blocks[0] - if len(first_block[0]) == 1: - if first_block[0] == '0': - spelling = ['nol'] - else: - spelling = self.BASE[int(first_block[0])] - elif len(first_block[0]) == 2: - spelling = self.puluh(first_block[0]) - else: - spelling = ( - self.ratus(first_block[0][0]) + self.puluh(first_block[0][1:3]) - ) - - word_blocks += (first_block[0], spelling), - - for block in blocks[1:]: - spelling = self.ratus(block[0][0]) + self.puluh(block[0][1:3]) - block += spelling, - word_blocks += block, - - return word_blocks - - def ratus(self, number): - # it is used to spell - if number == '1': - return ['seratus'] - elif number == '0': - return [] - else: - return self.BASE[int(number)] + ['ratus'] - - def puluh(self, number): - # it is used to spell - if number[0] == '1': - if number[1] == '0': - return ['sepuluh'] - elif number[1] == '1': - return ['sebelas'] - else: - return self.BASE[int(number[1])] + ['belas'] - elif number[0] == '0': - return self.BASE[int(number[1])] - else: - return ( - self.BASE[int(number[0])] + ['puluh'] - + self.BASE[int(number[1])] - ) - - def spell_float(self, float_part): - # spell the float number - word_list = [] - for n in float_part: - if n == '0': - word_list += ['nol'] - continue - word_list += self.BASE[int(n)] - return ' '.join(['', 'koma'] + word_list) - - def join(self, word_blocks, float_part): - """ - join the words by first join lists in the tuple - :param word_blocks: tuple - :rtype: str - """ - word_list = [] - length = len(word_blocks) - 1 - first_block = word_blocks[0], - start = 0 - - if length == 1 and first_block[0][0] == '1': - word_list += ['seribu'] - start = 1 - - for i in range(start, length + 1, 1): - word_list += word_blocks[i][1] - if not word_blocks[i][1]: - continue - if i == length: - break - word_list += [self.TENS_TO[(length - i) * 3]] - - return ' '.join(word_list) + float_part - - def to_cardinal(self, number): - if number >= self.max_num: - raise OverflowError(self.errmsg_toobig % (number, self.max_num)) - minus = '' - if number < 0: - minus = 'min ' - float_word = '' - n = self.split_by_koma(abs(number)) - if len(n) == 2: - float_word = self.spell_float(n[1]) - return minus + self.join(self.spell(self.split_by_3(n[0])), float_word) - - def to_ordinal(self, number): - self.verify_ordinal(number) - out_word = self.to_cardinal(number) - if out_word == "satu": - return "pertama" - return "ke" + out_word - - def to_ordinal_num(self, number): - self.verify_ordinal(number) - return "ke-" + str(number) - - def to_currency(self, value): - return self.to_cardinal(value) + " rupiah" - - def to_year(self, value): - return self.to_cardinal(value) - - def verify_ordinal(self, value): - if not value == int(value): - raise TypeError(self.errmsg_floatord % value) - if not abs(value) == value: - raise TypeError(self.errmsg_negord % value) diff --git a/fun_text_processing/num2words/num2words/lang_IT.py b/fun_text_processing/num2words/num2words/lang_IT.py deleted file mode 100644 index 6966d7300..000000000 --- a/fun_text_processing/num2words/num2words/lang_IT.py +++ /dev/null @@ -1,208 +0,0 @@ -# -*- coding: utf-8 -*- -# Copyright (c) 2018-2019, Filippo Costa. All Rights Reserved. - -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, -# MA 02110-1301 USA - -from __future__ import unicode_literals - -# Globals -# ------- - -ZERO = "zero" - -CARDINAL_WORDS = [ - ZERO, "uno", "due", "tre", "quattro", "cinque", "sei", "sette", "otto", - "nove", "dieci", "undici", "dodici", "tredici", "quattordici", "quindici", - "sedici", "diciassette", "diciotto", "diciannove" -] - -ORDINAL_WORDS = [ - ZERO, "primo", "secondo", "terzo", "quarto", "quinto", "sesto", "settimo", - "ottavo", "nono", "decimo", "undicesimo", "dodicesimo", "tredicesimo", - "quattordicesimo", "quindicesimo", "sedicesimo", "diciassettesimo", - "diciottesimo", "diciannovesimo" -] - -# The script can extrapolate the missing numbers from the base forms. -STR_TENS = {2: "venti", 3: "trenta", 4: "quaranta", 6: "sessanta"} - -# These prefixes are used for extremely big numbers. -EXPONENT_PREFIXES = [ - ZERO, "m", "b", "tr", "quadr", "quint", "sest", "sett", "ott", "nov", "dec" -] - - -# Utils -# ===== - -def phonetic_contraction(string): - return (string - .replace("oo", "o") # ex. "centootto" - .replace("ao", "o") # ex. "settantaotto" - .replace("io", "o") # ex. "ventiotto" - .replace("au", "u") # ex. "trentauno" - .replace("iu", "u") # ex. "ventiunesimo" - ) - - -def exponent_length_to_string(exponent_length): - # We always assume `exponent` to be a multiple of 3. If it's not true, then - # Num2Word_IT.big_number_to_cardinal did something wrong. - prefix = EXPONENT_PREFIXES[exponent_length // 6] - if exponent_length % 6 == 0: - return prefix + "ilione" - else: - return prefix + "iliardo" - - -def accentuate(string): - # This is inefficient: it may do several rewritings when deleting - # half-sentence accents. However, it is the easiest method and speed is - # not crucial (duh), so... - return " ".join( - # Deletes half-sentence accents and accentuates the last "tre" - [w.replace("tré", "tre")[:-3] + "tré" - # We shouldn't accentuate a single "tre": is has to be a composite - # word. ~~~~~~~~~~ - if w[-3:] == "tre" and len(w) > 3 - # Deletes half-sentence accents anyway - # ~~~~~~~~~~~~~~~~~~~~~~ - else w.replace("tré", "tre") - for w in string.split() - ]) - - -def omitt_if_zero(number_to_string): - return "" if number_to_string == ZERO else number_to_string - - -# Main class -# ========== - -class Num2Word_IT: - MINUS_PREFIX_WORD = "meno " - FLOAT_INFIX_WORD = " virgola " - - def __init__(self): - pass - - def float_to_words(self, float_number, ordinal=False): - if ordinal: - prefix = self.to_ordinal(int(float_number)) - else: - prefix = self.to_cardinal(int(float_number)) - float_part = str(float_number).split('.')[1] - postfix = " ".join( - # Drops the trailing zero and comma - [self.to_cardinal(int(c)) for c in float_part] - ) - return prefix + Num2Word_IT.FLOAT_INFIX_WORD + postfix - - def tens_to_cardinal(self, number): - tens = number // 10 - units = number % 10 - if tens in STR_TENS: - prefix = STR_TENS[tens] - else: - prefix = CARDINAL_WORDS[tens][:-1] + "anta" - postfix = omitt_if_zero(CARDINAL_WORDS[units]) - return phonetic_contraction(prefix + postfix) - - def hundreds_to_cardinal(self, number): - hundreds = number // 100 - prefix = "cento" - if hundreds != 1: - prefix = CARDINAL_WORDS[hundreds] + prefix - postfix = omitt_if_zero(self.to_cardinal(number % 100)) - return phonetic_contraction(prefix + postfix) - - def thousands_to_cardinal(self, number): - thousands = number // 1000 - if thousands == 1: - prefix = "mille" - else: - prefix = self.to_cardinal(thousands) + "mila" - postfix = omitt_if_zero(self.to_cardinal(number % 1000)) - # "mille" and "mila" don't need any phonetic contractions - return prefix + postfix - - def big_number_to_cardinal(self, number): - digits = [c for c in str(number)] - length = len(digits) - if length >= 66: - raise NotImplementedError("The given number is too large.") - # This is how many digits come before the "illion" term. - # cento miliardi => 3 - # dieci milioni => 2 - # un miliardo => 1 - predigits = length % 3 or 3 - multiplier = digits[:predigits] - exponent = digits[predigits:] - # Default infix string: "milione", "biliardo", "sestilione", ecc. - infix = exponent_length_to_string(len(exponent)) - if multiplier == ["1"]: - prefix = "un " - else: - prefix = self.to_cardinal(int("".join(multiplier))) - # Plural form ~~~~~~~~~~~ - infix = " " + infix[:-1] + "i" - # Read as: Does the value of exponent equal 0? - if set(exponent) != set("0"): - postfix = self.to_cardinal(int("".join(exponent))) - if " e " in postfix: - infix += ", " - else: - infix += " e " - else: - postfix = "" - return prefix + infix + postfix - - def to_cardinal(self, number): - if number < 0: - string = Num2Word_IT.MINUS_PREFIX_WORD + self.to_cardinal(-number) - elif isinstance(number, float): - string = self.float_to_words(number) - elif number < 20: - string = CARDINAL_WORDS[number] - elif number < 100: - string = self.tens_to_cardinal(number) - elif number < 1000: - string = self.hundreds_to_cardinal(number) - elif number < 1000000: - string = self.thousands_to_cardinal(number) - else: - string = self.big_number_to_cardinal(number) - return accentuate(string) - - def to_ordinal(self, number): - tens = number % 100 - # Italian grammar is poorly defined here ¯\_(ツ)_/¯: - # centodecimo VS centodieciesimo VS centesimo decimo? - is_outside_teens = not 10 < tens < 20 - if number < 0: - return Num2Word_IT.MINUS_PREFIX_WORD + self.to_ordinal(-number) - elif number % 1 != 0: - return self.float_to_words(number, ordinal=True) - elif number < 20: - return ORDINAL_WORDS[number] - elif is_outside_teens and tens % 10 == 3: - # Gets ride of the accent ~~~~~~~~~~ - return self.to_cardinal(number)[:-1] + "eesimo" - elif is_outside_teens and tens % 10 == 6: - return self.to_cardinal(number) + "esimo" - else: - string = self.to_cardinal(number)[:-1] - if string[-3:] == "mil": - string += "l" - return string + "esimo" diff --git a/fun_text_processing/num2words/num2words/lang_IW.py b/fun_text_processing/num2words/num2words/lang_IW.py deleted file mode 100644 index c43941fa4..000000000 --- a/fun_text_processing/num2words/num2words/lang_IW.py +++ /dev/null @@ -1,91 +0,0 @@ -# -*- coding: utf-8 -*- - -# Hebrew num2words - -from __future__ import division, print_function, unicode_literals - -from . import lang_EU - - -class Num2Word_IW(lang_EU.Num2Word_EU): - def set_high_numwords(self, high): - max = 3 + 3 * len(high) - for word, n in zip(high, range(max, 3, -3)): - self.cards[10 ** n] = word + " מִילִיוֹן" - - def setup(self): - super(Num2Word_IW, self).setup() - - self.negword = "פחות " - self.pointword = "נְקוּדָה" - self.exclude_title = ["ו", "נְקוּדָה", "פחות"] - - self.mid_numwords = [(1000, "אלף"), (100, "מאה"), - (90, "תִשׁעִים"), (80, "שמונים"), (70, "שִׁבעִים"), - (60, "שִׁשִׁים"), (50, "חמישים"), (40, "ארבעים"), - (30, "שְׁלוֹשִׁים")] - self.low_numwords = ['עשרים', 'תשע עשרה', 'שמונה עשרה', 'שבע עשרה', 'שש עשרה', 'חֲמֵשׁ עֶשׂרֵה', 'ארבעה עשר', 'שְׁלוֹשׁ עֶשׂרֵה', 'שתיים עשרה', 'אחד עשר', 'עשר', 'תֵשַׁע', 'שמונה', 'שבע', 'שֵׁשׁ', 'חָמֵשׁ', 'ארבע', 'שְׁלוֹשָׁה', 'שתיים', 'אחד', 'אֶפֶס'] - self.ords = {'אחד': 'ראשון', -'שתיים': 'שְׁנִיָה', -'שְׁלוֹשָׁה': 'שְׁלִישִׁי', -'ארבע': 'רביעי', -'חָמֵשׁ': 'חמישי', -'שֵׁשׁ': 'שִׁשִׁית', -'שבע': 'שְׁבִיעִית', -'שמונה': 'שמונה', -'תֵשַׁע': 'ט', -'עשר': 'עֲשִׂירִית'} - - def merge(self, lpair, rpair): - ltext, lnum = lpair - rtext, rnum = rpair - if lnum == 1 and rnum < 100: - return (rtext, rnum) - elif 100 > lnum > rnum: - return ("%s-%s" % (ltext, rtext), lnum + rnum) - elif lnum >= 100 > rnum: - return ("%s ו %s" % (ltext, rtext), lnum + rnum) - elif rnum > lnum: - return ("%s %s" % (ltext, rtext), lnum * rnum) - return ("%s, %s" % (ltext, rtext), lnum + rnum) - - # def to_ordinal(self, value): - # self.verify_ordinal(value) - # outwords = self.to_cardinal(value).split(" ") - # lastwords = outwords[-1].split("-") - # lastword = lastwords[-1].lower() - # try: - # lastword = self.ords[lastword] - # except KeyError: - # if lastword[-1] == "y": - # lastword = lastword[:-1] + "ie" - # lastword += "th" - # lastwords[-1] = self.title(lastword) - # outwords[-1] = "-".join(lastwords) - # return " ".join(outwords) - - # def to_ordinal_num(self, value): - # self.verify_ordinal(value) - # return "%s%s" % (value, self.to_ordinal(value)[-2:]) - - def to_year(self, val, suffix=None, longval=True): - if val < 0: - val = abs(val) - suffix = 'לִפנֵי הַסְפִירָה' if not suffix else suffix - high, low = (val // 100, val % 100) - # If year is 00XX, X00X, or beyond 9999, go cardinal. - if (high == 0 - or (high % 10 == 0 and low < 10) - or high >= 100): - valtext = self.to_cardinal(val) - else: - hightext = self.to_cardinal(high) - if low == 0: - lowtext = "מֵאָה" - elif low < 10: - lowtext = "%s" % self.to_cardinal(low) - else: - lowtext = self.to_cardinal(low) - valtext = "%s %s" % (hightext, lowtext) - return (valtext if not suffix - else "%s %s" % (valtext, suffix)) diff --git a/fun_text_processing/num2words/num2words/lang_JA.py b/fun_text_processing/num2words/num2words/lang_JA.py deleted file mode 100644 index b3ade3e88..000000000 --- a/fun_text_processing/num2words/num2words/lang_JA.py +++ /dev/null @@ -1,574 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import division, print_function, unicode_literals - -from .base import Num2Word_Base -from .compat import strtype, to_s -from .currency import parse_currency_parts, prefix_currency - - -def select_text(text, reading=False, prefer=None): - """Select the correct text from the Japanese number, reading and - alternatives""" - # select kanji number or kana reading - if reading: - text = text[1] - else: - text = text[0] - - # select the preferred one or the first one from multiple alternatives - if not isinstance(text, strtype): - common = set(text) & set(prefer or set()) - if len(common) == 1: - text = common.pop() - else: - text = text[0] - - return text - - -def rendaku_merge_pairs(lpair, rpair): - """Merge lpair < rpair while applying semi-irregular rendaku rules""" - ltext, lnum = lpair - rtext, rnum = rpair - if lnum > rnum: - raise ValueError - - if rpair == ("ひゃく", 100): - if lpair == ("さん", 3): - rtext = "びゃく" - elif lpair == ("ろく", 6): - ltext = "ろっ" - rtext = "ぴゃく" - elif lpair == ("はち", 8): - ltext = "はっ" - rtext = "ぴゃく" - elif rpair == ("せん", 1000): - if lpair == ("さん", 3): - rtext = "ぜん" - elif lpair == ("はち", 8): - ltext = "はっ" - elif rpair == ("ちょう", 10**12): - if lpair == ("いち", 1): - ltext = "いっ" - elif lpair == ("はち", 8): - ltext = "はっ" - elif lpair == ("じゅう", 10): - ltext = "じゅっ" - elif rpair == ("けい", 10**16): - if lpair == ("いち", 1): - ltext = "いっ" - elif lpair == ("ろく", 6): - ltext = "ろっ" - elif lpair == ("はち", 8): - ltext = "はっ" - elif lpair == ("じゅう", 10): - ltext = "じゅっ" - elif lpair == ("ひゃく", 100): - ltext = "ひゃっ" - - return ("%s%s" % (ltext, rtext), lnum * rnum) - - -# Source: https://www.sljfaq.org/afaq/era-list.html -# if there are multiple eras for the same year, use the last one -ERA_START = [ - (645, ("大化", "たいか")), - (650, ("白雉", "はくち")), - (686, ("朱鳥", "しゅちょう")), - (701, ("大宝", "たいほう")), - (704, ("慶雲", "けいうん")), - (708, ("和銅", "わどう")), - (715, ("霊亀", "れいき")), - (717, ("養老", "ようろう")), - (724, ("神亀", "じんき")), - (729, ("天平", "てんぴょう")), - (749, ("天平感宝", "てんぴょうかんぽう")), - (749, ("天平勝宝", "てんぴょうしょうほう")), - (757, ("天平宝字", "てんぴょうじょうじ")), - (765, ("天平神護", "てんぴょうじんご")), - (767, ("神護景雲", "じんごけいうん")), - (770, ("宝亀", "ほうき")), - (781, ("天応", "てんおう")), - (782, ("延暦", "えんりゃく")), - (806, ("大同", "だいどう")), - (810, ("弘仁", "こうにん")), - (823, ("天長", "てんちょう")), - (834, ("承和", "じょうわ")), - (848, ("嘉祥", "かしょう")), - (851, ("仁寿", "にんじゅ")), - (855, ("斉衡", "さいこう")), - (857, ("天安", "てんあん")), - (859, ("貞観", "じょうがん")), - (877, ("元慶", "がんぎょう")), - (885, ("仁和", "にんな")), - (889, ("寛平", "かんぴょう")), - (898, ("昌泰", "しょうたい")), - (901, ("延喜", "えんぎ")), - (923, ("延長", "えんちょう")), - (931, ("承平", "じょうへい")), - (938, ("天慶", "てんぎょう")), - (947, ("天暦", "てんりゃく")), - (957, ("天徳", "てんとく")), - (961, ("応和", "おうわ")), - (964, ("康保", "こうほう")), - (968, ("安和", "あんな")), - (970, ("天禄", "てんろく")), - (974, ("天延", "てんえん")), - (976, ("貞元", "じょうげん")), - (979, ("天元", "てんげん")), - (983, ("永観", "えいかん")), - (985, ("寛和", "かんな")), - (987, ("永延", "えいえん")), - (989, ("永祚", "えいそ")), - (990, ("正暦", "しょうりゃく")), - (995, ("長徳", "ちょうとく")), - (999, ("長保", "ちょうほう")), - (1004, ("寛弘", "かんこう")), - (1013, ("長和", "ちょうわ")), - (1017, ("寛仁", "かんにん")), - (1021, ("治安", "じあん")), - (1024, ("万寿", "まんじゅ")), - (1028, ("長元", "ちょうげん")), - (1037, ("長暦", "ちょうりゃく")), - (1040, ("長久", "ちょうきゅう")), - (1045, ("寛徳", "かんとく")), - (1046, ("永承", "えいしょう")), - (1053, ("天喜", "てんぎ")), - (1058, ("康平", "こうへい")), - (1065, ("治暦", "じりゃく")), - (1069, ("延久", "えんきゅう")), - (1074, ("承保", "じょうほう")), - (1078, ("承暦", "じょうりゃく")), - (1081, ("永保", "えいほう")), - (1084, ("応徳", "おうとく")), - (1087, ("寛治", "かんじ")), - (1095, ("嘉保", "かほう")), - (1097, ("永長", "えいちょう")), - (1098, ("承徳", "じょうとく")), - (1099, ("康和", "こうわ")), - (1104, ("長治", "ちょうじ")), - (1106, ("嘉承", "かじょう")), - (1108, ("天仁", "てんにん")), - (1110, ("天永", "てんねい")), - (1113, ("永久", "えいきゅう")), - (1118, ("元永", "げんえい")), - (1120, ("保安", "ほうあん")), - (1124, ("天治", "てんじ")), - (1126, ("大治", "だいじ")), - (1131, ("天承", "てんしょう")), - (1132, ("長承", "ちょうしょう")), - (1135, ("保延", "ほうえん")), - (1141, ("永治", "えいじ")), - (1142, ("康治", "こうじ")), - (1144, ("天養", "てんよう")), - (1145, ("久安", "きゅうあん")), - (1151, ("仁平", "にんぺい")), - (1154, ("久寿", "きゅうじゅ")), - (1156, ("保元", "ほうげん")), - (1159, ("平治", "へいじ")), - (1160, ("永暦", "えいりゃく")), - (1161, ("応保", "おうほう")), - (1163, ("長寛", "ちょうかん")), - (1165, ("永万", "えいまん")), - (1166, ("仁安", "にんあん")), - (1169, ("嘉応", "かおう")), - (1171, ("承安", "しょうあん")), - (1175, ("安元", "あんげん")), - (1177, ("治承", "じしょう")), - (1181, ("養和", "ようわ")), - (1182, ("寿永", "じゅえい")), - (1184, ("元暦", "げんりゃく")), - (1185, ("文治", "ぶんじ")), - (1190, ("建久", "けんきゅう")), - (1199, ("正治", "しょうじ")), - (1201, ("建仁", "けんにん")), - (1204, ("元久", "げんきゅう")), - (1206, ("建永", "けんえい")), - (1207, ("承元", "じょうげん")), - (1211, ("建暦", "けんりゃく")), - (1214, ("建保", "けんぽう")), - (1219, ("承久", "じょうきゅう")), - (1222, ("貞応", "じょうおう")), - (1225, ("元仁", "げんにん")), - (1225, ("嘉禄", "かろく")), - (1228, ("安貞", "あんてい")), - (1229, ("寛喜", "かんき")), - (1232, ("貞永", "じょうえい")), - (1233, ("天福", "てんぷく")), - (1235, ("文暦", "ぶんりゃく")), - (1235, ("嘉禎", "かてい")), - (1239, ("暦仁", "りゃくにん")), - (1239, ("延応", "えんおう")), - (1240, ("仁治", "にんじ")), - (1243, ("寛元", "かんげん")), - (1247, ("宝治", "ほうじ")), - (1249, ("建長", "けんちょう")), - (1256, ("康元", "こうげん")), - (1257, ("正嘉", "しょうか")), - (1259, ("正元", "しょうげん")), - (1260, ("文応", "ぶんおう")), - (1261, ("弘長", "こうちょう")), - (1264, ("文永", "ぶんえい")), - (1275, ("健治", "けんじ")), - (1278, ("弘安", "こうあん")), - (1288, ("正応", "しょうおう")), - (1293, ("永仁", "えいにん")), - (1299, ("正安", "しょうあん")), - (1303, ("乾元", "けんげん")), - (1303, ("嘉元", "かげん")), - (1307, ("徳治", "とくじ")), - (1308, ("延慶", "えんきょう")), - (1311, ("応長", "おうちょう")), - (1312, ("正和", "しょうわ")), - (1317, ("文保", "ぶんぽう")), - (1319, ("元応", "げんおう")), - (1321, ("元亨", "げんこう")), - (1325, ("正中", "しょうちゅ")), - (1326, ("嘉暦", "かりゃく")), - (1329, ("元徳", "げんとく")), - (1331, ("元弘", "げんこう")), - (1332, ("正慶", "しょうけい")), - (1334, ("建武", "けんむ")), - (1336, ("延元", "えいげん")), - (1338, ("暦応", "りゃくおう")), - (1340, ("興国", "こうこく")), - (1342, ("康永", "こうえい")), - (1345, ("貞和", "じょうわ")), - (1347, ("正平", "しょうへい")), - (1350, ("観応", "かんおう")), - (1352, ("文和", "ぶんな")), - (1356, ("延文", "えんぶん")), - (1361, ("康安", "こうあん")), - (1362, ("貞治", "じょうじ")), - (1368, ("応安", "おうあん")), - (1370, ("建徳", "けんとく")), - (1372, ("文中", "ぶんちゅう")), - (1375, ("永和", "えいわ")), - (1375, ("天授", "てんじゅ")), - (1379, ("康暦", "こうりゃく")), - (1381, ("永徳", "えいとく")), - (1381, ("弘和", "こうわ")), - (1384, ("至徳", "しとく")), - (1384, ("元中", "げんちゅう")), - (1387, ("嘉慶", "かけい")), - (1389, ("康応", "こうおう")), - (1390, ("明徳", "めいとく")), - (1394, ("応永", "おうえい")), - (1428, ("正長", "しょうちょう")), - (1429, ("永享", "えいきょう")), - (1441, ("嘉吉", "かきつ")), - (1444, ("文安", "ぶんあん")), - (1449, ("宝徳", "ほうとく")), - (1452, ("享徳", "きょうとく")), - (1455, ("康正", "こうしょう")), - (1457, ("長禄", "ちょうろく")), - (1461, ("寛正", "かんしょう")), - (1466, ("文正", "ぶんしょう")), - (1467, ("応仁", "おうにん")), - (1469, ("文明", "ぶんめい")), - (1487, ("長享", "ちょうきょう")), - (1489, ("延徳", "えんとく")), - (1492, ("明応", "めいおう")), - (1501, ("文亀", "ぶんき")), - (1504, ("永正", "えいしょう")), - (1521, ("大永", "だいえい")), - (1528, ("享禄", "きょうろく")), - (1532, ("天文", "てんぶん")), - (1555, ("弘治", "こうじ")), - (1558, ("永禄", "えいろく")), - (1570, ("元亀", "げんき")), - (1573, ("天正", "てんしょう")), - (1593, ("文禄", "ぶんろく")), - (1596, ("慶長", "けいちょう")), - (1615, ("元和", "げんな")), - (1624, ("寛永", "かんえい")), - (1645, ("正保", "しょうほう")), - (1648, ("慶安", "けいあん")), - (1652, ("承応", "じょうおう")), - (1655, ("明暦", "めいれき")), - (1658, ("万治", "まんじ")), - (1661, ("寛文", "かんぶん")), - (1673, ("延宝", "えんぽう")), - (1681, ("天和", "てんな")), - (1684, ("貞享", "じょうきょう")), - (1688, ("元禄", "げんろく")), - (1704, ("宝永", "ほうえい")), - (1711, ("正徳", "しょうとく")), - (1716, ("享保", "きょうほう")), - (1736, ("元文", "げんぶん")), - (1741, ("寛保", "かんぽう")), - (1744, ("延享", "えんきょう")), - (1748, ("寛延", "かんえん")), - (1751, ("宝暦", "ほうれき")), - (1764, ("明和", "めいわ")), - (1773, ("安永", "あんえい")), - (1781, ("天明", "てんめい")), - (1801, ("寛政", "かんせい")), - (1802, ("享和", "きょうわ")), - (1804, ("文化", "ぶんか")), - (1818, ("文政", "ぶんせい")), - (1831, ("天保", "てんぽう")), - (1845, ("弘化", "こうか")), - (1848, ("嘉永", "かえい")), - (1855, ("安政", "あんせい")), - (1860, ("万延", "まんえい")), - (1861, ("文久", "ぶんきゅう")), - (1864, ("元治", "げんじ")), - (1865, ("慶応", "けいおう")), - (1868, ("明治", "めいじ")), - (1912, ("大正", "たいしょう")), - (1926, ("昭和", "しょうわ")), - (1989, ("平成", "へいせい")), -] - - -class Num2Word_JA(Num2Word_Base): - CURRENCY_FORMS = { - 'JPY': (('円', 'えん'), ()), - } - - def set_high_numwords(self, high): - max = 4 * len(high) - for word, n in zip(high, range(max, 0, -4)): - self.cards[10 ** n] = word - - def setup(self): - self.negword = "マイナス" - self.pointword = ("点", "てん") - self.exclude_title = ["点", "マイナス"] - - self.high_numwords = [ - ("万", "まん"), # 10**4 man - ("億", "おく"), # 10**8 oku - ("兆", "ちょう"), # 10**12 chō - ("京", "けい"), # 10**16 kei - ("垓", "がい"), # 10**20 gai - ("秭", "し"), # 10**24 shi - ("穣", "じょう"), # 10**28 jō - ("溝", "こう"), # 10**32 kō - ("澗", "かん"), # 10**36 kan - ("正", "せい"), # 10**40 sei - ("載", "さい"), # 10**44 sai - ("極", "ごく"), # 10**48 goku - ] - - self.high_numwords.reverse() - - self.mid_numwords = [ - (1000, ("千", "せん")), - (100, ("百", "ひゃく")), - ] - - self.low_numwords = [ - ("十", "じゅう"), # 10 jū - ("九", "きゅう"), # 9 kyū - ("八", "はち"), # 8 hachi - ("七", ("なな", "しち")), # 7 nana, shichi - ("六", "ろく"), # 6 roku - ("五", "ご"), # 5 go - ("四", ("よん", "し")), # 4 yon, shi - ("三", "さん"), # 3 san - ("二", "に"), # 2 ni - ("一", "いち"), # 1 ichi - # both are alternatives, 零 doesn't map to ゼロ or 〇 to れい - (("零", "〇"), ("ゼロ", "れい")), # 0 ZERO, rei - ] - - def merge(self, lpair, rpair): - ltext, lnum = lpair - rtext, rnum = rpair - - fmt = "%s%s" - # ignore lpair if lnum is 1 and rnum is less than 10000 - if lnum == 1 and rnum < 10000: - return rpair - # rnum is added to lnum - elif lnum > rnum: - return (fmt % (ltext, rtext), lnum + rnum) - # rnum is multiplied by lnum - elif lnum < rnum: - return rendaku_merge_pairs(lpair, rpair) - - def _ordinal_suffix(self, reading, counter): - if reading: - if counter == "番": - return "ばんめ" - else: - raise NotImplementedError( - "Reading not implemented for %s" % counter) - else: - return counter + "目" - - def to_ordinal(self, value, reading=False, prefer=None, counter="番"): - self.verify_ordinal(value) - base = self.to_cardinal(value, reading=reading, prefer=prefer) - return "%s%s" % (base, self._ordinal_suffix(reading, counter)) - - def to_ordinal_num(self, value, reading=False, counter="番"): - return "%s%s" % (value, self._ordinal_suffix(reading, counter)) - - def to_year(self, val, suffix=None, longval=True, reading=False, - prefer=None, era=True): - year = val - # Gregorian calendar - if not era: - prefix = "" - if year < 0: - year = abs(year) - prefix = "きげんぜん" if reading else "紀元前" - - year_words = self.to_cardinal(year, reading=reading, prefer=prefer) - if reading and year % 10 == 9: - year_words = year_words[:-3] + "く" - - return "%s%s%s" % (prefix, year_words, "ねん" if reading else "年") - - # Era calendar (default) - min_year = ERA_START[0][0] - last_era_idx = len(ERA_START) - 1 - if year < min_year: - raise ValueError( - "Can't convert years less than %s to era" % min_year) - - first = 0 - last = last_era_idx - era_idx = None - while era_idx is None: - mid = (first + last) // 2 - if mid == last_era_idx or (ERA_START[mid][0] <= year and - ERA_START[mid + 1][0] > year): - era_idx = mid - # if an era lasting less than a year is preferred, choose it - if prefer: - i = mid - 1 - while i >= 0 and ERA_START[i][0] == year: - # match kanji or hiragana - if set(ERA_START[i][1]) & set(prefer): - era_idx = i - break - i -= 1 - - # ends up at the last index where year >= ERA_START[mid][0] - if year < ERA_START[mid][0]: - last = mid - 1 - else: - first = mid + 1 - - era = ERA_START[era_idx] - era_name = era[1][0] - era_year = year - era[0] + 1 - fmt = "%s%s年" - if reading == "arabic": - era_year_words = str(era_year) - elif reading: - era_name = era[1][1] - era_year_words = (self.to_cardinal(era_year, reading=True, - prefer=prefer) - if era_year != 1 else "がん") - if era_year % 10 == 9: - era_year_words = era_year_words[:-3] + "く" - fmt = "%s%sねん" - else: - era_year_words = (self.to_cardinal(era_year, reading=False, - prefer=prefer) - if era_year != 1 else "元") - - return fmt % (era_name, era_year_words) - - def to_currency(self, val, currency="JPY", cents=False, separator="", - adjective=False, reading=False, prefer=None): - left, right, is_negative = parse_currency_parts( - val, is_int_with_cents=cents) - - try: - cr1, cr2 = self.CURRENCY_FORMS[currency] - if (cents or abs(val) != left) and not cr2: - raise ValueError('Decimals not supported for "%s"' % currency) - except KeyError: - raise NotImplementedError( - 'Currency code "%s" not implemented for "%s"' % - (currency, self.__class__.__name__)) - - if adjective and currency in self.CURRENCY_ADJECTIVES: - cr1 = prefix_currency(self.CURRENCY_ADJECTIVES[currency], cr1) - - minus_str = self.negword if is_negative else "" - - return '%s%s%s%s%s' % ( - minus_str, - self.to_cardinal(left, reading=reading, prefer=prefer), - cr1[1] if reading else cr1[0], - self.to_cardinal(right, reading=reading, prefer=prefer) - if cr2 else '', - (cr2[1] if reading else cr2[0]) if cr2 else '', - ) - - def splitnum(self, value, reading, prefer): - for elem in self.cards: - if elem > value: - continue - - out = [] - if value == 0: - div, mod = 1, 0 - else: - div, mod = divmod(value, elem) - - if div == 1: - out.append((select_text(self.cards[1], reading, prefer), 1)) - else: - if div == value: # The system tallies, eg Roman Numerals - return [( - div * select_text(self.cards[elem], reading, prefer), - div * elem)] - out.append(self.splitnum(div, reading, prefer)) - - out.append((select_text(self.cards[elem], reading, prefer), elem)) - - if mod: - out.append(self.splitnum(mod, reading, prefer)) - - return out - - def to_cardinal(self, value, reading=False, prefer=None): - try: - assert int(value) == value - except (ValueError, TypeError, AssertionError): - return self.to_cardinal_float(value, reading=reading, - prefer=prefer) - - out = "" - if value < 0: - value = abs(value) - out = self.negword - - if value >= self.MAXVAL: - raise OverflowError(self.errmsg_toobig % (value, self.MAXVAL)) - - val = self.splitnum(value, reading, prefer) - words, _ = self.clean(val) - return self.title(out + words) - - def to_cardinal_float(self, value, reading=False, prefer=None): - prefer = prefer or ["れい"] - try: - float(value) == value - except (ValueError, TypeError, AssertionError): - raise TypeError(self.errmsg_nonnum % value) - - pre, post = self.float2tuple(float(value)) - - post = str(post) - post = '0' * (self.precision - len(post)) + post - - out = [self.to_cardinal(pre, reading=reading, prefer=prefer)] - if self.precision: - out.append(self.title(self.pointword[1 if reading else 0])) - - for i in range(self.precision): - curr = int(post[i]) - out.append(to_s( - self.to_cardinal(curr, reading=reading, prefer=prefer))) - - return "".join(out) diff --git a/fun_text_processing/num2words/num2words/lang_KH.py b/fun_text_processing/num2words/num2words/lang_KH.py deleted file mode 100644 index d7345dc3b..000000000 --- a/fun_text_processing/num2words/num2words/lang_KH.py +++ /dev/null @@ -1,136 +0,0 @@ -# -*- coding: utf-8 -*- -# Khmer - -from __future__ import division, print_function, unicode_literals - -from . import lang_EU - -# Khmer - -class Num2Word_KH(lang_EU.Num2Word_EU): - def set_high_numwords(self, high): - max = 3 * len(high) - for word, n in zip(high, range(max, 3, -1)): - # print(word[0],word[1],n) - self.cards[10 ** n] = word[1] - # self.cards[10**n] = word + "លាន" - # try: - # ordinal_word = self.high_numwords[high] - # except KeyError: - # max = 3 + 3 * len(high) - # for word, n in zip(high, range(max, 3, -3)): - # print(word) - # print(n) - # self.cards[10 ** n] = word - - def gen_high_numwords(self, units, tens, lows): - out = [u + t for t in tens for u in units] - out.reverse() - return out + lows - - def setup(self): - super(Num2Word_KH, self).setup() - - self.negword = "ដក " - self.pointword = "ចំណុច" - self.exclude_title = ["និង", "ចំណុច", "ដក"] - lows = [""] - units = [""] - tens = ["ពាន់លាន","រយលាន","កោដិ","លាន","សែន","ម៉ឺន"] - self.high_numwords = self.gen_high_numwords(units, tens, lows) - - self.high_numwords = [(1000000000, "មួយពាន់លាន"), - (100000000, "មួយរយលាន"), - (10000000, "មួយកោដិ"), - (1000000, "មួយលាន"), - (100000, "មួយសែន"), - (10000, "មួយម៉ឺន")] - # (2400000, "ពីរលានបួនម៉ឺន") - - self.mid_numwords = [(1000, "មួយពាន់"), (100, "មួយរយ"), - (90, "កៅសិប"), (80, "ប៉ែតសិប"), (70, "ចិតសិប"), - (60, "ហុកសិប"), (50, "ហាសិប"), (40, "សែសិប"), - (30, "សាមសិប")] - - self.low_numwords = ["ម្ភៃ", - "ដប់ប្រាំបួន", - "ដប់ប្រាំបី", - "ដប់ប្រាំពីរ", - "ដប់ប្រាំមួយ", - "ដប់ប្រាំ", - "ដប់បួន", - "ដប់បី", - "ដប់ពីរ", - "ដប់មួយ", - "ដប់", - "ប្រាំបួន", - "ប្រាំបី", - "ប្រាំពីរ", - "ប្រាំមួយ", - "ប្រាំ", - "បួន", - "បី", - "ពីរ", - "មួយ", - "សូន្យ"] - - self.ords = {"មួយ": "ទីមួយ", - "ពីរ": "ទីពីរ", - "បី": "ទីបី", - "បួន": "ទីបួន", - "ប្រាំ": "ទីប្រាំ", - "ប្រាំមួយ": "ទីប្រាំមួយ", - "ប្រាំពីរ": "ទីប្រាំពីរ", - "ប្រាំបី": "ទីប្រាំបី", - "ប្រាំបួន": "ទីប្រាំបួន", - "ដប់": "ទីដប់", - "ដប់មួយ": "ទីដប់មួយ", - "ដប់ពីរ": "ទីដប់ពីរ"} - - def merge(self, lpair, rpair): - ltext, lnum = lpair - rtext, rnum = rpair - if lnum == 1 and rnum < 100: - return (rtext, rnum) - elif 100 > lnum > rnum: - return ("%s%s" % (ltext, rtext), lnum + rnum) - elif lnum >= 100 > rnum: - return ("%s%s" % (ltext, rtext), lnum + rnum) - elif rnum > lnum: - return ("%s%s" % (ltext, rtext), lnum * rnum) - return ("%s%s" % (ltext, rtext), lnum + rnum) - - def to_ordinal_num(self, value): - self.verify_ordinal(value) - return "%s%s" % (value, self.to_ordinal(value)) - - def to_ordinal(self, value): - self.verify_ordinal(value) - try: - ordinal_word = self.ords[value] - except KeyError: - ordinal_word = "ទី" + value - return ordinal_word - - - def to_year(self, val, suffix=None, longval=True): - if val < 0: - val = abs(val) - suffix = 'មុនគ' if not suffix else suffix - high, low = (val // 100, val % 100) - # If year is 00XX, X00X, or beyond 9999, go cardinal. - if (high == 0 - or (high % 10 == 0 and low < 10) - or high >= 100): - valtext = self.to_cardinal(val) - else: - hightext = self.to_cardinal(high) - if low == 0: - lowtext = "រយ" - elif low < 10: - lowtext = "%s" % self.to_cardinal(low) - else: - lowtext = self.to_cardinal(low) - valtext = "%s %s" % (hightext, lowtext) - return (valtext if not suffix - else "%s %s" % (valtext, suffix)) diff --git a/fun_text_processing/num2words/num2words/lang_KN.py b/fun_text_processing/num2words/num2words/lang_KN.py deleted file mode 100644 index 2f89872fa..000000000 --- a/fun_text_processing/num2words/num2words/lang_KN.py +++ /dev/null @@ -1,167 +0,0 @@ -# -*- encoding: utf-8 -*- - -from __future__ import unicode_literals - -from .lang_EU import Num2Word_EU - - -class Num2Word_KN(Num2Word_EU): - def set_high_numwords(self, high): - for n, word in self.high_numwords: - self.cards[10**n] = word - - def setup(self): - self.low_numwords = [ - "ತೊಂಬತ್ತೊಂಬತ್ತು", - "ತೊಂಬತ್ತೆಂಟು", - "ತೊಂಬತ್ತೇಳು", - "ತೊಂಬತ್ತಾರು", - "ತೊಂಬತ್ತೈದು", - "ತೊಂಬತ್ತ ನಾಲ್ಕು", - "ತೊಂಬತ್ತ ಮೂರು", - "ತೊಂಬತ್ತೆರಡು", - "ತೊಂಬತ್ತೊಂದು", - "ತೊಂಬತ್ತು", - "ಎಂಬತ್ತೊಂಬತ್ತು", - "ಎಂಬತ್ತೆಂಟು", - "ಎಂಬತ್ತೇಳು", - "ಎಂಬತ್ತಾರು", - "ಎಂಬತ್ತೈದು", - "ಎಂಬತ್ತ್ ನಾಲ್ಕು", - "ಎಂಬತ್ತ್ ಮೂರು", - "ಎಂಬತ್ತೆರಡು", - "ಎಂಬತ್ತೊಂದು", - "ಎಂಬತ್ತು", - "ಎಪ್ಪತ್ತೊಂಬತ್ತು", - "ಎಪ್ಪತ್ತೆಂಟು", - "ಎಪ್ಪತ್ತೇಳು", - "ಎಪ್ಪತ್ತಾರು", - "ಎಪ್ಪತ್ತೈದು", - "ಎಪ್ಪತ್ತ್ ನಾಲ್ಕು", - "ಎಪ್ಪತ್ತ್ ಮೂರು", - "ಎಪ್ಪತ್ತೆರಡು", - "ಎಪ್ಪತ್ತೊಂದು", - "ಎಪ್ಪತ್ತು", - "ಅರವತ್ತೊಂಬತ್ತು", - "ಅರವತ್ತೆಂಟು", - "ಅರವತ್ತೇಳು", - "ಅರವತ್ತಾರು", - "ಅರವತ್ತೈದು", - "ಅರವತ್ತ್ ನಾಲ್ಕು", - "ಅರವತ್ತ್ ಮೂರು", - "ಅರವತ್ತೆರಡು", - "ಅರವತ್ತೊಂದು", - "ಅರವತ್ತು", - "ಐವತ್ತೊಂಬತ್ತು", - "ಐವತ್ತೆಂಟು", - "ಐವತ್ತೇಳು", - "ಐವತ್ತಾರು", - "ಐವತ್ತೈದು", - "ಐವತ್ತ್ನಾಲ್ಕು", - "ಐವತ್ತಮೂರು", - "ಐವತ್ತೆರಡು", - "ಐವತ್ತೊಂದು", - "ಐವತ್ತು", - "ನಲವತ್ತೊಂಬತ್ತು", - "ನಲವತ್ತೆಂಟು", - "ನಲವತ್ತೇಳು", - "ನಲವತ್ತಾರು", - "ನಲವತ್ತೈದು", - "ನಲವತ್ತ್ ನಾಲ್ಕು", - "ನಲವತ್ತ್ ಮೂರು", - "ನಲವತ್ತ್ ಎರಡು", - "ನಲವತ್ತೊಂದು", - "ನಲವತ್ತು", - "ಮೂವತ್ತ್ ಒಂಬತ್ತು", - "ಮೂವತ್ಎಂಟು", - "ಮೂವತ್ಏಳು", - "ಮೂವತ್ಆರು", - "ಮೂವತ್ತ್ ಐದು", - "ಮೂವತ್ತ್ ನಾಲ್ಕು", - "ಮೂವತ್ತ್ ಮೂರು", - "ಮೂವತ್ತ್ಎರಡು", - "ಮೂವತ್ತ್ಒಂದು", - "ಮೂವತ್ತು", - "ಇಪ್ಪತ್ತ್ಒಂಬತ್ತು", - "ಇಪ್ಪತ್ತ್ಎಂಟು", - "ಇಪ್ಪತ್ತ್ಏಳು", - "ಇಪ್ಪತ್ತ್ಆರು", - "ಇಪ್ಪತ್ತ್ ಐದು", - "ಇಪ್ಪತ್ತ್ ನಾಲ್ಕು", - "ಇಪ್ಪತ್ತ್ ಮೂರು", - "ಇಪ್ಪತ್ತ್ ಎರಡು", - "ಇಪ್ಪತ್ತ್ ಒಂದು", - "ಇಪ್ಪತ್ತು", - "ಹತ್ತೊಂಬತ್ತು", - "ಹದಿನೆಂಟು", - "ಹದಿನೇಳು", - "ಹದಿನಾರು", - "ಹದಿನೈದು", - "ಹದಿನಾಲ್ಕು", - "ಹದಿಮೂರು", - "ಹನ್ನೆರಡು", - "ಹನ್ನೊಂದು", - "ಹತ್ತು", - "ಒಂಬತ್ತು", - "ಎಂಟು", - "ಏಳು", - "ಆರು", - "ಐದು", - "ನಾಲ್ಕು", - "ಮೂರು", - "ಎರಡು", - "ಒಂದು", - "ಸೊನ್ನೆ", - ] - - self.mid_numwords = [(100, "ನೂರು")] - - self.high_numwords = [(7, "ಕೋಟಿ"), (5, "ಒಂದು ಲಕ್ಷ"), (3, "ಸಾವಿರ")] - - self.pointword = "ಬಿಂದು" - - self.modifiers = [ - "್", - "ಾ", - "ಿ", - "ೀ", - "ೀ", - "ು", - "ೂ", - "ೃ", - "ೆ", - "ೇ", - "ೈ", - "ೊ", - "ೋ", - "ೌ", - "ೕ", - ] - - def merge(self, lpair, rpair): - ltext, lnum = lpair - rtext, rnum = rpair - if lnum == 1 and rnum < 100: - return (rtext, rnum) - elif 100 > lnum > rnum: - return ("%s-%s" % (ltext, rtext), lnum + rnum) - elif lnum >= 100 > rnum: - if ltext[-1] in self.modifiers: - return ("%s %s" % (ltext[:-1], rtext), lnum + rnum) - else: - return ("%s %s" % (ltext + "ದ", rtext), lnum + rnum) - elif rnum > lnum: - return ("%s %s" % (ltext, rtext), lnum * rnum) - return ("%s %s" % (ltext, rtext), lnum + rnum) - - def to_ordinal_num(self, value): - self.verify_ordinal(value) - return "%s%s" % (value, self.to_ordinal(value)) - - def to_ordinal(self, value): - self.verify_ordinal(value) - outwords = self.to_cardinal(value) - if outwords[-1] in self.modifiers: - outwords = outwords[:-1] - ordinal_num = outwords + "ನೇ" - return ordinal_num diff --git a/fun_text_processing/num2words/num2words/lang_KO.py b/fun_text_processing/num2words/num2words/lang_KO.py deleted file mode 100644 index eeafaa2ed..000000000 --- a/fun_text_processing/num2words/num2words/lang_KO.py +++ /dev/null @@ -1,133 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import division, print_function, unicode_literals - -from .base import Num2Word_Base -from .currency import parse_currency_parts - - -class Num2Word_KO(Num2Word_Base): - CURRENCY_FORMS = { - 'KRW': ('원', None), - 'USD': ('달러', '센트'), - 'JPY': ('엔', None) - } - - def set_high_numwords(self, high): - max = 4 * len(high) - for word, n in zip(high, range(max, 0, -4)): - self.cards[10 ** n] = word - - def setup(self): - super(Num2Word_KO, self).setup() - - self.negword = "마이너스 " - self.pointword = "점" - - self.high_numwords = [ - '무량대수', - '불가사의', - '나유타', - '아승기', - '항하사', - '극', - '재', - '정', - '간', - '구', - '양', - '자', - '해', - '경', - '조', - '억', - '만'] - self.mid_numwords = [(1000, "천"), (100, "백")] - self.low_numwords = ["십", "구", "팔", "칠", "육", "오", "사", "삼", "이", - "일", "영"] - self.ords = {"일": "한", - "이": "두", - "삼": "세", - "사": "네", - "오": "다섯", - "육": "여섯", - "칠": "일곱", - "팔": "여덟", - "구": "아홉", - "십": "열", - "이십": "스물", - "삼십": "서른", - "사십": "마흔", - "오십": "쉰", - "육십": "예순", - "칠십": "일흔", - "팔십": "여든", - "구십": "아흔"} - - def merge(self, lpair, rpair): - ltext, lnum = lpair - rtext, rnum = rpair - if lnum == 1 and rnum <= 10000: - return rpair - elif 10000 > lnum > rnum: - return ("%s%s" % (ltext, rtext), lnum + rnum) - elif lnum >= 10000 and lnum > rnum: - return ("%s %s" % (ltext, rtext), lnum + rnum) - else: - return ("%s%s" % (ltext, rtext), lnum * rnum) - - def to_ordinal(self, value): - self.verify_ordinal(value) - if(value == 1): - return "첫 번째" - outwords = self.to_cardinal(value).split(" ") - lastwords = outwords[-1].split("백") - if "십" in lastwords[-1]: - ten_one = lastwords[-1].split("십") - ten_one[0] = self.ords[ten_one[0] + "십"] - try: - ten_one[1] = self.ords[ten_one[1]] - ten_one[0] = ten_one[0].replace("스무", "스물") - except KeyError: - pass - lastwords[-1] = ''.join(ten_one) - else: - lastwords[-1] = self.ords[lastwords[-1]] - outwords[-1] = "백 ".join(lastwords) - return " ".join(outwords) + " 번째" - - def to_ordinal_num(self, value): - self.verify_ordinal(value) - return "%s 번째" % (value) - - def to_year(self, val, suffix=None, longval=True): - if val < 0: - val = abs(val) - suffix = '기원전' if not suffix else suffix - valtext = self.to_cardinal(val) - return ("%s년" % valtext if not suffix - else "%s %s년" % (suffix, valtext)) - - def to_currency(self, val, currency="KRW", cents=False, separator="", - adjective=False): - left, right, is_negative = parse_currency_parts( - val, is_int_with_cents=cents) - - try: - cr1, cr2 = self.CURRENCY_FORMS[currency] - if (cents or right) and not cr2: - raise ValueError('Decimals not supported for "%s"' % currency) - except KeyError: - raise NotImplementedError( - 'Currency code "%s" not implemented for "%s"' % - (currency, self.__class__.__name__)) - - minus_str = self.negword if is_negative else "" - return '%s%s%s%s%s' % ( - minus_str, - ''.join(self.to_cardinal(left).split()), - cr1, - ' ' + self.to_cardinal(right) - if cr2 else '', - cr2 if cr2 else '', - ) diff --git a/fun_text_processing/num2words/num2words/lang_KZ.py b/fun_text_processing/num2words/num2words/lang_KZ.py deleted file mode 100644 index 18f4a3b42..000000000 --- a/fun_text_processing/num2words/num2words/lang_KZ.py +++ /dev/null @@ -1,118 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import unicode_literals - -from .base import Num2Word_Base -from .utils import get_digits, splitbyx - -ZERO = 'нөл' - -ONES = { - 1: 'бір', - 2: 'екі', - 3: 'үш', - 4: 'төрт', - 5: 'бес', - 6: 'алты', - 7: 'жеті', - 8: 'сегіз', - 9: 'тоғыз', -} - -TEN = 'он' - -TWENTIES = { - 2: 'жиырма', - 3: 'отыз', - 4: 'қырық', - 5: 'елу', - 6: 'алпыс', - 7: 'жетпіс', - 8: 'сексен', - 9: 'тоқсан', -} - -HUNDRED = 'жүз' - -THOUSANDS = { - 1: 'мың', - 2: 'миллион', - 3: 'миллиард', - 4: 'триллион', - 5: 'квадриллион', - 6: 'квинтиллион', - 7: 'секстиллион', - 8: 'септиллион', - 9: 'октиллион', - 10: 'нониллион', -} - - -class Num2Word_KZ(Num2Word_Base): - CURRENCY_FORMS = { - 'USD': ('доллар', 'цент'), - 'KZT': ('теңге', 'тиын'), - } - - def setup(self): - self.negword = "минус" - self.pointword = "бүтін" - - def to_cardinal(self, number): - n = str(number).replace(',', '.') - if '.' in n: - left, right = n.split('.') - leading_zero_count = len(right) - len(right.lstrip('0')) - return u'%s %s %s' % ( - self._int2word(int(left)), - self.pointword, - (ZERO + ' ') * leading_zero_count + self._int2word(int(right)) - ) - else: - return self._int2word(int(n)) - - def pluralize(self, n, form): - return form - - def _cents_verbose(self, number, currency): - return self._int2word(number, currency == 'KZT') - - def _int2word(self, n, feminine=False): - if n < 0: - return ' '.join([self.negword, self._int2word(abs(n))]) - - if n == 0: - return ZERO - - words = [] - chunks = list(splitbyx(str(n), 3)) - i = len(chunks) - for x in chunks: - i -= 1 - - if x == 0: - continue - - n1, n2, n3 = get_digits(x) - - if n3 > 0: - if n3 > 1: - words.append(ONES[n3]) - words.append(HUNDRED) - - if n2 == 1: - words.append(TEN) - elif n2 > 1: - words.append(TWENTIES[n2]) - - if n1 > 0: - words.append(ONES[n1]) - - if i > 0: - words.append(THOUSANDS[i]) - - return ' '.join(words) - - def to_ordinal(self, number): - # TODO: Implement to_ordinal - raise NotImplementedError() diff --git a/fun_text_processing/num2words/num2words/lang_LO.py b/fun_text_processing/num2words/num2words/lang_LO.py deleted file mode 100644 index a85c4735b..000000000 --- a/fun_text_processing/num2words/num2words/lang_LO.py +++ /dev/null @@ -1,110 +0,0 @@ -# -*- coding: utf-8 -*- - -# Lao number to words - -from __future__ import division, print_function, unicode_literals - -from . import lang_EU - - -class Num2Word_LO(lang_EU.Num2Word_EU): - # GIGA_SUFFIX = "iljarder" - # MEGA_SUFFIX = "iljoner" - - # def set_high_numwords(self, high): - # cap = 3 + 6 * len(high) - - # for word, n in zip(high, range(cap, 3, -6)): - # if self.GIGA_SUFFIX: - # self.cards[10 ** n] = word + self.GIGA_SUFFIX - - # if self.MEGA_SUFFIX: - # self.cards[10 ** (n - 3)] = word + self.MEGA_SUFFIX - def set_high_numwords(self, high): - max = 3 * len(high) - for word, n in zip(high, range(max, 0, -3)): - # print(word[0],word[1],n) - self.cards[10 ** n] = word[1] - - def setup(self): - super(Num2Word_LO, self).setup() - - self.negword = "ລົບ" - self.pointword = "ຈຸດ" - self.exclude_title = ["ແລະ", "ຈຸດ", "ລົບ"] - - self.high_numwords = [(10000000000, 'ສິບຕື້'), - (1000000000, 'ຫນຶ່ງຕື້'), - (100000000, 'ຫນຶ່ງຮ້ອຍລ້ານ'), - (10000000,'ສິບລ້ານ'), - (1000000, 'ຫນຶ່ງລ້ານ'), - (100000, 'ຫນຶ່ງແສນ'), - (1000, 'ຫນຶ່ງພັນ') ] - self.mid_numwords = [(1000, "ຫນຶ່ງພັນ"), (100, "ຫນຶ່ງຮ້ອຍ"), - (90, "ເກົ້າສິບ"), (80, "ແປດສິບ"), (70, "ເຈັດສິບ"), - (60, "ຫົກສິບ"), (50, "ຫ້າສິບ"), (40, "ສີ່ສິບ"), - (30, "ສາມສິບ")] - - self.low_numwords = ['ຊາວ', 'ສິບເກົ້າ', 'ສິບແປດ', 'ສິບເຈັດ', 'ສິບຫົກ', 'ສິບຫ້າ', 'ສິບສີ່', 'ສິບສາມ', 'ສິບສອງ', 'ສິບເອັດ', 'ສິບc', 'ເກົ້າ', 'ແປດ', 'ເຈັດ', 'ຫົກ', 'ຫ້າ', 'ສີ່', 'ສາມ', 'ສອງ', 'ຫນຶ່ງ', 'ສູນ'] - - - self.ords = {"ຫນຶ່ງ": "ທໍາອິດ", - "ສອງ": "ທີສອງ", - "ສາມ": "ທີສາມ", - "ສີ່": "ທີສີ່", - "ຫ້າ": "ທີຫ້າ", - "ຫົກ":"ທີຫົກ", - "ເຈັດ": "ທີເຈັດ", - "ແປດ": "ທີແປດ", - "ເກົ້າ": "ເກົ້າ", - "ສິບc": "ທີສິບ" - } - - - def merge(self, lpair, rpair): - ltext, lnum = lpair - rtext, rnum = rpair - if lnum == 1 and rnum < 100: - return (rtext, rnum) - elif 100 > lnum > rnum: - return ("%s%s" % (ltext, rtext), lnum + rnum) - elif lnum >= 100 > rnum: - return ("%s%s" % (ltext, rtext), lnum + rnum) - elif rnum > lnum: - return ("%s%s" % (ltext, rtext), lnum * rnum) - return ("%s%s" % (ltext, rtext), lnum + rnum) - - def to_ordinal_num(self, value): - self.verify_ordinal(value) - return "%s %s" % (value, self.to_ordinal(value)) - - def to_ordinal(self, value): - self.verify_ordinal(value) - try: - ordinal_word = self.ords[value] - except KeyError: - ordinal_word = value #TODO: check nepali ordinal word suffix - return ordinal_word - - - def to_year(self, val, suffix=None, longval=True): - if val < 0: - val = abs(val) - suffix = 'ກ່ອນສະຕະວັດ ' if not suffix else suffix - high, low = (val // 100, val % 100) - # If year is 00XX, X00X, or beyond 9999, go cardinal. - if (high == 0 - or (high % 10 == 0 and low < 10) - or high >= 100): - valtext = self.to_cardinal(val) - else: - hightext = self.to_cardinal(high) - if low == 0: - lowtext = "ຮ້ອຍ" - elif low < 10: - lowtext = "%s" % self.to_cardinal(low) - else: - lowtext = self.to_cardinal(low) - valtext = "%s %s" % (hightext, lowtext) - return (valtext if not suffix - else "%s %s" % (valtext, suffix)) diff --git a/fun_text_processing/num2words/num2words/lang_LT.py b/fun_text_processing/num2words/num2words/lang_LT.py deleted file mode 100644 index d04dedda4..000000000 --- a/fun_text_processing/num2words/num2words/lang_LT.py +++ /dev/null @@ -1,167 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import unicode_literals - -from .base import Num2Word_Base -from .utils import get_digits, splitbyx - -ZERO = ('nulis',) - -ONES_FEMININE = { - 1: ('viena',), - 2: ('dvi',), - 3: ('trys',), - 4: ('keturios',), - 5: ('penkios',), - 6: ('šešios',), - 7: ('septynios',), - 8: ('aštuonios',), - 9: ('devynios',), -} - -ONES = { - 1: ('vienas',), - 2: ('du',), - 3: ('trys',), - 4: ('keturi',), - 5: ('penki',), - 6: ('šeši',), - 7: ('septyni',), - 8: ('aštuoni',), - 9: ('devyni',), -} - -TENS = { - 0: ('dešimt',), - 1: ('vienuolika',), - 2: ('dvylika',), - 3: ('trylika',), - 4: ('keturiolika',), - 5: ('penkiolika',), - 6: ('šešiolika',), - 7: ('septyniolika',), - 8: ('aštuoniolika',), - 9: ('devyniolika',), -} - -TWENTIES = { - 2: ('dvidešimt',), - 3: ('trisdešimt',), - 4: ('keturiasdešimt',), - 5: ('penkiasdešimt',), - 6: ('šešiasdešimt',), - 7: ('septyniasdešimt',), - 8: ('aštuoniasdešimt',), - 9: ('devyniasdešimt',), -} - -HUNDRED = ('šimtas', 'šimtai') - -THOUSANDS = { - 1: ('tūkstantis', 'tūkstančiai', 'tūkstančių'), - 2: ('milijonas', 'milijonai', 'milijonų'), - 3: ('milijardas', 'milijardai', 'milijardų'), - 4: ('trilijonas', 'trilijonai', 'trilijonų'), - 5: ('kvadrilijonas', 'kvadrilijonai', 'kvadrilijonų'), - 6: ('kvintilijonas', 'kvintilijonai', 'kvintilijonų'), - 7: ('sikstilijonas', 'sikstilijonai', 'sikstilijonų'), - 8: ('septilijonas', 'septilijonai', 'septilijonų'), - 9: ('oktilijonas', 'oktilijonai', 'oktilijonų'), - 10: ('naintilijonas', 'naintilijonai', 'naintilijonų'), -} - -GENERIC_CENTS = ('centas', 'centai', 'centų') - - -class Num2Word_LT(Num2Word_Base): - CURRENCY_FORMS = { - 'LTL': (('litas', 'litai', 'litų'), GENERIC_CENTS), - 'EUR': (('euras', 'eurai', 'eurų'), GENERIC_CENTS), - 'USD': (('doleris', 'doleriai', 'dolerių'), GENERIC_CENTS), - 'GBP': ( - ('svaras sterlingų', 'svarai sterlingų', 'svarų sterlingų'), - ('pensas', 'pensai', 'pensų') - ), - 'PLN': ( - ('zlotas', 'zlotai', 'zlotų'), - ('grašis', 'grašiai', 'grašių')), - 'RUB': ( - ('rublis', 'rubliai', 'rublių'), - ('kapeika', 'kapeikos', 'kapeikų') - ), - } - - def setup(self): - self.negword = "minus" - self.pointword = "kablelis" - - def pluralize(self, n, forms): - n1, n2, n3 = get_digits(n) - if n2 == 1 or n1 == 0 or n == 0: - return forms[2] - elif n1 == 1: - return forms[0] - else: - return forms[1] - - def to_cardinal(self, number): - n = str(number).replace(',', '.') - base_str, n = self.parse_minus(n) - if '.' in n: - left, right = n.split('.') - leading_zero_count = len(right) - len(right.lstrip('0')) - decimal_part = ((ZERO[0] + ' ') * leading_zero_count + - self._int2word(int(right))) - return '%s%s %s %s' % ( - base_str, - self._int2word(int(left)), - self.pointword, - decimal_part - ) - else: - return "%s%s" % (base_str, self._int2word(int(n))) - - def to_ordinal(self, number): - raise NotImplementedError() - - def _cents_verbose(self, number, currency): - return self._int2word(number, currency == 'RUB') - - def _int2word(self, n, feminine=False): - if n == 0: - return ZERO[0] - - words = [] - chunks = list(splitbyx(str(n), 3)) - i = len(chunks) - - for x in chunks: - i -= 1 - - if x == 0: - continue - - n1, n2, n3 = get_digits(x) - - if n3 > 0: - words.append(ONES[n3][0]) - if n3 > 1: - words.append(HUNDRED[1]) - else: - words.append(HUNDRED[0]) - - if n2 > 1: - words.append(TWENTIES[n2][0]) - - if n2 == 1: - words.append(TENS[n1][0]) - elif n1 > 0: - if (i == 1 or feminine and i == 0) and n < 1000: - words.append(ONES_FEMININE[n1][0]) - else: - words.append(ONES[n1][0]) - - if i > 0: - words.append(self.pluralize(x, THOUSANDS[i])) - - return ' '.join(words) diff --git a/fun_text_processing/num2words/num2words/lang_LV.py b/fun_text_processing/num2words/num2words/lang_LV.py deleted file mode 100644 index 7a833d33c..000000000 --- a/fun_text_processing/num2words/num2words/lang_LV.py +++ /dev/null @@ -1,174 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import unicode_literals - -from .base import Num2Word_Base -from .utils import get_digits, splitbyx - -ZERO = ('nulle',) - -ONES = { - 1: ('viens',), - 2: ('divi',), - 3: ('trīs',), - 4: ('četri',), - 5: ('pieci',), - 6: ('seši',), - 7: ('septiņi',), - 8: ('astoņi',), - 9: ('deviņi',), -} - -TENS = { - 0: ('desmit',), - 1: ('vienpadsmit',), - 2: ('divpadsmit',), - 3: ('trīspadsmit',), - 4: ('četrpadsmit',), - 5: ('piecpadsmit',), - 6: ('sešpadsmit',), - 7: ('septiņpadsmit',), - 8: ('astoņpadsmit',), - 9: ('deviņpadsmit',), -} - -TWENTIES = { - 2: ('divdesmit',), - 3: ('trīsdesmit',), - 4: ('četrdesmit',), - 5: ('piecdesmit',), - 6: ('sešdesmit',), - 7: ('septiņdesmit',), - 8: ('astoņdesmit',), - 9: ('deviņdesmit',), -} - -HUNDRED = ('simts', 'simti', 'simtu') - -THOUSANDS = { - 1: ('tūkstotis', 'tūkstoši', 'tūkstošu'), - 2: ('miljons', 'miljoni', 'miljonu'), - 3: ('miljards', 'miljardi', 'miljardu'), - 4: ('triljons', 'triljoni', 'triljonu'), - 5: ('kvadriljons', 'kvadriljoni', 'kvadriljonu'), - 6: ('kvintiljons', 'kvintiljoni', 'kvintiljonu'), - 7: ('sikstiljons', 'sikstiljoni', 'sikstiljonu'), - 8: ('septiljons', 'septiljoni', 'septiljonu'), - 9: ('oktiljons', 'oktiljoni', 'oktiljonu'), - 10: ('nontiljons', 'nontiljoni', 'nontiljonu'), -} - -GENERIC_DOLLARS = ('dolārs', 'dolāri', 'dolāru') -GENERIC_CENTS = ('cents', 'centi', 'centu') - -GENERIC_KRONA = ('krona', 'kronas', 'kronu') -GENERIC_ERA = ('ēre', 'ēras', 'ēru') - - -class Num2Word_LV(Num2Word_Base): - """ - Sadly we have a legal form (used in legal and finance documents): - http://www.eiro.lv/files/upload/files/Eiro_rakstiba-1.pdf - https://likumi.lv/doc.php?id=254741 - http://eur-lex.europa.eu/legal-content/LV/TXT/HTML/?uri=CELEX:31998R0974&from=LV - - Source: http://publications.europa.eu/code/lv/lv-5000500.htm - """ - CURRENCY_FORMS = { - 'AUD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'CAD': (GENERIC_DOLLARS, GENERIC_CENTS), - # repalced by EUR - 'EEK': (GENERIC_KRONA, GENERIC_CENTS), - 'EUR': (('eiro', 'eiro', 'eiro'), GENERIC_CENTS), - 'EUR_LEGAL': (('euro', 'euro', 'euro'), GENERIC_CENTS), - 'GBP': ( - ('sterliņu mārciņa', 'sterliņu mārciņas', 'sterliņu mārciņu'), - ('penss', 'pensi', 'pensu')), - # replaced by EUR - 'LTL': (('lits', 'liti', 'litu'), GENERIC_CENTS), - # replaced by EUR - 'LVL': (('lats', 'lati', 'latu'), - ('santīms', 'santīmi', 'santīmu')), - 'USD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'RUB': (('rublis', 'rubļi', 'rubļu'), - ('kapeika', 'kapeikas', 'kapeiku')), - 'SEK': (GENERIC_KRONA, GENERIC_ERA), - 'NOK': (GENERIC_KRONA, GENERIC_ERA), - 'PLN': (('zlots', 'zloti', 'zlotu'), - ('grasis', 'graši', 'grašu')), - } - - CURRENCY_ADJECTIVES = { - 'AUD': 'Austrālijas', - 'CAD': 'Kanādas', - 'EEK': 'Igaunijas', - 'USD': 'ASV', - 'RUB': 'Kreivijas', - 'SEK': 'Zviedrijas', - 'NOK': 'Norvēģijas', - } - - def setup(self): - self.negword = "mīnus" - self.pointword = "komats" - - def to_cardinal(self, number): - n = str(number).replace(',', '.') - base_str, n = self.parse_minus(n) - if '.' in n: - left, right = n.split('.') - leading_zero_count = len(right) - len(right.lstrip('0')) - decimal_part = ((ZERO[0] + ' ') * leading_zero_count + - self._int2word(int(right))) - return '%s%s %s %s' % ( - base_str, - self._int2word(int(left)), - self.pointword, - decimal_part - ) - else: - return "%s%s" % (base_str, self._int2word(int(n))) - - def pluralize(self, n, forms): - form = 0 if (n % 10 == 1 and n % 100 != 11) else 1 if n != 0 else 2 - return forms[form] - - def to_ordinal(self, number): - raise NotImplementedError() - - def _int2word(self, n): - if n == 0: - return ZERO[0] - - words = [] - chunks = list(splitbyx(str(n), 3)) - i = len(chunks) - for x in chunks: - i -= 1 - - if x == 0: - continue - - n1, n2, n3 = get_digits(x) - - if n3 > 0: - if n3 == 1 and n2 == 0 and n1 > 0: - words.append(HUNDRED[2]) - elif n3 > 1: - words.append(ONES[n3][0]) - words.append(HUNDRED[1]) - else: - words.append(HUNDRED[0]) - - if n2 > 1: - words.append(TWENTIES[n2][0]) - - if n2 == 1: - words.append(TENS[n1][0]) - elif n1 > 0 and not (i > 0 and x == 1): - words.append(ONES[n1][0]) - - if i > 0: - words.append(self.pluralize(x, THOUSANDS[i])) - - return ' '.join(words) diff --git a/fun_text_processing/num2words/num2words/lang_MN.py b/fun_text_processing/num2words/num2words/lang_MN.py deleted file mode 100644 index ca2fdca26..000000000 --- a/fun_text_processing/num2words/num2words/lang_MN.py +++ /dev/null @@ -1,98 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import division, print_function, unicode_literals - -from . import lang_EU - - -class Num2Word_MN(lang_EU.Num2Word_EU): - # def set_high_numwords(self, high): - # max = 3 + 3 * len(high) - # for word, n in zip(high, range(max, 3, -3)): - # self.cards[10 ** n] = word - def set_high_numwords(self, high): - max = 3 * len(high) - for word, n in zip(high, range(max, 3, -1)): - # print(word[0],word[1],n) - self.cards[10 ** n] = word[1] - # self.cards[10**n] = word + "លាន" - # try: - # ordinal_word = self.high_numwords[high] - # except KeyError: - # max = 3 + 3 * len(high) - # for word, n in zip(high, range(max, 3, -3)): - # print(word) - # print(n) - # self.cards[10 ** n] = word - - def setup(self): - super(Num2Word_MN, self).setup() - - self.negword = 'хасах' - self.pointword = 'цэг' - - self.high_numwords = [(1000000,'нэг сая'), (100000, 'нэг зуун мянга'), (10000, 'арван мянга'), (1000, 'нэг мянга') ] - - self.mid_numwords = [(100,'нэг зуу'), (90,'ерэн'), (80,'наян'), (70,'далан'), (60,'жаран'), (50,'тавин'), (40,'дөчин'), (30,'гучин')] - - self.low_numwords = ['хорин','арван есөн','арван найман','арван долоон','арван зургаа','арван тав','арван дөрөв','арван гурав','арван хоёр','арван нэгэн','арав','есөн','найм','Долоо','зургаа','тав','дөрөв','гурав','хоёр','нэг','тэг'] - - self.ords = {'нэг': 'эхлээд', - 'хоёр': 'хоёрдугаарт', - 'гурав': 'гурав дахь', - 'дөрөв': 'урагш', - 'тав': 'тав дахь', - 'зургаа': 'зургаа дахь', - 'Долоо': 'долоо дахь', - 'найм': 'найм дахь', - 'есөн': 'ес дэх', - 'арав': 'Аравдугаар'} - - - def merge(self, lpair, rpair): - ltext, lnum = lpair - rtext, rnum = rpair - if lnum == 1 and rnum < 100: - return (rtext, rnum) - elif 100 > lnum > rnum: - return ("%s %s" % (ltext, rtext), lnum + rnum) - elif lnum >= 100 > rnum: - return ("%s болон %s" % (ltext, rtext), lnum + rnum) - elif rnum > lnum: - return ("%s %s" % (ltext, rtext), lnum * rnum) - return ("%s %s" % (ltext, rtext), lnum + rnum) - - def to_ordinal_num(self, value): - self.verify_ordinal(value) - return "%s%s" % (value, self.to_ordinal(value)) - - def to_ordinal(self, value): - self.verify_ordinal(value) - try: - ordinal_word = self.ords[value] - except KeyError: - ordinal_word = value #TODO: check nepali ordinal word suffix - return ordinal_word - - - def to_year(self, val, suffix=None, longval=True): - if val < 0: - val = abs(val) - suffix = 'BC' if not suffix else suffix - high, low = (val // 100, val % 100) - # If year is 00XX, X00X, or beyond 9999, go cardinal. - if (high == 0 - or (high % 10 == 0 and low < 10) - or high >= 100): - valtext = self.to_cardinal(val) - else: - hightext = self.to_cardinal(high) - if low == 0: - lowtext = "зуун" - elif low < 10: - lowtext = "%s" % self.to_cardinal(low) - else: - lowtext = self.to_cardinal(low) - valtext = "%s%s" % (hightext, lowtext) - return (valtext if not suffix - else "%s%s" % (valtext, suffix)) diff --git a/fun_text_processing/num2words/num2words/lang_MY.py b/fun_text_processing/num2words/num2words/lang_MY.py deleted file mode 100644 index b54f245e6..000000000 --- a/fun_text_processing/num2words/num2words/lang_MY.py +++ /dev/null @@ -1,97 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import division, print_function, unicode_literals - -from . import lang_EU - - -class Num2Word_MY(lang_EU.Num2Word_EU): - # def set_high_numwords(self, high): - # max = 3 + 3 * len(high) - # for word, n in zip(high, range(max, 3, -3)): - # self.cards[10 ** n] = word - def set_high_numwords(self, high): - max = 3 * len(high) - for word, n in zip(high, range(max, 3, -1)): - # print(word[0],word[1],n) - self.cards[10 ** n] = word[1] - # self.cards[10**n] = word + "លាន" - # try: - # ordinal_word = self.high_numwords[high] - # except KeyError: - # max = 3 + 3 * len(high) - # for word, n in zip(high, range(max, 3, -3)): - # print(word) - # print(n) - # self.cards[10 ** n] = word - - def setup(self): - super(Num2Word_MY, self).setup() - - self.negword = 'အမှတ်' - self.pointword = 'အမှတ်' - - self.high_numwords = [(1000000,'သန်း'), (100000, 'သိန်း'), (10000, 'သောင်း'), (1000, 'ထောင်') ] - - self.mid_numwords = [(100,'ရာ'), (90,'ကိုးဆယ်'), (80,'ရှစ်ဆယ်'), (70,'ခုနှစ်ဆယ်'), (60,'ခြောက်ဆယ့်'), (50,'ငါးဆယ်'), (40,'လေးဆယ်'), (30,'သုံးဆယ်')] - - # self.low_numwords = ['शून्य', 'एक', 'दुई', 'तीन', 'चार', 'पाँच', 'छ', 'सात', 'आठ', 'नौ', 'दस', 'एघार', 'बाह्र', 'तेह्र', 'चौध', 'पन्ध्र', 'सोह्र', 'सत्रह', 'अठार', 'उन्नीस'] - self.low_numwords = ['နှစ်ဆယ်','ဆယ့်ကိုး','ဆယ့်ရှစ်','ဆယ့်ခုနှစ်','ဆယ့်ခြောက်','ဆယ့်ငါး','ဆယ့်လေး','ဆယ့်သုံး','ဆယ့်နှစ်','ဆယ့်တစ်','ဆယ်','ကိုး','ရှစ်','ခုနှစ်','ခြောက်','ငါး','လေး','သုံး','နှစ်','တစ်','သုည'] - - self.ords = {'တစ်':'ပထမ', - 'နှစ်':'ဒုတိယ', - 'သုံး':'တတိယ', - 'လေး':'စတုတ္ထ', - 'ငါး':'ပဉ္စမ', - 'ခြောက်':'ဆဋ္ဌမ', - 'ခုနှစ်':'သတ္တမ', - 'ရှစ်':'အဋ္ဌမ', - 'ကိုး':'နဝမ', - 'ဆယ်':'ဒသမ'} - def merge(self, lpair, rpair): - ltext, lnum = lpair - rtext, rnum = rpair - if lnum == 1 and rnum < 100: - return (rtext, rnum) - elif 100 > lnum > rnum: - return ("%s %s" % (ltext, rtext), lnum + rnum) - elif lnum >= 100 > rnum: - return ("%s နှင့် %s" % (ltext, rtext), lnum + rnum) - elif rnum > lnum: - return ("%s %s" % (ltext, rtext), lnum * rnum) - return ("%s %s" % (ltext, rtext), lnum + rnum) - - def to_ordinal_num(self, value): - self.verify_ordinal(value) - return "%s%s" % (value, self.to_ordinal(value)) - - def to_ordinal(self, value): - self.verify_ordinal(value) - try: - ordinal_word = self.ords[value] - except KeyError: - ordinal_word = value #TODO: check nepali ordinal word suffix - return ordinal_word - - - def to_year(self, val, suffix=None, longval=True): - if val < 0: - val = abs(val) - suffix = 'BC' if not suffix else suffix - high, low = (val // 100, val % 100) - # If year is 00XX, X00X, or beyond 9999, go cardinal. - if (high == 0 - or (high % 10 == 0 and low < 10) - or high >= 100): - valtext = self.to_cardinal(val) - else: - hightext = self.to_cardinal(high) - if low == 0: - lowtext = "ရာ့" - elif low < 10: - lowtext = "%s" % self.to_cardinal(low) - else: - lowtext = self.to_cardinal(low) - valtext = "%s%s" % (hightext, lowtext) - return (valtext if not suffix - else "%s%s" % (valtext, suffix)) diff --git a/fun_text_processing/num2words/num2words/lang_NE.py b/fun_text_processing/num2words/num2words/lang_NE.py deleted file mode 100644 index f82c27a34..000000000 --- a/fun_text_processing/num2words/num2words/lang_NE.py +++ /dev/null @@ -1,250 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import division, print_function, unicode_literals - -from . import lang_EU - - -class Num2Word_NE(lang_EU.Num2Word_EU): - # def set_high_numwords(self, high): - # max = 3 + 3 * len(high) - # for word, n in zip(high, range(max, 3, -3)): - # self.cards[10 ** n] = word - def set_high_numwords(self, high): - max = 3 * len(high) - for word, n in zip(high, range(max, 0, -3)): - # print(word[0],word[1],n) - self.cards[10 ** n] = word[1] - # self.cards[10**n] = word + "លាន" - # try: - # ordinal_word = self.high_numwords[high] - # except KeyError: - # max = 3 + 3 * len(high) - # for word, n in zip(high, range(max, 3, -3)): - # print(word) - # print(n) - # self.cards[10 ** n] = word - - def setup(self): - super(Num2Word_NE, self).setup() - - self.negword = 'माइनस' - self.pointword = 'अंक' - - self.high_numwords = [(1000000000000000, 'सय ट्रिलियन'), (1000000000000,'ट्रिलियन'), (1000000000, 'बिलियन'), (1000000, 'मिलियन'), (1000, 'हजार') ] - - self.mid_numwords = [(100, 'एक सय'), (90, 'नब्बे'), (80, 'अस्सी'), (70, 'सत्तरी'), (60, 'साठ'), (50, 'पचास'), (40, 'चालीस'), (30, 'तीस'), (20, 'विस')] - - # self.low_numwords = ['शून्य', 'एक', 'दुई', 'तीन', 'चार', 'पाँच', 'छ', 'सात', 'आठ', 'नौ', 'दस', 'एघार', 'बाह्र', 'तेह्र', 'चौध', 'पन्ध्र', 'सोह्र', 'सत्रह', 'अठार', 'उन्नीस'] - self.low_numwords = ["उन्नाइस", - "अठार", - "सत्र", - "सोह्र", - "पन्ध्र", - "चौध", - "तेह्र", - "बाह्र", - "एघार", - "दश", - "नौ", - "आठ", - "सात", - "छ", - "पाँच", - "चार", - "तीन", - "दुई", - "एक", - "शून्य"] - - self.ords = {'एक': 'पहिलो', - 'दुई': 'द्दोस्रो', - 'तीन': 'तेस्रो', - 'चार': 'चौथुर्त', - 'पाँच': 'पाँचौ', - 'छ': 'छैटौ', - 'सात': 'सातौ', - 'आठ': 'आठौ', - 'नौ': 'नवौ', - 'दश': 'दसौ'} - - self.labeled_numbers = {'0': 'शुन्य', - '1': 'एक', - '2': 'दुई', - '3': 'तीन', - '4': 'चार', - '5': 'पाँच', - '6': 'छ', - '7': 'सात', - '8': 'आठ', - '9': 'नौ', - '10': 'दश', - '11': 'एघार', - '12': 'बाह्र', - '13': 'तेह्र', - '14': 'चौध', - '15': 'पन्ध्र', - '16': 'सोह्र', - '17': 'सत्र', - '18': 'अठार', - '19': 'उन्नाइस', - '20': 'विस', - '21': 'एक्काइस', - '22': 'बाइस', - '23': 'तेईस', - '24': 'चौविस', - '25': 'पच्चिस', - '26': 'छब्बिस', - '27': 'सत्ताइस', - '28': 'अठ्ठाईस', - '29': 'उनन्तिस', - '30': 'तिस', - '31': 'एकत्तिस', - '32': 'बत्तिस', - '33': 'तेत्तिस', - '34': 'चौँतिस', - '35': 'पैँतिस', - '36': 'छत्तिस', - '37': 'सैँतीस', - '38': 'अठतीस', - '39': 'उनन्चालीस', - '40': 'चालीस', - '41': 'एकचालीस', - '42': 'बयालीस', - '43': 'त्रियालीस', - '44': 'चवालीस', - '45': 'पैँतालीस', - '46': 'छयालीस', - '47': 'सच्चालीस', - '48': 'अठचालीस', - '49': 'उनन्चास', - '50': 'पचास', - '51': 'एकाउन्न', - '52': 'बाउन्न', - '53': 'त्रिपन्न', - '54': 'चउन्न', - '55': 'पचपन्न', - '56': 'छपन्न', - '57': 'सन्ताउन्न', - '58': 'अन्ठाउन्न', - '59': 'उनन्साठी', - '60': 'साठी', - '61': 'एकसट्ठी', - '62': 'बयसट्ठी', - '63': 'त्रिसट्ठी', - '64': 'चौंसट्ठी', - '65': 'पैंसट्ठी', - '66': 'छयसट्ठी', - '67': 'सतसट्ठी', - '68': 'अठसट्ठी', - '69': 'उनन्सत्तरी', - '70': 'सत्तरी', - '71': 'एकहत्तर', - '72': 'बहत्तर', - '73': 'त्रिहत्तर', - '74': 'चौहत्तर', - '75': 'पचहत्तर', - '76': 'छयहत्तर', - '77': 'सतहत्तर', - '78': 'अठहत्तर', - '79': 'उनासी', - '80': 'असी', - '81': 'एकासी', - '82': 'बयासी', - '83': 'त्रियासी', - '84': 'चौरासी', - '85': 'पचासी', - '86': 'छयासी', - '87': 'सतासी', - '88': 'अठासी', - '89': 'उनान्नब्बे', - '90': 'नब्बे', - '91': 'एकान्नब्बे', - '92': 'बयानब्बे', - '93': 'त्रियान्नब्बे', - '94': 'चौरान्नब्बे', - '95': 'पन्चानब्बे', - '96': 'छयान्नब्बे', - '97': 'सन्तान्नब्बे', - '98': 'अन्ठान्नब्बे', - '99': 'उनान्सय', - '100': 'एक सय', - '200': 'दुई सय', - '300': 'तीन सय', - '400': 'चार सय', - '500': 'पाँच सय', - '600': 'छ सय', - '700': 'सात सय', - '800': 'आठ सय', - '900': 'नौ सय', - '1,000': 'एक हजार', - '2,000': 'दुई हजार', - '3,000': 'तीन हजार', - '4,000': 'चार हजार', - '5,000': 'पाँच हजार', - '6,000': 'छ हजार', - '7,000': 'सात हजार', - '8,000': 'आठ हजार', - '9,000': 'नौ हजार', - '10,000': 'दश हजार', - '20,000': 'बिस हजार', - '30,000': 'तीस हजार', - '40,000': 'चालिस हजार', - '50,000': 'पचास हजार', - '60,000': 'साठी हजार', - '70,000': 'सत्तरी हजार', - '80,000': 'असी हजार', - '90,000': 'नब्बे हजार', - '100,000': 'एक लाख', - '1,000,000': 'दश लाख', - '10,000,000': 'एक करोड', - '1,000,000,000': 'एक अर्ब', - '100,000,000,000': 'एक खर्ब'} - - def merge(self, lpair, rpair): - ltext, lnum = lpair - rtext, rnum = rpair - if lnum == 1 and rnum < 100: - return (rtext, rnum) - elif 100 > lnum > rnum: - return ("%s %s" % (ltext, rtext), lnum + rnum) - elif lnum >= 100 > rnum: - return ("%s %s" % (ltext, rtext), lnum + rnum) - elif rnum > lnum: - return ("%s %s" % (ltext, rtext), lnum * rnum) - return ("%s %s" % (ltext, rtext), lnum + rnum) - - def to_ordinal_num(self, value): - self.verify_ordinal(value) - return "%s%s" % (value, self.to_ordinal(value)) - - def to_ordinal(self, value): - self.verify_ordinal(value) - try: - ordinal_word = self.ords[value] - except KeyError: - ordinal_word = value #TODO: check nepali ordinal word suffix - return ordinal_word - - - def to_year(self, val, suffix=None, longval=True): - if val < 0: - val = abs(val) - suffix = 'BC' if not suffix else suffix - high, low = (val // 100, val % 100) - # If year is 00XX, X00X, or beyond 9999, go cardinal. - if (high == 0 - or (high % 10 == 0 and low < 10) - or high >= 100): - valtext = self.to_cardinal(val) - else: - hightext = self.to_cardinal(high) - if low == 0: - lowtext = "सय" - elif low < 10: - lowtext = "%s" % self.to_cardinal(low) - else: - lowtext = self.to_cardinal(low) - valtext = "%s %s" % (hightext, lowtext) - return (valtext if not suffix - else "%s %s" % (valtext, suffix)) diff --git a/fun_text_processing/num2words/num2words/lang_NL.py b/fun_text_processing/num2words/num2words/lang_NL.py deleted file mode 100644 index 97a5aa244..000000000 --- a/fun_text_processing/num2words/num2words/lang_NL.py +++ /dev/null @@ -1,140 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import print_function, unicode_literals - -from .lang_EU import Num2Word_EU - - -class Num2Word_NL(Num2Word_EU): - CURRENCY_FORMS = { - 'EUR': (('euro', 'euro'), ('cent', 'cent')), - 'GBP': (('pond', 'pond'), ('penny', 'pence')), - 'USD': (('dollar', 'dollar'), ('cent', 'cent')), - 'CNY': (('yuan', 'yuan'), ('jiao', 'fen')), - } - - GIGA_SUFFIX = "iljard" - MEGA_SUFFIX = "iljoen" - - def setup(self): - super(Num2Word_NL, self).setup() - - self.negword = "min " - self.pointword = "komma" - # "Cannot treat float %s as ordinal." - self.errmsg_floatord = ( - "Het zwevende puntnummer %s kan niet omgezet worden " + - "naar een ordernummer." - ) - # "type(((type(%s)) ) not in [long, int, float]" - self.errmsg_nonnum = ( - "Alleen nummers (type (%s)) kunnen naar " + - "woorden omgezet worden." - ) - # "Cannot treat negative num %s as ordinal." - self.errmsg_negord = ( - "Het negatieve getal %s kan niet omgezet " + - "worden naar een ordernummer." - ) - # "abs(%s) must be less than %s." - self.errmsg_toobig = "Het getal %s moet minder zijn dan %s." - self.exclude_title = [] - - lows = ["non", "okt", "sept", "sext", "quint", "quadr", "tr", "b", "m"] - units = ["", "un", "duo", "tre", "quattuor", "quin", "sex", "sept", - "okto", "novem"] - tens = ["dez", "vigint", "trigint", "quadragint", "quinquagint", - "sexagint", "septuagint", "oktogint", "nonagint"] - - self.high_numwords = ( - ["zend"] + self.gen_high_numwords(units, tens, lows)) - - self.mid_numwords = [(1000, "duizend"), (100, "honderd"), - (90, "negentig"), (80, "tachtig"), - (70, "zeventig"), (60, "zestig"), - (50, "vijftig"), (40, "veertig"), - (30, "dertig")] - - self.low_numwords = ["twintig", "negentien", "achttien", "zeventien", - "zestien", "vijftien", "veertien", "dertien", - "twaalf", "elf", "tien", "negen", "acht", "zeven", - "zes", "vijf", "vier", "drie", "twee", "één", - "nul"] - - self.ords = {"één": "eerst", - "twee": "tweed", - "drie": "derd", - "vier": "vierd", - "vijf": "vijfd", - "zes": "zesd", - "zeven": "zevend", - "acht": "achtst", - "negen": "negend", - "tien": "tiend", - "elf": "elfd", - "twaalf": "twaalfd", - - "ig": "igst", - "erd": "erdst", - "end": "endst", - "joen": "joenst", - "rd": "rdst"} - - def merge(self, curr, next): - ctext, cnum, ntext, nnum = curr + next - - if cnum == 1: - if nnum < 10 ** 6: - return next - ctext = "een" - - if nnum > cnum: - if nnum >= 10 ** 6: - ctext += " " - val = cnum * nnum - else: - if nnum < 10 < cnum < 100: - if nnum == 1: - ntext = "een" - - if ntext.endswith("e"): - ntext += "ën" # "n" - else: - ntext += "en" - ntext, ctext = ctext, ntext # + "en" - elif cnum >= 10 ** 6: - ctext += " " - val = cnum + nnum - - word = ctext + ntext - return word, val - - def to_ordinal(self, value): - self.verify_ordinal(value) - outword = self.to_cardinal(value) - for key in self.ords: - if outword.endswith(key): - outword = outword[:len(outword) - len(key)] + self.ords[key] - break - return outword + "e" - - def to_ordinal_num(self, value): - self.verify_ordinal(value) - return str(value) + "e" - - def pluralize(self, n, forms): - """ - :param n: - :param forms: - :return: - - gettext form is nplurals=2; plural=(n != 1); - but this claims https://onzetaal.nl/taaladvies/euro-euros/ - not sure if it's applied only to euro - """ - return forms[0] - - def to_year(self, val, longval=True): - if not (val // 100) % 10: - return self.to_cardinal(val) - return self.to_splitnum(val, hightxt="honderd", longval=longval) diff --git a/fun_text_processing/num2words/num2words/lang_NO.py b/fun_text_processing/num2words/num2words/lang_NO.py deleted file mode 100644 index 5852abbe6..000000000 --- a/fun_text_processing/num2words/num2words/lang_NO.py +++ /dev/null @@ -1,89 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import division, print_function, unicode_literals - -from . import lang_EU - - -class Num2Word_NO(lang_EU.Num2Word_EU): - GIGA_SUFFIX = "illard" - MEGA_SUFFIX = "illion" - - def set_high_numwords(self, high): - cap = 3 + 6 * len(high) - - for word, n in zip(high, range(cap, 3, -6)): - if self.GIGA_SUFFIX: - self.cards[10 ** n] = word + self.GIGA_SUFFIX - - if self.MEGA_SUFFIX: - self.cards[10 ** (n - 3)] = word + self.MEGA_SUFFIX - - def setup(self): - super(Num2Word_NO, self).setup() - - self.negword = "minus " - self.pointword = "komma" - self.exclude_title = ["og", "komma", "minus"] - - self.mid_numwords = [(1000, "tusen"), (100, "hundre"), - (90, "nitti"), (80, "åtti"), (70, "sytti"), - (60, "seksti"), (50, "femti"), (40, "førti"), - (30, "tretti")] - self.low_numwords = ["tjue", "nitten", "atten", "sytten", "seksten", "femten", "fjorten", "tretten", "tolv", "elleve", "ti", "ni", "åtte", "syv", "seks", "fem", "fire", "tre", "to", "en", "null"] - self.ords = {"en": "først", - "to": "andre", - "tre": "tredje", - "fire": "fjerde", - "fem": "femte", - "seks": "sjette", - "syv": "syvende", - "åtte": "åttende", - "ni": "niende", - "ti": "tiende", - "elleve": "ellevte", - "tolv": "tolvte", - "tjue": "tjuende"} - - def merge(self, lpair, rpair): - ltext, lnum = lpair - rtext, rnum = rpair - if lnum == 1 and rnum < 100: - return (rtext, rnum) - elif 100 > lnum > rnum: - return ("%s-%s" % (ltext, rtext), lnum + rnum) - elif lnum >= 100 > rnum: - return ("%s og %s" % (ltext, rtext), lnum + rnum) - elif rnum > lnum: - return ("%s %s" % (ltext, rtext), lnum * rnum) - return ("%s, %s" % (ltext, rtext), lnum + rnum) - - def to_ordinal(self, value): - self.verify_ordinal(value) - outwords = self.to_cardinal(value).split(" ") - lastwords = outwords[-1].split("-") - lastword = lastwords[-1].lower() - try: - lastword = self.ords[lastword] - except KeyError: - if lastword[-2:] == "ti": - lastword = lastword + "ende" - else: - lastword += "de" - lastwords[-1] = self.title(lastword) - outwords[-1] = "".join(lastwords) - return " ".join(outwords) - - def to_ordinal_num(self, value): - self.verify_ordinal(value) - return "%s%s" % (value, self.to_ordinal(value)[-2:]) - - def to_year(self, val, longval=True): - if not (val // 100) % 10: - return self.to_cardinal(val) - return self.to_splitnum(val, hightxt="hundre", jointxt="og", - longval=longval) - - def to_currency(self, val, longval=True): - return self.to_splitnum(val, hightxt="krone/r", lowtxt="\xf8re/r", - jointxt="og", longval=longval, cents=True) diff --git a/fun_text_processing/num2words/num2words/lang_PL.py b/fun_text_processing/num2words/num2words/lang_PL.py deleted file mode 100644 index f7760aa2b..000000000 --- a/fun_text_processing/num2words/num2words/lang_PL.py +++ /dev/null @@ -1,245 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import unicode_literals - -import itertools - -from .base import Num2Word_Base -from .utils import get_digits, splitbyx - -ZERO = ('zero',) - -ONES = { - 1: ('jeden',), - 2: ('dwa',), - 3: ('trzy',), - 4: ('cztery',), - 5: ('pięć',), - 6: ('sześć',), - 7: ('siedem',), - 8: ('osiem',), - 9: ('dziewięć',), -} - -ONES_ORDINALS = { - 1: ('pierwszy', "pierwszo"), - 2: ('drugi', "dwu"), - 3: ('trzeci', "trzy"), - 4: ('czwarty', "cztero"), - 5: ('piąty', "pięcio"), - 6: ('szósty', "sześcio"), - 7: ('siódmy', "siedmio"), - 8: ('ósmy', "ośmio"), - 9: ('dziewiąty', "dziewięcio"), - 10: ('dziesiąty', "dziesięcio"), - 11: ('jedenasty', "jedenasto"), - 12: ('dwunasty', "dwunasto"), - 13: ('trzynasty', "trzynasto"), - 14: ('czternasty', "czternasto"), - 15: ('piętnasty', "piętnasto"), - 16: ('szesnasty', "szesnasto"), - 17: ('siedemnasty', "siedemnasto"), - 18: ('osiemnasty', "osiemnasto"), - 19: ('dziewiętnasty', "dziewiętnasto"), -} - -TENS = { - 0: ('dziesięć',), - 1: ('jedenaście',), - 2: ('dwanaście',), - 3: ('trzynaście',), - 4: ('czternaście',), - 5: ('piętnaście',), - 6: ('szesnaście',), - 7: ('siedemnaście',), - 8: ('osiemnaście',), - 9: ('dziewiętnaście',), -} - - -TWENTIES = { - 2: ('dwadzieścia',), - 3: ('trzydzieści',), - 4: ('czterdzieści',), - 5: ('pięćdziesiąt',), - 6: ('sześćdziesiąt',), - 7: ('siedemdziesiąt',), - 8: ('osiemdziesiąt',), - 9: ('dziewięćdzisiąt',), -} - -TWENTIES_ORDINALS = { - 2: ('dwudziesty', "dwudziesto"), - 3: ('trzydziesty', "trzydiesto"), - 4: ('czterdziesty', "czterdziesto"), - 5: ('pięćdziesiąty', "pięćdziesięcio"), - 6: ('sześćdziesiąty', "sześćdziesięcio"), - 7: ('siedemdziesiąty', "siedemdziesięcio"), - 8: ('osiemdziesiąty', "osiemdziesięcio"), - 9: ('dziewięćdzisiąty', "dziewięćdziesięcio"), -} - -HUNDREDS = { - 1: ('sto',), - 2: ('dwieście',), - 3: ('trzysta',), - 4: ('czterysta',), - 5: ('pięćset',), - 6: ('sześćset',), - 7: ('siedemset',), - 8: ('osiemset',), - 9: ('dziewięćset',), -} - -HUNDREDS_ORDINALS = { - 1: ('setny', "stu"), - 2: ('dwusetny', "dwustu"), - 3: ('trzysetny', "trzystu"), - 4: ('czterysetny', "czterystu"), - 5: ('pięćsetny', "pięcset"), - 6: ('sześćsetny', "sześćset"), - 7: ('siedemsetny', "siedemset"), - 8: ('osiemsetny', "ośiemset"), - 9: ('dziewięćsetny', "dziewięćset"), -} - -THOUSANDS = { - 1: ('tysiąc', 'tysiące', 'tysięcy'), # 10^3 -} - -prefixes_ordinal = { - 1: "tysięczny", - 2: "milionowy", - 3: "milairdowy" -} - -prefixes = ( # 10^(6*x) - "mi", # 10^6 - "bi", # 10^12 - "try", # 10^18 - "kwadry", # 10^24 - "kwinty", # 10^30 - "seksty", # 10^36 - "septy", # 10^42 - "okty", # 10^48 - "nony", # 10^54 - "decy" # 10^60 -) -suffixes = ("lion", "liard") # 10^x or 10^(x+3) - -for idx, (p, s) in enumerate(itertools.product(prefixes, suffixes)): - name = p + s - THOUSANDS[idx+2] = (name, name + 'y', name + 'ów') - - -class Num2Word_PL(Num2Word_Base): - CURRENCY_FORMS = { - 'PLN': ( - ('złoty', 'złote', 'złotych'), ('grosz', 'grosze', 'groszy') - ), - 'EUR': ( - ('euro', 'euro', 'euro'), ('cent', 'centy', 'centów') - ), - } - - def setup(self): - self.negword = "minus" - self.pointword = "przecinek" - - def to_cardinal(self, number): - n = str(number).replace(',', '.') - if '.' in n: - left, right = n.split('.') - leading_zero_count = len(right) - len(right.lstrip('0')) - decimal_part = ((ZERO[0] + ' ') * leading_zero_count + - self._int2word(int(right))) - return u'%s %s %s' % ( - self._int2word(int(left)), - self.pointword, - decimal_part - ) - else: - return self._int2word(int(n)) - - def pluralize(self, n, forms): - if n == 1: - form = 0 - elif 5 > n % 10 > 1 and (n % 100 < 10 or n % 100 > 20): - form = 1 - else: - form = 2 - return forms[form] - - def last_fragment_to_ordinal(self, last, words, level): - n1, n2, n3 = get_digits(last) - last_two = n2*10+n1 - if last_two == 0: - words.append(HUNDREDS_ORDINALS[n3][level]) - elif level == 1 and last == 1: - return - elif last_two < 20: - if n3 > 0: - words.append(HUNDREDS[n3][level]) - words.append(ONES_ORDINALS[last_two][level]) - elif last_two % 10 == 0: - if n3 > 0: - words.append(HUNDREDS[n3][level]) - words.append(TWENTIES_ORDINALS[n2][level]) - else: - if n3 > 0: - words.append(HUNDREDS[n3][0]) - words.append(TWENTIES_ORDINALS[n2][0]) - words.append(ONES_ORDINALS[n1][0]) - - def to_ordinal(self, number): - if number % 1 != 0: - raise NotImplementedError() - words = [] - fragments = list(splitbyx(str(number), 3)) - level = 0 - last = fragments[-1] - while last == 0: - level = level+1 - fragments.pop() - last = fragments[-1] - if len(fragments) > 1: - pre_part = self._int2word(number-(last*1000**level)) - words.append(pre_part) - self.last_fragment_to_ordinal(last, words, 0 if level == 0 else 1) - output = " ".join(words) - if last == 1 and level > 0 and output != "": - output = output + " " - if level > 0: - output = output + prefixes_ordinal[level] - return output - - def _int2word(self, n): - if n == 0: - return ZERO[0] - - words = [] - chunks = list(splitbyx(str(n), 3)) - i = len(chunks) - for x in chunks: - i -= 1 - - if x == 0: - continue - - n1, n2, n3 = get_digits(x) - - if n3 > 0: - words.append(HUNDREDS[n3][0]) - - if n2 > 1: - words.append(TWENTIES[n2][0]) - - if n2 == 1: - words.append(TENS[n1][0]) - elif n1 > 0 and not (i > 0 and x == 1): - words.append(ONES[n1][0]) - - if i > 0: - words.append(self.pluralize(x, THOUSANDS[i])) - - return ' '.join(words) diff --git a/fun_text_processing/num2words/num2words/lang_PT.py b/fun_text_processing/num2words/num2words/lang_PT.py deleted file mode 100644 index 72c306f64..000000000 --- a/fun_text_processing/num2words/num2words/lang_PT.py +++ /dev/null @@ -1,226 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import division, unicode_literals - -import re - -from .lang_EU import Num2Word_EU - -DOLLAR = ('dólar', 'dólares') -CENTS = ('cêntimo', 'cêntimos') - - -class Num2Word_PT(Num2Word_EU): - - CURRENCY_FORMS = { - 'AUD': (DOLLAR, CENTS), - 'CAD': (DOLLAR, CENTS), - 'EUR': (('euro', 'euros'), CENTS), - 'GBP': (('libra', 'libras'), ('péni', 'pence')), - 'USD': (DOLLAR, CENTS), - } - - GIGA_SUFFIX = None - MEGA_SUFFIX = "ilião" - - def setup(self): - super(Num2Word_PT, self).setup() - lows = ["quatr", "tr", "b", "m"] - self.high_numwords = self.gen_high_numwords([], [], lows) - self.negword = "menos " - self.pointword = "vírgula" - self.exclude_title = ["e", "vírgula", "menos"] - - self.mid_numwords = [ - (1000, "mil"), (100, "cem"), (90, "noventa"), - (80, "oitenta"), (70, "setenta"), (60, "sessenta"), - (50, "cinquenta"), (40, "quarenta"), (30, "trinta") - ] - self.low_numwords = [ - "vinte", "dezanove", "dezoito", "dezassete", "dezasseis", - "quinze", "catorze", "treze", "doze", "onze", "dez", - "nove", "oito", "sete", "seis", "cinco", "quatro", "três", "dois", - "um", "zero" - ] - self.ords = [ - { - 0: "", - 1: "primeiro", - 2: "segundo", - 3: "terceiro", - 4: "quarto", - 5: "quinto", - 6: "sexto", - 7: "sétimo", - 8: "oitavo", - 9: "nono", - }, - { - 0: "", - 1: "décimo", - 2: "vigésimo", - 3: "trigésimo", - 4: "quadragésimo", - 5: "quinquagésimo", - 6: "sexagésimo", - 7: "septuagésimo", - 8: "octogésimo", - 9: "nonagésimo", - }, - { - 0: "", - 1: "centésimo", - 2: "ducentésimo", - 3: "tricentésimo", - 4: "quadrigentésimo", - 5: "quingentésimo", - 6: "seiscentésimo", - 7: "septigentésimo", - 8: "octigentésimo", - 9: "nongentésimo", - }, - ] - self.thousand_separators = { - 3: "milésimo", - 6: "milionésimo", - 9: "milésimo milionésimo", - 12: "bilionésimo", - 15: "milésimo bilionésimo" - } - self.hundreds = { - 1: "cento", - 2: "duzentos", - 3: "trezentos", - 4: "quatrocentos", - 5: "quinhentos", - 6: "seiscentos", - 7: "setecentos", - 8: "oitocentos", - 9: "novecentos", - } - - def merge(self, curr, next): - ctext, cnum, ntext, nnum = curr + next - - if cnum == 1: - if nnum < 1000000: - return next - ctext = "um" - elif cnum == 100 and not nnum % 1000 == 0: - ctext = "cento" - - if nnum < cnum: - if cnum < 100: - return ("%s e %s" % (ctext, ntext), cnum + nnum) - return ("%s e %s" % (ctext, ntext), cnum + nnum) - - elif (not nnum % 1000000000) and cnum > 1: - ntext = ntext[:-4] + "liões" - elif (not nnum % 1000000) and cnum > 1: - ntext = ntext[:-4] + "lhões" - # correct "milião" to "milhão" - if ntext == 'milião': - ntext = 'milhão' - if nnum == 100: - ctext = self.hundreds[cnum] - ntext = "" - - else: - ntext = " " + ntext - - return (ctext + ntext, cnum * nnum) - - def to_cardinal(self, value): - result = super(Num2Word_PT, self).to_cardinal(value) - - # Transforms "mil e cento e catorze" into "mil cento e catorze" - # Transforms "cem milhões e duzentos mil e duzentos e dez" em "cem - # milhões duzentos mil duzentos e dez" but "cem milhões e duzentos - # mil e duzentos" in "cem milhões duzentos mil e duzentos" and not in - # "cem milhões duzentos mil duzentos" - for ext in ( - 'mil', 'milhão', 'milhões', 'mil milhões', - 'bilião', 'biliões', 'mil biliões'): - if re.match('.*{} e \\w*entos? (?=.*e)'.format(ext), result): - result = result.replace( - '{} e'.format(ext), '{}'.format(ext) - ) - - return result - - # for the ordinal conversion the code is similar to pt_BR code, - # although there are other rules that are probably more correct in - # Portugal. Concerning numbers from 2000th on, saying "dois - # milésimos" instead of "segundo milésimo" (the first number - # would be used in the cardinal form instead of the ordinal) is better. - # This was not implemented. - # source: - # https://ciberduvidas.iscte-iul.pt/consultorio/perguntas/a-forma-por-extenso-de-2000-e-de-outros-ordinais/16428 - def to_ordinal(self, value): - # Before changing this function remember this is used by pt-BR - # so act accordingly - self.verify_ordinal(value) - - result = [] - value = str(value) - thousand_separator = '' - - for idx, char in enumerate(value[::-1]): - if idx and idx % 3 == 0: - thousand_separator = self.thousand_separators[idx] - - if char != '0' and thousand_separator: - # avoiding "segundo milionésimo milésimo" for 6000000, - # for instance - result.append(thousand_separator) - thousand_separator = '' - - result.append(self.ords[idx % 3][int(char)]) - - result = ' '.join(result[::-1]) - result = result.strip() - result = re.sub('\\s+', ' ', result) - - if result.startswith('primeiro') and value != '1': - # avoiding "primeiro milésimo", "primeiro milionésimo" and so on - result = result[9:] - - return result - - def to_ordinal_num(self, value): - # Before changing this function remember this is used by pt-BR - # so act accordingly - self.verify_ordinal(value) - return "%sº" % (value) - - def to_year(self, val, longval=True): - # Before changing this function remember this is used by pt-BR - # so act accordingly - if val < 0: - return self.to_cardinal(abs(val)) + ' antes de Cristo' - return self.to_cardinal(val) - - def to_currency(self, val, currency='EUR', cents=True, separator=' e', - adjective=False): - # change negword because base.to_currency() does not need space after - backup_negword = self.negword - self.negword = self.negword[:-1] - result = super(Num2Word_PT, self).to_currency( - val, currency=currency, cents=cents, separator=separator, - adjective=adjective) - # undo the change on negword - self.negword = backup_negword - - # transforms "milhões euros" em "milhões de euros" - cr1, _ = self.CURRENCY_FORMS[currency] - - for ext in ( - 'milhão', 'milhões', 'bilião', - 'biliões', 'trilião', 'triliões'): - if re.match('.*{} (?={})'.format(ext, cr1[1]), result): - result = result.replace( - '{}'.format(ext), '{} de'.format(ext), 1 - ) - # do not print "e zero cêntimos" - result = result.replace(' e zero cêntimos', '') - return result diff --git a/fun_text_processing/num2words/num2words/lang_PT_BR.py b/fun_text_processing/num2words/num2words/lang_PT_BR.py deleted file mode 100644 index 834d11ce3..000000000 --- a/fun_text_processing/num2words/num2words/lang_PT_BR.py +++ /dev/null @@ -1,98 +0,0 @@ -# -*- coding: utf-8 -*- -from __future__ import division, unicode_literals - -import re - -from . import lang_PT - - -class Num2Word_PT_BR(lang_PT.Num2Word_PT): - def set_high_numwords(self, high): - max = 3 + 3*len(high) - for word, n in zip(high, range(max, 3, -3)): - self.cards[10**n] = word + "ilhão" - - def setup(self): - super(Num2Word_PT_BR, self).setup() - - self.low_numwords[1] = 'dezenove' - self.low_numwords[3] = 'dezessete' - self.low_numwords[4] = 'dezesseis' - - self.thousand_separators = { - 3: "milésimo", - 6: "milionésimo", - 9: "bilionésimo", - 12: "trilionésimo", - 15: "quadrilionésimo" - } - - def merge(self, curr, next): - ctext, cnum, ntext, nnum = curr + next - - if cnum == 1: - if nnum < 1000000: - return next - ctext = "um" - elif cnum == 100 and not nnum == 1000: - ctext = "cento" - - if nnum < cnum: - return ("%s e %s" % (ctext, ntext), cnum + nnum) - - elif (not nnum % 1000000) and cnum > 1: - ntext = ntext[:-4] + "lhões" - - if nnum == 100: - ctext = self.hundreds[cnum] - ntext = "" - - else: - ntext = " " + ntext - - return (ctext + ntext, cnum * nnum) - - def to_cardinal(self, value): - result = lang_PT.Num2Word_EU.to_cardinal(self, value) - - # Transforms "mil E cento e catorze reais" into "mil, cento e catorze - # reais" - for ext in ( - 'mil', 'milhão', 'milhões', 'bilhão', 'bilhões', - 'trilhão', 'trilhões', 'quatrilhão', 'quatrilhões'): - if re.match('.*{} e \\w*ento'.format(ext), result): - result = result.replace( - '{} e'.format(ext), '{},'.format(ext), 1 - ) - - return result - - def to_currency(self, val, longval=True): - integer_part, decimal_part = ('%.2f' % val).split('.') - - result = self.to_cardinal(int(integer_part)) - - appended_currency = False - for ext in ( - 'milhão', 'milhões', 'bilhão', 'bilhões', - 'trilhão', 'trilhões', 'quatrilhão', 'quatrilhões'): - if result.endswith(ext): - result += ' de reais' - appended_currency = True - - if result in ['um', 'menos um']: - result += ' real' - appended_currency = True - if not appended_currency: - result += ' reais' - - if int(decimal_part): - cents = self.to_cardinal(int(decimal_part)) - result += ' e ' + cents - - if cents == 'um': - result += ' centavo' - else: - result += ' centavos' - - return result diff --git a/fun_text_processing/num2words/num2words/lang_RO.py b/fun_text_processing/num2words/num2words/lang_RO.py deleted file mode 100644 index 8b539c301..000000000 --- a/fun_text_processing/num2words/num2words/lang_RO.py +++ /dev/null @@ -1,158 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import division, print_function, unicode_literals - -from . import lang_EU - - -class Num2Word_RO(lang_EU.Num2Word_EU): - GIGA_SUFFIX = "iliard/e" - MEGA_SUFFIX = "ilion" - # inflection for mi/billion follows different rule - MEGA_SUFFIX_I = "ilioane" - GIGA_SUFFIX_I = "iliarde" - - def setup(self): - super(Num2Word_RO, self).setup() - - self.negword = "minus " - self.pointword = "virgulă" - self.exclude_title = ["și", "virgulă", "minus"] - self.errmsg_toobig = ( - "Numărul e prea mare pentru a fi convertit în cuvinte." - ) - self.mid_numwords = [(1000, "mie/i"), (100, "sută/e"), - (90, "nouăzeci"), (80, "optzeci"), - (70, "șaptezeci"), (60, "șaizeci"), - (50, "cincizeci"), (40, "patruzeci"), - (30, "treizeci")] - self.low_numwords = ["douăzeci", "nouăsprezece", "optsprezece", - "șaptesprezece", "șaisprezece", "cincisprezece", - "paisprezece", "treisprezece", "doisprezece", - "unsprezece", "zece", "nouă", "opt", "șapte", - "șase", "cinci", "patru", "trei", "doi", - "unu", "zero"] - self.gen_numwords = ["", "o", "două", "trei", "patru", "cinci", - "șase", "șapte", "opt", "nouă"] - self.gen_numwords_n = ["", "un", "două", "trei", "patru", "cinci", - "șase", "șapte", "opt", "nouă"] - self.numwords_inflections = { - 100: self.gen_numwords, - 1000: self.gen_numwords, - 1000000: self.gen_numwords_n, - 1000000000: self.gen_numwords_n - } - self.ords = {"unu": "primul", - "doi": "al doilea", - "three": "al treilea", - "cinci": "al cincilea", - "opt": "al optulea", - "nouă": "al nouălea", - "doisprezece": "al doisprezecelea"} - - def merge(self, lpair, rpair): - ltext, lnum = lpair - rtext, rnum = rpair - rtext_i = self.inflect(rnum, rtext, lnum) - if 1 <= lnum < 10: - if rnum not in self.numwords_inflections: - return (rtext, rnum) - else: - rtext_i = self.inflect(lnum * rnum, rtext, lnum) - lresult = (self.numwords_inflections[rnum][lnum], rtext_i) - return ("%s %s" % lresult, rnum) - elif 10 < lnum < 100: - if lnum % 10 == 0: - if rnum in self.numwords_inflections: - rtext_i = self.inflect(lnum * rnum, rtext, lnum) - return ("%s %s" % (ltext, rtext_i), lnum * rnum) - else: - return ("%s și %s" % (ltext, rtext), lnum + rnum) - else: - rtext_i = self.inflect(lnum * rnum, rtext, lnum) - ltext_i = ltext if lnum % 10 != 2 \ - else ltext.replace("doi", "două") - return ("%s %s" % (ltext_i, rtext_i), lnum * rnum) - else: - if rnum in self.numwords_inflections: - rtext_i = self.inflect(lnum * rnum, rtext, lnum) - return ("%s %s" % (ltext, rtext_i), lnum * rnum) - - def to_ordinal(self, value): - self.verify_ordinal(value) - if value == 1: - return "primul" - else: - value = self.to_cardinal(value) - return "al %slea" % (value) - - def to_ordinal_num(self, value): - self.verify_ordinal(value) - if value == 1: - return "1-ul" - return "al %s-lea" % (value) - - def pluralize(self, n, forms): - if n == 1: - form = 0 - elif n == 0 or (n % 100 > 0 and n % 100 < 20): - form = 1 - else: - form = 2 - return forms[form] - - def inflect(self, value, text, side_effect=-1): - text = text.split("/") - result = text[0] - if len(text) > 1: - forms = [ - text[0], - text[0][:-1] + text[1], - "de " + text[0][:-1] + text[1] - ] - result = self.pluralize(side_effect, forms) - # mega inflections are different - if side_effect > 1 and result.endswith(self.MEGA_SUFFIX): - result = result.replace(self.MEGA_SUFFIX, self.MEGA_SUFFIX_I) - elif side_effect > 1 and result.endswith("iliare"): - result = result.replace("iliare", self.GIGA_SUFFIX_I) - return result - - def to_currency(self, val, currency="RON", cents=False, separator=" și", - adjective=False): - # romanian currency has a particularity for numeral: one - self.gen_numwords[1] = "una" - result = super(Num2Word_RO, self).to_currency( - int(round(val*100)), - currency, - True, - separator, - adjective - ) - self.gen_numwords[1] = "o" # revert numeral - return result.replace( - "unu leu", "un leu" - ).replace( - "unu ban", "un ban" - ).replace( - # if the romanian low text is 0, it is not usually printed - separator + " zero bani", "" - ) - - def to_year(self, val, suffix=None, longval=True): - result = super(Num2Word_RO, self).to_year( - val, - longval=longval - ) - # for years we want the era negation e.g. B.C., in our case - # it's î.Hr. or î.e.n - if result.startswith(self.negword): - result = result.replace(self.negword, "") - suffix = "î.Hr." if not suffix else suffix - if suffix: - result = "".join([ - result, - " ", - suffix - ]) - return result diff --git a/fun_text_processing/num2words/num2words/lang_RU.py b/fun_text_processing/num2words/num2words/lang_RU.py deleted file mode 100644 index 8cc9b15d8..000000000 --- a/fun_text_processing/num2words/num2words/lang_RU.py +++ /dev/null @@ -1,231 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import unicode_literals - -from .base import Num2Word_Base -from .utils import get_digits, splitbyx - -ZERO = ('ноль',) - -ONES_FEMININE = { - 1: ('одна',), - 2: ('две',), - 3: ('три',), - 4: ('четыре',), - 5: ('пять',), - 6: ('шесть',), - 7: ('семь',), - 8: ('восемь',), - 9: ('девять',), -} - -ONES = { - 1: ('один',), - 2: ('два',), - 3: ('три',), - 4: ('четыре',), - 5: ('пять',), - 6: ('шесть',), - 7: ('семь',), - 8: ('восемь',), - 9: ('девять',), -} - -TENS = { - 0: ('десять',), - 1: ('одиннадцать',), - 2: ('двенадцать',), - 3: ('тринадцать',), - 4: ('четырнадцать',), - 5: ('пятнадцать',), - 6: ('шестнадцать',), - 7: ('семнадцать',), - 8: ('восемнадцать',), - 9: ('девятнадцать',), -} - -TWENTIES = { - 2: ('двадцать',), - 3: ('тридцать',), - 4: ('сорок',), - 5: ('пятьдесят',), - 6: ('шестьдесят',), - 7: ('семьдесят',), - 8: ('восемьдесят',), - 9: ('девяносто',), -} - -HUNDREDS = { - 1: ('сто',), - 2: ('двести',), - 3: ('триста',), - 4: ('четыреста',), - 5: ('пятьсот',), - 6: ('шестьсот',), - 7: ('семьсот',), - 8: ('восемьсот',), - 9: ('девятьсот',), -} - -THOUSANDS = { - 1: ('тысяча', 'тысячи', 'тысяч'), # 10^3 - 2: ('миллион', 'миллиона', 'миллионов'), # 10^6 - 3: ('миллиард', 'миллиарда', 'миллиардов'), # 10^9 - 4: ('триллион', 'триллиона', 'триллионов'), # 10^12 - 5: ('квадриллион', 'квадриллиона', 'квадриллионов'), # 10^15 - 6: ('квинтиллион', 'квинтиллиона', 'квинтиллионов'), # 10^18 - 7: ('секстиллион', 'секстиллиона', 'секстиллионов'), # 10^21 - 8: ('септиллион', 'септиллиона', 'септиллионов'), # 10^24 - 9: ('октиллион', 'октиллиона', 'октиллионов'), # 10^27 - 10: ('нониллион', 'нониллиона', 'нониллионов'), # 10^30 -} - - -class Num2Word_RU(Num2Word_Base): - CURRENCY_FORMS = { - 'RUB': ( - ('рубль', 'рубля', 'рублей'), ('копейка', 'копейки', 'копеек') - ), - 'EUR': ( - ('евро', 'евро', 'евро'), ('цент', 'цента', 'центов') - ), - 'USD': ( - ('доллар', 'доллара', 'долларов'), ('цент', 'цента', 'центов') - ), - 'UAH': ( - ('гривна', 'гривны', 'гривен'), ('копейка', 'копейки', 'копеек') - ), - 'KZT': ( - ('тенге', 'тенге', 'тенге'), ('тиын', 'тиына', 'тиынов') - ), - } - - def setup(self): - self.negword = "минус" - self.pointword = "запятая" - self.ords = {"ноль": "нулевой", - "один": "первый", - "два": "второй", - "три": "третий", - "четыре": "четвертый", - "пять": "пятый", - "шесть": "шестой", - "семь": "седьмой", - "восемь": "восьмой", - "девять": "девятый", - "сто": "сотый"} - self.ords_feminine = {"один": "", - "одна": "", - "две": "двух", - "три": "трёх", - "четыре": "четырёх", - "пять": "пяти", - "шесть": "шести", - "семь": "семи", - "восемь": "восьми", - "девять": "девяти"} - - def to_cardinal(self, number): - n = str(number).replace(',', '.') - if '.' in n: - left, right = n.split('.') - leading_zero_count = len(right) - len(right.lstrip('0')) - decimal_part = ((ZERO[0] + ' ') * leading_zero_count + - self._int2word(int(right))) - return u'%s %s %s' % ( - self._int2word(int(left)), - self.pointword, - decimal_part - ) - else: - return self._int2word(int(n)) - - def pluralize(self, n, forms): - if n % 100 < 10 or n % 100 > 20: - if n % 10 == 1: - form = 0 - elif 5 > n % 10 > 1: - form = 1 - else: - form = 2 - else: - form = 2 - return forms[form] - - def to_ordinal(self, number): - self.verify_ordinal(number) - outwords = self.to_cardinal(number).split(" ") - lastword = outwords[-1].lower() - try: - if len(outwords) > 1: - if outwords[-2] in self.ords_feminine: - outwords[-2] = self.ords_feminine.get( - outwords[-2], outwords[-2]) - elif outwords[-2] == 'десять': - outwords[-2] = outwords[-2][:-1] + 'и' - if len(outwords) == 3: - if outwords[-3] in ['один', 'одна']: - outwords[-3] = '' - lastword = self.ords[lastword] - except KeyError: - if lastword[:-3] in self.ords_feminine: - lastword = self.ords_feminine.get( - lastword[:-3], lastword) + "сотый" - elif lastword[-1] == "ь" or lastword[-2] == "т": - lastword = lastword[:-1] + "ый" - elif lastword[-1] == "к": - lastword = lastword + "овой" - elif lastword[-5:] == "десят": - lastword = lastword.replace('ь', 'и') + 'ый' - elif lastword[-2] == "ч" or lastword[-1] == "ч": - if lastword[-2] == "ч": - lastword = lastword[:-1] + "ный" - if lastword[-1] == "ч": - lastword = lastword + "ный" - elif lastword[-1] == "н" or lastword[-2] == "н": - lastword = lastword[:lastword.rfind('н') + 1] + "ный" - elif lastword[-1] == "д" or lastword[-2] == "д": - lastword = lastword[:lastword.rfind('д') + 1] + "ный" - outwords[-1] = self.title(lastword) - return " ".join(outwords).strip() - - def _money_verbose(self, number, currency): - return self._int2word(number, currency == 'UAH') - - def _cents_verbose(self, number, currency): - return self._int2word(number, currency in ('UAH', 'RUB')) - - def _int2word(self, n, feminine=False): - if n < 0: - return ' '.join([self.negword, self._int2word(abs(n))]) - - if n == 0: - return ZERO[0] - - words = [] - chunks = list(splitbyx(str(n), 3)) - i = len(chunks) - for x in chunks: - i -= 1 - - if x == 0: - continue - - n1, n2, n3 = get_digits(x) - - if n3 > 0: - words.append(HUNDREDS[n3][0]) - - if n2 > 1: - words.append(TWENTIES[n2][0]) - - if n2 == 1: - words.append(TENS[n1][0]) - elif n1 > 0: - ones = ONES_FEMININE if i == 1 or feminine and i == 0 else ONES - words.append(ones[n1][0]) - - if i > 0: - words.append(self.pluralize(x, THOUSANDS[i])) - - return ' '.join(words) diff --git a/fun_text_processing/num2words/num2words/lang_SI.py b/fun_text_processing/num2words/num2words/lang_SI.py deleted file mode 100644 index a18c608b2..000000000 --- a/fun_text_processing/num2words/num2words/lang_SI.py +++ /dev/null @@ -1,108 +0,0 @@ -# -*- coding: utf-8 -*- - -# Sinhala number to words - -from __future__ import division, print_function, unicode_literals - -from . import lang_EU - - -class Num2Word_SI(lang_EU.Num2Word_EU): - # GIGA_SUFFIX = "iljarder" - # MEGA_SUFFIX = "iljoner" - - # def set_high_numwords(self, high): - # cap = 3 + 6 * len(high) - - # for word, n in zip(high, range(cap, 3, -6)): - # if self.GIGA_SUFFIX: - # self.cards[10 ** n] = word + self.GIGA_SUFFIX - - # if self.MEGA_SUFFIX: - # self.cards[10 ** (n - 3)] = word + self.MEGA_SUFFIX - def set_high_numwords(self, high): - max = 3 * len(high) - for word, n in zip(high, range(max, 0, -3)): - # print(word[0],word[1],n) - self.cards[10 ** n] = word[1] - - def setup(self): - super(Num2Word_SI, self).setup() - - self.negword = "අඩු " - self.pointword = "ලක්ෂ්යය" - self.exclude_title = ["හා", "ලක්ෂ්යය", "අඩු"] - - self.high_numwords = [(10000000000, 'බිලියන දහයකි'), - (1000000000, 'බිලියනයක්'), - (100000000, 'මිලියන සියයක්'), - (10000000,'දස මිලියන'), - (1000000, 'මිලියනයක්'), - (100000, 'එක් ලක්ෂයක්'), - (1000, 'දහසක්') ] - self.mid_numwords = [(1000, "දහසක්"), (100, "සියය"), - (90, "අනූවක්"), (80, "අසූව"), (70, "හැත්තෑ"), - (60, "හැට"), (50, "පනස්"), (40, "හතළිහක්"), - (30, "තිස්")] - - self.low_numwords = ['විස්සක්','දහනවය','දහඅට','දාහත','දහසය','පහළොව','දහහතර','දහතුන','දොළොස්','එකොළොස්','දස','නවය','අට','හත','හය','පහ','හතර','තුන්','දෙක','එක','ශුන්ය'] - - self.ords = {'එක': 'පළමුවන', - 'දෙක': 'දෙවැනි', - 'තුන්': 'තෙවන', - 'හතර': 'හතරවන', - 'පහ': 'පස්වන', - 'හය': 'හය වන', - 'හත': 'හත්වන', - 'අට': 'අටවැනි', - 'නවය': 'නවවැනි', - 'දස': 'දහවන'} - - - def merge(self, lpair, rpair): - ltext, lnum = lpair - rtext, rnum = rpair - if lnum == 1 and rnum < 100: - return (rtext, rnum) - elif 100 > lnum > rnum: - return ("%s %s" % (ltext, rtext), lnum + rnum) - elif lnum >= 100 > rnum: - return ("%s %s" % (ltext, rtext), lnum + rnum) - elif rnum > lnum: - return ("%s %s" % (ltext, rtext), lnum * rnum) - return ("%s %s" % (ltext, rtext), lnum + rnum) - - def to_ordinal_num(self, value): - self.verify_ordinal(value) - return "%s %s" % (value, self.to_ordinal(value)) - - def to_ordinal(self, value): - self.verify_ordinal(value) - try: - ordinal_word = self.ords[value] - except KeyError: - ordinal_word = value #TODO: check nepali ordinal word suffix - return ordinal_word - - - def to_year(self, val, suffix=None, longval=True): - if val < 0: - val = abs(val) - suffix = 'සියවසට පෙර ' if not suffix else suffix - high, low = (val // 100, val % 100) - # If year is 00XX, X00X, or beyond 9999, go cardinal. - if (high == 0 - or (high % 10 == 0 and low < 10) - or high >= 100): - valtext = self.to_cardinal(val) - else: - hightext = self.to_cardinal(high) - if low == 0: - lowtext = "සියයක්" #hundred - elif low < 10: - lowtext = "%s" % self.to_cardinal(low) - else: - lowtext = self.to_cardinal(low) - valtext = "%s %s" % (hightext, lowtext) - return (valtext if not suffix - else "%s %s" % (valtext, suffix)) diff --git a/fun_text_processing/num2words/num2words/lang_SL.py b/fun_text_processing/num2words/num2words/lang_SL.py deleted file mode 100644 index fb0e28766..000000000 --- a/fun_text_processing/num2words/num2words/lang_SL.py +++ /dev/null @@ -1,169 +0,0 @@ -# -*- coding: utf-8 -*- -# Copyright (c) 2003, Taro Ogawa. All Rights Reserved. -# Copyright (c) 2013, Savoir-faire Linux inc. All Rights Reserved. -# Copyright (c) 2015, Blaz Bregar. All Rights Reserved. - -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, -# MA 02110-1301 USA - -from __future__ import unicode_literals - -from .lang_EU import Num2Word_EU - - -class Num2Word_SL(Num2Word_EU): - GIGA_SUFFIX = "ilijard" - MEGA_SUFFIX = "ilijon" - - def setup(self): - super(Num2Word_SL, self).setup() - - self.negword = "minus " - self.pointword = "celih" - self.errmsg_nonnum = "Only numbers may be converted to words." - self.errmsg_toobig = "Number is too large to convert to words." - self.exclude_title = [] - - self.mid_numwords = [(1000, "tisoč"), (900, "devetsto"), - (800, "osemsto"), (700, "sedemsto"), - (600, "šeststo"), (500, "petsto"), - (400, "štiristo"), (300, "tristo"), - (200, "dvesto"), (100, "sto"), - (90, "devetdeset"), (80, "osemdeset"), - (70, "sedemdeset"), (60, "šestdeset"), - (50, "petdeset"), (40, "štirideset"), - (30, "trideset")] - self.low_numwords = ["dvajset", "devetnajst", "osemnajst", - "sedemnajst", "šestnajst", "petnajst", - "štirinajst", "trinajst", "dvanajst", - "enajst", "deset", "devet", "osem", "sedem", - "šest", "pet", "štiri", "tri", "dve", "ena", - "nič"] - self.ords = {"ena": "prv", - "dve": "drug", - "tri": "tretj", - "štiri": "četrt", - "sedem": "sedm", - "osem": "osm", - "sto": "stot", - "tisoč": "tisoč", - "milijon": "milijont" - } - self.ordflag = False - - def merge(self, curr, next): - ctext, cnum, ntext, nnum = curr + next - - if ctext.endswith("dve") and self.ordflag and nnum <= 1000000: - ctext = ctext[:len(ctext)-1] + "a" - - if ctext == "dve" and not self.ordflag and nnum < 1000000000: - ctext = "dva" - - if (ctext.endswith("tri") or ctext.endswith("štiri")) and\ - nnum == 1000000 and not self.ordflag: - if ctext.endswith("štiri"): - ctext = ctext[:-1] - ctext = ctext + "je" - - if cnum >= 20 and cnum < 100 and nnum == 2: - ntext = "dva" - - if ctext.endswith("ena") and nnum >= 1000: - ctext = ctext[0:-1] - - if cnum == 1: - if nnum < 10**6 or self.ordflag: - return next - ctext = "" - - if nnum > cnum: - if nnum >= 10**6: - if self.ordflag: - ntext += "t" - - elif cnum == 2: - if ntext.endswith("d"): - ntext += "i" - else: - ntext += "a" - - elif 2 < cnum < 5: - if ntext.endswith("d"): - ntext += "e" - elif not ntext.endswith("d"): - ntext += "i" - - elif ctext.endswith("en"): - if ntext.endswith("d") or ntext.endswith("n"): - ntext += "" - - elif ctext.endswith("dve") and ntext.endswith("n"): - ctext = ctext[:-1] + "a" - ntext += "a" - - elif ctext.endswith("je") and ntext.endswith("n"): - ntext += "i" - - else: - if ntext.endswith("d"): - ntext += "a" - elif ntext.endswith("n"): - ntext += "" - elif ntext.endswith("d"): - ntext += "e" - else: - ntext += "ov" - - if nnum >= 10**2 and self.ordflag is False and ctext: - ctext += " " - - val = cnum * nnum - else: - if nnum < 10 < cnum < 100: - ntext, ctext = ctext, ntext + "in" - elif cnum >= 10**2 and self.ordflag is False: - ctext += " " - val = cnum + nnum - - word = ctext + ntext - return (word, val) - - def to_ordinal(self, value): - self.verify_ordinal(value) - self.ordflag = True - outword = self.to_cardinal(value) - self.ordflag = False - for key in self.ords: - if outword.endswith(key): - outword = outword[:len(outword) - len(key)] + self.ords[key] - break - return outword + "i" - - # Is this correct?? - def to_ordinal_num(self, value): - self.verify_ordinal(value) - return str(value) + "." - - def to_currency(self, val, longval=True, old=False): - if old: - return self.to_splitnum(val, hightxt="evro/a/v", - lowtxt="stotin/a/i/ov", - jointxt="in", longval=longval) - return super(Num2Word_SL, self).to_currency(val, jointxt="in", - longval=longval) - - def to_year(self, val, longval=True): - if not (val//100) % 10: - return self.to_cardinal(val) - return self.to_splitnum(val, hightxt="hundert", longval=longval) diff --git a/fun_text_processing/num2words/num2words/lang_SR.py b/fun_text_processing/num2words/num2words/lang_SR.py deleted file mode 100644 index cd294897f..000000000 --- a/fun_text_processing/num2words/num2words/lang_SR.py +++ /dev/null @@ -1,204 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import unicode_literals - -from .base import Num2Word_Base -from .currency import parse_currency_parts, prefix_currency -from .utils import get_digits, splitbyx - -ZERO = ('nula',) - -ONES = { - 1: ('jedan', 'jedna'), - 2: ('dva', 'dve'), - 3: ('tri', 'tri'), - 4: ('četiri', 'četiri'), - 5: ('pet', 'pet'), - 6: ('šest', 'šest'), - 7: ('sedam', 'sedam'), - 8: ('osam', 'osam'), - 9: ('devet', 'devet'), -} - -TENS = { - 0: ('deset',), - 1: ('jedanaest',), - 2: ('dvanaest',), - 3: ('trinaest',), - 4: ('četrnaest',), - 5: ('petnaest',), - 6: ('šesnaest',), - 7: ('sedamnaest',), - 8: ('osamnaest',), - 9: ('devetnaest',), -} - -TWENTIES = { - 2: ('dvadeset',), - 3: ('trideset',), - 4: ('četrdeset',), - 5: ('pedeset',), - 6: ('šezdeset',), - 7: ('sedamdeset',), - 8: ('osamdeset',), - 9: ('devedeset',), -} - -HUNDREDS = { - 1: ('sto',), - 2: ('dvesta',), - 3: ('trista',), - 4: ('četristo',), - 5: ('petsto',), - 6: ('šesto',), - 7: ('sedamsto',), - 8: ('osamsto',), - 9: ('devetsto',), -} - -SCALE = { - 0: ('', '', '', False), - 1: ('hiljada', 'hiljade', 'hiljada', True), # 10^3 - 2: ('milion', 'miliona', 'miliona', False), # 10^6 - 3: ('bilion', 'biliona', 'biliona', False), # 10^9 - 4: ('trilion', 'triliona', 'triliona', False), # 10^12 - 5: ('kvadrilion', 'kvadriliona', 'kvadriliona', False), # 10^15 - 6: ('kvintilion', 'kvintiliona', 'kvintiliona', False), # 10^18 - 7: ('sekstilion', 'sekstiliona', 'sekstiliona', False), # 10^21 - 8: ('septilion', 'septiliona', 'septiliona', False), # 10^24 - 9: ('oktilion', 'oktiliona', 'oktiliona', False), # 10^27 - 10: ('nonilion', 'noniliona', 'noniliona', False), # 10^30 -} - - -class Num2Word_SR(Num2Word_Base): - CURRENCY_FORMS = { - 'RUB': ( - ('rublja', 'rublje', 'rublji', True), - ('kopejka', 'kopejke', 'kopejki', True) - ), - 'EUR': ( - ('evro', 'evra', 'evra', False), - ('cent', 'centa', 'centi', False) - ), - 'RSD': ( - ('dinar', 'dinara', 'dinara', False), - ('para', 'pare', 'para', True) - ), - } - - def setup(self): - self.negword = "minus" - self.pointword = "zapeta" - - def to_cardinal(self, number, feminine=False): - n = str(number).replace(',', '.') - if '.' in n: - left, right = n.split('.') - leading_zero_count = len(right) - len(right.lstrip('0')) - decimal_part = ((ZERO[0] + ' ') * leading_zero_count + - self._int2word(int(right), feminine)) - return u'%s %s %s' % ( - self._int2word(int(left), feminine), - self.pointword, - decimal_part - ) - else: - return self._int2word(int(n), feminine) - - def pluralize(self, number, forms): - if number % 100 < 10 or number % 100 > 20: - if number % 10 == 1: - form = 0 - elif 1 < number % 10 < 5: - form = 1 - else: - form = 2 - else: - form = 2 - return forms[form] - - def to_ordinal(self, number): - raise NotImplementedError() - - def _cents_verbose(self, number, currency): - return self._int2word( - number, - self.CURRENCY_FORMS[currency][1][-1] - ) - - def _int2word(self, number, feminine=False): - if number < 0: - return ' '.join([self.negword, self._int2word(abs(number))]) - - if number == 0: - return ZERO[0] - - words = [] - chunks = list(splitbyx(str(number), 3)) - chunk_len = len(chunks) - for chunk in chunks: - chunk_len -= 1 - digit_right, digit_mid, digit_left = get_digits(chunk) - - if digit_left > 0: - words.append(HUNDREDS[digit_left][0]) - - if digit_mid > 1: - words.append(TWENTIES[digit_mid][0]) - - if digit_mid == 1: - words.append(TENS[digit_right][0]) - elif digit_right > 0: - is_feminine = feminine or SCALE[chunk_len][-1] - gender_idx = int(is_feminine) - words.append( - ONES[digit_right][gender_idx] - ) - - if chunk_len > 0 and chunk != 0: - words.append(self.pluralize(chunk, SCALE[chunk_len])) - - return ' '.join(words) - - def to_currency(self, val, currency='EUR', cents=True, separator=',', - adjective=False): - """ - Args: - val: Numeric value - currency (str): Currency code - cents (bool): Verbose cents - separator (str): Cent separator - adjective (bool): Prefix currency name with adjective - Returns: - str: Formatted string - - """ - left, right, is_negative = parse_currency_parts(val) - - try: - cr1, cr2 = self.CURRENCY_FORMS[currency] - - except KeyError: - raise NotImplementedError( - 'Currency code "%s" not implemented for "%s"' % - (currency, self.__class__.__name__)) - - if adjective and currency in self.CURRENCY_ADJECTIVES: - cr1 = prefix_currency( - self.CURRENCY_ADJECTIVES[currency], - cr1 - ) - - minus_str = "%s " % self.negword if is_negative else "" - cents_str = self._cents_verbose(right, currency) \ - if cents else self._cents_terse(right, currency) - - return u'%s%s %s%s %s %s' % ( - minus_str, - self.to_cardinal(left, feminine=cr1[-1]), - self.pluralize(left, cr1), - separator, - cents_str, - self.pluralize(right, cr2) - ) diff --git a/fun_text_processing/num2words/num2words/lang_SV.py b/fun_text_processing/num2words/num2words/lang_SV.py deleted file mode 100644 index f07093340..000000000 --- a/fun_text_processing/num2words/num2words/lang_SV.py +++ /dev/null @@ -1,102 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import division, print_function, unicode_literals - -from . import lang_EU - - -class Num2Word_SV(lang_EU.Num2Word_EU): - GIGA_SUFFIX = "iljarder" - MEGA_SUFFIX = "iljoner" - - def set_high_numwords(self, high): - cap = 3 + 6 * len(high) - - for word, n in zip(high, range(cap, 3, -6)): - if self.GIGA_SUFFIX: - self.cards[10 ** n] = word + self.GIGA_SUFFIX - - if self.MEGA_SUFFIX: - self.cards[10 ** (n - 3)] = word + self.MEGA_SUFFIX - - def setup(self): - super(Num2Word_SV, self).setup() - - self.negword = "minus " - self.pointword = "komma" - self.exclude_title = ["och", "komma", "minus"] - - self.mid_numwords = [(1000, "tusen"), (100, "hundra"), - (90, "nittio"), (80, "åttio"), (70, "sjuttio"), - (60, "sextio"), (50, "femtio"), (40, "förtio"), - (30, "trettio")] - self.low_numwords = ["tjugo", "nitton", "arton", "sjutton", - "sexton", "femton", "fjorton", "tretton", - "tolv", "elva", "tio", "nio", "åtta", - "sju", "sex", "fem", "fyra", "tre", "två", - "ett", "noll"] - self.ords = {"noll": "nollte", - "ett": "första", - "två": "andra", - "tre": "tredje", - "fyra": "fjärde", - "fem": "femte", - "sex": "sjätte", - "sju": "sjunde", - "åtta": "åttonde", - "nio": "nionde", - "tio": "tionde", - "elva": "elfte", - "tolv": "tolfte", - "tjugo": "tjugonde"} - - def merge(self, lpair, rpair): - ltext, lnum = lpair - rtext, rnum = rpair - if lnum == 1 and rnum < 100: - return (rtext, rnum) - elif 100 > lnum > rnum: - return ("%s%s" % (ltext, rtext), lnum + rnum) - elif lnum >= 100 > rnum: - return ("%s%s" % (ltext, rtext), lnum + rnum) - elif rnum >= 1000000 and lnum == 1: - return ("%s %s" % ('en', rtext[:-2]), lnum + rnum) - elif rnum >= 1000000 and lnum > 1: - return ("%s %s" % (ltext, rtext), lnum + rnum) - elif rnum > lnum: - return ("%s%s" % (ltext, rtext), lnum * rnum) - return ("%s %s" % (ltext, rtext), lnum + rnum) - - def to_ordinal(self, value): - self.verify_ordinal(value) - outwords = self.to_cardinal(value).split(" ") - lastword = outwords[-1] - ending_length = 0 - try: - lastword_ending = self.ords[lastword[-4:]] - ending_length = 4 - except KeyError: - try: - lastword_ending = self.ords[lastword[-3:]] - ending_length = 3 - except KeyError: - lastword_ending = "de" - if lastword_ending == 'de': - lastword_first_part = self.title(lastword)[:] - else: - lastword_first_part = self.title(lastword)[:-ending_length] - lastword_correct = lastword_first_part + lastword_ending - outwords[-1] = lastword_correct - return " ".join(outwords) - - def to_ordinal_num(self, value): - raise NotImplementedError( - "'ordinal_num' is not implemented for swedish language") - - def to_year(self, val, longval=True): - raise NotImplementedError( - "'year' is not implemented for swedish language") - - def to_currency(self, val, longval=True): - raise NotImplementedError( - "'currency' is not implemented for swedish language") diff --git a/fun_text_processing/num2words/num2words/lang_SW.py b/fun_text_processing/num2words/num2words/lang_SW.py deleted file mode 100644 index 22c58fa25..000000000 --- a/fun_text_processing/num2words/num2words/lang_SW.py +++ /dev/null @@ -1,108 +0,0 @@ -# -*- coding: utf-8 -*- - -# Swahili number to words - -from __future__ import division, print_function, unicode_literals - -from . import lang_EU - - -class Num2Word_SW(lang_EU.Num2Word_EU): - # GIGA_SUFFIX = "iljarder" - # MEGA_SUFFIX = "iljoner" - - # def set_high_numwords(self, high): - # cap = 3 + 6 * len(high) - - # for word, n in zip(high, range(cap, 3, -6)): - # if self.GIGA_SUFFIX: - # self.cards[10 ** n] = word + self.GIGA_SUFFIX - - # if self.MEGA_SUFFIX: - # self.cards[10 ** (n - 3)] = word + self.MEGA_SUFFIX - def set_high_numwords(self, high): - max = 3 * len(high) - for word, n in zip(high, range(max, 0, -3)): - # print(word[0],word[1],n) - self.cards[10 ** n] = word[1] - - def setup(self): - super(Num2Word_SW, self).setup() - - self.negword = "kuondoa " - self.pointword = "hatua" - self.exclude_title = ["na", "hatua", "kuondoa"] - - self.high_numwords = [(10000000000, 'bilioni kumi'), - (1000000000, 'bilioni moja'), - (100000000, 'milioni mia moja'), - (10000000,'milioni kumi'), - (1000000, 'milioni moja'), - (100000, 'laki moja'), - (1000, 'elfu moja') ] - self.mid_numwords = [(1000, "elfu moja"), (100, "mia moja"), - (90, "tisini"), (80, "themanini"), (70, "sabini"), - (60, "sitini"), (50, "hamsini"), (40, "arobaini"), - (30, "thelathini")] - - self.low_numwords = ['sufuri', 'moja', 'mbili', 'tatu', 'nne', 'tano', 'sita', 'saba', 'nane', 'tisa', 'kumi', 'kumi na moja', 'kumi na mbili', 'kumi na tatu', 'kumi na nne', 'kumi na tano', 'kumi na sita', 'kumi na saba', 'kumi na nane', 'kumi na tisa', 'ishirini'] - - self.ords = {"moja":"kwanza", - "mbili":"pili", - "tatu":"cha tatu", - "nne":"nne", - "tano":"tano", - "sita":"ya sita", - "saba":"ya saba", - "nane":"ya nane", - "tisa":"ya tisa", - "kumi":"ya kumi"} - - - def merge(self, lpair, rpair): - ltext, lnum = lpair - rtext, rnum = rpair - if lnum == 1 and rnum < 100: - return (rtext, rnum) - elif 100 > lnum > rnum: - return ("%s %s" % (ltext, rtext), lnum + rnum) - elif lnum >= 100 > rnum: - return ("%s %s" % (ltext, rtext), lnum + rnum) - elif rnum > lnum: - return ("%s %s" % (ltext, rtext), lnum * rnum) - return ("%s %s" % (ltext, rtext), lnum + rnum) - - def to_ordinal_num(self, value): - self.verify_ordinal(value) - return "%s %s" % (value, self.to_ordinal(value)) - - def to_ordinal(self, value): - self.verify_ordinal(value) - try: - ordinal_word = self.ords[value] - except KeyError: - ordinal_word = value #TODO: check nepali ordinal word suffix - return ordinal_word - - - def to_year(self, val, suffix=None, longval=True): - if val < 0: - val = abs(val) - suffix = 'kabla ya karne ' if not suffix else suffix - high, low = (val // 100, val % 100) - # If year is 00XX, X00X, or beyond 9999, go cardinal. - if (high == 0 - or (high % 10 == 0 and low < 10) - or high >= 100): - valtext = self.to_cardinal(val) - else: - hightext = self.to_cardinal(high) - if low == 0: - lowtext = "mia" - elif low < 10: - lowtext = "%s" % self.to_cardinal(low) - else: - lowtext = self.to_cardinal(low) - valtext = "%s %s" % (hightext, lowtext) - return (valtext if not suffix - else "%s %s" % (valtext, suffix)) diff --git a/fun_text_processing/num2words/num2words/lang_TE.py b/fun_text_processing/num2words/num2words/lang_TE.py deleted file mode 100644 index 274d31254..000000000 --- a/fun_text_processing/num2words/num2words/lang_TE.py +++ /dev/null @@ -1,169 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import unicode_literals - -from .lang_EU import Num2Word_EU - - -class Num2Word_TE(Num2Word_EU): - def set_high_numwords(self, high): - for n, word in self.high_numwords: - self.cards[10**n] = word - - def setup(self): - self.low_numwords = [ - "తొంభై తొమ్మిది", - "తొంభై ఎనిమిది", - "తొంభై ఏడు", - "తొంభై ఆరు", - "తొంభై అయిదు", - "తొంభై నాలుగు", - "తొంభై మూడు", - "తొంభై రెండు", - "తొంభై ఒకటి", - "తొంభై", - "ఎనభై తొమ్మిది", - "ఎనభై ఎనిమిది", - "ఎనభై ఏడు", - "ఎనభై ఆరు", - "ఎనభై అయిదు", - "ఎనభై నాలుగు", - "ఎనభై మూడు", - "ఎనభై రెండు", - "ఎనభై ఒకటి", - "ఎనభై", - "డెబ్బై తొమ్మిది", - "డెబ్బై ఎనిమిది", - "డెబ్బై ఏడు", - "డెబ్బై ఆరు", - "డెబ్బై అయిదు", - "డెబ్బై నాలుగు", - "డెబ్బై మూడు", - "డెబ్బై రెండు", - "డెబ్బై ఒకటి", - "డెబ్బై", - "అరవై తొమ్మిది", - "అరవై ఎనిమిది", - "అరవై ఏడు", - "అరవై ఆరు", - "అరవై అయిదు", - "అరవై నాలుగు", - "అరవై మూడు", - "అరవై రెండు", - "అరవై ఒకటి", - "అరవై", - "యాభై తొమ్మిది", - "యాభై ఎనిమిది", - "యాభై ఏడు", - "యాభై ఆరు", - "యాభై అయిదు", - "యాభై నాలుగు", - "యాభై మూడు", - "యాభై రెండు", - "యాభై ఒకటి", - "యాభై ", - "నలభై తొమ్మిది", - "నలభై ఎనిమిది", - "నలభై ఏడు", - "నలభై ఆరు", - "నలభై అయిదు", - "నలభై నాలుగు", - "నలభై మూడు", - "నలభై రెండు", - "నలభై ఒకటి", - "నలభై", - "ముప్పై తొమ్మిది", - "ముప్పై ఎనిమిది", - "ముప్పై ఏడు", - "ముప్పై ఆరు", - "ముప్పై ఐదు", - "ముప్పై నాలుగు", - "ముప్పై మూడు", - "ముప్పై రెండు", - "ముప్పై ఒకటి", - "ముప్పై", - "ఇరవై తొమ్మిది", - "ఇరవై ఎనిమిది", - "ఇరవై ఏడు", - "ఇరవై ఆరు", - "ఇరవై అయిదు", - "ఇరవై నాలుగు", - "ఇరవై మూడు", - "ఇరవై రెండు", - "ఇరవై ఒకటి", - "ఇరవై", - "పందొమ్మిది", - "పధ్ధెనిమిది", - "పదిహేడు", - "పదహారు", - "పదునయిదు", - "పధ్నాలుగు", - "పదమూడు", - "పన్నెండు", - "పదకొండు", - "పది", - "తొమ్మిది", - "ఎనిమిది", - "ఏడు", - "ఆరు", - "అయిదు", - "నాలుగు", - "మూడు", - "రెండు", - "ఒకటి", - "సున్న", - ] - - self.mid_numwords = [(100, "వంద")] - - self.high_numwords = [(7, "కోట్ల"), (5, "లక్ష"), (3, "వేయి")] - - self.pointword = "బిందువు " - - self.modifiers = [ - " ్ ", - "ా ", - " ి ", - " ీ ", - " ు ", - " ూ ", - " ృ ", - " ౄ ", - " ె ", - " ే ", - " ై ", - " ొ", - " ో ", - " ౌ ", - " ఁ ", - " ం ", - " ః ", - ] - - def merge(self, lpair, rpair): - ltext, lnum = lpair - rtext, rnum = rpair - if lnum == 1 and rnum < 100: - return (rtext, rnum) - elif 100 > lnum > rnum: - return ("%s-%s" % (ltext, rtext), lnum + rnum) - elif lnum >= 100 > rnum: - if ltext[-1] in self.modifiers: - return ("%s %s" % (ltext[:-1], rtext), lnum + rnum) - else: - return ("%s %s" % (ltext+"ల", rtext), lnum + rnum) - elif rnum > lnum: - return ("%s %s" % (ltext, rtext), lnum * rnum) - return ("%s %s" % (ltext, rtext), lnum + rnum) - - def to_ordinal_num(self, value): - self.verify_ordinal(value) - return "%s%s" % (value, self.to_ordinal(value)[-1:]) - - def to_ordinal(self, value): - self.verify_ordinal(value) - outwords = self.to_cardinal(value) - if outwords[-1] in self.modifiers: - outwords = outwords[:-1] - ordinal_num = outwords + "వ" - return ordinal_num diff --git a/fun_text_processing/num2words/num2words/lang_TH.py b/fun_text_processing/num2words/num2words/lang_TH.py deleted file mode 100644 index cd60ee9d0..000000000 --- a/fun_text_processing/num2words/num2words/lang_TH.py +++ /dev/null @@ -1,168 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import unicode_literals - -from num2words.base import Num2Word_Base -from num2words.currency import parse_currency_parts -from num2words.utils import splitbyx - - -class Num2Word_TH(Num2Word_Base): - - def setup(self): - self.negword = 'ติดลบ' - self.pointword = 'จุด' - - self.CURRENCY_FORMS = { - 'THB': (('บาท', 'บาท'), ('สตางค์', 'สตางค์')), - 'USD': (('ดอลลาร์', 'ดอลลาร์'), ('เซนต์', 'เซนต์')), - 'EUR': (('ยูโร', 'ยูโร'), ('เซนต์', 'เซนต์')), - } - - self.high_numwords = [] - - self.mid_numwords = ['', 'สิบ', 'ร้อย', 'พัน', 'หมื่น', 'แสน', 'ล้าน'] - - self.low_numwords = [ - 'ศูนย์', 'หนึ่ง', 'สอง', 'สาม', 'สี่', - 'ห้า', 'หก', 'เจ็ด', 'แปด', 'เก้า' - ] - - def set_high_numwords(self, high_numwords): - pass - - def set_mid_numwords(self, mid_numwords): - pass - - def splitnum(self, six_num): - length = len(six_num) > 1 - word_num = '' - - for index, num in enumerate(map(int, six_num)): - if num: - if index: - word_num = self.mid_numwords[index] + word_num - - if length and num == 1 and index == 0: - word_num += 'เอ็ด' - elif index == 1 and num == 2: - word_num = 'ยี่' + word_num - elif index != 1 or num != 1: - word_num = self.low_numwords[num] + word_num - - elif num == 0 and index == 0 and length == 0: - word_num = self.low_numwords[0] - - return word_num - - def split_six(self, num_txt): - result = splitbyx(num_txt, 6, format_int=False) - result = list(result)[::-1] - number_list = [] - for i in result: - number_list.append(i[::-1]) - return number_list - - def add_text_million(self, word_num): - result = '' - - for index, t in enumerate(reversed(word_num)): - if index == 0: - result = t - else: - result = result + 'ล้าน' + t - - return result - - def round_2_decimal(self, number): - integer, cents, negative = parse_currency_parts( - number, is_int_with_cents=False - ) - integer = '{}'.format(integer) - cents = '{}'.format(cents) - - if len(cents) < 2: - add_zero = 2 - len(cents) - cents = ('0' * add_zero) + cents - - text_num = integer + '.' + cents - - return text_num, negative - - def left_num_to_text(self, number): - - left_num_list = self.split_six(number) - - left_text_list = [] - for i in left_num_list: - left_text_list.append(self.splitnum(i)) - - left_text = self.add_text_million(left_text_list) - return left_text - - def to_cardinal(self, number): - negative = number < 0 - - pre, post = self.float2tuple(number) - precision = self.precision - pre = '{}'.format(pre) - post = '{}'.format(post) - - if negative: - pre = pre.lstrip('-') - - if len(post) < precision: - add_zero = precision - len(post) - post = ('0' * add_zero) + post - - result = self.left_num_to_text(pre) - - right_text = '' - if not post == '0': - for i in map(int, post): - right_text = right_text + self.low_numwords[i] - result = result + 'จุด' + right_text - - if negative: - result = 'ติดลบ' + result - - return result - - def to_ordinal(self, number): - return self.to_cardinal(number) - - def to_currency(self, number, currency='THB'): - - number, negative = self.round_2_decimal(number) - - split_num = number.split('.') - - left_num = split_num[0] - left_text = self.left_num_to_text(left_num) - - right_num = split_num[1] - right_text = self.splitnum(right_num[::-1].rstrip('0')) - - try: - cr1, cr2 = self.CURRENCY_FORMS[currency] - - except KeyError: - raise NotImplementedError( - 'Currency code "%s" not implemented for "%s"' % - (currency, self.__class__.__name__)) - - if right_num == '00': - if currency == 'THB': - result = left_text + cr1[0] + 'ถ้วน' - else: - result = left_text + cr1[0] - else: - if left_num == '0': - result = right_text + cr2[0] - else: - result = left_text + cr1[0] + right_text + cr2[0] - - if negative: - result = self.negword + result - - return result diff --git a/fun_text_processing/num2words/num2words/lang_TR.py b/fun_text_processing/num2words/num2words/lang_TR.py deleted file mode 100644 index 7fba87116..000000000 --- a/fun_text_processing/num2words/num2words/lang_TR.py +++ /dev/null @@ -1,846 +0,0 @@ -# -*- coding: utf-8 -*- -# Copyright (c) 2003, Taro Ogawa. All Rights Reserved. -# Copyright (c) 2013, Savoir-faire Linux inc. All Rights Reserved. -# Copyright (c) 2017, Tufan Kaynak, Framras. All Rights Reserved. - -# This library is free software; you can redistribute it and/or -# modify it under the terms of the GNU Lesser General Public -# License as published by the Free Software Foundation; either -# version 2.1 of the License, or (at your option) any later version. -# This library is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU -# Lesser General Public License for more details. -# You should have received a copy of the GNU Lesser General Public -# License along with this library; if not, write to the Free Software -# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, -# MA 02110-1301 USA - -from __future__ import unicode_literals - -from .base import Num2Word_Base - - -class Num2Word_TR(Num2Word_Base): - def __init__(self): - self.precision = 2 - self.negword = u"eksi" - self.pointword = u"virgül" - self.CURRENCY_UNIT = u"lira" - self.CURRENCY_SUBUNIT = u"kuruş" - self.errmsg_nonnum = u"Sadece sayılar yazıya çevrilebilir." - self.errmsg_floatord = u"Tam sayı olmayan {} sıralamada kullanılamaz." - self.errmsg_negord = u"Pozitif olmayan {} sıralamada kullanılamaz." - self.errmsg_toobig = u"abs({}) sayı yazıya çevirmek için çok büyük. " \ - u"Yazıya çevrilebilecek en büyük rakam {}." - self.exclude_title = [] - self.DECIMAL_SIGN = "," - self.ORDINAL_SIGN = "." - self.ZERO = u"sıfır" - self.CARDINAL_ONES = { - "1": u"bir", - "2": u"iki", - "3": u"üç", - "4": u"dört", - "5": u"beş", - "6": u"altı", - "7": u"yedi", - "8": u"sekiz", - "9": u"dokuz" - } - self.ORDINAL_ONES = { - "1": u"birinci", - "2": u"ikinci", - "3": u"üçüncü", - "4": u"dördüncü", - "5": u"beşinci", - "6": u"altıncı", - "7": u"yedinci", - "8": u"sekizinci", - "9": u"dokuzuncu" - } - self.CARDINAL_TENS = { - "1": u"on", - "2": u"yirmi", - "3": u"otuz", - "4": u"kırk", - "5": u"elli", - "6": u"altmış", - "7": u"yetmiş", - "8": u"seksen", - "9": u"doksan" - } - self.ORDINAL_TENS = { - "1": u"onuncu", - "2": u"yirminci", - "3": u"otuzuncu", - "4": u"kırkıncı", - "5": u"ellinci", - "6": u"altmışıncı", - "7": u"yetmişinci", - "8": u"sekseninci", - "9": u"doksanıncı" - } - self.HUNDREDS = { - "2": u"iki", - "3": u"üç", - "4": u"dört", - "5": u"beş", - "6": u"altı", - "7": u"yedi", - "8": u"sekiz", - "9": u"dokuz" - } - self.CARDINAL_HUNDRED = (u"yüz",) - self.ORDINAL_HUNDRED = (u"yüzüncü",) - self.CARDINAL_TRIPLETS = { - 1: u"bin", - 2: u"milyon", - 3: u"milyar", - 4: u"trilyon", - 5: u"katrilyon", - 6: u"kentilyon" - } - self.ORDINAL_TRIPLETS = { - 1: u"bininci", - 2: u"milyonuncu", - 3: u"milyarıncı", - 4: u"trilyonuncu", - 5: u"katrilyonuncu", - 6: u"kentilyon" - } - self.MAXVAL = (10 ** ((len(self.CARDINAL_TRIPLETS) + 1) * 3)) - 1 - - self.integers_to_read = [] - self.total_triplets_to_read = 0 - self.total_digits_outside_triplets = 0 - self.order_of_last_zero_digit = 0 - - def to_cardinal(self, value): - wrd = "" - is_cardinal = self.verify_cardinal(value) - if not is_cardinal: - return wrd - - if not int(value) == value: - return self.to_cardinal_float(value) - self.to_splitnum(value) - - if self.order_of_last_zero_digit >= len(self.integers_to_read[0]): - # number like 00 and all 0s and even more, raise error - return wrd - - if self.total_triplets_to_read == 1: - if self.total_digits_outside_triplets == 2: - if self.order_of_last_zero_digit == 1: - # number like x0, read cardinal x0 and return - wrd += self.CARDINAL_TENS.get( - self.integers_to_read[0][0], "" - ) - return wrd - if self.order_of_last_zero_digit == 0: - # number like xy, read cardinal xy and return - wrd += self.CARDINAL_TENS.get( - self.integers_to_read[0][0], "" - ) - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][1], "" - ) - return wrd - - if self.total_digits_outside_triplets == 1: - if self.order_of_last_zero_digit == 0: - # number like x, read cardinal x and return - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][0], "" - ) - if self.integers_to_read[0][0] == "0": - return self.ZERO - return wrd - - if self.total_digits_outside_triplets == 0: - if self.order_of_last_zero_digit == 2: - # number like x00, read cardinal x00 and return - wrd += self.HUNDREDS.get( - self.integers_to_read[0][0], "" - ) - wrd += self.CARDINAL_HUNDRED[0] - return wrd - if self.order_of_last_zero_digit == 1: - # number like xy0, read cardinal xy0 and return - wrd += self.HUNDREDS.get( - self.integers_to_read[0][0], "" - ) - wrd += self.CARDINAL_HUNDRED[0] - wrd += self.CARDINAL_TENS.get( - self.integers_to_read[0][1], "" - ) - return wrd - if self.order_of_last_zero_digit == 0: - # number like xyz, read cardinal xyz and return - wrd += self.HUNDREDS.get( - self.integers_to_read[0][0], "" - ) - wrd += self.CARDINAL_HUNDRED[0] - wrd += self.CARDINAL_TENS.get( - self.integers_to_read[0][1], "" - ) - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][2], "" - ) - return wrd - - if self.total_triplets_to_read >= 2: - if self.total_digits_outside_triplets == 2: - if self.order_of_last_zero_digit == len( - self.integers_to_read[0]) - 1: - # number like x0 and all 0s, read cardinal x0 0..0 - # and return - wrd += self.CARDINAL_TENS.get( - self.integers_to_read[0][0], "" - ) - wrd += self.CARDINAL_TRIPLETS[ - self.total_triplets_to_read - 1 - ] - return wrd - if self.order_of_last_zero_digit == len( - self.integers_to_read[0]) - 2: - # number like xy and all 0s, read cardinal xy 0..0 - # and return - wrd += self.CARDINAL_TENS.get( - self.integers_to_read[0][0], "" - ) - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][1], "" - ) - wrd += self.CARDINAL_TRIPLETS[ - self.total_triplets_to_read - 1 - ] - return wrd - if self.order_of_last_zero_digit < len( - self.integers_to_read[0]) - 2: - # number like xy and others, read cardinal xy n..n - # and return - wrd += self.CARDINAL_TENS.get( - self.integers_to_read[0][0], "" - ) - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][1], "" - ) - wrd += self.CARDINAL_TRIPLETS[ - self.total_triplets_to_read - 1 - ] - - if self.total_digits_outside_triplets == 1: - if self.order_of_last_zero_digit == len( - self.integers_to_read[0]) - 1: - # number like x and all 0s, read cardinal x 0..0 - # and return - if not (self.total_triplets_to_read == 2 and - self.integers_to_read[0][0] == "1"): - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][0], "" - ) - wrd += self.CARDINAL_TRIPLETS[ - self.total_triplets_to_read - 1 - ] - return wrd - if self.order_of_last_zero_digit < len( - self.integers_to_read[0]) - 1: - # number like x and others, read cardinal x n..n - # and return - if not (self.total_triplets_to_read == 2 and - self.integers_to_read[0][0] == "1"): - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][0], "" - ) - wrd += self.CARDINAL_TRIPLETS[ - self.total_triplets_to_read - 1 - ] - - if self.total_digits_outside_triplets == 0: - if self.order_of_last_zero_digit == len( - self.integers_to_read[0]) - 1: - # number like x00 and all 0s, read cardinal x00 0..0 - # and return - wrd += self.HUNDREDS.get(self.integers_to_read[0][0], "") - wrd += self.CARDINAL_HUNDRED[0] - wrd += self.CARDINAL_TRIPLETS[ - self.total_triplets_to_read - 1 - ] - return wrd - if self.order_of_last_zero_digit == len( - self.integers_to_read[0]) - 2: - # number like xy0 and all 0s, read cardinal xy0 0..0 - # and return - wrd += self.HUNDREDS.get( - self.integers_to_read[0][0], "" - ) - wrd += self.CARDINAL_HUNDRED[0] - wrd += self.CARDINAL_TENS.get( - self.integers_to_read[0][1], "" - ) - wrd += self.CARDINAL_TRIPLETS[ - self.total_triplets_to_read - 1 - ] - return wrd - if self.order_of_last_zero_digit == len( - self.integers_to_read[0]) - 3: - # number like xyz and all 0s, read cardinal xyz 0..0 - # and return - wrd += self.HUNDREDS.get(self.integers_to_read[0][0], "") - wrd += self.CARDINAL_HUNDRED[0] - wrd += self.CARDINAL_TENS.get( - self.integers_to_read[0][1], "" - ) - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][2], "" - ) - wrd += self.CARDINAL_TRIPLETS[ - self.total_triplets_to_read - 1 - ] - return wrd - if self.order_of_last_zero_digit < len( - self.integers_to_read[0]) - 3: - # number like xyz and all others, read cardinal xyz n..n - wrd += self.HUNDREDS.get(self.integers_to_read[0][0], "") - wrd += self.CARDINAL_HUNDRED[0] - wrd += self.CARDINAL_TENS.get( - self.integers_to_read[0][1], "" - ) - if not (self.total_triplets_to_read == 2 and - self.integers_to_read[0][2] == "1"): - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][2], "" - ) - wrd += self.CARDINAL_TRIPLETS[ - self.total_triplets_to_read - 1 - ] - - for i in list(range(self.total_triplets_to_read - 1, 0, -1)): - reading_triplet_order = self.total_triplets_to_read - i - if self.total_digits_outside_triplets == 0: - last_read_digit_order = reading_triplet_order * 3 - else: - last_read_digit_order = (reading_triplet_order - 1) * 3 +\ - self.total_digits_outside_triplets - - if not self.integers_to_read[0][ - last_read_digit_order: last_read_digit_order + 3 - ] == "000": - if not self.integers_to_read[0][ - last_read_digit_order - ] == "0": - wrd += self.HUNDREDS.get( - self.integers_to_read[0][last_read_digit_order], "" - ) - if self.order_of_last_zero_digit == len( - self.integers_to_read[0]) - ( - last_read_digit_order) - 1: - if i == 1: - wrd += self.CARDINAL_HUNDRED[0] - return wrd - elif i > 1: - wrd += self.CARDINAL_HUNDRED[0] - wrd += self.CARDINAL_TRIPLETS[i - 1] - return wrd - else: - wrd += self.CARDINAL_HUNDRED[0] - - if not self.integers_to_read[0][ - last_read_digit_order + 1] == "0": - if self.order_of_last_zero_digit == len( - self.integers_to_read[0]) - ( - last_read_digit_order) - 2: - if i == 1: - wrd += self.CARDINAL_TENS.get( - self.integers_to_read[0][ - last_read_digit_order + 1], "" - ) - return wrd - elif i > 1: - wrd += self.CARDINAL_TENS.get( - self.integers_to_read[0][ - last_read_digit_order + 1], "" - ) - wrd += self.CARDINAL_TRIPLETS[i - 1] - return wrd - else: - wrd += self.CARDINAL_TENS.get( - self.integers_to_read[0][ - last_read_digit_order + 1], "" - ) - - if not self.integers_to_read[0][ - last_read_digit_order + 2] == "0": - if self.order_of_last_zero_digit == len( - self.integers_to_read[0]) - ( - last_read_digit_order) - 3: - if i == 1: - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][ - last_read_digit_order + 2], "" - ) - return wrd - if i == 2: - if not self.integers_to_read[0][ - last_read_digit_order: - last_read_digit_order + 2 - ] == "00": - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][ - last_read_digit_order + 2], "" - ) - elif not self.integers_to_read[0][ - last_read_digit_order + 2] == "1": - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][ - last_read_digit_order + 2], "" - ) - wrd += self.CARDINAL_TRIPLETS[i - 1] - return wrd - if i > 2: - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][ - last_read_digit_order + 2], "" - ) - wrd += self.CARDINAL_TRIPLETS[i - 1] - return wrd - else: - if not self.integers_to_read[0][ - last_read_digit_order: - last_read_digit_order + 2 - ] == "00": - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][ - last_read_digit_order + 2], "" - ) - else: - if i == 2: - if not self.integers_to_read[0][ - last_read_digit_order: - last_read_digit_order + 2 - ] == "00": - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][ - last_read_digit_order + 2], "" - ) - elif not self.integers_to_read[0][ - last_read_digit_order + 2 - ] == "1": - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][ - last_read_digit_order + 2], "" - ) - - wrd += self.CARDINAL_TRIPLETS[i - 1] - - return wrd - - def to_cardinal_float(self, value): - self.to_splitnum(value) - wrd = "" - wrd += self.pointword - if len(self.integers_to_read[1]) >= 1: - wrd += self.CARDINAL_TENS.get(self.integers_to_read[1][0], "") - - if len(self.integers_to_read[1]) == 2: - wrd += self.CARDINAL_ONES.get(self.integers_to_read[1][1], "") - - if self.integers_to_read[0] == "0": - wrd = self.ZERO + wrd - else: - wrd = self.to_cardinal(int(self.integers_to_read[0])) + wrd - return wrd - - def verify_cardinal(self, value): - iscardinal = True - try: - if not float(value) == value: - iscardinal = False - except (ValueError, TypeError): - raise TypeError(self.errmsg_nonnum) - if abs(value) >= self.MAXVAL: - raise OverflowError(self.errmsg_toobig.format(value, self.MAXVAL)) - return iscardinal - - def verify_ordinal(self, value): - isordinal = True - try: - if not int(value) == value: - isordinal = False - if not abs(value) == value: - raise TypeError(self.errmsg_negord.format(value)) - except (ValueError, TypeError): - raise TypeError(self.errmsg_nonnum) - if abs(value) >= self.MAXVAL: - raise OverflowError(self.errmsg_toobig.format(value, self.MAXVAL)) - return isordinal - - def to_ordinal(self, value): - wrd = "" - isordinal = self.verify_ordinal(value) - if isordinal: - self.to_splitnum(value) - - if self.order_of_last_zero_digit >= len(self.integers_to_read[0]): - # number like 00 and all 0s and even more, raise error - return wrd - - if self.total_triplets_to_read == 1: - if self.total_digits_outside_triplets == 2: - if self.order_of_last_zero_digit == 1: - # number like x0, read ordinal x0 and return - wrd += self.ORDINAL_TENS.get( - self.integers_to_read[0][0], "" - ) - return wrd - if self.order_of_last_zero_digit == 0: - # number like xy, read ordinal xy and return - wrd += self.CARDINAL_TENS.get( - self.integers_to_read[0][0], "" - ) - wrd += self.ORDINAL_ONES.get( - self.integers_to_read[0][1], "" - ) - return wrd - - if self.total_digits_outside_triplets == 1: - if self.order_of_last_zero_digit == 0: - # number like x, read ordinal x and return - wrd += self.ORDINAL_ONES.get( - self.integers_to_read[0][0], "" - ) - if self.integers_to_read[0][0] == "0": - return u"sıfırıncı" - return wrd - - if self.total_digits_outside_triplets == 0: - if self.order_of_last_zero_digit == 2: - # number like x00, read ordinal x00 and return - wrd += self.HUNDREDS.get( - self.integers_to_read[0][0], "" - ) - wrd += self.ORDINAL_HUNDRED[0] - return wrd - if self.order_of_last_zero_digit == 1: - # number like xy0, read ordinal xy0 and return - wrd += self.HUNDREDS.get( - self.integers_to_read[0][0], "" - ) - wrd += self.CARDINAL_HUNDRED[0] - wrd += self.ORDINAL_TENS.get( - self.integers_to_read[0][1], "" - ) - return wrd - if self.order_of_last_zero_digit == 0: - # number like xyz, read ordinal xyz and return - wrd += self.HUNDREDS.get( - self.integers_to_read[0][0], "" - ) - wrd += self.CARDINAL_HUNDRED[0] - wrd += self.CARDINAL_TENS.get( - self.integers_to_read[0][1], "" - ) - if not self.integers_to_read[0][2] == "0": - wrd += self.ORDINAL_ONES.get( - self.integers_to_read[0][2], "" - ) - return wrd - - if self.total_triplets_to_read >= 2: - if self.total_digits_outside_triplets == 2: - if self.order_of_last_zero_digit == len( - self.integers_to_read[0]) - 1: - # number like x0 and all 0s, read ordinal x0 0..0 - # and return - wrd += self.CARDINAL_TENS.get( - self.integers_to_read[0][0], "" - ) - wrd += self.ORDINAL_TRIPLETS[ - self.total_triplets_to_read - 1 - ] - return wrd - if self.order_of_last_zero_digit == len( - self.integers_to_read[0]) - 2: - # number like xy and all 0s, read ordinal xy 0..0 - # and return - wrd += self.CARDINAL_TENS.get( - self.integers_to_read[0][0], "" - ) - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][1], "" - ) - wrd += self.ORDINAL_TRIPLETS[ - self.total_triplets_to_read - 1 - ] - return wrd - if self.order_of_last_zero_digit < len( - self.integers_to_read[0]) - 2: - # number like xy and others, read cardinal xy n..n - # and return - wrd += self.CARDINAL_TENS.get( - self.integers_to_read[0][0], "" - ) - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][1], "" - ) - wrd += self.CARDINAL_TRIPLETS[ - self.total_triplets_to_read - 1 - ] - - if self.total_digits_outside_triplets == 1: - if self.order_of_last_zero_digit == len( - self.integers_to_read[0]) - 1: - # number like x and all 0s, read ordinal x 0..0 - # and return - if not (self.total_triplets_to_read == 2 and - self.integers_to_read[0][0] == "1"): - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][0], "" - ) - wrd += self.ORDINAL_TRIPLETS[ - self.total_triplets_to_read - 1 - ] - return wrd - if self.order_of_last_zero_digit < len( - self.integers_to_read[0]) - 1: - # number like x and others, read cardinal x n..n - # and return - if not (self.total_triplets_to_read == 2 and - self.integers_to_read[0][0] == "1"): - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][0], "" - ) - wrd += self.CARDINAL_TRIPLETS[ - self.total_triplets_to_read - 1 - ] - - if self.total_digits_outside_triplets == 0: - if self.order_of_last_zero_digit == len( - self.integers_to_read[0]) - 1: - # number like x00 and all 0s, read ordinal x00 0..0 - # and return - wrd += self.HUNDREDS.get( - self.integers_to_read[0][0], "" - ) - wrd += self.CARDINAL_HUNDRED[0] - wrd += self.ORDINAL_TRIPLETS[ - self.total_triplets_to_read - 1 - ] - return wrd - if self.order_of_last_zero_digit == len( - self.integers_to_read[0]) - 2: - # number like xy0 and all 0s, read ordinal xy0 0..0 - # and return - wrd += self.HUNDREDS.get( - self.integers_to_read[0][0], "" - ) - wrd += self.CARDINAL_HUNDRED[0] - wrd += self.CARDINAL_TENS.get( - self.integers_to_read[0][1], "" - ) - wrd += self.ORDINAL_TRIPLETS[ - self.total_triplets_to_read - 1 - ] - return wrd - if self.order_of_last_zero_digit == len( - self.integers_to_read[0]) - 3: - # number like xyz and all 0s, read ordinal xyz 0..0 - # and return - wrd += self.HUNDREDS.get( - self.integers_to_read[0][0], "" - ) - wrd += self.CARDINAL_HUNDRED[0] - wrd += self.CARDINAL_TENS.get( - self.integers_to_read[0][1], "" - ) - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][2], "" - ) - wrd += self.ORDINAL_TRIPLETS[ - self.total_triplets_to_read - 1 - ] - return wrd - if self.order_of_last_zero_digit < len( - self.integers_to_read[0]) - 3: - # number like xyz and all others, read cardinal - # xyz n..n - wrd += self.HUNDREDS.get( - self.integers_to_read[0][0], "" - ) - wrd += self.CARDINAL_HUNDRED[0] - wrd += self.CARDINAL_TENS.get( - self.integers_to_read[0][1], "" - ) - if not (self.total_triplets_to_read == 2 and - self.integers_to_read[0][2] == "1"): - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][2], "" - ) - wrd += self.CARDINAL_TRIPLETS[ - self.total_triplets_to_read - 1 - ] - - for i in list(range(self.total_triplets_to_read - 1, 0, -1)): - reading_triplet_order = self.total_triplets_to_read - i - if self.total_digits_outside_triplets == 0: - last_read_digit_order = reading_triplet_order * 3 - else: - last_read_digit_order = \ - (reading_triplet_order - 1) * 3 + \ - self.total_digits_outside_triplets - - if not self.integers_to_read[0][ - last_read_digit_order: last_read_digit_order + 3 - ] == "000": - if not self.integers_to_read[0][ - last_read_digit_order - ] == "0": - if not self.integers_to_read[0][ - last_read_digit_order - ] == "1": - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][ - last_read_digit_order - ], "" - ) - if self.order_of_last_zero_digit == len( - self.integers_to_read[0]) - ( - last_read_digit_order) - 1: - if i == 1: - wrd += self.ORDINAL_HUNDRED[0] - return wrd - elif i > 1: - wrd += self.CARDINAL_HUNDRED[0] - wrd += self.ORDINAL_TRIPLETS[i - 1] - return wrd - else: - wrd += self.CARDINAL_HUNDRED[0] - - if not self.integers_to_read[0][ - last_read_digit_order + 1 - ] == "0": - if self.order_of_last_zero_digit == len( - self.integers_to_read[0]) - ( - last_read_digit_order) - 2: - if i == 1: - wrd += self.ORDINAL_TENS.get( - self.integers_to_read[0][ - last_read_digit_order + 1], "" - ) - return wrd - elif i > 1: - wrd += self.CARDINAL_TENS.get( - self.integers_to_read[0][ - last_read_digit_order + 1], "" - ) - wrd += self.ORDINAL_TRIPLETS[i - 1] - return wrd - else: - wrd += self.CARDINAL_TENS.get( - self.integers_to_read[0][ - last_read_digit_order + 1], "" - ) - - if not self.integers_to_read[0][ - last_read_digit_order + 2 - ] == "0": - if self.order_of_last_zero_digit == len( - self.integers_to_read[0]) - ( - last_read_digit_order) - 3: - if i == 1: - wrd += self.ORDINAL_ONES.get( - self.integers_to_read[0][ - last_read_digit_order + 2], "" - ) - return wrd - if i == 2: - if not self.integers_to_read[0][ - last_read_digit_order: - last_read_digit_order + 2] == "00": - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][ - last_read_digit_order + 2], "" - ) - elif not self.integers_to_read[0][ - last_read_digit_order + 2 - ] == "1": - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][ - last_read_digit_order + 2], "" - ) - wrd += self.ORDINAL_TRIPLETS[i - 1] - return wrd - if i > 2: - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][ - last_read_digit_order + 2], "" - ) - wrd += self.ORDINAL_TRIPLETS[i - 1] - return wrd - else: - if not self.integers_to_read[0][ - last_read_digit_order: - last_read_digit_order + 2] == "00": - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][ - last_read_digit_order + 2], "" - ) - else: - if not self.integers_to_read[0][ - last_read_digit_order: - last_read_digit_order + 2] == "00": - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][ - last_read_digit_order + 2], "" - ) - elif not self.integers_to_read[0][ - last_read_digit_order + 2] == "1": - wrd += self.CARDINAL_ONES.get( - self.integers_to_read[0][ - last_read_digit_order + 2], "" - ) - - wrd += self.CARDINAL_TRIPLETS[i - 1] - - return wrd - - def to_splitnum(self, val): - float_digits = str(int(val * 10 ** self.precision)) - if not int(val) == 0: - self.integers_to_read = [ - str(int(val)), - float_digits[len(float_digits) - self.precision:] - ] - else: - self.integers_to_read = [ - "0", - "0" * (self.precision - len(float_digits)) + - float_digits[len(float_digits) - self.precision:] - ] - if len(self.integers_to_read[0]) % 3 > 0: - self.total_triplets_to_read = (len(self.integers_to_read[0]) // 3)\ - + 1 - elif len(self.integers_to_read[0]) % 3 == 0: - self.total_triplets_to_read = len(self.integers_to_read[0]) // 3 - self.total_digits_outside_triplets = len(self.integers_to_read[0]) % 3 - - okunacak = list(self.integers_to_read[0][::-1]) - self.order_of_last_zero_digit = 0 - found = 0 - for i in range(len(okunacak) - 1): - if int(okunacak[i]) == 0 and found == 0: - self.order_of_last_zero_digit = i + 1 - else: - found = 1 - - def to_currency(self, value): - if int(value) == 0: - return u"bedelsiz" - valueparts = self.to_cardinal(value).split(self.pointword) - if len(valueparts) == 1: - return valueparts[0] + self.CURRENCY_UNIT - if len(valueparts) == 2: - return self.CURRENCY_UNIT.join(valueparts) + \ - self.CURRENCY_SUBUNIT diff --git a/fun_text_processing/num2words/num2words/lang_UK.py b/fun_text_processing/num2words/num2words/lang_UK.py deleted file mode 100644 index 835e728d4..000000000 --- a/fun_text_processing/num2words/num2words/lang_UK.py +++ /dev/null @@ -1,837 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import unicode_literals - -from .base import Num2Word_Base -from .utils import get_digits, splitbyx - -ZERO = ('нуль',) - -ONES_FEMININE = { - 1: ('одна',), - 2: ('дві',), - 3: ('три',), - 4: ('чотири',), - 5: ('п\'ять',), - 6: ('шість',), - 7: ('сім',), - 8: ('вісім',), - 9: ('дев\'ять',), -} - -ONES = { - 1: ('один',), - 2: ('два',), - 3: ('три',), - 4: ('чотири',), - 5: ('п\'ять',), - 6: ('шість',), - 7: ('сім',), - 8: ('вісім',), - 9: ('дев\'ять',), -} - -ONES_ORDINALS = { - 1: ("перший", "одно"), - 2: ("другий", "двох"), - 3: ("третій", "трьох"), - 4: ("четвертий", "чотирьох"), - 5: ("п'ятий", "п'яти"), - 6: ("шостий", "шести"), - 7: ("сьомий", "семи"), - 8: ("восьмий", "восьми"), - 9: ("дев'ятий", "дев'яти"), - 10: ("десятий", "десяти"), - 11: ("одинадцятий", "одинадцяти"), - 12: ("дванадцятий", "дванадцяти"), - 13: ("тринадцятий", "тринадцяти"), - 14: ("чотирнадцятий", "чотирнадцяти"), - 15: ("п'ятнадцятий", "п'ятнадцяти"), - 16: ("шістнадцятий", "шістнадцяти"), - 17: ("сімнадцятий", "сімнадцяти"), - 18: ("вісімнадцятий", "вісімнадцяти"), - 19: ("дев'ятнадцятий", "дев'ятнадцяти"), -} - -TENS = { - 0: ('десять',), - 1: ('одинадцять',), - 2: ('дванадцять',), - 3: ('тринадцять',), - 4: ('чотирнадцять',), - 5: ('п\'ятнадцять',), - 6: ('шістнадцять',), - 7: ('сімнадцять',), - 8: ('вісімнадцять',), - 9: ('дев\'ятнадцять',), -} - -TWENTIES = { - 2: ('двадцять',), - 3: ('тридцять',), - 4: ('сорок',), - 5: ('п\'ятдесят',), - 6: ('шістдесят',), - 7: ('сімдесят',), - 8: ('вісімдесят',), - 9: ('дев\'яносто',), -} - -TWENTIES_ORDINALS = { - 2: ("двадцятий", "двадцяти"), - 3: ("тридцятий", "тридцяти"), - 4: ("сороковий", "сорока"), - 5: ("п'ятдесятий", "п'ятдесяти"), - 6: ("шістдесятий", "шістдесяти"), - 7: ("сімдесятий", "сімдесяти"), - 8: ("вісімдесятий", "вісімдесяти"), - 9: ("дев'яностий", "дев'яности"), -} - -HUNDREDS = { - 1: ('сто',), - 2: ('двісті',), - 3: ('триста',), - 4: ('чотириста',), - 5: ('п\'ятсот',), - 6: ('шістсот',), - 7: ('сімсот',), - 8: ('вісімсот',), - 9: ('дев\'ятсот',), -} - -HUNDREDS_ORDINALS = { - 1: ("сотий", "сто"), - 2: ("двохсотий", "двохсот"), - 3: ("трьохсотий", "трьохсот"), - 4: ("чотирьохсотий", "чотирьохсот"), - 5: ("п'ятисотий", "п'ятсот"), - 6: ("шестисотий", "шістсот"), - 7: ("семисотий", "сімсот"), - 8: ("восьмисотий", "вісімсот"), - 9: ("дев'ятисотий", "дев'ятсот"), -} - -THOUSANDS = { - 1: ('тисяча', 'тисячі', 'тисяч'), # 10^3 - 2: ('мільйон', 'мільйони', 'мільйонів'), # 10^6 - 3: ('мільярд', 'мільярди', 'мільярдів'), # 10^9 - 4: ('трильйон', 'трильйони', 'трильйонів'), # 10^12 - 5: ('квадрильйон', 'квадрильйони', 'квадрильйонів'), # 10^15 - 6: ('квінтильйон', 'квінтильйони', 'квінтильйонів'), # 10^18 - 7: ('секстильйон', 'секстильйони', 'секстильйонів'), # 10^21 - 8: ('септильйон', 'септильйони', 'септильйонів'), # 10^24 - 9: ('октильйон', 'октильйони', 'октильйонів'), # 10^27 - 10: ('нонільйон', 'нонільйони', 'нонільйонів'), # 10^30 -} - -prefixes_ordinal = { - 1: "тисячний", - 2: "мільйонний", - 3: "мільярдний", - 4: "трильйонний", - 5: "квадрильйонний", - 6: "квінтильйонний", - 7: "секстильйонний", - 8: "септильйонний", - 9: "октильйонний", - 10: "нонільйонний", -} - -FEMININE_MONEY = ('AOA', 'BAM', 'BDT', 'BWP', 'CZK', 'DKK', - 'ERN', 'HNL', 'HRK', 'IDR', 'INR', 'ISK', - 'JPY', 'KPW', 'KRW', 'LKR', 'MOP', 'MRU', - 'MUR', 'MVR', 'MWK', 'NGN', 'NIO', 'NOK', - 'NPR', 'PKR', 'SCR', 'SEK', 'STN', 'TRY', - 'WST', 'UAH', 'ZMW') -FEMININE_CENTS = ('ALL', 'BDT', 'BGN', 'BYN', 'GHS', 'HRK', - 'ILS', 'INR', 'NPR', 'OMR', 'OMR', 'PKR', - 'RSD', 'RUB', 'UAH') - -GENERIC_DOLLARS = ('долар', 'долари', 'доларів') -GENERIC_CENTS = ('цент', 'центи', 'центів') - - -class Num2Word_UK(Num2Word_Base): - CURRENCY_FORMS = { - 'AED': ( - ('дирхам', 'дирхами', 'дирхамів'), - ('філс', 'філси', 'філсів') - ), - 'AFN': ( - ('афгані', 'афгані', 'афгані'), - ('пул', 'пули', 'пулів') - ), - 'ALL': ( - ('лек', 'леки', 'леків'), - ('кіндарка', 'кіндарки', 'кіндарок') - ), - 'AMD': ( - ('драм', 'драми', 'драмів'), - ('лум', 'лум', 'лум') - ), - 'ANG': ( - ('гульден', 'гульдени', 'гульденів'), - GENERIC_CENTS - ), - 'AOA': ( - ('кванза', 'кванзи', 'кванз'), - ('сентимо', 'сентимо', 'сентимо') - ), - 'ARS': ( - ('песо', 'песо', 'песо'), - ('сентаво', 'сентаво', 'сентаво') - ), - 'AUD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'AWG': ( - ('флорин', 'флорини', 'флоринів'), - GENERIC_CENTS - ), - 'AZN': ( - ('манат', 'манати', 'манатів'), - ('гяпік', 'гяпіки', 'гяпіків') - ), - 'BAM': ( - ('марка', 'марки', 'марок'), - ('фенінг', 'фенінги', 'фенінгів') - ), - 'BBD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'BDT': ( - ('така', 'таки', 'так'), - ('пойша', 'пойші', 'пойш') - ), - 'BGN': ( - ('лев', 'леви', 'левів'), - ('стотинка', 'стотинки', 'стотинок') - ), - 'BHD': ( - ('динар', 'динари', 'динарів'), - ('філс', 'філси', 'філсів') - ), - 'BIF': ( - ('франк', 'франки', 'франків'), - ('сантим', 'сантими', 'сантимів') - ), - 'BMD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'BND': (GENERIC_DOLLARS, GENERIC_CENTS), - 'BOB': ( - ('болівіано', 'болівіано', 'болівіано'), - ('сентаво', 'сентаво', 'сентаво') - ), - 'BRL': ( - ('реал', 'реали', 'реалів'), - ('сентаво', 'сентаво', 'сентаво') - ), - 'BSD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'BTN': ( - ('нгултрум', 'нгултруми', 'нгултрумів'), - ('четрум', 'четруми', 'четрумів') - ), - 'BWP': ( - ('пула', 'пули', 'пул'), - ('тхебе', 'тхебе', 'тхебе') - ), - 'BYN': ( - ('рубель', 'рублі', 'рублів'), - ('копійка', 'копійки', 'копійок') - ), - 'BZD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'CAD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'CDF': ( - ('франк', 'франки', 'франків'), - ('сантим', 'сантими', 'сантимів') - ), - 'CHF': ( - ('франк', 'франки', 'франків'), - ('сантим', 'сантими', 'сантимів') - ), - 'CLP': ( - ('песо', 'песо', 'песо'), - ('сентаво', 'сентаво', 'сентаво') - ), - 'CNY': ( - ('юань', 'юані', 'юанів'), - ('финь', 'фині', 'финів') - ), - 'COP': ( - ('песо', 'песо', 'песо'), - ('сентаво', 'сентаво', 'сентаво') - ), - 'CRC': ( - ('колон', 'колони', 'колонів'), - ('сентімо', 'сентімо', 'сентімо') - ), - 'CUC': ( - ('песо', 'песо', 'песо'), - ('сентаво', 'сентаво', 'сентаво') - ), - 'CUP': ( - ('песо', 'песо', 'песо'), - ('сентаво', 'сентаво', 'сентаво') - ), - 'CVE': ( - ('ескудо', 'ескудо', 'ескудо'), - ('сентаво', 'сентаво', 'сентаво') - ), - 'CZK': ( - ('крона', 'крони', 'крон'), - ('гелер', 'гелери', 'гелерів') - ), - 'DJF': ( - ('франк', 'франки', 'франків'), - ('сантим', 'сантими', 'сантимів') - ), - 'DKK': ( - ('крона', 'крони', 'крон'), - ('ере', 'ере', 'ере') - ), - 'DOP': ( - ('песо', 'песо', 'песо'), - ('сентаво', 'сентаво', 'сентаво') - ), - 'DZD': ( - ('динар', 'динари', 'динарів'), - ('сантим', 'сантими', 'сантимів') - ), - 'EGP': ( - ('фунт', 'фунти', 'фунтів'), - ('піастр', 'піастри', 'піастрів') - ), - 'ERN': ( - ('накфа', 'накфи', 'накф'), - GENERIC_CENTS - ), - 'ETB': ( - ('бир', 'бири', 'бирів'), - GENERIC_CENTS - ), - 'EUR': ( - ('євро', 'євро', 'євро'), - GENERIC_CENTS - ), - 'FJD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'FKP': ( - ('фунт', 'фунти', 'фунтів'), - ('пенс', 'пенси', 'пенсів') - ), - 'GBP': ( - ('фунт', 'фунти', 'фунтів'), - ('пенс', 'пенси', 'пенсів') - ), - 'GEL': ( - ('ларі', 'ларі', 'ларі'), - ('тетрі', 'тетрі', 'тетрі') - ), - 'GHS': ( - ('седі', 'седі', 'седі'), - ('песева', 'песеви', 'песев') - ), - 'GIP': ( - ('фунт', 'фунти', 'фунтів'), - ('пенс', 'пенси', 'пенсів') - ), - 'GMD': ( - ('даласі', 'даласі', 'даласі'), - ('бутут', 'бутути', 'бутутів') - ), - 'GNF': ( - ('франк', 'франки', 'франків'), - ('сантим', 'сантими', 'сантимів') - ), - 'GTQ': ( - ('кетсаль', 'кетсалі', 'кетсалів'), - ('сентаво', 'сентаво', 'сентаво') - ), - 'GYD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'HKD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'HNL': ( - ('лемпіра', 'лемпіри', 'лемпір'), - ('сентаво', 'сентаво', 'сентаво') - ), - 'HRK': ( - ('куна', 'куни', 'кун'), - ('ліпа', 'ліпи', 'ліп') - ), - 'HTG': ( - ('гурд', 'гурди', 'гурдів'), - ('сантим', 'сантими', 'сантимів') - ), - 'HUF': ( - ('форинт', 'форинти', 'форинтів'), - ('філлер', 'філлери', 'філлерів') - ), - 'IDR': ( - ('рупія', 'рупії', 'рупій'), - GENERIC_CENTS - ), - 'ILS': ( - ('шекель', 'шекелі', 'шекелів'), - ('агора', 'агори', 'агор') - ), - 'INR': ( - ('рупія', 'рупії', 'рупій'), - ('пайса', 'пайси', 'пайс') - ), - 'IQD': ( - ('динар', 'динари', 'динарів'), - ('філс', 'філси', 'філсів') - ), - 'IRR': ( - ('ріал', 'ріали', 'ріалів'), - ('динар', 'динари', 'динарів') - ), - 'ISK': ( - ('крона', 'крони', 'крон'), - ('ейре', 'ейре', 'ейре') - ), - 'JMD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'JOD': ( - ('динар', 'динари', 'динарів'), - ('філс', 'філси', 'філсів') - ), - 'JPY': ( - ('єна', 'єни', 'єн'), - ('сен', 'сен', 'сен') - ), - 'KES': ( - ('шилінг', 'шилінги', 'шилінгів'), - GENERIC_CENTS - ), - 'KGS': ( - ('сом', 'соми', 'сомів'), - ('тиїн', 'тиїни', 'тиїнів') - ), - 'KHR': ( - ('рієль', 'рієлі', 'рієлів'), - ('су', 'су', 'су') - ), - 'KMF': ( - ('франк', 'франки', 'франків'), - ('сантим', 'сантими', 'сантимів') - ), - 'KPW': ( - ('вона', 'вони', 'вон'), - ('чон', 'чони', 'чонів') - ), - 'KRW': ( - ('вона', 'вони', 'вон'), - ('джеон', 'джеони', 'джеонів') - ), - 'KWD': ( - ('динар', 'динари', 'динарів'), - ('філс', 'філси', 'філсів') - ), - 'KYD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'KZT': ( - ('теньге', 'теньге', 'теньге'), - ('тиїн', 'тиїни', 'тиїнів')), - 'LAK': ( - ('кіп', 'кіпи', 'кіпів'), - ('ат', 'ати', 'атів') - ), - 'LBP': ( - ('фунт', 'фунти', 'фунтів'), - ('піастр', 'піастри', 'піастрів') - ), - 'LKR': ( - ('рупія', 'рупії', 'рупій'), - GENERIC_CENTS - ), - 'LRD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'LSL': ( - ('лоті', 'малоті', 'малоті'), - ('сенте', 'лісенте', 'лісенте') - ), - 'LYD': ( - ('динар', 'динари', 'динарів'), - ('дирхам', 'дирхами', 'дирхамів') - ), - 'MAD': ( - ('дирхам', 'дирхами', 'дирхамів'), - ('сантим', 'сантими', 'сантимів') - ), - 'MDL': ( - ('лей', 'леї', 'леї'), - ('бан', 'бані', 'бані') - ), - 'MGA': ( - ('аріарі', 'аріарі', 'аріарі'), - ('іраймбіланья', 'іраймбіланья', 'іраймбіланья') - ), - 'MKD': ( - ('денар', 'денари', 'денарів'), - ('дені', 'дені', 'дені') - ), - 'MMK': ( - ('к\'ят', 'к\'ят', 'к\'ят'), - ('п\'я', 'п\'я', 'п\'я') - ), - 'MNT': ( - ('тугрик', 'тугрики', 'тугриків'), - ('мунгу', 'мунгу', 'мунгу') - ), - 'MOP': ( - ('патака', 'патакі', 'патак'), - ('аво', 'аво', 'аво') - ), - 'MRU': ( - ('угія', 'угії', 'угій'), - ('хумс', 'хумс', 'хумс') - ), - 'MUR': ( - ('рупія', 'рупії', 'рупій'), - GENERIC_CENTS - ), - 'MVR': ( - ('руфія', 'руфії', 'руфій'), - ('ларі', 'ларі', 'ларі') - ), - 'MWK': ( - ('квача', 'квачі', 'квач'), - ('тамбала', 'тамбала', 'тамбала') - ), - 'MXN': ( - ('песо', 'песо', 'песо'), - ('сентаво', 'сентаво', 'сентаво') - ), - 'MYR': ( - ('рингіт', 'рингіти', 'рингітів'), - GENERIC_CENTS - ), - 'MZN': ( - ('метікал', 'метікали', 'метікалів'), - ('сентаво', 'сентаво', 'сентаво') - ), - 'NAD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'NGN': ( - ('найра', 'найри', 'найр'), - ('кобо', 'кобо', 'кобо') - ), - 'NIO': ( - ('кордоба', 'кордоби', 'кордоб'), - ('сентаво', 'сентаво', 'сентаво') - ), - 'NOK': ( - ('крона', 'крони', 'крон'), - ('ере', 'ере', 'ере') - ), - 'NPR': ( - ('рупія', 'рупії', 'рупій'), - ('пайса', 'пайси', 'пайс') - ), - 'NZD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'OMR': ( - ('ріал', 'ріали', 'ріалів'), - ('байза', 'байзи', 'байз') - ), - 'PAB': ( - ('бальбоа', 'бальбоа', 'бальбоа'), - ('сентесімо', 'сентесімо', 'сентесімо') - ), - 'PEN': ( - ('соль', 'соль', 'соль'), - ('сентімо', 'сентімо', 'сентімо') - ), - 'PGK': ( - ('кіна', 'кіна', 'кіна'), - ('тойя', 'тойя', 'тойя') - ), - 'PHP': ( - ('песо', 'песо', 'песо'), - ('сентаво', 'сентаво', 'сентаво') - ), - 'PKR': ( - ('рупія', 'рупії', 'рупій'), - ('пайса', 'пайси', 'пайс') - ), - 'PLN': ( - ('злотий', 'злоті', 'злотих'), - ('грош', 'гроші', 'грошів') - ), - 'PYG': ( - ('гуарані', 'гуарані', 'гуарані'), - ('сентімо', 'сентімо', 'сентімо') - ), - 'QAR': ( - ('ріал', 'ріали', 'ріалів'), - ('дирхам', 'дирхами', 'дирхамів') - ), - 'RON': ( - ('лей', 'леї', 'леї'), - ('бан', 'бані', 'бані') - ), - 'RSD': ( - ('динар', 'динари', 'динарів'), - ('пара', 'пари', 'пар') - ), - 'RUB': ( - ('рубль', 'рублі', 'рублів'), - ('копійка', 'копійки', 'копійок') - ), - 'RWF': ( - ('франк', 'франки', 'франків'), - ('сантим', 'сантими', 'сантимів') - ), - 'SAR': ( - ('ріал', 'ріали', 'ріалів'), - ('халал', 'халали', 'халалів') - ), - 'SBD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'SCR': ( - ('рупія', 'рупії', 'рупій'), - GENERIC_CENTS - ), - 'SDG': ( - ('фунт', 'фунти', 'фунтів'), - ('піастр', 'піастри', 'піастрів') - ), - 'SEK': ( - ('крона', 'крони', 'крон'), - ('ере', 'ере', 'ере') - ), - 'SGD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'SHP': ( - ('фунт', 'фунти', 'фунтів'), - ('пенс', 'пенси', 'пенсів') - ), - 'SLL': ( - ('леоне', 'леоне', 'леоне'), - GENERIC_CENTS - ), - 'SOS': ( - ('шилінг', 'шилінги', 'шилінгів'), - GENERIC_CENTS - ), - 'SRD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'SSP': ( - ('фунт', 'фунти', 'фунтів'), - ('піастр', 'піастри', 'піастрів') - ), - 'STN': ( - ('добра', 'добри', 'добр'), - ('сентімо', 'сентімо', 'сентімо') - ), - 'SYP': ( - ('фунт', 'фунти', 'фунтів'), - ('піастр', 'піастри', 'піастрів') - ), - 'SZL': ( - ('ліланґені', 'ліланґені', 'ліланґені'), - GENERIC_CENTS - ), - 'THB': ( - ('бат', 'бати', 'батів'), - ('сатанг', 'сатанги', 'сатангів') - ), - 'TJS': ( - ('сомоні', 'сомоні', 'сомоні'), - ('дірам', 'дірами', 'дірамів') - ), - 'TMT': ( - ('манат', 'манати', 'манатів'), - ('тенге', 'тенге', 'тенге') - ), - 'TND': ( - ('динар', 'динари', 'динарів'), - ('міллім', 'мілліми', 'міллімів') - ), - 'TOP': ( - ('паанга', 'паанга', 'паанга'), - ('сеніті', 'сеніті', 'сеніті') - ), - 'TRY': ( - ('ліра', 'ліри', 'лір'), - ('куруш', 'куруші', 'курушів') - ), - 'TTD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'TWD': ( - ('новий долар', 'нові долари', 'нових доларів'), - GENERIC_CENTS - ), - 'TZS': ( - ('шилінг', 'шилінги', 'шилінгів'), - GENERIC_CENTS - ), - 'UAH': ( - ('гривня', 'гривні', 'гривень'), - ('копійка', 'копійки', 'копійок') - ), - 'UGX': ( - ('шилінг', 'шилінги', 'шилінгів'), - GENERIC_CENTS - ), - 'USD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'UYU': ( - ('песо', 'песо', 'песо'), - ('сентесімо', 'сентесімо', 'сентесімо') - ), - 'UZS': ( - ('сум', 'суми', 'сумів'), - ('тиїн', 'тиїни', 'тиїнів') - ), - 'VND': ( - ('донг', 'донги', 'донгів'), - ('су', 'су', 'су') - ), - 'WST': ( - ('тала', 'тали', 'тал'), - ('сене', 'сене', 'сене') - ), - 'XCD': (GENERIC_DOLLARS, GENERIC_CENTS), - 'YER': ( - ('ріал', 'ріали', 'ріалів'), - ('філс', 'філси', 'філсів') - ), - 'ZAR': ( - ('ранд', 'ранди', 'рандів'), - GENERIC_CENTS - ), - 'ZMW': ( - ('квача', 'квачі', 'квач'), - ('нгве', 'нгве', 'нгве') - ), - } - - def setup(self): - self.negword = "мінус" - self.pointword = "кома" - - def to_cardinal(self, number): - n = str(number).replace(',', '.') - if '.' in n: - left, right = n.split('.') - leading_zero_count = len(right) - len(right.lstrip('0')) - decimal_part = ((ZERO[0] + ' ') * leading_zero_count + - self._int2word(int(right))) - return u'%s %s %s' % ( - self._int2word(int(left)), - self.pointword, - decimal_part - ) - else: - return self._int2word(int(n)) - - def pluralize(self, n, forms): - if n % 100 < 10 or n % 100 > 20: - if n % 10 == 1: - form = 0 - elif 5 > n % 10 > 1: - form = 1 - else: - form = 2 - else: - form = 2 - - return forms[form] - - def _int2word(self, n, feminine=False): - if n < 0: - return ' '.join([self.negword, self._int2word(abs(n))]) - - if n == 0: - return ZERO[0] - - words = [] - chunks = list(splitbyx(str(n), 3)) - i = len(chunks) - for x in chunks: - i -= 1 - - if x == 0: - continue - - n1, n2, n3 = get_digits(x) - - if n3 > 0: - words.append(HUNDREDS[n3][0]) - - if n2 > 1: - words.append(TWENTIES[n2][0]) - - if n2 == 1: - words.append(TENS[n1][0]) - # elif n1 > 0 and not (i > 0 and x == 1): - elif n1 > 0: - ones = ONES_FEMININE if i == 1 or feminine and i == 0 else ONES - words.append(ones[n1][0]) - - if i > 0: - words.append(self.pluralize(x, THOUSANDS[i])) - - return ' '.join(words) - - def _money_verbose(self, number, currency): - return self._int2word(number, currency in FEMININE_MONEY) - - def _cents_verbose(self, number, currency): - return self._int2word(number, currency in FEMININE_CENTS) - - @staticmethod - def last_fragment_to_ordinal(last, words, level): - n1, n2, n3 = get_digits(last) - last_two = n2*10+n1 - if last_two == 0: - words.append(HUNDREDS_ORDINALS[n3][level]) - elif level == 1 and last == 1: - return - elif last_two < 20: - if level == 0: - if n3 > 0: - words.append(HUNDREDS[n3][0]) - words.append(ONES_ORDINALS[last_two][0]) - else: - last_fragment_string = '' - if n3 > 0: - last_fragment_string += HUNDREDS_ORDINALS[n3][1] - last_fragment_string += ONES_ORDINALS[last_two][1] - words.append(last_fragment_string) - elif last_two % 10 == 0: - if level == 0: - if n3 > 0: - words.append(HUNDREDS[n3][0]) - words.append(TWENTIES_ORDINALS[n2][0]) - else: - last_fragment_string = '' - if n3 > 0: - last_fragment_string += HUNDREDS_ORDINALS[n3][1] - last_fragment_string += TWENTIES_ORDINALS[n2][1] - words.append(last_fragment_string) - else: - if level == 0: - if n3 > 0: - words.append(HUNDREDS[n3][0]) - words.append(TWENTIES[n2][0]) - words.append(ONES_ORDINALS[n1][0]) - else: - last_fragment_string = '' - if n3 > 0: - last_fragment_string += HUNDREDS_ORDINALS[n3][1] - last_fragment_string += TWENTIES_ORDINALS[n2][1] - last_fragment_string += ONES_ORDINALS[n1][1] - words.append(last_fragment_string) - - def to_ordinal(self, number): - self.verify_ordinal(number) - - words = [] - fragments = list(splitbyx(str(number), 3)) - level = 0 - last = fragments[-1] - while last == 0: - level = level + 1 - fragments.pop() - last = fragments[-1] - if len(fragments) > 1: - pre_part = self._int2word(number - (last * 1000 ** level)) - words.append(pre_part) - Num2Word_UK.last_fragment_to_ordinal( - last, - words, - 0 if level == 0 else 1 - ) - output = " ".join(words) - if last == 1 and level > 0 and output != "": - output = output + " " - if level > 0: - output = output + prefixes_ordinal[level] - return output diff --git a/fun_text_processing/num2words/num2words/lang_UR.py b/fun_text_processing/num2words/num2words/lang_UR.py deleted file mode 100644 index dab8dbfaa..000000000 --- a/fun_text_processing/num2words/num2words/lang_UR.py +++ /dev/null @@ -1,232 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import division, print_function, unicode_literals - -from . import lang_EU - - -class Num2Word_UR(lang_EU.Num2Word_EU): - # def set_high_numwords(self, high): - # max = 3 + 3 * len(high) - # for word, n in zip(high, range(max, 3, -3)): - # self.cards[10 ** n] = word - def reverse_text(text): - return ''.join(reversed(text)) - - def set_high_numwords(self, high): - max = 3 * len(high) - for word, n in zip(high, range(max, 3, -3)): - # print(word[0],word[1],n) - self.cards[10 ** n] = word[1] - # self.cards[10**n] = word + "លាន" - # try: - # ordinal_word = self.high_numwords[high] - # except KeyError: - # max = 3 + 3 * len(high) - # for word, n in zip(high, range(max, 3, -3)): - # print(word) - # print(n) - # self.cards[10 ** n] = word - - def setup(self): - super(Num2Word_UR, self).setup() - - self.negword = 'مائنس' - self.pointword = 'پوائنٹ' - - self.high_numwords = [(1000000000000,'ٹریلین'), (1000000000, 'ایک ارب'), (1000000, 'دس لاکھ') ] - - self.mid_numwords = [(1000, 'ہزار'), (100,'سو'),(90, 'نوے'), (80, 'اسی'), (70, 'ستر'), (60, 'ساٹھ'), (50, 'پچاس'), (40, 'چالیس'), (30, 'تیس')] - - self.low_numwords = ["بیس", - "انیس", - "اٹھارہ", - "سترہ", - "سولہ", - "پندرہ", - "چودہ", - "تیرہ", - "بارہ", - "گیارہ", - "دس", - "نو", - "آٹھ", - "سات", - "چھ", - "پانچ", - "چار", - "تین", - "دو", - "ایک", - "صفر"] - - self.ords = {'ایک':'پہلا', - 'دو':'دوسرا', - 'تین':'تیسرے', - 'چار':'چوتھا', - 'پانچ':'پانچویں', - 'چھ':'چھٹا', - 'سات':'ساتویں', - 'آٹھ':'آٹھویں', - 'نو':'نویں', - 'دس':'دسویں'} - - self.labeled_numbers = {'0':'صفر', - '1':'ایک', - '2':'دو', - '3':'تین', - '4':'چار', - '5':'پانچ', - '6':'چھ', - '7':'سات', - '8':'آٹھ', - '9':'نو', - '10':'دس', - '11':'گیارہ', - '12':'بارہ', - '13':'تیرہ', - '14':'چودہ', - '15':'پندرہ', - '16':'سولہ', - '17':'سترہ', - '18':'اٹھارہ', - '19':'انیس', - '20':'بیس', - '21':'اکیس', - '22':'بائیس', - '23':'تئیس', - '24':'چوبیس', - '25':'پچیس', - '26':'چھببیس', - '27':'ستائیس', - '28':'اٹھائیس', - '29':'انتیس', - '30':'تیس', - '31':'اکتیس', - '32':'بتیس', - '33':'تینتیس', - '34':'چونتیس', - '35':'پینتیس', - '36':'چھتیس', - '37':'سینتیس', - '38':'اڑتیس', - '39':'انتالیس', - '40':'چالیس', - '41':'اکتالیس', - '42':'بیالیس', - '43':'تینتالیس', - '44':'چوالیس', - '45':'پینتالیس', - '46':'چھیالیس', - '47':'سینتالیس', - '48':'اڑتالیس', - '49':'انچاس', - '50':'پچاس', - '51':'اکیاون', - '52':'باون', - '53':'ترپن', - '54':'چون', - '55':'پچپن', - '56':'چھپن', - '57':'ستاون', - '58':'اٹھاون', - '59':'انسٹھ', - '60':'ساٹھ', - '61':'اکسٹھ', - '62':'باسٹھ', - '63':'ترسٹھ', - '64':'چوسٹھ', - '65':'پینسٹھ', - '66':'چھیاسٹھ', - '67':'سڑسٹھ', - '68':'اٹھسٹھ', - '69':'انہتر', - '70':'ستر', - '71':'اکہتر', - '72':'بہتر', - '73':'تہتر', - '74':'چوہتر', - '75':'پچہتر', - '76':'چھہتر', - '77':'ستتر', - '78':'اٹھہتر', - '79':'اناسی', - '80':'اسی', - '81':'اکیاسی', - '82':'بیاسی', - '83':'تراسی', - '84':'چوراسی', - '85':'پچاسی', - '86':'چھیاسی', - '87':'ستاسی', - '88':'اٹھاسی', - '89':'نواسی', - '90':'نوے', - '91':'اکانوے', - '92':'بانوے', - '93':'ترانوے', - '94':'چورانوے', - '95':'پچانوے', - '96':'چھیانوے', - '97':'ستانوے', - '98':'اٹھانوے', - '99':'ننانوے', - '100':'سو', - '200':'دو سو', - '1000':'ایک ہزار', - '2000':'دو ہزار', - '100000':'ایک لاکھ', - '1000000':'دس لاکھ', - '2000000':'بیس لاکھ', - '10000000':'ایک کروڑ', - '100000000':'دس کروڑ', - '1000000000':'ایک ارب', - '10000000000':'دس ارب'} - - def merge(self, lpair, rpair): - ltext, lnum = lpair - rtext, rnum = rpair - if lnum == 1 and rnum < 100: - return (rtext, rnum) - elif 100 > lnum > rnum: - return ("%s%s" % (ltext, rtext), lnum + rnum) - elif lnum >= 100 > rnum: - return ("%s%s" % (ltext, rtext), lnum + rnum) - elif rnum > lnum: - return ("%s%s" % (ltext, rtext), lnum * rnum) - return ("%s%s" % (ltext, rtext), lnum + rnum) - - def to_ordinal_num(self, value): - self.verify_ordinal(value) - return "%s%s" % (value, self.to_ordinal(value)) - - def to_ordinal(self, value): - self.verify_ordinal(value) - try: - ordinal_word = self.ords[value] - except KeyError: - ordinal_word = value #TODO: check nepali ordinal word suffix - return ordinal_word - - - def to_year(self, val, suffix=None, longval=True): - if val < 0: - val = abs(val) - suffix = 'BC' if not suffix else suffix - high, low = (val // 100, val % 100) - # If year is 00XX, X00X, or beyond 9999, go cardinal. - if (high == 0 - or (high % 10 == 0 and low < 10) - or high >= 100): - valtext = self.to_cardinal(val) - else: - hightext = self.to_cardinal(high) - if low == 0: - lowtext = "सय" - elif low < 10: - lowtext = "%s" % self.to_cardinal(low) - else: - lowtext = self.to_cardinal(low) - valtext = "%s%s" % (hightext, lowtext) - return (valtext if not suffix - else "%s%s" % (valtext, suffix)) diff --git a/fun_text_processing/num2words/num2words/lang_VI.py b/fun_text_processing/num2words/num2words/lang_VI.py deleted file mode 100644 index 199994931..000000000 --- a/fun_text_processing/num2words/num2words/lang_VI.py +++ /dev/null @@ -1,86 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import unicode_literals - -to_19 = (u'không', u'một', u'hai', u'ba', u'bốn', u'năm', u'sáu', - u'bảy', u'tám', u'chín', u'mười', u'mười một', u'mười hai', - u'mười ba', u'mười bốn', u'mười lăm', u'mười sáu', u'mười bảy', - u'mười tám', u'mười chín') -tens = (u'hai mươi', u'ba mươi', u'bốn mươi', u'năm mươi', - u'sáu mươi', u'bảy mươi', u'tám mươi', u'chín mươi') -denom = ('', - u'nghìn', u'triệu', u'tỷ', u'nghìn tỷ', u'trăm nghìn tỷ', - 'Quintillion', 'Sextillion', 'Septillion', 'Octillion', 'Nonillion', - 'Decillion', 'Undecillion', 'Duodecillion', 'Tredecillion', - 'Quattuordecillion', 'Sexdecillion', 'Septendecillion', - 'Octodecillion', 'Novemdecillion', 'Vigintillion') - - -class Num2Word_VI(object): - - def _convert_nn(self, val): - if val < 20: - return to_19[val] - for (dcap, dval) in ((k, 20 + (10 * v)) for (v, k) in enumerate(tens)): - if dval + 10 > val: - if val % 10: - a = u'lăm' - if to_19[val % 10] == u'một': - a = u'mốt' - else: - a = to_19[val % 10] - if to_19[val % 10] == u'năm': - a = u'lăm' - return dcap + ' ' + a - return dcap - - def _convert_nnn(self, val): - word = '' - (mod, rem) = (val % 100, val // 100) - if rem > 0: - word = to_19[rem] + u' trăm' - if mod > 0: - word = word + ' ' - if mod > 0 and mod < 10: - if mod == 5: - word = word != '' and word + u'lẻ năm' or word + u'năm' - else: - word = word != '' and word + u'lẻ ' \ - + self._convert_nn(mod) or word + self._convert_nn(mod) - if mod >= 10: - word = word + self._convert_nn(mod) - return word - - def vietnam_number(self, val): - if val < 100: - return self._convert_nn(val) - if val < 1000: - return self._convert_nnn(val) - for (didx, dval) in ((v - 1, 1000 ** v) for v in range(len(denom))): - if dval > val: - mod = 1000 ** didx - lval = val // mod - r = val - (lval * mod) - - ret = self._convert_nnn(lval) + u' ' + denom[didx] - if 99 >= r > 0: - ret = self._convert_nnn(lval) + u' ' + denom[didx] + u' lẻ' - if r > 0: - ret = ret + ' ' + self.vietnam_number(r) - return ret - - def number_to_text(self, number): - number = '%.2f' % number - the_list = str(number).split('.') - start_word = self.vietnam_number(int(the_list[0])) - final_result = start_word - if len(the_list) > 1 and int(the_list[1]) > 0: - end_word = self.vietnam_number(int(the_list[1])) - final_result = final_result + ' phẩy ' + end_word - return final_result - - def to_cardinal(self, number): - return self.number_to_text(number) - - def to_ordinal(self, number): - return self.to_cardinal(number) diff --git a/fun_text_processing/num2words/num2words/utils.py b/fun_text_processing/num2words/num2words/utils.py deleted file mode 100644 index 5b7cd8ce3..000000000 --- a/fun_text_processing/num2words/num2words/utils.py +++ /dev/null @@ -1,20 +0,0 @@ -# -*- coding: utf-8 -*- - - -def splitbyx(n, x, format_int=True): - length = len(n) - if length > x: - start = length % x - if start > 0: - result = n[:start] - yield int(result) if format_int else result - for i in range(start, length, x): - result = n[i:i+x] - yield int(result) if format_int else result - else: - yield int(n) if format_int else n - - -def get_digits(n): - a = [int(x) for x in reversed(list(('%03d' % n)[-3:]))] - return a diff --git a/fun_text_processing/num2words/requirements-test.txt b/fun_text_processing/num2words/requirements-test.txt deleted file mode 100644 index 7165ac7e4..000000000 --- a/fun_text_processing/num2words/requirements-test.txt +++ /dev/null @@ -1,6 +0,0 @@ -flake8 -flake8-copyright -isort -pep8<1.6 -coverage -delegator.py diff --git a/fun_text_processing/num2words/setup.py b/fun_text_processing/num2words/setup.py deleted file mode 100644 index a764d1d0f..000000000 --- a/fun_text_processing/num2words/setup.py +++ /dev/null @@ -1,64 +0,0 @@ -# -*- coding: utf-8 -*- - -import re -from io import open - -from setuptools import find_packages, setup - -PACKAGE_NAME = "num2words" - -CLASSIFIERS = [ - 'Development Status :: 5 - Production/Stable', - 'Intended Audience :: Developers', - 'License :: OSI Approved :: GNU Library or Lesser General Public License ' - '(LGPL)', - 'Programming Language :: Python :: 2.7', - 'Programming Language :: Python :: 3', - 'Topic :: Software Development :: Internationalization', - 'Topic :: Software Development :: Libraries :: Python Modules', - 'Topic :: Software Development :: Localization', - 'Topic :: Text Processing :: Linguistic', -] - -LONG_DESC = open('README.rst', 'rt', encoding="utf-8").read() + '\n\n' + \ - open('CHANGES.rst', 'rt', encoding="utf-8").read() - - -def find_version(fname): - """Parse file & return version number matching 0.0.1 regex - Returns str or raises RuntimeError - """ - version = '' - with open(fname, 'r', encoding="utf-8") as fp: - reg = re.compile(r'__version__ = [\'"]([^\'"]*)[\'"]') - for line in fp: - m = reg.match(line) - if m: - version = m.group(1) - break - if not version: - raise RuntimeError('Cannot find version information') - return version - - -setup( - name=PACKAGE_NAME, - version=find_version("bin/num2words"), - description='Modules to convert numbers to multilingual words.', - long_description=LONG_DESC, - license='Alibaba Group', - author='Zhang Chong, Alibaba DAMO Academy', - author_email='', - maintainer='Zhang Chong', - maintainer_email='', - keywords=' number word numbers words convert conversion ' - 'localisation localization internationalisation ' - 'internationalization', - url='https://github.com/num2words', - packages=find_packages(exclude=['tests']), - test_suite='tests', - classifiers=CLASSIFIERS, - scripts=['bin/num2words'], - install_requires=["docopt>=0.6.2"], - tests_require=['delegator.py'], -) diff --git a/fun_text_processing/num2words/tox.ini b/fun_text_processing/num2words/tox.ini deleted file mode 100644 index 11c8f61b0..000000000 --- a/fun_text_processing/num2words/tox.ini +++ /dev/null @@ -1,32 +0,0 @@ -[tox] -envlist = flake8,isort,py27,py34,py35,py36,py37 - -[testenv] -passenv = TRAVIS TRAVIS_* -deps = - coverage - delegator.py -commands = - coverage run -m unittest discover - coverage report --fail-under=75 --omit=.tox/*,tests/*,/usr/* - coverage report --fail-under=100 --include=tests/* --skip-covered - -[testenv:flake8] -changedir = {toxinidir} -deps = - flake8 - flake8-copyright -commands = - flake8 - -[testenv:isort] -changedir = {toxinidir} -deps = - isort - delegator.py -commands = - isort --check-only --recursive --diff num2words tests - -[testenv:py27] -setenv = - PYTHONIOENCODING = UTF-8 diff --git a/fun_text_processing/text_normalization/data_loader_utils.py b/fun_text_processing/text_normalization/data_loader_utils.py index 0b9de5237..9b145405f 100644 --- a/fun_text_processing/text_normalization/data_loader_utils.py +++ b/fun_text_processing/text_normalization/data_loader_utils.py @@ -1,6 +1,3 @@ - - - import json import re import string @@ -14,7 +11,7 @@ import logging EOS_TYPE = "EOS" PUNCT_TYPE = "PUNCT" PLAIN_TYPE = "PLAIN" -Instance = namedtuple('Instance', 'token_type un_normalized normalized') +Instance = namedtuple("Instance", "token_type un_normalized normalized") known_types = [ "PLAIN", "DATE", @@ -38,7 +35,7 @@ def _load_kaggle_text_norm_file(file_path: str) -> List[Instance]: """ https://www.kaggle.com/richardwilliamsproat/text-normalization-for-english-russian-and-polish Loads text file in the Kaggle Google text normalization file format: \t\t<`self` if trivial class or normalized text> - E.g. + E.g. PLAIN Brillantaisia PLAIN is PLAIN a @@ -54,10 +51,10 @@ def _load_kaggle_text_norm_file(file_path: str) -> List[Instance]: Args: file_path: file path to text file - Returns: flat list of instances + Returns: flat list of instances """ res = [] - with open(file_path, 'r') as fp: + with open(file_path, "r") as fp: for line in fp: parts = line.strip().split("\t") if parts[0] == "": @@ -68,9 +65,13 @@ def _load_kaggle_text_norm_file(file_path: str) -> List[Instance]: l_normalized = l_normalized.lower() if l_type == PLAIN_TYPE: - res.append(Instance(token_type=l_type, un_normalized=l_token, normalized=l_token)) + res.append( + Instance(token_type=l_type, un_normalized=l_token, normalized=l_token) + ) elif l_type != PUNCT_TYPE: - res.append(Instance(token_type=l_type, un_normalized=l_token, normalized=l_normalized)) + res.append( + Instance(token_type=l_type, un_normalized=l_token, normalized=l_normalized) + ) return res @@ -78,7 +79,7 @@ def load_files(file_paths: List[str], load_func=_load_kaggle_text_norm_file) -> """ Load given list of text files using the `load_func` function. - Args: + Args: file_paths: list of file paths load_func: loading function @@ -104,9 +105,11 @@ def clean_generic(text: str) -> str: return text -def evaluate(preds: List[str], labels: List[str], input: Optional[List[str]] = None, verbose: bool = True) -> float: +def evaluate( + preds: List[str], labels: List[str], input: Optional[List[str]] = None, verbose: bool = True +) -> float: """ - Evaluates accuracy given predictions and labels. + Evaluates accuracy given predictions and labels. Args: preds: predictions @@ -174,8 +177,12 @@ def training_data_to_sentences(data: List[Instance]) -> Tuple[List[str], List[st else: sentence.append(instance) sentence_categories.update([instance.token_type]) - un_normalized = [" ".join([instance.un_normalized for instance in sentence]) for sentence in sentences] - normalized = [" ".join([instance.normalized for instance in sentence]) for sentence in sentences] + un_normalized = [ + " ".join([instance.un_normalized for instance in sentence]) for sentence in sentences + ] + normalized = [ + " ".join([instance.normalized for instance in sentence]) for sentence in sentences + ] return un_normalized, normalized, categories @@ -189,14 +196,14 @@ def post_process_punctuation(text: str) -> str: Returns: text with normalized spaces and quotes """ text = ( - text.replace('( ', '(') - .replace(' )', ')') - .replace('{ ', '{') - .replace(' }', '}') - .replace('[ ', '[') - .replace(' ]', ']') - .replace(' ', ' ') - .replace('”', '"') + text.replace("( ", "(") + .replace(" )", ")") + .replace("{ ", "{") + .replace(" }", "}") + .replace("[ ", "[") + .replace(" ]", "]") + .replace(" ", " ") + .replace("”", '"') .replace("’", "'") .replace("»", '"') .replace("«", '"') @@ -204,14 +211,14 @@ def post_process_punctuation(text: str) -> str: .replace("„", '"') .replace("´", "'") .replace("’", "'") - .replace('“', '"') + .replace("“", '"') .replace("‘", "'") - .replace('`', "'") - .replace('- -', "--") + .replace("`", "'") + .replace("- -", "--") ) for punct in "!,.:;?": - text = text.replace(f' {punct}', punct) + text = text.replace(f" {punct}", punct) return text.strip() @@ -224,12 +231,12 @@ def pre_process(text: str) -> str: Returns: text with spaces around punctuation marks """ - space_both = '[]' + space_both = "[]" for punct in space_both: - text = text.replace(punct, ' ' + punct + ' ') + text = text.replace(punct, " " + punct + " ") # remove extra space - text = re.sub(r' +', ' ', text) + text = re.sub(r" +", " ", text) return text @@ -237,13 +244,13 @@ def load_file(file_path: str) -> List[str]: """ Loads given text file with separate lines into list of string. - Args: + Args: file_path: file path Returns: flat list of string """ res = [] - with open(file_path, 'r') as fp: + with open(file_path, "r") as fp: for line in fp: res.append(line) return res @@ -256,11 +263,11 @@ def write_file(file_path: str, data: List[str]): Args: file_path: file path data: list of string - + """ - with open(file_path, 'w') as fp: + with open(file_path, "w") as fp: for line in data: - fp.write(line + '\n') + fp.write(line + "\n") def post_process_punct(input: str, normalized_text: str, add_unicode_punct: bool = False): @@ -309,7 +316,11 @@ def post_process_punct(input: str, normalized_text: str, add_unicode_punct: bool def _is_valid(idx_out, idx_in, normalized_text, input): """Check if previous or next word match (for cases when punctuation marks are part of semiotic token, i.e. some punctuation can be missing in the normalized text)""" - return (idx_out > 0 and idx_in > 0 and normalized_text[idx_out - 1] == input[idx_in - 1]) or ( + return ( + idx_out > 0 + and idx_in > 0 + and normalized_text[idx_out - 1] == input[idx_in - 1] + ) or ( idx_out < len(normalized_text) - 1 and idx_in < len(input) - 1 and normalized_text[idx_out + 1] == input[idx_in + 1] @@ -336,4 +347,4 @@ def post_process_punct(input: str, normalized_text: str, add_unicode_punct: bool logging.debug(f"Skipping post-processing of {''.join(normalized_text)} for '{punct}'") normalized_text = "".join(normalized_text) - return re.sub(r' +', ' ', normalized_text) + return re.sub(r" +", " ", normalized_text) diff --git a/fun_text_processing/text_normalization/de/__init__.py b/fun_text_processing/text_normalization/de/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/text_normalization/de/__init__.py +++ b/fun_text_processing/text_normalization/de/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/text_normalization/de/data/__init__.py b/fun_text_processing/text_normalization/de/data/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/text_normalization/de/data/__init__.py +++ b/fun_text_processing/text_normalization/de/data/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/text_normalization/de/data/electronic/__init__.py b/fun_text_processing/text_normalization/de/data/electronic/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/text_normalization/de/data/electronic/__init__.py +++ b/fun_text_processing/text_normalization/de/data/electronic/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/text_normalization/de/data/measure/__init__.py b/fun_text_processing/text_normalization/de/data/measure/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/text_normalization/de/data/measure/__init__.py +++ b/fun_text_processing/text_normalization/de/data/measure/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/text_normalization/de/data/money/__init__.py b/fun_text_processing/text_normalization/de/data/money/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/text_normalization/de/data/money/__init__.py +++ b/fun_text_processing/text_normalization/de/data/money/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/text_normalization/de/data/months/__init__.py b/fun_text_processing/text_normalization/de/data/months/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/text_normalization/de/data/months/__init__.py +++ b/fun_text_processing/text_normalization/de/data/months/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/text_normalization/de/data/numbers/__init__.py b/fun_text_processing/text_normalization/de/data/numbers/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/text_normalization/de/data/numbers/__init__.py +++ b/fun_text_processing/text_normalization/de/data/numbers/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/text_normalization/de/data/ordinals/__init__.py b/fun_text_processing/text_normalization/de/data/ordinals/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/text_normalization/de/data/ordinals/__init__.py +++ b/fun_text_processing/text_normalization/de/data/ordinals/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/text_normalization/de/data/time/__init__.py b/fun_text_processing/text_normalization/de/data/time/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/text_normalization/de/data/time/__init__.py +++ b/fun_text_processing/text_normalization/de/data/time/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/text_normalization/de/taggers/__init__.py b/fun_text_processing/text_normalization/de/taggers/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/text_normalization/de/taggers/__init__.py +++ b/fun_text_processing/text_normalization/de/taggers/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/text_normalization/de/taggers/cardinal.py b/fun_text_processing/text_normalization/de/taggers/cardinal.py index cf53ba65f..ebe52e385 100644 --- a/fun_text_processing/text_normalization/de/taggers/cardinal.py +++ b/fun_text_processing/text_normalization/de/taggers/cardinal.py @@ -1,4 +1,3 @@ - from collections import defaultdict import pynini @@ -15,7 +14,7 @@ from pynini.lib import pynutil AND = "und" -def get_ties_digit(digit_path: str, tie_path: str) -> 'pynini.FstLike': +def get_ties_digit(digit_path: str, tie_path: str) -> "pynini.FstLike": """ getting all inverse normalizations for numbers between 21 - 100 @@ -52,7 +51,7 @@ def get_ties_digit(digit_path: str, tie_path: str) -> 'pynini.FstLike': class CardinalFst(GraphFst): """ - Finite state transducer for classifying cardinals, e.g. + Finite state transducer for classifying cardinals, e.g. "101" -> cardinal { integer: "ein hundert und zehn" } Args: @@ -84,45 +83,70 @@ class CardinalFst(GraphFst): ) def hundred_non_zero(): - return (graph_digit_no_one + insert_space | pynini.cross("1", "ein ")) + pynutil.insert("hundert") + ( - pynini.closure(insert_space + pynutil.insert(AND, weight=0.0001), 0, 1) + insert_space + tens_no_zero() + return (graph_digit_no_one + insert_space | pynini.cross("1", "ein ")) + pynutil.insert( + "hundert" + ) + ( + pynini.closure(insert_space + pynutil.insert(AND, weight=0.0001), 0, 1) + + insert_space + + tens_no_zero() | pynutil.delete("00") - ) | pynutil.delete("0") + tens_no_zero() + ) | pynutil.delete( + "0" + ) + tens_no_zero() def thousand(): - return (hundred_non_zero() + insert_space + pynutil.insert("tausend") | pynutil.delete("000")) + ( - insert_space + hundred_non_zero() | pynutil.delete("000") - ) + return ( + hundred_non_zero() + insert_space + pynutil.insert("tausend") + | pynutil.delete("000") + ) + (insert_space + hundred_non_zero() | pynutil.delete("000")) optional_plural_quantity_en = pynini.closure(pynutil.insert("en", weight=-0.0001), 0, 1) optional_plural_quantity_n = pynini.closure(pynutil.insert("n", weight=-0.0001), 0, 1) graph_million = pynini.union( - hundred_non_zero() + insert_space + pynutil.insert("million") + optional_plural_quantity_en, + hundred_non_zero() + + insert_space + + pynutil.insert("million") + + optional_plural_quantity_en, pynutil.delete("000"), ) graph_billion = pynini.union( - hundred_non_zero() + insert_space + pynutil.insert("milliarde") + optional_plural_quantity_n, + hundred_non_zero() + + insert_space + + pynutil.insert("milliarde") + + optional_plural_quantity_n, pynutil.delete("000"), ) graph_trillion = pynini.union( - hundred_non_zero() + insert_space + pynutil.insert("billion") + optional_plural_quantity_en, + hundred_non_zero() + + insert_space + + pynutil.insert("billion") + + optional_plural_quantity_en, pynutil.delete("000"), ) graph_quadrillion = pynini.union( - hundred_non_zero() + insert_space + pynutil.insert("billiarde") + optional_plural_quantity_n, + hundred_non_zero() + + insert_space + + pynutil.insert("billiarde") + + optional_plural_quantity_n, pynutil.delete("000"), ) graph_quintillion = pynini.union( - hundred_non_zero() + insert_space + pynutil.insert("trillion") + optional_plural_quantity_en, + hundred_non_zero() + + insert_space + + pynutil.insert("trillion") + + optional_plural_quantity_en, pynutil.delete("000"), ) graph_sextillion = pynini.union( - hundred_non_zero() + insert_space + pynutil.insert("trilliarde") + optional_plural_quantity_n, + hundred_non_zero() + + insert_space + + pynutil.insert("trilliarde") + + optional_plural_quantity_n, pynutil.delete("000"), ) graph = pynini.union( @@ -152,7 +176,7 @@ class CardinalFst(GraphFst): self.graph = ( ((DAMO_DIGIT - "0" + pynini.closure(DAMO_DIGIT, 0)) - "0" - "1") @ pynini.cdrewrite(pynini.closure(pynutil.insert("0")), "[BOS]", "", DAMO_SIGMA) - @ DAMO_DIGIT ** 24 + @ DAMO_DIGIT**24 @ graph @ pynini.cdrewrite(delete_space, "[BOS]", "", DAMO_SIGMA) @ pynini.cdrewrite(delete_space, "", "[EOS]", DAMO_SIGMA) @@ -167,7 +191,7 @@ class CardinalFst(GraphFst): self.graph_hundred_component_at_least_one_none_zero_digit = ( ((DAMO_DIGIT - "0" + pynini.closure(DAMO_DIGIT, 0)) - "0" - "1") @ pynini.cdrewrite(pynini.closure(pynutil.insert("0")), "[BOS]", "", DAMO_SIGMA) - @ DAMO_DIGIT ** 3 + @ DAMO_DIGIT**3 @ hundred_non_zero() ) | pynini.cross("1", "eins") @@ -176,11 +200,16 @@ class CardinalFst(GraphFst): ) self.two_digit_non_zero = ( - pynini.closure(DAMO_DIGIT, 1, 2) @ self.graph_hundred_component_at_least_one_none_zero_digit + pynini.closure(DAMO_DIGIT, 1, 2) + @ self.graph_hundred_component_at_least_one_none_zero_digit ) - optional_minus_graph = pynini.closure(pynutil.insert("negative: ") + pynini.cross("-", "\"true\" "), 0, 1) + optional_minus_graph = pynini.closure( + pynutil.insert("negative: ") + pynini.cross("-", '"true" '), 0, 1 + ) - final_graph = optional_minus_graph + pynutil.insert("integer: \"") + self.graph + pynutil.insert("\"") + final_graph = ( + optional_minus_graph + pynutil.insert('integer: "') + self.graph + pynutil.insert('"') + ) final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/text_normalization/de/taggers/date.py b/fun_text_processing/text_normalization/de/taggers/date.py index cec6fd191..1c882aa30 100644 --- a/fun_text_processing/text_normalization/de/taggers/date.py +++ b/fun_text_processing/text_normalization/de/taggers/date.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.de.utils import get_abs_path, load_labels from fun_text_processing.text_normalization.en.graph_utils import ( @@ -16,7 +15,7 @@ ties_graph = pynini.invert(pynini.string_file(get_abs_path("data/numbers/ties.ts delete_leading_zero = (pynutil.delete("0") | (DAMO_DIGIT - "0")) + DAMO_DIGIT -def get_year_graph(cardinal: GraphFst) -> 'pynini.FstLike': +def get_year_graph(cardinal: GraphFst) -> "pynini.FstLike": """ Returns year verbalizations as fst @@ -28,7 +27,7 @@ def get_year_graph(cardinal: GraphFst) -> 'pynini.FstLike': cardinal: cardinal GraphFst """ - year_gt_2000 = (pynini.union("21", "20") + DAMO_DIGIT ** 2) @ cardinal.graph + year_gt_2000 = (pynini.union("21", "20") + DAMO_DIGIT**2) @ cardinal.graph graph_two_digit = delete_leading_zero @ cardinal.two_digit_non_zero hundred = pynutil.insert("hundert") @@ -42,7 +41,10 @@ def get_year_graph(cardinal: GraphFst) -> 'pynini.FstLike': graph_double_double |= pynini.accep("20") @ graph_two_digit + insert_space + graph_two_digit graph = ( graph_double_double - | (pynini.accep("1") + DAMO_DIGIT) @ graph_two_digit + insert_space + pynutil.delete("00") + hundred + | (pynini.accep("1") + DAMO_DIGIT) @ graph_two_digit + + insert_space + + pynutil.delete("00") + + hundred | year_gt_2000 ) return graph @@ -50,7 +52,7 @@ def get_year_graph(cardinal: GraphFst) -> 'pynini.FstLike': class DateFst(GraphFst): """ - Finite state transducer for classifying date, e.g. + Finite state transducer for classifying date, e.g. "01.04.2010" -> date { day: "erster" month: "april" year: "zwei tausend zehn" preserve_order: true } "1994" -> date { year: "neunzehn vier und neuzig" } "1900" -> date { year: "neunzehn hundert" } @@ -82,24 +84,24 @@ class DateFst(GraphFst): optional_leading_zero = delete_leading_zero | DAMO_DIGIT # 01, 31, 1 digit_day = optional_leading_zero @ pynini.union(*[str(x) for x in range(1, 32)]) @ numbers - day = (pynutil.insert("day: \"") + digit_day + pynutil.insert("\"")).optimize() + day = (pynutil.insert('day: "') + digit_day + pynutil.insert('"')).optimize() digit_month = optional_leading_zero @ pynini.union(*[str(x) for x in range(1, 13)]) number_to_month = digit_month @ number_to_month digit_month @= numbers - month_name = (pynutil.insert("month: \"") + month_graph + pynutil.insert("\"")).optimize() + month_name = (pynutil.insert('month: "') + month_graph + pynutil.insert('"')).optimize() month_number = ( - pynutil.insert("month: \"") + pynutil.insert('month: "') + (pynutil.add_weight(digit_month, weight=0.0001) | number_to_month) - + pynutil.insert("\"") + + pynutil.insert('"') ).optimize() # prefer cardinal over year year = pynutil.add_weight(get_year_graph(cardinal=cardinal), weight=0.001) self.year = year - year_only = pynutil.insert("year: \"") + year + pynutil.insert("\"") + year_only = pynutil.insert('year: "') + year + pynutil.insert('"') graph_dmy = ( day diff --git a/fun_text_processing/text_normalization/de/taggers/decimal.py b/fun_text_processing/text_normalization/de/taggers/decimal.py index 923775317..834629a71 100644 --- a/fun_text_processing/text_normalization/de/taggers/decimal.py +++ b/fun_text_processing/text_normalization/de/taggers/decimal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.de.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import GraphFst, insert_space @@ -7,34 +6,42 @@ from pynini.lib import pynutil quantities = pynini.string_file(get_abs_path("data/numbers/quantities.tsv")) -def get_quantity(decimal: 'pynini.FstLike', cardinal_up_to_hundred: 'pynini.FstLike') -> 'pynini.FstLike': +def get_quantity( + decimal: "pynini.FstLike", cardinal_up_to_hundred: "pynini.FstLike" +) -> "pynini.FstLike": """ Returns FST that transforms either a cardinal or decimal followed by a quantity into a numeral, e.g. 1 million -> integer_part: "eine" quantity: "million" e.g. 1.4 million -> integer_part: "eins" fractional_part: "vier" quantity: "million" - Args: + Args: decimal: decimal FST cardinal_up_to_hundred: cardinal FST """ numbers = cardinal_up_to_hundred res = ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + numbers - + pynutil.insert("\"") + + pynutil.insert('"') + pynini.accep(" ") - + pynutil.insert("quantity: \"") + + pynutil.insert('quantity: "') + quantities - + pynutil.insert("\"") + + pynutil.insert('"') + ) + res |= ( + decimal + + pynini.accep(" ") + + pynutil.insert('quantity: "') + + quantities + + pynutil.insert('"') ) - res |= decimal + pynini.accep(" ") + pynutil.insert("quantity: \"") + quantities + pynutil.insert("\"") return res class DecimalFst(GraphFst): """ - Finite state transducer for classifying decimal, e.g. + Finite state transducer for classifying decimal, e.g. -11,4006 billion -> decimal { negative: "true" integer_part: "elf" fractional_part: "vier null null sechs" quantity: "billion" preserve_order: true } 1 billion -> decimal { integer_part: "eins" quantity: "billion" preserve_order: true } Args: @@ -52,10 +59,16 @@ class DecimalFst(GraphFst): self.graph = graph_digit + pynini.closure(insert_space + graph_digit).optimize() point = pynutil.delete(",") - optional_graph_negative = pynini.closure(pynutil.insert("negative: ") + pynini.cross("-", "\"true\" "), 0, 1) + optional_graph_negative = pynini.closure( + pynutil.insert("negative: ") + pynini.cross("-", '"true" '), 0, 1 + ) - self.graph_fractional = pynutil.insert("fractional_part: \"") + self.graph + pynutil.insert("\"") - self.graph_integer = pynutil.insert("integer_part: \"") + cardinal.graph + pynutil.insert("\"") + self.graph_fractional = ( + pynutil.insert('fractional_part: "') + self.graph + pynutil.insert('"') + ) + self.graph_integer = ( + pynutil.insert('integer_part: "') + cardinal.graph + pynutil.insert('"') + ) final_graph_wo_sign = self.graph_integer + point + insert_space + self.graph_fractional self.final_graph_wo_negative = final_graph_wo_sign | get_quantity( diff --git a/fun_text_processing/text_normalization/de/taggers/electronic.py b/fun_text_processing/text_normalization/de/taggers/electronic.py index 187ac92fb..62df59104 100644 --- a/fun_text_processing/text_normalization/de/taggers/electronic.py +++ b/fun_text_processing/text_normalization/de/taggers/electronic.py @@ -1,8 +1,11 @@ - - import pynini from fun_text_processing.text_normalization.de.utils import get_abs_path, load_labels -from fun_text_processing.text_normalization.en.graph_utils import DAMO_ALPHA, DAMO_DIGIT, GraphFst, insert_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_ALPHA, + DAMO_DIGIT, + GraphFst, + insert_space, +) from pynini.lib import pynutil @@ -21,22 +24,31 @@ class ElectronicFst(GraphFst): super().__init__(name="electronic", kind="classify", deterministic=deterministic) dot = pynini.accep(".") - accepted_common_domains = [x[0] for x in load_labels(get_abs_path("data/electronic/domain.tsv"))] + accepted_common_domains = [ + x[0] for x in load_labels(get_abs_path("data/electronic/domain.tsv")) + ] accepted_common_domains = pynini.union(*accepted_common_domains) accepted_symbols = [x[0] for x in load_labels(get_abs_path("data/electronic/symbols.tsv"))] accepted_symbols = pynini.union(*accepted_symbols) - dot accepted_characters = pynini.closure(DAMO_ALPHA | DAMO_DIGIT | accepted_symbols) # email - username = pynutil.insert("username: \"") + accepted_characters + pynutil.insert("\"") + pynini.cross('@', ' ') + username = ( + pynutil.insert('username: "') + + accepted_characters + + pynutil.insert('"') + + pynini.cross("@", " ") + ) domain_graph = accepted_characters + dot + accepted_characters - domain_graph = pynutil.insert("domain: \"") + domain_graph + pynutil.insert("\"") + domain_graph = pynutil.insert('domain: "') + domain_graph + pynutil.insert('"') domain_common_graph = ( - pynutil.insert("domain: \"") + pynutil.insert('domain: "') + accepted_characters + accepted_common_domains - + pynini.closure((accepted_symbols | dot) + pynini.closure(accepted_characters, 1), 0, 1) - + pynutil.insert("\"") + + pynini.closure( + (accepted_symbols | dot) + pynini.closure(accepted_characters, 1), 0, 1 + ) + + pynutil.insert('"') ) graph = (username + domain_graph) | domain_common_graph @@ -44,7 +56,7 @@ class ElectronicFst(GraphFst): protocol_start = pynini.accep("https://") | pynini.accep("http://") protocol_end = pynini.accep("www.") protocol = protocol_start | protocol_end | (protocol_start + protocol_end) - protocol = pynutil.insert("protocol: \"") + protocol + pynutil.insert("\"") + protocol = pynutil.insert('protocol: "') + protocol + pynutil.insert('"') graph |= protocol + insert_space + (domain_graph | domain_common_graph) self.graph = graph diff --git a/fun_text_processing/text_normalization/de/taggers/fraction.py b/fun_text_processing/text_normalization/de/taggers/fraction.py index 920a82d1a..fca1af31e 100644 --- a/fun_text_processing/text_normalization/de/taggers/fraction.py +++ b/fun_text_processing/text_normalization/de/taggers/fraction.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import GraphFst from pynini.lib import pynutil @@ -21,13 +20,15 @@ class FractionFst(GraphFst): cardinal_graph = cardinal.graph self.optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("-", "\"true\" "), 0, 1 + pynutil.insert("negative: ") + pynini.cross("-", '"true" '), 0, 1 ) - self.integer = pynutil.insert("integer_part: \"") + cardinal_graph + pynutil.insert("\"") + self.integer = pynutil.insert('integer_part: "') + cardinal_graph + pynutil.insert('"') self.numerator = ( - pynutil.insert("numerator: \"") + cardinal_graph + pynini.cross(pynini.union("/", " / "), "\" ") + pynutil.insert('numerator: "') + + cardinal_graph + + pynini.cross(pynini.union("/", " / "), '" ') ) - self.denominator = pynutil.insert("denominator: \"") + cardinal_graph + pynutil.insert("\"") + self.denominator = pynutil.insert('denominator: "') + cardinal_graph + pynutil.insert('"') self.graph = ( self.optional_graph_negative diff --git a/fun_text_processing/text_normalization/de/taggers/measure.py b/fun_text_processing/text_normalization/de/taggers/measure.py index 9ad177e5e..314234573 100644 --- a/fun_text_processing/text_normalization/de/taggers/measure.py +++ b/fun_text_processing/text_normalization/de/taggers/measure.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.de.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import ( @@ -22,7 +21,9 @@ def singular_to_plural(): _n = DAMO_SIGMA + pynini.union("e") + pynutil.insert("n") _en = ( DAMO_SIGMA - + pynini.union("ent", "and", "ant", "ist", "or", "ion", "ik", "heit", "keit", "schaft", "tät", "ung") + + pynini.union( + "ent", "and", "ant", "ist", "or", "ion", "ik", "heit", "keit", "schaft", "tät", "ung" + ) + pynutil.insert("en") ) _nen = DAMO_SIGMA + pynini.union("in") + (pynutil.insert("e") | pynutil.insert("nen")) @@ -55,12 +56,16 @@ class MeasureFst(GraphFst): for False multiple transduction are generated (used for audio-based normalization) """ - def __init__(self, cardinal: GraphFst, decimal: GraphFst, fraction: GraphFst, deterministic: bool = True): + def __init__( + self, cardinal: GraphFst, decimal: GraphFst, fraction: GraphFst, deterministic: bool = True + ): super().__init__(name="measure", kind="classify", deterministic=deterministic) cardinal_graph = cardinal.graph graph_unit_singular = convert_space(unit_singular) - graph_unit_plural = graph_unit_singular @ pynini.cdrewrite(convert_space(suppletive), "", "[EOS]", DAMO_SIGMA) + graph_unit_plural = graph_unit_singular @ pynini.cdrewrite( + convert_space(suppletive), "", "[EOS]", DAMO_SIGMA + ) optional_graph_negative = pynini.closure("-", 0, 1) graph_unit_denominator = ( @@ -68,22 +73,26 @@ class MeasureFst(GraphFst): ) optional_unit_denominator = pynini.closure( - pynutil.insert(DAMO_NON_BREAKING_SPACE) + graph_unit_denominator, 0, 1, + pynutil.insert(DAMO_NON_BREAKING_SPACE) + graph_unit_denominator, + 0, + 1, ) unit_plural = ( - pynutil.insert("units: \"") + pynutil.insert('units: "') + (graph_unit_plural + (optional_unit_denominator) | graph_unit_denominator) - + pynutil.insert("\"") + + pynutil.insert('"') ) unit_singular_graph = ( - pynutil.insert("units: \"") + pynutil.insert('units: "') + ((graph_unit_singular + optional_unit_denominator) | graph_unit_denominator) - + pynutil.insert("\"") + + pynutil.insert('"') ) - subgraph_decimal = decimal.fst + insert_space + pynini.closure(pynutil.delete(" "), 0, 1) + unit_plural + subgraph_decimal = ( + decimal.fst + insert_space + pynini.closure(pynutil.delete(" "), 0, 1) + unit_plural + ) subgraph_cardinal = ( (optional_graph_negative + (pynini.closure(DAMO_DIGIT) - "1")) @ cardinal.fst @@ -101,57 +110,59 @@ class MeasureFst(GraphFst): + unit_singular_graph ) - subgraph_fraction = fraction.fst + insert_space + pynini.closure(pynutil.delete(" "), 0, 1) + unit_plural + subgraph_fraction = ( + fraction.fst + insert_space + pynini.closure(pynutil.delete(" "), 0, 1) + unit_plural + ) cardinal_dash_alpha = ( - pynutil.insert("cardinal { integer: \"") + pynutil.insert('cardinal { integer: "') + cardinal_graph - + pynutil.delete('-') - + pynutil.insert("\" } units: \"") + + pynutil.delete("-") + + pynutil.insert('" } units: "') + pynini.closure(DAMO_ALPHA, 1) - + pynutil.insert("\"") + + pynutil.insert('"') ) alpha_dash_cardinal = ( - pynutil.insert("units: \"") + pynutil.insert('units: "') + pynini.closure(DAMO_ALPHA, 1) - + pynutil.delete('-') - + pynutil.insert("\"") - + pynutil.insert(" cardinal { integer: \"") + + pynutil.delete("-") + + pynutil.insert('"') + + pynutil.insert(' cardinal { integer: "') + cardinal_graph - + pynutil.insert("\" }") + + pynutil.insert('" }') ) decimal_dash_alpha = ( pynutil.insert("decimal { ") + decimal.final_graph_wo_negative - + pynutil.delete('-') - + pynutil.insert(" } units: \"") + + pynutil.delete("-") + + pynutil.insert(' } units: "') + pynini.closure(DAMO_ALPHA, 1) - + pynutil.insert("\"") + + pynutil.insert('"') ) decimal_times = ( pynutil.insert("decimal { ") + decimal.final_graph_wo_negative - + pynutil.insert(" } units: \"") - + pynini.union('x', 'X') - + pynutil.insert("\"") + + pynutil.insert(' } units: "') + + pynini.union("x", "X") + + pynutil.insert('"') ) cardinal_times = ( - pynutil.insert("cardinal { integer: \"") + pynutil.insert('cardinal { integer: "') + cardinal_graph - + pynutil.insert("\" } units: \"") - + pynini.union('x', 'X') - + pynutil.insert("\"") + + pynutil.insert('" } units: "') + + pynini.union("x", "X") + + pynutil.insert('"') ) alpha_dash_decimal = ( - pynutil.insert("units: \"") + pynutil.insert('units: "') + pynini.closure(DAMO_ALPHA, 1) - + pynutil.delete('-') - + pynutil.insert("\"") + + pynutil.delete("-") + + pynutil.insert('"') + pynutil.insert(" decimal { ") + decimal.final_graph_wo_negative + pynutil.insert(" }") diff --git a/fun_text_processing/text_normalization/de/taggers/money.py b/fun_text_processing/text_normalization/de/taggers/money.py index 18b07f91f..71ae2a3ae 100644 --- a/fun_text_processing/text_normalization/de/taggers/money.py +++ b/fun_text_processing/text_normalization/de/taggers/money.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.de.utils import get_abs_path, load_labels from fun_text_processing.text_normalization.en.graph_utils import ( @@ -43,13 +42,19 @@ class MoneyFst(GraphFst): maj_singular_graph = convert_space(maj_singular) maj_plural_graph = maj_singular_graph - graph_maj_singular = pynutil.insert("currency_maj: \"") + maj_singular_graph + pynutil.insert("\"") - graph_maj_plural = pynutil.insert("currency_maj: \"") + maj_plural_graph + pynutil.insert("\"") + graph_maj_singular = ( + pynutil.insert('currency_maj: "') + maj_singular_graph + pynutil.insert('"') + ) + graph_maj_plural = ( + pynutil.insert('currency_maj: "') + maj_plural_graph + pynutil.insert('"') + ) optional_delete_fractional_zeros = pynini.closure( pynutil.delete(",") + pynini.closure(pynutil.delete("0"), 1), 0, 1 ) - graph_integer_one = pynutil.insert("integer_part: \"") + pynini.cross("1", "ein") + pynutil.insert("\"") + graph_integer_one = ( + pynutil.insert('integer_part: "') + pynini.cross("1", "ein") + pynutil.insert('"') + ) # only for decimals where third decimal after comma is non-zero or with quantity decimal_delete_last_zeros = ( @@ -61,11 +66,15 @@ class MoneyFst(GraphFst): ) decimal_with_quantity = DAMO_SIGMA + DAMO_ALPHA graph_decimal = ( - graph_maj_plural + insert_space + (decimal_delete_last_zeros | decimal_with_quantity) @ graph_decimal_final + graph_maj_plural + + insert_space + + (decimal_delete_last_zeros | decimal_with_quantity) @ graph_decimal_final ) graph_integer = ( - pynutil.insert("integer_part: \"") + ((DAMO_SIGMA - "1") @ cardinal_graph) + pynutil.insert("\"") + pynutil.insert('integer_part: "') + + ((DAMO_SIGMA - "1") @ cardinal_graph) + + pynutil.insert('"') ) graph_integer_only = graph_maj_singular + insert_space + graph_integer_one @@ -84,37 +93,53 @@ class MoneyFst(GraphFst): | ((DAMO_DIGIT - "0") + DAMO_DIGIT) ) - graph_min_singular = pynutil.insert(" currency_min: \"") + min_singular + pynutil.insert("\"") - graph_min_plural = pynutil.insert(" currency_min: \"") + min_plural + pynutil.insert("\"") + graph_min_singular = pynutil.insert(' currency_min: "') + min_singular + pynutil.insert('"') + graph_min_plural = pynutil.insert(' currency_min: "') + min_plural + pynutil.insert('"') # format ** euro ** cent decimal_graph_with_minor = None for curr_symbol, _ in maj_singular_labels: preserve_order = pynutil.insert(" preserve_order: true") - integer_plus_maj = graph_integer + insert_space + pynutil.insert(curr_symbol) @ graph_maj_plural - integer_plus_maj |= graph_integer_one + insert_space + pynutil.insert(curr_symbol) @ graph_maj_singular + integer_plus_maj = ( + graph_integer + insert_space + pynutil.insert(curr_symbol) @ graph_maj_plural + ) + integer_plus_maj |= ( + graph_integer_one + insert_space + pynutil.insert(curr_symbol) @ graph_maj_singular + ) # non zero integer part integer_plus_maj = (pynini.closure(DAMO_DIGIT) - "0") @ integer_plus_maj graph_fractional_one = two_digits_fractional_part @ pynini.cross("1", "ein") - graph_fractional_one = pynutil.insert("fractional_part: \"") + graph_fractional_one + pynutil.insert("\"") + graph_fractional_one = ( + pynutil.insert('fractional_part: "') + graph_fractional_one + pynutil.insert('"') + ) graph_fractional = ( - two_digits_fractional_part @ (pynini.closure(DAMO_DIGIT, 1, 2) - "1") @ cardinal.two_digit_non_zero + two_digits_fractional_part + @ (pynini.closure(DAMO_DIGIT, 1, 2) - "1") + @ cardinal.two_digit_non_zero + ) + graph_fractional = ( + pynutil.insert('fractional_part: "') + graph_fractional + pynutil.insert('"') ) - graph_fractional = pynutil.insert("fractional_part: \"") + graph_fractional + pynutil.insert("\"") - fractional_plus_min = graph_fractional + insert_space + pynutil.insert(curr_symbol) @ graph_min_plural + fractional_plus_min = ( + graph_fractional + insert_space + pynutil.insert(curr_symbol) @ graph_min_plural + ) fractional_plus_min |= ( - graph_fractional_one + insert_space + pynutil.insert(curr_symbol) @ graph_min_singular + graph_fractional_one + + insert_space + + pynutil.insert(curr_symbol) @ graph_min_singular ) - decimal_graph_with_minor_curr = integer_plus_maj + pynini.cross(",", " ") + fractional_plus_min + decimal_graph_with_minor_curr = ( + integer_plus_maj + pynini.cross(",", " ") + fractional_plus_min + ) decimal_graph_with_minor_curr |= pynutil.add_weight( integer_plus_maj + pynini.cross(",", " ") - + pynutil.insert("fractional_part: \"") + + pynutil.insert('fractional_part: "') + two_digits_fractional_part @ cardinal.two_digit_non_zero - + pynutil.insert("\""), + + pynutil.insert('"'), weight=0.0001, ) diff --git a/fun_text_processing/text_normalization/de/taggers/ordinal.py b/fun_text_processing/text_normalization/de/taggers/ordinal.py index 9a26a9c09..6481dc81f 100644 --- a/fun_text_processing/text_normalization/de/taggers/ordinal.py +++ b/fun_text_processing/text_normalization/de/taggers/ordinal.py @@ -1,4 +1,3 @@ - # Adapted from https://github.com/google/TextNormalizationCoveringGrammars # Russian minimally supervised number grammar. @@ -9,7 +8,7 @@ from pynini.lib import pynutil class OrdinalFst(GraphFst): """ - Finite state transducer for classifying cardinals, e.g. + Finite state transducer for classifying cardinals, e.g. "2." -> ordinal { integer: "zwei" } } "2tes" -> ordinal { integer: "zwei" } } @@ -27,10 +26,12 @@ class OrdinalFst(GraphFst): self.graph = ( ( pynini.closure(DAMO_DIGIT | pynini.accep(".")) - + pynutil.delete(pynutil.add_weight(pynini.union(*endings), weight=0.0001) | pynini.accep(".")) + + pynutil.delete( + pynutil.add_weight(pynini.union(*endings), weight=0.0001) | pynini.accep(".") + ) ) @ cardinal_graph ).optimize() - final_graph = pynutil.insert("integer: \"") + self.graph + pynutil.insert("\"") + final_graph = pynutil.insert('integer: "') + self.graph + pynutil.insert('"') final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/text_normalization/de/taggers/telephone.py b/fun_text_processing/text_normalization/de/taggers/telephone.py index bf769c9e8..018d6f060 100644 --- a/fun_text_processing/text_normalization/de/taggers/telephone.py +++ b/fun_text_processing/text_normalization/de/taggers/telephone.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.de.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import DAMO_DIGIT, GraphFst, insert_space @@ -7,9 +6,9 @@ from pynini.lib import pynutil class TelephoneFst(GraphFst): """ - Finite state transducer for classifying telephone, which includes country code, number part and extension + Finite state transducer for classifying telephone, which includes country code, number part and extension - E.g + E.g "+49 1234-1233" -> telephone { country_code: "plus neun und vierzig" number_part: "eins zwei drei vier eins zwei drei drei" preserve_order: true } "(012) 1234-1233" -> telephone { country_code: "null eins zwei" number_part: "eins zwei drei vier eins zwei drei drei" preserve_order: true } (0**) @@ -23,7 +22,9 @@ class TelephoneFst(GraphFst): def __init__(self, cardinal: GraphFst, deterministic: bool = True): super().__init__(name="telephone", kind="classify", deterministic=deterministic) - graph_zero = pynini.invert(pynini.string_file(get_abs_path("data/numbers/zero.tsv"))).optimize() + graph_zero = pynini.invert( + pynini.string_file(get_abs_path("data/numbers/zero.tsv")) + ).optimize() graph_digit_no_zero = pynini.invert( pynini.string_file(get_abs_path("data/numbers/digit.tsv")) ).optimize() | pynini.cross("1", "eins") @@ -31,7 +32,7 @@ class TelephoneFst(GraphFst): numbers_with_single_digits = pynini.closure(graph_digit + insert_space) + graph_digit - two_digit_and_zero = (DAMO_DIGIT ** 2 @ cardinal.two_digit_non_zero) | graph_zero + two_digit_and_zero = (DAMO_DIGIT**2 @ cardinal.two_digit_non_zero) | graph_zero # def add_space_after_two_digit(): # return pynini.closure(two_digit_and_zero + insert_space) + ( # two_digit_and_zero @@ -39,19 +40,25 @@ class TelephoneFst(GraphFst): country_code = pynini.closure(pynini.cross("+", "plus "), 0, 1) + two_digit_and_zero country_code |= ( - pynutil.delete("(") + graph_zero + insert_space + numbers_with_single_digits + pynutil.delete(")") + pynutil.delete("(") + + graph_zero + + insert_space + + numbers_with_single_digits + + pynutil.delete(")") ) country_code |= graph_zero + insert_space + numbers_with_single_digits - country_code = pynutil.insert("country_code: \"") + country_code + pynutil.insert("\"") + country_code = pynutil.insert('country_code: "') + country_code + pynutil.insert('"') del_separator = pynini.cross(pynini.union("-", " "), " ") # numbers_with_two_digits = pynini.closure(graph_digit + insert_space) + add_space_after_two_digit() + pynini.closure(insert_space + graph_digit) # numbers = numbers_with_two_digits + pynini.closure(del_separator + numbers_with_two_digits, 0, 1) - numbers = numbers_with_single_digits + pynini.closure(del_separator + numbers_with_single_digits, 0, 1) + numbers = numbers_with_single_digits + pynini.closure( + del_separator + numbers_with_single_digits, 0, 1 + ) number_length = pynini.closure((DAMO_DIGIT | pynini.union("-", " ", ")", "(")), 7) number_part = pynini.compose(number_length, numbers) - number = pynutil.insert("number_part: \"") + number_part + pynutil.insert("\"") + number = pynutil.insert('number_part: "') + number_part + pynutil.insert('"') graph = country_code + pynini.accep(" ") + number self.graph = graph diff --git a/fun_text_processing/text_normalization/de/taggers/time.py b/fun_text_processing/text_normalization/de/taggers/time.py index 8880654ad..9d5768ce3 100644 --- a/fun_text_processing/text_normalization/de/taggers/time.py +++ b/fun_text_processing/text_normalization/de/taggers/time.py @@ -1,8 +1,11 @@ - - import pynini from fun_text_processing.text_normalization.de.utils import get_abs_path -from fun_text_processing.text_normalization.en.graph_utils import DAMO_DIGIT, GraphFst, convert_space, insert_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_DIGIT, + GraphFst, + convert_space, + insert_space, +) from pynini.lib import pynutil @@ -13,7 +16,7 @@ class TimeFst(GraphFst): "2 Uhr" -> time { hours: "2" } "09:00 Uhr" -> time { hours: "2" } "02:15:10 Uhr" -> time { hours: "2" minutes: "15" seconds: "10"} - + Args: deterministic: if True will provide a single transduction option, for False multiple transduction are generated (used for audio-based normalization) @@ -29,29 +32,36 @@ class TimeFst(GraphFst): labels_minute_single = [str(x) for x in range(1, 10)] labels_minute_double = [str(x) for x in range(10, 60)] - delete_leading_zero_to_double_digit = (pynutil.delete("0") | (DAMO_DIGIT - "0")) + DAMO_DIGIT + delete_leading_zero_to_double_digit = ( + pynutil.delete("0") | (DAMO_DIGIT - "0") + ) + DAMO_DIGIT graph_hour = pynini.union(*labels_hour) graph_minute_single = pynini.union(*labels_minute_single) graph_minute_double = pynini.union(*labels_minute_double) - final_graph_hour_only = pynutil.insert("hours: \"") + graph_hour + pynutil.insert("\"") + final_graph_hour_only = pynutil.insert('hours: "') + graph_hour + pynutil.insert('"') final_graph_hour = ( - pynutil.insert("hours: \"") + delete_leading_zero_to_double_digit @ graph_hour + pynutil.insert("\"") + pynutil.insert('hours: "') + + delete_leading_zero_to_double_digit @ graph_hour + + pynutil.insert('"') ) final_graph_minute = ( - pynutil.insert("minutes: \"") + pynutil.insert('minutes: "') + (pynutil.delete("0") + graph_minute_single | graph_minute_double) - + pynutil.insert("\"") + + pynutil.insert('"') ) final_graph_second = ( - pynutil.insert("seconds: \"") + pynutil.insert('seconds: "') + (pynutil.delete("0") + graph_minute_single | graph_minute_double) - + pynutil.insert("\"") + + pynutil.insert('"') ) final_time_zone_optional = pynini.closure( - pynini.accep(" ") + pynutil.insert("zone: \"") + convert_space(time_zone_graph) + pynutil.insert("\""), + pynini.accep(" ") + + pynutil.insert('zone: "') + + convert_space(time_zone_graph) + + pynutil.insert('"'), 0, 1, ) @@ -69,9 +79,9 @@ class TimeFst(GraphFst): graph_hms = ( final_graph_hour + pynutil.delete(":") - + (pynini.cross("00", " minutes: \"0\"") | (insert_space + final_graph_minute)) + + (pynini.cross("00", ' minutes: "0"') | (insert_space + final_graph_minute)) + pynutil.delete(":") - + (pynini.cross("00", " seconds: \"0\"") | (insert_space + final_graph_second)) + + (pynini.cross("00", ' seconds: "0"') | (insert_space + final_graph_second)) + final_suffix + final_time_zone_optional + pynutil.insert(" preserve_order: true") diff --git a/fun_text_processing/text_normalization/de/taggers/tokenize_and_classify.py b/fun_text_processing/text_normalization/de/taggers/tokenize_and_classify.py index 04c4e017f..52a3c925d 100644 --- a/fun_text_processing/text_normalization/de/taggers/tokenize_and_classify.py +++ b/fun_text_processing/text_normalization/de/taggers/tokenize_and_classify.py @@ -1,4 +1,3 @@ - import os import pynini @@ -31,7 +30,7 @@ import logging class ClassifyFst(GraphFst): """ Final class that composes all other classification grammars. This class can process an entire sentence, that is lower cased. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. Args: @@ -79,7 +78,10 @@ class ClassifyFst(GraphFst): self.fraction = FractionFst(cardinal=self.cardinal, deterministic=deterministic) fraction_graph = self.fraction.fst self.measure = MeasureFst( - cardinal=self.cardinal, decimal=self.decimal, fraction=self.fraction, deterministic=deterministic + cardinal=self.cardinal, + decimal=self.decimal, + fraction=self.fraction, + deterministic=deterministic, ) measure_graph = self.measure.fst self.date = DateFst(cardinal=self.cardinal, deterministic=deterministic) @@ -91,9 +93,13 @@ class ClassifyFst(GraphFst): telephone_graph = self.telephone.fst self.electronic = ElectronicFst(deterministic=deterministic) electronic_graph = self.electronic.fst - self.money = MoneyFst(cardinal=self.cardinal, decimal=self.decimal, deterministic=deterministic) + self.money = MoneyFst( + cardinal=self.cardinal, decimal=self.decimal, deterministic=deterministic + ) money_graph = self.money.fst - self.whitelist = WhiteListFst(input_case=input_case, deterministic=deterministic, input_file=whitelist) + self.whitelist = WhiteListFst( + input_case=input_case, deterministic=deterministic, input_file=whitelist + ) whitelist_graph = self.whitelist.fst punct_graph = PunctuationFst(deterministic=deterministic).fst @@ -113,13 +119,21 @@ class ClassifyFst(GraphFst): classify |= pynutil.add_weight(word_graph, 100) - punct = pynutil.insert("tokens { ") + pynutil.add_weight(punct_graph, weight=1.1) + pynutil.insert(" }") + punct = ( + pynutil.insert("tokens { ") + + pynutil.add_weight(punct_graph, weight=1.1) + + pynutil.insert(" }") + ) token = pynutil.insert("tokens { ") + classify + pynutil.insert(" }") token_plus_punct = ( - pynini.closure(punct + pynutil.insert(" ")) + token + pynini.closure(pynutil.insert(" ") + punct) + pynini.closure(punct + pynutil.insert(" ")) + + token + + pynini.closure(pynutil.insert(" ") + punct) ) - graph = token_plus_punct + pynini.closure(pynutil.add_weight(delete_extra_space, 1.1) + token_plus_punct) + graph = token_plus_punct + pynini.closure( + pynutil.add_weight(delete_extra_space, 1.1) + token_plus_punct + ) graph = delete_space + graph + delete_space self.fst = graph.optimize() diff --git a/fun_text_processing/text_normalization/de/taggers/whitelist.py b/fun_text_processing/text_normalization/de/taggers/whitelist.py index 58d4f7b9a..3253a6b08 100644 --- a/fun_text_processing/text_normalization/de/taggers/whitelist.py +++ b/fun_text_processing/text_normalization/de/taggers/whitelist.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.de.utils import get_abs_path, load_labels from fun_text_processing.text_normalization.en.graph_utils import GraphFst, convert_space @@ -31,7 +30,8 @@ class WhiteListFst(GraphFst): graph = _get_whitelist_graph(input_case, get_abs_path("data/whitelist.tsv")) if not deterministic and input_case != "lower_cased": graph |= pynutil.add_weight( - _get_whitelist_graph("lower_cased", get_abs_path("data/whitelist.tsv")), weight=0.0001 + _get_whitelist_graph("lower_cased", get_abs_path("data/whitelist.tsv")), + weight=0.0001, ) if input_file: @@ -42,9 +42,11 @@ class WhiteListFst(GraphFst): graph = whitelist_provided if not deterministic: - units_graph = _get_whitelist_graph(input_case, file=get_abs_path("data/measure/measurements.tsv")) + units_graph = _get_whitelist_graph( + input_case, file=get_abs_path("data/measure/measurements.tsv") + ) graph |= units_graph self.graph = graph self.final_graph = convert_space(self.graph).optimize() - self.fst = (pynutil.insert("name: \"") + self.final_graph + pynutil.insert("\"")).optimize() + self.fst = (pynutil.insert('name: "') + self.final_graph + pynutil.insert('"')).optimize() diff --git a/fun_text_processing/text_normalization/de/taggers/word.py b/fun_text_processing/text_normalization/de/taggers/word.py index 362517d92..318ba5da4 100644 --- a/fun_text_processing/text_normalization/de/taggers/word.py +++ b/fun_text_processing/text_normalization/de/taggers/word.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_SPACE, GraphFst from pynini.lib import pynutil @@ -17,5 +15,5 @@ class WordFst(GraphFst): def __init__(self, deterministic: bool = True): super().__init__(name="word", kind="classify") - word = pynutil.insert("name: \"") + pynini.closure(DAMO_NOT_SPACE, 1) + pynutil.insert("\"") + word = pynutil.insert('name: "') + pynini.closure(DAMO_NOT_SPACE, 1) + pynutil.insert('"') self.fst = word.optimize() diff --git a/fun_text_processing/text_normalization/de/utils.py b/fun_text_processing/text_normalization/de/utils.py index e233075dd..1a1f30f2e 100644 --- a/fun_text_processing/text_normalization/de/utils.py +++ b/fun_text_processing/text_normalization/de/utils.py @@ -1,4 +1,3 @@ - import csv import os @@ -11,13 +10,13 @@ def get_abs_path(rel_path): Args: rel_path: relative path to this file - + Returns absolute path """ abs_path = os.path.dirname(os.path.abspath(__file__)) + os.sep + rel_path if not os.path.exists(abs_path): - logging.warning(f'{abs_path} does not exist') + logging.warning(f"{abs_path} does not exist") return abs_path diff --git a/fun_text_processing/text_normalization/de/verbalizers/__init__.py b/fun_text_processing/text_normalization/de/verbalizers/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/text_normalization/de/verbalizers/__init__.py +++ b/fun_text_processing/text_normalization/de/verbalizers/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/text_normalization/de/verbalizers/cardinal.py b/fun_text_processing/text_normalization/de/verbalizers/cardinal.py index aef0c5619..13e181096 100644 --- a/fun_text_processing/text_normalization/de/verbalizers/cardinal.py +++ b/fun_text_processing/text_normalization/de/verbalizers/cardinal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst from pynini.lib import pynutil @@ -16,11 +15,11 @@ class CardinalFst(GraphFst): def __init__(self, deterministic: bool = True): super().__init__(name="cardinal", kind="verbalize", deterministic=deterministic) - optional_sign = pynini.closure(pynini.cross("negative: \"true\" ", "minus "), 0, 1) + optional_sign = pynini.closure(pynini.cross('negative: "true" ', "minus "), 0, 1) self.optional_sign = optional_sign integer = pynini.closure(DAMO_NOT_QUOTE, 1) - self.integer = pynutil.delete(" \"") + integer + pynutil.delete("\"") + self.integer = pynutil.delete(' "') + integer + pynutil.delete('"') integer = pynutil.delete("integer:") + self.integer self.numbers = integer diff --git a/fun_text_processing/text_normalization/de/verbalizers/date.py b/fun_text_processing/text_normalization/de/verbalizers/date.py index c3413a916..e05525f06 100644 --- a/fun_text_processing/text_normalization/de/verbalizers/date.py +++ b/fun_text_processing/text_normalization/de/verbalizers/date.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.de.utils import get_abs_path, load_labels from fun_text_processing.text_normalization.en.graph_utils import ( @@ -25,20 +24,28 @@ class DateFst(GraphFst): def __init__(self, ordinal: GraphFst, deterministic: bool = True): super().__init__(name="date", kind="verbalize", deterministic=deterministic) - day_cardinal = pynutil.delete("day: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") - day = day_cardinal @ pynini.cdrewrite(ordinal.ordinal_stem, "", "[EOS]", DAMO_SIGMA) + pynutil.insert("ter") + day_cardinal = ( + pynutil.delete('day: "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') + ) + day = day_cardinal @ pynini.cdrewrite( + ordinal.ordinal_stem, "", "[EOS]", DAMO_SIGMA + ) + pynutil.insert("ter") - months_names = pynini.union(*[x[1] for x in load_labels(get_abs_path("data/months/abbr_to_name.tsv"))]) - month = pynutil.delete("month: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + months_names = pynini.union( + *[x[1] for x in load_labels(get_abs_path("data/months/abbr_to_name.tsv"))] + ) + month = pynutil.delete('month: "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') final_month = month @ months_names final_month |= month @ pynini.difference(DAMO_SIGMA, months_names) @ pynini.cdrewrite( ordinal.ordinal_stem, "", "[EOS]", DAMO_SIGMA ) + pynutil.insert("ter") - year = pynutil.delete("year: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + year = pynutil.delete('year: "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') # day month year - graph_dmy = day + pynini.accep(" ") + final_month + pynini.closure(pynini.accep(" ") + year, 0, 1) + graph_dmy = ( + day + pynini.accep(" ") + final_month + pynini.closure(pynini.accep(" ") + year, 0, 1) + ) graph_dmy |= final_month + pynini.accep(" ") + year self.graph = graph_dmy | year diff --git a/fun_text_processing/text_normalization/de/verbalizers/decimal.py b/fun_text_processing/text_normalization/de/verbalizers/decimal.py index ee330db6a..db4100af6 100644 --- a/fun_text_processing/text_normalization/de/verbalizers/decimal.py +++ b/fun_text_processing/text_normalization/de/verbalizers/decimal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.de.taggers.decimal import quantities from fun_text_processing.text_normalization.en.graph_utils import ( @@ -12,8 +11,8 @@ from pynini.lib import pynutil class DecimalFst(GraphFst): """ - Finite state transducer for classifying decimal, e.g. - decimal { negative: "true" integer_part: "elf" fractional_part: "vier null sechs" quantity: "billionen" } -> minus elf komma vier null sechs billionen + Finite state transducer for classifying decimal, e.g. + decimal { negative: "true" integer_part: "elf" fractional_part: "vier null sechs" quantity: "billionen" } -> minus elf komma vier null sechs billionen decimal { integer_part: "eins" quantity: "billion" } -> eins billion """ @@ -22,21 +21,34 @@ class DecimalFst(GraphFst): super().__init__(name="decimal", kind="classify", deterministic=deterministic) delete_space = pynutil.delete(" ") - self.optional_sign = pynini.closure(pynini.cross("negative: \"true\"", "minus ") + delete_space, 0, 1) - self.integer = pynutil.delete("integer_part: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + self.optional_sign = pynini.closure( + pynini.cross('negative: "true"', "minus ") + delete_space, 0, 1 + ) + self.integer = ( + pynutil.delete('integer_part: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + ) self.fractional_default = ( - pynutil.delete("fractional_part: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + pynutil.delete('fractional_part: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') ) self.fractional = pynutil.insert(" komma ") + self.fractional_default self.quantity = ( - delete_space + insert_space + pynutil.delete("quantity: \"") + quantities + pynutil.delete("\"") + delete_space + + insert_space + + pynutil.delete('quantity: "') + + quantities + + pynutil.delete('"') ) self.optional_quantity = pynini.closure(self.quantity, 0, 1) graph = self.optional_sign + ( - self.integer + self.quantity | self.integer + delete_space + self.fractional + self.optional_quantity + self.integer + self.quantity + | self.integer + delete_space + self.fractional + self.optional_quantity ) self.numbers = graph diff --git a/fun_text_processing/text_normalization/de/verbalizers/electronic.py b/fun_text_processing/text_normalization/de/verbalizers/electronic.py index ca0faacc2..e8cbe5d9b 100644 --- a/fun_text_processing/text_normalization/de/verbalizers/electronic.py +++ b/fun_text_processing/text_normalization/de/verbalizers/electronic.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.de.utils import get_abs_path from fun_text_processing.text_normalization.en.graph_utils import ( @@ -28,7 +27,9 @@ class ElectronicFst(GraphFst): graph_digit_no_zero = pynini.invert( pynini.string_file(get_abs_path("data/numbers/digit.tsv")) ).optimize() | pynini.cross("1", "eins") - graph_zero = pynini.invert(pynini.string_file(get_abs_path("data/numbers/zero.tsv"))).optimize() + graph_zero = pynini.invert( + pynini.string_file(get_abs_path("data/numbers/zero.tsv")) + ).optimize() graph_digit = graph_digit_no_zero | graph_zero graph_symbols = pynini.string_file(get_abs_path("data/electronic/symbols.tsv")).optimize() server_common = pynini.string_file(get_abs_path("data/electronic/server_name.tsv")) @@ -41,18 +42,20 @@ class ElectronicFst(GraphFst): verbalize_characters = pynini.cdrewrite(graph_symbols | graph_digit, "", "", DAMO_SIGMA) - user_name = pynutil.delete("username: \"") + add_space_after_char() + pynutil.delete("\"") + user_name = pynutil.delete('username: "') + add_space_after_char() + pynutil.delete('"') user_name @= verbalize_characters - convert_defaults = pynutil.add_weight(DAMO_NOT_QUOTE, weight=0.0001) | domain_common | server_common + convert_defaults = ( + pynutil.add_weight(DAMO_NOT_QUOTE, weight=0.0001) | domain_common | server_common + ) domain = convert_defaults + pynini.closure(insert_space + convert_defaults) domain @= verbalize_characters - domain = pynutil.delete("domain: \"") + domain + pynutil.delete("\"") + domain = pynutil.delete('domain: "') + domain + pynutil.delete('"') protocol = ( - pynutil.delete("protocol: \"") + pynutil.delete('protocol: "') + add_space_after_char() @ pynini.cdrewrite(graph_symbols, "", "", DAMO_SIGMA) - + pynutil.delete("\"") + + pynutil.delete('"') ) self.graph = (pynini.closure(protocol + pynini.accep(" "), 0, 1) + domain) | ( user_name + pynini.accep(" ") + pynutil.insert("at ") + domain diff --git a/fun_text_processing/text_normalization/de/verbalizers/fraction.py b/fun_text_processing/text_normalization/de/verbalizers/fraction.py index a77e26dcd..e926d4bae 100644 --- a/fun_text_processing/text_normalization/de/verbalizers/fraction.py +++ b/fun_text_processing/text_normalization/de/verbalizers/fraction.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_NOT_QUOTE, @@ -26,27 +25,38 @@ class FractionFst(GraphFst): def __init__(self, ordinal: GraphFst, deterministic: bool = True): super().__init__(name="fraction", kind="verbalize", deterministic=deterministic) - optional_sign = pynini.closure(pynini.cross("negative: \"true\"", "minus ") + pynutil.delete(" "), 0, 1) - change_one = pynini.cdrewrite( - pynutil.add_weight(pynini.cross("eins", "ein"), weight=-0.0001), "[BOS]", "[EOS]", DAMO_SIGMA + optional_sign = pynini.closure( + pynini.cross('negative: "true"', "minus ") + pynutil.delete(" "), 0, 1 ) - change_numerator_two = pynini.cdrewrite(pynini.cross("zweitel", "halbe"), "[BOS]", "[EOS]", DAMO_SIGMA) - integer = pynutil.delete("integer_part: \"") + change_one + pynutil.delete("\" ") - numerator = pynutil.delete("numerator: \"") + change_one + pynutil.delete("\" ") + change_one = pynini.cdrewrite( + pynutil.add_weight(pynini.cross("eins", "ein"), weight=-0.0001), + "[BOS]", + "[EOS]", + DAMO_SIGMA, + ) + change_numerator_two = pynini.cdrewrite( + pynini.cross("zweitel", "halbe"), "[BOS]", "[EOS]", DAMO_SIGMA + ) + integer = pynutil.delete('integer_part: "') + change_one + pynutil.delete('" ') + numerator = pynutil.delete('numerator: "') + change_one + pynutil.delete('" ') denominator = ( - pynutil.delete("denominator: \"") + pynutil.delete('denominator: "') + pynini.closure(DAMO_NOT_QUOTE) @ ( - pynini.cdrewrite(pynini.closure(ordinal.ordinal_stem, 0, 1), "", "[EOS]", DAMO_SIGMA) + pynini.cdrewrite( + pynini.closure(ordinal.ordinal_stem, 0, 1), "", "[EOS]", DAMO_SIGMA + ) + pynutil.insert("tel") ) @ change_numerator_two - + pynutil.delete("\"") + + pynutil.delete('"') ) integer += insert_space + pynini.closure(pynutil.insert("und ", weight=0.001), 0, 1) - denominator_one_half = pynini.cdrewrite(pynini.cross("ein halbe", "ein halb"), "[BOS]", "[EOS]", DAMO_SIGMA) + denominator_one_half = pynini.cdrewrite( + pynini.cross("ein halbe", "ein halb"), "[BOS]", "[EOS]", DAMO_SIGMA + ) fraction_default = (numerator + insert_space + denominator) @ denominator_one_half diff --git a/fun_text_processing/text_normalization/de/verbalizers/measure.py b/fun_text_processing/text_normalization/de/verbalizers/measure.py index e802ce80d..90156a5b3 100644 --- a/fun_text_processing/text_normalization/de/verbalizers/measure.py +++ b/fun_text_processing/text_normalization/de/verbalizers/measure.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_NOT_QUOTE, @@ -14,7 +13,7 @@ class MeasureFst(GraphFst): Finite state transducer for verbalizing measure, e.g. measure { cardinal { integer: "zwei" units: "unzen" } } -> "zwei unzen" measure { cardinal { integer_part: "zwei" quantity: "millionen" units: "unzen" } } -> "zwei millionen unzen" - + Args: decimal: decimal GraphFst cardinal: cardinal GraphFst @@ -23,9 +22,11 @@ class MeasureFst(GraphFst): for False multiple transduction are generated (used for audio-based normalization) """ - def __init__(self, decimal: GraphFst, cardinal: GraphFst, fraction: GraphFst, deterministic: bool): + def __init__( + self, decimal: GraphFst, cardinal: GraphFst, fraction: GraphFst, deterministic: bool + ): super().__init__(name="measure", kind="verbalize", deterministic=deterministic) - unit = pynutil.delete("units: \"") + pynini.closure(DAMO_NOT_QUOTE) + pynutil.delete("\"") + unit = pynutil.delete('units: "') + pynini.closure(DAMO_NOT_QUOTE) + pynutil.delete('"') graph_decimal = decimal.fst graph_cardinal = cardinal.fst diff --git a/fun_text_processing/text_normalization/de/verbalizers/money.py b/fun_text_processing/text_normalization/de/verbalizers/money.py index a562b2c2d..f92b3928f 100644 --- a/fun_text_processing/text_normalization/de/verbalizers/money.py +++ b/fun_text_processing/text_normalization/de/verbalizers/money.py @@ -1,6 +1,9 @@ - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_preserve_order +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_preserve_order, +) from pynini.lib import pynutil @@ -25,14 +28,28 @@ class MoneyFst(GraphFst): keep_space = pynini.accep(" ") - maj = pynutil.delete("currency_maj: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") - min = pynutil.delete("currency_min: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") - - fractional_part = ( - pynutil.delete("fractional_part: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + maj = ( + pynutil.delete('currency_maj: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + ) + min = ( + pynutil.delete('currency_min: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') ) - integer_part = pynutil.delete("integer_part: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + fractional_part = ( + pynutil.delete('fractional_part: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + ) + + integer_part = ( + pynutil.delete('integer_part: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + ) optional_add_and = pynini.closure(pynutil.insert("und "), 0, 1) # *** currency_maj diff --git a/fun_text_processing/text_normalization/de/verbalizers/ordinal.py b/fun_text_processing/text_normalization/de/verbalizers/ordinal.py index 4c18c66a1..8b37e3314 100644 --- a/fun_text_processing/text_normalization/de/verbalizers/ordinal.py +++ b/fun_text_processing/text_normalization/de/verbalizers/ordinal.py @@ -1,7 +1,10 @@ - import pynini from fun_text_processing.text_normalization.de.utils import get_abs_path -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, DAMO_SIGMA, GraphFst +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + DAMO_SIGMA, + GraphFst, +) from pynini.lib import pynutil @@ -22,14 +25,19 @@ class OrdinalFst(GraphFst): graph_ties = pynini.string_file(get_abs_path("data/ordinals/ties.tsv")).invert() graph_thousands = pynini.string_file(get_abs_path("data/ordinals/thousands.tsv")).invert() - graph = pynutil.delete("integer: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + graph = ( + pynutil.delete('integer: "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') + ) suffixes = pynini.union("ten", "tem", "ter", "tes", "te") convert_rest = pynutil.insert(suffixes, weight=0.01) self.ordinal_stem = graph_digit | graph_ties | graph_thousands suffix = pynini.cdrewrite( - pynini.closure(self.ordinal_stem, 0, 1) + convert_rest, "", "[EOS]", DAMO_SIGMA, + pynini.closure(self.ordinal_stem, 0, 1) + convert_rest, + "", + "[EOS]", + DAMO_SIGMA, ).optimize() self.graph = pynini.compose(graph, suffix) self.suffix = suffix diff --git a/fun_text_processing/text_normalization/de/verbalizers/telephone.py b/fun_text_processing/text_normalization/de/verbalizers/telephone.py index 417dbb068..43a6d1d88 100644 --- a/fun_text_processing/text_normalization/de/verbalizers/telephone.py +++ b/fun_text_processing/text_normalization/de/verbalizers/telephone.py @@ -1,13 +1,16 @@ - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_preserve_order +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_preserve_order, +) from pynini.lib import pynutil class TelephoneFst(GraphFst): """ Finite state transducer for verbalizing telephone, e.g. - telephone { country_code: "plus neun und vierzig" number_part: "null eins eins eins null null null" } + telephone { country_code: "plus neun und vierzig" number_part: "null eins eins eins null null null" } -> "plus neun und vierzig null eins eins eins null null null" Args: @@ -18,9 +21,17 @@ class TelephoneFst(GraphFst): def __init__(self, deterministic: bool = True): super().__init__(name="telephone", kind="verbalize", deterministic=deterministic) - country_code = pynutil.delete("country_code: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + country_code = ( + pynutil.delete('country_code: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + ) - number_part = pynutil.delete("number_part: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + number_part = ( + pynutil.delete('number_part: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + ) self.graph = country_code + pynini.accep(" ") + number_part diff --git a/fun_text_processing/text_normalization/de/verbalizers/time.py b/fun_text_processing/text_normalization/de/verbalizers/time.py index 482a4c2d8..f00502108 100644 --- a/fun_text_processing/text_normalization/de/verbalizers/time.py +++ b/fun_text_processing/text_normalization/de/verbalizers/time.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.de.utils import get_abs_path, load_labels from fun_text_processing.text_normalization.en.graph_utils import ( @@ -33,24 +32,35 @@ class TimeFst(GraphFst): # add weight so when using inverse text normalization this conversion is depriotized night_to_early = pynutil.add_weight( - pynini.invert(pynini.string_file(get_abs_path("data/time/hour_to_night.tsv"))).optimize(), weight=0.0001 + pynini.invert( + pynini.string_file(get_abs_path("data/time/hour_to_night.tsv")) + ).optimize(), + weight=0.0001, ) - hour_to = pynini.invert(pynini.string_file(get_abs_path("data/time/hour_to.tsv"))).optimize() - minute_to = pynini.invert(pynini.string_file(get_abs_path("data/time/minute_to.tsv"))).optimize() + hour_to = pynini.invert( + pynini.string_file(get_abs_path("data/time/hour_to.tsv")) + ).optimize() + minute_to = pynini.invert( + pynini.string_file(get_abs_path("data/time/minute_to.tsv")) + ).optimize() time_zone_graph = pynini.invert( - convert_space(pynini.union(*[x[1] for x in load_labels(get_abs_path("data/time/time_zone.tsv"))])) + convert_space( + pynini.union(*[x[1] for x in load_labels(get_abs_path("data/time/time_zone.tsv"))]) + ) ) - graph_zero = pynini.invert(pynini.string_file(get_abs_path("data/numbers/zero.tsv"))).optimize() + graph_zero = pynini.invert( + pynini.string_file(get_abs_path("data/numbers/zero.tsv")) + ).optimize() number_verbalization = graph_zero | cardinal_tagger.two_digit_non_zero - hour = pynutil.delete("hours: \"") + pynini.closure(DAMO_DIGIT, 1) + pynutil.delete("\"") + hour = pynutil.delete('hours: "') + pynini.closure(DAMO_DIGIT, 1) + pynutil.delete('"') hour_verbalized = hour @ number_verbalization @ pynini.cdrewrite( pynini.cross("eins", "ein"), "[BOS]", "[EOS]", DAMO_SIGMA ) + pynutil.insert(" uhr") - minute = pynutil.delete("minutes: \"") + pynini.closure(DAMO_DIGIT, 1) + pynutil.delete("\"") - zone = pynutil.delete("zone: \"") + time_zone_graph + pynutil.delete("\"") + minute = pynutil.delete('minutes: "') + pynini.closure(DAMO_DIGIT, 1) + pynutil.delete('"') + zone = pynutil.delete('zone: "') + time_zone_graph + pynutil.delete('"') optional_zone = pynini.closure(pynini.accep(" ") + zone, 0, 1) - second = pynutil.delete("seconds: \"") + pynini.closure(DAMO_DIGIT, 1) + pynutil.delete("\"") + second = pynutil.delete('seconds: "') + pynini.closure(DAMO_DIGIT, 1) + pynutil.delete('"') graph_hms = ( hour_verbalized + pynini.accep(" ") @@ -62,7 +72,8 @@ class TimeFst(GraphFst): + optional_zone ) graph_hms @= pynini.cdrewrite( - pynini.cross("eins minuten", "eine minute") | pynini.cross("eins sekunden", "eine sekunde"), + pynini.cross("eins minuten", "eine minute") + | pynini.cross("eins sekunden", "eine sekunde"), pynini.union(" ", "[BOS]"), "", DAMO_SIGMA, @@ -81,18 +92,26 @@ class TimeFst(GraphFst): + pynini.accep(" ") + pynutil.insert("nach ") # + hour @ number_verbalization - + hour @ pynini.cdrewrite(night_to_early, "[BOS]", "[EOS]", DAMO_SIGMA) @ number_verbalization + + hour + @ pynini.cdrewrite(night_to_early, "[BOS]", "[EOS]", DAMO_SIGMA) + @ number_verbalization ) graph_m30_h = ( minute @ pynini.cross("30", "halb") + pynini.accep(" ") - + hour @ pynini.cdrewrite(night_to_early, "[BOS]", "[EOS]", DAMO_SIGMA) @ hour_to @ number_verbalization + + hour + @ pynini.cdrewrite(night_to_early, "[BOS]", "[EOS]", DAMO_SIGMA) + @ hour_to + @ number_verbalization ) graph_m_to_h = ( minute @ minute_to @ min_29 @ (number_verbalization | pynini.cross("15", "viertel")) + pynini.accep(" ") + pynutil.insert("vor ") - + hour @ pynini.cdrewrite(night_to_early, "[BOS]", "[EOS]", DAMO_SIGMA) @ hour_to @ number_verbalization + + hour + @ pynini.cdrewrite(night_to_early, "[BOS]", "[EOS]", DAMO_SIGMA) + @ hour_to + @ number_verbalization ) self.graph = ( diff --git a/fun_text_processing/text_normalization/de/verbalizers/verbalize.py b/fun_text_processing/text_normalization/de/verbalizers/verbalize.py index 1533327b7..dfd555549 100644 --- a/fun_text_processing/text_normalization/de/verbalizers/verbalize.py +++ b/fun_text_processing/text_normalization/de/verbalizers/verbalize.py @@ -1,4 +1,3 @@ - from fun_text_processing.text_normalization.de.taggers.cardinal import CardinalFst as CardinalTagger from fun_text_processing.text_normalization.de.verbalizers.cardinal import CardinalFst from fun_text_processing.text_normalization.de.verbalizers.date import DateFst @@ -17,7 +16,7 @@ from fun_text_processing.text_normalization.en.verbalizers.whitelist import Whit class VerbalizeFst(GraphFst): """ Composes other verbalizer grammars. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. Args: @@ -38,7 +37,9 @@ class VerbalizeFst(GraphFst): fraction_graph = fraction.fst date = DateFst(ordinal=ordinal) date_graph = date.fst - measure = MeasureFst(cardinal=cardinal, decimal=decimal, fraction=fraction, deterministic=deterministic) + measure = MeasureFst( + cardinal=cardinal, decimal=decimal, fraction=fraction, deterministic=deterministic + ) measure_graph = measure.fst electronic = ElectronicFst(deterministic=deterministic) electronic_graph = electronic.fst diff --git a/fun_text_processing/text_normalization/de/verbalizers/verbalize_final.py b/fun_text_processing/text_normalization/de/verbalizers/verbalize_final.py index 6587412a2..810a1c0e1 100644 --- a/fun_text_processing/text_normalization/de/verbalizers/verbalize_final.py +++ b/fun_text_processing/text_normalization/de/verbalizers/verbalize_final.py @@ -1,4 +1,3 @@ - import os import pynini @@ -18,7 +17,7 @@ import logging class VerbalizeFinalFst(GraphFst): """ Finite state transducer that verbalizes an entire sentence - + Args: deterministic: if True will provide a single transduction option, for False multiple options (used for audio-based normalization) @@ -26,16 +25,20 @@ class VerbalizeFinalFst(GraphFst): overwrite_cache: set to True to overwrite .far files """ - def __init__(self, deterministic: bool = True, cache_dir: str = None, overwrite_cache: bool = False): + def __init__( + self, deterministic: bool = True, cache_dir: str = None, overwrite_cache: bool = False + ): super().__init__(name="verbalize_final", kind="verbalize", deterministic=deterministic) far_file = None if cache_dir is not None and cache_dir != "None": os.makedirs(cache_dir, exist_ok=True) - far_file = os.path.join(cache_dir, f"de_tn_{deterministic}_deterministic_verbalizer.far") + far_file = os.path.join( + cache_dir, f"de_tn_{deterministic}_deterministic_verbalizer.far" + ) if not overwrite_cache and far_file and os.path.exists(far_file): self.fst = pynini.Far(far_file, mode="r")["verbalize"] - logging.info(f'VerbalizeFinalFst graph was restored from {far_file}.') + logging.info(f"VerbalizeFinalFst graph was restored from {far_file}.") else: verbalize = VerbalizeFst(deterministic=deterministic).fst word = WordFst(deterministic=deterministic).fst diff --git a/fun_text_processing/text_normalization/en/__init__.py b/fun_text_processing/text_normalization/en/__init__.py index 40ce4aa38..27a262458 100644 --- a/fun_text_processing/text_normalization/en/__init__.py +++ b/fun_text_processing/text_normalization/en/__init__.py @@ -1,5 +1,3 @@ - - from fun_text_processing.text_normalization.en.taggers.tokenize_and_classify import ClassifyFst from fun_text_processing.text_normalization.en.verbalizers.verbalize import VerbalizeFst from fun_text_processing.text_normalization.en.verbalizers.verbalize_final import VerbalizeFinalFst diff --git a/fun_text_processing/text_normalization/en/clean_eval_data.py b/fun_text_processing/text_normalization/en/clean_eval_data.py index da9d7a94a..8a782e9d9 100644 --- a/fun_text_processing/text_normalization/en/clean_eval_data.py +++ b/fun_text_processing/text_normalization/en/clean_eval_data.py @@ -1,5 +1,3 @@ - - from argparse import ArgumentParser from typing import List @@ -55,7 +53,7 @@ class Filter: Args: processes given instance with process function - + Returns: processed instance if instance belongs to expected class type or original instance """ if instance.token_type != self.class_type: @@ -73,7 +71,9 @@ def process_cardinal_1(instance: Instance) -> Instance: normalized = instance.normalized un_normalized = re.sub(r"[^0-9]", "", un_normalized) normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_ordinal_1(instance: Instance) -> bool: @@ -86,7 +86,9 @@ def process_ordinal_1(instance: Instance) -> Instance: normalized = instance.normalized un_normalized = re.sub(r"[,\s]", "", un_normalized) normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_decimal_1(instance: Instance) -> bool: @@ -99,7 +101,9 @@ def process_decimal_1(instance: Instance) -> Instance: un_normalized = re.sub(r",", "", un_normalized) normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_measure_1(instance: Instance) -> bool: @@ -116,7 +120,9 @@ def process_measure_1(instance: Instance) -> Instance: normalized = re.sub(r"[^a-z\s]", "", normalized) normalized = re.sub(r"per ([a-z\s]*)s$", r"per \1", normalized) normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_money_1(instance: Instance) -> bool: @@ -133,7 +139,9 @@ def process_money_1(instance: Instance) -> Instance: un_normalized = re.sub(r"(\d)m\s*$", r"\1 million", un_normalized) un_normalized = re.sub(r"(\d)bn?\s*$", r"\1 billion", un_normalized) normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_time_1(instance: Instance) -> bool: @@ -148,7 +156,9 @@ def process_time_1(instance: Instance) -> Instance: un_normalized = re.sub(r"(\d)\s?p\s?m\s?", r"\1 p.m.", un_normalized) normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_plain_1(instance: Instance) -> bool: @@ -159,7 +169,9 @@ def filter_plain_1(instance: Instance) -> bool: def process_plain_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_punct_1(instance: Instance) -> bool: @@ -170,7 +182,9 @@ def filter_punct_1(instance: Instance) -> bool: def process_punct_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_date_1(instance: Instance) -> bool: @@ -183,7 +197,9 @@ def process_date_1(instance: Instance) -> Instance: un_normalized = re.sub(r",", "", un_normalized) normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_letters_1(instance: Instance) -> bool: @@ -195,7 +211,9 @@ def process_letters_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_verbatim_1(instance: Instance) -> bool: @@ -206,7 +224,9 @@ def filter_verbatim_1(instance: Instance) -> bool: def process_verbatim_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_digit_1(instance: Instance) -> bool: @@ -218,7 +238,9 @@ def process_digit_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_telephone_1(instance: Instance) -> bool: @@ -230,7 +252,9 @@ def process_telephone_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_electronic_1(instance: Instance) -> bool: @@ -242,7 +266,9 @@ def process_electronic_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_fraction_1(instance: Instance) -> bool: @@ -254,7 +280,9 @@ def process_fraction_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) def filter_address_1(instance: Instance) -> bool: @@ -266,27 +294,51 @@ def process_address_1(instance: Instance) -> Instance: un_normalized = instance.un_normalized normalized = instance.normalized normalized = re.sub(r"[^a-z ]", "", normalized) - return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + return Instance( + token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized + ) filters = [] -filters.append(Filter(class_type="CARDINAL", process_func=process_cardinal_1, filter_func=filter_cardinal_1)) -filters.append(Filter(class_type="ORDINAL", process_func=process_ordinal_1, filter_func=filter_ordinal_1)) -filters.append(Filter(class_type="DECIMAL", process_func=process_decimal_1, filter_func=filter_decimal_1)) -filters.append(Filter(class_type="MEASURE", process_func=process_measure_1, filter_func=filter_measure_1)) +filters.append( + Filter(class_type="CARDINAL", process_func=process_cardinal_1, filter_func=filter_cardinal_1) +) +filters.append( + Filter(class_type="ORDINAL", process_func=process_ordinal_1, filter_func=filter_ordinal_1) +) +filters.append( + Filter(class_type="DECIMAL", process_func=process_decimal_1, filter_func=filter_decimal_1) +) +filters.append( + Filter(class_type="MEASURE", process_func=process_measure_1, filter_func=filter_measure_1) +) filters.append(Filter(class_type="MONEY", process_func=process_money_1, filter_func=filter_money_1)) filters.append(Filter(class_type="TIME", process_func=process_time_1, filter_func=filter_time_1)) filters.append(Filter(class_type="DATE", process_func=process_date_1, filter_func=filter_date_1)) filters.append(Filter(class_type="PLAIN", process_func=process_plain_1, filter_func=filter_plain_1)) filters.append(Filter(class_type="PUNCT", process_func=process_punct_1, filter_func=filter_punct_1)) -filters.append(Filter(class_type="LETTERS", process_func=process_letters_1, filter_func=filter_letters_1)) -filters.append(Filter(class_type="VERBATIM", process_func=process_verbatim_1, filter_func=filter_verbatim_1)) +filters.append( + Filter(class_type="LETTERS", process_func=process_letters_1, filter_func=filter_letters_1) +) +filters.append( + Filter(class_type="VERBATIM", process_func=process_verbatim_1, filter_func=filter_verbatim_1) +) filters.append(Filter(class_type="DIGIT", process_func=process_digit_1, filter_func=filter_digit_1)) -filters.append(Filter(class_type="TELEPHONE", process_func=process_telephone_1, filter_func=filter_telephone_1)) -filters.append(Filter(class_type="ELECTRONIC", process_func=process_electronic_1, filter_func=filter_electronic_1)) -filters.append(Filter(class_type="FRACTION", process_func=process_fraction_1, filter_func=filter_fraction_1)) -filters.append(Filter(class_type="ADDRESS", process_func=process_address_1, filter_func=filter_address_1)) +filters.append( + Filter(class_type="TELEPHONE", process_func=process_telephone_1, filter_func=filter_telephone_1) +) +filters.append( + Filter( + class_type="ELECTRONIC", process_func=process_electronic_1, filter_func=filter_electronic_1 + ) +) +filters.append( + Filter(class_type="FRACTION", process_func=process_fraction_1, filter_func=filter_fraction_1) +) +filters.append( + Filter(class_type="ADDRESS", process_func=process_address_1, filter_func=filter_address_1) +) filters.append(Filter(class_type=EOS_TYPE, process_func=lambda x: x, filter_func=lambda x: True)) @@ -315,8 +367,10 @@ def filter_loaded_data(data: List[Instance], verbose: bool = False) -> List[Inst def parse_args(): parser = ArgumentParser() - parser.add_argument("--input", help="input file path", type=str, default='./en_with_types/output-00001-of-00100') - parser.add_argument("--verbose", help="print filtered instances", action='store_true') + parser.add_argument( + "--input", help="input file path", type=str, default="./en_with_types/output-00001-of-00100" + ) + parser.add_argument("--verbose", help="print filtered instances", action="store_true") return parser.parse_args() diff --git a/fun_text_processing/text_normalization/en/data/address/__init__.py b/fun_text_processing/text_normalization/en/data/address/__init__.py index b28b04f64..8b1378917 100644 --- a/fun_text_processing/text_normalization/en/data/address/__init__.py +++ b/fun_text_processing/text_normalization/en/data/address/__init__.py @@ -1,3 +1 @@ - - diff --git a/fun_text_processing/text_normalization/en/data/electronic/__init__.py b/fun_text_processing/text_normalization/en/data/electronic/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/en/data/electronic/__init__.py +++ b/fun_text_processing/text_normalization/en/data/electronic/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/en/data/telephone/__init__.py b/fun_text_processing/text_normalization/en/data/telephone/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/en/data/telephone/__init__.py +++ b/fun_text_processing/text_normalization/en/data/telephone/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/en/graph_utils.py b/fun_text_processing/text_normalization/en/graph_utils.py index 2ee3536ab..2cab9ee90 100644 --- a/fun_text_processing/text_normalization/en/graph_utils.py +++ b/fun_text_processing/text_normalization/en/graph_utils.py @@ -1,4 +1,3 @@ - import os import string from pathlib import Path @@ -19,9 +18,9 @@ DAMO_UPPER = pynini.union(*string.ascii_uppercase).optimize() DAMO_ALPHA = pynini.union(DAMO_LOWER, DAMO_UPPER).optimize() DAMO_ALNUM = pynini.union(DAMO_DIGIT, DAMO_ALPHA).optimize() DAMO_HEX = pynini.union(*string.hexdigits).optimize() -DAMO_NON_BREAKING_SPACE = u"\u00A0" +DAMO_NON_BREAKING_SPACE = "\u00A0" DAMO_SPACE = " " -DAMO_WHITE_SPACE = pynini.union(" ", "\t", "\n", "\r", u"\u00A0").optimize() +DAMO_WHITE_SPACE = pynini.union(" ", "\t", "\n", "\r", "\u00A0").optimize() DAMO_NOT_SPACE = pynini.difference(DAMO_CHAR, DAMO_WHITE_SPACE).optimize() DAMO_NOT_QUOTE = pynini.difference(DAMO_CHAR, r'"').optimize() @@ -36,31 +35,55 @@ insert_space = pynutil.insert(" ") delete_extra_space = pynini.cross(pynini.closure(DAMO_WHITE_SPACE, 1), " ") delete_preserve_order = pynini.closure( pynutil.delete(" preserve_order: true") - | (pynutil.delete(" field_order: \"") + DAMO_NOT_QUOTE + pynutil.delete("\"")) + | (pynutil.delete(' field_order: "') + DAMO_NOT_QUOTE + pynutil.delete('"')) ) suppletive = pynini.string_file(get_abs_path("data/suppletive.tsv")) # _v = pynini.union("a", "e", "i", "o", "u") _c = pynini.union( - "b", "c", "d", "f", "g", "h", "j", "k", "l", "m", "n", "p", "q", "r", "s", "t", "v", "w", "x", "y", "z" + "b", + "c", + "d", + "f", + "g", + "h", + "j", + "k", + "l", + "m", + "n", + "p", + "q", + "r", + "s", + "t", + "v", + "w", + "x", + "y", + "z", ) _ies = DAMO_SIGMA + _c + pynini.cross("y", "ies") _es = DAMO_SIGMA + pynini.union("s", "sh", "ch", "x", "z") + pynutil.insert("es") _s = DAMO_SIGMA + pynutil.insert("s") graph_plural = plurals._priority_union( - suppletive, plurals._priority_union(_ies, plurals._priority_union(_es, _s, DAMO_SIGMA), DAMO_SIGMA), DAMO_SIGMA + suppletive, + plurals._priority_union(_ies, plurals._priority_union(_es, _s, DAMO_SIGMA), DAMO_SIGMA), + DAMO_SIGMA, ).optimize() SINGULAR_TO_PLURAL = graph_plural PLURAL_TO_SINGULAR = pynini.invert(graph_plural) -TO_LOWER = pynini.union(*[pynini.cross(x, y) for x, y in zip(string.ascii_uppercase, string.ascii_lowercase)]) +TO_LOWER = pynini.union( + *[pynini.cross(x, y) for x, y in zip(string.ascii_uppercase, string.ascii_lowercase)] +) TO_UPPER = pynini.invert(TO_LOWER) MIN_NEG_WEIGHT = -0.0001 MIN_POS_WEIGHT = 0.0001 -def generator_main(file_name: str, graphs: Dict[str, 'pynini.FstLike']): +def generator_main(file_name: str, graphs: Dict[str, "pynini.FstLike"]): """ Exports graph as OpenFst finite state archive (FAR) file with given file name and rule name. @@ -72,7 +95,7 @@ def generator_main(file_name: str, graphs: Dict[str, 'pynini.FstLike']): for rule, graph in graphs.items(): exporter[rule] = graph.optimize() exporter.close() - print(f'Created {file_name}') + print(f"Created {file_name}") def get_plurals(fst): @@ -99,7 +122,7 @@ def get_singulars(fst): return PLURAL_TO_SINGULAR @ fst -def convert_space(fst) -> 'pynini.FstLike': +def convert_space(fst) -> "pynini.FstLike": """ Converts space to nonbreaking space. Used only in tagger grammars for transducing token values within quotes, e.g. name: "hello kitty" @@ -110,7 +133,9 @@ def convert_space(fst) -> 'pynini.FstLike': Returns output fst where breaking spaces are converted to non breaking spaces """ - return fst @ pynini.cdrewrite(pynini.cross(DAMO_SPACE, DAMO_NON_BREAKING_SPACE), "", "", DAMO_SIGMA) + return fst @ pynini.cdrewrite( + pynini.cross(DAMO_SPACE, DAMO_NON_BREAKING_SPACE), "", "", DAMO_SIGMA + ) class GraphFst: @@ -130,9 +155,11 @@ class GraphFst: self._fst = None self.deterministic = deterministic - self.far_path = Path(os.path.dirname(__file__) + '/grammars/' + kind + '/' + name + '.far') + self.far_path = Path(os.path.dirname(__file__) + "/grammars/" + kind + "/" + name + ".far") if self.far_exist(): - self._fst = Far(self.far_path, mode="r", arc_type="standard", far_type="default").get_fst() + self._fst = Far( + self.far_path, mode="r", arc_type="standard", far_type="default" + ).get_fst() def far_exist(self) -> bool: """ @@ -141,14 +168,14 @@ class GraphFst: return self.far_path.exists() @property - def fst(self) -> 'pynini.FstLike': + def fst(self) -> "pynini.FstLike": return self._fst @fst.setter def fst(self, fst): self._fst = fst - def add_tokens(self, fst) -> 'pynini.FstLike': + def add_tokens(self, fst) -> "pynini.FstLike": """ Wraps class name around to given fst @@ -160,7 +187,7 @@ class GraphFst: """ return pynutil.insert(f"{self.name} {{ ") + fst + pynutil.insert(" }") - def delete_tokens(self, fst) -> 'pynini.FstLike': + def delete_tokens(self, fst) -> "pynini.FstLike": """ Deletes class name wrap around output of given fst @@ -179,4 +206,4 @@ class GraphFst: + delete_space + pynutil.delete("}") ) - return res @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + return res @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) diff --git a/fun_text_processing/text_normalization/en/taggers/abbreviation.py b/fun_text_processing/text_normalization/en/taggers/abbreviation.py index a0dee2445..5748aa6fb 100644 --- a/fun_text_processing/text_normalization/en/taggers/abbreviation.py +++ b/fun_text_processing/text_normalization/en/taggers/abbreviation.py @@ -1,6 +1,3 @@ - - - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_UPPER, GraphFst, insert_space from pynini.lib import pynutil @@ -17,7 +14,7 @@ class AbbreviationFst(GraphFst): for False multiple transduction are generated (used for audio-based normalization) """ - def __init__(self, whitelist: 'pynini.FstLike', deterministic: bool = True): + def __init__(self, whitelist: "pynini.FstLike", deterministic: bool = True): super().__init__(name="abbreviation", kind="classify", deterministic=deterministic) dot = pynini.accep(".") @@ -30,9 +27,12 @@ class AbbreviationFst(GraphFst): # exclude words that are included in the whitelist graph = pynini.compose( - pynini.difference(pynini.project(graph, "input"), pynini.project(whitelist.graph, "input")), graph + pynini.difference( + pynini.project(graph, "input"), pynini.project(whitelist.graph, "input") + ), + graph, ) - graph = pynutil.insert("value: \"") + graph.optimize() + pynutil.insert("\"") + graph = pynutil.insert('value: "') + graph.optimize() + pynutil.insert('"') graph = self.add_tokens(graph) self.fst = graph.optimize() diff --git a/fun_text_processing/text_normalization/en/taggers/cardinal.py b/fun_text_processing/text_normalization/en/taggers/cardinal.py index c406fe49a..0a3e1a40d 100644 --- a/fun_text_processing/text_normalization/en/taggers/cardinal.py +++ b/fun_text_processing/text_normalization/en/taggers/cardinal.py @@ -1,6 +1,3 @@ - - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_DIGIT, @@ -17,7 +14,7 @@ from pynini.lib import pynutil class CardinalFst(GraphFst): """ - Finite state transducer for classifying cardinals, e.g. + Finite state transducer for classifying cardinals, e.g. -23 -> cardinal { negative: "true" integer: "twenty three" } } Args: @@ -40,7 +37,9 @@ class CardinalFst(GraphFst): graph_zero = pynini.string_file(get_abs_path("data/number/zero.tsv")) single_digits_graph = pynini.invert(graph_digit | graph_zero) - self.single_digits_graph = single_digits_graph + pynini.closure(insert_space + single_digits_graph) + self.single_digits_graph = single_digits_graph + pynini.closure( + insert_space + single_digits_graph + ) if not deterministic: # for a single token allow only the same normalization @@ -51,7 +50,9 @@ class CardinalFst(GraphFst): self.single_digits_graph = single_digits_graph_zero + pynini.closure( insert_space + single_digits_graph_zero ) - self.single_digits_graph |= single_digits_graph_oh + pynini.closure(insert_space + single_digits_graph_oh) + self.single_digits_graph |= single_digits_graph_oh + pynini.closure( + insert_space + single_digits_graph_oh + ) single_digits_graph_with_commas = pynini.closure( self.single_digits_graph + insert_space, 1, 3 @@ -65,27 +66,37 @@ class CardinalFst(GraphFst): 1, ) - optional_minus_graph = pynini.closure(pynutil.insert("negative: ") + pynini.cross("-", "\"true\" "), 0, 1) + optional_minus_graph = pynini.closure( + pynutil.insert("negative: ") + pynini.cross("-", '"true" '), 0, 1 + ) graph = ( pynini.closure(DAMO_DIGIT, 1, 3) - + (pynini.closure(pynutil.delete(",") + DAMO_DIGIT ** 3) | pynini.closure(DAMO_DIGIT ** 3)) + + (pynini.closure(pynutil.delete(",") + DAMO_DIGIT**3) | pynini.closure(DAMO_DIGIT**3)) ) @ graph self.graph = graph self.graph_with_and = self.add_optional_and(graph) if deterministic: - long_numbers = pynini.compose(DAMO_DIGIT ** (5, ...), self.single_digits_graph).optimize() - final_graph = plurals._priority_union(long_numbers, self.graph_with_and, DAMO_SIGMA).optimize() + long_numbers = pynini.compose( + DAMO_DIGIT ** (5, ...), self.single_digits_graph + ).optimize() + final_graph = plurals._priority_union( + long_numbers, self.graph_with_and, DAMO_SIGMA + ).optimize() cardinal_with_leading_zeros = pynini.compose( pynini.accep("0") + pynini.closure(DAMO_DIGIT), self.single_digits_graph ) final_graph |= cardinal_with_leading_zeros else: - leading_zeros = pynini.compose(pynini.closure(pynini.accep("0"), 1), self.single_digits_graph) + leading_zeros = pynini.compose( + pynini.closure(pynini.accep("0"), 1), self.single_digits_graph + ) cardinal_with_leading_zeros = ( - leading_zeros + pynutil.insert(" ") + pynini.compose(pynini.closure(DAMO_DIGIT), self.graph_with_and) + leading_zeros + + pynutil.insert(" ") + + pynini.compose(pynini.closure(DAMO_DIGIT), self.graph_with_and) ) # add small weight to non-default graphs to make sure the deterministic option is listed first @@ -97,7 +108,9 @@ class CardinalFst(GraphFst): | cardinal_with_leading_zeros ) - final_graph = optional_minus_graph + pynutil.insert("integer: \"") + final_graph + pynutil.insert("\"") + final_graph = ( + optional_minus_graph + pynutil.insert('integer: "') + final_graph + pynutil.insert('"') + ) final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() @@ -111,16 +124,26 @@ class CardinalFst(GraphFst): not_quote, not_quote + pynini.union("thousand", "million") + not_quote ).optimize() integer = ( - not_quote + pynutil.add_weight(pynini.cross("hundred ", "hundred and ") + no_thousand_million, -0.0001) + not_quote + + pynutil.add_weight( + pynini.cross("hundred ", "hundred and ") + no_thousand_million, -0.0001 + ) ).optimize() - no_hundred = pynini.difference(DAMO_SIGMA, not_quote + pynini.accep("hundred") + not_quote).optimize() + no_hundred = pynini.difference( + DAMO_SIGMA, not_quote + pynini.accep("hundred") + not_quote + ).optimize() integer |= ( - not_quote + pynutil.add_weight(pynini.cross("thousand ", "thousand and ") + no_hundred, -0.0001) + not_quote + + pynutil.add_weight( + pynini.cross("thousand ", "thousand and ") + no_hundred, -0.0001 + ) ).optimize() optional_hundred = pynini.compose((DAMO_DIGIT - "0") ** 3, graph).optimize() - optional_hundred = pynini.compose(optional_hundred, DAMO_SIGMA + pynini.cross(" hundred", "") + DAMO_SIGMA) + optional_hundred = pynini.compose( + optional_hundred, DAMO_SIGMA + pynini.cross(" hundred", "") + DAMO_SIGMA + ) graph_with_and |= pynini.compose(graph, integer).optimize() graph_with_and |= optional_hundred return graph_with_and diff --git a/fun_text_processing/text_normalization/en/taggers/date.py b/fun_text_processing/text_normalization/en/taggers/date.py index 897d05da0..b5acfabc8 100644 --- a/fun_text_processing/text_normalization/en/taggers/date.py +++ b/fun_text_processing/text_normalization/en/taggers/date.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_CHAR, @@ -30,7 +28,7 @@ year_suffix = pynini.string_map(year_suffix).optimize() def get_ties_graph(deterministic: bool = True): """ - Returns two digit transducer, e.g. + Returns two digit transducer, e.g. 03 -> o three 12 -> thirteen 20 -> twenty @@ -40,7 +38,10 @@ def get_ties_graph(deterministic: bool = True): if deterministic: graph = graph | pynini.cross("0", "oh") + insert_space + graph_digit else: - graph = graph | (pynini.cross("0", "oh") | pynini.cross("0", "zero")) + insert_space + graph_digit + graph = ( + graph + | (pynini.cross("0", "oh") | pynini.cross("0", "zero")) + insert_space + graph_digit + ) return graph.optimize() @@ -59,7 +60,12 @@ def get_four_digit_year_graph(deterministic: bool = True): | (graph_teen + insert_space + (ties_graph | pynini.cross("1", "ten"))) ) + pynutil.delete("0s") - graph_with_s |= (graph_teen | graph_ties) + insert_space + pynini.cross("00", "hundred") + pynutil.delete("s") + graph_with_s |= ( + (graph_teen | graph_ties) + + insert_space + + pynini.cross("00", "hundred") + + pynutil.delete("s") + ) graph_with_s = graph_with_s @ pynini.cdrewrite( pynini.cross("y", "ies") | pynutil.insert("s"), "", "[EOS]", DAMO_SIGMA ) @@ -95,7 +101,8 @@ def _get_two_digit_year_with_s_graph(): graph = ( pynini.closure(pynutil.delete("'"), 0, 1) + pynini.compose( - ties_graph + pynutil.delete("0s"), pynini.cdrewrite(pynini.cross("y", "ies"), "", "[EOS]", DAMO_SIGMA) + ties_graph + pynutil.delete("0s"), + pynini.cdrewrite(pynini.cross("y", "ies"), "", "[EOS]", DAMO_SIGMA), ) ).optimize() return graph @@ -106,33 +113,44 @@ def _get_year_graph(cardinal_graph, deterministic: bool = True): Transducer for year, only from 1000 - 2999 e.g. 1290 -> twelve nineteen 2000 - 2009 will be verbalized as two thousand. - + Transducer for 3 digit year, e.g. 123-> one twenty three - + Transducer for year with suffix 123 A.D., 4200 B.C """ graph = get_four_digit_year_graph(deterministic) - graph = (pynini.union("1", "2") + (DAMO_DIGIT ** 3) + pynini.closure(pynini.cross(" s", "s") | "s", 0, 1)) @ graph + graph = ( + pynini.union("1", "2") + + (DAMO_DIGIT**3) + + pynini.closure(pynini.cross(" s", "s") | "s", 0, 1) + ) @ graph graph |= _get_two_digit_year_with_s_graph() - three_digit_year = (DAMO_DIGIT @ cardinal_graph) + insert_space + (DAMO_DIGIT ** 2) @ cardinal_graph + three_digit_year = ( + (DAMO_DIGIT @ cardinal_graph) + insert_space + (DAMO_DIGIT**2) @ cardinal_graph + ) year_with_suffix = ( - (get_four_digit_year_graph(deterministic=True) | three_digit_year) + delete_space + insert_space + year_suffix + (get_four_digit_year_graph(deterministic=True) | three_digit_year) + + delete_space + + insert_space + + year_suffix ) graph |= year_with_suffix return graph.optimize() def _get_two_digit_year(cardinal_graph, single_digits_graph): - wo_digit_year = DAMO_DIGIT ** (2) @ plurals._priority_union(cardinal_graph, single_digits_graph, DAMO_SIGMA) + wo_digit_year = DAMO_DIGIT ** (2) @ plurals._priority_union( + cardinal_graph, single_digits_graph, DAMO_SIGMA + ) return wo_digit_year class DateFst(GraphFst): """ - Finite state transducer for classifying date, e.g. + Finite state transducer for classifying date, e.g. jan. 5, 2012 -> date { month: "january" day: "five" year: "twenty twelve" preserve_order: true } jan. 5 -> date { month: "january" day: "five" preserve_order: true } 5 january 2012 -> date { day: "five" month: "january" year: "twenty twelve" preserve_order: true } @@ -153,9 +171,9 @@ class DateFst(GraphFst): # january month_graph = pynini.string_file(get_abs_path("data/date/month_name.tsv")).optimize() # January, JANUARY - month_graph |= pynini.compose(TO_LOWER + pynini.closure(DAMO_CHAR), month_graph) | pynini.compose( - TO_LOWER ** (2, ...), month_graph - ) + month_graph |= pynini.compose( + TO_LOWER + pynini.closure(DAMO_CHAR), month_graph + ) | pynini.compose(TO_LOWER ** (2, ...), month_graph) # jan month_abbr_graph = pynini.string_file(get_abs_path("data/date/month_abbr.tsv")).optimize() @@ -167,7 +185,9 @@ class DateFst(GraphFst): ) + pynini.closure(pynutil.delete("."), 0, 1) month_graph |= month_abbr_graph.optimize() - month_numbers_labels = pynini.string_file(get_abs_path("data/date/month_number.tsv")).optimize() + month_numbers_labels = pynini.string_file( + get_abs_path("data/date/month_number.tsv") + ).optimize() cardinal_graph = cardinal.graph_hundred_component_at_least_one_none_zero_digit year_graph = _get_year_graph(cardinal_graph=cardinal_graph, deterministic=deterministic) @@ -175,45 +195,53 @@ class DateFst(GraphFst): # three_digit_year = (DAMO_DIGIT @ cardinal_graph) + insert_space + (DAMO_DIGIT ** 2) @ cardinal_graph # year_graph |= three_digit_year - month_graph = pynutil.insert("month: \"") + month_graph + pynutil.insert("\"") - month_numbers_graph = pynutil.insert("month: \"") + month_numbers_labels + pynutil.insert("\"") + month_graph = pynutil.insert('month: "') + month_graph + pynutil.insert('"') + month_numbers_graph = ( + pynutil.insert('month: "') + month_numbers_labels + pynutil.insert('"') + ) endings = ["rd", "th", "st", "nd"] endings += [x.upper() for x in endings] endings = pynini.union(*endings) day_graph = ( - pynutil.insert("day: \"") + pynutil.insert('day: "') + pynini.closure(pynutil.delete("the "), 0, 1) + ( - ((pynini.union("1", "2") + DAMO_DIGIT) | DAMO_DIGIT | (pynini.accep("3") + pynini.union("0", "1"))) + ( + (pynini.union("1", "2") + DAMO_DIGIT) + | DAMO_DIGIT + | (pynini.accep("3") + pynini.union("0", "1")) + ) + pynini.closure(pynutil.delete(endings), 0, 1) ) @ cardinal_graph - + pynutil.insert("\"") + + pynutil.insert('"') ) two_digit_year = _get_two_digit_year( cardinal_graph=cardinal_graph, single_digits_graph=cardinal.single_digits_graph ) - two_digit_year = pynutil.insert("year: \"") + two_digit_year + pynutil.insert("\"") + two_digit_year = pynutil.insert('year: "') + two_digit_year + pynutil.insert('"') # if lm: # two_digit_year = pynini.compose(pynini.difference(DAMO_DIGIT, "0") + DAMO_DIGIT ** (3), two_digit_year) # year_graph = pynini.compose(pynini.difference(DAMO_DIGIT, "0") + DAMO_DIGIT ** (2), year_graph) # year_graph |= pynini.compose(pynini.difference(DAMO_DIGIT, "0") + DAMO_DIGIT ** (4, ...), year_graph) - graph_year = pynutil.insert(" year: \"") + pynutil.delete(" ") + year_graph + pynutil.insert("\"") + graph_year = ( + pynutil.insert(' year: "') + pynutil.delete(" ") + year_graph + pynutil.insert('"') + ) graph_year |= ( - pynutil.insert(" year: \"") + pynutil.insert(' year: "') + pynini.accep(",") + pynini.closure(pynini.accep(" "), 0, 1) + year_graph - + pynutil.insert("\"") + + pynutil.insert('"') ) optional_graph_year = pynini.closure(graph_year, 0, 1) - year_graph = pynutil.insert("year: \"") + year_graph + pynutil.insert("\"") + year_graph = pynutil.insert('year: "') + year_graph + pynutil.insert('"') graph_mdy = month_graph + ( (delete_extra_space + day_graph) @@ -243,7 +271,7 @@ class DateFst(GraphFst): ) graph_dmy = day_graph + delete_extra_space + month_graph + optional_graph_year - day_ex_month = (DAMO_DIGIT ** 2 - pynini.project(month_numbers_graph, "input")) @ day_graph + day_ex_month = (DAMO_DIGIT**2 - pynini.project(month_numbers_graph, "input")) @ day_graph for x in ["-", "/", "."]: delete_sep = pynutil.delete(x) graph_dmy |= ( @@ -296,10 +324,10 @@ class DateFst(GraphFst): for month in [x[0] for x in load_labels(get_abs_path("data/date/month_name.tsv"))]: for day in [x[0] for x in load_labels(get_abs_path("data/date/day.tsv"))]: ymd_to_mdy_curr = ( - pynutil.insert("month: \"" + month + "\" day: \"" + day + "\" ") - + pynini.accep('year:') + pynutil.insert('month: "' + month + '" day: "' + day + '" ') + + pynini.accep("year:") + DAMO_SIGMA - + pynutil.delete(" month: \"" + month + "\" day: \"" + day + "\"") + + pynutil.delete(' month: "' + month + '" day: "' + day + '"') ) # YY-MM-DD -> MM-DD-YY @@ -311,10 +339,10 @@ class DateFst(GraphFst): ) ymd_to_dmy_curr = ( - pynutil.insert("day: \"" + day + "\" month: \"" + month + "\" ") - + pynini.accep('year:') + pynutil.insert('day: "' + day + '" month: "' + month + '" ') + + pynini.accep("year:") + DAMO_SIGMA - + pynutil.delete(" month: \"" + month + "\" day: \"" + day + "\"") + + pynutil.delete(' month: "' + month + '" day: "' + day + '"') ) # YY-MM-DD -> MM-DD-YY @@ -326,9 +354,9 @@ class DateFst(GraphFst): ) mdy_to_dmy_curr = ( - pynutil.insert("day: \"" + day + "\" month: \"" + month + "\" ") - + pynutil.delete("month: \"" + month + "\" day: \"" + day + "\" ") - + pynini.accep('year:') + pynutil.insert('day: "' + day + '" month: "' + month + '" ') + + pynutil.delete('month: "' + month + '" day: "' + day + '" ') + + pynini.accep("year:") + DAMO_SIGMA ).optimize() # MM-DD-YY -> verbalize as MM-DD-YY (February fourth 1991) or DD-MM-YY (the fourth of February 1991) @@ -339,9 +367,9 @@ class DateFst(GraphFst): else pynini.union(mdy_to_dmy_curr, mdy_to_dmy_graph).optimize() ).optimize() - md_to_dm_curr = pynutil.insert("day: \"" + day + "\" month: \"" + month + "\"") + pynutil.delete( - "month: \"" + month + "\" day: \"" + day + "\"" - ) + md_to_dm_curr = pynutil.insert( + 'day: "' + day + '" month: "' + month + '"' + ) + pynutil.delete('month: "' + month + '" day: "' + day + '"') md_to_dm_curr = pynini.compose(m_sep_d, md_to_dm_curr).optimize() md_to_dm_graph = ( diff --git a/fun_text_processing/text_normalization/en/taggers/decimal.py b/fun_text_processing/text_normalization/en/taggers/decimal.py index 19f0ccb18..839e9f5a7 100644 --- a/fun_text_processing/text_normalization/en/taggers/decimal.py +++ b/fun_text_processing/text_normalization/en/taggers/decimal.py @@ -1,7 +1,10 @@ - - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_SIGMA, TO_UPPER, GraphFst, get_abs_path +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_SIGMA, + TO_UPPER, + GraphFst, + get_abs_path, +) from pynini.lib import pynutil delete_space = pynutil.delete(" ") @@ -11,28 +14,30 @@ quantities_abbr |= TO_UPPER @ quantities_abbr def get_quantity( - decimal: 'pynini.FstLike', cardinal_up_to_hundred: 'pynini.FstLike', include_abbr: bool -) -> 'pynini.FstLike': + decimal: "pynini.FstLike", cardinal_up_to_hundred: "pynini.FstLike", include_abbr: bool +) -> "pynini.FstLike": """ Returns FST that transforms either a cardinal or decimal followed by a quantity into a numeral, e.g. 1 million -> integer_part: "one" quantity: "million" e.g. 1.5 million -> integer_part: "one" fractional_part: "five" quantity: "million" - Args: + Args: decimal: decimal FST cardinal_up_to_hundred: cardinal FST """ quantity_wo_thousand = pynini.project(quantities, "input") - pynini.union("k", "K", "thousand") if include_abbr: - quantity_wo_thousand |= pynini.project(quantities_abbr, "input") - pynini.union("k", "K", "thousand") + quantity_wo_thousand |= pynini.project(quantities_abbr, "input") - pynini.union( + "k", "K", "thousand" + ) res = ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + cardinal_up_to_hundred - + pynutil.insert("\"") + + pynutil.insert('"') + pynini.closure(pynutil.delete(" "), 0, 1) - + pynutil.insert(" quantity: \"") + + pynutil.insert(' quantity: "') + (quantity_wo_thousand @ (quantities | quantities_abbr)) - + pynutil.insert("\"") + + pynutil.insert('"') ) if include_abbr: quantity = quantities | quantities_abbr @@ -41,16 +46,16 @@ def get_quantity( res |= ( decimal + pynini.closure(pynutil.delete(" "), 0, 1) - + pynutil.insert("quantity: \"") + + pynutil.insert('quantity: "') + quantity - + pynutil.insert("\"") + + pynutil.insert('"') ) return res class DecimalFst(GraphFst): """ - Finite state transducer for classifying decimal, e.g. + Finite state transducer for classifying decimal, e.g. -12.5006 billion -> decimal { negative: "true" integer_part: "12" fractional_part: "five o o six" quantity: "billion" } 1 billion -> decimal { integer_part: "one" quantity: "billion" } @@ -71,10 +76,16 @@ class DecimalFst(GraphFst): self.graph = self.graph | cardinal_graph point = pynutil.delete(".") - optional_graph_negative = pynini.closure(pynutil.insert("negative: ") + pynini.cross("-", "\"true\" "), 0, 1) + optional_graph_negative = pynini.closure( + pynutil.insert("negative: ") + pynini.cross("-", '"true" '), 0, 1 + ) - self.graph_fractional = pynutil.insert("fractional_part: \"") + self.graph + pynutil.insert("\"") - self.graph_integer = pynutil.insert("integer_part: \"") + cardinal_graph + pynutil.insert("\"") + self.graph_fractional = ( + pynutil.insert('fractional_part: "') + self.graph + pynutil.insert('"') + ) + self.graph_integer = ( + pynutil.insert('integer_part: "') + cardinal_graph + pynutil.insert('"') + ) final_graph_wo_sign = ( pynini.closure(self.graph_integer + pynutil.insert(" "), 0, 1) + point @@ -83,10 +94,14 @@ class DecimalFst(GraphFst): ) quantity_w_abbr = get_quantity( - final_graph_wo_sign, cardinal_graph_hundred_component_at_least_one_none_zero_digit, include_abbr=True + final_graph_wo_sign, + cardinal_graph_hundred_component_at_least_one_none_zero_digit, + include_abbr=True, ) quantity_wo_abbr = get_quantity( - final_graph_wo_sign, cardinal_graph_hundred_component_at_least_one_none_zero_digit, include_abbr=False + final_graph_wo_sign, + cardinal_graph_hundred_component_at_least_one_none_zero_digit, + include_abbr=False, ) self.final_graph_wo_negative_w_abbr = final_graph_wo_sign | quantity_w_abbr self.final_graph_wo_negative = final_graph_wo_sign | quantity_wo_abbr @@ -99,17 +114,25 @@ class DecimalFst(GraphFst): | (DAMO_SIGMA + "zero" + DAMO_SIGMA + "oh" + DAMO_SIGMA), ).optimize() no_zero_oh = pynini.difference( - DAMO_SIGMA, DAMO_SIGMA + pynini.accep("zero") + DAMO_SIGMA + pynini.accep("oh") + DAMO_SIGMA + DAMO_SIGMA, + DAMO_SIGMA + pynini.accep("zero") + DAMO_SIGMA + pynini.accep("oh") + DAMO_SIGMA, ).optimize() self.final_graph_wo_negative |= pynini.compose( self.final_graph_wo_negative, pynini.cdrewrite( - pynini.cross("integer_part: \"zero\"", "integer_part: \"oh\""), DAMO_SIGMA, DAMO_SIGMA, DAMO_SIGMA + pynini.cross('integer_part: "zero"', 'integer_part: "oh"'), + DAMO_SIGMA, + DAMO_SIGMA, + DAMO_SIGMA, ), ) - self.final_graph_wo_negative = pynini.compose(self.final_graph_wo_negative, no_oh_zero).optimize() - self.final_graph_wo_negative = pynini.compose(self.final_graph_wo_negative, no_zero_oh).optimize() + self.final_graph_wo_negative = pynini.compose( + self.final_graph_wo_negative, no_oh_zero + ).optimize() + self.final_graph_wo_negative = pynini.compose( + self.final_graph_wo_negative, no_zero_oh + ).optimize() final_graph = optional_graph_negative + self.final_graph_wo_negative diff --git a/fun_text_processing/text_normalization/en/taggers/electronic.py b/fun_text_processing/text_normalization/en/taggers/electronic.py index 269cdb12c..26ac25f81 100644 --- a/fun_text_processing/text_normalization/en/taggers/electronic.py +++ b/fun_text_processing/text_normalization/en/taggers/electronic.py @@ -1,6 +1,3 @@ - - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_ALPHA, @@ -26,43 +23,60 @@ class ElectronicFst(GraphFst): def __init__(self, deterministic: bool = True): super().__init__(name="electronic", kind="classify", deterministic=deterministic) - accepted_symbols = pynini.project(pynini.string_file(get_abs_path("data/electronic/symbol.tsv")), "input") + accepted_symbols = pynini.project( + pynini.string_file(get_abs_path("data/electronic/symbol.tsv")), "input" + ) accepted_common_domains = pynini.project( pynini.string_file(get_abs_path("data/electronic/domain.tsv")), "input" ) - all_accepted_symbols = DAMO_ALPHA + pynini.closure(DAMO_ALPHA | DAMO_DIGIT | accepted_symbols) + all_accepted_symbols = DAMO_ALPHA + pynini.closure( + DAMO_ALPHA | DAMO_DIGIT | accepted_symbols + ) graph_symbols = pynini.string_file(get_abs_path("data/electronic/symbol.tsv")).optimize() username = ( - pynutil.insert("username: \"") + all_accepted_symbols + pynutil.insert("\"") + pynini.cross('@', ' ') + pynutil.insert('username: "') + + all_accepted_symbols + + pynutil.insert('"') + + pynini.cross("@", " ") + ) + domain_graph = all_accepted_symbols + pynini.accep(".") + all_accepted_symbols + DAMO_ALPHA + protocol_symbols = pynini.closure( + (graph_symbols | pynini.cross(":", "semicolon")) + pynutil.insert(" ") ) - domain_graph = all_accepted_symbols + pynini.accep('.') + all_accepted_symbols + DAMO_ALPHA - protocol_symbols = pynini.closure((graph_symbols | pynini.cross(":", "semicolon")) + pynutil.insert(" ")) protocol_start = (pynini.cross("https", "HTTPS ") | pynini.cross("http", "HTTP ")) + ( pynini.accep("://") @ protocol_symbols ) - protocol_file_start = pynini.accep("file") + insert_space + (pynini.accep(":///") @ protocol_symbols) + protocol_file_start = ( + pynini.accep("file") + insert_space + (pynini.accep(":///") @ protocol_symbols) + ) protocol_end = pynini.cross("www", "WWW ") + pynini.accep(".") @ protocol_symbols - protocol = protocol_file_start | protocol_start | protocol_end | (protocol_start + protocol_end) + protocol = ( + protocol_file_start | protocol_start | protocol_end | (protocol_start + protocol_end) + ) domain_graph = ( - pynutil.insert("domain: \"") + pynutil.insert('domain: "') + pynini.difference(domain_graph, pynini.project(protocol, "input") + DAMO_SIGMA) - + pynutil.insert("\"") + + pynutil.insert('"') ) domain_common_graph = ( - pynutil.insert("domain: \"") + pynutil.insert('domain: "') + pynini.difference( all_accepted_symbols + accepted_common_domains - + pynini.closure(accepted_symbols + pynini.closure(DAMO_ALPHA | DAMO_DIGIT | accepted_symbols), 0, 1), + + pynini.closure( + accepted_symbols + pynini.closure(DAMO_ALPHA | DAMO_DIGIT | accepted_symbols), + 0, + 1, + ), pynini.project(protocol, "input") + DAMO_SIGMA, ) - + pynutil.insert("\"") + + pynutil.insert('"') ) - protocol = pynutil.insert("protocol: \"") + protocol + pynutil.insert("\"") + protocol = pynutil.insert('protocol: "') + protocol + pynutil.insert('"') # email graph = username + domain_graph # abc.com, abc.com/123-sm diff --git a/fun_text_processing/text_normalization/en/taggers/fraction.py b/fun_text_processing/text_normalization/en/taggers/fraction.py index d7ee09117..08155ae13 100644 --- a/fun_text_processing/text_normalization/en/taggers/fraction.py +++ b/fun_text_processing/text_normalization/en/taggers/fraction.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import GraphFst, get_abs_path from pynini.lib import pynutil @@ -22,19 +20,25 @@ class FractionFst(GraphFst): super().__init__(name="fraction", kind="classify", deterministic=deterministic) cardinal_graph = cardinal.graph - integer = pynutil.insert("integer_part: \"") + cardinal_graph + pynutil.insert("\"") + integer = pynutil.insert('integer_part: "') + cardinal_graph + pynutil.insert('"') numerator = ( - pynutil.insert("numerator: \"") + cardinal_graph + (pynini.cross("/", "\" ") | pynini.cross(" / ", "\" ")) + pynutil.insert('numerator: "') + + cardinal_graph + + (pynini.cross("/", '" ') | pynini.cross(" / ", '" ')) ) endings = ["rd", "th", "st", "nd"] endings += [x.upper() for x in endings] optional_end = pynini.closure(pynini.cross(pynini.union(*endings), ""), 0, 1) - denominator = pynutil.insert("denominator: \"") + cardinal_graph + optional_end + pynutil.insert("\"") + denominator = ( + pynutil.insert('denominator: "') + cardinal_graph + optional_end + pynutil.insert('"') + ) graph = pynini.closure(integer + pynini.accep(" "), 0, 1) + (numerator + denominator) - graph |= pynini.closure(integer + (pynini.accep(" ") | pynutil.insert(" ")), 0, 1) + pynini.compose( + graph |= pynini.closure( + integer + (pynini.accep(" ") | pynutil.insert(" ")), 0, 1 + ) + pynini.compose( pynini.string_file(get_abs_path("data/number/fraction.tsv")), (numerator + denominator) ) diff --git a/fun_text_processing/text_normalization/en/taggers/measure.py b/fun_text_processing/text_normalization/en/taggers/measure.py index 30f6cb928..703ce8ed1 100644 --- a/fun_text_processing/text_normalization/en/taggers/measure.py +++ b/fun_text_processing/text_normalization/en/taggers/measure.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_ALPHA, @@ -19,14 +17,16 @@ from fun_text_processing.text_normalization.en.graph_utils import ( from fun_text_processing.text_normalization.en.taggers.ordinal import OrdinalFst as OrdinalTagger from fun_text_processing.text_normalization.en.taggers.whitelist import get_formats from fun_text_processing.text_normalization.en.utils import get_abs_path, load_labels -from fun_text_processing.text_normalization.en.verbalizers.ordinal import OrdinalFst as OrdinalVerbalizer +from fun_text_processing.text_normalization.en.verbalizers.ordinal import ( + OrdinalFst as OrdinalVerbalizer, +) from pynini.examples import plurals from pynini.lib import pynutil class MeasureFst(GraphFst): """ - Finite state transducer for classifying measure, suppletive aware, e.g. + Finite state transducer for classifying measure, suppletive aware, e.g. -12kg -> measure { negative: "true" cardinal { integer: "twelve" } units: "kilograms" } 1kg -> measure { cardinal { integer: "one" } units: "kilogram" } .5kg -> measure { decimal { fractional_part: "five" } units: "kilograms" } @@ -39,7 +39,9 @@ class MeasureFst(GraphFst): for False multiple transduction are generated (used for audio-based normalization) """ - def __init__(self, cardinal: GraphFst, decimal: GraphFst, fraction: GraphFst, deterministic: bool = True): + def __init__( + self, cardinal: GraphFst, decimal: GraphFst, fraction: GraphFst, deterministic: bool = True + ): super().__init__(name="measure", kind="classify", deterministic=deterministic) cardinal_graph = cardinal.graph_with_and | self.get_range(cardinal.graph_with_and) @@ -48,30 +50,42 @@ class MeasureFst(GraphFst): graph_unit |= pynini.string_file(get_abs_path("data/measure/unit_alternatives.tsv")) graph_unit |= pynini.compose( - pynini.closure(TO_LOWER, 1) + (DAMO_ALPHA | TO_LOWER) + pynini.closure(DAMO_ALPHA | TO_LOWER), graph_unit + pynini.closure(TO_LOWER, 1) + + (DAMO_ALPHA | TO_LOWER) + + pynini.closure(DAMO_ALPHA | TO_LOWER), + graph_unit, ).optimize() graph_unit_plural = convert_space(graph_unit @ SINGULAR_TO_PLURAL) graph_unit = convert_space(graph_unit) - optional_graph_negative = pynini.closure(pynutil.insert("negative: ") + pynini.cross("-", "\"true\" "), 0, 1) + optional_graph_negative = pynini.closure( + pynutil.insert("negative: ") + pynini.cross("-", '"true" '), 0, 1 + ) graph_unit2 = ( - pynini.cross("/", "per") + delete_zero_or_one_space + pynutil.insert(DAMO_NON_BREAKING_SPACE) + graph_unit + pynini.cross("/", "per") + + delete_zero_or_one_space + + pynutil.insert(DAMO_NON_BREAKING_SPACE) + + graph_unit ) optional_graph_unit2 = pynini.closure( - delete_zero_or_one_space + pynutil.insert(DAMO_NON_BREAKING_SPACE) + graph_unit2, 0, 1, + delete_zero_or_one_space + pynutil.insert(DAMO_NON_BREAKING_SPACE) + graph_unit2, + 0, + 1, ) unit_plural = ( - pynutil.insert("units: \"") + pynutil.insert('units: "') + (graph_unit_plural + optional_graph_unit2 | graph_unit2) - + pynutil.insert("\"") + + pynutil.insert('"') ) unit_singular = ( - pynutil.insert("units: \"") + (graph_unit + optional_graph_unit2 | graph_unit2) + pynutil.insert("\"") + pynutil.insert('units: "') + + (graph_unit + optional_graph_unit2 | graph_unit2) + + pynutil.insert('"') ) subgraph_decimal = ( @@ -89,18 +103,18 @@ class MeasureFst(GraphFst): + decimal.final_graph_wo_negative + delete_space + pynutil.insert(" } ") - + pynutil.insert("units: \"") + + pynutil.insert('units: "') + pynini.union("AM", "FM") - + pynutil.insert("\"") + + pynutil.insert('"') ) subgraph_cardinal = ( pynutil.insert("cardinal { ") + optional_graph_negative - + pynutil.insert("integer: \"") + + pynutil.insert('integer: "') + ((DAMO_SIGMA - "1") @ cardinal_graph) + delete_space - + pynutil.insert("\"") + + pynutil.insert('"') + pynutil.insert(" } ") + unit_plural ) @@ -108,59 +122,63 @@ class MeasureFst(GraphFst): subgraph_cardinal |= ( pynutil.insert("cardinal { ") + optional_graph_negative - + pynutil.insert("integer: \"") + + pynutil.insert('integer: "') + pynini.cross("1", "one") + delete_space - + pynutil.insert("\"") + + pynutil.insert('"') + pynutil.insert(" } ") + unit_singular ) unit_graph = ( - pynutil.insert("cardinal { integer: \"-\" } units: \"") + pynutil.insert('cardinal { integer: "-" } units: "') + pynini.cross(pynini.union("/", "per"), "per") + delete_zero_or_one_space + pynutil.insert(DAMO_NON_BREAKING_SPACE) + graph_unit - + pynutil.insert("\" preserve_order: true") + + pynutil.insert('" preserve_order: true') ) decimal_dash_alpha = ( pynutil.insert("decimal { ") + decimal.final_graph_wo_negative - + pynini.cross('-', '') - + pynutil.insert(" } units: \"") + + pynini.cross("-", "") + + pynutil.insert(' } units: "') + pynini.closure(DAMO_ALPHA, 1) - + pynutil.insert("\"") + + pynutil.insert('"') ) decimal_times = ( pynutil.insert("decimal { ") + decimal.final_graph_wo_negative - + pynutil.insert(" } units: \"") - + pynini.cross(pynini.union('x', "X"), 'x') - + pynutil.insert("\"") + + pynutil.insert(' } units: "') + + pynini.cross(pynini.union("x", "X"), "x") + + pynutil.insert('"') ) alpha_dash_decimal = ( - pynutil.insert("units: \"") + pynutil.insert('units: "') + pynini.closure(DAMO_ALPHA, 1) - + pynini.accep('-') - + pynutil.insert("\"") + + pynini.accep("-") + + pynutil.insert('"') + pynutil.insert(" decimal { ") + decimal.final_graph_wo_negative + pynutil.insert(" } preserve_order: true") ) subgraph_fraction = ( - pynutil.insert("fraction { ") + fraction.graph + delete_space + pynutil.insert(" } ") + unit_plural + pynutil.insert("fraction { ") + + fraction.graph + + delete_space + + pynutil.insert(" } ") + + unit_plural ) address = self.get_address_graph(cardinal) address = ( - pynutil.insert("units: \"address\" cardinal { integer: \"") + pynutil.insert('units: "address" cardinal { integer: "') + address - + pynutil.insert("\" } preserve_order: true") + + pynutil.insert('" } preserve_order: true') ) math_operations = pynini.string_file(get_abs_path("data/measure/math_operation.tsv")) @@ -191,9 +209,9 @@ class MeasureFst(GraphFst): ) math = ( - pynutil.insert("units: \"math\" cardinal { integer: \"") + pynutil.insert('units: "math" cardinal { integer: "') + math - + pynutil.insert("\" } preserve_order: true") + + pynutil.insert('" } preserve_order: true') ) final_graph = ( subgraph_decimal @@ -240,11 +258,13 @@ class MeasureFst(GraphFst): ordinal_verbalizer = OrdinalVerbalizer().graph ordinal_tagger = OrdinalTagger(cardinal=cardinal).graph ordinal_num = pynini.compose( - pynutil.insert("integer: \"") + ordinal_tagger + pynutil.insert("\""), ordinal_verbalizer + pynutil.insert('integer: "') + ordinal_tagger + pynutil.insert('"'), ordinal_verbalizer ) - address_num = DAMO_DIGIT ** (1, 2) @ cardinal.graph_hundred_component_at_least_one_none_zero_digit - address_num += insert_space + DAMO_DIGIT ** 2 @ ( + address_num = ( + DAMO_DIGIT ** (1, 2) @ cardinal.graph_hundred_component_at_least_one_none_zero_digit + ) + address_num += insert_space + DAMO_DIGIT**2 @ ( pynini.closure(pynini.cross("0", "zero "), 0, 1) + cardinal.graph_hundred_component_at_least_one_none_zero_digit ) @@ -282,11 +302,19 @@ class MeasureFst(GraphFst): state = pynini.invert(state_graph) state = pynini.closure(pynini.accep(",") + pynini.accep(DAMO_SPACE) + state, 0, 1) - zip_code = pynini.compose(DAMO_DIGIT ** 5, cardinal.single_digits_graph) - zip_code = pynini.closure(pynini.closure(pynini.accep(","), 0, 1) + pynini.accep(DAMO_SPACE) + zip_code, 0, 1,) + zip_code = pynini.compose(DAMO_DIGIT**5, cardinal.single_digits_graph) + zip_code = pynini.closure( + pynini.closure(pynini.accep(","), 0, 1) + pynini.accep(DAMO_SPACE) + zip_code, + 0, + 1, + ) - address = address_num + direction + address_words + pynini.closure(city + state + zip_code, 0, 1) + address = ( + address_num + direction + address_words + pynini.closure(city + state + zip_code, 0, 1) + ) - address |= address_num + direction + address_words + pynini.closure(pynini.cross(".", ""), 0, 1) + address |= ( + address_num + direction + address_words + pynini.closure(pynini.cross(".", ""), 0, 1) + ) return address diff --git a/fun_text_processing/text_normalization/en/taggers/money.py b/fun_text_processing/text_normalization/en/taggers/money.py index e2525c1c6..f975bdb05 100644 --- a/fun_text_processing/text_normalization/en/taggers/money.py +++ b/fun_text_processing/text_normalization/en/taggers/money.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_ALPHA, @@ -20,7 +18,7 @@ maj_singular = pynini.string_file((get_abs_path("data/money/currency_major.tsv") class MoneyFst(GraphFst): """ - Finite state transducer for classifying money, suppletive aware, e.g. + Finite state transducer for classifying money, suppletive aware, e.g. $12.05 -> money { integer_part: "twelve" currency_maj: "dollars" fractional_part: "five" currency_min: "cents" preserve_order: true } $12.0500 -> money { integer_part: "twelve" currency_maj: "dollars" fractional_part: "five" currency_min: "cents" preserve_order: true } $1 -> money { currency_maj: "dollar" integer_part: "one" } @@ -46,14 +44,18 @@ class MoneyFst(GraphFst): maj_unit_plural = convert_space(maj_singular @ SINGULAR_TO_PLURAL) maj_unit_singular = convert_space(maj_singular) - graph_maj_singular = pynutil.insert("currency_maj: \"") + maj_unit_singular + pynutil.insert("\"") - graph_maj_plural = pynutil.insert("currency_maj: \"") + maj_unit_plural + pynutil.insert("\"") + graph_maj_singular = ( + pynutil.insert('currency_maj: "') + maj_unit_singular + pynutil.insert('"') + ) + graph_maj_plural = pynutil.insert('currency_maj: "') + maj_unit_plural + pynutil.insert('"') optional_delete_fractional_zeros = pynini.closure( pynutil.delete(".") + pynini.closure(pynutil.delete("0"), 1), 0, 1 ) - graph_integer_one = pynutil.insert("integer_part: \"") + pynini.cross("1", "one") + pynutil.insert("\"") + graph_integer_one = ( + pynutil.insert('integer_part: "') + pynini.cross("1", "one") + pynutil.insert('"') + ) # only for decimals where third decimal after comma is non-zero or with quantity decimal_delete_last_zeros = ( pynini.closure(DAMO_DIGIT | pynutil.delete(",")) @@ -65,11 +67,15 @@ class MoneyFst(GraphFst): decimal_with_quantity = DAMO_SIGMA + DAMO_ALPHA graph_decimal = ( - graph_maj_plural + insert_space + (decimal_delete_last_zeros | decimal_with_quantity) @ graph_decimal_final + graph_maj_plural + + insert_space + + (decimal_delete_last_zeros | decimal_with_quantity) @ graph_decimal_final ) graph_integer = ( - pynutil.insert("integer_part: \"") + ((DAMO_SIGMA - "1") @ cardinal_graph) + pynutil.insert("\"") + pynutil.insert('integer_part: "') + + ((DAMO_SIGMA - "1") @ cardinal_graph) + + pynutil.insert('"') ) graph_integer_only = graph_maj_singular + insert_space + graph_integer_one @@ -88,51 +94,71 @@ class MoneyFst(GraphFst): | ((DAMO_DIGIT - "0") + DAMO_DIGIT) ) - graph_min_singular = pynutil.insert(" currency_min: \"") + min_singular + pynutil.insert("\"") - graph_min_plural = pynutil.insert(" currency_min: \"") + min_plural + pynutil.insert("\"") + graph_min_singular = pynutil.insert(' currency_min: "') + min_singular + pynutil.insert('"') + graph_min_plural = pynutil.insert(' currency_min: "') + min_plural + pynutil.insert('"') # format ** dollars ** cent decimal_graph_with_minor = None integer_graph_reordered = None decimal_default_reordered = None for curr_symbol, _ in maj_singular_labels: preserve_order = pynutil.insert(" preserve_order: true") - integer_plus_maj = graph_integer + insert_space + pynutil.insert(curr_symbol) @ graph_maj_plural - integer_plus_maj |= graph_integer_one + insert_space + pynutil.insert(curr_symbol) @ graph_maj_singular + integer_plus_maj = ( + graph_integer + insert_space + pynutil.insert(curr_symbol) @ graph_maj_plural + ) + integer_plus_maj |= ( + graph_integer_one + insert_space + pynutil.insert(curr_symbol) @ graph_maj_singular + ) integer_plus_maj_with_comma = pynini.compose( - DAMO_DIGIT - "0" + pynini.closure(DAMO_DIGIT | pynutil.delete(",")), integer_plus_maj + DAMO_DIGIT - "0" + pynini.closure(DAMO_DIGIT | pynutil.delete(",")), + integer_plus_maj, ) integer_plus_maj = pynini.compose(pynini.closure(DAMO_DIGIT) - "0", integer_plus_maj) integer_plus_maj |= integer_plus_maj_with_comma graph_fractional_one = two_digits_fractional_part @ pynini.cross("1", "one") - graph_fractional_one = pynutil.insert("fractional_part: \"") + graph_fractional_one + pynutil.insert("\"") + graph_fractional_one = ( + pynutil.insert('fractional_part: "') + graph_fractional_one + pynutil.insert('"') + ) graph_fractional = ( two_digits_fractional_part @ (pynini.closure(DAMO_DIGIT, 1, 2) - "1") @ cardinal.graph_hundred_component_at_least_one_none_zero_digit ) - graph_fractional = pynutil.insert("fractional_part: \"") + graph_fractional + pynutil.insert("\"") - - fractional_plus_min = graph_fractional + insert_space + pynutil.insert(curr_symbol) @ graph_min_plural - fractional_plus_min |= ( - graph_fractional_one + insert_space + pynutil.insert(curr_symbol) @ graph_min_singular + graph_fractional = ( + pynutil.insert('fractional_part: "') + graph_fractional + pynutil.insert('"') ) - decimal_graph_with_minor_curr = integer_plus_maj + pynini.cross(".", " ") + fractional_plus_min + fractional_plus_min = ( + graph_fractional + insert_space + pynutil.insert(curr_symbol) @ graph_min_plural + ) + fractional_plus_min |= ( + graph_fractional_one + + insert_space + + pynutil.insert(curr_symbol) @ graph_min_singular + ) + + decimal_graph_with_minor_curr = ( + integer_plus_maj + pynini.cross(".", " ") + fractional_plus_min + ) if not deterministic: decimal_graph_with_minor_curr |= pynutil.add_weight( integer_plus_maj + pynini.cross(".", " ") - + pynutil.insert("fractional_part: \"") - + two_digits_fractional_part @ cardinal.graph_hundred_component_at_least_one_none_zero_digit - + pynutil.insert("\""), + + pynutil.insert('fractional_part: "') + + two_digits_fractional_part + @ cardinal.graph_hundred_component_at_least_one_none_zero_digit + + pynutil.insert('"'), weight=0.0001, ) - default_fraction_graph = (decimal_delete_last_zeros | decimal_with_quantity) @ graph_decimal_final + default_fraction_graph = ( + decimal_delete_last_zeros | decimal_with_quantity + ) @ graph_decimal_final decimal_graph_with_minor_curr |= ( - pynini.closure(pynutil.delete("0"), 0, 1) + pynutil.delete(".") + fractional_plus_min + pynini.closure(pynutil.delete("0"), 0, 1) + + pynutil.delete(".") + + fractional_plus_min ) decimal_graph_with_minor_curr = ( pynutil.delete(curr_symbol) + decimal_graph_with_minor_curr + preserve_order @@ -141,7 +167,9 @@ class MoneyFst(GraphFst): decimal_graph_with_minor = ( decimal_graph_with_minor_curr if decimal_graph_with_minor is None - else pynini.union(decimal_graph_with_minor, decimal_graph_with_minor_curr).optimize() + else pynini.union( + decimal_graph_with_minor, decimal_graph_with_minor_curr + ).optimize() ) if not deterministic: @@ -152,7 +180,9 @@ class MoneyFst(GraphFst): integer_graph_reordered = ( integer_graph_reordered_curr if integer_graph_reordered is None - else pynini.union(integer_graph_reordered, integer_graph_reordered_curr).optimize() + else pynini.union( + integer_graph_reordered, integer_graph_reordered_curr + ).optimize() ) decimal_default_reordered_curr = ( pynutil.delete(curr_symbol) @@ -174,7 +204,8 @@ class MoneyFst(GraphFst): final_graph |= integer_graph_reordered | decimal_default_reordered # to handle "$2.00" cases final_graph |= pynini.compose( - DAMO_SIGMA + pynutil.delete(".") + pynini.closure(pynutil.delete("0"), 1), integer_graph_reordered + DAMO_SIGMA + pynutil.delete(".") + pynini.closure(pynutil.delete("0"), 1), + integer_graph_reordered, ) final_graph = self.add_tokens(final_graph.optimize()) self.fst = final_graph.optimize() diff --git a/fun_text_processing/text_normalization/en/taggers/ordinal.py b/fun_text_processing/text_normalization/en/taggers/ordinal.py index a44e2aad1..79dcefbc7 100644 --- a/fun_text_processing/text_normalization/en/taggers/ordinal.py +++ b/fun_text_processing/text_normalization/en/taggers/ordinal.py @@ -1,6 +1,3 @@ - - - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_DIGIT, GraphFst from pynini.lib import pynutil @@ -10,7 +7,7 @@ class OrdinalFst(GraphFst): """ Finite state transducer for classifying ordinal, e.g. 13th -> ordinal { integer: "thirteen" } - + Args: cardinal: CardinalFst deterministic: if True will provide a single transduction option, @@ -44,6 +41,6 @@ class OrdinalFst(GraphFst): 1, ) + pynutil.delete(pynini.union("th", "TH")) self.graph = (st_format | nd_format | rd_format | th_format) @ cardinal_graph - final_graph = pynutil.insert("integer: \"") + self.graph + pynutil.insert("\"") + final_graph = pynutil.insert('integer: "') + self.graph + pynutil.insert('"') final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/text_normalization/en/taggers/punctuation.py b/fun_text_processing/text_normalization/en/taggers/punctuation.py index 9b5794339..193dca07c 100644 --- a/fun_text_processing/text_normalization/en/taggers/punctuation.py +++ b/fun_text_processing/text_normalization/en/taggers/punctuation.py @@ -1,10 +1,12 @@ - - import sys from unicodedata import category import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_SPACE, DAMO_SIGMA, GraphFst +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_SPACE, + DAMO_SIGMA, + GraphFst, +) from fun_text_processing.text_normalization.en.utils import get_abs_path, load_labels from pynini.examples import plurals from pynini.lib import pynutil @@ -23,7 +25,7 @@ class PunctuationFst(GraphFst): def __init__(self, deterministic: bool = True): super().__init__(name="punctuation", kind="classify", deterministic=deterministic) - s = "!#%&\'()*+,-./:;<=>?@^_`{|}~\"" + s = "!#%&'()*+,-./:;<=>?@^_`{|}~\"" punct_symbols_to_exclude = ["[", "]"] punct_unicode = [ @@ -42,7 +44,10 @@ class PunctuationFst(GraphFst): emphasis = ( pynini.accep("<") + ( - (pynini.closure(DAMO_NOT_SPACE - pynini.union("<", ">"), 1) + pynini.closure(pynini.accep("/"), 0, 1)) + ( + pynini.closure(DAMO_NOT_SPACE - pynini.union("<", ">"), 1) + + pynini.closure(pynini.accep("/"), 0, 1) + ) | (pynini.accep("/") + pynini.closure(DAMO_NOT_SPACE - pynini.union("<", ">"), 1)) ) + pynini.accep(">") @@ -50,4 +55,4 @@ class PunctuationFst(GraphFst): punct = plurals._priority_union(emphasis, punct, DAMO_SIGMA) self.graph = punct - self.fst = (pynutil.insert("name: \"") + self.graph + pynutil.insert("\"")).optimize() + self.fst = (pynutil.insert('name: "') + self.graph + pynutil.insert('"')).optimize() diff --git a/fun_text_processing/text_normalization/en/taggers/range.py b/fun_text_processing/text_normalization/en/taggers/range.py index 0c76993c8..21fce605c 100644 --- a/fun_text_processing/text_normalization/en/taggers/range.py +++ b/fun_text_processing/text_normalization/en/taggers/range.py @@ -1,13 +1,16 @@ - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_DIGIT, GraphFst, convert_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_DIGIT, + GraphFst, + convert_space, +) from pynini.lib import pynutil class RangeFst(GraphFst): """ This class is a composite class of two other class instances - + Args: time: composed tagger and verbalizer date: composed tagger and verbalizer @@ -18,7 +21,12 @@ class RangeFst(GraphFst): """ def __init__( - self, time: GraphFst, date: GraphFst, cardinal: GraphFst, deterministic: bool = True, lm: bool = False, + self, + time: GraphFst, + date: GraphFst, + cardinal: GraphFst, + deterministic: bool = True, + lm: bool = False, ): super().__init__(name="range", kind="classify", deterministic=deterministic) @@ -32,16 +40,20 @@ class RangeFst(GraphFst): cardinal = cardinal.graph_with_and # YEAR - date_year_four_digit = (DAMO_DIGIT ** 4 + pynini.closure(pynini.accep("s"), 0, 1)) @ date - date_year_two_digit = (DAMO_DIGIT ** 2 + pynini.closure(pynini.accep("s"), 0, 1)) @ date + date_year_four_digit = (DAMO_DIGIT**4 + pynini.closure(pynini.accep("s"), 0, 1)) @ date + date_year_two_digit = (DAMO_DIGIT**2 + pynini.closure(pynini.accep("s"), 0, 1)) @ date year_to_year_graph = ( date_year_four_digit + delete_space + pynini.cross("-", " to ") + delete_space - + (date_year_four_digit | date_year_two_digit | (DAMO_DIGIT ** 2 @ cardinal)) + + (date_year_four_digit | date_year_two_digit | (DAMO_DIGIT**2 @ cardinal)) + ) + mid_year_graph = ( + pynini.accep("mid") + + pynini.cross("-", " ") + + (date_year_four_digit | date_year_two_digit) ) - mid_year_graph = pynini.accep("mid") + pynini.cross("-", " ") + (date_year_four_digit | date_year_two_digit) self.graph |= year_to_year_graph self.graph |= mid_year_graph @@ -55,7 +67,11 @@ class RangeFst(GraphFst): if not deterministic or lm: # cardinal ---- cardinal_to_cardinal_graph = ( - cardinal + delete_space + pynini.cross("-", pynini.union(" to ", " minus ")) + delete_space + cardinal + cardinal + + delete_space + + pynini.cross("-", pynini.union(" to ", " minus ")) + + delete_space + + cardinal ) range_graph |= cardinal_to_cardinal_graph | ( @@ -79,10 +95,14 @@ class RangeFst(GraphFst): ) for x in ["/", " / "]: - range_graph |= cardinal + pynini.closure(pynini.cross(x, " divided by ") + cardinal, 1) + range_graph |= cardinal + pynini.closure( + pynini.cross(x, " divided by ") + cardinal, 1 + ) self.graph |= range_graph self.graph = self.graph.optimize() - graph = pynutil.insert("name: \"") + convert_space(self.graph).optimize() + pynutil.insert("\"") + graph = ( + pynutil.insert('name: "') + convert_space(self.graph).optimize() + pynutil.insert('"') + ) self.fst = graph.optimize() diff --git a/fun_text_processing/text_normalization/en/taggers/roman.py b/fun_text_processing/text_normalization/en/taggers/roman.py index fbd27ca79..3e71cca32 100644 --- a/fun_text_processing/text_normalization/en/taggers/roman.py +++ b/fun_text_processing/text_normalization/en/taggers/roman.py @@ -1,6 +1,3 @@ - - - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_ALPHA, DAMO_SIGMA, GraphFst from fun_text_processing.text_normalization.en.utils import get_abs_path, load_labels @@ -22,7 +19,7 @@ class RomanFst(GraphFst): roman_dict = load_labels(get_abs_path("data/roman/roman_to_spoken.tsv")) default_graph = pynini.string_map(roman_dict).optimize() - default_graph = pynutil.insert("integer: \"") + default_graph + pynutil.insert("\"") + default_graph = pynutil.insert('integer: "') + default_graph + pynutil.insert('"') ordinal_limit = 19 if deterministic: @@ -31,14 +28,16 @@ class RomanFst(GraphFst): else: start_idx = 0 - graph_teens = pynini.string_map([x[0] for x in roman_dict[start_idx:ordinal_limit]]).optimize() + graph_teens = pynini.string_map( + [x[0] for x in roman_dict[start_idx:ordinal_limit]] + ).optimize() # roman numerals up to ordinal_limit with a preceding name are converted to ordinal form names = get_names() graph = ( - pynutil.insert("key_the_ordinal: \"") + pynutil.insert('key_the_ordinal: "') + names - + pynutil.insert("\"") + + pynutil.insert('"') + pynini.accep(" ") + graph_teens @ default_graph ).optimize() @@ -52,7 +51,11 @@ class RomanFst(GraphFst): key_words = pynini.string_map(key_words).optimize() graph |= ( - pynutil.insert("key_cardinal: \"") + key_words + pynutil.insert("\"") + pynini.accep(" ") + default_graph + pynutil.insert('key_cardinal: "') + + key_words + + pynutil.insert('"') + + pynini.accep(" ") + + default_graph ).optimize() if deterministic or lm: @@ -60,8 +63,9 @@ class RomanFst(GraphFst): roman_to_cardinal = pynini.compose( pynini.closure(DAMO_ALPHA, 2), ( - pynutil.insert("default_cardinal: \"default\" ") - + (pynini.string_map([x[0] for x in roman_dict[:50]]).optimize()) @ default_graph + pynutil.insert('default_cardinal: "default" ') + + (pynini.string_map([x[0] for x in roman_dict[:50]]).optimize()) + @ default_graph ), ) graph |= roman_to_cardinal @@ -70,9 +74,9 @@ class RomanFst(GraphFst): roman_to_cardinal = pynini.compose( pynini.difference(DAMO_SIGMA, "I"), ( - pynutil.insert("default_cardinal: \"default\" integer: \"") + pynutil.insert('default_cardinal: "default" integer: "') + pynini.string_map(roman_dict).optimize() - + pynutil.insert("\"") + + pynutil.insert('"') ), ).optimize() graph |= roman_to_cardinal @@ -80,7 +84,11 @@ class RomanFst(GraphFst): # convert three digit roman or up with suffix to ordinal roman_to_ordinal = pynini.compose( pynini.closure(DAMO_ALPHA, 3), - (pynutil.insert("default_ordinal: \"default\" ") + graph_teens @ default_graph + pynutil.delete("th")), + ( + pynutil.insert('default_ordinal: "default" ') + + graph_teens @ default_graph + + pynutil.delete("th") + ), ) graph |= roman_to_ordinal diff --git a/fun_text_processing/text_normalization/en/taggers/serial.py b/fun_text_processing/text_normalization/en/taggers/serial.py index f039f6aec..ee31f6e3c 100644 --- a/fun_text_processing/text_normalization/en/taggers/serial.py +++ b/fun_text_processing/text_normalization/en/taggers/serial.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_ALPHA, @@ -16,7 +15,7 @@ from pynini.lib import pynutil class SerialFst(GraphFst): """ This class is a composite class of two other class instances - + Args: time: composed tagger and verbalizer date: composed tagger and verbalizer @@ -26,7 +25,9 @@ class SerialFst(GraphFst): lm: whether to use for hybrid LM """ - def __init__(self, cardinal: GraphFst, ordinal: GraphFst, deterministic: bool = True, lm: bool = False): + def __init__( + self, cardinal: GraphFst, ordinal: GraphFst, deterministic: bool = True, lm: bool = False + ): super().__init__(name="integer", kind="classify", deterministic=deterministic) """ @@ -42,16 +43,17 @@ class SerialFst(GraphFst): pynini.accep("0") + pynini.closure(DAMO_DIGIT), cardinal.single_digits_graph ).optimize() # TODO: "#" doesn't work from the file - symbols_graph = pynini.string_file(get_abs_path("data/whitelist/symbol.tsv")).optimize() | pynini.cross( - "#", "hash" - ) + symbols_graph = pynini.string_file( + get_abs_path("data/whitelist/symbol.tsv") + ).optimize() | pynini.cross("#", "hash") num_graph |= symbols_graph if not self.deterministic and not lm: num_graph |= cardinal.single_digits_graph # also allow double digits to be pronounced as integer in serial number num_graph |= pynutil.add_weight( - DAMO_DIGIT ** 2 @ cardinal.graph_hundred_component_at_least_one_none_zero_digit, weight=0.0001 + DAMO_DIGIT**2 @ cardinal.graph_hundred_component_at_least_one_none_zero_digit, + weight=0.0001, ) # add space between letter and digit/symbol @@ -84,10 +86,17 @@ class SerialFst(GraphFst): serial_graph |= num_letter + next_alpha_or_num # numbers only with 2+ delimiters serial_graph |= ( - num_graph + delimiter + num_graph + delimiter + num_graph + pynini.closure(delimiter + num_graph) + num_graph + + delimiter + + num_graph + + delimiter + + num_graph + + pynini.closure(delimiter + num_graph) ) # 2+ symbols - serial_graph |= pynini.compose(DAMO_SIGMA + symbols + DAMO_SIGMA, num_graph + delimiter + num_graph) + serial_graph |= pynini.compose( + DAMO_SIGMA + symbols + DAMO_SIGMA, num_graph + delimiter + num_graph + ) # exclude ordinal numbers from serial options serial_graph = pynini.compose( @@ -113,10 +122,13 @@ class SerialFst(GraphFst): # this is not to verbolize "/" as "slash" in cases like "import/export" serial_graph = pynini.compose( pynini.difference( - DAMO_SIGMA, pynini.closure(DAMO_ALPHA, 1) + pynini.accep("/") + pynini.closure(DAMO_ALPHA, 1) + DAMO_SIGMA, + pynini.closure(DAMO_ALPHA, 1) + pynini.accep("/") + pynini.closure(DAMO_ALPHA, 1), ), serial_graph, ) self.graph = serial_graph.optimize() - graph = pynutil.insert("name: \"") + convert_space(self.graph).optimize() + pynutil.insert("\"") + graph = ( + pynutil.insert('name: "') + convert_space(self.graph).optimize() + pynutil.insert('"') + ) self.fst = graph.optimize() diff --git a/fun_text_processing/text_normalization/en/taggers/telephone.py b/fun_text_processing/text_normalization/en/taggers/telephone.py index 8ac473fca..3b9920f27 100644 --- a/fun_text_processing/text_normalization/en/taggers/telephone.py +++ b/fun_text_processing/text_normalization/en/taggers/telephone.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_ALPHA, @@ -17,11 +15,11 @@ from pynini.lib import pynutil class TelephoneFst(GraphFst): """ - Finite state transducer for classifying telephone, and IP, and SSN which includes country code, number part and extension - country code optional: +*** + Finite state transducer for classifying telephone, and IP, and SSN which includes country code, number part and extension + country code optional: +*** number part: ***-***-****, or (***) ***-**** extension optional: 1-9999 - E.g + E.g +1 123-123-5678-1 -> telephone { country_code: "one" number_part: "one two three, one two three, five six seven eight" extension: "one" } 1-800-GO-U-HAUL -> telephone { country_code: "one" number_part: "one, eight hundred GO U HAUL" } Args: @@ -36,7 +34,10 @@ class TelephoneFst(GraphFst): zero = pynini.cross("0", "zero") if not deterministic: zero |= pynini.cross("0", pynini.union("o", "oh")) - digit = pynini.invert(pynini.string_file(get_abs_path("data/number/digit.tsv"))).optimize() | zero + digit = ( + pynini.invert(pynini.string_file(get_abs_path("data/number/digit.tsv"))).optimize() + | zero + ) telephone_prompts = pynini.string_file(get_abs_path("data/telephone/telephone_prompt.tsv")) country_code = ( @@ -47,8 +48,10 @@ class TelephoneFst(GraphFst): + pynutil.insert(",") ) country_code |= telephone_prompts - country_code = pynutil.insert("country_code: \"") + country_code + pynutil.insert("\"") - country_code = country_code + pynini.closure(pynutil.delete("-"), 0, 1) + delete_space + insert_space + country_code = pynutil.insert('country_code: "') + country_code + pynutil.insert('"') + country_code = ( + country_code + pynini.closure(pynutil.delete("-"), 0, 1) + delete_space + insert_space + ) area_part_default = pynini.closure(digit + insert_space, 2, 2) + digit area_part = pynini.cross("800", "eight hundred") | pynini.compose( @@ -60,32 +63,42 @@ class TelephoneFst(GraphFst): | ( pynutil.delete("(") + area_part - + ((pynutil.delete(")") + pynini.closure(pynutil.delete(" "), 0, 1)) | pynutil.delete(")-")) + + ( + (pynutil.delete(")") + pynini.closure(pynutil.delete(" "), 0, 1)) + | pynutil.delete(")-") + ) ) ) + add_separator del_separator = pynini.closure(pynini.union("-", " ", "."), 0, 1) number_length = ((DAMO_DIGIT + del_separator) | (DAMO_ALPHA + del_separator)) ** 7 number_words = pynini.closure( - (DAMO_DIGIT @ digit) + (insert_space | (pynini.cross("-", ', '))) + (DAMO_DIGIT @ digit) + (insert_space | (pynini.cross("-", ", "))) | DAMO_ALPHA - | (DAMO_ALPHA + pynini.cross("-", ' ')) + | (DAMO_ALPHA + pynini.cross("-", " ")) ) number_words |= pynini.closure( - (DAMO_DIGIT @ digit) + (insert_space | (pynini.cross(".", ', '))) + (DAMO_DIGIT @ digit) + (insert_space | (pynini.cross(".", ", "))) | DAMO_ALPHA - | (DAMO_ALPHA + pynini.cross(".", ' ')) + | (DAMO_ALPHA + pynini.cross(".", " ")) ) number_words = pynini.compose(number_length, number_words) number_part = area_part + number_words - number_part = pynutil.insert("number_part: \"") + number_part + pynutil.insert("\"") + number_part = pynutil.insert('number_part: "') + number_part + pynutil.insert('"') extension = ( - pynutil.insert("extension: \"") + pynini.closure(digit + insert_space, 0, 3) + digit + pynutil.insert("\"") + pynutil.insert('extension: "') + + pynini.closure(digit + insert_space, 0, 3) + + digit + + pynutil.insert('"') ) extension = pynini.closure(insert_space + extension, 0, 1) - graph = plurals._priority_union(country_code + number_part, number_part, DAMO_SIGMA).optimize() - graph = plurals._priority_union(country_code + number_part + extension, graph, DAMO_SIGMA).optimize() + graph = plurals._priority_union( + country_code + number_part, number_part, DAMO_SIGMA + ).optimize() + graph = plurals._priority_union( + country_code + number_part + extension, graph, DAMO_SIGMA + ).optimize() graph = plurals._priority_union(number_part + extension, graph, DAMO_SIGMA).optimize() # ip @@ -94,11 +107,16 @@ class TelephoneFst(GraphFst): ip_graph = digit_to_str_graph + (pynini.cross(".", " dot ") + digit_to_str_graph) ** 3 graph |= ( pynini.closure( - pynutil.insert("country_code: \"") + ip_prompts + pynutil.insert("\"") + delete_extra_space, 0, 1 + pynutil.insert('country_code: "') + + ip_prompts + + pynutil.insert('"') + + delete_extra_space, + 0, + 1, ) - + pynutil.insert("number_part: \"") + + pynutil.insert('number_part: "') + ip_graph.optimize() - + pynutil.insert("\"") + + pynutil.insert('"') ) # ssn ssn_prompts = pynini.string_file(get_abs_path("data/telephone/ssn_prompt.tsv")) @@ -106,15 +124,22 @@ class TelephoneFst(GraphFst): two_digit_part = digit + pynutil.insert(" ") + digit four_digit_part = digit + (pynutil.insert(" ") + digit) ** 3 ssn_separator = pynini.cross("-", ", ") - ssn_graph = three_digit_part + ssn_separator + two_digit_part + ssn_separator + four_digit_part + ssn_graph = ( + three_digit_part + ssn_separator + two_digit_part + ssn_separator + four_digit_part + ) graph |= ( pynini.closure( - pynutil.insert("country_code: \"") + ssn_prompts + pynutil.insert("\"") + delete_extra_space, 0, 1 + pynutil.insert('country_code: "') + + ssn_prompts + + pynutil.insert('"') + + delete_extra_space, + 0, + 1, ) - + pynutil.insert("number_part: \"") + + pynutil.insert('number_part: "') + ssn_graph.optimize() - + pynutil.insert("\"") + + pynutil.insert('"') ) final_graph = self.add_tokens(graph) diff --git a/fun_text_processing/text_normalization/en/taggers/time.py b/fun_text_processing/text_normalization/en/taggers/time.py index be2297ff8..60c6d97c4 100644 --- a/fun_text_processing/text_normalization/en/taggers/time.py +++ b/fun_text_processing/text_normalization/en/taggers/time.py @@ -1,6 +1,3 @@ - - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_DIGIT, @@ -28,7 +25,7 @@ class TimeFst(GraphFst): 02:00 -> time { hours: "two" } 2:00 -> time { hours: "two" } 10:00:05 a.m. -> time { hours: "ten" minutes: "zero" seconds: "five" suffix: "a m" } - + Args: cardinal: CardinalFst deterministic: if True will provide a single transduction option, @@ -59,25 +56,27 @@ class TimeFst(GraphFst): graph_minute_single = pynini.union(*labels_minute_single) @ cardinal graph_minute_double = pynini.union(*labels_minute_double) @ cardinal - final_graph_hour = pynutil.insert("hours: \"") + graph_hour + pynutil.insert("\"") + final_graph_hour = pynutil.insert('hours: "') + graph_hour + pynutil.insert('"') final_graph_minute = ( - pynutil.insert("minutes: \"") + pynutil.insert('minutes: "') + (pynini.cross("0", "o") + insert_space + graph_minute_single | graph_minute_double) - + pynutil.insert("\"") + + pynutil.insert('"') ) final_graph_second = ( - pynutil.insert("seconds: \"") + pynutil.insert('seconds: "') + (pynini.cross("0", "o") + insert_space + graph_minute_single | graph_minute_double) - + pynutil.insert("\"") + + pynutil.insert('"') + ) + final_suffix = ( + pynutil.insert('suffix: "') + convert_space(suffix_graph) + pynutil.insert('"') ) - final_suffix = pynutil.insert("suffix: \"") + convert_space(suffix_graph) + pynutil.insert("\"") final_suffix_optional = pynini.closure(delete_space + insert_space + final_suffix, 0, 1) final_time_zone_optional = pynini.closure( delete_space + insert_space - + pynutil.insert("zone: \"") + + pynutil.insert('zone: "') + convert_space(time_zone_graph) - + pynutil.insert("\""), + + pynutil.insert('"'), 0, 1, ) @@ -95,9 +94,9 @@ class TimeFst(GraphFst): graph_hms = ( final_graph_hour + pynutil.delete(":") - + (pynini.cross("00", " minutes: \"zero\"") | insert_space + final_graph_minute) + + (pynini.cross("00", ' minutes: "zero"') | insert_space + final_graph_minute) + pynutil.delete(":") - + (pynini.cross("00", " seconds: \"zero\"") | insert_space + final_graph_second) + + (pynini.cross("00", ' seconds: "zero"') | insert_space + final_graph_second) + final_suffix_optional + final_time_zone_optional ) @@ -113,7 +112,9 @@ class TimeFst(GraphFst): + final_time_zone_optional ) # 2 pm est - graph_h = final_graph_hour + delete_space + insert_space + final_suffix + final_time_zone_optional + graph_h = ( + final_graph_hour + delete_space + insert_space + final_suffix + final_time_zone_optional + ) final_graph = (graph_hm | graph_h | graph_hm2 | graph_hms).optimize() final_graph = self.add_tokens(final_graph) diff --git a/fun_text_processing/text_normalization/en/taggers/tokenize_and_classify.py b/fun_text_processing/text_normalization/en/taggers/tokenize_and_classify.py index 3b5590dd4..0fb8221bf 100644 --- a/fun_text_processing/text_normalization/en/taggers/tokenize_and_classify.py +++ b/fun_text_processing/text_normalization/en/taggers/tokenize_and_classify.py @@ -1,5 +1,3 @@ - - import os import time @@ -39,9 +37,9 @@ import logging class ClassifyFst(GraphFst): """ Final class that composes all other classification grammars. This class can process an entire sentence including punctuation. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. - + Args: input_case: accepting either "lower_cased" or "cased" input. deterministic: if True will provide a single transduction option, @@ -66,77 +64,112 @@ class ClassifyFst(GraphFst): os.makedirs(cache_dir, exist_ok=True) whitelist_file = os.path.basename(whitelist) if whitelist else "" far_file = os.path.join( - cache_dir, f"en_tn_{deterministic}_deterministic_{input_case}_{whitelist_file}_tokenize.far" + cache_dir, + f"en_tn_{deterministic}_deterministic_{input_case}_{whitelist_file}_tokenize.far", ) if not overwrite_cache and far_file and os.path.exists(far_file): self.fst = pynini.Far(far_file, mode="r")["tokenize_and_classify"] - logging.info(f'ClassifyFst.fst was restored from {far_file}.') + logging.info(f"ClassifyFst.fst was restored from {far_file}.") else: logging.info(f"Creating ClassifyFst grammars.") start_time = time.time() cardinal = CardinalFst(deterministic=deterministic) cardinal_graph = cardinal.fst - logging.debug(f"cardinal: {time.time() - start_time: .2f}s -- {cardinal_graph.num_states()} nodes") + logging.debug( + f"cardinal: {time.time() - start_time: .2f}s -- {cardinal_graph.num_states()} nodes" + ) start_time = time.time() ordinal = OrdinalFst(cardinal=cardinal, deterministic=deterministic) ordinal_graph = ordinal.fst - logging.debug(f"ordinal: {time.time() - start_time: .2f}s -- {ordinal_graph.num_states()} nodes") + logging.debug( + f"ordinal: {time.time() - start_time: .2f}s -- {ordinal_graph.num_states()} nodes" + ) start_time = time.time() decimal = DecimalFst(cardinal=cardinal, deterministic=deterministic) decimal_graph = decimal.fst - logging.debug(f"decimal: {time.time() - start_time: .2f}s -- {decimal_graph.num_states()} nodes") + logging.debug( + f"decimal: {time.time() - start_time: .2f}s -- {decimal_graph.num_states()} nodes" + ) start_time = time.time() fraction = FractionFst(deterministic=deterministic, cardinal=cardinal) fraction_graph = fraction.fst - logging.debug(f"fraction: {time.time() - start_time: .2f}s -- {fraction_graph.num_states()} nodes") + logging.debug( + f"fraction: {time.time() - start_time: .2f}s -- {fraction_graph.num_states()} nodes" + ) start_time = time.time() - measure = MeasureFst(cardinal=cardinal, decimal=decimal, fraction=fraction, deterministic=deterministic) + measure = MeasureFst( + cardinal=cardinal, decimal=decimal, fraction=fraction, deterministic=deterministic + ) measure_graph = measure.fst - logging.debug(f"measure: {time.time() - start_time: .2f}s -- {measure_graph.num_states()} nodes") + logging.debug( + f"measure: {time.time() - start_time: .2f}s -- {measure_graph.num_states()} nodes" + ) start_time = time.time() date_graph = DateFst(cardinal=cardinal, deterministic=deterministic).fst - logging.debug(f"date: {time.time() - start_time: .2f}s -- {date_graph.num_states()} nodes") + logging.debug( + f"date: {time.time() - start_time: .2f}s -- {date_graph.num_states()} nodes" + ) start_time = time.time() time_graph = TimeFst(cardinal=cardinal, deterministic=deterministic).fst - logging.debug(f"time: {time.time() - start_time: .2f}s -- {time_graph.num_states()} nodes") + logging.debug( + f"time: {time.time() - start_time: .2f}s -- {time_graph.num_states()} nodes" + ) start_time = time.time() telephone_graph = TelephoneFst(deterministic=deterministic).fst - logging.debug(f"telephone: {time.time() - start_time: .2f}s -- {telephone_graph.num_states()} nodes") + logging.debug( + f"telephone: {time.time() - start_time: .2f}s -- {telephone_graph.num_states()} nodes" + ) start_time = time.time() electonic_graph = ElectronicFst(deterministic=deterministic).fst - logging.debug(f"electronic: {time.time() - start_time: .2f}s -- {electonic_graph.num_states()} nodes") + logging.debug( + f"electronic: {time.time() - start_time: .2f}s -- {electonic_graph.num_states()} nodes" + ) start_time = time.time() - money_graph = MoneyFst(cardinal=cardinal, decimal=decimal, deterministic=deterministic).fst - logging.debug(f"money: {time.time() - start_time: .2f}s -- {money_graph.num_states()} nodes") + money_graph = MoneyFst( + cardinal=cardinal, decimal=decimal, deterministic=deterministic + ).fst + logging.debug( + f"money: {time.time() - start_time: .2f}s -- {money_graph.num_states()} nodes" + ) start_time = time.time() whitelist_graph = WhiteListFst( input_case=input_case, deterministic=deterministic, input_file=whitelist ).fst - logging.debug(f"whitelist: {time.time() - start_time: .2f}s -- {whitelist_graph.num_states()} nodes") + logging.debug( + f"whitelist: {time.time() - start_time: .2f}s -- {whitelist_graph.num_states()} nodes" + ) start_time = time.time() punctuation = PunctuationFst(deterministic=deterministic) punct_graph = punctuation.fst - logging.debug(f"punct: {time.time() - start_time: .2f}s -- {punct_graph.num_states()} nodes") + logging.debug( + f"punct: {time.time() - start_time: .2f}s -- {punct_graph.num_states()} nodes" + ) start_time = time.time() word_graph = WordFst(punctuation=punctuation, deterministic=deterministic).fst - logging.debug(f"word: {time.time() - start_time: .2f}s -- {word_graph.num_states()} nodes") + logging.debug( + f"word: {time.time() - start_time: .2f}s -- {word_graph.num_states()} nodes" + ) start_time = time.time() - serial_graph = SerialFst(cardinal=cardinal, ordinal=ordinal, deterministic=deterministic).fst - logging.debug(f"serial: {time.time() - start_time: .2f}s -- {serial_graph.num_states()} nodes") + serial_graph = SerialFst( + cardinal=cardinal, ordinal=ordinal, deterministic=deterministic + ).fst + logging.debug( + f"serial: {time.time() - start_time: .2f}s -- {serial_graph.num_states()} nodes" + ) start_time = time.time() v_time_graph = vTimeFst(deterministic=deterministic).fst @@ -147,7 +180,9 @@ class ClassifyFst(GraphFst): range_graph = RangeFst( time=time_final, date=date_final, cardinal=cardinal, deterministic=deterministic ).fst - logging.debug(f"range: {time.time() - start_time: .2f}s -- {range_graph.num_states()} nodes") + logging.debug( + f"range: {time.time() - start_time: .2f}s -- {range_graph.num_states()} nodes" + ) classify = ( pynutil.add_weight(whitelist_graph, 1.01) @@ -162,7 +197,9 @@ class ClassifyFst(GraphFst): | pynutil.add_weight(electonic_graph, 1.1) | pynutil.add_weight(fraction_graph, 1.1) | pynutil.add_weight(range_graph, 1.1) - | pynutil.add_weight(serial_graph, 1.1001) # should be higher than the rest of the classes + | pynutil.add_weight( + serial_graph, 1.1001 + ) # should be higher than the rest of the classes ) roman_graph = RomanFst(deterministic=deterministic).fst @@ -172,7 +209,11 @@ class ClassifyFst(GraphFst): abbreviation_graph = AbbreviationFst(deterministic=deterministic).fst classify |= pynutil.add_weight(abbreviation_graph, 100) - punct = pynutil.insert("tokens { ") + pynutil.add_weight(punct_graph, weight=2.1) + pynutil.insert(" }") + punct = ( + pynutil.insert("tokens { ") + + pynutil.add_weight(punct_graph, weight=2.1) + + pynutil.insert(" }") + ) punct = pynini.closure( pynini.compose(pynini.closure(DAMO_WHITE_SPACE, 1), delete_extra_space) | (pynutil.insert(" ") + punct), @@ -182,7 +223,9 @@ class ClassifyFst(GraphFst): classify |= pynutil.add_weight(word_graph, 100) token = pynutil.insert("tokens { ") + classify + pynutil.insert(" }") token_plus_punct = ( - pynini.closure(punct + pynutil.insert(" ")) + token + pynini.closure(pynutil.insert(" ") + punct) + pynini.closure(punct + pynutil.insert(" ")) + + token + + pynini.closure(pynutil.insert(" ") + punct) ) graph = token_plus_punct + pynini.closure( diff --git a/fun_text_processing/text_normalization/en/taggers/tokenize_and_classify_lm.py b/fun_text_processing/text_normalization/en/taggers/tokenize_and_classify_lm.py index cdfa85470..497fa6b67 100644 --- a/fun_text_processing/text_normalization/en/taggers/tokenize_and_classify_lm.py +++ b/fun_text_processing/text_normalization/en/taggers/tokenize_and_classify_lm.py @@ -1,5 +1,3 @@ - - import os import pynini @@ -33,13 +31,17 @@ from fun_text_processing.text_normalization.en.taggers.word import WordFst from fun_text_processing.text_normalization.en.verbalizers.cardinal import CardinalFst as vCardinal from fun_text_processing.text_normalization.en.verbalizers.date import DateFst as vDate from fun_text_processing.text_normalization.en.verbalizers.decimal import DecimalFst as vDecimal -from fun_text_processing.text_normalization.en.verbalizers.electronic import ElectronicFst as vElectronic +from fun_text_processing.text_normalization.en.verbalizers.electronic import ( + ElectronicFst as vElectronic, +) from fun_text_processing.text_normalization.en.verbalizers.fraction import FractionFst as vFraction from fun_text_processing.text_normalization.en.verbalizers.measure import MeasureFst as vMeasure from fun_text_processing.text_normalization.en.verbalizers.money import MoneyFst as vMoney from fun_text_processing.text_normalization.en.verbalizers.ordinal import OrdinalFst as vOrdinal from fun_text_processing.text_normalization.en.verbalizers.roman import RomanFst as vRoman -from fun_text_processing.text_normalization.en.verbalizers.telephone import TelephoneFst as vTelephone +from fun_text_processing.text_normalization.en.verbalizers.telephone import ( + TelephoneFst as vTelephone, +) from fun_text_processing.text_normalization.en.verbalizers.time import TimeFst as vTime from fun_text_processing.text_normalization.en.verbalizers.word import WordFst as vWord from pynini.examples import plurals @@ -53,7 +55,7 @@ class ClassifyFst(GraphFst): Final class that composes all other classification grammars. This class can process an entire sentence including punctuation. For deployment, this grammar will be compiled and exported to OpenFst Finite State Archive (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. - + Args: input_case: accepting either "lower_cased" or "cased" input. deterministic: if True will provide a single transduction option, @@ -74,19 +76,20 @@ class ClassifyFst(GraphFst): super().__init__(name="tokenize_and_classify", kind="classify", deterministic=deterministic) far_file = None - if cache_dir is not None and cache_dir != 'None': + if cache_dir is not None and cache_dir != "None": os.makedirs(cache_dir, exist_ok=True) whitelist_file = os.path.basename(whitelist) if whitelist else "" far_file = os.path.join( - cache_dir, f"_{input_case}_en_tn_{deterministic}_deterministic{whitelist_file}_lm.far" + cache_dir, + f"_{input_case}_en_tn_{deterministic}_deterministic{whitelist_file}_lm.far", ) if not overwrite_cache and far_file and os.path.exists(far_file): - self.fst = pynini.Far(far_file, mode='r')['tokenize_and_classify'] + self.fst = pynini.Far(far_file, mode="r")["tokenize_and_classify"] no_digits = pynini.closure(pynini.difference(DAMO_CHAR, DAMO_DIGIT)) self.fst_no_digits = pynini.compose(self.fst, no_digits).optimize() - logging.info(f'ClassifyFst.fst was restored from {far_file}.') + logging.info(f"ClassifyFst.fst was restored from {far_file}.") else: - logging.info(f'Creating ClassifyFst grammars. This might take some time...') + logging.info(f"Creating ClassifyFst grammars. This might take some time...") # TAGGERS cardinal = CardinalFst(deterministic=True, lm=True) cardinal_tagger = cardinal @@ -100,7 +103,9 @@ class ClassifyFst(GraphFst): fraction = FractionFst(deterministic=True, cardinal=cardinal) fraction_graph = fraction.fst - measure = MeasureFst(cardinal=cardinal, decimal=decimal, fraction=fraction, deterministic=True) + measure = MeasureFst( + cardinal=cardinal, decimal=decimal, fraction=fraction, deterministic=True + ) measure_graph = measure.fst date = DateFst(cardinal=cardinal, deterministic=True, lm=True) date_graph = date.fst @@ -111,9 +116,13 @@ class ClassifyFst(GraphFst): telephone_graph = TelephoneFst(deterministic=True).fst electronic_graph = ElectronicFst(deterministic=True).fst money_graph = MoneyFst(cardinal=cardinal, decimal=decimal, deterministic=False).fst - whitelist = WhiteListFst(input_case=input_case, deterministic=False, input_file=whitelist) + whitelist = WhiteListFst( + input_case=input_case, deterministic=False, input_file=whitelist + ) whitelist_graph = whitelist.graph - serial_graph = SerialFst(cardinal=cardinal, ordinal=ordinal, deterministic=deterministic, lm=True).fst + serial_graph = SerialFst( + cardinal=cardinal, ordinal=ordinal, deterministic=deterministic, lm=True + ).fst # VERBALIZERS cardinal = vCardinal(deterministic=True) @@ -126,7 +135,9 @@ class ClassifyFst(GraphFst): v_fraction_graph = fraction.fst v_telephone_graph = vTelephone(deterministic=True).fst v_electronic_graph = vElectronic(deterministic=True).fst - measure = vMeasure(decimal=decimal, cardinal=cardinal, fraction=fraction, deterministic=False) + measure = vMeasure( + decimal=decimal, cardinal=cardinal, fraction=fraction, deterministic=False + ) v_measure_graph = measure.fst v_time_graph = vTime(deterministic=True).fst v_date_graph = vDate(ordinal=ordinal, deterministic=deterministic, lm=True).fst @@ -135,7 +146,9 @@ class ClassifyFst(GraphFst): v_word_graph = vWord(deterministic=deterministic).fst cardinal_or_date_final = plurals._priority_union(date_graph, cardinal_graph, DAMO_SIGMA) - cardinal_or_date_final = pynini.compose(cardinal_or_date_final, (v_cardinal_graph | v_date_graph)) + cardinal_or_date_final = pynini.compose( + cardinal_or_date_final, (v_cardinal_graph | v_date_graph) + ) time_final = pynini.compose(time_graph, v_time_graph) ordinal_final = pynini.compose(ordinal_graph, v_ordinal_graph) @@ -160,14 +173,23 @@ class ClassifyFst(GraphFst): roman_graph = RomanFst(deterministic=deterministic, lm=True).fst # the weight matches the word_graph weight for "I" cases in long sentences with multiple semiotic tokens - classify_and_verbalize |= pynutil.add_weight(pynini.compose(roman_graph, v_roman_graph), sem_w) + classify_and_verbalize |= pynutil.add_weight( + pynini.compose(roman_graph, v_roman_graph), sem_w + ) date_final = pynini.compose(date_graph, v_date_graph) range_graph = RangeFst( - time=time_final, cardinal=cardinal_tagger, date=date_final, deterministic=deterministic + time=time_final, + cardinal=cardinal_tagger, + date=date_final, + deterministic=deterministic, ).fst - classify_and_verbalize |= pynutil.add_weight(pynini.compose(range_graph, v_word_graph), sem_w) - classify_and_verbalize = pynutil.insert("< ") + classify_and_verbalize + pynutil.insert(" >") + classify_and_verbalize |= pynutil.add_weight( + pynini.compose(range_graph, v_word_graph), sem_w + ) + classify_and_verbalize = ( + pynutil.insert("< ") + classify_and_verbalize + pynutil.insert(" >") + ) classify_and_verbalize |= pynutil.add_weight(word_graph, word_w) punct_only = pynutil.add_weight(punct_graph, weight=punct_w) @@ -213,4 +235,4 @@ class ClassifyFst(GraphFst): if far_file: generator_main(far_file, {"tokenize_and_classify": self.fst}) - logging.info(f'ClassifyFst grammars are saved to {far_file}.') + logging.info(f"ClassifyFst grammars are saved to {far_file}.") diff --git a/fun_text_processing/text_normalization/en/taggers/tokenize_and_classify_with_audio.py b/fun_text_processing/text_normalization/en/taggers/tokenize_and_classify_with_audio.py index 853704b91..244adbb38 100644 --- a/fun_text_processing/text_normalization/en/taggers/tokenize_and_classify_with_audio.py +++ b/fun_text_processing/text_normalization/en/taggers/tokenize_and_classify_with_audio.py @@ -1,5 +1,3 @@ - - import os import pynini @@ -30,17 +28,23 @@ from fun_text_processing.text_normalization.en.taggers.telephone import Telephon from fun_text_processing.text_normalization.en.taggers.time import TimeFst from fun_text_processing.text_normalization.en.taggers.whitelist import WhiteListFst from fun_text_processing.text_normalization.en.taggers.word import WordFst -from fun_text_processing.text_normalization.en.verbalizers.abbreviation import AbbreviationFst as vAbbreviation +from fun_text_processing.text_normalization.en.verbalizers.abbreviation import ( + AbbreviationFst as vAbbreviation, +) from fun_text_processing.text_normalization.en.verbalizers.cardinal import CardinalFst as vCardinal from fun_text_processing.text_normalization.en.verbalizers.date import DateFst as vDate from fun_text_processing.text_normalization.en.verbalizers.decimal import DecimalFst as vDecimal -from fun_text_processing.text_normalization.en.verbalizers.electronic import ElectronicFst as vElectronic +from fun_text_processing.text_normalization.en.verbalizers.electronic import ( + ElectronicFst as vElectronic, +) from fun_text_processing.text_normalization.en.verbalizers.fraction import FractionFst as vFraction from fun_text_processing.text_normalization.en.verbalizers.measure import MeasureFst as vMeasure from fun_text_processing.text_normalization.en.verbalizers.money import MoneyFst as vMoney from fun_text_processing.text_normalization.en.verbalizers.ordinal import OrdinalFst as vOrdinal from fun_text_processing.text_normalization.en.verbalizers.roman import RomanFst as vRoman -from fun_text_processing.text_normalization.en.verbalizers.telephone import TelephoneFst as vTelephone +from fun_text_processing.text_normalization.en.verbalizers.telephone import ( + TelephoneFst as vTelephone, +) from fun_text_processing.text_normalization.en.verbalizers.time import TimeFst as vTime from fun_text_processing.text_normalization.en.verbalizers.word import WordFst as vWord from pynini.lib import pynutil @@ -53,7 +57,7 @@ class ClassifyFst(GraphFst): Final class that composes all other classification grammars. This class can process an entire sentence including punctuation. For deployment, this grammar will be compiled and exported to OpenFst Finite State Archive (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. - + Args: input_case: accepting either "lower_cased" or "cased" input. deterministic: if True will provide a single transduction option, @@ -74,19 +78,19 @@ class ClassifyFst(GraphFst): super().__init__(name="tokenize_and_classify", kind="classify", deterministic=deterministic) far_file = None - if cache_dir is not None and cache_dir != 'None': + if cache_dir is not None and cache_dir != "None": os.makedirs(cache_dir, exist_ok=True) whitelist_file = os.path.basename(whitelist) if whitelist else "" far_file = os.path.join( cache_dir, f"_{input_case}_en_tn_{deterministic}_deterministic{whitelist_file}.far" ) if not overwrite_cache and far_file and os.path.exists(far_file): - self.fst = pynini.Far(far_file, mode='r')['tokenize_and_classify'] + self.fst = pynini.Far(far_file, mode="r")["tokenize_and_classify"] no_digits = pynini.closure(pynini.difference(DAMO_CHAR, DAMO_DIGIT)) self.fst_no_digits = pynini.compose(self.fst, no_digits).optimize() - logging.info(f'ClassifyFst.fst was restored from {far_file}.') + logging.info(f"ClassifyFst.fst was restored from {far_file}.") else: - logging.info(f'Creating ClassifyFst grammars. This might take some time...') + logging.info(f"Creating ClassifyFst grammars. This might take some time...") # TAGGERS cardinal = CardinalFst(deterministic=deterministic) cardinal_graph = cardinal.fst @@ -100,7 +104,9 @@ class ClassifyFst(GraphFst): fraction = FractionFst(deterministic=deterministic, cardinal=cardinal) fraction_graph = fraction.fst - measure = MeasureFst(cardinal=cardinal, decimal=decimal, fraction=fraction, deterministic=deterministic) + measure = MeasureFst( + cardinal=cardinal, decimal=decimal, fraction=fraction, deterministic=deterministic + ) measure_graph = measure.fst date_graph = DateFst(cardinal=cardinal, deterministic=deterministic).fst punctuation = PunctuationFst(deterministic=True) @@ -109,10 +115,16 @@ class ClassifyFst(GraphFst): time_graph = TimeFst(cardinal=cardinal, deterministic=deterministic).fst telephone_graph = TelephoneFst(deterministic=deterministic).fst electronic_graph = ElectronicFst(deterministic=deterministic).fst - money_graph = MoneyFst(cardinal=cardinal, decimal=decimal, deterministic=deterministic).fst - whitelist = WhiteListFst(input_case=input_case, deterministic=deterministic, input_file=whitelist) + money_graph = MoneyFst( + cardinal=cardinal, decimal=decimal, deterministic=deterministic + ).fst + whitelist = WhiteListFst( + input_case=input_case, deterministic=deterministic, input_file=whitelist + ) whitelist_graph = whitelist.graph - serial_graph = SerialFst(cardinal=cardinal, ordinal=deterministic_ordinal, deterministic=deterministic).fst + serial_graph = SerialFst( + cardinal=cardinal, ordinal=deterministic_ordinal, deterministic=deterministic + ).fst # VERBALIZERS cardinal = vCardinal(deterministic=deterministic) @@ -125,7 +137,9 @@ class ClassifyFst(GraphFst): v_fraction_graph = fraction.fst v_telephone_graph = vTelephone(deterministic=deterministic).fst v_electronic_graph = vElectronic(deterministic=deterministic).fst - measure = vMeasure(decimal=decimal, cardinal=cardinal, fraction=fraction, deterministic=deterministic) + measure = vMeasure( + decimal=decimal, cardinal=cardinal, fraction=fraction, deterministic=deterministic + ) v_measure_graph = measure.fst v_time_graph = vTime(deterministic=deterministic).fst v_date_graph = vDate(ordinal=ordinal, deterministic=deterministic).fst @@ -138,7 +152,10 @@ class ClassifyFst(GraphFst): time_final = pynini.compose(time_graph, det_v_time_graph) date_final = pynini.compose(date_graph, det_v_date_graph) range_graph = RangeFst( - time=time_final, date=date_final, cardinal=CardinalFst(deterministic=True), deterministic=deterministic + time=time_final, + date=date_final, + cardinal=CardinalFst(deterministic=True), + deterministic=deterministic, ).fst v_word_graph = vWord(deterministic=deterministic).fst @@ -167,9 +184,13 @@ class ClassifyFst(GraphFst): if not deterministic: roman_graph = RomanFst(deterministic=deterministic).fst # the weight matches the word_graph weight for "I" cases in long sentences with multiple semiotic tokens - classify_and_verbalize |= pynutil.add_weight(pynini.compose(roman_graph, v_roman_graph), word_w) + classify_and_verbalize |= pynutil.add_weight( + pynini.compose(roman_graph, v_roman_graph), word_w + ) - abbreviation_graph = AbbreviationFst(whitelist=whitelist, deterministic=deterministic).fst + abbreviation_graph = AbbreviationFst( + whitelist=whitelist, deterministic=deterministic + ).fst classify_and_verbalize |= pynutil.add_weight( pynini.compose(abbreviation_graph, v_abbreviation), word_w ) @@ -214,4 +235,4 @@ class ClassifyFst(GraphFst): if far_file: generator_main(far_file, {"tokenize_and_classify": self.fst}) - logging.info(f'ClassifyFst grammars are saved to {far_file}.') + logging.info(f"ClassifyFst grammars are saved to {far_file}.") diff --git a/fun_text_processing/text_normalization/en/taggers/whitelist.py b/fun_text_processing/text_normalization/en/taggers/whitelist.py index 568c12fde..d5a8f79ff 100644 --- a/fun_text_processing/text_normalization/en/taggers/whitelist.py +++ b/fun_text_processing/text_normalization/en/taggers/whitelist.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_CHAR, @@ -79,14 +77,18 @@ class WhiteListFst(GraphFst): ) if not deterministic: - multiple_forms_whitelist_graph = get_formats(get_abs_path("data/whitelist/alternatives_all_format.tsv")) + multiple_forms_whitelist_graph = get_formats( + get_abs_path("data/whitelist/alternatives_all_format.tsv") + ) graph |= multiple_forms_whitelist_graph - graph_unit = pynini.string_file(get_abs_path("data/measure/unit.tsv")) | pynini.string_file( - get_abs_path("data/measure/unit_alternatives.tsv") - ) + graph_unit = pynini.string_file( + get_abs_path("data/measure/unit.tsv") + ) | pynini.string_file(get_abs_path("data/measure/unit_alternatives.tsv")) graph_unit_plural = graph_unit @ SINGULAR_TO_PLURAL - units_graph = pynini.compose(DAMO_CHAR ** (3, ...), convert_space(graph_unit | graph_unit_plural)) + units_graph = pynini.compose( + DAMO_CHAR ** (3, ...), convert_space(graph_unit | graph_unit_plural) + ) graph |= units_graph # convert to states only if comma is present before the abbreviation to avoid converting all caps words, @@ -103,7 +105,11 @@ class WhiteListFst(GraphFst): states.extend(additional_options) state_graph = pynini.string_map(states) - graph |= pynini.closure(DAMO_NOT_SPACE, 1) + pynini.union(", ", ",") + pynini.invert(state_graph).optimize() + graph |= ( + pynini.closure(DAMO_NOT_SPACE, 1) + + pynini.union(", ", ",") + + pynini.invert(state_graph).optimize() + ) if input_file: whitelist_provided = _get_whitelist_graph(input_case, input_file) @@ -114,7 +120,7 @@ class WhiteListFst(GraphFst): self.graph = (convert_space(graph)).optimize() - self.fst = (pynutil.insert("name: \"") + self.graph + pynutil.insert("\"")).optimize() + self.fst = (pynutil.insert('name: "') + self.graph + pynutil.insert('"')).optimize() def get_formats(input_f, input_case="cased", is_default=True): @@ -126,13 +132,21 @@ def get_formats(input_f, input_case="cased", is_default=True): for x, y in multiple_formats: if input_case == "lower_cased": x = x.lower() - additional_options.append((f"{x}.", y)) # default "dr" -> doctor, this includes period "dr." -> doctor - additional_options.append((f"{x[0].upper() + x[1:]}", f"{y[0].upper() + y[1:]}")) # "Dr" -> Doctor - additional_options.append((f"{x[0].upper() + x[1:]}.", f"{y[0].upper() + y[1:]}")) # "Dr." -> Doctor + additional_options.append( + (f"{x}.", y) + ) # default "dr" -> doctor, this includes period "dr." -> doctor + additional_options.append( + (f"{x[0].upper() + x[1:]}", f"{y[0].upper() + y[1:]}") + ) # "Dr" -> Doctor + additional_options.append( + (f"{x[0].upper() + x[1:]}.", f"{y[0].upper() + y[1:]}") + ) # "Dr." -> Doctor multiple_formats.extend(additional_options) if not is_default: - multiple_formats = [(x, f"|raw_start|{x}|raw_end||norm_start|{y}|norm_end|") for (x, y) in multiple_formats] + multiple_formats = [ + (x, f"|raw_start|{x}|raw_end||norm_start|{y}|norm_end|") for (x, y) in multiple_formats + ] multiple_formats = pynini.string_map(multiple_formats) return multiple_formats diff --git a/fun_text_processing/text_normalization/en/taggers/word.py b/fun_text_processing/text_normalization/en/taggers/word.py index 9409e1511..6468f1e13 100644 --- a/fun_text_processing/text_normalization/en/taggers/word.py +++ b/fun_text_processing/text_normalization/en/taggers/word.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( MIN_NEG_WEIGHT, @@ -32,7 +30,9 @@ class WordFst(GraphFst): punct = PunctuationFst().graph default_graph = pynini.closure(pynini.difference(DAMO_NOT_SPACE, punct.project("input")), 1) - symbols_to_exclude = (pynini.union("$", "€", "₩", "£", "¥", "#", "%") | DAMO_DIGIT).optimize() + symbols_to_exclude = ( + pynini.union("$", "€", "₩", "£", "¥", "#", "%") | DAMO_DIGIT + ).optimize() graph = pynini.closure(pynini.difference(DAMO_NOT_SPACE, symbols_to_exclude), 1) graph = pynutil.add_weight(graph, MIN_NEG_WEIGHT) | default_graph @@ -57,8 +57,12 @@ class WordFst(GraphFst): ) # allow sentences of words in IPA format separated with spaces or punct marks delim = (punct_marks | pynini.accep(" ")) ** (1, ...) - ipa_phonemes = ipa_phonemes + pynini.closure(delim + ipa_phonemes) + pynini.closure(delim, 0, 1) - ipa_phonemes = (pynini.accep(pynini.escape("[")) + ipa_phonemes + pynini.accep(pynini.escape("]"))).optimize() + ipa_phonemes = ( + ipa_phonemes + pynini.closure(delim + ipa_phonemes) + pynini.closure(delim, 0, 1) + ) + ipa_phonemes = ( + pynini.accep(pynini.escape("[")) + ipa_phonemes + pynini.accep(pynini.escape("]")) + ).optimize() if not deterministic: phoneme = ( @@ -75,4 +79,4 @@ class WordFst(GraphFst): phoneme |= ipa_phonemes self.graph = plurals._priority_union(convert_space(phoneme.optimize()), graph, DAMO_SIGMA) - self.fst = (pynutil.insert("name: \"") + self.graph + pynutil.insert("\"")).optimize() + self.fst = (pynutil.insert('name: "') + self.graph + pynutil.insert('"')).optimize() diff --git a/fun_text_processing/text_normalization/en/utils.py b/fun_text_processing/text_normalization/en/utils.py index 829d7476b..aa3cbdcb9 100644 --- a/fun_text_processing/text_normalization/en/utils.py +++ b/fun_text_processing/text_normalization/en/utils.py @@ -1,5 +1,3 @@ - - import csv import os @@ -10,10 +8,10 @@ def get_abs_path(rel_path): Args: rel_path: relative path to this file - + Returns absolute path """ - return os.path.dirname(os.path.abspath(__file__)) + '/' + rel_path + return os.path.dirname(os.path.abspath(__file__)) + "/" + rel_path def load_labels(abs_path): @@ -32,7 +30,7 @@ def load_labels(abs_path): def augment_labels_with_punct_at_end(labels): """ - augments labels: if key ends on a punctuation that value does not have, add a new label + augments labels: if key ends on a punctuation that value does not have, add a new label where the value maintains the punctuation Args: diff --git a/fun_text_processing/text_normalization/en/verbalizers/abbreviation.py b/fun_text_processing/text_normalization/en/verbalizers/abbreviation.py index 1c5a0cab4..c3812aa33 100644 --- a/fun_text_processing/text_normalization/en/verbalizers/abbreviation.py +++ b/fun_text_processing/text_normalization/en/verbalizers/abbreviation.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst from pynini.lib import pynutil @@ -18,6 +16,6 @@ class AbbreviationFst(GraphFst): def __init__(self, deterministic: bool = True): super().__init__(name="abbreviation", kind="verbalize", deterministic=deterministic) - graph = pynutil.delete("value: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + graph = pynutil.delete('value: "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') delete_tokens = self.delete_tokens(graph) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/text_normalization/en/verbalizers/cardinal.py b/fun_text_processing/text_normalization/en/verbalizers/cardinal.py index 70bf20221..0f0c08ed6 100644 --- a/fun_text_processing/text_normalization/en/verbalizers/cardinal.py +++ b/fun_text_processing/text_normalization/en/verbalizers/cardinal.py @@ -1,7 +1,9 @@ - - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -18,14 +20,14 @@ class CardinalFst(GraphFst): def __init__(self, deterministic: bool = True): super().__init__(name="cardinal", kind="verbalize", deterministic=deterministic) - self.optional_sign = pynini.cross("negative: \"true\"", "minus ") + self.optional_sign = pynini.cross('negative: "true"', "minus ") if not deterministic: - self.optional_sign |= pynini.cross("negative: \"true\"", "negative ") + self.optional_sign |= pynini.cross('negative: "true"', "negative ") self.optional_sign = pynini.closure(self.optional_sign + delete_space, 0, 1) integer = pynini.closure(DAMO_NOT_QUOTE) - self.integer = delete_space + pynutil.delete("\"") + integer + pynutil.delete("\"") + self.integer = delete_space + pynutil.delete('"') + integer + pynutil.delete('"') integer = pynutil.delete("integer:") + self.integer self.numbers = self.optional_sign + integer diff --git a/fun_text_processing/text_normalization/en/verbalizers/date.py b/fun_text_processing/text_normalization/en/verbalizers/date.py index f3ad374cf..9be8f9a21 100644 --- a/fun_text_processing/text_normalization/en/verbalizers/date.py +++ b/fun_text_processing/text_normalization/en/verbalizers/date.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_NOT_QUOTE, @@ -31,26 +29,34 @@ class DateFst(GraphFst): day_cardinal = ( pynutil.delete("day:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) day = day_cardinal @ ordinal.suffix - month = pynutil.delete("month:") + delete_space + pynutil.delete("\"") + month + pynutil.delete("\"") + month = ( + pynutil.delete("month:") + + delete_space + + pynutil.delete('"') + + month + + pynutil.delete('"') + ) year = ( pynutil.delete("year:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') ) # month (day) year graph_mdy = ( - month + pynini.closure(delete_extra_space + day, 0, 1) + pynini.closure(delete_extra_space + year, 0, 1) + month + + pynini.closure(delete_extra_space + day, 0, 1) + + pynini.closure(delete_extra_space + year, 0, 1) ) # may 5 -> may five if not deterministic and not lm: @@ -74,14 +80,19 @@ class DateFst(GraphFst): pynutil.delete("preserve_order:") + delete_space + pynutil.delete("true") + delete_space | pynutil.delete("field_order:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + DAMO_NOT_QUOTE - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space ) final_graph = ( - (plurals._priority_union(graph_mdy, pynutil.add_weight(graph_dmy, 0.0001), DAMO_SIGMA) | year) + ( + plurals._priority_union( + graph_mdy, pynutil.add_weight(graph_dmy, 0.0001), DAMO_SIGMA + ) + | year + ) + delete_space + optional_preserve_order ) diff --git a/fun_text_processing/text_normalization/en/verbalizers/decimal.py b/fun_text_processing/text_normalization/en/verbalizers/decimal.py index 27dd772fb..766424f61 100644 --- a/fun_text_processing/text_normalization/en/verbalizers/decimal.py +++ b/fun_text_processing/text_normalization/en/verbalizers/decimal.py @@ -1,7 +1,10 @@ - - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space, insert_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, + insert_space, +) from pynini.lib import pynutil @@ -17,18 +20,18 @@ class DecimalFst(GraphFst): def __init__(self, cardinal, deterministic: bool = True): super().__init__(name="decimal", kind="verbalize", deterministic=deterministic) - self.optional_sign = pynini.cross("negative: \"true\"", "minus ") + self.optional_sign = pynini.cross('negative: "true"', "minus ") if not deterministic: - self.optional_sign |= pynini.cross("negative: \"true\"", "negative ") + self.optional_sign |= pynini.cross('negative: "true"', "negative ") self.optional_sign = pynini.closure(self.optional_sign + delete_space, 0, 1) self.integer = pynutil.delete("integer_part:") + cardinal.integer self.optional_integer = pynini.closure(self.integer + delete_space + insert_space, 0, 1) self.fractional_default = ( pynutil.delete("fractional_part:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) self.fractional = pynutil.insert("point ") + self.fractional_default @@ -38,9 +41,9 @@ class DecimalFst(GraphFst): + insert_space + pynutil.delete("quantity:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) self.optional_quantity = pynini.closure(self.quantity, 0, 1) diff --git a/fun_text_processing/text_normalization/en/verbalizers/electronic.py b/fun_text_processing/text_normalization/en/verbalizers/electronic.py index 3a30dc2c3..142a9dd26 100644 --- a/fun_text_processing/text_normalization/en/verbalizers/electronic.py +++ b/fun_text_processing/text_normalization/en/verbalizers/electronic.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_NOT_QUOTE, @@ -28,7 +26,9 @@ class ElectronicFst(GraphFst): def __init__(self, deterministic: bool = True): super().__init__(name="electronic", kind="verbalize", deterministic=deterministic) - graph_digit_no_zero = pynini.invert(pynini.string_file(get_abs_path("data/number/digit.tsv"))).optimize() + graph_digit_no_zero = pynini.invert( + pynini.string_file(get_abs_path("data/number/digit.tsv")) + ).optimize() graph_zero = pynini.cross("0", "zero") if not deterministic: @@ -38,7 +38,10 @@ class ElectronicFst(GraphFst): graph_symbols = pynini.string_file(get_abs_path("data/electronic/symbol.tsv")).optimize() default_chars_symbols = pynini.cdrewrite( - pynutil.insert(" ") + (graph_symbols | graph_digit) + pynutil.insert(" "), "", "", DAMO_SIGMA + pynutil.insert(" ") + (graph_symbols | graph_digit) + pynutil.insert(" "), + "", + "", + DAMO_SIGMA, ) default_chars_symbols = pynini.compose( pynini.closure(DAMO_NOT_SPACE), default_chars_symbols.optimize() @@ -47,9 +50,9 @@ class ElectronicFst(GraphFst): user_name = ( pynutil.delete("username:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + default_chars_symbols - + pynutil.delete("\"") + + pynutil.delete('"') ) domain_common = pynini.string_file(get_abs_path("data/electronic/domain.tsv")) @@ -58,22 +61,29 @@ class ElectronicFst(GraphFst): default_chars_symbols + insert_space + plurals._priority_union( - domain_common, pynutil.add_weight(pynini.cross(".", "dot"), weight=0.0001), DAMO_SIGMA + domain_common, + pynutil.add_weight(pynini.cross(".", "dot"), weight=0.0001), + DAMO_SIGMA, ) + pynini.closure( - insert_space + (pynini.cdrewrite(TO_UPPER, "", "", DAMO_SIGMA) @ default_chars_symbols), 0, 1 + insert_space + + (pynini.cdrewrite(TO_UPPER, "", "", DAMO_SIGMA) @ default_chars_symbols), + 0, + 1, ) ) domain = ( pynutil.delete("domain:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + domain + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') ).optimize() - protocol = pynutil.delete("protocol: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + protocol = ( + pynutil.delete('protocol: "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') + ) graph = ( pynini.closure(protocol + delete_space, 0, 1) + pynini.closure(user_name + delete_space + pynutil.insert(" at ") + delete_space, 0, 1) diff --git a/fun_text_processing/text_normalization/en/verbalizers/fraction.py b/fun_text_processing/text_normalization/en/verbalizers/fraction.py index fc5279def..03c4683cd 100644 --- a/fun_text_processing/text_normalization/en/verbalizers/fraction.py +++ b/fun_text_processing/text_normalization/en/verbalizers/fraction.py @@ -1,7 +1,10 @@ - - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, DAMO_SIGMA, GraphFst, insert_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + DAMO_SIGMA, + GraphFst, + insert_space, +) from fun_text_processing.text_normalization.en.verbalizers.ordinal import OrdinalFst from pynini.examples import plurals from pynini.lib import pynutil @@ -22,13 +25,19 @@ class FractionFst(GraphFst): super().__init__(name="fraction", kind="verbalize", deterministic=deterministic) suffix = OrdinalFst().suffix - integer = pynutil.delete("integer_part: \"") + pynini.closure(DAMO_NOT_QUOTE) + pynutil.delete("\" ") - denominator_one = pynini.cross("denominator: \"one\"", "over one") - denominator_half = pynini.cross("denominator: \"two\"", "half") - denominator_quarter = pynini.cross("denominator: \"four\"", "quarter") + integer = ( + pynutil.delete('integer_part: "') + + pynini.closure(DAMO_NOT_QUOTE) + + pynutil.delete('" ') + ) + denominator_one = pynini.cross('denominator: "one"', "over one") + denominator_half = pynini.cross('denominator: "two"', "half") + denominator_quarter = pynini.cross('denominator: "four"', "quarter") denominator_rest = ( - pynutil.delete("denominator: \"") + pynini.closure(DAMO_NOT_QUOTE) @ suffix + pynutil.delete("\"") + pynutil.delete('denominator: "') + + pynini.closure(DAMO_NOT_QUOTE) @ suffix + + pynutil.delete('"') ) denominators = plurals._priority_union( @@ -41,18 +50,24 @@ class FractionFst(GraphFst): DAMO_SIGMA, ).optimize() if not deterministic: - denominators |= pynutil.delete("denominator: \"") + (pynini.accep("four") @ suffix) + pynutil.delete("\"") + denominators |= ( + pynutil.delete('denominator: "') + + (pynini.accep("four") @ suffix) + + pynutil.delete('"') + ) - numerator_one = pynutil.delete("numerator: \"") + pynini.accep("one") + pynutil.delete("\" ") + numerator_one = pynutil.delete('numerator: "') + pynini.accep("one") + pynutil.delete('" ') numerator_one = numerator_one + insert_space + denominators numerator_rest = ( - pynutil.delete("numerator: \"") + pynutil.delete('numerator: "') + (pynini.closure(DAMO_NOT_QUOTE) - pynini.accep("one")) - + pynutil.delete("\" ") + + pynutil.delete('" ') ) numerator_rest = numerator_rest + insert_space + denominators numerator_rest @= pynini.cdrewrite( - plurals._priority_union(pynini.cross("half", "halves"), pynutil.insert("s"), DAMO_SIGMA), + plurals._priority_union( + pynini.cross("half", "halves"), pynutil.insert("s"), DAMO_SIGMA + ), "", "[EOS]", DAMO_SIGMA, @@ -68,7 +83,10 @@ class FractionFst(GraphFst): graph = integer + graph graph @= pynini.cdrewrite( - pynini.cross("and one half", "and a half") | pynini.cross("over ones", "over one"), "", "[EOS]", DAMO_SIGMA + pynini.cross("and one half", "and a half") | pynini.cross("over ones", "over one"), + "", + "[EOS]", + DAMO_SIGMA, ) self.graph = graph diff --git a/fun_text_processing/text_normalization/en/verbalizers/measure.py b/fun_text_processing/text_normalization/en/verbalizers/measure.py index 29f89398b..a4308b6ba 100644 --- a/fun_text_processing/text_normalization/en/verbalizers/measure.py +++ b/fun_text_processing/text_normalization/en/verbalizers/measure.py @@ -1,7 +1,10 @@ - - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space, insert_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, + insert_space, +) from pynini.lib import pynutil @@ -11,7 +14,7 @@ class MeasureFst(GraphFst): measure { negative: "true" cardinal { integer: "twelve" } units: "kilograms" } -> minus twelve kilograms measure { decimal { integer_part: "twelve" fractional_part: "five" } units: "kilograms" } -> twelve point five kilograms tokens { measure { units: "covid" decimal { integer_part: "nineteen" fractional_part: "five" } } } -> covid nineteen point five - + Args: decimal: DecimalFst cardinal: CardinalFst @@ -20,18 +23,20 @@ class MeasureFst(GraphFst): for False multiple transduction are generated (used for audio-based normalization) """ - def __init__(self, decimal: GraphFst, cardinal: GraphFst, fraction: GraphFst, deterministic: bool = True): + def __init__( + self, decimal: GraphFst, cardinal: GraphFst, fraction: GraphFst, deterministic: bool = True + ): super().__init__(name="measure", kind="verbalize", deterministic=deterministic) optional_sign = cardinal.optional_sign unit = ( - pynutil.delete("units: \"") + pynutil.delete('units: "') + pynini.difference(pynini.closure(DAMO_NOT_QUOTE, 1), pynini.union("address", "math")) - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space ) if not deterministic: - unit |= pynini.compose(unit, pynini.cross(pynini.union("inch", "inches"), "\"")) + unit |= pynini.compose(unit, pynini.cross(pynini.union("inch", "inches"), '"')) graph_decimal = ( pynutil.delete("decimal {") @@ -53,17 +58,31 @@ class MeasureFst(GraphFst): ) graph_fraction = ( - pynutil.delete("fraction {") + delete_space + fraction.graph + delete_space + pynutil.delete("}") + pynutil.delete("fraction {") + + delete_space + + fraction.graph + + delete_space + + pynutil.delete("}") ) - graph = (graph_cardinal | graph_decimal | graph_fraction) + delete_space + insert_space + unit + graph = ( + (graph_cardinal | graph_decimal | graph_fraction) + delete_space + insert_space + unit + ) # SH adds "preserve_order: true" by default - preserve_order = pynutil.delete("preserve_order:") + delete_space + pynutil.delete("true") + delete_space - graph |= unit + insert_space + (graph_cardinal | graph_decimal) + delete_space + pynini.closure(preserve_order) + preserve_order = ( + pynutil.delete("preserve_order:") + delete_space + pynutil.delete("true") + delete_space + ) + graph |= ( + unit + + insert_space + + (graph_cardinal | graph_decimal) + + delete_space + + pynini.closure(preserve_order) + ) # for only unit graph |= ( - pynutil.delete("cardinal { integer: \"-\"") + pynutil.delete('cardinal { integer: "-"') + delete_space + pynutil.delete("}") + delete_space @@ -71,14 +90,14 @@ class MeasureFst(GraphFst): + pynini.closure(preserve_order) ) address = ( - pynutil.delete("units: \"address\" ") + pynutil.delete('units: "address" ') + delete_space + graph_cardinal + delete_space + pynini.closure(preserve_order) ) math = ( - pynutil.delete("units: \"math\" ") + pynutil.delete('units: "math" ') + delete_space + graph_cardinal + delete_space diff --git a/fun_text_processing/text_normalization/en/verbalizers/money.py b/fun_text_processing/text_normalization/en/verbalizers/money.py index 66f7fce0e..241fbb78d 100644 --- a/fun_text_processing/text_normalization/en/verbalizers/money.py +++ b/fun_text_processing/text_normalization/en/verbalizers/money.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_NOT_QUOTE, @@ -24,11 +22,21 @@ class MoneyFst(GraphFst): def __init__(self, decimal: GraphFst, deterministic: bool = True): super().__init__(name="money", kind="verbalize", deterministic=deterministic) keep_space = pynini.accep(" ") - maj = pynutil.delete("currency_maj: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") - min = pynutil.delete("currency_min: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + maj = ( + pynutil.delete('currency_maj: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + ) + min = ( + pynutil.delete('currency_min: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + ) fractional_part = ( - pynutil.delete("fractional_part: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + pynutil.delete('fractional_part: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') ) integer_part = decimal.integer @@ -42,7 +50,9 @@ class MoneyFst(GraphFst): if not deterministic: fractional |= pynutil.insert("and ") + fractional - graph_integer_with_minor = integer_part + keep_space + maj + keep_space + fractional + delete_preserve_order + graph_integer_with_minor = ( + integer_part + keep_space + maj + keep_space + fractional + delete_preserve_order + ) # *** point *** currency_maj graph_decimal = decimal.numbers + keep_space + maj diff --git a/fun_text_processing/text_normalization/en/verbalizers/ordinal.py b/fun_text_processing/text_normalization/en/verbalizers/ordinal.py index 2a8c3fb80..e8e6c4a1f 100644 --- a/fun_text_processing/text_normalization/en/verbalizers/ordinal.py +++ b/fun_text_processing/text_normalization/en/verbalizers/ordinal.py @@ -1,8 +1,10 @@ - - - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, DAMO_SIGMA, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + DAMO_SIGMA, + GraphFst, + delete_space, +) from fun_text_processing.text_normalization.en.utils import get_abs_path from pynini.lib import pynutil @@ -26,14 +28,17 @@ class OrdinalFst(GraphFst): graph = ( pynutil.delete("integer:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) convert_rest = pynutil.insert("th") suffix = pynini.cdrewrite( - graph_digit | graph_teens | pynini.cross("ty", "tieth") | convert_rest, "", "[EOS]", DAMO_SIGMA, + graph_digit | graph_teens | pynini.cross("ty", "tieth") | convert_rest, + "", + "[EOS]", + DAMO_SIGMA, ).optimize() self.graph = pynini.compose(graph, suffix) self.suffix = suffix diff --git a/fun_text_processing/text_normalization/en/verbalizers/post_processing.py b/fun_text_processing/text_normalization/en/verbalizers/post_processing.py index e0da98f05..5d604f013 100644 --- a/fun_text_processing/text_normalization/en/verbalizers/post_processing.py +++ b/fun_text_processing/text_normalization/en/verbalizers/post_processing.py @@ -1,6 +1,3 @@ - - - import os import pynini @@ -36,7 +33,7 @@ class PostProcessingFst: far_file = os.path.join(cache_dir, "en_tn_post_processing.far") if not overwrite_cache and far_file and os.path.exists(far_file): self.fst = pynini.Far(far_file, mode="r")["post_process_graph"] - logging.info(f'Post processing graph was restored from {far_file}.') + logging.info(f"Post processing graph was restored from {far_file}.") else: self.set_punct_dict() self.fst = self.get_punct_postprocess_graph() @@ -48,55 +45,55 @@ class PostProcessingFst: self.punct_marks = { "'": [ "'", - '´', - 'ʹ', - 'ʻ', - 'ʼ', - 'ʽ', - 'ʾ', - 'ˈ', - 'ˊ', - 'ˋ', - '˴', - 'ʹ', - '΄', - '՚', - '՝', - 'י', - '׳', - 'ߴ', - 'ߵ', - 'ᑊ', - 'ᛌ', - '᾽', - '᾿', - '`', - '´', - '῾', - '‘', - '’', - '‛', - '′', - '‵', - 'ꞌ', - '＇', - '｀', - '𖽑', - '𖽒', + "´", + "ʹ", + "ʻ", + "ʼ", + "ʽ", + "ʾ", + "ˈ", + "ˊ", + "ˋ", + "˴", + "ʹ", + "΄", + "՚", + "՝", + "י", + "׳", + "ߴ", + "ߵ", + "ᑊ", + "ᛌ", + "᾽", + "᾿", + "`", + "´", + "῾", + "‘", + "’", + "‛", + "′", + "‵", + "ꞌ", + "＇", + "｀", + "𖽑", + "𖽒", ], } def get_punct_postprocess_graph(self): """ - Returns graph to post process punctuation marks. + Returns graph to post process punctuation marks. - {``} quotes are converted to {"}. Note, if there are spaces around single quote {'}, they will be kept. - By default, a space is added after a punctuation mark, and spaces are removed before punctuation marks. + {``} quotes are converted to {"}. Note, if there are spaces around single quote {'}, they will be kept. + By default, a space is added after a punctuation mark, and spaces are removed before punctuation marks. """ punct_marks_all = PunctuationFst().punct_marks # no_space_before_punct assume no space before them - quotes = ["'", "\"", "``", "«"] + quotes = ["'", '"', "``", "«"] dashes = ["-", "—"] brackets = ["<", "{", "("] open_close_single_quotes = [ @@ -105,7 +102,9 @@ class PostProcessingFst: open_close_double_quotes = [('"', '"'), ("``", "``"), ("“", "”")] open_close_symbols = open_close_single_quotes + open_close_double_quotes - allow_space_before_punct = ["&"] + quotes + dashes + brackets + [k[0] for k in open_close_symbols] + allow_space_before_punct = ( + ["&"] + quotes + dashes + brackets + [k[0] for k in open_close_symbols] + ) no_space_before_punct = [m for m in punct_marks_all if m not in allow_space_before_punct] no_space_before_punct = pynini.union(*no_space_before_punct) @@ -119,7 +118,8 @@ class PostProcessingFst: graph = ( pynini.closure(non_punct) + pynini.closure( - no_space_before_punct | pynutil.add_weight(delete_space + no_space_before_punct, MIN_NEG_WEIGHT) + no_space_before_punct + | pynutil.add_weight(delete_space + no_space_before_punct, MIN_NEG_WEIGHT) ) + pynini.closure(non_punct) ) @@ -129,14 +129,21 @@ class PostProcessingFst: ).optimize() # remove space after no_space_after_punct (even if there are no matching closing brackets) - no_space_after_punct = pynini.cdrewrite(delete_space, no_space_after_punct, DAMO_SIGMA, DAMO_SIGMA).optimize() + no_space_after_punct = pynini.cdrewrite( + delete_space, no_space_after_punct, DAMO_SIGMA, DAMO_SIGMA + ).optimize() graph = pynini.compose(graph, no_space_after_punct).optimize() # remove space around text in quotes single_quote = pynutil.add_weight(pynini.accep("`"), MIN_NEG_WEIGHT) double_quotes = pynutil.add_weight(pynini.accep('"'), MIN_NEG_WEIGHT) quotes_graph = ( - single_quote + delete_space_optional + DAMO_ALPHA + DAMO_SIGMA + delete_space_optional + single_quote + single_quote + + delete_space_optional + + DAMO_ALPHA + + DAMO_SIGMA + + delete_space_optional + + single_quote ).optimize() # this is to make sure multiple quotes are tagged from right to left without skipping any quotes in the left diff --git a/fun_text_processing/text_normalization/en/verbalizers/roman.py b/fun_text_processing/text_normalization/en/verbalizers/roman.py index 172c69793..18fec223b 100644 --- a/fun_text_processing/text_normalization/en/verbalizers/roman.py +++ b/fun_text_processing/text_normalization/en/verbalizers/roman.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst from fun_text_processing.text_normalization.en.verbalizers.ordinal import OrdinalFst @@ -24,32 +22,34 @@ class RomanFst(GraphFst): ordinal = pynini.compose(cardinal, suffix) graph = ( - pynutil.delete("key_cardinal: \"") + pynutil.delete('key_cardinal: "') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.accep(" ") - + pynutil.delete("integer: \"") + + pynutil.delete('integer: "') + cardinal - + pynutil.delete("\"") + + pynutil.delete('"') ).optimize() graph |= ( - pynutil.delete("default_cardinal: \"default\" integer: \"") + cardinal + pynutil.delete("\"") + pynutil.delete('default_cardinal: "default" integer: "') + + cardinal + + pynutil.delete('"') ).optimize() graph |= ( - pynutil.delete("default_ordinal: \"default\" integer: \"") + ordinal + pynutil.delete("\"") + pynutil.delete('default_ordinal: "default" integer: "') + ordinal + pynutil.delete('"') ).optimize() graph |= ( - pynutil.delete("key_the_ordinal: \"") + pynutil.delete('key_the_ordinal: "') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.accep(" ") - + pynutil.delete("integer: \"") + + pynutil.delete('integer: "') + pynini.closure(pynutil.insert("the "), 0, 1) + ordinal - + pynutil.delete("\"") + + pynutil.delete('"') ).optimize() delete_tokens = self.delete_tokens(graph) diff --git a/fun_text_processing/text_normalization/en/verbalizers/telephone.py b/fun_text_processing/text_normalization/en/verbalizers/telephone.py index cd7cde637..5c4d0ef68 100644 --- a/fun_text_processing/text_normalization/en/verbalizers/telephone.py +++ b/fun_text_processing/text_normalization/en/verbalizers/telephone.py @@ -1,7 +1,10 @@ - - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space, insert_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, + insert_space, +) from pynini.lib import pynutil @@ -20,9 +23,9 @@ class TelephoneFst(GraphFst): super().__init__(name="telephone", kind="verbalize", deterministic=deterministic) optional_country_code = pynini.closure( - pynutil.delete("country_code: \"") + pynutil.delete('country_code: "') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space + insert_space, 0, @@ -30,18 +33,18 @@ class TelephoneFst(GraphFst): ) number_part = ( - pynutil.delete("number_part: \"") + pynutil.delete('number_part: "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynini.closure(pynutil.add_weight(pynutil.delete(" "), -0.0001), 0, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_extension = pynini.closure( delete_space + insert_space - + pynutil.delete("extension: \"") + + pynutil.delete('extension: "') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\""), + + pynutil.delete('"'), 0, 1, ) diff --git a/fun_text_processing/text_normalization/en/verbalizers/time.py b/fun_text_processing/text_normalization/en/verbalizers/time.py index a52b75d8a..715435eb9 100644 --- a/fun_text_processing/text_normalization/en/verbalizers/time.py +++ b/fun_text_processing/text_normalization/en/verbalizers/time.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_NOT_QUOTE, @@ -27,39 +25,39 @@ class TimeFst(GraphFst): hour = ( pynutil.delete("hours:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) minute = ( pynutil.delete("minutes:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) suffix = ( pynutil.delete("suffix:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_suffix = pynini.closure(delete_space + insert_space + suffix, 0, 1) zone = ( pynutil.delete("zone:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_zone = pynini.closure(delete_space + insert_space + zone, 0, 1) second = ( pynutil.delete("seconds:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) graph_hms = ( hour diff --git a/fun_text_processing/text_normalization/en/verbalizers/verbalize.py b/fun_text_processing/text_normalization/en/verbalizers/verbalize.py index 027668308..0d0fddc20 100644 --- a/fun_text_processing/text_normalization/en/verbalizers/verbalize.py +++ b/fun_text_processing/text_normalization/en/verbalizers/verbalize.py @@ -1,5 +1,3 @@ - - from fun_text_processing.text_normalization.en.graph_utils import GraphFst from fun_text_processing.text_normalization.en.verbalizers.abbreviation import AbbreviationFst from fun_text_processing.text_normalization.en.verbalizers.cardinal import CardinalFst @@ -19,7 +17,7 @@ from fun_text_processing.text_normalization.en.verbalizers.whitelist import Whit class VerbalizeFst(GraphFst): """ Composes other verbalizer grammars. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. Args: @@ -39,7 +37,9 @@ class VerbalizeFst(GraphFst): fraction_graph = fraction.fst telephone_graph = TelephoneFst(deterministic=deterministic).fst electronic_graph = ElectronicFst(deterministic=deterministic).fst - measure = MeasureFst(decimal=decimal, cardinal=cardinal, fraction=fraction, deterministic=deterministic) + measure = MeasureFst( + decimal=decimal, cardinal=cardinal, fraction=fraction, deterministic=deterministic + ) measure_graph = measure.fst time_graph = TimeFst(deterministic=deterministic).fst date_graph = DateFst(ordinal=ordinal, deterministic=deterministic).fst diff --git a/fun_text_processing/text_normalization/en/verbalizers/verbalize_final.py b/fun_text_processing/text_normalization/en/verbalizers/verbalize_final.py index a63fcdf12..14a07220d 100644 --- a/fun_text_processing/text_normalization/en/verbalizers/verbalize_final.py +++ b/fun_text_processing/text_normalization/en/verbalizers/verbalize_final.py @@ -1,5 +1,3 @@ - - import os import pynini @@ -28,16 +26,20 @@ class VerbalizeFinalFst(GraphFst): overwrite_cache: set to True to overwrite .far files """ - def __init__(self, deterministic: bool = True, cache_dir: str = None, overwrite_cache: bool = False): + def __init__( + self, deterministic: bool = True, cache_dir: str = None, overwrite_cache: bool = False + ): super().__init__(name="verbalize_final", kind="verbalize", deterministic=deterministic) far_file = None if cache_dir is not None and cache_dir != "None": os.makedirs(cache_dir, exist_ok=True) - far_file = os.path.join(cache_dir, f"en_tn_{deterministic}_deterministic_verbalizer.far") + far_file = os.path.join( + cache_dir, f"en_tn_{deterministic}_deterministic_verbalizer.far" + ) if not overwrite_cache and far_file and os.path.exists(far_file): self.fst = pynini.Far(far_file, mode="r")["verbalize"] - logging.info(f'VerbalizeFinalFst graph was restored from {far_file}.') + logging.info(f"VerbalizeFinalFst graph was restored from {far_file}.") else: verbalize = VerbalizeFst(deterministic=deterministic).fst word = WordFst(deterministic=deterministic).fst diff --git a/fun_text_processing/text_normalization/en/verbalizers/whitelist.py b/fun_text_processing/text_normalization/en/verbalizers/whitelist.py index 4f0d5f24f..396e86ce2 100644 --- a/fun_text_processing/text_normalization/en/verbalizers/whitelist.py +++ b/fun_text_processing/text_normalization/en/verbalizers/whitelist.py @@ -1,6 +1,10 @@ - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_CHAR, DAMO_SIGMA, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_CHAR, + DAMO_SIGMA, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -19,9 +23,9 @@ class WhiteListFst(GraphFst): graph = ( pynutil.delete("name:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_CHAR - " ", 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) - graph = graph @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + graph = graph @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) self.fst = graph.optimize() diff --git a/fun_text_processing/text_normalization/en/verbalizers/word.py b/fun_text_processing/text_normalization/en/verbalizers/word.py index e835b2fe9..c0744a9bc 100644 --- a/fun_text_processing/text_normalization/en/verbalizers/word.py +++ b/fun_text_processing/text_normalization/en/verbalizers/word.py @@ -1,6 +1,10 @@ - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_CHAR, DAMO_SIGMA, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_CHAR, + DAMO_SIGMA, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -17,7 +21,13 @@ class WordFst(GraphFst): def __init__(self, deterministic: bool = True): super().__init__(name="word", kind="verbalize", deterministic=deterministic) chars = pynini.closure(DAMO_CHAR - " ", 1) - char = pynutil.delete("name:") + delete_space + pynutil.delete("\"") + chars + pynutil.delete("\"") - graph = char @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", DAMO_SIGMA) + char = ( + pynutil.delete("name:") + + delete_space + + pynutil.delete('"') + + chars + + pynutil.delete('"') + ) + graph = char @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", DAMO_SIGMA) self.fst = graph.optimize() diff --git a/fun_text_processing/text_normalization/es/__init__.py b/fun_text_processing/text_normalization/es/__init__.py index 0b7d10cb0..749ab5396 100644 --- a/fun_text_processing/text_normalization/es/__init__.py +++ b/fun_text_processing/text_normalization/es/__init__.py @@ -1,4 +1 @@ - - - LOCALIZATION = "eu" # Set to am for alternate formatting diff --git a/fun_text_processing/text_normalization/es/data/__init__.py b/fun_text_processing/text_normalization/es/data/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/es/data/__init__.py +++ b/fun_text_processing/text_normalization/es/data/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/es/data/dates/__init__.py b/fun_text_processing/text_normalization/es/data/dates/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/es/data/dates/__init__.py +++ b/fun_text_processing/text_normalization/es/data/dates/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/es/data/electronic/__init__.py b/fun_text_processing/text_normalization/es/data/electronic/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/es/data/electronic/__init__.py +++ b/fun_text_processing/text_normalization/es/data/electronic/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/es/data/fractions/__init__.py b/fun_text_processing/text_normalization/es/data/fractions/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/es/data/fractions/__init__.py +++ b/fun_text_processing/text_normalization/es/data/fractions/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/es/data/measures/__init__.py b/fun_text_processing/text_normalization/es/data/measures/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/es/data/measures/__init__.py +++ b/fun_text_processing/text_normalization/es/data/measures/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/es/data/money/__init__.py b/fun_text_processing/text_normalization/es/data/money/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/es/data/money/__init__.py +++ b/fun_text_processing/text_normalization/es/data/money/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/es/data/numbers/__init__.py b/fun_text_processing/text_normalization/es/data/numbers/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/es/data/numbers/__init__.py +++ b/fun_text_processing/text_normalization/es/data/numbers/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/es/data/ordinals/__init__.py b/fun_text_processing/text_normalization/es/data/ordinals/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/es/data/ordinals/__init__.py +++ b/fun_text_processing/text_normalization/es/data/ordinals/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/es/data/roman/__init__.py b/fun_text_processing/text_normalization/es/data/roman/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/es/data/roman/__init__.py +++ b/fun_text_processing/text_normalization/es/data/roman/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/es/data/time/__init__.py b/fun_text_processing/text_normalization/es/data/time/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/es/data/time/__init__.py +++ b/fun_text_processing/text_normalization/es/data/time/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/es/graph_utils.py b/fun_text_processing/text_normalization/es/graph_utils.py index 733ee0c52..2fc5f8fa2 100644 --- a/fun_text_processing/text_normalization/es/graph_utils.py +++ b/fun_text_processing/text_normalization/es/graph_utils.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_SIGMA, DAMO_SPACE from fun_text_processing.text_normalization.es import LOCALIZATION @@ -22,12 +20,16 @@ else: decimal_separator = pynini.accep(",") ones = pynini.union("un", "ún") -fem_ones = pynini.union(pynini.cross("un", "una"), pynini.cross("ún", "una"), pynini.cross("uno", "una")) -one_to_one_hundred = pynini.union(digits, "uno", tens, teens, twenties, tens + pynini.accep(" y ") + digits) +fem_ones = pynini.union( + pynini.cross("un", "una"), pynini.cross("ún", "una"), pynini.cross("uno", "una") +) +one_to_one_hundred = pynini.union( + digits, "uno", tens, teens, twenties, tens + pynini.accep(" y ") + digits +) fem_hundreds = hundreds @ pynini.cdrewrite(pynini.cross("ientos", "ientas"), "", "", DAMO_SIGMA) -def strip_accent(fst: 'pynini.FstLike') -> 'pynini.FstLike': +def strip_accent(fst: "pynini.FstLike") -> "pynini.FstLike": """ Converts all accented vowels to non-accented equivalents @@ -37,7 +39,7 @@ def strip_accent(fst: 'pynini.FstLike') -> 'pynini.FstLike': return fst @ pynini.cdrewrite(accents, "", "", DAMO_SIGMA) -def shift_cardinal_gender(fst: 'pynini.FstLike') -> 'pynini.FstLike': +def shift_cardinal_gender(fst: "pynini.FstLike") -> "pynini.FstLike": """ Applies gender conversion rules to a cardinal string. These include: rendering all masculine forms of "uno" (including apocopated forms) as "una" and Converting all gendered numbers in the hundreds series (200,300,400...) to feminine equivalent (e.g. "doscientos" -> "doscientas"). Conversion only applies @@ -59,23 +61,27 @@ def shift_cardinal_gender(fst: 'pynini.FstLike') -> 'pynini.FstLike': + (pynini.accep("mil") | pynini.accep("milésimo")) + pynini.closure(DAMO_SPACE + hundreds, 0, 1) + pynini.closure(DAMO_SPACE + one_to_one_hundred, 0, 1) - + pynini.union(pynini.accep("[EOS]"), pynini.accep("\""), decimal_separator) + + pynini.union(pynini.accep("[EOS]"), pynini.accep('"'), decimal_separator) ) before_double_digits = pynini.closure(DAMO_SPACE + one_to_one_hundred, 0, 1) + pynini.union( - pynini.accep("[EOS]"), pynini.accep("\"") + pynini.accep("[EOS]"), pynini.accep('"') ) - fem_allign = pynini.cdrewrite(fem_hundreds, "", before_mil, DAMO_SIGMA) # doscientas mil dosciento - fem_allign @= pynini.cdrewrite(fem_hundreds, "", before_double_digits, DAMO_SIGMA) # doscientas mil doscienta + fem_allign = pynini.cdrewrite( + fem_hundreds, "", before_mil, DAMO_SIGMA + ) # doscientas mil dosciento + fem_allign @= pynini.cdrewrite( + fem_hundreds, "", before_double_digits, DAMO_SIGMA + ) # doscientas mil doscienta fem_allign @= pynini.cdrewrite( - fem_ones, "", pynini.union("[EOS]", "\"", decimal_separator), DAMO_SIGMA + fem_ones, "", pynini.union("[EOS]", '"', decimal_separator), DAMO_SIGMA ) # If before a quote or EOS, we know it's the end of a string return fst @ fem_allign -def shift_number_gender(fst: 'pynini.FstLike') -> 'pynini.FstLike': +def shift_number_gender(fst: "pynini.FstLike") -> "pynini.FstLike": """ Performs gender conversion on all verbalized numbers in output. All values in the hundreds series (200,300,400) are changed to feminine gender (e.g. "doscientos" -> "doscientas") and all forms of "uno" (including apocopated forms) are converted to "una". @@ -90,13 +96,13 @@ def shift_number_gender(fst: 'pynini.FstLike') -> 'pynini.FstLike': """ fem_allign = pynini.cdrewrite(fem_hundreds, "", "", DAMO_SIGMA) fem_allign @= pynini.cdrewrite( - fem_ones, "", pynini.union(DAMO_SPACE, pynini.accep("[EOS]"), pynini.accep("\"")), DAMO_SIGMA + fem_ones, "", pynini.union(DAMO_SPACE, pynini.accep("[EOS]"), pynini.accep('"')), DAMO_SIGMA ) # If before a quote or EOS, we know it's the end of a string return fst @ fem_allign -def strip_cardinal_apocope(fst: 'pynini.FstLike') -> 'pynini.FstLike': +def strip_cardinal_apocope(fst: "pynini.FstLike") -> "pynini.FstLike": """ Reverts apocope on cardinal strings in line with formation rules. e.g. "un" -> "uno". Due to cardinal formation rules, this in effect only affects strings where the final value is a variation of "un". @@ -109,11 +115,11 @@ def strip_cardinal_apocope(fst: 'pynini.FstLike') -> 'pynini.FstLike': """ # Since cardinals use apocope by default for large values (e.g. "millón"), this only needs to act on the last instance of one strip = pynini.cross("un", "uno") | pynini.cross("ún", "uno") - strip = pynini.cdrewrite(strip, "", pynini.union("[EOS]", "\""), DAMO_SIGMA) + strip = pynini.cdrewrite(strip, "", pynini.union("[EOS]", '"'), DAMO_SIGMA) return fst @ strip -def add_cardinal_apocope_fem(fst: 'pynini.FstLike') -> 'pynini.FstLike': +def add_cardinal_apocope_fem(fst: "pynini.FstLike") -> "pynini.FstLike": """ Adds apocope on cardinal strings in line with stressing rules. e.g. "una" -> "un". This only occurs when "una" precedes a stressed "a" sound in formal speech. This is not predictable with text string, so is included for non-deterministic cases. @@ -126,11 +132,11 @@ def add_cardinal_apocope_fem(fst: 'pynini.FstLike') -> 'pynini.FstLike': """ # Since the stress trigger follows the cardinal string and only affects the preceding sound, this only needs to act on the last instance of one strip = pynini.cross("una", "un") | pynini.cross("veintiuna", "veintiún") - strip = pynini.cdrewrite(strip, "", pynini.union("[EOS]", "\""), DAMO_SIGMA) + strip = pynini.cdrewrite(strip, "", pynini.union("[EOS]", '"'), DAMO_SIGMA) return fst @ strip -def roman_to_int(fst: 'pynini.FstLike') -> 'pynini.FstLike': +def roman_to_int(fst: "pynini.FstLike") -> "pynini.FstLike": """ Alters given fst to convert Roman integers (lower and upper cased) into Arabic numerals. Valid for values up to 1000. e.g. diff --git a/fun_text_processing/text_normalization/es/taggers/__init__.py b/fun_text_processing/text_normalization/es/taggers/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/es/taggers/__init__.py +++ b/fun_text_processing/text_normalization/es/taggers/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/es/taggers/cardinal.py b/fun_text_processing/text_normalization/es/taggers/cardinal.py index df18ff408..a5b79be91 100644 --- a/fun_text_processing/text_normalization/es/taggers/cardinal.py +++ b/fun_text_processing/text_normalization/es/taggers/cardinal.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_ALPHA, @@ -23,7 +21,7 @@ twenties = pynini.invert(pynini.string_file(get_abs_path("data/numbers/twenties. hundreds = pynini.invert(pynini.string_file(get_abs_path("data/numbers/hundreds.tsv"))) -def filter_punctuation(fst: 'pynini.FstLike') -> 'pynini.FstLike': +def filter_punctuation(fst: "pynini.FstLike") -> "pynini.FstLike": """ Helper function for parsing number strings. Converts common cardinal strings (groups of three digits delineated by 'cardinal_separator' - see graph_utils) and converts to a string of digits: @@ -35,7 +33,7 @@ def filter_punctuation(fst: 'pynini.FstLike') -> 'pynini.FstLike': Returns: fst: A pynini.FstLike object """ - exactly_three_digits = DAMO_DIGIT ** 3 # for blocks of three + exactly_three_digits = DAMO_DIGIT**3 # for blocks of three up_to_three_digits = pynini.closure(DAMO_DIGIT, 1, 3) # for start of string cardinal_string = pynini.closure( @@ -72,7 +70,7 @@ class CardinalFst(GraphFst): # Any double digit graph_tens = teen - graph_tens |= ties + (pynutil.delete('0') | (pynutil.insert(" y ") + graph_digit)) + graph_tens |= ties + (pynutil.delete("0") | (pynutil.insert(" y ") + graph_digit)) graph_tens |= twenties self.tens = graph_tens.optimize() @@ -83,11 +81,15 @@ class CardinalFst(GraphFst): # Three digit strings graph_hundreds = hundreds + pynini.union( - pynutil.delete("00"), (insert_space + graph_tens), (pynini.cross("0", DAMO_SPACE) + graph_digit) + pynutil.delete("00"), + (insert_space + graph_tens), + (pynini.cross("0", DAMO_SPACE) + graph_digit), ) graph_hundreds |= pynini.cross("100", "cien") graph_hundreds |= ( - pynini.cross("1", "ciento") + insert_space + pynini.union(graph_tens, pynutil.delete("0") + graph_digit) + pynini.cross("1", "ciento") + + insert_space + + pynini.union(graph_tens, pynutil.delete("0") + graph_digit) ) self.hundreds = graph_hundreds.optimize() @@ -106,32 +108,53 @@ class CardinalFst(GraphFst): pynutil.delete("000") + graph_hundreds_component_at_least_one_none_zero_digit, graph_hundreds_component_at_least_one_none_zero_digit_no_one + pynutil.insert(" mil") - + ((insert_space + graph_hundreds_component_at_least_one_none_zero_digit) | pynutil.delete("000")), + + ( + (insert_space + graph_hundreds_component_at_least_one_none_zero_digit) + | pynutil.delete("000") + ), pynini.cross("001", "mil") - + ((insert_space + graph_hundreds_component_at_least_one_none_zero_digit) | pynutil.delete("000")), + + ( + (insert_space + graph_hundreds_component_at_least_one_none_zero_digit) + | pynutil.delete("000") + ), ) graph_thousands_component_at_least_one_none_zero_digit_no_one = pynini.union( pynutil.delete("000") + graph_hundreds_component_at_least_one_none_zero_digit_no_one, graph_hundreds_component_at_least_one_none_zero_digit_no_one + pynutil.insert(" mil") - + ((insert_space + graph_hundreds_component_at_least_one_none_zero_digit) | pynutil.delete("000")), + + ( + (insert_space + graph_hundreds_component_at_least_one_none_zero_digit) + | pynutil.delete("000") + ), pynini.cross("001", "mil") - + ((insert_space + graph_hundreds_component_at_least_one_none_zero_digit) | pynutil.delete("000")), + + ( + (insert_space + graph_hundreds_component_at_least_one_none_zero_digit) + | pynutil.delete("000") + ), ) graph_million = pynutil.add_weight(pynini.cross("000001", "un millón"), -0.001) - graph_million |= graph_thousands_component_at_least_one_none_zero_digit_no_one + pynutil.insert(" millones") + graph_million |= ( + graph_thousands_component_at_least_one_none_zero_digit_no_one + + pynutil.insert(" millones") + ) graph_million |= pynutil.delete("000000") graph_million += insert_space graph_billion = pynutil.add_weight(pynini.cross("000001", "un billón"), -0.001) - graph_billion |= graph_thousands_component_at_least_one_none_zero_digit_no_one + pynutil.insert(" billones") + graph_billion |= ( + graph_thousands_component_at_least_one_none_zero_digit_no_one + + pynutil.insert(" billones") + ) graph_billion |= pynutil.delete("000000") graph_billion += insert_space graph_trillion = pynutil.add_weight(pynini.cross("000001", "un trillón"), -0.001) - graph_trillion |= graph_thousands_component_at_least_one_none_zero_digit_no_one + pynutil.insert(" trillones") + graph_trillion |= ( + graph_thousands_component_at_least_one_none_zero_digit_no_one + + pynutil.insert(" trillones") + ) graph_trillion |= pynutil.delete("000000") graph_trillion += insert_space @@ -145,21 +168,28 @@ class CardinalFst(GraphFst): self.graph = ( ((DAMO_DIGIT - "0") + pynini.closure(DAMO_DIGIT, 0)) @ pynini.cdrewrite(pynini.closure(pynutil.insert("0")), "[BOS]", "", DAMO_SIGMA) - @ DAMO_DIGIT ** 24 + @ DAMO_DIGIT**24 @ graph @ pynini.cdrewrite(delete_space, "[BOS]", "", DAMO_SIGMA) @ pynini.cdrewrite(delete_space, "", "[EOS]", DAMO_SIGMA) @ pynini.cdrewrite( - pynini.cross(pynini.closure(DAMO_WHITE_SPACE, 2), DAMO_SPACE), DAMO_ALPHA, DAMO_ALPHA, DAMO_SIGMA + pynini.cross(pynini.closure(DAMO_WHITE_SPACE, 2), DAMO_SPACE), + DAMO_ALPHA, + DAMO_ALPHA, + DAMO_SIGMA, ) ) self.graph |= zero self.graph = filter_punctuation(self.graph).optimize() - optional_minus_graph = pynini.closure(pynutil.insert("negative: ") + pynini.cross("-", "\"true\" "), 0, 1) + optional_minus_graph = pynini.closure( + pynutil.insert("negative: ") + pynini.cross("-", '"true" '), 0, 1 + ) - final_graph = optional_minus_graph + pynutil.insert("integer: \"") + self.graph + pynutil.insert("\"") + final_graph = ( + optional_minus_graph + pynutil.insert('integer: "') + self.graph + pynutil.insert('"') + ) final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/text_normalization/es/taggers/date.py b/fun_text_processing/text_normalization/es/taggers/date.py index df09d838b..0b82d4a25 100644 --- a/fun_text_processing/text_normalization/es/taggers/date.py +++ b/fun_text_processing/text_normalization/es/taggers/date.py @@ -1,7 +1,10 @@ - - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_DIGIT, DAMO_SPACE, GraphFst, delete_extra_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_DIGIT, + DAMO_SPACE, + GraphFst, + delete_extra_space, +) from fun_text_processing.text_normalization.es.utils import get_abs_path from pynini.lib import pynutil @@ -34,22 +37,27 @@ class DateFst(GraphFst): # 01, 31, 1 digit_day = optional_leading_zero @ pynini.union(*[str(x) for x in range(1, 32)]) @ numbers - day = (pynutil.insert("day: \"") + digit_day + pynutil.insert("\"")).optimize() + day = (pynutil.insert('day: "') + digit_day + pynutil.insert('"')).optimize() digit_month = optional_leading_zero @ pynini.union(*[str(x) for x in range(1, 13)]) number_to_month = digit_month @ number_to_month - month_name = (pynutil.insert("month: \"") + month_graph + pynutil.insert("\"")).optimize() - month_number = (pynutil.insert("month: \"") + number_to_month + pynutil.insert("\"")).optimize() + month_name = (pynutil.insert('month: "') + month_graph + pynutil.insert('"')).optimize() + month_number = ( + pynutil.insert('month: "') + number_to_month + pynutil.insert('"') + ).optimize() # prefer cardinal over year year = (DAMO_DIGIT - "0") + pynini.closure(DAMO_DIGIT, 1, 3) # 90, 990, 1990 year @= numbers self.year = year - year_only = pynutil.insert("year: \"") + year + pynutil.insert("\"") + year_only = pynutil.insert('year: "') + year + pynutil.insert('"') year_with_articles = ( - pynutil.insert("year: \"") + pynini.closure(articles + DAMO_SPACE, 0, 1) + year + pynutil.insert("\"") + pynutil.insert('year: "') + + pynini.closure(articles + DAMO_SPACE, 0, 1) + + year + + pynutil.insert('"') ) graph_dmy = ( @@ -61,7 +69,10 @@ class DateFst(GraphFst): ) graph_mdy = ( # English influences on language - month_name + delete_extra_space + day + pynini.closure(DAMO_SPACE + year_with_articles, 0, 1) + month_name + + delete_extra_space + + day + + pynini.closure(DAMO_SPACE + year_with_articles, 0, 1) ) separators = [".", "-", "/"] @@ -75,7 +86,9 @@ class DateFst(GraphFst): dash = "-" day_optional = pynini.closure(pynini.cross(dash, DAMO_SPACE) + day, 0, 1) - graph_ymd = DAMO_DIGIT ** 4 @ year_only + pynini.cross(dash, DAMO_SPACE) + month_number + day_optional + graph_ymd = ( + DAMO_DIGIT**4 @ year_only + pynini.cross(dash, DAMO_SPACE) + month_number + day_optional + ) final_graph = graph_dmy + pynutil.insert(" preserve_order: true") final_graph |= graph_ymd diff --git a/fun_text_processing/text_normalization/es/taggers/decimals.py b/fun_text_processing/text_normalization/es/taggers/decimals.py index 0032c2733..a68e0ccf9 100644 --- a/fun_text_processing/text_normalization/es/taggers/decimals.py +++ b/fun_text_processing/text_normalization/es/taggers/decimals.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_DIGIT, @@ -22,7 +20,9 @@ digit = pynini.invert(pynini.string_file(get_abs_path("data/numbers/digit.tsv")) zero = pynini.invert(pynini.string_file(get_abs_path("data/numbers/zero.tsv"))) -def get_quantity(decimal_graph: 'pynini.FstLike', cardinal_graph: 'pynini.FstLike') -> 'pynini.FstLike': +def get_quantity( + decimal_graph: "pynini.FstLike", cardinal_graph: "pynini.FstLike" +) -> "pynini.FstLike": """ Returns FST that transforms either a cardinal or decimal followed by a quantity into a numeral, e.g. 2 millones -> integer_part: "dos" quantity: "millones" @@ -37,15 +37,21 @@ def get_quantity(decimal_graph: 'pynini.FstLike', cardinal_graph: 'pynini.FstLik numbers = pynini.cdrewrite(pynutil.delete(cardinal_separator), "", "", DAMO_SIGMA) @ numbers res = ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + numbers # The cardinal we're passing only produces 'un' for one, so gender agreement is safe (all quantities are masculine). Limit to 10^6 power. - + pynutil.insert("\"") + + pynutil.insert('"') + DAMO_SPACE - + pynutil.insert("quantity: \"") + + pynutil.insert('quantity: "') + quantities - + pynutil.insert("\"") + + pynutil.insert('"') + ) + res |= ( + decimal_graph + + DAMO_SPACE + + pynutil.insert('quantity: "') + + quantities + + pynutil.insert('"') ) - res |= decimal_graph + DAMO_SPACE + pynutil.insert("quantity: \"") + quantities + pynutil.insert("\"") return res @@ -92,9 +98,13 @@ class DecimalFst(GraphFst): self.graph = graph.optimize() graph_separator = pynutil.delete(decimal_separator) - optional_graph_negative = pynini.closure(pynutil.insert("negative: ") + pynini.cross("-", "\"true\" "), 0, 1) + optional_graph_negative = pynini.closure( + pynutil.insert("negative: ") + pynini.cross("-", '"true" '), 0, 1 + ) - self.graph_fractional = pynutil.insert("fractional_part: \"") + self.graph + pynutil.insert("\"") + self.graph_fractional = ( + pynutil.insert('fractional_part: "') + self.graph + pynutil.insert('"') + ) # Integer graph maintains apocope except for ones place graph_integer = ( @@ -102,8 +112,10 @@ class DecimalFst(GraphFst): if deterministic else pynini.union(cardinal.graph, strip_cardinal_apocope(cardinal.graph)) ) # Gives us forms w/ and w/o apocope - self.graph_integer = pynutil.insert("integer_part: \"") + graph_integer + pynutil.insert("\"") - final_graph_wo_sign = self.graph_integer + graph_separator + insert_space + self.graph_fractional + self.graph_integer = pynutil.insert('integer_part: "') + graph_integer + pynutil.insert('"') + final_graph_wo_sign = ( + self.graph_integer + graph_separator + insert_space + self.graph_fractional + ) self.final_graph_wo_negative = ( final_graph_wo_sign | get_quantity(final_graph_wo_sign, cardinal.graph).optimize() diff --git a/fun_text_processing/text_normalization/es/taggers/electronic.py b/fun_text_processing/text_normalization/es/taggers/electronic.py index 13853cc4c..37c8983fd 100644 --- a/fun_text_processing/text_normalization/es/taggers/electronic.py +++ b/fun_text_processing/text_normalization/es/taggers/electronic.py @@ -1,7 +1,10 @@ - - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_ALPHA, DAMO_DIGIT, GraphFst, insert_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_ALPHA, + DAMO_DIGIT, + GraphFst, + insert_space, +) from fun_text_processing.text_normalization.es.utils import get_abs_path, load_labels from pynini.lib import pynutil @@ -27,23 +30,27 @@ class ElectronicFst(GraphFst): accepted_common_domains = pynini.union(*common_domains) accepted_symbols = pynini.union(*symbols) - dot accepted_characters = pynini.closure(DAMO_ALPHA | DAMO_DIGIT | accepted_symbols) - acceepted_characters_with_dot = pynini.closure(DAMO_ALPHA | DAMO_DIGIT | accepted_symbols | dot) + acceepted_characters_with_dot = pynini.closure( + DAMO_ALPHA | DAMO_DIGIT | accepted_symbols | dot + ) # email username = ( - pynutil.insert("username: \"") + pynutil.insert('username: "') + acceepted_characters_with_dot - + pynutil.insert("\"") - + pynini.cross('@', ' ') + + pynutil.insert('"') + + pynini.cross("@", " ") ) domain_graph = accepted_characters + dot + accepted_characters - domain_graph = pynutil.insert("domain: \"") + domain_graph + pynutil.insert("\"") + domain_graph = pynutil.insert('domain: "') + domain_graph + pynutil.insert('"') domain_common_graph = ( - pynutil.insert("domain: \"") + pynutil.insert('domain: "') + accepted_characters + accepted_common_domains - + pynini.closure((accepted_symbols | dot) + pynini.closure(accepted_characters, 1), 0, 1) - + pynutil.insert("\"") + + pynini.closure( + (accepted_symbols | dot) + pynini.closure(accepted_characters, 1), 0, 1 + ) + + pynutil.insert('"') ) graph = (username + domain_graph) | domain_common_graph @@ -55,7 +62,7 @@ class ElectronicFst(GraphFst): else pynini.accep("www.") | pynini.cross("www.", "doble ve doble ve doble ve.") ) protocol = protocol_start | protocol_end | (protocol_start + protocol_end) - protocol = pynutil.insert("protocol: \"") + protocol + pynutil.insert("\"") + protocol = pynutil.insert('protocol: "') + protocol + pynutil.insert('"') graph |= protocol + insert_space + (domain_graph | domain_common_graph) self.graph = graph diff --git a/fun_text_processing/text_normalization/es/taggers/fraction.py b/fun_text_processing/text_normalization/es/taggers/fraction.py index 6ec552641..125a79ceb 100644 --- a/fun_text_processing/text_normalization/es/taggers/fraction.py +++ b/fun_text_processing/text_normalization/es/taggers/fraction.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_CHAR, @@ -34,15 +32,50 @@ class FractionFst(GraphFst): ordinal_graph = ordinal.graph # 2-10 are all ordinals - three_to_ten = pynini.string_map(["2", "3", "4", "5", "6", "7", "8", "9", "10",]) + three_to_ten = pynini.string_map( + [ + "2", + "3", + "4", + "5", + "6", + "7", + "8", + "9", + "10", + ] + ) block_three_to_ten = pynutil.delete(three_to_ten) # To block cardinal productions if not deterministic: # Multiples of tens are sometimes rendered as ordinals - three_to_ten |= pynini.string_map(["20", "30", "40", "50", "60", "70", "80", "90",]) + three_to_ten |= pynini.string_map( + [ + "20", + "30", + "40", + "50", + "60", + "70", + "80", + "90", + ] + ) graph_three_to_ten = three_to_ten @ ordinal_graph graph_three_to_ten @= pynini.cdrewrite(ordinal_exceptions, "", "", DAMO_SIGMA) # Higher powers of tens (and multiples) are converted to ordinals. - hundreds = pynini.string_map(["100", "200", "300", "400", "500", "600", "700", "800", "900",]) + hundreds = pynini.string_map( + [ + "100", + "200", + "300", + "400", + "500", + "600", + "700", + "800", + "900", + ] + ) graph_hundreds = hundreds @ ordinal_graph multiples_of_thousand = ordinal.multiples_of_thousand # So we can have X milésimos @@ -55,29 +88,37 @@ class FractionFst(GraphFst): graph_higher_powers_of_ten += higher_powers_of_ten graph_higher_powers_of_ten = cardinal_graph @ graph_higher_powers_of_ten graph_higher_powers_of_ten @= pynini.cdrewrite( - pynutil.delete("un "), pynini.accep("[BOS]"), pynini.project(higher_powers_of_ten, "output"), DAMO_SIGMA + pynutil.delete("un "), + pynini.accep("[BOS]"), + pynini.project(higher_powers_of_ten, "output"), + DAMO_SIGMA, ) # we drop 'un' from these ordinals (millionths, not one-millionths) - graph_higher_powers_of_ten = multiples_of_thousand | graph_hundreds | graph_higher_powers_of_ten + graph_higher_powers_of_ten = ( + multiples_of_thousand | graph_hundreds | graph_higher_powers_of_ten + ) block_higher_powers_of_ten = pynutil.delete( pynini.project(graph_higher_powers_of_ten, "input") ) # For cardinal graph graph_fractions_ordinals = graph_higher_powers_of_ten | graph_three_to_ten graph_fractions_ordinals += pynutil.insert( - "\" morphosyntactic_features: \"ordinal\"" + '" morphosyntactic_features: "ordinal"' ) # We note the root for processing later # Blocking the digits and hundreds from Cardinal graph graph_fractions_cardinals = pynini.cdrewrite( - block_three_to_ten | block_higher_powers_of_ten, pynini.accep("[BOS]"), pynini.accep("[EOS]"), DAMO_SIGMA + block_three_to_ten | block_higher_powers_of_ten, + pynini.accep("[BOS]"), + pynini.accep("[EOS]"), + DAMO_SIGMA, ) graph_fractions_cardinals @= DAMO_CHAR.plus @ pynini.cdrewrite( pynutil.delete("0"), pynini.accep("[BOS]"), pynini.accep("[EOS]"), DAMO_SIGMA ) # Empty characters become '0' for DAMO_CHAR fst, so need to block graph_fractions_cardinals @= cardinal_graph graph_fractions_cardinals += pynutil.insert( - "\" morphosyntactic_features: \"add_root\"" + '" morphosyntactic_features: "add_root"' ) # blocking these entries to reduce erroneous possibilities in debugging if deterministic: @@ -88,14 +129,18 @@ class FractionFst(GraphFst): graph_denominator = pynini.union( graph_fractions_ordinals, graph_fractions_cardinals, - pynutil.add_weight(cardinal_graph + pynutil.insert("\""), 0.001), + pynutil.add_weight(cardinal_graph + pynutil.insert('"'), 0.001), ) # Last form is simply recording the cardinal. Weighting so last resort - integer = pynutil.insert("integer_part: \"") + cardinal_graph + pynutil.insert("\"") + DAMO_SPACE - numerator = ( - pynutil.insert("numerator: \"") + cardinal_graph + (pynini.cross("/", "\" ") | pynini.cross(" / ", "\" ")) + integer = ( + pynutil.insert('integer_part: "') + cardinal_graph + pynutil.insert('"') + DAMO_SPACE ) - denominator = pynutil.insert("denominator: \"") + graph_denominator + numerator = ( + pynutil.insert('numerator: "') + + cardinal_graph + + (pynini.cross("/", '" ') | pynini.cross(" / ", '" ')) + ) + denominator = pynutil.insert('denominator: "') + graph_denominator self.graph = pynini.closure(integer, 0, 1) + numerator + denominator diff --git a/fun_text_processing/text_normalization/es/taggers/measure.py b/fun_text_processing/text_normalization/es/taggers/measure.py index bb69667df..d1a3003ea 100644 --- a/fun_text_processing/text_normalization/es/taggers/measure.py +++ b/fun_text_processing/text_normalization/es/taggers/measure.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_ALPHA, @@ -41,7 +39,9 @@ class MeasureFst(GraphFst): for False multiple transduction are generated (used for audio-based normalization) """ - def __init__(self, cardinal: GraphFst, decimal: GraphFst, fraction: GraphFst, deterministic: bool = True): + def __init__( + self, cardinal: GraphFst, decimal: GraphFst, fraction: GraphFst, deterministic: bool = True + ): super().__init__(name="measure", kind="classify", deterministic=deterministic) cardinal_graph = cardinal.graph @@ -58,22 +58,26 @@ class MeasureFst(GraphFst): ) optional_unit_denominator = pynini.closure( - pynutil.insert(DAMO_NON_BREAKING_SPACE) + graph_unit_denominator, 0, 1, + pynutil.insert(DAMO_NON_BREAKING_SPACE) + graph_unit_denominator, + 0, + 1, ) unit_plural = ( - pynutil.insert("units: \"") + pynutil.insert('units: "') + ((graph_unit_plural + optional_unit_denominator) | graph_unit_denominator) - + pynutil.insert("\"") + + pynutil.insert('"') ) unit_singular_graph = ( - pynutil.insert("units: \"") + pynutil.insert('units: "') + ((graph_unit_singular + optional_unit_denominator) | graph_unit_denominator) - + pynutil.insert("\"") + + pynutil.insert('"') ) - subgraph_decimal = decimal.fst + insert_space + pynini.closure(DAMO_SPACE, 0, 1) + unit_plural + subgraph_decimal = ( + decimal.fst + insert_space + pynini.closure(DAMO_SPACE, 0, 1) + unit_plural + ) subgraph_cardinal = ( (optional_graph_negative + (DAMO_SIGMA - "1")) @ cardinal.fst @@ -89,57 +93,59 @@ class MeasureFst(GraphFst): + unit_singular_graph ) - subgraph_fraction = fraction.fst + insert_space + pynini.closure(delete_space, 0, 1) + unit_singular_graph + subgraph_fraction = ( + fraction.fst + insert_space + pynini.closure(delete_space, 0, 1) + unit_singular_graph + ) decimal_times = ( pynutil.insert("decimal { ") + decimal.final_graph_wo_negative - + pynutil.insert(" } units: \"") - + pynini.union('x', 'X') - + pynutil.insert("\"") + + pynutil.insert(' } units: "') + + pynini.union("x", "X") + + pynutil.insert('"') ) cardinal_times = ( - pynutil.insert("cardinal { integer: \"") + pynutil.insert('cardinal { integer: "') + strip_cardinal_apocope(cardinal_graph) - + pynutil.insert("\" } units: \"") - + pynini.union('x', 'X') - + pynutil.insert("\"") + + pynutil.insert('" } units: "') + + pynini.union("x", "X") + + pynutil.insert('"') ) cardinal_dash_alpha = ( - pynutil.insert("cardinal { integer: \"") + pynutil.insert('cardinal { integer: "') + strip_cardinal_apocope(cardinal_graph) - + pynutil.delete('-') - + pynutil.insert("\" } units: \"") + + pynutil.delete("-") + + pynutil.insert('" } units: "') + pynini.closure(DAMO_ALPHA, 1) - + pynutil.insert("\"") + + pynutil.insert('"') ) decimal_dash_alpha = ( pynutil.insert("decimal { ") + decimal.final_graph_wo_negative - + pynutil.delete('-') - + pynutil.insert(" } units: \"") + + pynutil.delete("-") + + pynutil.insert(' } units: "') + pynini.closure(DAMO_ALPHA, 1) - + pynutil.insert("\"") + + pynutil.insert('"') ) alpha_dash_cardinal = ( - pynutil.insert("units: \"") + pynutil.insert('units: "') + pynini.closure(DAMO_ALPHA, 1) - + pynutil.delete('-') - + pynutil.insert("\"") - + pynutil.insert(" cardinal { integer: \"") + + pynutil.delete("-") + + pynutil.insert('"') + + pynutil.insert(' cardinal { integer: "') + cardinal_graph - + pynutil.insert("\" } preserve_order: true") + + pynutil.insert('" } preserve_order: true') ) alpha_dash_decimal = ( - pynutil.insert("units: \"") + pynutil.insert('units: "') + pynini.closure(DAMO_ALPHA, 1) - + pynutil.delete('-') - + pynutil.insert("\"") + + pynutil.delete("-") + + pynutil.insert('"') + pynutil.insert(" decimal { ") + decimal.final_graph_wo_negative + pynutil.insert(" } preserve_order: true") diff --git a/fun_text_processing/text_normalization/es/taggers/money.py b/fun_text_processing/text_normalization/es/taggers/money.py index a7f54dc64..9d78d867c 100644 --- a/fun_text_processing/text_normalization/es/taggers/money.py +++ b/fun_text_processing/text_normalization/es/taggers/money.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_ALPHA, @@ -50,40 +48,71 @@ class MoneyFst(GraphFst): maj_plural_graph = maj_singular @ (fem_plural | masc_plural) min_plural_graph = min_singular @ (fem_plural | masc_plural) - graph_maj_singular = pynutil.insert("currency_maj: \"") + maj_singular_graph + pynutil.insert("\"") - graph_maj_plural = pynutil.insert("currency_maj: \"") + maj_plural_graph + pynutil.insert("\"") + graph_maj_singular = ( + pynutil.insert('currency_maj: "') + maj_singular_graph + pynutil.insert('"') + ) + graph_maj_plural = ( + pynutil.insert('currency_maj: "') + maj_plural_graph + pynutil.insert('"') + ) - graph_integer_one = pynutil.insert("integer_part: \"") + pynini.cross("1", "un") + pynutil.insert("\"") + graph_integer_one = ( + pynutil.insert('integer_part: "') + pynini.cross("1", "un") + pynutil.insert('"') + ) decimal_with_quantity = (DAMO_SIGMA + DAMO_ALPHA) @ graph_decimal_final graph_decimal_plural = pynini.union( - graph_maj_plural + pynini.closure(delete_space, 0, 1) + insert_space + graph_decimal_final, # $1,05 - graph_decimal_final + pynini.closure(delete_space, 0, 1) + insert_space + graph_maj_plural, # 1,05 $ + graph_maj_plural + + pynini.closure(delete_space, 0, 1) + + insert_space + + graph_decimal_final, # $1,05 + graph_decimal_final + + pynini.closure(delete_space, 0, 1) + + insert_space + + graph_maj_plural, # 1,05 $ ) graph_decimal_plural = ( (DAMO_SIGMA - "1") + decimal_separator + DAMO_SIGMA ) @ graph_decimal_plural # Can't have "un euros" graph_decimal_singular = pynini.union( - graph_maj_singular + pynini.closure(delete_space, 0, 1) + insert_space + graph_decimal_final, # $1,05 - graph_decimal_final + pynini.closure(delete_space, 0, 1) + insert_space + graph_maj_singular, # 1,05 $ + graph_maj_singular + + pynini.closure(delete_space, 0, 1) + + insert_space + + graph_decimal_final, # $1,05 + graph_decimal_final + + pynini.closure(delete_space, 0, 1) + + insert_space + + graph_maj_singular, # 1,05 $ ) - graph_decimal_singular = (pynini.accep("1") + decimal_separator + DAMO_SIGMA) @ graph_decimal_singular + graph_decimal_singular = ( + pynini.accep("1") + decimal_separator + DAMO_SIGMA + ) @ graph_decimal_singular graph_decimal = pynini.union( graph_decimal_singular, graph_decimal_plural, - graph_maj_plural + pynini.closure(delete_space, 0, 1) + insert_space + decimal_with_quantity, + graph_maj_plural + + pynini.closure(delete_space, 0, 1) + + insert_space + + decimal_with_quantity, ) graph_integer = ( - pynutil.insert("integer_part: \"") + ((DAMO_SIGMA - "1") @ cardinal_graph) + pynutil.insert("\"") + pynutil.insert('integer_part: "') + + ((DAMO_SIGMA - "1") @ cardinal_graph) + + pynutil.insert('"') ) graph_integer_only = pynini.union( - graph_maj_singular + pynini.closure(delete_space, 0, 1) + insert_space + graph_integer_one, - graph_integer_one + pynini.closure(delete_space, 0, 1) + insert_space + graph_maj_singular, + graph_maj_singular + + pynini.closure(delete_space, 0, 1) + + insert_space + + graph_integer_one, + graph_integer_one + + pynini.closure(delete_space, 0, 1) + + insert_space + + graph_maj_singular, ) graph_integer_only |= pynini.union( graph_maj_plural + pynini.closure(delete_space, 0, 1) + insert_space + graph_integer, @@ -103,8 +132,12 @@ class MoneyFst(GraphFst): | ((DAMO_DIGIT - "0") + DAMO_DIGIT) ) - graph_min_singular = pynutil.insert("currency_min: \"") + min_singular_graph + pynutil.insert("\"") - graph_min_plural = pynutil.insert("currency_min: \"") + min_plural_graph + pynutil.insert("\"") + graph_min_singular = ( + pynutil.insert('currency_min: "') + min_singular_graph + pynutil.insert('"') + ) + graph_min_plural = ( + pynutil.insert('currency_min: "') + min_plural_graph + pynutil.insert('"') + ) # format ** euro ** cent decimal_graph_with_minor = None @@ -119,19 +152,25 @@ class MoneyFst(GraphFst): integer_plus_maj = (pynini.closure(DAMO_DIGIT) - "0") @ integer_plus_maj graph_fractional_one = ( - pynutil.insert("fractional_part: \"") + pynutil.insert('fractional_part: "') + two_digits_fractional_part @ pynini.cross("1", "un") - + pynutil.insert("\"") + + pynutil.insert('"') ) graph_fractional = ( - two_digits_fractional_part @ (pynini.closure(DAMO_DIGIT, 1, 2) - "1") @ cardinal.two_digit_non_zero + two_digits_fractional_part + @ (pynini.closure(DAMO_DIGIT, 1, 2) - "1") + @ cardinal.two_digit_non_zero + ) + graph_fractional = ( + pynutil.insert('fractional_part: "') + graph_fractional + pynutil.insert('"') ) - graph_fractional = pynutil.insert("fractional_part: \"") + graph_fractional + pynutil.insert("\"") fractional_plus_min = pynini.union( graph_fractional + insert_space + pynutil.insert(curr_symbol) @ graph_min_plural, - graph_fractional_one + insert_space + pynutil.insert(curr_symbol) @ graph_min_singular, + graph_fractional_one + + insert_space + + pynutil.insert(curr_symbol) @ graph_min_singular, ) decimal_graph_with_minor_curr = ( @@ -140,9 +179,9 @@ class MoneyFst(GraphFst): decimal_graph_with_minor_curr |= pynutil.add_weight( integer_plus_maj + pynini.cross(decimal_separator, DAMO_SPACE) - + pynutil.insert("fractional_part: \"") + + pynutil.insert('fractional_part: "') + two_digits_fractional_part @ cardinal.two_digit_non_zero - + pynutil.insert("\""), + + pynutil.insert('"'), weight=0.0001, ) diff --git a/fun_text_processing/text_normalization/es/taggers/ordinal.py b/fun_text_processing/text_normalization/es/taggers/ordinal.py index eecfe33d9..a4e0920b0 100644 --- a/fun_text_processing/text_normalization/es/taggers/ordinal.py +++ b/fun_text_processing/text_normalization/es/taggers/ordinal.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_CHAR, @@ -19,7 +17,7 @@ ties = pynini.invert(pynini.string_file(get_abs_path("data/ordinals/ties.tsv"))) hundreds = pynini.invert(pynini.string_file(get_abs_path("data/ordinals/hundreds.tsv"))) -def get_one_to_one_thousand(cardinal: 'pynini.FstLike') -> 'pynini.FstLike': +def get_one_to_one_thousand(cardinal: "pynini.FstLike") -> "pynini.FstLike": """ Produces an acceptor for verbalizations of all numbers from 1 to 1000. Needed for ordinals and fractions. @@ -36,7 +34,7 @@ def get_one_to_one_thousand(cardinal: 'pynini.FstLike') -> 'pynini.FstLike': class OrdinalFst(GraphFst): """ Finite state transducer for classifying ordinal - "21.º" -> ordinal { integer: "vigésimo primero" morphosyntactic_features: "gender_masc" } + "21.º" -> ordinal { integer: "vigésimo primero" morphosyntactic_features: "gender_masc" } This class converts ordinal up to the millionth (millonésimo) order (exclusive). This FST also records the ending of the ordinal (called "morphosyntactic_features"): @@ -75,7 +73,8 @@ class OrdinalFst(GraphFst): ) graph_hundred_component = pynini.union( - graph_hundreds + pynini.closure(DAMO_SPACE + pynini.union(graph_tens_component, graph_digit), 0, 1), + graph_hundreds + + pynini.closure(DAMO_SPACE + pynini.union(graph_tens_component, graph_digit), 0, 1), graph_tens_component, graph_digit, ) @@ -88,7 +87,9 @@ class OrdinalFst(GraphFst): graph_thousands = ( strip_accent(self.one_to_one_thousand) + DAMO_SPACE + thousands ) # Cardinals become prefix for thousands series. Snce accent on the powers of ten we strip accent from leading words - graph_thousands @= pynini.cdrewrite(delete_space, "", "milésimo", DAMO_SIGMA) # merge as a prefix + graph_thousands @= pynini.cdrewrite( + delete_space, "", "milésimo", DAMO_SIGMA + ) # merge as a prefix graph_thousands |= thousands self.multiples_of_thousand = (cardinal_graph @ graph_thousands).optimize() @@ -96,7 +97,9 @@ class OrdinalFst(GraphFst): if ( not deterministic ): # Formally the words preceding the power of ten should be a prefix, but some maintain word boundaries. - graph_thousands |= (self.one_to_one_thousand @ graph_hundred_component) + DAMO_SPACE + thousands + graph_thousands |= ( + (self.one_to_one_thousand @ graph_hundred_component) + DAMO_SPACE + thousands + ) graph_thousands += pynini.closure(DAMO_SPACE + graph_hundred_component, 0, 1) @@ -118,14 +121,18 @@ class OrdinalFst(GraphFst): fem = pynini.accep("gender_fem") apocope = pynini.accep("apocope") - delete_period = pynini.closure(pynutil.delete("."), 0, 1) # Sometimes the period is omitted f + delete_period = pynini.closure( + pynutil.delete("."), 0, 1 + ) # Sometimes the period is omitted f accept_masc = delete_period + pynini.cross("º", masc) accep_fem = delete_period + pynini.cross("ª", fem) accep_apocope = delete_period + pynini.cross("ᵉʳ", apocope) # Managing Romanization - graph_roman = pynutil.insert("integer: \"") + roman_to_int(ordinal_graph) + pynutil.insert("\"") + graph_roman = ( + pynutil.insert('integer: "') + roman_to_int(ordinal_graph) + pynutil.insert('"') + ) if not deterministic: # Introduce plural plural = pynini.closure(pynutil.insert("/plural"), 0, 1) @@ -136,27 +143,29 @@ class OrdinalFst(GraphFst): insert_morphology = pynutil.insert(pynini.union(masc, fem)) + plural insert_morphology |= pynutil.insert(apocope) insert_morphology = ( - pynutil.insert(" morphosyntactic_features: \"") + insert_morphology + pynutil.insert("\"") + pynutil.insert(' morphosyntactic_features: "') + + insert_morphology + + pynutil.insert('"') ) graph_roman += insert_morphology else: # We insert both genders as default - graph_roman += pynutil.insert(" morphosyntactic_features: \"gender_masc\"") | pynutil.insert( - " morphosyntactic_features: \"gender_fem\"" - ) + graph_roman += pynutil.insert( + ' morphosyntactic_features: "gender_masc"' + ) | pynutil.insert(' morphosyntactic_features: "gender_fem"') # Rest of graph convert_abbreviation = accept_masc | accep_fem | accep_apocope graph = ( - pynutil.insert("integer: \"") + pynutil.insert('integer: "') + ordinal_graph - + pynutil.insert("\"") - + pynutil.insert(" morphosyntactic_features: \"") + + pynutil.insert('"') + + pynutil.insert(' morphosyntactic_features: "') + convert_abbreviation - + pynutil.insert("\"") + + pynutil.insert('"') ) graph = pynini.union(graph, graph_roman) diff --git a/fun_text_processing/text_normalization/es/taggers/telephone.py b/fun_text_processing/text_normalization/es/taggers/telephone.py index 77a46d470..9c4ca5b64 100644 --- a/fun_text_processing/text_normalization/es/taggers/telephone.py +++ b/fun_text_processing/text_normalization/es/taggers/telephone.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_SIGMA, GraphFst, insert_space from fun_text_processing.text_normalization.es.graph_utils import ones @@ -25,8 +23,8 @@ class TelephoneFst(GraphFst): (we ignore more complicated cases such as "doscientos y dos" or "tres nueves"). Args: - deterministic: if True will provide a single transduction option, - for False multiple transduction are generated (used for audio-based normalization) + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) """ def __init__(self, deterministic: bool = True): @@ -126,7 +124,7 @@ class TelephoneFst(GraphFst): number_part = pynini.union(ten_digit_graph, nine_digit_graph, eight_digit_graph) number_part @= pynini.cdrewrite(pynini.cross(ones, "uno"), "", "", DAMO_SIGMA) - number_part = pynutil.insert("number_part: \"") + number_part + pynutil.insert("\"") + number_part = pynutil.insert('number_part: "') + number_part + pynutil.insert('"') graph = number_part final_graph = self.add_tokens(graph) diff --git a/fun_text_processing/text_normalization/es/taggers/time.py b/fun_text_processing/text_normalization/es/taggers/time.py index d0b644176..f8f77e9ff 100644 --- a/fun_text_processing/text_normalization/es/taggers/time.py +++ b/fun_text_processing/text_normalization/es/taggers/time.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_DIGIT, @@ -38,7 +36,7 @@ class TimeFst(GraphFst): change_one = pynini.cdrewrite(one, "", "", DAMO_SIGMA) cardinal_graph = cardinal.graph @ change_one - day_suffix = pynutil.insert("suffix: \"") + suffix + pynutil.insert("\"") + day_suffix = pynutil.insert('suffix: "') + suffix + pynutil.insert('"') day_suffix = delete_space + insert_space + day_suffix delete_hora_suffix = delete_space + insert_space + pynutil.delete("h") @@ -57,32 +55,46 @@ class TimeFst(GraphFst): ) graph_24 = ( - pynini.closure(DAMO_DIGIT, 1, 2) @ delete_leading_zero_to_double_digit @ pynini.union(*labels_hour_24) + pynini.closure(DAMO_DIGIT, 1, 2) + @ delete_leading_zero_to_double_digit + @ pynini.union(*labels_hour_24) ) graph_12 = ( - pynini.closure(DAMO_DIGIT, 1, 2) @ delete_leading_zero_to_double_digit @ pynini.union(*labels_hour_12) + pynini.closure(DAMO_DIGIT, 1, 2) + @ delete_leading_zero_to_double_digit + @ pynini.union(*labels_hour_12) ) graph_hour_24 = graph_24 @ cardinal_graph graph_hour_12 = graph_12 @ cardinal_graph - graph_minute_single = delete_leading_zero_to_double_digit @ pynini.union(*labels_minute_single) + graph_minute_single = delete_leading_zero_to_double_digit @ pynini.union( + *labels_minute_single + ) graph_minute_double = pynini.union(*labels_minute_double) graph_minute = pynini.union(graph_minute_single, graph_minute_double) @ cardinal_graph final_graph_hour_only_24 = ( - pynutil.insert("hours: \"") + graph_hour_24 + pynutil.insert("\"") + delete_hora_suffix + pynutil.insert('hours: "') + graph_hour_24 + pynutil.insert('"') + delete_hora_suffix + ) + final_graph_hour_only_12 = ( + pynutil.insert('hours: "') + graph_hour_12 + pynutil.insert('"') + day_suffix ) - final_graph_hour_only_12 = pynutil.insert("hours: \"") + graph_hour_12 + pynutil.insert("\"") + day_suffix - final_graph_hour_24 = pynutil.insert("hours: \"") + graph_hour_24 + pynutil.insert("\"") - final_graph_hour_12 = pynutil.insert("hours: \"") + graph_hour_12 + pynutil.insert("\"") + final_graph_hour_24 = pynutil.insert('hours: "') + graph_hour_24 + pynutil.insert('"') + final_graph_hour_12 = pynutil.insert('hours: "') + graph_hour_12 + pynutil.insert('"') - final_graph_minute = pynutil.insert("minutes: \"") + graph_minute + pynutil.insert("\"") - final_graph_second = pynutil.insert("seconds: \"") + graph_minute + pynutil.insert("\"") + final_graph_minute = pynutil.insert('minutes: "') + graph_minute + pynutil.insert('"') + final_graph_second = pynutil.insert('seconds: "') + graph_minute + pynutil.insert('"') final_time_zone_optional = pynini.closure( - delete_space + insert_space + pynutil.insert("zone: \"") + time_zone_graph + pynutil.insert("\""), 0, 1, + delete_space + + insert_space + + pynutil.insert('zone: "') + + time_zone_graph + + pynutil.insert('"'), + 0, + 1, ) # 02.30 h @@ -91,7 +103,8 @@ class TimeFst(GraphFst): + delete_time_delimiter + (pynutil.delete("00") | (insert_space + final_graph_minute)) + pynini.closure( - delete_time_delimiter + (pynini.cross("00", " seconds: \"0\"") | (insert_space + final_graph_second)), + delete_time_delimiter + + (pynini.cross("00", ' seconds: "0"') | (insert_space + final_graph_second)), 0, 1, ) # For seconds 2.30.35 h @@ -108,7 +121,7 @@ class TimeFst(GraphFst): + delete_minute_suffix + pynini.closure( delete_space - + (pynini.cross("00", " seconds: \"0\"") | (insert_space + final_graph_second)) + + (pynini.cross("00", ' seconds: "0"') | (insert_space + final_graph_second)) + delete_second_suffix, 0, 1, @@ -122,7 +135,8 @@ class TimeFst(GraphFst): + delete_time_delimiter + (pynutil.delete("00") | (insert_space + final_graph_minute)) + pynini.closure( - delete_time_delimiter + (pynini.cross("00", " seconds: \"0\"") | (insert_space + final_graph_second)), + delete_time_delimiter + + (pynini.cross("00", ' seconds: "0"') | (insert_space + final_graph_second)), 0, 1, ) # For seconds 2.30.35 a. m. @@ -131,27 +145,42 @@ class TimeFst(GraphFst): ) graph_h = ( - pynini.union(final_graph_hour_only_24, final_graph_hour_only_12) + final_time_zone_optional + pynini.union(final_graph_hour_only_24, final_graph_hour_only_12) + + final_time_zone_optional ) # Should always have a time indicator, else we'll pass to cardinals if not deterministic: # This includes alternate vocalization (hour menos min, min para hour), here we shift the times and indicate a `style` tag - hour_shift_24 = pynini.invert(pynini.string_file(get_abs_path("data/time/hour_to_24.tsv"))) - hour_shift_12 = pynini.invert(pynini.string_file(get_abs_path("data/time/hour_to_12.tsv"))) + hour_shift_24 = pynini.invert( + pynini.string_file(get_abs_path("data/time/hour_to_24.tsv")) + ) + hour_shift_12 = pynini.invert( + pynini.string_file(get_abs_path("data/time/hour_to_12.tsv")) + ) minute_shift = pynini.string_file(get_abs_path("data/time/minute_to.tsv")) graph_hour_to_24 = graph_24 @ hour_shift_24 @ cardinal_graph graph_hour_to_12 = graph_12 @ hour_shift_12 @ cardinal_graph - graph_minute_to = pynini.union(graph_minute_single, graph_minute_double) @ minute_shift @ cardinal_graph + graph_minute_to = ( + pynini.union(graph_minute_single, graph_minute_double) + @ minute_shift + @ cardinal_graph + ) - final_graph_hour_to_24 = pynutil.insert("hours: \"") + graph_hour_to_24 + pynutil.insert("\"") - final_graph_hour_to_12 = pynutil.insert("hours: \"") + graph_hour_to_12 + pynutil.insert("\"") + final_graph_hour_to_24 = ( + pynutil.insert('hours: "') + graph_hour_to_24 + pynutil.insert('"') + ) + final_graph_hour_to_12 = ( + pynutil.insert('hours: "') + graph_hour_to_12 + pynutil.insert('"') + ) - final_graph_minute_to = pynutil.insert("minutes: \"") + graph_minute_to + pynutil.insert("\"") + final_graph_minute_to = ( + pynutil.insert('minutes: "') + graph_minute_to + pynutil.insert('"') + ) - graph_menos = pynutil.insert(" style: \"1\"") - graph_para = pynutil.insert(" style: \"2\"") + graph_menos = pynutil.insert(' style: "1"') + graph_para = pynutil.insert(' style: "2"') final_graph_style = graph_menos | graph_para diff --git a/fun_text_processing/text_normalization/es/taggers/tokenize_and_classify.py b/fun_text_processing/text_normalization/es/taggers/tokenize_and_classify.py index cd4c36174..fae83e913 100644 --- a/fun_text_processing/text_normalization/es/taggers/tokenize_and_classify.py +++ b/fun_text_processing/text_normalization/es/taggers/tokenize_and_classify.py @@ -1,6 +1,3 @@ - - - import os import pynini @@ -75,10 +72,15 @@ class ClassifyFst(GraphFst): self.decimal = DecimalFst(cardinal=self.cardinal, deterministic=deterministic) decimal_graph = self.decimal.fst - self.fraction = FractionFst(cardinal=self.cardinal, ordinal=self.ordinal, deterministic=deterministic) + self.fraction = FractionFst( + cardinal=self.cardinal, ordinal=self.ordinal, deterministic=deterministic + ) fraction_graph = self.fraction.fst self.measure = MeasureFst( - cardinal=self.cardinal, decimal=self.decimal, fraction=self.fraction, deterministic=deterministic + cardinal=self.cardinal, + decimal=self.decimal, + fraction=self.fraction, + deterministic=deterministic, ) measure_graph = self.measure.fst self.date = DateFst(cardinal=self.cardinal, deterministic=deterministic) @@ -90,9 +92,13 @@ class ClassifyFst(GraphFst): telephone_graph = self.telephone.fst self.electronic = ElectronicFst(deterministic=deterministic) electronic_graph = self.electronic.fst - self.money = MoneyFst(cardinal=self.cardinal, decimal=self.decimal, deterministic=deterministic) + self.money = MoneyFst( + cardinal=self.cardinal, decimal=self.decimal, deterministic=deterministic + ) money_graph = self.money.fst - self.whitelist = WhiteListFst(input_case=input_case, deterministic=deterministic, input_file=whitelist) + self.whitelist = WhiteListFst( + input_case=input_case, deterministic=deterministic, input_file=whitelist + ) whitelist_graph = self.whitelist.fst punct_graph = PunctuationFst(deterministic=deterministic).fst @@ -110,7 +116,11 @@ class ClassifyFst(GraphFst): | pynutil.add_weight(electronic_graph, 1.1) | pynutil.add_weight(word_graph, 200) ) - punct = pynutil.insert("tokens { ") + pynutil.add_weight(punct_graph, weight=2.1) + pynutil.insert(" }") + punct = ( + pynutil.insert("tokens { ") + + pynutil.add_weight(punct_graph, weight=2.1) + + pynutil.insert(" }") + ) punct = pynini.closure( pynini.compose(pynini.closure(DAMO_WHITE_SPACE, 1), delete_extra_space) | (pynutil.insert(" ") + punct), @@ -118,7 +128,9 @@ class ClassifyFst(GraphFst): ) token = pynutil.insert("tokens { ") + classify + pynutil.insert(" }") token_plus_punct = ( - pynini.closure(punct + pynutil.insert(" ")) + token + pynini.closure(pynutil.insert(" ") + punct) + pynini.closure(punct + pynutil.insert(" ")) + + token + + pynini.closure(pynutil.insert(" ") + punct) ) graph = token_plus_punct + pynini.closure( diff --git a/fun_text_processing/text_normalization/es/taggers/whitelist.py b/fun_text_processing/text_normalization/es/taggers/whitelist.py index 11ff0ff37..66638f14e 100644 --- a/fun_text_processing/text_normalization/es/taggers/whitelist.py +++ b/fun_text_processing/text_normalization/es/taggers/whitelist.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import GraphFst, convert_space from fun_text_processing.text_normalization.es.utils import get_abs_path, load_labels @@ -32,7 +30,8 @@ class WhiteListFst(GraphFst): graph = _get_whitelist_graph(input_case, get_abs_path("data/whitelist.tsv")) if not deterministic and input_case != "lower_cased": graph |= pynutil.add_weight( - _get_whitelist_graph("lower_cased", get_abs_path("data/whitelist.tsv")), weight=0.0001 + _get_whitelist_graph("lower_cased", get_abs_path("data/whitelist.tsv")), + weight=0.0001, ) if input_file: @@ -43,9 +42,11 @@ class WhiteListFst(GraphFst): graph = whitelist_provided if not deterministic: - units_graph = _get_whitelist_graph(input_case, file=get_abs_path("data/measures/measurements.tsv")) + units_graph = _get_whitelist_graph( + input_case, file=get_abs_path("data/measures/measurements.tsv") + ) graph |= units_graph self.graph = graph self.final_graph = convert_space(self.graph).optimize() - self.fst = (pynutil.insert("name: \"") + self.final_graph + pynutil.insert("\"")).optimize() + self.fst = (pynutil.insert('name: "') + self.final_graph + pynutil.insert('"')).optimize() diff --git a/fun_text_processing/text_normalization/es/taggers/word.py b/fun_text_processing/text_normalization/es/taggers/word.py index b2ba79c4b..0b0a04780 100644 --- a/fun_text_processing/text_normalization/es/taggers/word.py +++ b/fun_text_processing/text_normalization/es/taggers/word.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_SPACE, GraphFst from pynini.lib import pynutil @@ -17,5 +15,5 @@ class WordFst(GraphFst): def __init__(self, deterministic: bool = True): super().__init__(name="word", kind="classify") - word = pynutil.insert("name: \"") + pynini.closure(DAMO_NOT_SPACE, 1) + pynutil.insert("\"") + word = pynutil.insert('name: "') + pynini.closure(DAMO_NOT_SPACE, 1) + pynutil.insert('"') self.fst = word.optimize() diff --git a/fun_text_processing/text_normalization/es/utils.py b/fun_text_processing/text_normalization/es/utils.py index 29875f79b..10af0c66f 100644 --- a/fun_text_processing/text_normalization/es/utils.py +++ b/fun_text_processing/text_normalization/es/utils.py @@ -1,6 +1,3 @@ - - - import csv import os @@ -14,7 +11,7 @@ def get_abs_path(rel_path): Returns absolute path """ - return os.path.dirname(os.path.abspath(__file__)) + '/' + rel_path + return os.path.dirname(os.path.abspath(__file__)) + "/" + rel_path def load_labels(abs_path): diff --git a/fun_text_processing/text_normalization/es/verbalizers/__init__.py b/fun_text_processing/text_normalization/es/verbalizers/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/es/verbalizers/__init__.py +++ b/fun_text_processing/text_normalization/es/verbalizers/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/es/verbalizers/cardinal.py b/fun_text_processing/text_normalization/es/verbalizers/cardinal.py index cdb2540dd..4a12d5eb7 100644 --- a/fun_text_processing/text_normalization/es/verbalizers/cardinal.py +++ b/fun_text_processing/text_normalization/es/verbalizers/cardinal.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst from fun_text_processing.text_normalization.es.graph_utils import ( @@ -12,21 +10,21 @@ from pynini.lib import pynutil class CardinalFst(GraphFst): """ - Finite state transducer for verbalizing cardinals - e.g. cardinal { integer: "dos" } -> "dos" + Finite state transducer for verbalizing cardinals + e.g. cardinal { integer: "dos" } -> "dos" - Args: - deterministic: if True will provide a single transduction option, - for False multiple transduction are generated (used for audio-based normalization) - """ + Args: + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ def __init__(self, deterministic: bool = True): super().__init__(name="cardinal", kind="verbalize", deterministic=deterministic) - optional_sign = pynini.closure(pynini.cross("negative: \"true\" ", "menos "), 0, 1) + optional_sign = pynini.closure(pynini.cross('negative: "true" ', "menos "), 0, 1) self.optional_sign = optional_sign integer = pynini.closure(DAMO_NOT_QUOTE, 1) - self.integer = pynutil.delete(" \"") + integer + pynutil.delete("\"") + self.integer = pynutil.delete(' "') + integer + pynutil.delete('"') integer = pynutil.delete("integer:") + self.integer diff --git a/fun_text_processing/text_normalization/es/verbalizers/date.py b/fun_text_processing/text_normalization/es/verbalizers/date.py index 005803082..b4750ddb8 100644 --- a/fun_text_processing/text_normalization/es/verbalizers/date.py +++ b/fun_text_processing/text_normalization/es/verbalizers/date.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_NOT_QUOTE, @@ -27,7 +25,9 @@ class DateFst(GraphFst): def __init__(self, deterministic: bool = True): super().__init__(name="date", kind="verbalize", deterministic=deterministic) - day_cardinal = pynutil.delete("day: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + day_cardinal = ( + pynutil.delete('day: "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') + ) day = strip_cardinal_apocope(day_cardinal) primero = pynini.cdrewrite(pynini.cross("uno", "primero"), "[BOS]", "[EOS]", DAMO_SIGMA) @@ -35,31 +35,38 @@ class DateFst(GraphFst): (day @ primero) if deterministic else pynini.union(day, day @ primero) ) # Primero for first day is traditional, but will vary depending on region - month = pynutil.delete("month: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + month = pynutil.delete('month: "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') year = ( - pynutil.delete("year: \"") + pynutil.delete('year: "') + articles + DAMO_SPACE + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) # Insert preposition if wasn't originally with the year. This would mean a space was present year = pynutil.add_weight(year, -0.001) year |= ( - pynutil.delete("year: \"") + pynutil.delete('year: "') + pynutil.insert("de ") + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) # day month year - graph_dmy = day + pynini.cross(DAMO_SPACE, " de ") + month + pynini.closure(pynini.accep(" ") + year, 0, 1) + graph_dmy = ( + day + + pynini.cross(DAMO_SPACE, " de ") + + month + + pynini.closure(pynini.accep(" ") + year, 0, 1) + ) graph_mdy = month + DAMO_SPACE + day + pynini.closure(DAMO_SPACE + year, 0, 1) if deterministic: - graph_mdy += pynutil.delete(" preserve_order: true") # Only accepts this if was explicitly passed + graph_mdy += pynutil.delete( + " preserve_order: true" + ) # Only accepts this if was explicitly passed self.graph = graph_dmy | graph_mdy final_graph = self.graph + delete_preserve_order diff --git a/fun_text_processing/text_normalization/es/verbalizers/decimals.py b/fun_text_processing/text_normalization/es/verbalizers/decimals.py index 43c9d2bb3..ccc624642 100644 --- a/fun_text_processing/text_normalization/es/verbalizers/decimals.py +++ b/fun_text_processing/text_normalization/es/verbalizers/decimals.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_NOT_QUOTE, @@ -20,25 +18,35 @@ from pynini.lib import pynutil class DecimalFst(GraphFst): """ - Finite state transducer for classifying decimal, e.g. - decimal { negative: "true" integer_part: "dos" fractional_part: "cuatro cero" quantity: "billones" } -> menos dos coma quatro cero billones - decimal { integer_part: "un" quantity: "billón" } -> un billón + Finite state transducer for classifying decimal, e.g. + decimal { negative: "true" integer_part: "dos" fractional_part: "cuatro cero" quantity: "billones" } -> menos dos coma quatro cero billones + decimal { integer_part: "un" quantity: "billón" } -> un billón Args: - deterministic: if True will provide a single transduction option, - for False multiple transduction are generated (used for audio-based normalization) - """ + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ def __init__(self, deterministic: bool = True): super().__init__(name="decimal", kind="classify", deterministic=deterministic) - optional_sign = pynini.closure(pynini.cross("negative: \"true\"", "menos ") + delete_space, 0, 1) - integer = pynutil.delete("integer_part: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + optional_sign = pynini.closure( + pynini.cross('negative: "true"', "menos ") + delete_space, 0, 1 + ) + integer = ( + pynutil.delete('integer_part: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + ) fractional_default = ( - pynutil.delete("fractional_part: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + pynutil.delete('fractional_part: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') ) - conjunction = pynutil.insert(" punto ") if LOCALIZATION == "am" else pynutil.insert(" coma ") + conjunction = ( + pynutil.insert(" punto ") if LOCALIZATION == "am" else pynutil.insert(" coma ") + ) if not deterministic: conjunction |= pynutil.insert(pynini.union(" con ", " y ")) fractional_default |= strip_cardinal_apocope(fractional_default) @@ -47,9 +55,9 @@ class DecimalFst(GraphFst): quantity = ( delete_space + insert_space - + pynutil.delete("quantity: \"") + + pynutil.delete('quantity: "') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_quantity = pynini.closure(quantity, 0, 1) @@ -58,13 +66,17 @@ class DecimalFst(GraphFst): ) # Allowing permutation for fem gender, don't include quantity since "million","billion", etc.. are masculine - graph_fem = optional_sign + (shift_cardinal_gender(integer) + delete_space + shift_number_gender(fractional)) + graph_fem = optional_sign + ( + shift_cardinal_gender(integer) + delete_space + shift_number_gender(fractional) + ) if not deterministic: # "una" will drop to "un" in certain cases graph_fem |= add_cardinal_apocope_fem(graph_fem) self.numbers_only_quantity = ( optional_sign - + pynini.union((integer + quantity), (integer + delete_space + fractional + quantity)).optimize() + + pynini.union( + (integer + quantity), (integer + delete_space + fractional + quantity) + ).optimize() ) self.graph_masc = (graph_masc + delete_preserve_order).optimize() diff --git a/fun_text_processing/text_normalization/es/verbalizers/electronic.py b/fun_text_processing/text_normalization/es/verbalizers/electronic.py index 6104ff6f1..1b4d51d1c 100644 --- a/fun_text_processing/text_normalization/es/verbalizers/electronic.py +++ b/fun_text_processing/text_normalization/es/verbalizers/electronic.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_NOT_QUOTE, @@ -35,7 +33,8 @@ class ElectronicFst(GraphFst): super().__init__(name="electronic", kind="verbalize", deterministic=deterministic) graph_digit_no_zero = ( - digit_no_zero @ pynini.cdrewrite(pynini.cross("un", "uno"), "", "", DAMO_SIGMA).optimize() + digit_no_zero + @ pynini.cdrewrite(pynini.cross("un", "uno"), "", "", DAMO_SIGMA).optimize() ) graph_digit = graph_digit_no_zero | zero @@ -46,18 +45,20 @@ class ElectronicFst(GraphFst): verbalize_characters = pynini.cdrewrite(graph_symbols | graph_digit, "", "", DAMO_SIGMA) - user_name = pynutil.delete("username: \"") + add_space_after_char() + pynutil.delete("\"") + user_name = pynutil.delete('username: "') + add_space_after_char() + pynutil.delete('"') user_name @= verbalize_characters - convert_defaults = pynutil.add_weight(DAMO_NOT_QUOTE, weight=0.0001) | domain_common | server_common + convert_defaults = ( + pynutil.add_weight(DAMO_NOT_QUOTE, weight=0.0001) | domain_common | server_common + ) domain = convert_defaults + pynini.closure(insert_space + convert_defaults) domain @= verbalize_characters - domain = pynutil.delete("domain: \"") + domain + pynutil.delete("\"") + domain = pynutil.delete('domain: "') + domain + pynutil.delete('"') protocol = ( - pynutil.delete("protocol: \"") + pynutil.delete('protocol: "') + add_space_after_char() @ pynini.cdrewrite(graph_symbols, "", "", DAMO_SIGMA) - + pynutil.delete("\"") + + pynutil.delete('"') ) self.graph = (pynini.closure(protocol + pynini.accep(" "), 0, 1) + domain) | ( user_name + pynini.accep(" ") + pynutil.insert("arroba ") + domain diff --git a/fun_text_processing/text_normalization/es/verbalizers/fraction.py b/fun_text_processing/text_normalization/es/verbalizers/fraction.py index 3f99d830e..7e157af3c 100644 --- a/fun_text_processing/text_normalization/es/verbalizers/fraction.py +++ b/fun_text_processing/text_normalization/es/verbalizers/fraction.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_CHAR, @@ -21,15 +19,15 @@ from pynini.lib import pynutil class FractionFst(GraphFst): """ - Finite state transducer for verbalizing fraction - e.g. tokens { fraction { integer: "treinta y tres" numerator: "cuatro" denominator: "quinto" } } -> - treinta y tres y cuatro quintos + Finite state transducer for verbalizing fraction + e.g. tokens { fraction { integer: "treinta y tres" numerator: "cuatro" denominator: "quinto" } } -> + treinta y tres y cuatro quintos - Args: - deterministic: if True will provide a single transduction option, - for False multiple transduction are generated (used for audio-based normalization) - """ + Args: + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ def __init__(self, deterministic: bool = True): super().__init__(name="fraction", kind="verbalize", deterministic=deterministic) @@ -40,28 +38,28 @@ class FractionFst(GraphFst): conjunction = pynutil.insert(" y ") integer = ( - pynutil.delete("integer_part: \"") + pynutil.delete('integer_part: "') + strip_cardinal_apocope(pynini.closure(DAMO_NOT_QUOTE)) - + pynutil.delete("\"") + + pynutil.delete('"') ) - numerator_one = pynutil.delete("numerator: \"") + pynini.accep("un") + pynutil.delete("\" ") + numerator_one = pynutil.delete('numerator: "') + pynini.accep("un") + pynutil.delete('" ') numerator = ( - pynutil.delete("numerator: \"") + pynutil.delete('numerator: "') + pynini.difference(pynini.closure(DAMO_NOT_QUOTE), "un") - + pynutil.delete("\" ") + + pynutil.delete('" ') ) - denominator_add_stem = pynutil.delete("denominator: \"") + ( + denominator_add_stem = pynutil.delete('denominator: "') + ( pynini.closure(DAMO_NOT_QUOTE) + fraction_stem - + pynutil.delete("\" morphosyntactic_features: \"add_root\"") + + pynutil.delete('" morphosyntactic_features: "add_root"') ) - denominator_ordinal = pynutil.delete("denominator: \"") + ( - pynini.closure(DAMO_NOT_QUOTE) + pynutil.delete("\" morphosyntactic_features: \"ordinal\"") + denominator_ordinal = pynutil.delete('denominator: "') + ( + pynini.closure(DAMO_NOT_QUOTE) + pynutil.delete('" morphosyntactic_features: "ordinal"') ) - denominator_cardinal = pynutil.delete("denominator: \"") + ( - pynini.closure(DAMO_NOT_QUOTE) + pynutil.delete("\"") + denominator_cardinal = pynutil.delete('denominator: "') + ( + pynini.closure(DAMO_NOT_QUOTE) + pynutil.delete('"') ) denominator_singular = pynini.union(denominator_add_stem, denominator_ordinal) @@ -76,7 +74,9 @@ class FractionFst(GraphFst): merge = pynini.cdrewrite( pynini.cross(" y ", "i"), "", "", DAMO_SIGMA ) # The denominator must be a single word, with the conjunction "y" replaced by i - merge @= pynini.cdrewrite(delete_space, "", pynini.difference(DAMO_CHAR, "parte"), DAMO_SIGMA) + merge @= pynini.cdrewrite( + delete_space, "", pynini.difference(DAMO_CHAR, "parte"), DAMO_SIGMA + ) # The merger can produce duplicate vowels. This is not allowed in orthography delete_duplicates = pynini.string_map([("aa", "a"), ("oo", "o")]) # Removes vowels @@ -90,10 +90,15 @@ class FractionFst(GraphFst): ) merge_into_single_word = merge @ remove_accents @ delete_duplicates - fraction_default = numerator + delete_space + insert_space + (denominator_plural @ merge_into_single_word) + fraction_default = ( + numerator + delete_space + insert_space + (denominator_plural @ merge_into_single_word) + ) fraction_with_one = ( - numerator_one + delete_space + insert_space + (denominator_singular @ merge_into_single_word) + numerator_one + + delete_space + + insert_space + + (denominator_singular @ merge_into_single_word) ) fraction_with_cardinal = strip_cardinal_apocope(numerator | numerator_one) @@ -106,9 +111,13 @@ class FractionFst(GraphFst): # Other rules will manage use of "un" at end, so just worry about endings exceptions = pynini.string_map([("tercia", "tercera")]) apply_exceptions = pynini.cdrewrite(exceptions, "", "", DAMO_SIGMA) - vowel_change = pynini.cdrewrite(pynini.cross("o", "a"), "", pynini.accep("[EOS]"), DAMO_SIGMA) + vowel_change = pynini.cdrewrite( + pynini.cross("o", "a"), "", pynini.accep("[EOS]"), DAMO_SIGMA + ) - denominator_singular_fem = shift_cardinal_gender(denominator_singular) @ vowel_change @ apply_exceptions + denominator_singular_fem = ( + shift_cardinal_gender(denominator_singular) @ vowel_change @ apply_exceptions + ) denominator_plural_fem = denominator_singular_fem + plural numerator_one_fem = shift_cardinal_gender(numerator_one) @@ -129,7 +138,8 @@ class FractionFst(GraphFst): fraction_with_one_fem = numerator_one_fem + delete_space + insert_space fraction_with_one_fem += pynini.union( - denominator_singular_fem @ merge_stem, denominator_singular_fem @ merge_into_single_word + denominator_singular_fem @ merge_stem, + denominator_singular_fem @ merge_into_single_word, ) # Both forms exists fraction_with_one_fem += pynutil.insert(" parte") fraction_with_one_fem @= pynini.cdrewrite( @@ -147,7 +157,9 @@ class FractionFst(GraphFst): ) # Case of no merger fraction_default |= fraction_default_fem - fraction_with_one |= numerator_one + delete_space + insert_space + denominator_singular @ merge_stem + fraction_with_one |= ( + numerator_one + delete_space + insert_space + denominator_singular @ merge_stem + ) fraction_with_one |= fraction_with_one_fem fraction_with_one @= pynini.cdrewrite( @@ -166,7 +178,9 @@ class FractionFst(GraphFst): + (denominator_plural @ pynini.cross("medios", "medias")) ) fraction_with_one |= ( - pynutil.delete(numerator_one) + delete_space + (denominator_singular @ pynini.cross("medio", "media")) + pynutil.delete(numerator_one) + + delete_space + + (denominator_singular @ pynini.cross("medio", "media")) ) fraction_fem = fraction_with_one | fraction_default | fraction_with_cardinal diff --git a/fun_text_processing/text_normalization/es/verbalizers/measure.py b/fun_text_processing/text_normalization/es/verbalizers/measure.py index 0d27ba382..09f9c884d 100644 --- a/fun_text_processing/text_normalization/es/verbalizers/measure.py +++ b/fun_text_processing/text_normalization/es/verbalizers/measure.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_NOT_QUOTE, @@ -38,7 +36,9 @@ class MeasureFst(GraphFst): for False multiple transduction are generated (used for audio-based normalization) """ - def __init__(self, decimal: GraphFst, cardinal: GraphFst, fraction: GraphFst, deterministic: bool): + def __init__( + self, decimal: GraphFst, cardinal: GraphFst, fraction: GraphFst, deterministic: bool + ): super().__init__(name="measure", kind="verbalize", deterministic=deterministic) graph_decimal_masc = decimal.delete_tokens(decimal.graph_masc) @@ -52,38 +52,59 @@ class MeasureFst(GraphFst): DAMO_WHITE_SPACE + "por" + pynini.closure(DAMO_NOT_QUOTE, 1), 0, 1 ) unit_masc |= "por" + pynini.closure(DAMO_NOT_QUOTE, 1) - unit_masc = pynutil.delete("units: \"") + (pynini.closure(DAMO_NOT_QUOTE) @ unit_masc) + pynutil.delete("\"") + unit_masc = ( + pynutil.delete('units: "') + + (pynini.closure(DAMO_NOT_QUOTE) @ unit_masc) + + pynutil.delete('"') + ) unit_fem = (unit_plural_fem | unit_singular_fem) + pynini.closure( DAMO_WHITE_SPACE + "por" + pynini.closure(DAMO_NOT_QUOTE, 1), 0, 1 ) - unit_fem = pynutil.delete("units: \"") + (pynini.closure(DAMO_NOT_QUOTE) @ unit_fem) + pynutil.delete("\"") + unit_fem = ( + pynutil.delete('units: "') + + (pynini.closure(DAMO_NOT_QUOTE) @ unit_fem) + + pynutil.delete('"') + ) graph_masc = (graph_cardinal_masc | graph_decimal_masc) + DAMO_WHITE_SPACE + unit_masc graph_masc |= graph_fraction_masc + DAMO_WHITE_SPACE + pynutil.insert("de ") + unit_masc graph_masc |= pynutil.add_weight( - graph_fraction_masc @ (DAMO_SIGMA + pynini.union("medio", "medios")) + DAMO_WHITE_SPACE + unit_masc, -0.001 + graph_fraction_masc @ (DAMO_SIGMA + pynini.union("medio", "medios")) + + DAMO_WHITE_SPACE + + unit_masc, + -0.001, ) # "medio litro" not "medio de litro" graph_fem = (graph_cardinal_fem | graph_decimal_fem) + DAMO_WHITE_SPACE + unit_fem graph_fem |= graph_fraction_fem + DAMO_WHITE_SPACE + pynutil.insert("de ") + unit_fem graph_fem |= pynutil.add_weight( - graph_fraction_fem @ (DAMO_SIGMA + pynini.union("media", "medias")) + DAMO_WHITE_SPACE + unit_fem, -0.001 + graph_fraction_fem @ (DAMO_SIGMA + pynini.union("media", "medias")) + + DAMO_WHITE_SPACE + + unit_fem, + -0.001, ) graph = graph_masc | graph_fem graph = ( pynini.cdrewrite( - pynutil.insert(" de"), "quantity: \"" + pynini.closure(DAMO_NOT_QUOTE, 1), "\"", DAMO_SIGMA + pynutil.insert(" de"), + 'quantity: "' + pynini.closure(DAMO_NOT_QUOTE, 1), + '"', + DAMO_SIGMA, ) @ graph ) # billones de xyz - graph @= pynini.cdrewrite(pynini.cross(ones, "uno"), "", DAMO_WHITE_SPACE + "por", DAMO_SIGMA) + graph @= pynini.cdrewrite( + pynini.cross(ones, "uno"), "", DAMO_WHITE_SPACE + "por", DAMO_SIGMA + ) # To manage alphanumeric combonations ("a-8, 5x"), we let them use a weighted default path. - alpha_num_unit = pynutil.delete("units: \"") + pynini.closure(DAMO_NOT_QUOTE) + pynutil.delete("\"") + alpha_num_unit = ( + pynutil.delete('units: "') + pynini.closure(DAMO_NOT_QUOTE) + pynutil.delete('"') + ) graph_alpha_num = pynini.union( (graph_cardinal_masc | graph_decimal_masc) + DAMO_SPACE + alpha_num_unit, alpha_num_unit + delete_extra_space + (graph_cardinal_masc | graph_decimal_masc), diff --git a/fun_text_processing/text_normalization/es/verbalizers/money.py b/fun_text_processing/text_normalization/es/verbalizers/money.py index 4bb9a8378..96800c6b4 100644 --- a/fun_text_processing/text_normalization/es/verbalizers/money.py +++ b/fun_text_processing/text_normalization/es/verbalizers/money.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_NOT_QUOTE, @@ -8,7 +6,10 @@ from fun_text_processing.text_normalization.en.graph_utils import ( GraphFst, delete_preserve_order, ) -from fun_text_processing.text_normalization.es.graph_utils import shift_cardinal_gender, strip_cardinal_apocope +from fun_text_processing.text_normalization.es.graph_utils import ( + shift_cardinal_gender, + strip_cardinal_apocope, +) from fun_text_processing.text_normalization.es.utils import get_abs_path from pynini.lib import pynutil @@ -41,60 +42,66 @@ class MoneyFst(GraphFst): super().__init__(name="money", kind="verbalize", deterministic=deterministic) maj_singular_masc = ( - pynutil.delete("currency_maj: \"") + pynutil.delete('currency_maj: "') + (pynini.closure(DAMO_NOT_QUOTE, 1) @ masc_singular) - + pynutil.delete("\"") + + pynutil.delete('"') ) maj_singular_fem = ( - pynutil.delete("currency_maj: \"") + pynutil.delete('currency_maj: "') + (pynini.closure(DAMO_NOT_QUOTE, 1) @ fem_singular) - + pynutil.delete("\"") + + pynutil.delete('"') ) maj_plural_masc = ( - pynutil.delete("currency_maj: \"") + pynutil.delete('currency_maj: "') + (pynini.closure(DAMO_NOT_QUOTE, 1) @ masc_plural) - + pynutil.delete("\"") + + pynutil.delete('"') ) maj_plural_fem = ( - pynutil.delete("currency_maj: \"") + pynutil.delete('currency_maj: "') + (pynini.closure(DAMO_NOT_QUOTE, 1) @ fem_plural) - + pynutil.delete("\"") + + pynutil.delete('"') ) maj_masc = maj_plural_masc | maj_singular_masc # Tagger kept quantity resolution stable maj_fem = maj_plural_fem | maj_singular_fem min_singular_masc = ( - pynutil.delete("currency_min: \"") + pynutil.delete('currency_min: "') + (pynini.closure(DAMO_NOT_QUOTE, 1) @ masc_singular) - + pynutil.delete("\"") + + pynutil.delete('"') ) min_singular_fem = ( - pynutil.delete("currency_min: \"") + pynutil.delete('currency_min: "') + (pynini.closure(DAMO_NOT_QUOTE, 1) @ fem_singular) - + pynutil.delete("\"") + + pynutil.delete('"') ) min_plural_masc = ( - pynutil.delete("currency_min: \"") + pynutil.delete('currency_min: "') + (pynini.closure(DAMO_NOT_QUOTE, 1) @ masc_plural) - + pynutil.delete("\"") + + pynutil.delete('"') ) min_plural_fem = ( - pynutil.delete("currency_min: \"") + pynutil.delete('currency_min: "') + (pynini.closure(DAMO_NOT_QUOTE, 1) @ fem_plural) - + pynutil.delete("\"") + + pynutil.delete('"') ) min_masc = min_plural_masc | min_singular_masc min_fem = min_plural_fem | min_singular_fem fractional_part = ( - pynutil.delete("fractional_part: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + pynutil.delete('fractional_part: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') ) - integer_part = pynutil.delete("integer_part: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + integer_part = ( + pynutil.delete('integer_part: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + ) optional_add_and = pynini.closure(pynutil.insert(pynini.union("con ", "y ")), 0, 1) # *** currency_maj @@ -132,20 +139,27 @@ class MoneyFst(GraphFst): graph_decimal_masc = decimal.graph_masc + DAMO_SPACE + maj_masc graph_decimal_fem = decimal.graph_fem - graph_decimal_fem |= decimal.numbers_only_quantity # can still have "x billions" with fem currency + graph_decimal_fem |= ( + decimal.numbers_only_quantity + ) # can still have "x billions" with fem currency graph_decimal_fem += DAMO_SPACE + maj_fem graph_decimal = graph_decimal_fem | graph_decimal_masc graph_decimal = ( pynini.cdrewrite( - pynutil.insert(" de"), "quantity: \"" + pynini.closure(DAMO_NOT_QUOTE, 1), "\"", DAMO_SIGMA + pynutil.insert(" de"), + 'quantity: "' + pynini.closure(DAMO_NOT_QUOTE, 1), + '"', + DAMO_SIGMA, ) @ graph_decimal ) # formally it's millones/billones de *** # *** current_min graph_minor_masc = fractional_part + DAMO_SPACE + min_masc + delete_preserve_order - graph_minor_fem = shift_cardinal_gender(fractional_part) + DAMO_SPACE + min_fem + delete_preserve_order + graph_minor_fem = ( + shift_cardinal_gender(fractional_part) + DAMO_SPACE + min_fem + delete_preserve_order + ) graph_minor = graph_minor_fem | graph_minor_masc diff --git a/fun_text_processing/text_normalization/es/verbalizers/ordinal.py b/fun_text_processing/text_normalization/es/verbalizers/ordinal.py index 6f1786199..1c97473da 100644 --- a/fun_text_processing/text_normalization/es/verbalizers/ordinal.py +++ b/fun_text_processing/text_normalization/es/verbalizers/ordinal.py @@ -1,7 +1,10 @@ - - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, DAMO_SIGMA, DAMO_SPACE, GraphFst +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + DAMO_SIGMA, + DAMO_SPACE, + GraphFst, +) from fun_text_processing.text_normalization.es.graph_utils import shift_number_gender from pynini.lib import pynutil @@ -11,7 +14,7 @@ class OrdinalFst(GraphFst): Finite state transducer for verbalizing ordinals e.g. ordinal { integer: "tercer" } } -> "tercero" -> "tercera" - -> "tercer" + -> "tercer" Args: deterministic: if True will provide a single transduction option, @@ -21,16 +24,20 @@ class OrdinalFst(GraphFst): def __init__(self, deterministic: bool = True): super().__init__(name="ordinal", kind="verbalize", deterministic=deterministic) - graph = pynutil.delete("integer: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + graph = ( + pynutil.delete('integer: "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') + ) # masculne gender we leave as is - graph_masc = graph + pynutil.delete(" morphosyntactic_features: \"gender_masc") + graph_masc = graph + pynutil.delete(' morphosyntactic_features: "gender_masc') # shift gender graph_fem_ending = graph @ pynini.cdrewrite( pynini.cross("o", "a"), "", DAMO_SPACE | pynini.accep("[EOS]"), DAMO_SIGMA ) - graph_fem = shift_number_gender(graph_fem_ending) + pynutil.delete(" morphosyntactic_features: \"gender_fem") + graph_fem = shift_number_gender(graph_fem_ending) + pynutil.delete( + ' morphosyntactic_features: "gender_fem' + ) # Apocope just changes tercero and primero. May occur if someone wrote 11.er (uncommon) graph_apocope = ( @@ -38,21 +45,24 @@ class OrdinalFst(GraphFst): | pynini.cross("primero", "primer") | pynini.cross("undécimo", "decimoprimer") ) # In case someone wrote 11.er with deterministic - graph_apocope = (graph @ pynini.cdrewrite(graph_apocope, "", "", DAMO_SIGMA)) + pynutil.delete( - " morphosyntactic_features: \"apocope" - ) + graph_apocope = ( + graph @ pynini.cdrewrite(graph_apocope, "", "", DAMO_SIGMA) + ) + pynutil.delete(' morphosyntactic_features: "apocope') graph = graph_apocope | graph_masc | graph_fem if not deterministic: # Plural graph graph_plural = pynini.cdrewrite( - pynutil.insert("s"), pynini.union("o", "a"), DAMO_SPACE | pynini.accep("[EOS]"), DAMO_SIGMA + pynutil.insert("s"), + pynini.union("o", "a"), + DAMO_SPACE | pynini.accep("[EOS]"), + DAMO_SIGMA, ) graph |= (graph @ graph_plural) + pynutil.delete("/plural") - self.graph = graph + pynutil.delete("\"") + self.graph = graph + pynutil.delete('"') delete_tokens = self.delete_tokens(self.graph) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/text_normalization/es/verbalizers/telephone.py b/fun_text_processing/text_normalization/es/verbalizers/telephone.py index 75acee934..e347be650 100644 --- a/fun_text_processing/text_normalization/es/verbalizers/telephone.py +++ b/fun_text_processing/text_normalization/es/verbalizers/telephone.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst from pynini.lib import pynutil @@ -19,6 +17,10 @@ class TelephoneFst(GraphFst): def __init__(self, deterministic: bool = True): super().__init__(name="telephone", kind="verbalize") - number_part = pynutil.delete("number_part: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + number_part = ( + pynutil.delete('number_part: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + ) delete_tokens = self.delete_tokens(number_part) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/text_normalization/es/verbalizers/time.py b/fun_text_processing/text_normalization/es/verbalizers/time.py index e209c8799..448b9cb63 100644 --- a/fun_text_processing/text_normalization/es/verbalizers/time.py +++ b/fun_text_processing/text_normalization/es/verbalizers/time.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_NOT_QUOTE, @@ -33,7 +31,9 @@ class TimeFst(GraphFst): def __init__(self, deterministic: bool = True): super().__init__(name="time", kind="verbalize", deterministic=deterministic) - change_minutes = pynini.cdrewrite(alt_minutes, pynini.accep("[BOS]"), pynini.accep("[EOS]"), DAMO_SIGMA) + change_minutes = pynini.cdrewrite( + alt_minutes, pynini.accep("[BOS]"), pynini.accep("[EOS]"), DAMO_SIGMA + ) morning_phrases = pynini.cross("am", "de la mañana") afternoon_phrases = pynini.cross("pm", "de la tarde") @@ -58,40 +58,44 @@ class TimeFst(GraphFst): hour = ( pynutil.delete("hours:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) minute = ( pynutil.delete("minutes:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') + ) + minute = ( + (minute @ change_minutes) + if deterministic + else pynini.union(minute, minute @ change_minutes) ) - minute = (minute @ change_minutes) if deterministic else pynini.union(minute, minute @ change_minutes) suffix = ( pynutil.delete("suffix:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) zone = ( pynutil.delete("zone:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) optional_zone = pynini.closure(delete_space + insert_space + zone, 0, 1) second = ( pynutil.delete("seconds:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) graph_hms = ( @@ -148,8 +152,8 @@ class TimeFst(GraphFst): graph = (graph_hms | graph_hm | graph_h) + optional_zone if not deterministic: - graph_style_1 = pynutil.delete(" style: \"1\"") - graph_style_2 = pynutil.delete(" style: \"2\"") + graph_style_1 = pynutil.delete(' style: "1"') + graph_style_2 = pynutil.delete(' style: "2"') graph_menos = hour + delete_space + pynutil.insert(" menos ") + minute + graph_style_1 graph_menos |= ( diff --git a/fun_text_processing/text_normalization/es/verbalizers/verbalize.py b/fun_text_processing/text_normalization/es/verbalizers/verbalize.py index 512beed96..075f88958 100644 --- a/fun_text_processing/text_normalization/es/verbalizers/verbalize.py +++ b/fun_text_processing/text_normalization/es/verbalizers/verbalize.py @@ -1,5 +1,3 @@ - - from fun_text_processing.text_normalization.en.graph_utils import GraphFst from fun_text_processing.text_normalization.en.verbalizers.whitelist import WhiteListFst from fun_text_processing.text_normalization.es.verbalizers.cardinal import CardinalFst @@ -37,7 +35,9 @@ class VerbalizeFst(GraphFst): fraction_graph = fraction.fst date = DateFst(deterministic=deterministic) date_graph = date.fst - measure = MeasureFst(cardinal=cardinal, decimal=decimal, fraction=fraction, deterministic=deterministic) + measure = MeasureFst( + cardinal=cardinal, decimal=decimal, fraction=fraction, deterministic=deterministic + ) measure_graph = measure.fst electronic = ElectronicFst(deterministic=deterministic) electronic_graph = electronic.fst diff --git a/fun_text_processing/text_normalization/es/verbalizers/verbalize_final.py b/fun_text_processing/text_normalization/es/verbalizers/verbalize_final.py index c9589e94d..93c671e27 100644 --- a/fun_text_processing/text_normalization/es/verbalizers/verbalize_final.py +++ b/fun_text_processing/text_normalization/es/verbalizers/verbalize_final.py @@ -1,6 +1,3 @@ - - - import os import pynini @@ -28,16 +25,20 @@ class VerbalizeFinalFst(GraphFst): overwrite_cache: set to True to overwrite .far files """ - def __init__(self, deterministic: bool = True, cache_dir: str = None, overwrite_cache: bool = False): + def __init__( + self, deterministic: bool = True, cache_dir: str = None, overwrite_cache: bool = False + ): super().__init__(name="verbalize_final", kind="verbalize", deterministic=deterministic) far_file = None if cache_dir is not None and cache_dir != "None": os.makedirs(cache_dir, exist_ok=True) - far_file = os.path.join(cache_dir, f"es_tn_{deterministic}_deterministic_verbalizer.far") + far_file = os.path.join( + cache_dir, f"es_tn_{deterministic}_deterministic_verbalizer.far" + ) if not overwrite_cache and far_file and os.path.exists(far_file): self.fst = pynini.Far(far_file, mode="r")["verbalize"] - logging.info(f'VerbalizeFinalFst graph was restored from {far_file}.') + logging.info(f"VerbalizeFinalFst graph was restored from {far_file}.") else: verbalize = VerbalizeFst(deterministic=deterministic).fst diff --git a/fun_text_processing/text_normalization/export_models.py b/fun_text_processing/text_normalization/export_models.py index 72c3ff8a7..09fab4527 100644 --- a/fun_text_processing/text_normalization/export_models.py +++ b/fun_text_processing/text_normalization/export_models.py @@ -3,14 +3,23 @@ from time import perf_counter from argparse import ArgumentParser from fun_text_processing.text_normalization.en.graph_utils import generator_main + def parse_args(): parser = ArgumentParser() - + parser.add_argument( - "--language", help="language", choices=['de', 'en', 'es', 'ru', 'zh'], default="en", type=str + "--language", + help="language", + choices=["de", "en", "es", "ru", "zh"], + default="en", + type=str, ) parser.add_argument( - "--input_case", help="input capitalization", choices=["lower_cased", "cased"], default="cased", type=str + "--input_case", + help="input capitalization", + choices=["lower_cased", "cased"], + default="cased", + type=str, ) parser.add_argument( "--export_dir", @@ -20,27 +29,53 @@ def parse_args(): ) return parser.parse_args() -def get_grammars(lang: str="en", input_case: str="cased"): - if lang=='de': - from fun_text_processing.text_normalization.de.taggers.tokenize_and_classify import ClassifyFst - from fun_text_processing.text_normalization.de.verbalizers.verbalize_final import VerbalizeFinalFst - elif lang=='en': - from fun_text_processing.text_normalization.en.taggers.tokenize_and_classify import ClassifyFst - from fun_text_processing.text_normalization.en.verbalizers.verbalize_final import VerbalizeFinalFst - elif lang=='es': - from fun_text_processing.text_normalization.es.taggers.tokenize_and_classify import ClassifyFst - from fun_text_processing.text_normalization.es.verbalizers.verbalize_final import VerbalizeFinalFst - elif lang=='ru': - from fun_text_processing.text_normalization.ru.taggers.tokenize_and_classify import ClassifyFst - from fun_text_processing.text_normalization.ru.verbalizers.verbalize_final import VerbalizeFinalFst - elif lang=='zh': - from fun_text_processing.text_normalization.zh.taggers.tokenize_and_classify import ClassifyFst - from fun_text_processing.text_normalization.zh.verbalizers.verbalize_final import VerbalizeFinalFst - else: - from fun_text_processing.text_normalization.en.taggers.tokenize_and_classify import ClassifyFst - from fun_text_processing.text_normalization.en.verbalizers.verbalize_final import VerbalizeFinalFst - return ClassifyFst(input_case=input_case).fst, VerbalizeFinalFst().fst +def get_grammars(lang: str = "en", input_case: str = "cased"): + if lang == "de": + from fun_text_processing.text_normalization.de.taggers.tokenize_and_classify import ( + ClassifyFst, + ) + from fun_text_processing.text_normalization.de.verbalizers.verbalize_final import ( + VerbalizeFinalFst, + ) + elif lang == "en": + from fun_text_processing.text_normalization.en.taggers.tokenize_and_classify import ( + ClassifyFst, + ) + from fun_text_processing.text_normalization.en.verbalizers.verbalize_final import ( + VerbalizeFinalFst, + ) + elif lang == "es": + from fun_text_processing.text_normalization.es.taggers.tokenize_and_classify import ( + ClassifyFst, + ) + from fun_text_processing.text_normalization.es.verbalizers.verbalize_final import ( + VerbalizeFinalFst, + ) + elif lang == "ru": + from fun_text_processing.text_normalization.ru.taggers.tokenize_and_classify import ( + ClassifyFst, + ) + from fun_text_processing.text_normalization.ru.verbalizers.verbalize_final import ( + VerbalizeFinalFst, + ) + elif lang == "zh": + from fun_text_processing.text_normalization.zh.taggers.tokenize_and_classify import ( + ClassifyFst, + ) + from fun_text_processing.text_normalization.zh.verbalizers.verbalize_final import ( + VerbalizeFinalFst, + ) + else: + from fun_text_processing.text_normalization.en.taggers.tokenize_and_classify import ( + ClassifyFst, + ) + from fun_text_processing.text_normalization.en.verbalizers.verbalize_final import ( + VerbalizeFinalFst, + ) + + return ClassifyFst(input_case=input_case).fst, VerbalizeFinalFst().fst + if __name__ == "__main__": args = parse_args() @@ -54,5 +89,4 @@ if __name__ == "__main__": tagger_fst, verbalizer_fst = get_grammars(args.language, args.input_case) generator_main(tagger_far_file, {"tokenize_and_classify": tagger_fst}) generator_main(verbalizer_far_file, {"verbalize": verbalizer_fst}) - print(f'Time to generate graph: {round(perf_counter() - start_time, 2)} sec') - + print(f"Time to generate graph: {round(perf_counter() - start_time, 2)} sec") diff --git a/fun_text_processing/text_normalization/normalize.py b/fun_text_processing/text_normalization/normalize.py index d87db4a46..c1cebdd09 100644 --- a/fun_text_processing/text_normalization/normalize.py +++ b/fun_text_processing/text_normalization/normalize.py @@ -1,5 +1,3 @@ - - import itertools import os import re @@ -29,7 +27,7 @@ try: except (ModuleNotFoundError, ImportError) as e: NLP_AVAILABLE = False -SPACE_DUP = re.compile(' {2,}') +SPACE_DUP = re.compile(" {2,}") class Normalizer: @@ -50,7 +48,7 @@ class Normalizer: def __init__( self, input_case: str, - lang: str = 'en', + lang: str = "en", deterministic: bool = True, cache_dir: str = None, overwrite_cache: bool = False, @@ -63,36 +61,62 @@ class Normalizer: self.post_processor = None if lang == "en": - from fun_text_processing.text_normalization.en.verbalizers.verbalize_final import VerbalizeFinalFst - from fun_text_processing.text_normalization.en.verbalizers.post_processing import PostProcessingFst + from fun_text_processing.text_normalization.en.verbalizers.verbalize_final import ( + VerbalizeFinalFst, + ) + from fun_text_processing.text_normalization.en.verbalizers.post_processing import ( + PostProcessingFst, + ) if post_process: - self.post_processor = PostProcessingFst(cache_dir=cache_dir, overwrite_cache=overwrite_cache) + self.post_processor = PostProcessingFst( + cache_dir=cache_dir, overwrite_cache=overwrite_cache + ) if deterministic: - from fun_text_processing.text_normalization.en.taggers.tokenize_and_classify import ClassifyFst + from fun_text_processing.text_normalization.en.taggers.tokenize_and_classify import ( + ClassifyFst, + ) else: if lm: - from fun_text_processing.text_normalization.en.taggers.tokenize_and_classify_lm import ClassifyFst + from fun_text_processing.text_normalization.en.taggers.tokenize_and_classify_lm import ( + ClassifyFst, + ) else: from fun_text_processing.text_normalization.en.taggers.tokenize_and_classify_with_audio import ( ClassifyFst, ) - elif lang == 'ru': + elif lang == "ru": # Ru TN only support non-deterministic cases and produces multiple normalization options # use normalize_with_audio.py - from fun_text_processing.text_normalization.ru.taggers.tokenize_and_classify import ClassifyFst - from fun_text_processing.text_normalization.ru.verbalizers.verbalize_final import VerbalizeFinalFst - elif lang == 'de': - from fun_text_processing.text_normalization.de.taggers.tokenize_and_classify import ClassifyFst - from fun_text_processing.text_normalization.de.verbalizers.verbalize_final import VerbalizeFinalFst - elif lang == 'es': - from fun_text_processing.text_normalization.es.taggers.tokenize_and_classify import ClassifyFst - from fun_text_processing.text_normalization.es.verbalizers.verbalize_final import VerbalizeFinalFst - elif lang == 'zh': - from fun_text_processing.text_normalization.zh.taggers.tokenize_and_classify import ClassifyFst - from fun_text_processing.text_normalization.zh.verbalizers.verbalize_final import VerbalizeFinalFst + from fun_text_processing.text_normalization.ru.taggers.tokenize_and_classify import ( + ClassifyFst, + ) + from fun_text_processing.text_normalization.ru.verbalizers.verbalize_final import ( + VerbalizeFinalFst, + ) + elif lang == "de": + from fun_text_processing.text_normalization.de.taggers.tokenize_and_classify import ( + ClassifyFst, + ) + from fun_text_processing.text_normalization.de.verbalizers.verbalize_final import ( + VerbalizeFinalFst, + ) + elif lang == "es": + from fun_text_processing.text_normalization.es.taggers.tokenize_and_classify import ( + ClassifyFst, + ) + from fun_text_processing.text_normalization.es.verbalizers.verbalize_final import ( + VerbalizeFinalFst, + ) + elif lang == "zh": + from fun_text_processing.text_normalization.zh.taggers.tokenize_and_classify import ( + ClassifyFst, + ) + from fun_text_processing.text_normalization.zh.verbalizers.verbalize_final import ( + VerbalizeFinalFst, + ) self.tagger = ClassifyFst( input_case=input_case, deterministic=deterministic, @@ -144,7 +168,9 @@ class Normalizer: try: normalized_texts = Parallel(n_jobs=n_jobs)( - delayed(self.process_batch)(texts[i : i + batch], verbose, punct_pre_process, punct_post_process) + delayed(self.process_batch)( + texts[i : i + batch], verbose, punct_pre_process, punct_post_process + ) for i in range(0, len(texts), batch) ) except BaseException as e: @@ -164,7 +190,10 @@ class Normalizer: """ normalized_lines = [ self.normalize( - text, verbose=verbose, punct_pre_process=punct_pre_process, punct_post_process=punct_post_process + text, + verbose=verbose, + punct_pre_process=punct_pre_process, + punct_post_process=punct_post_process, ) for text in tqdm(batch) ] @@ -238,7 +267,11 @@ class Normalizer: return splits def normalize( - self, text: str, verbose: bool = False, punct_pre_process: bool = False, punct_post_process: bool = False + self, + text: str, + verbose: bool = False, + punct_pre_process: bool = False, + punct_post_process: bool = False, ) -> str: """ Main function. Normalizes tokens from written to spoken form @@ -286,10 +319,10 @@ class Normalizer: break if verbalizer_lattice is None: raise ValueError(f"No permutations were generated from tokens {s}") - output += ' ' + self.select_verbalizer(verbalizer_lattice) - output = SPACE_DUP.sub(' ', output[1:]) + output += " " + self.select_verbalizer(verbalizer_lattice) + output = SPACE_DUP.sub(" ", output[1:]) - if self.lang == "en" and hasattr(self, 'post_processor'): + if self.lang == "en" and hasattr(self, "post_processor"): output = self.post_process(output) if punct_post_process: @@ -311,11 +344,11 @@ class Normalizer: Returns list of sentences """ - lower_case_unicode = '' - upper_case_unicode = '' + lower_case_unicode = "" + upper_case_unicode = "" if self.lang == "ru": - lower_case_unicode = '\u0430-\u04FF' - upper_case_unicode = '\u0410-\u042F' + lower_case_unicode = "\u0430-\u04FF" + upper_case_unicode = "\u0410-\u042F" # Read and split transcript by utterance (roughly, sentences) split_pattern = f"(? 'pynini.FstLike': + def find_tags(self, text: str) -> "pynini.FstLike": """ Given text use tagger Fst to tag text @@ -394,7 +429,7 @@ class Normalizer: lattice = text @ self.tagger.fst return lattice - def select_tag(self, lattice: 'pynini.FstLike') -> str: + def select_tag(self, lattice: "pynini.FstLike") -> str: """ Given tagged lattice return shortest path @@ -406,7 +441,7 @@ class Normalizer: tagged_text = pynini.shortestpath(lattice, nshortest=1, unique=True).string() return tagged_text - def find_verbalizer(self, tagged_text: str) -> 'pynini.FstLike': + def find_verbalizer(self, tagged_text: str) -> "pynini.FstLike": """ Given tagged text creates verbalization lattice This is context-independent. @@ -419,7 +454,7 @@ class Normalizer: lattice = tagged_text @ self.verbalizer.fst return lattice - def select_verbalizer(self, lattice: 'pynini.FstLike') -> str: + def select_verbalizer(self, lattice: "pynini.FstLike") -> str: """ Given verbalized lattice return shortest path @@ -433,7 +468,7 @@ class Normalizer: # output = pynini.shortestpath(lattice, nshortest=1, unique=True).string() return output - def post_process(self, normalized_text: 'pynini.FstLike') -> str: + def post_process(self, normalized_text: "pynini.FstLike") -> str: """ Runs post processing graph on normalized text @@ -457,22 +492,34 @@ def parse_args(): input = parser.add_mutually_exclusive_group() input.add_argument("--text", dest="input_string", help="input string", type=str) input.add_argument("--input_file", dest="input_file", help="input file path", type=str) - parser.add_argument('--output_file', dest="output_file", help="output file path", type=str) - parser.add_argument("--language", help="language", choices=["en", "de", "es", "zh"], default="en", type=str) + parser.add_argument("--output_file", dest="output_file", help="output file path", type=str) parser.add_argument( - "--input_case", help="input capitalization", choices=["lower_cased", "cased"], default="cased", type=str + "--language", help="language", choices=["en", "de", "es", "zh"], default="en", type=str ) - parser.add_argument("--verbose", help="print info for debugging", action='store_true') + parser.add_argument( + "--input_case", + help="input capitalization", + choices=["lower_cased", "cased"], + default="cased", + type=str, + ) + parser.add_argument("--verbose", help="print info for debugging", action="store_true") parser.add_argument( "--punct_post_process", help="set to True to enable punctuation post processing to match input.", action="store_true", ) parser.add_argument( - "--punct_pre_process", help="set to True to enable punctuation pre processing", action="store_true" + "--punct_pre_process", + help="set to True to enable punctuation pre processing", + action="store_true", + ) + parser.add_argument( + "--overwrite_cache", help="set to True to re-create .far grammar files", action="store_true" + ) + parser.add_argument( + "--whitelist", help="path to a file with with whitelist", default=None, type=str ) - parser.add_argument("--overwrite_cache", help="set to True to re-create .far grammar files", action="store_true") - parser.add_argument("--whitelist", help="path to a file with with whitelist", default=None, type=str) parser.add_argument( "--cache_dir", help="path to a dir with .far grammar file. Set to None to avoid using cache", diff --git a/fun_text_processing/text_normalization/normalize_with_audio.py b/fun_text_processing/text_normalization/normalize_with_audio.py index 61e32f2f7..973a1d225 100644 --- a/fun_text_processing/text_normalization/normalize_with_audio.py +++ b/fun_text_processing/text_normalization/normalize_with_audio.py @@ -1,4 +1,3 @@ - import json import os import time @@ -72,7 +71,7 @@ class NormalizerWithAudio(Normalizer): def __init__( self, input_case: str, - lang: str = 'en', + lang: str = "en", cache_dir: str = None, overwrite_cache: bool = False, whitelist: str = None, @@ -92,7 +91,13 @@ class NormalizerWithAudio(Normalizer): ) self.lm = lm - def normalize(self, text: str, n_tagged: int, punct_post_process: bool = True, verbose: bool = False,) -> str: + def normalize( + self, + text: str, + n_tagged: int, + punct_post_process: bool = True, + verbose: bool = False, + ) -> str: """ Main function. Normalizes tokens from written to spoken form e.g. 12 kg -> twelve kilograms @@ -161,7 +166,8 @@ class NormalizerWithAudio(Normalizer): if self.processor: normalized_texts = [self.processor.detokenize([t]) for t in normalized_texts] normalized_texts = [ - post_process_punct(input=original_text, normalized_text=t) for t in normalized_texts + post_process_punct(input=original_text, normalized_text=t) + for t in normalized_texts ] if self.lm: @@ -199,9 +205,13 @@ class NormalizerWithAudio(Normalizer): else: try: # try self.tagger graph that produces output without digits - tagged_texts = rewrite.top_rewrites(text, self.tagger.fst_no_digits, nshortest=n_tagged) + tagged_texts = rewrite.top_rewrites( + text, self.tagger.fst_no_digits, nshortest=n_tagged + ) except pynini.lib.rewrite.Error: - tagged_texts = rewrite.top_rewrites(text, self.tagger.fst, nshortest=n_tagged) + tagged_texts = rewrite.top_rewrites( + text, self.tagger.fst, nshortest=n_tagged + ) else: tagged_texts = rewrite.top_rewrites(text, self.tagger.fst, nshortest=n_tagged) return tagged_texts @@ -266,14 +276,16 @@ class NormalizerWithAudio(Normalizer): return input_text, cer if verbose: - print('-' * 30) + print("-" * 30) for option in normalized_texts: print(option) - print('-' * 30) + print("-" * 30) return normalized_text, cer -def calculate_cer(normalized_texts: List[str], pred_text: str, remove_punct=False) -> List[Tuple[str, float]]: +def calculate_cer( + normalized_texts: List[str], pred_text: str, remove_punct=False +) -> List[Tuple[str, float]]: """ Calculates character error rate (CER) @@ -285,7 +297,7 @@ def calculate_cer(normalized_texts: List[str], pred_text: str, remove_punct=Fals """ normalized_options = [] for text in normalized_texts: - text_clean = text.replace('-', ' ').lower() + text_clean = text.replace("-", " ").lower() if remove_punct: for punct in "!?:;,.-()*+-/<=>@^_": text_clean = text_clean.replace(punct, "") @@ -307,23 +319,38 @@ def get_asr_model(asr_model): asr_model = ASRModel.from_pretrained(asr_model) else: raise ValueError( - f'Provide path to the pretrained checkpoint or choose from {ASRModel.get_available_model_names()}' + f"Provide path to the pretrained checkpoint or choose from {ASRModel.get_available_model_names()}" ) return asr_model def parse_args(): parser = ArgumentParser() - parser.add_argument("--text", help="input string or path to a .txt file", default=None, type=str) parser.add_argument( - "--input_case", help="input capitalization", choices=["lower_cased", "cased"], default="cased", type=str + "--text", help="input string or path to a .txt file", default=None, type=str ) parser.add_argument( - "--language", help="Select target language", choices=["en", "ru", "de", "es"], default="en", type=str + "--input_case", + help="input capitalization", + choices=["lower_cased", "cased"], + default="cased", + type=str, ) - parser.add_argument("--audio_data", default=None, help="path to an audio file or .json manifest") parser.add_argument( - '--model', type=str, default='QuartzNet15x5Base-En', help='Pre-trained model name or path to model checkpoint' + "--language", + help="Select target language", + choices=["en", "ru", "de", "es"], + default="en", + type=str, + ) + parser.add_argument( + "--audio_data", default=None, help="path to an audio file or .json manifest" + ) + parser.add_argument( + "--model", + type=str, + default="QuartzNet15x5Base-En", + help="Pre-trained model name or path to model checkpoint", ) parser.add_argument( "--n_tagged", @@ -338,19 +365,29 @@ def parse_args(): action="store_true", ) parser.add_argument( - "--no_punct_post_process", help="set to True to disable punctuation post processing", action="store_true" + "--no_punct_post_process", + help="set to True to disable punctuation post processing", + action="store_true", + ) + parser.add_argument( + "--overwrite_cache", help="set to True to re-create .far grammar files", action="store_true" + ) + parser.add_argument( + "--whitelist", help="path to a file with with whitelist", default=None, type=str ) - parser.add_argument("--overwrite_cache", help="set to True to re-create .far grammar files", action="store_true") - parser.add_argument("--whitelist", help="path to a file with with whitelist", default=None, type=str) parser.add_argument( "--cache_dir", help="path to a dir with .far grammar file. Set to None to avoid using cache", default=None, type=str, ) - parser.add_argument("--n_jobs", default=-2, type=int, help="The maximum number of concurrently running jobs") parser.add_argument( - "--lm", action="store_true", help="Set to True for WFST+LM. Only available for English right now." + "--n_jobs", default=-2, type=int, help="The maximum number of concurrently running jobs" + ) + parser.add_argument( + "--lm", + action="store_true", + help="Set to True for WFST+LM. Only available for English right now.", ) parser.add_argument( "--cer_threshold", @@ -358,18 +395,29 @@ def parse_args(): type=int, help="if CER for pred_text is above the cer_threshold, no normalization will be performed", ) - parser.add_argument("--batch_size", default=200, type=int, help="Number of examples for each process") + parser.add_argument( + "--batch_size", default=200, type=int, help="Number of examples for each process" + ) return parser.parse_args() def _normalize_line( - normalizer: NormalizerWithAudio, n_tagged, verbose, line: str, remove_punct, punct_post_process, cer_threshold + normalizer: NormalizerWithAudio, + n_tagged, + verbose, + line: str, + remove_punct, + punct_post_process, + cer_threshold, ): line = json.loads(line) pred_text = line["pred_text"] normalized_texts = normalizer.normalize( - text=line["text"], verbose=verbose, n_tagged=n_tagged, punct_post_process=punct_post_process, + text=line["text"], + verbose=verbose, + n_tagged=n_tagged, + punct_post_process=punct_post_process, ) normalized_texts = set(normalized_texts) @@ -424,15 +472,15 @@ def normalize_manifest( with open(f"{dir_name}/{batch_idx:05}.json", "w") as f_out: for line in normalized_lines: - f_out.write(json.dumps(line, ensure_ascii=False) + '\n') + f_out.write(json.dumps(line, ensure_ascii=False) + "\n") print(f"Batch -- {batch_idx} -- is complete") - manifest_out = audio_data.replace('.json', '_normalized.json') - with open(audio_data, 'r') as f: + manifest_out = audio_data.replace(".json", "_normalized.json") + with open(audio_data, "r") as f: lines = f.readlines() - print(f'Normalizing {len(lines)} lines of {audio_data}...') + print(f"Normalizing {len(lines)} lines of {audio_data}...") # to save intermediate results to a file batch = min(len(lines), batch_size) @@ -452,7 +500,7 @@ def normalize_manifest( lines = f_in.read() f_out.write(lines) - print(f'Normalized version saved at {manifest_out}') + print(f"Normalized version saved at {manifest_out}") if __name__ == "__main__": @@ -473,7 +521,7 @@ if __name__ == "__main__": ) if os.path.exists(args.text): - with open(args.text, 'r') as f: + with open(args.text, "r") as f: args.text = f.read().strip() normalized_texts = normalizer.normalize( text=args.text, @@ -503,7 +551,7 @@ if __name__ == "__main__": print(norm_text) elif not os.path.exists(args.audio_data): raise ValueError(f"{args.audio_data} not found.") - elif args.audio_data.endswith('.json'): + elif args.audio_data.endswith(".json"): normalizer = NormalizerWithAudio( input_case=args.input_case, lang=args.language, @@ -527,4 +575,4 @@ if __name__ == "__main__": + "'--audio_data' path to audio file and '--text' path to a text file OR" "'--text' string text (for debugging without audio)" ) - print(f'Execution time: {round((time.time() - start)/60, 2)} min.') + print(f"Execution time: {round((time.time() - start)/60, 2)} min.") diff --git a/fun_text_processing/text_normalization/ru/alphabet.py b/fun_text_processing/text_normalization/ru/alphabet.py index 2c0850db1..d63f67b92 100644 --- a/fun_text_processing/text_normalization/ru/alphabet.py +++ b/fun_text_processing/text_normalization/ru/alphabet.py @@ -1,5 +1,3 @@ - - # Copyright 2017 Google Inc. @@ -7,7 +5,10 @@ # Russian minimally supervised number grammar. import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NON_BREAKING_SPACE, DAMO_SPACE +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NON_BREAKING_SPACE, + DAMO_SPACE, +) from fun_text_processing.text_normalization.ru.utils import get_abs_path RU_LOWER_ALPHA = "абвгдеёжзийклмнопрстуфхцчшщъыьэюя" @@ -41,7 +42,9 @@ RU_STRESSED_MAP = [ ("Ё", "Е"), ] -REWRITE_STRESSED = pynini.closure(pynini.string_map(RU_STRESSED_MAP).optimize() | RU_ALPHA).optimize() +REWRITE_STRESSED = pynini.closure( + pynini.string_map(RU_STRESSED_MAP).optimize() | RU_ALPHA +).optimize() TO_CYRILLIC = pynini.string_file(get_abs_path("data/latin_to_cyrillic.tsv")).optimize() TO_LATIN = pynini.invert(TO_CYRILLIC).optimize() RU_ALPHA_OR_SPACE = pynini.union(RU_ALPHA, DAMO_SPACE, DAMO_NON_BREAKING_SPACE).optimize() diff --git a/fun_text_processing/text_normalization/ru/taggers/cardinal.py b/fun_text_processing/text_normalization/ru/taggers/cardinal.py index 390968850..9264146d4 100644 --- a/fun_text_processing/text_normalization/ru/taggers/cardinal.py +++ b/fun_text_processing/text_normalization/ru/taggers/cardinal.py @@ -1,5 +1,3 @@ - - # Copyright 2017 Google Inc. @@ -21,7 +19,7 @@ from pynini.lib import pynutil class CardinalFst(GraphFst): """ - Finite state transducer for classifying cardinals, e.g. + Finite state transducer for classifying cardinals, e.g. "1 001" -> cardinal { integer: "тысяча один" } Args: @@ -34,59 +32,75 @@ class CardinalFst(GraphFst): def __init__(self, number_names: dict, alternative_formats: dict, deterministic: bool = False): super().__init__(name="cardinal", kind="classify", deterministic=deterministic) - self.cardinal_numbers_default = self.get_cardinal_numbers(number_names, alternative_formats, mode="all") + self.cardinal_numbers_default = self.get_cardinal_numbers( + number_names, alternative_formats, mode="all" + ) self.cardinal_numbers_nominative = self.get_cardinal_numbers( number_names, alternative_formats, mode="nominative" ) self.optional_graph_negative = pynini.closure( - pynutil.insert("negative: ") + pynini.cross("-", "\"true\"") + insert_space, 0, 1 + pynutil.insert("negative: ") + pynini.cross("-", '"true"') + insert_space, 0, 1 ) self.cardinal_numbers_with_optional_negative = ( self.optional_graph_negative - + pynutil.insert("integer: \"") + + pynutil.insert('integer: "') + self.cardinal_numbers_default - + pynutil.insert("\"") + + pynutil.insert('"') ) # "03" -> remove leading zeros and verbalize leading_zeros = pynini.closure(pynini.cross("0", "")) - self.cardinal_numbers_with_leading_zeros = (leading_zeros + self.cardinal_numbers_default).optimize() + self.cardinal_numbers_with_leading_zeros = ( + leading_zeros + self.cardinal_numbers_default + ).optimize() # "123" -> "один два три" - single_digits_graph = pynini.string_file(get_abs_path("data/numbers/cardinals_nominative_case.tsv")).optimize() + single_digits_graph = pynini.string_file( + get_abs_path("data/numbers/cardinals_nominative_case.tsv") + ).optimize() single_digits_graph = pynini.compose(DAMO_DIGIT, single_digits_graph) - self.single_digits_graph = single_digits_graph + pynini.closure(insert_space + single_digits_graph) + self.single_digits_graph = single_digits_graph + pynini.closure( + insert_space + single_digits_graph + ) optional_quantity = pynini.string_file(get_abs_path("data/numbers/quantity.tsv")).optimize() - optional_quantity = pynutil.insert("quantity: \"") + optional_quantity + pynutil.insert("\"") + optional_quantity = pynutil.insert('quantity: "') + optional_quantity + pynutil.insert('"') optional_quantity = pynini.closure( - (pynutil.add_weight(pynini.accep(DAMO_SPACE), -0.1) | insert_space) + optional_quantity, 0, 1 + (pynutil.add_weight(pynini.accep(DAMO_SPACE), -0.1) | insert_space) + optional_quantity, + 0, + 1, ) serial_graph = self.get_serial_graph() final_graph = ( self.optional_graph_negative - + pynutil.insert("integer: \"") + + pynutil.insert('integer: "') + self.cardinal_numbers_with_leading_zeros - + pynutil.insert("\"") + + pynutil.insert('"') + optional_quantity ).optimize() final_graph = pynutil.add_weight(final_graph, -0.1) - final_graph |= pynutil.insert("integer: \"") + pynutil.add_weight(serial_graph, 10) + pynutil.insert("\"") + final_graph |= ( + pynutil.insert('integer: "') + + pynutil.add_weight(serial_graph, 10) + + pynutil.insert('"') + ) self.final_graph = final_graph # to cover cases "2-х" -> "двух" (this is not covered by ordinal endings) final_graph |= pynini.compose( - pynini.compose(DAMO_DIGIT ** (1, ...) + pynini.cross('-х', ''), final_graph), - DAMO_SIGMA + pynini.accep("х\"") + DAMO_SIGMA, + pynini.compose(DAMO_DIGIT ** (1, ...) + pynini.cross("-х", ""), final_graph), + DAMO_SIGMA + pynini.accep('х"') + DAMO_SIGMA, ) final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() - def get_cardinal_numbers(self, number_names: dict, alternative_formats: dict, mode: str = "all"): + def get_cardinal_numbers( + self, number_names: dict, alternative_formats: dict, mode: str = "all" + ): """Returns cardinal numbers names graph. Args: @@ -95,13 +109,13 @@ class CardinalFst(GraphFst): mode: "all" - to return graph that includes all Ru cases, "nominative" to return only the nominative form """ if mode == "all": - cardinal_names = number_names['cardinal_number_names'] + cardinal_names = number_names["cardinal_number_names"] elif mode == "nominative": - cardinal_names = number_names['cardinal_names_nominative'] + cardinal_names = number_names["cardinal_names_nominative"] else: - raise ValueError(f'{mode} is not supported.') - one_thousand_alternative = alternative_formats['one_thousand_alternative'] - separators = alternative_formats['separators'] + raise ValueError(f"{mode} is not supported.") + one_thousand_alternative = alternative_formats["one_thousand_alternative"] + separators = alternative_formats["separators"] cardinal_numbers = cardinal_names | pynini.compose(cardinal_names, one_thousand_alternative) cardinal_numbers = pynini.compose(separators, cardinal_numbers) @@ -126,11 +140,26 @@ class CardinalFst(GraphFst): # at least 1 alpha and 1 digit is present at_least_one_alpha_num = ( - DAMO_SIGMA + (RU_ALPHA | pynini.project(TO_CYRILLIC, "input")) + DAMO_SIGMA + DAMO_DIGIT + DAMO_SIGMA - ) | (DAMO_SIGMA + DAMO_DIGIT + DAMO_SIGMA + (RU_ALPHA | pynini.project(TO_CYRILLIC, "input")) + DAMO_SIGMA) + DAMO_SIGMA + + (RU_ALPHA | pynini.project(TO_CYRILLIC, "input")) + + DAMO_SIGMA + + DAMO_DIGIT + + DAMO_SIGMA + ) | ( + DAMO_SIGMA + + DAMO_DIGIT + + DAMO_SIGMA + + (RU_ALPHA | pynini.project(TO_CYRILLIC, "input")) + + DAMO_SIGMA + ) serial_graph = pynini.compose(at_least_one_alpha_num, serial_graph.optimize()).optimize() # numbers only with 2+ delimiters serial_graph |= ( - num_graph + delimiter + num_graph + delimiter + num_graph + pynini.closure(delimiter + num_graph) + num_graph + + delimiter + + num_graph + + delimiter + + num_graph + + pynini.closure(delimiter + num_graph) ).optimize() return serial_graph.optimize() diff --git a/fun_text_processing/text_normalization/ru/taggers/date.py b/fun_text_processing/text_normalization/ru/taggers/date.py index 2cc9c6981..904308eec 100644 --- a/fun_text_processing/text_normalization/ru/taggers/date.py +++ b/fun_text_processing/text_normalization/ru/taggers/date.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_DIGIT, @@ -15,7 +14,7 @@ from pynini.lib import pynutil class DateFst(GraphFst): """ - Finite state transducer for classifying date, e.g. + Finite state transducer for classifying date, e.g. "01.05" -> tokens { date { day: "первое мая" } } Args: @@ -28,13 +27,15 @@ class DateFst(GraphFst): super().__init__(name="date", kind="classify", deterministic=deterministic) # Ru format: DD-MM-YYYY or DD-MM-YY - month_abbr_to_names = pynini.string_file(get_abs_path("data/months/abbr_to_name.tsv")).optimize() + month_abbr_to_names = pynini.string_file( + get_abs_path("data/months/abbr_to_name.tsv") + ).optimize() delete_sep = pynutil.add_weight(pynini.cross(".", " "), 1.09) | pynutil.add_weight( pynini.cross(pynini.union("/", "-"), " "), 1.1 ) - numbers = number_names['ordinal_number_names'] + numbers = number_names["ordinal_number_names"] zero = (pynutil.add_weight(pynini.cross("0", ""), -0.1)) | ( pynutil.add_weight(pynini.cross("0", "ноль "), 0.1) @@ -45,26 +46,33 @@ class DateFst(GraphFst): day = zero_digit | digit_day day = pynini.compose( - day, pynini.difference(DAMO_SIGMA, DAMO_SIGMA + pynini.union("ой", "ая", "ых", "ые", "ыми")) + day, + pynini.difference(DAMO_SIGMA, DAMO_SIGMA + pynini.union("ой", "ая", "ых", "ые", "ыми")), ) - day = (pynutil.insert("day: \"") + day + pynutil.insert("\"")).optimize() + day = (pynutil.insert('day: "') + day + pynutil.insert('"')).optimize() digit_month = zero_digit | pynini.compose(pynini.accep("1") + DAMO_DIGIT, numbers) - month_number_to_abbr = pynini.string_file(get_abs_path("data/months/numbers.tsv")).optimize() + month_number_to_abbr = pynini.string_file( + get_abs_path("data/months/numbers.tsv") + ).optimize() month_number_to_abbr = ( ( - ((pynutil.add_weight(pynini.cross("0", ""), -0.1) | pynini.accep("1")) + DAMO_DIGIT) | DAMO_DIGIT + ((pynutil.add_weight(pynini.cross("0", ""), -0.1) | pynini.accep("1")) + DAMO_DIGIT) + | DAMO_DIGIT ).optimize() @ month_number_to_abbr ).optimize() month_name = ( - (month_number_to_abbr @ month_abbr_to_names) | pynutil.add_weight(month_abbr_to_names, 0.1) + (month_number_to_abbr @ month_abbr_to_names) + | pynutil.add_weight(month_abbr_to_names, 0.1) ).optimize() month = ( - pynutil.insert("month: \"") + (month_name | pynutil.add_weight(digit_month, 0.1)) + pynutil.insert("\"") + pynutil.insert('month: "') + + (month_name | pynutil.add_weight(digit_month, 0.1)) + + pynutil.insert('"') ).optimize() - year = pynini.compose(((DAMO_DIGIT ** 4) | (DAMO_DIGIT ** 2)), numbers).optimize() + year = pynini.compose(((DAMO_DIGIT**4) | (DAMO_DIGIT**2)), numbers).optimize() year |= zero_digit # reduce year options @@ -74,8 +82,12 @@ class DateFst(GraphFst): year, pynini.difference(DAMO_SIGMA, DAMO_SIGMA + pynini.union("ой", "ого")) ) | pynutil.add_weight(pynini.compose(year, DAMO_SIGMA + pynini.union("ой", "ого")), -0.1) - year_restrict1 = pynini.difference(DAMO_SIGMA, pynini.union(*year_wrong_beginning) + DAMO_SIGMA) - year_restrict2 = pynini.difference(DAMO_SIGMA, DAMO_SIGMA + pynini.union(*year_wrong_endings)) + year_restrict1 = pynini.difference( + DAMO_SIGMA, pynini.union(*year_wrong_beginning) + DAMO_SIGMA + ) + year_restrict2 = pynini.difference( + DAMO_SIGMA, DAMO_SIGMA + pynini.union(*year_wrong_endings) + ) year = pynini.compose(pynini.compose(year, year_restrict1), year_restrict2) year_word_singular = ["год", "года", "году", "годом", "годе"] @@ -83,11 +95,15 @@ class DateFst(GraphFst): year_word = pynini.cross("г.", pynini.union(*year_word_singular)) year_word |= pynini.cross("гг.", pynini.union(*year_word_plural)) - year_word = (pynutil.add_weight(insert_space, -0.1) | pynutil.add_weight(pynini.accep(" "), 0.1)) + year_word + year_word = ( + pynutil.add_weight(insert_space, -0.1) | pynutil.add_weight(pynini.accep(" "), 0.1) + ) + year_word - year_optional = pynutil.insert("year: \"") + year + pynini.closure(year_word, 0, 1) + pynutil.insert("\"") + year_optional = ( + pynutil.insert('year: "') + year + pynini.closure(year_word, 0, 1) + pynutil.insert('"') + ) year_optional = pynini.closure(delete_sep + year_optional, 0, 1).optimize() - year_only = pynutil.insert("year: \"") + year + year_word + pynutil.insert("\"") + year_only = pynutil.insert('year: "') + year + year_word + pynutil.insert('"') tagger_graph = (day + delete_sep + month + year_optional) | year_only @@ -95,29 +111,29 @@ class DateFst(GraphFst): day = ( pynutil.delete("day:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) month = ( pynutil.delete("month:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) year = ( pynutil.delete("year:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') ) year_optional = pynini.closure(delete_extra_space + year, 0, 1) graph_dmy = day + delete_extra_space + month + year_optional verbalizer_graph = (graph_dmy | year) + delete_space self.final_graph = pynini.compose(tagger_graph, verbalizer_graph).optimize() - self.fst = pynutil.insert("day: \"") + self.final_graph + pynutil.insert("\"") + self.fst = pynutil.insert('day: "') + self.final_graph + pynutil.insert('"') self.fst = self.add_tokens(self.fst).optimize() diff --git a/fun_text_processing/text_normalization/ru/taggers/decimals.py b/fun_text_processing/text_normalization/ru/taggers/decimals.py index 95a654c0f..7d7d73f6d 100644 --- a/fun_text_processing/text_normalization/ru/taggers/decimals.py +++ b/fun_text_processing/text_normalization/ru/taggers/decimals.py @@ -1,8 +1,12 @@ - from collections import defaultdict import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_DIGIT, DAMO_SPACE, GraphFst, insert_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_DIGIT, + DAMO_SPACE, + GraphFst, + insert_space, +) from fun_text_processing.text_normalization.en.utils import load_labels from fun_text_processing.text_normalization.ru.utils import get_abs_path from pynini.lib import pynutil @@ -29,14 +33,16 @@ def prepare_labels_for_insertion(file_path: str): for k in mapping: mapping[k] = ( insert_space - + pynini.union(*[pynutil.add_weight(pynutil.insert(end), weight) for end, weight in mapping[k]]) + + pynini.union( + *[pynutil.add_weight(pynutil.insert(end), weight) for end, weight in mapping[k]] + ) ).optimize() return mapping class DecimalFst(GraphFst): """ - Finite state transducer for classifying decimal, e.g. + Finite state transducer for classifying decimal, e.g. "1,08" -> tokens { decimal { integer_part: "одно целая" fractional_part: "восемь сотых} } Args: @@ -51,43 +57,57 @@ class DecimalFst(GraphFst): integer_part = cardinal.cardinal_numbers_default cardinal_numbers_with_leading_zeros = cardinal.cardinal_numbers_with_leading_zeros - delimiter_map = prepare_labels_for_insertion(get_abs_path("data/numbers/decimal_delimiter.tsv")) + delimiter_map = prepare_labels_for_insertion( + get_abs_path("data/numbers/decimal_delimiter.tsv") + ) delimiter = ( pynini.cross(",", "") - + delimiter_map['@@decimal_delimiter@@'] + + delimiter_map["@@decimal_delimiter@@"] + pynini.closure(pynutil.add_weight(pynutil.insert(" и"), 0.5), 0, 1) ).optimize() - decimal_endings_map = prepare_labels_for_insertion(get_abs_path("data/numbers/decimal_endings.tsv")) + decimal_endings_map = prepare_labels_for_insertion( + get_abs_path("data/numbers/decimal_endings.tsv") + ) self.integer_part = integer_part + delimiter - graph_integer = pynutil.insert("integer_part: \"") + self.integer_part + pynutil.insert("\"") + graph_integer = pynutil.insert('integer_part: "') + self.integer_part + pynutil.insert('"') - graph_fractional = DAMO_DIGIT @ cardinal_numbers_with_leading_zeros + decimal_endings_map['10'] - graph_fractional |= (DAMO_DIGIT + DAMO_DIGIT) @ cardinal_numbers_with_leading_zeros + decimal_endings_map[ - '100' - ] + graph_fractional = ( + DAMO_DIGIT @ cardinal_numbers_with_leading_zeros + decimal_endings_map["10"] + ) + graph_fractional |= ( + DAMO_DIGIT + DAMO_DIGIT + ) @ cardinal_numbers_with_leading_zeros + decimal_endings_map["100"] graph_fractional |= ( DAMO_DIGIT + DAMO_DIGIT + DAMO_DIGIT - ) @ cardinal_numbers_with_leading_zeros + decimal_endings_map['1000'] + ) @ cardinal_numbers_with_leading_zeros + decimal_endings_map["1000"] graph_fractional |= ( DAMO_DIGIT + DAMO_DIGIT + DAMO_DIGIT + DAMO_DIGIT - ) @ cardinal_numbers_with_leading_zeros + decimal_endings_map['10000'] + ) @ cardinal_numbers_with_leading_zeros + decimal_endings_map["10000"] - self.optional_quantity = pynini.string_file(get_abs_path("data/numbers/quantity.tsv")).optimize() + self.optional_quantity = pynini.string_file( + get_abs_path("data/numbers/quantity.tsv") + ).optimize() self.graph_fractional = graph_fractional - graph_fractional = pynutil.insert("fractional_part: \"") + graph_fractional + pynutil.insert("\"") + graph_fractional = ( + pynutil.insert('fractional_part: "') + graph_fractional + pynutil.insert('"') + ) optional_quantity = pynini.closure( (pynutil.add_weight(pynini.accep(DAMO_SPACE), -0.1) | insert_space) - + pynutil.insert("quantity: \"") + + pynutil.insert('quantity: "') + self.optional_quantity - + pynutil.insert("\""), + + pynutil.insert('"'), 0, 1, ) self.final_graph = ( - cardinal.optional_graph_negative + graph_integer + insert_space + graph_fractional + optional_quantity + cardinal.optional_graph_negative + + graph_integer + + insert_space + + graph_fractional + + optional_quantity ) self.final_graph = self.add_tokens(self.final_graph) diff --git a/fun_text_processing/text_normalization/ru/taggers/electronic.py b/fun_text_processing/text_normalization/ru/taggers/electronic.py index fff14335b..f1d0eeb5d 100644 --- a/fun_text_processing/text_normalization/ru/taggers/electronic.py +++ b/fun_text_processing/text_normalization/ru/taggers/electronic.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_ALPHA, @@ -29,32 +27,34 @@ class ElectronicFst(GraphFst): # tagger accepted_symbols = [] - with open(get_abs_path("data/electronic/symbols.tsv"), 'r', encoding='utf-8') as f: + with open(get_abs_path("data/electronic/symbols.tsv"), "r", encoding="utf-8") as f: for line in f: - symbol, _ = line.split('\t') + symbol, _ = line.split("\t") accepted_symbols.append(pynini.accep(symbol)) username = ( - pynutil.insert("username: \"") + pynutil.insert('username: "') + DAMO_ALPHA + pynini.closure(DAMO_ALPHA | DAMO_DIGIT | pynini.union(*accepted_symbols)) - + pynutil.insert("\"") - + pynini.cross('@', ' ') + + pynutil.insert('"') + + pynini.cross("@", " ") ) domain_graph = ( DAMO_ALPHA - + (pynini.closure(DAMO_ALPHA | DAMO_DIGIT | pynini.accep('-') | pynini.accep('.'))) + + (pynini.closure(DAMO_ALPHA | DAMO_DIGIT | pynini.accep("-") | pynini.accep("."))) + (DAMO_ALPHA | DAMO_DIGIT) ) - domain_graph = pynutil.insert("domain: \"") + domain_graph + pynutil.insert("\"") + domain_graph = pynutil.insert('domain: "') + domain_graph + pynutil.insert('"') tagger_graph = (username + domain_graph).optimize() # verbalizer - graph_digit = pynini.string_file(get_abs_path("data/numbers/digits_nominative_case.tsv")).optimize() + graph_digit = pynini.string_file( + get_abs_path("data/numbers/digits_nominative_case.tsv") + ).optimize() graph_symbols = pynini.string_file(get_abs_path("data/electronic/symbols.tsv")).optimize() user_name = ( pynutil.delete("username:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + ( pynini.closure( pynutil.add_weight(graph_digit + insert_space, 1.09) @@ -62,7 +62,7 @@ class ElectronicFst(GraphFst): | pynutil.add_weight(DAMO_NOT_QUOTE + insert_space, 1.1) ) ) - + pynutil.delete("\"") + + pynutil.delete('"') ) domain_default = ( @@ -77,22 +77,37 @@ class ElectronicFst(GraphFst): + pynini.closure(graph_symbols + insert_space) + pynini.closure((graph_digit | DAMO_ALPHA) + insert_space, 1) ) - server_common = pynini.string_file(get_abs_path("data/electronic/server_name.tsv")) + insert_space - domain_common = pynini.cross(".", "точка ") + pynini.string_file(get_abs_path("data/electronic/domain.tsv")) + server_common = ( + pynini.string_file(get_abs_path("data/electronic/server_name.tsv")) + insert_space + ) + domain_common = pynini.cross(".", "точка ") + pynini.string_file( + get_abs_path("data/electronic/domain.tsv") + ) domain = ( pynutil.delete("domain:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + (pynutil.add_weight(server_common, 1.09) | pynutil.add_weight(server_default, 1.1)) + (pynutil.add_weight(domain_common, 1.09) | pynutil.add_weight(domain_default, 1.1)) + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') ) - graph = user_name + delete_space + pynutil.insert("собака ") + delete_space + domain + delete_space + graph = ( + user_name + + delete_space + + pynutil.insert("собака ") + + delete_space + + domain + + delete_space + ) # replace all latin letters with their Ru verbalization - verbalizer_graph = (graph.optimize() @ (pynini.closure(TO_CYRILLIC | RU_ALPHA | pynini.accep(" ")))).optimize() + verbalizer_graph = ( + graph.optimize() @ (pynini.closure(TO_CYRILLIC | RU_ALPHA | pynini.accep(" "))) + ).optimize() verbalizer_graph = verbalizer_graph.optimize() self.final_graph = (tagger_graph @ verbalizer_graph).optimize() - self.fst = self.add_tokens(pynutil.insert("username: \"") + self.final_graph + pynutil.insert("\"")).optimize() + self.fst = self.add_tokens( + pynutil.insert('username: "') + self.final_graph + pynutil.insert('"') + ).optimize() diff --git a/fun_text_processing/text_normalization/ru/taggers/measure.py b/fun_text_processing/text_normalization/ru/taggers/measure.py index 01f3c7b94..b1ff90df4 100644 --- a/fun_text_processing/text_normalization/ru/taggers/measure.py +++ b/fun_text_processing/text_normalization/ru/taggers/measure.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_NON_BREAKING_SPACE, @@ -32,7 +31,11 @@ class MeasureFst(GraphFst): # adding weight to make sure the space is preserved for ITN delete_space = pynini.closure( - pynutil.add_weight(pynutil.delete(pynini.union(DAMO_SPACE, DAMO_NON_BREAKING_SPACE)), -1), 0, 1 + pynutil.add_weight( + pynutil.delete(pynini.union(DAMO_SPACE, DAMO_NON_BREAKING_SPACE)), -1 + ), + 0, + 1, ) cardinal_graph = cardinal.cardinal_numbers_default @@ -44,56 +47,60 @@ class MeasureFst(GraphFst): pynutil.add_weight(pynutil.insert(DAMO_NON_BREAKING_SPACE), -0.1) | pynutil.add_weight(pynutil.insert(DAMO_SPACE), 0.1) ).optimize() - slash_unit = (pynini.cross("/", "в") | pynini.cross("/", "за")) + space_for_units + graph_unit + slash_unit = ( + (pynini.cross("/", "в") | pynini.cross("/", "за")) + space_for_units + graph_unit + ) unit_slash_unit = pynutil.add_weight(graph_unit + space_for_units + slash_unit, -0.1) - default_units = pynutil.insert("units: \"") + (graph_unit | unit_slash_unit) + pynutil.insert("\"") - slash_units = pynutil.insert("units: \"") + slash_unit + pynutil.insert("\"") + default_units = ( + pynutil.insert('units: "') + (graph_unit | unit_slash_unit) + pynutil.insert('"') + ) + slash_units = pynutil.insert('units: "') + slash_unit + pynutil.insert('"') subgraph_decimal = decimal.final_graph + ((delete_space + default_units) | slash_units) cardinal_space = ( pynutil.insert("cardinal { ") + optional_graph_negative - + pynutil.insert("integer: \"") + + pynutil.insert('integer: "') + cardinal_graph + ( - (delete_space + pynutil.insert("\"") + pynutil.insert(" } ") + default_units) - | (pynutil.insert("\"") + pynutil.insert(" } ") + slash_units) + (delete_space + pynutil.insert('"') + pynutil.insert(" } ") + default_units) + | (pynutil.insert('"') + pynutil.insert(" } ") + slash_units) ) ) cardinal_optional_dash_alpha = ( - pynutil.insert("cardinal { integer: \"") + pynutil.insert('cardinal { integer: "') + cardinal_graph - + pynini.closure(pynini.cross('-', ''), 0, 1) - + pynutil.insert("\" } units: \"") + + pynini.closure(pynini.cross("-", ""), 0, 1) + + pynutil.insert('" } units: "') + pynini.closure(RU_ALPHA, 1) - + pynutil.insert("\"") + + pynutil.insert('"') ) alpha_optional_dash_cardinal = ( - pynutil.insert("units: \"") + pynutil.insert('units: "') + pynini.closure(RU_ALPHA, 1) - + pynini.closure(pynini.cross('-', ''), 0, 1) - + pynutil.insert("\"") - + pynutil.insert(" cardinal { integer: \"") + + pynini.closure(pynini.cross("-", ""), 0, 1) + + pynutil.insert('"') + + pynutil.insert(' cardinal { integer: "') + cardinal_graph_nominative - + pynutil.insert("\" } preserve_order: true") + + pynutil.insert('" } preserve_order: true') ) decimal_dash_alpha = ( decimal.final_graph - + pynini.cross('-', '') - + pynutil.insert(" units: \"") + + pynini.cross("-", "") + + pynutil.insert(' units: "') + pynini.closure(RU_ALPHA, 1) - + pynutil.insert("\"") + + pynutil.insert('"') ) alpha_dash_decimal = ( - pynutil.insert("units: \"") + pynutil.insert('units: "') + pynini.closure(RU_ALPHA, 1) - + pynini.cross('-', '') - + pynutil.insert("\" ") + + pynini.cross("-", "") + + pynutil.insert('" ') + decimal.final_graph + pynutil.insert(" preserve_order: true") ) @@ -109,31 +116,48 @@ class MeasureFst(GraphFst): ).optimize() # verbalizer - unit = pynutil.delete("units: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + delete_space + unit = ( + pynutil.delete('units: "') + + pynini.closure(DAMO_NOT_QUOTE, 1) + + pynutil.delete('"') + + delete_space + ) - optional_sign = pynini.closure(pynini.cross("negative: \"true\" ", "минус "), 0, 1) - integer = pynutil.delete(" \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + optional_sign = pynini.closure(pynini.cross('negative: "true" ', "минус "), 0, 1) + integer = pynutil.delete(' "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') integer_part = pynutil.delete("integer_part:") + integer fractional_part = pynutil.delete("fractional_part:") + integer optional_quantity_part = pynini.closure( pynini.accep(" ") - + pynutil.delete("quantity: \"") + + pynutil.delete('quantity: "') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\""), + + pynutil.delete('"'), 0, 1, ) - graph_decimal = optional_sign + integer_part + pynini.accep(" ") + fractional_part + optional_quantity_part + graph_decimal = ( + optional_sign + + integer_part + + pynini.accep(" ") + + fractional_part + + optional_quantity_part + ) - graph_decimal = pynutil.delete("decimal {") + delete_space + graph_decimal + delete_space + pynutil.delete("}") + graph_decimal = ( + pynutil.delete("decimal {") + + delete_space + + graph_decimal + + delete_space + + pynutil.delete("}") + ) graph_cardinal = ( pynutil.delete("cardinal {") + delete_space + optional_sign - + pynutil.delete("integer: \"") + + pynutil.delete('integer: "') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space + pynutil.delete("}") ) @@ -141,7 +165,9 @@ class MeasureFst(GraphFst): verbalizer_graph = (graph_cardinal | graph_decimal) + delete_space + insert_space + unit # SH adds "preserve_order: true" by default - preserve_order = pynutil.delete("preserve_order:") + delete_space + pynutil.delete("true") + delete_space + preserve_order = ( + pynutil.delete("preserve_order:") + delete_space + pynutil.delete("true") + delete_space + ) verbalizer_graph |= ( unit + insert_space @@ -153,5 +179,5 @@ class MeasureFst(GraphFst): final_graph = (tagger_graph @ verbalizer_graph).optimize() self.fst = self.add_tokens( - pynutil.insert("cardinal { integer: \"") + final_graph + pynutil.insert("\" }") + pynutil.insert('cardinal { integer: "') + final_graph + pynutil.insert('" }') ).optimize() diff --git a/fun_text_processing/text_normalization/ru/taggers/money.py b/fun_text_processing/text_normalization/ru/taggers/money.py index 025a1d838..2003d4548 100644 --- a/fun_text_processing/text_normalization/ru/taggers/money.py +++ b/fun_text_processing/text_normalization/ru/taggers/money.py @@ -1,6 +1,10 @@ - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, DAMO_SIGMA, DAMO_SPACE, GraphFst +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + DAMO_SIGMA, + DAMO_SPACE, + GraphFst, +) from fun_text_processing.text_normalization.ru.utils import get_abs_path from pynini.lib import pynutil @@ -26,21 +30,30 @@ class MoneyFst(GraphFst): unit_plural = pynini.string_file(get_abs_path("data/currency/currency_plural.tsv")) # adding weight to make sure the space is preserved for ITN - optional_delimiter = pynini.closure(pynutil.add_weight(pynini.cross(DAMO_SPACE, ""), -100), 0, 1) + optional_delimiter = pynini.closure( + pynutil.add_weight(pynini.cross(DAMO_SPACE, ""), -100), 0, 1 + ) graph_unit_singular = ( - optional_delimiter + pynutil.insert(" currency: \"") + unit_singular + pynutil.insert("\"") + optional_delimiter + + pynutil.insert(' currency: "') + + unit_singular + + pynutil.insert('"') + ) + graph_unit_plural = ( + optional_delimiter + pynutil.insert(' currency: "') + unit_plural + pynutil.insert('"') ) - graph_unit_plural = optional_delimiter + pynutil.insert(" currency: \"") + unit_plural + pynutil.insert("\"") one = pynini.compose(pynini.accep("1"), cardinal_graph).optimize() - singular_graph = pynutil.insert("integer_part: \"") + one + pynutil.insert("\"") + graph_unit_singular + singular_graph = ( + pynutil.insert('integer_part: "') + one + pynutil.insert('"') + graph_unit_singular + ) graph_decimal = decimal_graph + graph_unit_plural graph_integer = ( - pynutil.insert("integer_part: \"") + pynutil.insert('integer_part: "') + ((DAMO_SIGMA - "1") @ cardinal_graph) - + pynutil.insert("\"") + + pynutil.insert('"') + (graph_unit_plural) ) @@ -48,24 +61,26 @@ class MoneyFst(GraphFst): tagger_graph = (graph_integer.optimize() | graph_decimal.optimize()).optimize() # verbalizer - integer = pynutil.delete("\"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + integer = pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') integer_part = pynutil.delete("integer_part: ") + integer unit = ( pynutil.delete("currency: ") - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) unit = pynini.accep(DAMO_SPACE) + unit verbalizer_graph_cardinal = (integer_part + unit).optimize() fractional_part = pynutil.delete("fractional_part: ") + integer - optional_quantity = pynini.closure(pynini.accep(DAMO_SPACE) + pynutil.delete("quantity: ") + integer, 0, 1) + optional_quantity = pynini.closure( + pynini.accep(DAMO_SPACE) + pynutil.delete("quantity: ") + integer, 0, 1 + ) verbalizer_graph_decimal = ( - pynutil.delete('decimal { ') + pynutil.delete("decimal { ") + integer_part + pynini.accep(" ") + fractional_part @@ -78,5 +93,5 @@ class MoneyFst(GraphFst): self.final_graph = (tagger_graph @ verbalizer_graph).optimize() self.fst = self.add_tokens( - pynutil.insert("integer_part: \"") + self.final_graph + pynutil.insert("\"") + pynutil.insert('integer_part: "') + self.final_graph + pynutil.insert('"') ).optimize() diff --git a/fun_text_processing/text_normalization/ru/taggers/number_names.py b/fun_text_processing/text_normalization/ru/taggers/number_names.py index 98fddafd4..f033e8977 100644 --- a/fun_text_processing/text_normalization/ru/taggers/number_names.py +++ b/fun_text_processing/text_normalization/ru/taggers/number_names.py @@ -1,5 +1,3 @@ - - # Copyright 2017 Google Inc. @@ -30,18 +28,26 @@ def get_number_names(): and 2) Ng, A. H., Gorman, K., and Sproat, R. 2017. Minimally supervised written-to-spoken text normalization. In ASRU, pages 665-670. """ - a = pynini.Far(get_abs_path('data/utils/util_arithmetic.far'), mode='r') - d = a['DELTA_STAR'] - f = a['IARITHMETIC_RESTRICTED'] - g = pynini.Fst.read(get_abs_path('data/utils/g.fst')) + a = pynini.Far(get_abs_path("data/utils/util_arithmetic.far"), mode="r") + d = a["DELTA_STAR"] + f = a["IARITHMETIC_RESTRICTED"] + g = pynini.Fst.read(get_abs_path("data/utils/g.fst")) fg = (d @ (f @ (f @ (f @ g).optimize()).optimize()).optimize()).optimize() assert rewrite.top_rewrite("230", fg) == "(+ 200 30 +)" # Compiles lexicon transducers (L). - cardinal_name_nominative = pynini.string_file(get_abs_path("data/numbers/1_cardinals_nominative.tsv")).optimize() - cardinal_name_genitive = pynini.string_file(get_abs_path("data/numbers/2_cardinals_genitive.tsv")).optimize() - cardinal_name_dative = pynini.string_file(get_abs_path("data/numbers/3_cardinals_dative.tsv")).optimize() - cardinal_name_accusative = pynini.string_file(get_abs_path("data/numbers/4_cardinals_accusative.tsv")).optimize() + cardinal_name_nominative = pynini.string_file( + get_abs_path("data/numbers/1_cardinals_nominative.tsv") + ).optimize() + cardinal_name_genitive = pynini.string_file( + get_abs_path("data/numbers/2_cardinals_genitive.tsv") + ).optimize() + cardinal_name_dative = pynini.string_file( + get_abs_path("data/numbers/3_cardinals_dative.tsv") + ).optimize() + cardinal_name_accusative = pynini.string_file( + get_abs_path("data/numbers/4_cardinals_accusative.tsv") + ).optimize() cardinal_name_instrumental = pynini.string_file( get_abs_path("data/numbers/5_cardinals_instrumental.tsv") ).optimize() @@ -77,18 +83,22 @@ def get_number_names(): complex_numbers |= pynini.cross(f"(* {number} 1000 *)", f"{number}000") complex_numbers = ( - DAMO_SIGMA + pynutil.add_weight(complex_numbers, -1) + pynini.closure(pynini.union(" ", ")", "(", "+", "*")) + DAMO_SIGMA + + pynutil.add_weight(complex_numbers, -1) + + pynini.closure(pynini.union(" ", ")", "(", "+", "*")) ).optimize() fg_ordinal = pynutil.add_weight(pynini.compose(fg, complex_numbers), -1) | fg ordinal_name = pynini.string_file(get_abs_path("data/numbers/ordinals.tsv")) - ordinal_l = (pynini.closure(cardinal_name_nominative + pynini.accep(" ")) + ordinal_name).optimize() + ordinal_l = ( + pynini.closure(cardinal_name_nominative + pynini.accep(" ")) + ordinal_name + ).optimize() # Composes L with the leaf transducer (P), then composes that with FG. - p = a['LEAVES'] + p = a["LEAVES"] number_names = {} - number_names['ordinal_number_names'] = (fg_ordinal @ (p @ ordinal_l)).optimize() - number_names['cardinal_number_names'] = (fg @ (p @ cardinal_l)).optimize() - number_names['cardinal_names_nominative'] = (fg @ (p @ cardinal_names_nominative_l)).optimize() + number_names["ordinal_number_names"] = (fg_ordinal @ (p @ ordinal_l)).optimize() + number_names["cardinal_number_names"] = (fg @ (p @ cardinal_l)).optimize() + number_names["cardinal_names_nominative"] = (fg @ (p @ cardinal_names_nominative_l)).optimize() return number_names @@ -96,7 +106,7 @@ def get_alternative_formats(): """ Utils to get alternative formats for numbers. """ - one_alternatives = load_labels(get_abs_path('data/numbers/cardinals_alternatives.tsv')) + one_alternatives = load_labels(get_abs_path("data/numbers/cardinals_alternatives.tsv")) one_thousand_map = [] for k in one_alternatives: default, alternative = k @@ -108,21 +118,21 @@ def get_alternative_formats(): # Adapted from # https://github.com/google/TextNormalizationCoveringGrammars/blob/master/src/universal/thousands_punct.grm # Specifies common ways of delimiting thousands in digit strings. - t = pynini.Far(get_abs_path('data/utils/universal_thousands_punct.far')) + t = pynini.Far(get_abs_path("data/utils/universal_thousands_punct.far")) separators = ( - pynutil.add_weight(t['dot_thousands'], 0.1) - | pynutil.add_weight(t['no_delimiter'], -0.1) - | pynutil.add_weight(t['space_thousands'], 0.1) + pynutil.add_weight(t["dot_thousands"], 0.1) + | pynutil.add_weight(t["no_delimiter"], -0.1) + | pynutil.add_weight(t["space_thousands"], 0.1) ) alternative_formats = {} - alternative_formats['one_thousand_alternative'] = one_thousand_alternative.optimize() - alternative_formats['separators'] = separators.optimize() + alternative_formats["one_thousand_alternative"] = one_thousand_alternative.optimize() + alternative_formats["separators"] = separators.optimize() return alternative_formats -if __name__ == '__main__': +if __name__ == "__main__": from fun_text_processing.text_normalization.en.graph_utils import generator_main numbers = get_number_names() for k, v in numbers.items(): - generator_main(f'{k}.far', {k: v}) + generator_main(f"{k}.far", {k: v}) diff --git a/fun_text_processing/text_normalization/ru/taggers/ordinal.py b/fun_text_processing/text_normalization/ru/taggers/ordinal.py index 00e3462bf..2e0b39372 100644 --- a/fun_text_processing/text_normalization/ru/taggers/ordinal.py +++ b/fun_text_processing/text_normalization/ru/taggers/ordinal.py @@ -1,4 +1,3 @@ - # Adapted from https://github.com/google/TextNormalizationCoveringGrammars # Russian minimally supervised number grammar. @@ -11,7 +10,7 @@ from pynini.lib import pynutil class OrdinalFst(GraphFst): """ - Finite state transducer for classifying cardinals, e.g. + Finite state transducer for classifying cardinals, e.g. "2" -> ordinal { integer: "второе" } } Args: @@ -24,10 +23,10 @@ class OrdinalFst(GraphFst): def __init__(self, number_names: dict, alternative_formats: dict, deterministic=False): super().__init__(name="ordinal", kind="classify", deterministic=deterministic) - one_thousand_alternative = alternative_formats['one_thousand_alternative'] - separators = alternative_formats['separators'] + one_thousand_alternative = alternative_formats["one_thousand_alternative"] + separators = alternative_formats["separators"] - ordinal = number_names['ordinal_number_names'] + ordinal = number_names["ordinal_number_names"] ordinal |= ordinal @ one_thousand_alternative ordinal_numbers = separators @ ordinal @@ -48,6 +47,6 @@ class OrdinalFst(GraphFst): self.ordinal_numbers_with_leading_zeros = (leading_zeros + ordinal_numbers).optimize() final_graph = (ordinal_numbers | ordinal_numbers_marked).optimize() - final_graph = pynutil.insert("integer: \"") + final_graph + pynutil.insert("\"") + final_graph = pynutil.insert('integer: "') + final_graph + pynutil.insert('"') final_graph = self.add_tokens(final_graph) self.fst = final_graph.optimize() diff --git a/fun_text_processing/text_normalization/ru/taggers/telephone.py b/fun_text_processing/text_normalization/ru/taggers/telephone.py index c884b0fe5..6a1563e34 100644 --- a/fun_text_processing/text_normalization/ru/taggers/telephone.py +++ b/fun_text_processing/text_normalization/ru/taggers/telephone.py @@ -1,15 +1,19 @@ - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_DIGIT, GraphFst, delete_space, insert_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_DIGIT, + GraphFst, + delete_space, + insert_space, +) from fun_text_processing.text_normalization.ru.alphabet import RU_ALPHA_OR_SPACE from pynini.lib import pynutil class TelephoneFst(GraphFst): """ - Finite state transducer for classifying telephone, which includes country code, number part and extension + Finite state transducer for classifying telephone, which includes country code, number part and extension - E.g + E.g "8-913-983-56-01" -> telephone { number_part: "восемь девятьсот тринадцать девятьсот восемьдесят три пятьдесят шесть ноль один" } Args: @@ -25,41 +29,43 @@ class TelephoneFst(GraphFst): number = number_names["cardinal_names_nominative"] country_code = ( - pynutil.insert("country_code: \"") + pynutil.insert('country_code: "') + pynini.closure(pynutil.add_weight(pynutil.delete("+"), 0.1), 0, 1) + number + separator - + pynutil.insert("\"") + + pynutil.insert('"') ) optional_country_code = pynini.closure(country_code + insert_space, 0, 1) number_part = ( - DAMO_DIGIT ** 3 @ number + DAMO_DIGIT**3 @ number + separator - + DAMO_DIGIT ** 3 @ number + + DAMO_DIGIT**3 @ number + separator - + DAMO_DIGIT ** 2 @ number + + DAMO_DIGIT**2 @ number + separator - + DAMO_DIGIT ** 2 @ (pynini.closure(pynini.cross("0", "ноль ")) + number) + + DAMO_DIGIT**2 @ (pynini.closure(pynini.cross("0", "ноль ")) + number) ) - number_part = pynutil.insert("number_part: \"") + number_part + pynutil.insert("\"") + number_part = pynutil.insert('number_part: "') + number_part + pynutil.insert('"') tagger_graph = (optional_country_code + number_part).optimize() # verbalizer verbalizer_graph = pynini.closure( - pynutil.delete("country_code: \"") + pynutil.delete('country_code: "') + pynini.closure(RU_ALPHA_OR_SPACE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space, 0, 1, ) verbalizer_graph += ( - pynutil.delete("number_part: \"") + pynini.closure(RU_ALPHA_OR_SPACE, 1) + pynutil.delete("\"") + pynutil.delete('number_part: "') + + pynini.closure(RU_ALPHA_OR_SPACE, 1) + + pynutil.delete('"') ) verbalizer_graph = verbalizer_graph.optimize() self.final_graph = (tagger_graph @ verbalizer_graph).optimize() self.fst = self.add_tokens( - pynutil.insert("number_part: \"") + self.final_graph + pynutil.insert("\"") + pynutil.insert('number_part: "') + self.final_graph + pynutil.insert('"') ).optimize() diff --git a/fun_text_processing/text_normalization/ru/taggers/time.py b/fun_text_processing/text_normalization/ru/taggers/time.py index 6ab728d02..2d5de5a43 100644 --- a/fun_text_processing/text_normalization/ru/taggers/time.py +++ b/fun_text_processing/text_normalization/ru/taggers/time.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import GraphFst from fun_text_processing.text_normalization.ru.utils import get_abs_path @@ -10,7 +8,7 @@ class TimeFst(GraphFst): """ Finite state transducer for classifying time, e.g. "02:15" -> time { hours: "два часа пятнадцать минут" } - + Args: number_names: number_names for cardinal and ordinal numbers deterministic: if True will provide a single transduction option, @@ -20,18 +18,26 @@ class TimeFst(GraphFst): def __init__(self, number_names: dict, deterministic: bool = True): super().__init__(name="time", kind="classify", deterministic=deterministic) - increment_hour_ordinal = pynini.string_file(get_abs_path("data/time/increment_hour_ordinal.tsv")) - increment_hour_cardinal = pynini.string_file(get_abs_path("data/time/increment_hour_cardinal.tsv")) + increment_hour_ordinal = pynini.string_file( + get_abs_path("data/time/increment_hour_ordinal.tsv") + ) + increment_hour_cardinal = pynini.string_file( + get_abs_path("data/time/increment_hour_cardinal.tsv") + ) convert_hour = pynini.string_file(get_abs_path("data/time/time_convert.tsv")) - number = pynini.closure(pynini.cross("0", ""), 0, 1) + number_names['cardinal_names_nominative'] + number = ( + pynini.closure(pynini.cross("0", ""), 0, 1) + number_names["cardinal_names_nominative"] + ) hour_options = pynini.project(increment_hour_ordinal, "input") hour_options = hour_options | pynini.project(convert_hour, "output") hour_exeption_ends_with_one = pynini.union(*["01", "21"]) hour_exeption_ends_rest = pynini.union(*["02", "03", "04", "22", "23"]) hour_other = ( - pynini.difference(hour_options, pynini.union(hour_exeption_ends_with_one, hour_exeption_ends_rest)) + pynini.difference( + hour_options, pynini.union(hour_exeption_ends_with_one, hour_exeption_ends_rest) + ) ).optimize() hour = hour_exeption_ends_with_one @ number + pynutil.insert(" час") @@ -56,41 +62,52 @@ class TimeFst(GraphFst): self.minutes = minutes.optimize() # 17:15 -> "семнадцать часов и пятнадцать минут" hm = ( - pynutil.insert("hours: \"") + pynutil.insert('hours: "') + hour.optimize() - + pynutil.insert("\"") - + (pynini.cross(":", " ") + pynutil.insert("minutes: \"") + optional_and + minutes.optimize()) - + pynutil.insert("\"") + + pynutil.insert('"') + + ( + pynini.cross(":", " ") + + pynutil.insert('minutes: "') + + optional_and + + minutes.optimize() + ) + + pynutil.insert('"') + pynutil.insert(" preserve_order: true") ) - h = pynutil.insert("hours: \"") + hour + pynutil.insert("\"") + pynutil.delete(":00") + h = pynutil.insert('hours: "') + hour + pynutil.insert('"') + pynutil.delete(":00") self.graph_preserve_order = (hm | h).optimize() # 17:15 -> "пятнадцать минут шестого" # Requires permutations for the correct verbalization - self.increment_hour_ordinal = pynini.compose(hour_options, increment_hour_ordinal).optimize() + self.increment_hour_ordinal = pynini.compose( + hour_options, increment_hour_ordinal + ).optimize() m_next_h = ( - pynutil.insert("hours: \"") + pynutil.insert('hours: "') + self.increment_hour_ordinal - + pynutil.insert("\"") + + pynutil.insert('"') + pynini.cross(":", " ") - + pynutil.insert("minutes: \"") + + pynutil.insert('minutes: "') + minutes - + pynutil.insert("\"") + + pynutil.insert('"') ) # 17:45 -> "без пятнадцати минут шесть" # Requires permutations for the correct verbalization - self.mins_to_h = pynini.string_file(get_abs_path("data/time/minutes_to_hour.tsv")).optimize() - self.increment_hour_cardinal = pynini.compose(hour_options, increment_hour_cardinal).optimize() + self.mins_to_h = pynini.string_file( + get_abs_path("data/time/minutes_to_hour.tsv") + ).optimize() + self.increment_hour_cardinal = pynini.compose( + hour_options, increment_hour_cardinal + ).optimize() m_to_h = ( - pynutil.insert("hours: \"") + pynutil.insert('hours: "') + self.increment_hour_cardinal - + pynutil.insert("\"") + + pynutil.insert('"') + pynini.cross(":", " ") - + pynutil.insert("minutes: \"без ") + + pynutil.insert('minutes: "без ') + self.mins_to_h - + pynutil.insert("\"") + + pynutil.insert('"') ) self.final_graph = m_next_h | self.graph_preserve_order | m_to_h diff --git a/fun_text_processing/text_normalization/ru/taggers/tokenize_and_classify.py b/fun_text_processing/text_normalization/ru/taggers/tokenize_and_classify.py index 2bc522718..43d7e6601 100644 --- a/fun_text_processing/text_normalization/ru/taggers/tokenize_and_classify.py +++ b/fun_text_processing/text_normalization/ru/taggers/tokenize_and_classify.py @@ -1,5 +1,3 @@ - - import os import pynini @@ -16,7 +14,10 @@ from fun_text_processing.text_normalization.ru.taggers.decimals import DecimalFs from fun_text_processing.text_normalization.ru.taggers.electronic import ElectronicFst from fun_text_processing.text_normalization.ru.taggers.measure import MeasureFst from fun_text_processing.text_normalization.ru.taggers.money import MoneyFst -from fun_text_processing.text_normalization.ru.taggers.number_names import get_alternative_formats, get_number_names +from fun_text_processing.text_normalization.ru.taggers.number_names import ( + get_alternative_formats, + get_number_names, +) from fun_text_processing.text_normalization.ru.taggers.ordinal import OrdinalFst from fun_text_processing.text_normalization.ru.taggers.telephone import TelephoneFst from fun_text_processing.text_normalization.ru.taggers.time import TimeFst @@ -53,7 +54,7 @@ class ClassifyFst(GraphFst): super().__init__(name="tokenize_and_classify", kind="classify", deterministic=deterministic) if deterministic: raise ValueError( - 'Ru TN only supports non-deterministic cases and produces multiple normalization options.' + "Ru TN only supports non-deterministic cases and produces multiple normalization options." ) far_file = None if cache_dir is not None and cache_dir != "None": @@ -71,19 +72,25 @@ class ClassifyFst(GraphFst): alternative_formats = get_alternative_formats() self.cardinal = CardinalFst( - number_names=number_names, alternative_formats=alternative_formats, deterministic=deterministic + number_names=number_names, + alternative_formats=alternative_formats, + deterministic=deterministic, ) cardinal_graph = self.cardinal.fst self.ordinal = OrdinalFst( - number_names=number_names, alternative_formats=alternative_formats, deterministic=deterministic + number_names=number_names, + alternative_formats=alternative_formats, + deterministic=deterministic, ) ordinal_graph = self.ordinal.fst self.decimal = DecimalFst(cardinal=self.cardinal, deterministic=deterministic) decimal_graph = self.decimal.fst - self.measure = MeasureFst(cardinal=self.cardinal, decimal=self.decimal, deterministic=deterministic) + self.measure = MeasureFst( + cardinal=self.cardinal, decimal=self.decimal, deterministic=deterministic + ) measure_graph = self.measure.fst self.date = DateFst(number_names=number_names, deterministic=deterministic) date_graph = self.date.fst @@ -94,9 +101,13 @@ class ClassifyFst(GraphFst): telephone_graph = self.telephone.fst self.electronic = ElectronicFst(deterministic=deterministic) electronic_graph = self.electronic.fst - self.money = MoneyFst(cardinal=self.cardinal, decimal=self.decimal, deterministic=deterministic) + self.money = MoneyFst( + cardinal=self.cardinal, decimal=self.decimal, deterministic=deterministic + ) money_graph = self.money.fst - self.whitelist = WhiteListFst(input_case=input_case, deterministic=deterministic, input_file=whitelist) + self.whitelist = WhiteListFst( + input_case=input_case, deterministic=deterministic, input_file=whitelist + ) whitelist_graph = self.whitelist.fst punct_graph = PunctuationFst(deterministic=deterministic).fst @@ -114,13 +125,21 @@ class ClassifyFst(GraphFst): | pynutil.add_weight(word_graph, 100) ) - punct = pynutil.insert("tokens { ") + pynutil.add_weight(punct_graph, weight=1.1) + pynutil.insert(" }") + punct = ( + pynutil.insert("tokens { ") + + pynutil.add_weight(punct_graph, weight=1.1) + + pynutil.insert(" }") + ) token = pynutil.insert("tokens { ") + classify + pynutil.insert(" }") token_plus_punct = ( - pynini.closure(punct + pynutil.insert(" ")) + token + pynini.closure(pynutil.insert(" ") + punct) + pynini.closure(punct + pynutil.insert(" ")) + + token + + pynini.closure(pynutil.insert(" ") + punct) ) - graph = token_plus_punct + pynini.closure(pynutil.add_weight(delete_extra_space, 1.1) + token_plus_punct) + graph = token_plus_punct + pynini.closure( + pynutil.add_weight(delete_extra_space, 1.1) + token_plus_punct + ) graph = delete_space + graph + delete_space self.fst = graph.optimize() diff --git a/fun_text_processing/text_normalization/ru/taggers/whitelist.py b/fun_text_processing/text_normalization/ru/taggers/whitelist.py index 60518c69f..287fff148 100644 --- a/fun_text_processing/text_normalization/ru/taggers/whitelist.py +++ b/fun_text_processing/text_normalization/ru/taggers/whitelist.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_CHAR, GraphFst, convert_space from fun_text_processing.text_normalization.ru.alphabet import RU_ALPHA, TO_CYRILLIC @@ -43,9 +41,11 @@ class WhiteListFst(GraphFst): units_graph = _get_whitelist_graph(input_case, file=get_abs_path("data/measurements.tsv")) # do not replace single letter units, like `м`, `°` and `%` will be replaced - units_graph = pynini.compose((DAMO_CHAR ** (2, ...) | pynini.difference(DAMO_CHAR, RU_ALPHA)), units_graph) + units_graph = pynini.compose( + (DAMO_CHAR ** (2, ...) | pynini.difference(DAMO_CHAR, RU_ALPHA)), units_graph + ) graph |= units_graph.optimize() graph |= TO_CYRILLIC + pynini.closure(pynutil.insert(" ") + TO_CYRILLIC) self.final_graph = convert_space(graph) - self.fst = (pynutil.insert("name: \"") + self.final_graph + pynutil.insert("\"")).optimize() + self.fst = (pynutil.insert('name: "') + self.final_graph + pynutil.insert('"')).optimize() diff --git a/fun_text_processing/text_normalization/ru/taggers/word.py b/fun_text_processing/text_normalization/ru/taggers/word.py index 362517d92..318ba5da4 100644 --- a/fun_text_processing/text_normalization/ru/taggers/word.py +++ b/fun_text_processing/text_normalization/ru/taggers/word.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_SPACE, GraphFst from pynini.lib import pynutil @@ -17,5 +15,5 @@ class WordFst(GraphFst): def __init__(self, deterministic: bool = True): super().__init__(name="word", kind="classify") - word = pynutil.insert("name: \"") + pynini.closure(DAMO_NOT_SPACE, 1) + pynutil.insert("\"") + word = pynutil.insert('name: "') + pynini.closure(DAMO_NOT_SPACE, 1) + pynutil.insert('"') self.fst = word.optimize() diff --git a/fun_text_processing/text_normalization/ru/utils.py b/fun_text_processing/text_normalization/ru/utils.py index 861f75bcb..1a1f30f2e 100644 --- a/fun_text_processing/text_normalization/ru/utils.py +++ b/fun_text_processing/text_normalization/ru/utils.py @@ -1,5 +1,3 @@ - - import csv import os @@ -12,13 +10,13 @@ def get_abs_path(rel_path): Args: rel_path: relative path to this file - + Returns absolute path """ abs_path = os.path.dirname(os.path.abspath(__file__)) + os.sep + rel_path if not os.path.exists(abs_path): - logging.warning(f'{abs_path} does not exist') + logging.warning(f"{abs_path} does not exist") return abs_path @@ -31,6 +29,6 @@ def load_labels(abs_path): Returns dictionary of mappings """ - label_tsv = open(abs_path, encoding='utf-8') + label_tsv = open(abs_path, encoding="utf-8") labels = list(csv.reader(label_tsv, delimiter="\t")) return labels diff --git a/fun_text_processing/text_normalization/ru/verbalizers/cardinal.py b/fun_text_processing/text_normalization/ru/verbalizers/cardinal.py index 37281481c..2693477f4 100644 --- a/fun_text_processing/text_normalization/ru/verbalizers/cardinal.py +++ b/fun_text_processing/text_normalization/ru/verbalizers/cardinal.py @@ -1,4 +1,3 @@ - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst from pynini.lib import pynutil @@ -16,16 +15,18 @@ class CardinalFst(GraphFst): def __init__(self, deterministic: bool = True): super().__init__(name="cardinal", kind="verbalize", deterministic=deterministic) - optional_sign = pynini.closure(pynini.cross("negative: \"true\" ", "минус "), 0, 1) + optional_sign = pynini.closure(pynini.cross('negative: "true" ', "минус "), 0, 1) optional_quantity_part = pynini.closure( pynini.accep(" ") - + pynutil.delete("quantity: \"") + + pynutil.delete('quantity: "') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\""), + + pynutil.delete('"'), 0, 1, ) - integer = pynutil.delete("integer: \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + integer = ( + pynutil.delete('integer: "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') + ) self.graph = optional_sign + integer + optional_quantity_part delete_tokens = self.delete_tokens(self.graph) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/text_normalization/ru/verbalizers/date.py b/fun_text_processing/text_normalization/ru/verbalizers/date.py index d959ea380..b03eca40a 100644 --- a/fun_text_processing/text_normalization/ru/verbalizers/date.py +++ b/fun_text_processing/text_normalization/ru/verbalizers/date.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import GraphFst from fun_text_processing.text_normalization.ru.alphabet import RU_ALPHA @@ -19,6 +17,6 @@ class DateFst(GraphFst): def __init__(self, deterministic: bool = True): super().__init__(name="date", kind="verbalize", deterministic=deterministic) - graph = pynutil.delete("day: \"") + pynini.closure(RU_ALPHA | " ", 1) + pynutil.delete("\"") + graph = pynutil.delete('day: "') + pynini.closure(RU_ALPHA | " ", 1) + pynutil.delete('"') delete_tokens = self.delete_tokens(graph) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/text_normalization/ru/verbalizers/decimal.py b/fun_text_processing/text_normalization/ru/verbalizers/decimal.py index 4fd5dbdf6..ee1f2b484 100644 --- a/fun_text_processing/text_normalization/ru/verbalizers/decimal.py +++ b/fun_text_processing/text_normalization/ru/verbalizers/decimal.py @@ -1,7 +1,9 @@ - - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil @@ -19,21 +21,26 @@ class DecimalFst(GraphFst): def __init__(self, deterministic: bool = True): super().__init__(name="decimal", kind="verbalize", deterministic=deterministic) - optional_sign = pynini.closure(pynini.cross("negative: \"true\" ", "минус "), 0, 1) - integer = pynutil.delete(" \"") + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete("\"") + optional_sign = pynini.closure(pynini.cross('negative: "true" ', "минус "), 0, 1) + integer = pynutil.delete(' "') + pynini.closure(DAMO_NOT_QUOTE, 1) + pynutil.delete('"') integer_part = pynutil.delete("integer_part:") + integer fractional_part = pynutil.delete("fractional_part:") + integer optional_quantity_part = pynini.closure( pynini.accep(" ") - + pynutil.delete("quantity: \"") + + pynutil.delete('quantity: "') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\""), + + pynutil.delete('"'), 0, 1, ) self.graph = ( - optional_sign + integer_part + pynini.accep(" ") + fractional_part + optional_quantity_part + delete_space + optional_sign + + integer_part + + pynini.accep(" ") + + fractional_part + + optional_quantity_part + + delete_space ) delete_tokens = self.delete_tokens(self.graph) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/text_normalization/ru/verbalizers/electronic.py b/fun_text_processing/text_normalization/ru/verbalizers/electronic.py index 1c67b08f0..5d4264828 100644 --- a/fun_text_processing/text_normalization/ru/verbalizers/electronic.py +++ b/fun_text_processing/text_normalization/ru/verbalizers/electronic.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import GraphFst from fun_text_processing.text_normalization.ru.alphabet import RU_ALPHA @@ -19,6 +17,6 @@ class ElectronicFst(GraphFst): def __init__(self, deterministic: bool = True): super().__init__(name="electronic", kind="verbalize", deterministic=deterministic) - graph = pynutil.delete("username: \"") + pynini.closure(RU_ALPHA | " ") + pynutil.delete("\"") + graph = pynutil.delete('username: "') + pynini.closure(RU_ALPHA | " ") + pynutil.delete('"') delete_tokens = self.delete_tokens(graph) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/text_normalization/ru/verbalizers/measure.py b/fun_text_processing/text_normalization/ru/verbalizers/measure.py index 951a40029..d3eb1ceb4 100644 --- a/fun_text_processing/text_normalization/ru/verbalizers/measure.py +++ b/fun_text_processing/text_normalization/ru/verbalizers/measure.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import ( DAMO_NON_BREAKING_SPACE, @@ -15,7 +13,7 @@ class MeasureFst(GraphFst): """ Finite state transducer for verbalizing measure, e.g. measure { cardinal { integer: "два килограма" } } -> "два килограма" - + Args: deterministic: if True will provide a single transduction option, for False multiple transduction are generated (used for audio-based normalization) @@ -25,9 +23,9 @@ class MeasureFst(GraphFst): super().__init__(name="measure", kind="verbalize", deterministic=deterministic) graph = ( - pynutil.delete(" cardinal { integer: \"") + pynutil.delete(' cardinal { integer: "') + pynini.closure(RU_ALPHA | DAMO_SPACE | DAMO_NON_BREAKING_SPACE) - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space + pynutil.delete("}") ) diff --git a/fun_text_processing/text_normalization/ru/verbalizers/money.py b/fun_text_processing/text_normalization/ru/verbalizers/money.py index 7f07de721..a6c257448 100644 --- a/fun_text_processing/text_normalization/ru/verbalizers/money.py +++ b/fun_text_processing/text_normalization/ru/verbalizers/money.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import GraphFst from fun_text_processing.text_normalization.ru.alphabet import RU_ALPHA @@ -20,5 +18,7 @@ class MoneyFst(GraphFst): super().__init__(name="money", kind="verbalize", deterministic=deterministic) graph = pynini.closure(RU_ALPHA | " ") - delete_tokens = self.delete_tokens(pynutil.delete("integer_part: \"") + graph + pynutil.delete("\"")) + delete_tokens = self.delete_tokens( + pynutil.delete('integer_part: "') + graph + pynutil.delete('"') + ) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/text_normalization/ru/verbalizers/ordinal.py b/fun_text_processing/text_normalization/ru/verbalizers/ordinal.py index ad6c914e3..564d361f0 100644 --- a/fun_text_processing/text_normalization/ru/verbalizers/ordinal.py +++ b/fun_text_processing/text_normalization/ru/verbalizers/ordinal.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst from pynini.lib import pynutil @@ -19,6 +17,6 @@ class OrdinalFst(GraphFst): super().__init__(name="ordinal", kind="verbalize", deterministic=deterministic) value = pynini.closure(DAMO_NOT_QUOTE) - graph = pynutil.delete("integer: \"") + value + pynutil.delete("\"") + graph = pynutil.delete('integer: "') + value + pynutil.delete('"') delete_tokens = self.delete_tokens(graph) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/text_normalization/ru/verbalizers/telephone.py b/fun_text_processing/text_normalization/ru/verbalizers/telephone.py index 5513add57..044e3e911 100644 --- a/fun_text_processing/text_normalization/ru/verbalizers/telephone.py +++ b/fun_text_processing/text_normalization/ru/verbalizers/telephone.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.en.graph_utils import GraphFst from fun_text_processing.text_normalization.ru.alphabet import RU_ALPHA @@ -19,6 +17,10 @@ class TelephoneFst(GraphFst): def __init__(self, deterministic: bool = True): super().__init__(name="telephone", kind="verbalize", deterministic=deterministic) - graph = pynutil.delete("number_part: \"") + pynini.closure(RU_ALPHA | " ", 1) + pynutil.delete("\"") + graph = ( + pynutil.delete('number_part: "') + + pynini.closure(RU_ALPHA | " ", 1) + + pynutil.delete('"') + ) delete_tokens = self.delete_tokens(graph) self.fst = delete_tokens.optimize() diff --git a/fun_text_processing/text_normalization/ru/verbalizers/time.py b/fun_text_processing/text_normalization/ru/verbalizers/time.py index 55289cc0c..cdabfaccc 100644 --- a/fun_text_processing/text_normalization/ru/verbalizers/time.py +++ b/fun_text_processing/text_normalization/ru/verbalizers/time.py @@ -1,7 +1,10 @@ - - import pynini -from fun_text_processing.text_normalization.en.graph_utils import DAMO_NOT_QUOTE, GraphFst, delete_space, insert_space +from fun_text_processing.text_normalization.en.graph_utils import ( + DAMO_NOT_QUOTE, + GraphFst, + delete_space, + insert_space, +) from pynini.lib import pynutil @@ -21,20 +24,25 @@ class TimeFst(GraphFst): hour = ( pynutil.delete("hours:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) minutes = ( pynutil.delete("minutes:") + delete_space - + pynutil.delete("\"") + + pynutil.delete('"') + pynini.closure(DAMO_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') ) self.graph = ( - hour + delete_space + insert_space + minutes + delete_space + pynutil.delete("preserve_order: true") + hour + + delete_space + + insert_space + + minutes + + delete_space + + pynutil.delete("preserve_order: true") ) self.graph |= hour + delete_space self.graph |= minutes + delete_space + insert_space + hour + delete_space diff --git a/fun_text_processing/text_normalization/ru/verbalizers/verbalize.py b/fun_text_processing/text_normalization/ru/verbalizers/verbalize.py index ce635aa2f..465ff0a39 100644 --- a/fun_text_processing/text_normalization/ru/verbalizers/verbalize.py +++ b/fun_text_processing/text_normalization/ru/verbalizers/verbalize.py @@ -1,5 +1,3 @@ - - from fun_text_processing.text_normalization.en.graph_utils import GraphFst from fun_text_processing.text_normalization.en.verbalizers.whitelist import WhiteListFst from fun_text_processing.text_normalization.ru.verbalizers.cardinal import CardinalFst @@ -16,7 +14,7 @@ from fun_text_processing.text_normalization.ru.verbalizers.time import TimeFst class VerbalizeFst(GraphFst): """ Composes other verbalizer grammars. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. Args: diff --git a/fun_text_processing/text_normalization/ru/verbalizers/verbalize_final.py b/fun_text_processing/text_normalization/ru/verbalizers/verbalize_final.py index 841a94ef0..1e95e86e4 100644 --- a/fun_text_processing/text_normalization/ru/verbalizers/verbalize_final.py +++ b/fun_text_processing/text_normalization/ru/verbalizers/verbalize_final.py @@ -1,4 +1,3 @@ - import os import pynini @@ -17,7 +16,7 @@ import logging class VerbalizeFinalFst(GraphFst): """ - Finite state transducer that verbalizes an entire sentence, e.g. + Finite state transducer that verbalizes an entire sentence, e.g. tokens { name: "its" } tokens { time { hours: "12" minutes: "30" } } tokens { name: "now" } -> its 12:30 now Args: @@ -27,16 +26,20 @@ class VerbalizeFinalFst(GraphFst): overwrite_cache: set to True to overwrite .far files """ - def __init__(self, deterministic: bool = True, cache_dir: str = None, overwrite_cache: bool = False): + def __init__( + self, deterministic: bool = True, cache_dir: str = None, overwrite_cache: bool = False + ): super().__init__(name="verbalize_final", kind="verbalize", deterministic=deterministic) far_file = None if cache_dir is not None and cache_dir != "None": os.makedirs(cache_dir, exist_ok=True) - far_file = os.path.join(cache_dir, f"ru_tn_{deterministic}_deterministic_verbalizer.far") + far_file = os.path.join( + cache_dir, f"ru_tn_{deterministic}_deterministic_verbalizer.far" + ) if not overwrite_cache and far_file and os.path.exists(far_file): self.fst = pynini.Far(far_file, mode="r")["verbalize"] - logging.info(f'VerbalizeFinalFst graph was restored from {far_file}.') + logging.info(f"VerbalizeFinalFst graph was restored from {far_file}.") else: verbalize = VerbalizeFst().fst diff --git a/fun_text_processing/text_normalization/run_evaluate.py b/fun_text_processing/text_normalization/run_evaluate.py index b20694098..470a961d0 100644 --- a/fun_text_processing/text_normalization/run_evaluate.py +++ b/fun_text_processing/text_normalization/run_evaluate.py @@ -1,5 +1,3 @@ - - from argparse import ArgumentParser from fun_text_processing.text_normalization.data_loader_utils import ( @@ -12,18 +10,22 @@ from fun_text_processing.text_normalization.data_loader_utils import ( from fun_text_processing.text_normalization.normalize import Normalizer -''' +""" Runs Evaluation on data in the format of : \t\t<`self` if trivial class or normalized text> like the Google text normalization data https://www.kaggle.com/richardwilliamsproat/text-normalization-for-english-russian-and-polish -''' +""" def parse_args(): parser = ArgumentParser() parser.add_argument("--input", help="input file path", type=str) - parser.add_argument("--lang", help="language", choices=['en'], default="en", type=str) + parser.add_argument("--lang", help="language", choices=["en"], default="en", type=str) parser.add_argument( - "--input_case", help="input capitalization", choices=["lower_cased", "cased"], default="cased", type=str + "--input_case", + help="input capitalization", + choices=["lower_cased", "cased"], + default="cased", + type=str, ) parser.add_argument( "--cat", @@ -33,7 +35,9 @@ def parse_args(): default=None, choices=known_types, ) - parser.add_argument("--filter", action='store_true', help="clean data for normalization purposes") + parser.add_argument( + "--filter", action="store_true", help="clean data for normalization purposes" + ) return parser.parse_args() @@ -41,7 +45,7 @@ if __name__ == "__main__": # Example usage: # python run_evaluate.py --input= --cat= --filter args = parse_args() - if args.lang == 'en': + if args.lang == "en": from fun_text_processing.text_normalization.en.clean_eval_data import filter_loaded_data file_path = args.input normalizer = Normalizer(input_case=args.input_case, lang=args.lang) @@ -76,30 +80,35 @@ if __name__ == "__main__": preds=tokens_prediction, labels=tokens_normalized, input=tokens_un_normalized ) print(" - Accuracy: " + str(token_accuracy[token_type])) - token_count_per_type = {token_type: len(tokens_per_type[token_type][0]) for token_type in tokens_per_type} + token_count_per_type = { + token_type: len(tokens_per_type[token_type][0]) for token_type in tokens_per_type + } token_weighted_accuracy = [ - token_count_per_type[token_type] * accuracy for token_type, accuracy in token_accuracy.items() + token_count_per_type[token_type] * accuracy + for token_type, accuracy in token_accuracy.items() ] print("- Accuracy: " + str(sum(token_weighted_accuracy) / sum(token_count_per_type.values()))) - print(" - Total: " + str(sum(token_count_per_type.values())), '\n') + print(" - Total: " + str(sum(token_count_per_type.values())), "\n") - print(" - Total: " + str(sum(token_count_per_type.values())), '\n') + print(" - Total: " + str(sum(token_count_per_type.values())), "\n") for token_type in token_accuracy: if token_type not in known_types: raise ValueError("Unexpected token type: " + token_type) if args.category is None: - c1 = ['Class', 'sent level'] + known_types - c2 = ['Num Tokens', len(sentences_normalized)] + [ - token_count_per_type[known_type] if known_type in tokens_per_type else '0' for known_type in known_types + c1 = ["Class", "sent level"] + known_types + c2 = ["Num Tokens", len(sentences_normalized)] + [ + token_count_per_type[known_type] if known_type in tokens_per_type else "0" + for known_type in known_types ] - c3 = ['Normalization', sentences_accuracy] + [ - token_accuracy[known_type] if known_type in token_accuracy else '0' for known_type in known_types + c3 = ["Normalization", sentences_accuracy] + [ + token_accuracy[known_type] if known_type in token_accuracy else "0" + for known_type in known_types ] for i in range(len(c1)): - print(f'{str(c1[i]):10s} | {str(c2[i]):10s} | {str(c3[i]):5s}') + print(f"{str(c1[i]):10s} | {str(c2[i]):10s} | {str(c3[i]):5s}") else: - print(f'numbers\t{token_count_per_type[args.category]}') - print(f'Normalization\t{token_accuracy[args.category]}') + print(f"numbers\t{token_count_per_type[args.category]}") + print(f"Normalization\t{token_accuracy[args.category]}") diff --git a/fun_text_processing/text_normalization/token_parser.py b/fun_text_processing/text_normalization/token_parser.py index 10b2ab32d..cb5dfe131 100644 --- a/fun_text_processing/text_normalization/token_parser.py +++ b/fun_text_processing/text_normalization/token_parser.py @@ -1,5 +1,3 @@ - - import string from collections import OrderedDict from typing import Dict, List, Union @@ -22,7 +20,7 @@ class TokenParser: Args: text: text to be parsed - + """ self.text = text self.len_text = len(text) @@ -76,9 +74,9 @@ class TokenParser: if self.char == ":": self.parse_char(":") self.parse_ws() - self.parse_char("\"") + self.parse_char('"') value_string = self.parse_string_value() - self.parse_char("\"") + self.parse_char('"') return value_string elif self.char == "{": d = OrderedDict() @@ -95,11 +93,11 @@ class TokenParser: def parse_char(self, exp) -> bool: """ - Parses character + Parses character Args: exp: character to read in - + Returns true if successful """ assert self.char == exp @@ -112,7 +110,7 @@ class TokenParser: Args: exp: characters to read in - + Returns true if successful """ ok = False @@ -147,7 +145,7 @@ class TokenParser: # assert self.char not in string.whitespace and self.char != EOS assert self.char != EOS l = [] - while self.char != "\"" or self.text[self.index + 1] != " ": + while self.char != '"' or self.text[self.index + 1] != " ": l.append(self.char) if not self.read(): raise ValueError() @@ -169,8 +167,8 @@ class TokenParser: def read(self): """ - Reads in next char. - + Reads in next char. + Returns true if not EOS """ if self.index < self.len_text - 1: # should be unique diff --git a/fun_text_processing/text_normalization/zh/__init__.py b/fun_text_processing/text_normalization/zh/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/zh/__init__.py +++ b/fun_text_processing/text_normalization/zh/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/zh/data/__init__.py b/fun_text_processing/text_normalization/zh/data/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/zh/data/__init__.py +++ b/fun_text_processing/text_normalization/zh/data/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/zh/data/char/__init__.py b/fun_text_processing/text_normalization/zh/data/char/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/zh/data/char/__init__.py +++ b/fun_text_processing/text_normalization/zh/data/char/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/zh/data/date/__init__.py b/fun_text_processing/text_normalization/zh/data/date/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/zh/data/date/__init__.py +++ b/fun_text_processing/text_normalization/zh/data/date/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/zh/data/denylist/__init__.py b/fun_text_processing/text_normalization/zh/data/denylist/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/zh/data/denylist/__init__.py +++ b/fun_text_processing/text_normalization/zh/data/denylist/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/zh/data/erhua/__init__.py b/fun_text_processing/text_normalization/zh/data/erhua/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/zh/data/erhua/__init__.py +++ b/fun_text_processing/text_normalization/zh/data/erhua/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/zh/data/math/__init__.py b/fun_text_processing/text_normalization/zh/data/math/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/zh/data/math/__init__.py +++ b/fun_text_processing/text_normalization/zh/data/math/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/zh/data/measure/__init__.py b/fun_text_processing/text_normalization/zh/data/measure/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/zh/data/measure/__init__.py +++ b/fun_text_processing/text_normalization/zh/data/measure/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/zh/data/money/__init__.py b/fun_text_processing/text_normalization/zh/data/money/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/zh/data/money/__init__.py +++ b/fun_text_processing/text_normalization/zh/data/money/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/zh/data/number/__init__.py b/fun_text_processing/text_normalization/zh/data/number/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/zh/data/number/__init__.py +++ b/fun_text_processing/text_normalization/zh/data/number/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/zh/data/time/__init__.py b/fun_text_processing/text_normalization/zh/data/time/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/zh/data/time/__init__.py +++ b/fun_text_processing/text_normalization/zh/data/time/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/zh/data/whitelist/__init__.py b/fun_text_processing/text_normalization/zh/data/whitelist/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/zh/data/whitelist/__init__.py +++ b/fun_text_processing/text_normalization/zh/data/whitelist/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/zh/graph_utils.py b/fun_text_processing/text_normalization/zh/graph_utils.py index 83f38ac30..6664de5bd 100644 --- a/fun_text_processing/text_normalization/zh/graph_utils.py +++ b/fun_text_processing/text_normalization/zh/graph_utils.py @@ -1,5 +1,3 @@ - - import os import string from pathlib import Path @@ -18,7 +16,7 @@ FUN_UPPER = pynini.union(*string.ascii_uppercase).optimize() FUN_ALPHA = pynini.union(FUN_LOWER, FUN_UPPER).optimize() FUN_SPACE = " " -FUN_WHITE_SPACE = pynini.union(" ", "\t", "\n", "\r", u"\u00A0").optimize() +FUN_WHITE_SPACE = pynini.union(" ", "\t", "\n", "\r", "\u00A0").optimize() FUN_NOT_SPACE = pynini.difference(FUN_CHAR, FUN_WHITE_SPACE).optimize() FUN_NOT_QUOTE = pynini.difference(FUN_CHAR, r'"').optimize() @@ -33,7 +31,7 @@ insert_space = pynutil.insert(" ") delete_extra_space = pynini.cross(pynini.closure(FUN_WHITE_SPACE, 1), " ") -def generator_main(file_name: str, graphs: Dict[str, 'pynini.FstLike']): +def generator_main(file_name: str, graphs: Dict[str, "pynini.FstLike"]): """ Exports graph as OpenFst finite state archive (FAR) file with given file name and rule name. @@ -45,7 +43,7 @@ def generator_main(file_name: str, graphs: Dict[str, 'pynini.FstLike']): for rule, graph in graphs.items(): exporter[rule] = graph.optimize() exporter.close() - print(f'Created {file_name}') + print(f"Created {file_name}") class GraphFst: @@ -65,9 +63,11 @@ class GraphFst: self._fst = None self.deterministic = deterministic - self.far_path = Path(os.path.dirname(__file__) + '/grammars/' + kind + '/' + name + '.far') + self.far_path = Path(os.path.dirname(__file__) + "/grammars/" + kind + "/" + name + ".far") if self.far_exist(): - self._fst = Far(self.far_path, mode="r", arc_type="standard", far_type="default").get_fst() + self._fst = Far( + self.far_path, mode="r", arc_type="standard", far_type="default" + ).get_fst() def far_exist(self) -> bool: """ @@ -76,14 +76,14 @@ class GraphFst: return self.far_path.exists() @property - def fst(self) -> 'pynini.FstLike': + def fst(self) -> "pynini.FstLike": return self._fst @fst.setter def fst(self, fst): self._fst = fst - def add_tokens(self, fst) -> 'pynini.FstLike': + def add_tokens(self, fst) -> "pynini.FstLike": """ Wraps class name around to given fst @@ -95,7 +95,7 @@ class GraphFst: """ return pynutil.insert(f"{self.name} {{ ") + fst + pynutil.insert(" }") - def delete_tokens(self, fst) -> 'pynini.FstLike': + def delete_tokens(self, fst) -> "pynini.FstLike": """ Deletes class name wrap around output of given fst @@ -114,4 +114,4 @@ class GraphFst: + delete_space + pynutil.delete("}") ) - return res @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", FUN_SIGMA) + return res @ pynini.cdrewrite(pynini.cross("\u00A0", " "), "", "", FUN_SIGMA) diff --git a/fun_text_processing/text_normalization/zh/taggers/__init__.py b/fun_text_processing/text_normalization/zh/taggers/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/zh/taggers/__init__.py +++ b/fun_text_processing/text_normalization/zh/taggers/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/zh/taggers/cardinal.py b/fun_text_processing/text_normalization/zh/taggers/cardinal.py index 1c14483fc..d30bf10a1 100644 --- a/fun_text_processing/text_normalization/zh/taggers/cardinal.py +++ b/fun_text_processing/text_normalization/zh/taggers/cardinal.py @@ -1,24 +1,28 @@ - - import pynini from fun_text_processing.text_normalization.zh.graph_utils import GraphFst -from fun_text_processing.text_normalization.zh.utils import UNIT_1e01, UNIT_1e02, UNIT_1e03, UNIT_1e04, get_abs_path +from fun_text_processing.text_normalization.zh.utils import ( + UNIT_1e01, + UNIT_1e02, + UNIT_1e03, + UNIT_1e04, + get_abs_path, +) from pynini.lib import pynutil class Cardinal(GraphFst): - ''' - self.graph_cardinal: - 5 -> 五 - 12 -> 十二 - 213 -> 二百一十三 - 3123 -> 三千一百二十三 - 51234 -> 五万一千二百三十四 - 51,234 -> 五万一千二百三十四 - 0.125 -> 零点一二五 - self.fst: - 123 -> cardinal{ integer: "一二三" } - ''' + """ + self.graph_cardinal: + 5 -> 五 + 12 -> 十二 + 213 -> 二百一十三 + 3123 -> 三千一百二十三 + 51234 -> 五万一千二百三十四 + 51,234 -> 五万一千二百三十四 + 0.125 -> 零点一二五 + self.fst: + 123 -> cardinal{ integer: "一二三" } + """ def __init__(self, deterministic: bool = True, lm: bool = False): super().__init__(name="cardinal", kind="classify", deterministic=deterministic) @@ -31,16 +35,20 @@ class Cardinal(GraphFst): graph_digit_with_zero = graph_digit | graph_zero # from 0(read out) to 9 graph_no_zero = pynini.cross("0", "") # cardinalzero,in some case we don't read it out graph_digit_no_zero = graph_digit | graph_no_zero # from 0(not read out) to 9 - insert_zero = pynutil.insert('零') + insert_zero = pynutil.insert("零") delete_punct = pynutil.delete(",") # to deal with 5,000,cardinalin english form # 15 in 215 - graph_ten_u = (graph_digit + pynutil.insert(UNIT_1e01) + graph_digit_no_zero) | (graph_zero + graph_digit) + graph_ten_u = (graph_digit + pynutil.insert(UNIT_1e01) + graph_digit_no_zero) | ( + graph_zero + graph_digit + ) # 15 only - graph_ten = (graph_teen + pynutil.insert(UNIT_1e01) + graph_digit_no_zero) | (graph_zero + graph_digit) + graph_ten = (graph_teen + pynutil.insert(UNIT_1e01) + graph_digit_no_zero) | ( + graph_zero + graph_digit + ) # 215 or 200 graph_hundred = (graph_digit + pynutil.insert(UNIT_1e02) + graph_ten_u) | ( - graph_digit + pynutil.insert(UNIT_1e02) + graph_no_zero ** 2 + graph_digit + pynutil.insert(UNIT_1e02) + graph_no_zero**2 ) # 3000 or 3002 or 3012 or 3123 # both with 3,000 or 3,002 or 3,012 or 3,123 @@ -55,8 +63,15 @@ class Cardinal(GraphFst): + pynutil.insert(UNIT_1e01) + graph_digit_no_zero ) - | (graph_digit + pynutil.insert(UNIT_1e03) + delete_punct.ques + graph_zero + graph_no_zero + graph_digit) - | (graph_digit + pynutil.insert(UNIT_1e03) + delete_punct.ques + graph_no_zero ** 3) + | ( + graph_digit + + pynutil.insert(UNIT_1e03) + + delete_punct.ques + + graph_zero + + graph_no_zero + + graph_digit + ) + | (graph_digit + pynutil.insert(UNIT_1e03) + delete_punct.ques + graph_no_zero**3) ) # 20001234 or 2001234 or 201234 or 21234 or 20234 or 20023 or 20002 or 20000 # 8 digits max supported,for 9 digits cardinal,often write with unit instead of read it in only cardinalform @@ -73,8 +88,8 @@ class Cardinal(GraphFst): + insert_zero + (graph_digit + pynutil.insert(UNIT_1e01) + graph_digit_no_zero) ) - | (graph_no_zero + delete_punct.ques + graph_no_zero ** 2 + insert_zero + graph_digit) - | (graph_no_zero ** 4) + | (graph_no_zero + delete_punct.ques + graph_no_zero**2 + insert_zero + graph_digit) + | (graph_no_zero**4) ) ) @@ -85,12 +100,20 @@ class Cardinal(GraphFst): + pynini.closure(graph_digit_with_zero, 1) ) - graph = graph_hundred | graph_ten | graph_digit_with_zero | graph_thousand_sign | graph_ten_thousand_sign + graph = ( + graph_hundred + | graph_ten + | graph_digit_with_zero + | graph_thousand_sign + | graph_ten_thousand_sign + ) # 456.123 - graph_decimal = graph + pynini.cross('.', '点') + pynini.closure(graph_digit_with_zero, 1) + graph_decimal = graph + pynini.cross(".", "点") + pynini.closure(graph_digit_with_zero, 1) graph_cardinal = graph | graph_decimal signed_cardinal = graph_sign + graph_cardinal self.graph_cardinal = (graph_cardinal | signed_cardinal).optimize() - graph_numstring = self.add_tokens(pynutil.insert("integer: \"") + graph_numstring + pynutil.insert("\"")) + graph_numstring = self.add_tokens( + pynutil.insert('integer: "') + graph_numstring + pynutil.insert('"') + ) self.fst = graph_numstring.optimize() diff --git a/fun_text_processing/text_normalization/zh/taggers/char.py b/fun_text_processing/text_normalization/zh/taggers/char.py index 0d702f35b..20bd2e2c0 100644 --- a/fun_text_processing/text_normalization/zh/taggers/char.py +++ b/fun_text_processing/text_normalization/zh/taggers/char.py @@ -1,16 +1,14 @@ - - from fun_text_processing.text_normalization.zh.graph_utils import FUN_CHAR, GraphFst from pynini.lib import pynutil class Char(GraphFst): - ''' - 你 -> char { name: "你" } - ''' + """ + 你 -> char { name: "你" } + """ def __init__(self, deterministic: bool = True, lm: bool = False): super().__init__(name="char", kind="classify", deterministic=deterministic) - graph = pynutil.insert("name: \"") + FUN_CHAR + pynutil.insert("\"") + graph = pynutil.insert('name: "') + FUN_CHAR + pynutil.insert('"') self.fst = graph.optimize() diff --git a/fun_text_processing/text_normalization/zh/taggers/date.py b/fun_text_processing/text_normalization/zh/taggers/date.py index 2e7b3897f..56f5672e1 100644 --- a/fun_text_processing/text_normalization/zh/taggers/date.py +++ b/fun_text_processing/text_normalization/zh/taggers/date.py @@ -1,20 +1,23 @@ - - import pynini -from fun_text_processing.text_normalization.zh.graph_utils import FUN_CHAR, FUN_DIGIT, GraphFst, insert_space +from fun_text_processing.text_normalization.zh.graph_utils import ( + FUN_CHAR, + FUN_DIGIT, + GraphFst, + insert_space, +) from fun_text_processing.text_normalization.zh.utils import get_abs_path from pynini.lib import pynutil class Date(GraphFst): - ''' - 2002年 -> tokens { date { year: "2002" } } - 2002-01-28 -> tokens { date { year: "2002" month: "01" day: "28"} } - 2002/01/28 -> tokens { date { year: "2002" month: "01" day: "28"} } - 2002.01.28 -> tokens { date { year: "2002" month: "01" day: "28"} } - 2002/02 -> tokens { date { year: "2002" month "02"} } - 02/11 -> tokens { date { year: "02" month "11"} } different with case "fraction 2/11" - ''' + """ + 2002年 -> tokens { date { year: "2002" } } + 2002-01-28 -> tokens { date { year: "2002" month: "01" day: "28"} } + 2002/01/28 -> tokens { date { year: "2002" month: "01" day: "28"} } + 2002.01.28 -> tokens { date { year: "2002" month: "01" day: "28"} } + 2002/02 -> tokens { date { year: "2002" month "02"} } + 02/11 -> tokens { date { year: "02" month "11"} } different with case "fraction 2/11" + """ def __init__(self, deterministic: bool = True, lm: bool = False): super().__init__(name="date", kind="classify", deterministic=deterministic) @@ -22,15 +25,15 @@ class Date(GraphFst): graph_zero = pynini.string_file(get_abs_path("data/number/zero.tsv")) year_whitelist = pynini.string_file(get_abs_path("data/date/year_suffix.tsv")) - delete_date_sign = pynutil.delete("/") | pynutil.delete('-') | pynutil.delete('.') + delete_date_sign = pynutil.delete("/") | pynutil.delete("-") | pynutil.delete(".") # 2012年 date_type0 = ( - pynutil.insert("year: \"") + pynutil.insert('year: "') + pynini.closure(graph_digit | graph_zero, 2, 4) + "年" + pynini.difference(FUN_CHAR, year_whitelist) - + pynutil.insert("\"") + + pynutil.insert('"') ) year_2_4_digit = pynini.closure(FUN_DIGIT, 2, 4) + delete_date_sign @@ -43,50 +46,50 @@ class Date(GraphFst): # 2012/01/28 date_type1 = ( - pynutil.insert("year: \"") + pynutil.insert('year: "') + year_2_4_digit - + pynutil.insert("\"") + + pynutil.insert('"') + insert_space - + pynutil.insert("month: \"") + + pynutil.insert('month: "') + month - + pynutil.insert("\"") + + pynutil.insert('"') + insert_space - + pynutil.insert("day: \"") + + pynutil.insert('day: "') + day - + pynutil.insert("\"") + + pynutil.insert('"') ) # 12/01 date_type2 = ( - pynutil.insert("year: \"") + pynutil.insert('year: "') + year_2_4_digit - + pynutil.insert("\"") + + pynutil.insert('"') + insert_space - + pynutil.insert("month: \"") + + pynutil.insert('month: "') + month_no_day_with_zero - + pynutil.insert("\"") + + pynutil.insert('"') ) # 2012/11 date_type3 = ( - pynutil.insert("year: \"") + pynutil.insert('year: "') + year_4_digit - + pynutil.insert("\"") + + pynutil.insert('"') + insert_space - + pynutil.insert("month: \"") + + pynutil.insert('month: "') + month_no_day - + pynutil.insert("\"") + + pynutil.insert('"') ) # 02/05 date_type4 = ( - pynutil.insert("year: \"") + pynutil.insert('year: "') + year_2_digit_with_zero - + pynutil.insert("\"") + + pynutil.insert('"') + insert_space - + pynutil.insert("month: \"") + + pynutil.insert('month: "') + month_no_day - + pynutil.insert("\"") + + pynutil.insert('"') ) # add your date type as date_typex here. graph = date_type0 | date_type1 | date_type2 | date_type3 | date_type4 diff --git a/fun_text_processing/text_normalization/zh/taggers/fraction.py b/fun_text_processing/text_normalization/zh/taggers/fraction.py index 059539f9a..2eb699870 100644 --- a/fun_text_processing/text_normalization/zh/taggers/fraction.py +++ b/fun_text_processing/text_normalization/zh/taggers/fraction.py @@ -1,28 +1,26 @@ - - import pynini from fun_text_processing.text_normalization.zh.graph_utils import FUN_DIGIT, GraphFst, insert_space from pynini.lib import pynutil class Fraction(GraphFst): - ''' - 1/5 -> tokens { fraction { numerator: "1" denominator: "5" } } - ''' + """ + 1/5 -> tokens { fraction { numerator: "1" denominator: "5" } } + """ def __init__(self, deterministic: bool = True, lm: bool = False): super().__init__(name="fraction", kind="classify", deterministic=deterministic) - numerator = pynini.closure(FUN_DIGIT, 1) + pynutil.delete('/') + numerator = pynini.closure(FUN_DIGIT, 1) + pynutil.delete("/") denominator = pynini.closure(FUN_DIGIT, 1) graph = ( - pynutil.insert("numerator: \"") + pynutil.insert('numerator: "') + numerator - + pynutil.insert("\"") + + pynutil.insert('"') + insert_space - + pynutil.insert("denominator: \"") + + pynutil.insert('denominator: "') + denominator - + pynutil.insert("\"") + + pynutil.insert('"') ) self.fst = self.add_tokens(graph).optimize() diff --git a/fun_text_processing/text_normalization/zh/taggers/math_symbol.py b/fun_text_processing/text_normalization/zh/taggers/math_symbol.py index 1dc8a7fda..7ca7d52e6 100644 --- a/fun_text_processing/text_normalization/zh/taggers/math_symbol.py +++ b/fun_text_processing/text_normalization/zh/taggers/math_symbol.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.zh.graph_utils import GraphFst from fun_text_processing.text_normalization.zh.taggers.cardinal import Cardinal @@ -8,23 +6,23 @@ from pynini.lib import pynutil class MathSymbol(GraphFst): - ''' - + -> tokens { sign: "加" } - ''' + """ + + -> tokens { sign: "加" } + """ def __init__(self, deterministic: bool = True, lm: bool = False): super().__init__(name="sign", kind="classify", deterministic=deterministic) - ''' + """ add your sign in data/math/symbol.tsv,this graph just convert sigh to character,you can add more cases with detailed cases - ''' + """ score_sign = pynini.string_file(get_abs_path("data/math/score.tsv")) score = ( - pynutil.insert("score: \"") + pynutil.insert('score: "') + Cardinal().graph_cardinal + score_sign + Cardinal().graph_cardinal - + pynutil.insert("\"") + + pynutil.insert('"') ) graph = score self.fst = graph.optimize() diff --git a/fun_text_processing/text_normalization/zh/taggers/measure.py b/fun_text_processing/text_normalization/zh/taggers/measure.py index 95b5760cf..eb1de500d 100644 --- a/fun_text_processing/text_normalization/zh/taggers/measure.py +++ b/fun_text_processing/text_normalization/zh/taggers/measure.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.zh.graph_utils import GraphFst, insert_space from fun_text_processing.text_normalization.zh.taggers.cardinal import Cardinal @@ -8,9 +6,9 @@ from pynini.lib import pynutil class Measure(GraphFst): - ''' - 1kg -> tokens { measure { cardinal { integer: "一" } units: "千克" } } - ''' + """ + 1kg -> tokens { measure { cardinal { integer: "一" } units: "千克" } } + """ def __init__(self, deterministic: bool = True, lm: bool = False): super().__init__(name="measure", kind="classify", deterministic=deterministic) @@ -19,21 +17,21 @@ class Measure(GraphFst): units_zh = pynini.string_file(get_abs_path("data/measure/units_zh.tsv")) graph = ( pynutil.insert("cardinal { ") - + pynutil.insert("integer: \"") + + pynutil.insert('integer: "') + Cardinal().graph_cardinal - + pynutil.insert("\"") + + pynutil.insert('"') + pynutil.insert(" }") + insert_space - + pynutil.insert("units: \"") + + pynutil.insert('units: "') + (units_en | units_zh) - + pynutil.insert("\"") + + pynutil.insert('"') ) percent_graph = ( pynutil.insert("decimal { ") - + pynutil.insert("integer_part: \"") + + pynutil.insert('integer_part: "') + Cardinal().graph_cardinal + pynutil.delete("%") - + pynutil.insert("\"") + + pynutil.insert('"') + pynutil.insert(" }") ) graph |= percent_graph diff --git a/fun_text_processing/text_normalization/zh/taggers/money.py b/fun_text_processing/text_normalization/zh/taggers/money.py index 804beaa73..a234c9f5c 100644 --- a/fun_text_processing/text_normalization/zh/taggers/money.py +++ b/fun_text_processing/text_normalization/zh/taggers/money.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.zh.graph_utils import GraphFst, insert_space from fun_text_processing.text_normalization.zh.taggers.cardinal import Cardinal @@ -8,9 +6,9 @@ from pynini.lib import pynutil class Money(GraphFst): - ''' - ￥1.25 -> tokens { money { fractional_part: "元" integer_part: "一点五" } } - ''' + """ + ￥1.25 -> tokens { money { fractional_part: "元" integer_part: "一点五" } } + """ def __init__(self, deterministic: bool = True, lm: bool = False): super().__init__(name="money", kind="classify", deterministic=deterministic) @@ -18,13 +16,13 @@ class Money(GraphFst): currency_code = pynini.string_file(get_abs_path("data/money/currency_code.tsv")) currency_symbol = pynini.string_file(get_abs_path("data/money/currency_symbol.tsv")) graph = ( - pynutil.insert("fractional_part: \"") + pynutil.insert('fractional_part: "') + (currency_code | currency_symbol) - + pynutil.insert("\"") + + pynutil.insert('"') + insert_space - + pynutil.insert("integer_part: \"") + + pynutil.insert('integer_part: "') + Cardinal().graph_cardinal - + pynutil.insert("\"") + + pynutil.insert('"') ) self.fst = self.add_tokens(graph).optimize() diff --git a/fun_text_processing/text_normalization/zh/taggers/preprocessor.py b/fun_text_processing/text_normalization/zh/taggers/preprocessor.py index fa296da1f..2e0074523 100644 --- a/fun_text_processing/text_normalization/zh/taggers/preprocessor.py +++ b/fun_text_processing/text_normalization/zh/taggers/preprocessor.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.zh.graph_utils import FUN_SIGMA, GraphFst from fun_text_processing.text_normalization.zh.utils import get_abs_path @@ -7,29 +5,35 @@ from pynini.lib import pynutil class PreProcessor(GraphFst): - ''' - Preprocessing of TN: - 1. interjections removal such as '啊, 呃' - 2. fullwidth -> halfwidth char conversion - 好啊 -> 好 - 呃对 -> 对 - ： -> : - ； -> ; - ''' + """ + Preprocessing of TN: + 1. interjections removal such as '啊, 呃' + 2. fullwidth -> halfwidth char conversion + 好啊 -> 好 + 呃对 -> 对 + ： -> : + ； -> ; + """ def __init__( - self, remove_interjections: bool = True, fullwidth_to_halfwidth: bool = True, + self, + remove_interjections: bool = True, + fullwidth_to_halfwidth: bool = True, ): super().__init__(name="PreProcessor", kind="processor") - graph = pynini.cdrewrite('', '', '', FUN_SIGMA) + graph = pynini.cdrewrite("", "", "", FUN_SIGMA) if remove_interjections: - remove_interjections_graph = pynutil.delete(pynini.string_file(get_abs_path('data/denylist/denylist.tsv'))) - graph @= pynini.cdrewrite(remove_interjections_graph, '', '', FUN_SIGMA) + remove_interjections_graph = pynutil.delete( + pynini.string_file(get_abs_path("data/denylist/denylist.tsv")) + ) + graph @= pynini.cdrewrite(remove_interjections_graph, "", "", FUN_SIGMA) if fullwidth_to_halfwidth: - fullwidth_to_halfwidth_graph = pynini.string_file(get_abs_path('data/char/fullwidth_to_halfwidth.tsv')) - graph @= pynini.cdrewrite(fullwidth_to_halfwidth_graph, '', '', FUN_SIGMA) + fullwidth_to_halfwidth_graph = pynini.string_file( + get_abs_path("data/char/fullwidth_to_halfwidth.tsv") + ) + graph @= pynini.cdrewrite(fullwidth_to_halfwidth_graph, "", "", FUN_SIGMA) self.fst = graph.optimize() diff --git a/fun_text_processing/text_normalization/zh/taggers/time.py b/fun_text_processing/text_normalization/zh/taggers/time.py index 98557071f..d5ee7ef78 100644 --- a/fun_text_processing/text_normalization/zh/taggers/time.py +++ b/fun_text_processing/text_normalization/zh/taggers/time.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.zh.graph_utils import GraphFst, insert_space from fun_text_processing.text_normalization.zh.utils import get_abs_path @@ -7,11 +5,11 @@ from pynini.lib import pynutil class Time(GraphFst): - ''' - 1:02 -> tokens { time { hours: "1" minitus: "02" } } - 1:02:36 -> tokens { time { hours: "1" minutes: "02" seconds: "36" } } - 1:02 am -> tokens { time { hours: "1" minutes: "02" seconds: "36" suffix "am" } } - ''' + """ + 1:02 -> tokens { time { hours: "1" minitus: "02" } } + 1:02:36 -> tokens { time { hours: "1" minutes: "02" seconds: "36" } } + 1:02 am -> tokens { time { hours: "1" minutes: "02" seconds: "36" suffix "am" } } + """ def __init__(self, deterministic: bool = True, lm: bool = False): super().__init__(name="time", kind="classify", deterministic=deterministic) @@ -26,35 +24,37 @@ class Time(GraphFst): m = time_tens + time_digit s = (time_tens + time_digit) | time_digit - delete_colon = pynini.cross(':', ' ') + delete_colon = pynini.cross(":", " ") # 5:05, 14:30 h_m = ( - pynutil.insert('hours: \"') + pynutil.insert('hours: "') + h - + pynutil.insert('\"') + + pynutil.insert('"') + delete_colon - + pynutil.insert('minutes: \"') + + pynutil.insert('minutes: "') + m - + pynutil.insert('\"') + + pynutil.insert('"') ) # 1:30:15 h_m_s = ( - pynutil.insert('hours: \"') + pynutil.insert('hours: "') + h - + pynutil.insert('\"') + + pynutil.insert('"') + delete_colon - + pynutil.insert('minutes: \"') + + pynutil.insert('minutes: "') + m - + pynutil.insert('\"') + + pynutil.insert('"') + delete_colon - + pynutil.insert('seconds: \"') + + pynutil.insert('seconds: "') + s - + pynutil.insert('\"') + + pynutil.insert('"') ) graph = h_m | h_m_s - graph_suffix = graph + insert_space + pynutil.insert('suffix: \"') + time_suffix + pynutil.insert('\"') + graph_suffix = ( + graph + insert_space + pynutil.insert('suffix: "') + time_suffix + pynutil.insert('"') + ) graph |= graph_suffix self.fst = self.add_tokens(graph).optimize() diff --git a/fun_text_processing/text_normalization/zh/taggers/tokenize_and_classify.py b/fun_text_processing/text_normalization/zh/taggers/tokenize_and_classify.py index 439357c4d..1cba12fc5 100644 --- a/fun_text_processing/text_normalization/zh/taggers/tokenize_and_classify.py +++ b/fun_text_processing/text_normalization/zh/taggers/tokenize_and_classify.py @@ -1,5 +1,3 @@ - - import os import pynini @@ -20,9 +18,9 @@ from pynini.lib import pynutil class ClassifyFst(GraphFst): """ Final class that composes all other classification grammars. This class can process an entire sentence including punctuation. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. - + Args: input_case: accepting either "lower_cased" or "cased" input. deterministic: if True will provide a single transduction option, @@ -47,7 +45,8 @@ class ClassifyFst(GraphFst): os.makedirs(cache_dir, exist_ok=True) whitelist_file = os.path.basename(whitelist) if whitelist else "" far_file = os.path.join( - cache_dir, f"zh_tn_{deterministic}_deterministic_{input_case}_{whitelist_file}_tokenize.far" + cache_dir, + f"zh_tn_{deterministic}_deterministic_{input_case}_{whitelist_file}_tokenize.far", ) if not overwrite_cache and far_file and os.path.exists(far_file): self.fst = pynini.Far(far_file, mode="r")["tokenize_and_classify"] @@ -77,5 +76,8 @@ class ClassifyFst(GraphFst): tagger = pynini.cdrewrite(token.optimize(), "", "", FUN_SIGMA).optimize() - preprocessor = PreProcessor(remove_interjections=True, fullwidth_to_halfwidth=True,) + preprocessor = PreProcessor( + remove_interjections=True, + fullwidth_to_halfwidth=True, + ) self.fst = preprocessor.fst @ tagger diff --git a/fun_text_processing/text_normalization/zh/taggers/whitelist.py b/fun_text_processing/text_normalization/zh/taggers/whitelist.py index 6880224c1..616c55b40 100644 --- a/fun_text_processing/text_normalization/zh/taggers/whitelist.py +++ b/fun_text_processing/text_normalization/zh/taggers/whitelist.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.zh.graph_utils import GraphFst from fun_text_processing.text_normalization.zh.utils import get_abs_path @@ -7,20 +5,20 @@ from pynini.lib import pynutil class Whitelist(GraphFst): - ''' - ATM -> tokens { whitelist: "ATM" } - ''' + """ + ATM -> tokens { whitelist: "ATM" } + """ def __init__(self, deterministic: bool = True, lm: bool = False): super().__init__(name="whitelist", kind="classify", deterministic=deterministic) whitelist = pynini.string_file(get_abs_path("data/whitelist/default.tsv")) - erhua = pynutil.insert("erhua: \"") + pynini.accep('儿') + pynutil.insert("\"") + erhua = pynutil.insert('erhua: "') + pynini.accep("儿") + pynutil.insert('"') sign = pynini.string_file(get_abs_path("data/math/symbol.tsv")) whitelist = ( - pynutil.insert("name: \"") + pynutil.insert('name: "') + (pynini.string_file(get_abs_path("data/erhua/whitelist.tsv")) | whitelist | sign) - + pynutil.insert("\"") + + pynutil.insert('"') ) graph = pynutil.add_weight(erhua, 0.1) | whitelist diff --git a/fun_text_processing/text_normalization/zh/utils.py b/fun_text_processing/text_normalization/zh/utils.py index b9793e4d6..237a43795 100644 --- a/fun_text_processing/text_normalization/zh/utils.py +++ b/fun_text_processing/text_normalization/zh/utils.py @@ -1,12 +1,10 @@ - - import csv import os -UNIT_1e01 = '十' -UNIT_1e02 = '百' -UNIT_1e03 = '千' -UNIT_1e04 = '万' +UNIT_1e01 = "十" +UNIT_1e02 = "百" +UNIT_1e03 = "千" +UNIT_1e04 = "万" def get_abs_path(rel_path): @@ -15,10 +13,10 @@ def get_abs_path(rel_path): Args: rel_path: relative path to this file - + Returns absolute path """ - return os.path.dirname(os.path.abspath(__file__)) + '/' + rel_path + return os.path.dirname(os.path.abspath(__file__)) + "/" + rel_path def load_labels(abs_path): @@ -37,7 +35,7 @@ def load_labels(abs_path): def augment_labels_with_punct_at_end(labels): """ - augments labels: if key ends on a punctuation that value does not have, add a new label + augments labels: if key ends on a punctuation that value does not have, add a new label where the value maintains the punctuation Args: diff --git a/fun_text_processing/text_normalization/zh/verbalizers/__init__.py b/fun_text_processing/text_normalization/zh/verbalizers/__init__.py index 139597f9c..8b1378917 100644 --- a/fun_text_processing/text_normalization/zh/verbalizers/__init__.py +++ b/fun_text_processing/text_normalization/zh/verbalizers/__init__.py @@ -1,2 +1 @@ - diff --git a/fun_text_processing/text_normalization/zh/verbalizers/cardinal.py b/fun_text_processing/text_normalization/zh/verbalizers/cardinal.py index 5606b02f1..6b4bcc7c8 100644 --- a/fun_text_processing/text_normalization/zh/verbalizers/cardinal.py +++ b/fun_text_processing/text_normalization/zh/verbalizers/cardinal.py @@ -4,13 +4,13 @@ from pynini.lib import pynutil class Cardinal(GraphFst): - ''' - tokens { cardinal { integer: "一二三" } } -> 一二三 - ''' + """ + tokens { cardinal { integer: "一二三" } } -> 一二三 + """ def __init__(self, deterministic: bool = True, lm: bool = False): super().__init__(name="cardinal", kind="verbalize", deterministic=deterministic) - graph = pynutil.delete('integer: \"') + pynini.closure(FUN_NOT_QUOTE) + pynutil.delete('\"') + graph = pynutil.delete('integer: "') + pynini.closure(FUN_NOT_QUOTE) + pynutil.delete('"') self.fst = self.delete_tokens(graph).optimize() diff --git a/fun_text_processing/text_normalization/zh/verbalizers/char.py b/fun_text_processing/text_normalization/zh/verbalizers/char.py index 4d693fcde..b04e44162 100644 --- a/fun_text_processing/text_normalization/zh/verbalizers/char.py +++ b/fun_text_processing/text_normalization/zh/verbalizers/char.py @@ -1,16 +1,14 @@ - - from fun_text_processing.text_normalization.zh.graph_utils import FUN_NOT_QUOTE, GraphFst from pynini.lib import pynutil class Char(GraphFst): - ''' - tokens { char: "你" } -> 你 - ''' + """ + tokens { char: "你" } -> 你 + """ def __init__(self, deterministic: bool = True, lm: bool = False): super().__init__(name="char", kind="verbalize", deterministic=deterministic) - graph = pynutil.delete("name: \"") + FUN_NOT_QUOTE + pynutil.delete("\"") + graph = pynutil.delete('name: "') + FUN_NOT_QUOTE + pynutil.delete('"') self.fst = graph.optimize() diff --git a/fun_text_processing/text_normalization/zh/verbalizers/date.py b/fun_text_processing/text_normalization/zh/verbalizers/date.py index dd0ae12de..cf0e604d2 100644 --- a/fun_text_processing/text_normalization/zh/verbalizers/date.py +++ b/fun_text_processing/text_normalization/zh/verbalizers/date.py @@ -1,20 +1,22 @@ - - import pynini -from fun_text_processing.text_normalization.zh.graph_utils import FUN_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.text_normalization.zh.graph_utils import ( + FUN_NOT_QUOTE, + GraphFst, + delete_space, +) from fun_text_processing.text_normalization.zh.utils import UNIT_1e01, get_abs_path from pynini.lib import pynutil class Date(GraphFst): - ''' - tokens { date { year: "2002" month: "01" day: "28"} } -> 二零零二年一月二十八日 - tokens { date { year: "2002" } } -> 二零零八年 - ''' + """ + tokens { date { year: "2002" month: "01" day: "28"} } -> 二零零二年一月二十八日 + tokens { date { year: "2002" } } -> 二零零八年 + """ def __init__(self, deterministic: bool = True, lm: bool = False): super().__init__(name="date", kind="verbalize", deterministic=deterministic) - date_type0 = pynutil.delete('year: \"') + pynini.closure(FUN_NOT_QUOTE) + pynutil.delete('\"') + date_type0 = pynutil.delete('year: "') + pynini.closure(FUN_NOT_QUOTE) + pynutil.delete('"') graph_digit = pynini.string_file(get_abs_path("data/number/digit.tsv")) graph_teen = pynini.string_file(get_abs_path("data/number/digit_teen.tsv")) graph_zero = pynini.string_file(get_abs_path("data/number/zero.tsv")) @@ -26,32 +28,32 @@ class Date(GraphFst): ) date_type1 = ( - pynutil.delete("year: \"") + pynutil.delete('year: "') + graph_year + pynutil.insert("年") - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space - + pynutil.delete("month: \"") + + pynutil.delete('month: "') + graph_2_digit_date + pynutil.insert("月") - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space - + pynutil.delete("day: \"") + + pynutil.delete('day: "') + graph_2_digit_date + pynutil.insert("日") - + pynutil.delete("\"") + + pynutil.delete('"') ) date_type2 = ( - pynutil.delete("year: \"") + pynutil.delete('year: "') + graph_year + pynutil.insert("年") - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space - + pynutil.delete("month: \"") + + pynutil.delete('month: "') + graph_2_digit_date + pynutil.insert("月") - + pynutil.delete("\"") + + pynutil.delete('"') ) graph = date_type0 | date_type1 | date_type2 diff --git a/fun_text_processing/text_normalization/zh/verbalizers/fraction.py b/fun_text_processing/text_normalization/zh/verbalizers/fraction.py index 111a79dca..414c470cb 100644 --- a/fun_text_processing/text_normalization/zh/verbalizers/fraction.py +++ b/fun_text_processing/text_normalization/zh/verbalizers/fraction.py @@ -1,20 +1,20 @@ - - from fun_text_processing.text_normalization.zh.graph_utils import GraphFst from fun_text_processing.text_normalization.zh.taggers.cardinal import Cardinal from pynini.lib import pynutil class Fraction(GraphFst): - ''' - tokens { fraction { denominator: "5" numerator: "1" } } -> 五分之一 - ''' + """ + tokens { fraction { denominator: "5" numerator: "1" } } -> 五分之一 + """ def __init__(self, deterministic: bool = True, lm: bool = False): super().__init__(name="fraction", kind="verbalize", deterministic=deterministic) - denominator = pynutil.delete("denominator: \"") + Cardinal().graph_cardinal + pynutil.delete("\"") - numerator = pynutil.delete("numerator: \"") + Cardinal().graph_cardinal + pynutil.delete("\"") + denominator = ( + pynutil.delete('denominator: "') + Cardinal().graph_cardinal + pynutil.delete('"') + ) + numerator = pynutil.delete('numerator: "') + Cardinal().graph_cardinal + pynutil.delete('"') graph = denominator + pynutil.delete(" ") + pynutil.insert("分之") + numerator self.fst = self.delete_tokens(graph).optimize() diff --git a/fun_text_processing/text_normalization/zh/verbalizers/math_symbol.py b/fun_text_processing/text_normalization/zh/verbalizers/math_symbol.py index 939f869ed..4a70fda27 100644 --- a/fun_text_processing/text_normalization/zh/verbalizers/math_symbol.py +++ b/fun_text_processing/text_normalization/zh/verbalizers/math_symbol.py @@ -1,18 +1,16 @@ - - import pynini from fun_text_processing.text_normalization.zh.graph_utils import FUN_NOT_QUOTE, GraphFst from pynini.lib import pynutil class MathSymbol(GraphFst): - ''' - tokens { sign: "加" } -> 加 - ''' + """ + tokens { sign: "加" } -> 加 + """ def __init__(self, deterministic: bool = True, lm: bool = False): super().__init__(name="sign", kind="verbalize", deterministic=deterministic) - graph = pynutil.delete('score: \"') + pynini.closure(FUN_NOT_QUOTE) + pynutil.delete('\"') + graph = pynutil.delete('score: "') + pynini.closure(FUN_NOT_QUOTE) + pynutil.delete('"') self.fst = graph.optimize() diff --git a/fun_text_processing/text_normalization/zh/verbalizers/measure.py b/fun_text_processing/text_normalization/zh/verbalizers/measure.py index 0f0156b28..2119044c1 100644 --- a/fun_text_processing/text_normalization/zh/verbalizers/measure.py +++ b/fun_text_processing/text_normalization/zh/verbalizers/measure.py @@ -1,14 +1,16 @@ - - import pynini -from fun_text_processing.text_normalization.zh.graph_utils import FUN_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.text_normalization.zh.graph_utils import ( + FUN_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil class Measure(GraphFst): - ''' - tokens { measure { cardinal: "一" } units: "千克" } } -> 一千克 - ''' + """ + tokens { measure { cardinal: "一" } units: "千克" } } -> 一千克 + """ def __init__(self, deterministic: bool = True, lm: bool = False): super().__init__(name="measure", kind="verbalize", deterministic=deterministic) @@ -16,22 +18,22 @@ class Measure(GraphFst): graph = ( pynutil.delete("cardinal {") + delete_space - + pynutil.delete("integer: \"") + + pynutil.delete('integer: "') + pynini.closure(FUN_NOT_QUOTE) - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space + pynutil.delete("}") + delete_space - + pynutil.delete("units: \"") + + pynutil.delete('units: "') + pynini.closure(FUN_NOT_QUOTE) - + pynutil.delete("\"") + + pynutil.delete('"') ) percent_graph = ( pynutil.delete("decimal { ") - + pynutil.delete("integer_part: \"") + + pynutil.delete('integer_part: "') + pynutil.insert("百分之") + pynini.closure(FUN_NOT_QUOTE, 1) - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space + pynutil.delete("}") ) diff --git a/fun_text_processing/text_normalization/zh/verbalizers/money.py b/fun_text_processing/text_normalization/zh/verbalizers/money.py index ad59b4694..7f1aa9e60 100644 --- a/fun_text_processing/text_normalization/zh/verbalizers/money.py +++ b/fun_text_processing/text_normalization/zh/verbalizers/money.py @@ -1,20 +1,31 @@ - - import pynini -from fun_text_processing.text_normalization.zh.graph_utils import FUN_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.text_normalization.zh.graph_utils import ( + FUN_NOT_QUOTE, + GraphFst, + delete_space, +) from pynini.lib import pynutil class Money(GraphFst): - ''' - tokens { money { integer_part: "一点五" fractional_part: "元" } } -> 一点五元 - ''' + """ + tokens { money { integer_part: "一点五" fractional_part: "元" } } -> 一点五元 + """ def __init__(self, deterministic: bool = True, lm: bool = False): super().__init__(name="money", kind="verbalize", deterministic=deterministic) - cur = pynutil.delete("fractional_part: \"") + pynini.closure(FUN_NOT_QUOTE) + pynutil.delete("\"") - num = pynutil.delete("integer_part: \"") + pynini.closure(FUN_NOT_QUOTE) + pynutil.delete("\"") + delete_space + cur = ( + pynutil.delete('fractional_part: "') + + pynini.closure(FUN_NOT_QUOTE) + + pynutil.delete('"') + ) + num = ( + pynutil.delete('integer_part: "') + + pynini.closure(FUN_NOT_QUOTE) + + pynutil.delete('"') + + delete_space + ) graph = num + cur self.fst = self.delete_tokens(graph).optimize() diff --git a/fun_text_processing/text_normalization/zh/verbalizers/postprocessor.py b/fun_text_processing/text_normalization/zh/verbalizers/postprocessor.py index 77dc9ff41..e87f28986 100644 --- a/fun_text_processing/text_normalization/zh/verbalizers/postprocessor.py +++ b/fun_text_processing/text_normalization/zh/verbalizers/postprocessor.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.zh.graph_utils import ( FUN_ALPHA, @@ -14,45 +12,57 @@ from pynini.lib import pynutil, utf8 class PostProcessor(GraphFst): - ''' - Postprocessing of TN, now contains: - 1. punctuation removal - 2. letter case conversion - 3. oov tagger - ''' + """ + Postprocessing of TN, now contains: + 1. punctuation removal + 2. letter case conversion + 3. oov tagger + """ def __init__( - self, remove_puncts: bool = False, to_upper: bool = False, to_lower: bool = False, tag_oov: bool = False, + self, + remove_puncts: bool = False, + to_upper: bool = False, + to_lower: bool = False, + tag_oov: bool = False, ): super().__init__(name="PostProcessor", kind="processor") - graph = pynini.cdrewrite('', '', '', FUN_SIGMA) + graph = pynini.cdrewrite("", "", "", FUN_SIGMA) if remove_puncts: remove_puncts_graph = pynutil.delete( - pynini.union(FUN_PUNCT, pynini.string_file(get_abs_path('data/char/punctuations_zh.tsv'))) + pynini.union( + FUN_PUNCT, pynini.string_file(get_abs_path("data/char/punctuations_zh.tsv")) + ) ) graph @= pynini.cdrewrite(remove_puncts_graph, "", "", FUN_SIGMA).optimize() if to_upper or to_lower: if to_upper: - conv_cases_graph = pynini.inverse(pynini.string_file(get_abs_path('data/char/upper_to_lower.tsv'))) + conv_cases_graph = pynini.inverse( + pynini.string_file(get_abs_path("data/char/upper_to_lower.tsv")) + ) else: - conv_cases_graph = pynini.string_file(get_abs_path('data/char/upper_to_lower.tsv')) + conv_cases_graph = pynini.string_file(get_abs_path("data/char/upper_to_lower.tsv")) graph @= pynini.cdrewrite(conv_cases_graph, "", "", FUN_SIGMA).optimize() if tag_oov: - zh_charset_std = pynini.string_file(get_abs_path("data/char/charset_national_standard_2013_8105.tsv")) + zh_charset_std = pynini.string_file( + get_abs_path("data/char/charset_national_standard_2013_8105.tsv") + ) zh_charset_ext = pynini.string_file(get_abs_path("data/char/charset_extension.tsv")) zh_charset = ( - zh_charset_std | zh_charset_ext | pynini.string_file(get_abs_path("data/char/punctuations_zh.tsv")) + zh_charset_std + | zh_charset_ext + | pynini.string_file(get_abs_path("data/char/punctuations_zh.tsv")) ) en_charset = FUN_DIGIT | FUN_ALPHA | FUN_PUNCT | FUN_WHITE_SPACE charset = zh_charset | en_charset with open(get_abs_path("data/char/oov_tags.tsv"), "r") as f: - tags = f.readline().strip().split('\t') + tags = f.readline().strip().split("\t") assert len(tags) == 2 ltag, rtag = tags diff --git a/fun_text_processing/text_normalization/zh/verbalizers/time.py b/fun_text_processing/text_normalization/zh/verbalizers/time.py index c7823757b..0c88ed85d 100644 --- a/fun_text_processing/text_normalization/zh/verbalizers/time.py +++ b/fun_text_processing/text_normalization/zh/verbalizers/time.py @@ -1,16 +1,18 @@ - - import pynini -from fun_text_processing.text_normalization.zh.graph_utils import FUN_NOT_QUOTE, GraphFst, delete_space +from fun_text_processing.text_normalization.zh.graph_utils import ( + FUN_NOT_QUOTE, + GraphFst, + delete_space, +) from fun_text_processing.text_normalization.zh.utils import UNIT_1e01, get_abs_path from pynini.lib import pynutil class Time(GraphFst): - ''' - tokens { time { h: "1" m: "02" s: "36" } } -> 一点零二分三十六秒 - tokens { time { suffix "am" hours: "1" minutes: "02" seconds: "36" } } -> 上午一点零二分三十六秒 - ''' + """ + tokens { time { h: "1" m: "02" s: "36" } } -> 一点零二分三十六秒 + tokens { time { suffix "am" hours: "1" minutes: "02" seconds: "36" } } -> 上午一点零二分三十六秒 + """ def __init__(self, deterministic: bool = True, lm: bool = False): super().__init__(name="time", kind="verbalize", deterministic=deterministic) @@ -33,50 +35,54 @@ class Time(GraphFst): # 6:25 h_m = ( - pynutil.delete("hours: \"") + pynutil.delete('hours: "') + h + pynutil.insert("点") - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space - + pynutil.delete("minutes: \"") + + pynutil.delete('minutes: "') + (graph_2_digit_time) + pynutil.insert("分") - + pynutil.delete("\"") + + pynutil.delete('"') ) # 23:00 h_00 = ( - pynutil.delete("hours: \"") + pynutil.delete('hours: "') + h + pynutil.insert("点") - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space - + pynutil.delete("minutes: \"") + + pynutil.delete('minutes: "') + (graph_2_digit_zero_none) - + pynutil.delete("\"") + + pynutil.delete('"') ) # 9:12:52 h_m_s = ( - pynutil.delete("hours: \"") + pynutil.delete('hours: "') + h + pynutil.insert("点") - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space - + pynutil.delete("minutes: \"") + + pynutil.delete('minutes: "') + m + pynutil.insert("分") - + pynutil.delete("\"") + + pynutil.delete('"') + delete_space - + pynutil.delete("seconds: \"") + + pynutil.delete('seconds: "') + s + pynutil.insert("秒") - + pynutil.delete("\"") + + pynutil.delete('"') ) graph = h_m | h_m_s | h_00 graph_suffix = ( - pynutil.delete("suffix: \"") + pynini.closure(FUN_NOT_QUOTE) + pynutil.delete("\"") + delete_space + graph + pynutil.delete('suffix: "') + + pynini.closure(FUN_NOT_QUOTE) + + pynutil.delete('"') + + delete_space + + graph ) graph |= graph_suffix self.fst = self.delete_tokens(graph).optimize() diff --git a/fun_text_processing/text_normalization/zh/verbalizers/verbalize.py b/fun_text_processing/text_normalization/zh/verbalizers/verbalize.py index d01c2e7c5..b58213779 100644 --- a/fun_text_processing/text_normalization/zh/verbalizers/verbalize.py +++ b/fun_text_processing/text_normalization/zh/verbalizers/verbalize.py @@ -1,5 +1,3 @@ - - import pynini from fun_text_processing.text_normalization.zh.graph_utils import GraphFst from fun_text_processing.text_normalization.zh.verbalizers.cardinal import Cardinal @@ -16,7 +14,7 @@ from fun_text_processing.text_normalization.zh.verbalizers.whitelist import Whit class VerbalizeFst(GraphFst): """ Composes other verbalizer grammars. - For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. More details to deployment at NeMo/tools/text_processing_deployment. Args: deterministic: if True will provide a single transduction option, diff --git a/fun_text_processing/text_normalization/zh/verbalizers/verbalize_final.py b/fun_text_processing/text_normalization/zh/verbalizers/verbalize_final.py index d03e9f3a9..1c7f77b00 100644 --- a/fun_text_processing/text_normalization/zh/verbalizers/verbalize_final.py +++ b/fun_text_processing/text_normalization/zh/verbalizers/verbalize_final.py @@ -1,9 +1,11 @@ - - import os import pynini -from fun_text_processing.text_normalization.zh.graph_utils import GraphFst, delete_space, generator_main +from fun_text_processing.text_normalization.zh.graph_utils import ( + GraphFst, + delete_space, + generator_main, +) from fun_text_processing.text_normalization.zh.verbalizers.postprocessor import PostProcessor from fun_text_processing.text_normalization.zh.verbalizers.verbalize import VerbalizeFst from pynini.lib import pynutil @@ -12,26 +14,37 @@ from pynini.lib import pynutil class VerbalizeFinalFst(GraphFst): - """ + """ """ - """ - - def __init__(self, deterministic: bool = True, cache_dir: str = None, overwrite_cache: bool = False): + def __init__( + self, deterministic: bool = True, cache_dir: str = None, overwrite_cache: bool = False + ): super().__init__(name="verbalize_final", kind="verbalize", deterministic=deterministic) far_file = None if cache_dir is not None and cache_dir != "None": os.makedirs(cache_dir, exist_ok=True) - far_file = os.path.join(cache_dir, f"zh_tn_{deterministic}_deterministic_verbalizer.far") + far_file = os.path.join( + cache_dir, f"zh_tn_{deterministic}_deterministic_verbalizer.far" + ) if not overwrite_cache and far_file and os.path.exists(far_file): self.fst = pynini.Far(far_file, mode="r")["verbalize"] else: token_graph = VerbalizeFst(deterministic=deterministic) token_verbalizer = ( - pynutil.delete("tokens {") + delete_space + token_graph.fst + delete_space + pynutil.delete(" }") + pynutil.delete("tokens {") + + delete_space + + token_graph.fst + + delete_space + + pynutil.delete(" }") ) verbalizer = pynini.closure(delete_space + token_verbalizer + delete_space) - postprocessor = PostProcessor(remove_puncts=False, to_upper=False, to_lower=False, tag_oov=False,) + postprocessor = PostProcessor( + remove_puncts=False, + to_upper=False, + to_lower=False, + tag_oov=False, + ) self.fst = (verbalizer @ postprocessor.fst).optimize() if far_file: diff --git a/fun_text_processing/text_normalization/zh/verbalizers/whitelist.py b/fun_text_processing/text_normalization/zh/verbalizers/whitelist.py index 212eb6b9f..4d0406662 100644 --- a/fun_text_processing/text_normalization/zh/verbalizers/whitelist.py +++ b/fun_text_processing/text_normalization/zh/verbalizers/whitelist.py @@ -1,18 +1,16 @@ - - import pynini from fun_text_processing.text_normalization.zh.graph_utils import FUN_NOT_QUOTE, GraphFst from pynini.lib import pynutil class Whitelist(GraphFst): - ''' - tokens { whitelist: "ATM" } -> A T M - ''' + """ + tokens { whitelist: "ATM" } -> A T M + """ def __init__(self, deterministic: bool = True, lm: bool = False): super().__init__(name="whitelist", kind="verbalize", deterministic=deterministic) - remove_erhua = pynutil.delete("erhua: \"") + pynutil.delete("儿") + pynutil.delete("\"") - whitelist = pynutil.delete("name: \"") + pynini.closure(FUN_NOT_QUOTE) + pynutil.delete("\"") + remove_erhua = pynutil.delete('erhua: "') + pynutil.delete("儿") + pynutil.delete('"') + whitelist = pynutil.delete('name: "') + pynini.closure(FUN_NOT_QUOTE) + pynutil.delete('"') graph = remove_erhua | whitelist self.fst = graph.optimize() diff --git a/funasr/__init__.py b/funasr/__init__.py index a649ac490..c7f7c219f 100644 --- a/funasr/__init__.py +++ b/funasr/__init__.py @@ -13,6 +13,7 @@ with open(version_file, "r") as f: import importlib import pkgutil + def import_submodules(package, recursive=True): if isinstance(package, str): try: @@ -23,7 +24,7 @@ def import_submodules(package, recursive=True): pass results = {} if not isinstance(package, str): - for loader, name, is_pkg in pkgutil.walk_packages(package.__path__, package.__name__ + '.'): + for loader, name, is_pkg in pkgutil.walk_packages(package.__path__, package.__name__ + "."): try: results[name] = importlib.import_module(name) except Exception as e: @@ -34,9 +35,10 @@ def import_submodules(package, recursive=True): results.update(import_submodules(name)) return results + import_submodules(__name__) from funasr.auto.auto_model import AutoModel from funasr.auto.auto_frontend import AutoFrontend -os.environ["HYDRA_FULL_ERROR"] = "1" \ No newline at end of file +os.environ["HYDRA_FULL_ERROR"] = "1" diff --git a/funasr/auto/auto_frontend.py b/funasr/auto/auto_frontend.py index b802b83a9..b8fa7e0b4 100644 --- a/funasr/auto/auto_frontend.py +++ b/funasr/auto/auto_frontend.py @@ -33,7 +33,7 @@ class AutoFrontend: if "model_conf" not in kwargs: logging.info("download models from model hub: {}".format(kwargs.get("hub", "ms"))) kwargs = download_model(**kwargs) - + # build frontend frontend = kwargs.get("frontend", None) if frontend is not None: @@ -45,25 +45,23 @@ class AutoFrontend: del kwargs["frontend"] self.kwargs = kwargs - def __call__(self, input, input_len=None, kwargs=None, **cfg): - + kwargs = self.kwargs if kwargs is None else kwargs kwargs.update(cfg) - key_list, data_list = prepare_data_iterator(input, input_len=input_len) batch_size = kwargs.get("batch_size", 1) device = kwargs.get("device", "cpu") if device == "cpu": batch_size = 1 - + meta_data = {} - + result_list = [] num_samples = len(data_list) pbar = tqdm(colour="blue", total=num_samples + 1, dynamic_ncols=True) - + time0 = time.perf_counter() for beg_idx in range(0, num_samples, batch_size): end_idx = min(num_samples, beg_idx + batch_size) @@ -72,27 +70,32 @@ class AutoFrontend: # extract fbank feats time1 = time.perf_counter() - audio_sample_list = load_audio_text_image_video(data_batch, fs=self.frontend.fs, audio_fs=kwargs.get("fs", 16000)) + audio_sample_list = load_audio_text_image_video( + data_batch, fs=self.frontend.fs, audio_fs=kwargs.get("fs", 16000) + ) time2 = time.perf_counter() meta_data["load_data"] = f"{time2 - time1:0.3f}" - speech, speech_lengths = extract_fbank(audio_sample_list, data_type=kwargs.get("data_type", "sound"), - frontend=self.frontend, **kwargs) + speech, speech_lengths = extract_fbank( + audio_sample_list, + data_type=kwargs.get("data_type", "sound"), + frontend=self.frontend, + **kwargs, + ) time3 = time.perf_counter() meta_data["extract_feat"] = f"{time3 - time2:0.3f}" - meta_data["batch_data_time"] = speech_lengths.sum().item() * self.frontend.frame_shift * self.frontend.lfr_n / 1000 - + meta_data["batch_data_time"] = ( + speech_lengths.sum().item() * self.frontend.frame_shift * self.frontend.lfr_n / 1000 + ) + speech.to(device=device), speech_lengths.to(device=device) batch = {"input": speech, "input_len": speech_lengths, "key": key_batch} result_list.append(batch) - + pbar.update(1) - description = ( - f"{meta_data}, " - ) + description = f"{meta_data}, " pbar.set_description(description) - + time_end = time.perf_counter() pbar.set_description(f"time escaped total: {time_end - time0:0.3f}") - - return result_list + return result_list diff --git a/funasr/auto/auto_model.py b/funasr/auto/auto_model.py index ba8881aa4..10d012dbe 100644 --- a/funasr/auto/auto_model.py +++ b/funasr/auto/auto_model.py @@ -36,7 +36,7 @@ except: def prepare_data_iterator(data_in, input_len=None, data_type=None, key=None): """ - + :param input: :param input_len: :param data_type: @@ -46,39 +46,43 @@ def prepare_data_iterator(data_in, input_len=None, data_type=None, key=None): data_list = [] key_list = [] filelist = [".scp", ".txt", ".json", ".jsonl", ".text"] - + chars = string.ascii_letters + string.digits - if isinstance(data_in, str) and data_in.startswith('http'): # url + if isinstance(data_in, str) and data_in.startswith("http"): # url data_in = download_from_url(data_in) - if isinstance(data_in, str) and os.path.exists(data_in): # wav_path; filelist: wav.scp, file.jsonl;text.txt; + if isinstance(data_in, str) and os.path.exists( + data_in + ): # wav_path; filelist: wav.scp, file.jsonl;text.txt; _, file_extension = os.path.splitext(data_in) file_extension = file_extension.lower() - if file_extension in filelist: #filelist: wav.scp, file.jsonl;text.txt; - with open(data_in, encoding='utf-8') as fin: + if file_extension in filelist: # filelist: wav.scp, file.jsonl;text.txt; + with open(data_in, encoding="utf-8") as fin: for line in fin: - key = "rand_key_" + ''.join(random.choice(chars) for _ in range(13)) - if data_in.endswith(".jsonl"): #file.jsonl: json.dumps({"source": data}) + key = "rand_key_" + "".join(random.choice(chars) for _ in range(13)) + if data_in.endswith(".jsonl"): # file.jsonl: json.dumps({"source": data}) lines = json.loads(line.strip()) data = lines["source"] key = data["key"] if "key" in data else key - else: # filelist, wav.scp, text.txt: id \t data or data + else: # filelist, wav.scp, text.txt: id \t data or data lines = line.strip().split(maxsplit=1) - data = lines[1] if len(lines)>1 else lines[0] - key = lines[0] if len(lines)>1 else key - + data = lines[1] if len(lines) > 1 else lines[0] + key = lines[0] if len(lines) > 1 else key + data_list.append(data) key_list.append(key) else: if key is None: - key = "rand_key_" + ''.join(random.choice(chars) for _ in range(13)) + key = "rand_key_" + "".join(random.choice(chars) for _ in range(13)) data_list = [data_in] key_list = [key] elif isinstance(data_in, (list, tuple)): - if data_type is not None and isinstance(data_type, (list, tuple)): # mutiple inputs + if data_type is not None and isinstance(data_type, (list, tuple)): # mutiple inputs data_list_tmp = [] for data_in_i, data_type_i in zip(data_in, data_type): - key_list, data_list_i = prepare_data_iterator(data_in=data_in_i, data_type=data_type_i) + key_list, data_list_i = prepare_data_iterator( + data_in=data_in_i, data_type=data_type_i + ) data_list_tmp.append(data_list_i) data_list = [] for item in zip(*data_list_tmp): @@ -86,26 +90,29 @@ def prepare_data_iterator(data_in, input_len=None, data_type=None, key=None): else: # [audio sample point, fbank, text] data_list = data_in - key_list = ["rand_key_" + ''.join(random.choice(chars) for _ in range(13)) for _ in range(len(data_in))] - else: # raw text; audio sample point, fbank; bytes - if isinstance(data_in, bytes): # audio bytes + key_list = [ + "rand_key_" + "".join(random.choice(chars) for _ in range(13)) + for _ in range(len(data_in)) + ] + else: # raw text; audio sample point, fbank; bytes + if isinstance(data_in, bytes): # audio bytes data_in = load_bytes(data_in) if key is None: - key = "rand_key_" + ''.join(random.choice(chars) for _ in range(13)) + key = "rand_key_" + "".join(random.choice(chars) for _ in range(13)) data_list = [data_in] key_list = [key] - + return key_list, data_list class AutoModel: - + def __init__(self, **kwargs): if not kwargs.get("disable_log", True): tables.print() model, kwargs = self.build_model(**kwargs) - + # if vad_model is not None, build vad model else None vad_model = kwargs.get("vad_model", None) vad_kwargs = {} if kwargs.get("vad_kwargs", {}) is None else kwargs.get("vad_kwargs", {}) @@ -136,11 +143,11 @@ class AutoModel: spk_kwargs["device"] = kwargs["device"] spk_model, spk_kwargs = self.build_model(**spk_kwargs) self.cb_model = ClusterBackend().to(kwargs["device"]) - spk_mode = kwargs.get("spk_mode", 'punc_segment') + spk_mode = kwargs.get("spk_mode", "punc_segment") if spk_mode not in ["default", "vad_segment", "punc_segment"]: logging.error("spk_mode should be one of default, vad_segment and punc_segment.") self.spk_mode = spk_mode - + self.kwargs = kwargs self.model = model self.vad_model = vad_model @@ -150,13 +157,13 @@ class AutoModel: self.spk_model = spk_model self.spk_kwargs = spk_kwargs self.model_path = kwargs.get("model_path") - + def build_model(self, **kwargs): assert "model" in kwargs if "model_conf" not in kwargs: logging.info("download models from model hub: {}".format(kwargs.get("hub", "ms"))) kwargs = download_model(**kwargs) - + set_all_random_seed(kwargs.get("seed", 0)) device = kwargs.get("device", "cuda") @@ -166,28 +173,34 @@ class AutoModel: kwargs["device"] = device torch.set_num_threads(kwargs.get("ncpu", 4)) - + # build tokenizer tokenizer = kwargs.get("tokenizer", None) if tokenizer is not None: tokenizer_class = tables.tokenizer_classes.get(tokenizer) tokenizer = tokenizer_class(**kwargs.get("tokenizer_conf", {})) - kwargs["token_list"] = tokenizer.token_list if hasattr(tokenizer, "token_list") else None - kwargs["token_list"] = tokenizer.get_vocab() if hasattr(tokenizer, "get_vocab") else kwargs["token_list"] + kwargs["token_list"] = ( + tokenizer.token_list if hasattr(tokenizer, "token_list") else None + ) + kwargs["token_list"] = ( + tokenizer.get_vocab() if hasattr(tokenizer, "get_vocab") else kwargs["token_list"] + ) vocab_size = len(kwargs["token_list"]) if kwargs["token_list"] is not None else -1 if vocab_size == -1 and hasattr(tokenizer, "get_vocab_size"): vocab_size = tokenizer.get_vocab_size() else: vocab_size = -1 kwargs["tokenizer"] = tokenizer - + # build frontend frontend = kwargs.get("frontend", None) kwargs["input_size"] = None if frontend is not None: frontend_class = tables.frontend_classes.get(frontend) frontend = frontend_class(**kwargs.get("frontend_conf", {})) - kwargs["input_size"] = frontend.output_size() if hasattr(frontend, "output_size") else None + kwargs["input_size"] = ( + frontend.output_size() if hasattr(frontend, "output_size") else None + ) kwargs["frontend"] = frontend # build model model_class = tables.model_classes.get(kwargs["model"]) @@ -196,7 +209,7 @@ class AutoModel: deep_update(model_conf, kwargs) model = model_class(**model_conf, vocab_size=vocab_size) model.to(device) - + # init_param init_param = kwargs.get("init_param", None) if init_param is not None: @@ -212,12 +225,12 @@ class AutoModel: ) else: print(f"error, init_param does not exist!: {init_param}") - + # fp16 if kwargs.get("fp16", False): model.to(torch.float16) return model, kwargs - + def __call__(self, *args, **cfg): kwargs = self.kwargs deep_update(kwargs, cfg) @@ -227,10 +240,10 @@ class AutoModel: def generate(self, input, input_len=None, **cfg): if self.vad_model is None: return self.inference(input, input_len=input_len, **cfg) - + else: return self.inference_with_vad(input, input_len=input_len, **cfg) - + def inference(self, input, input_len=None, model=None, kwargs=None, key=None, **cfg): kwargs = self.kwargs if kwargs is None else kwargs deep_update(kwargs, cfg) @@ -241,13 +254,17 @@ class AutoModel: # if kwargs.get("device", "cpu") == "cpu": # batch_size = 1 - key_list, data_list = prepare_data_iterator(input, input_len=input_len, data_type=kwargs.get("data_type", None), key=key) + key_list, data_list = prepare_data_iterator( + input, input_len=input_len, data_type=kwargs.get("data_type", None), key=key + ) speed_stats = {} asr_result_list = [] num_samples = len(data_list) disable_pbar = self.kwargs.get("disable_pbar", False) - pbar = tqdm(colour="blue", total=num_samples, dynamic_ncols=True) if not disable_pbar else None + pbar = ( + tqdm(colour="blue", total=num_samples, dynamic_ncols=True) if not disable_pbar else None + ) time_speech_total = 0.0 time_escape_total = 0.0 for beg_idx in range(0, num_samples, batch_size): @@ -256,14 +273,14 @@ class AutoModel: key_batch = key_list[beg_idx:end_idx] batch = {"data_in": data_batch, "key": key_batch} - if (end_idx - beg_idx) == 1 and kwargs.get("data_type", None) == "fbank": # fbank + if (end_idx - beg_idx) == 1 and kwargs.get("data_type", None) == "fbank": # fbank batch["data_in"] = data_batch[0] batch["data_lengths"] = input_len time1 = time.perf_counter() with torch.no_grad(): - res = model.inference(**batch, **kwargs) - if isinstance(res, (list, tuple)): + res = model.inference(**batch, **kwargs) + if isinstance(res, (list, tuple)): results = res[0] meta_data = res[1] if len(res) > 1 else {} time2 = time.perf_counter() @@ -278,9 +295,7 @@ class AutoModel: speed_stats["forward"] = f"{time_escape:0.3f}" speed_stats["batch_size"] = f"{len(results)}" speed_stats["rtf"] = f"{(time_escape) / batch_data_time:0.3f}" - description = ( - f"{speed_stats}, " - ) + description = f"{speed_stats}, " if pbar: pbar.update(1) pbar.set_description(description) @@ -298,27 +313,35 @@ class AutoModel: # step.1: compute the vad model deep_update(self.vad_kwargs, cfg) beg_vad = time.time() - res = self.inference(input, input_len=input_len, model=self.vad_model, kwargs=self.vad_kwargs, **cfg) + res = self.inference( + input, input_len=input_len, model=self.vad_model, kwargs=self.vad_kwargs, **cfg + ) end_vad = time.time() # FIX(gcf): concat the vad clips for sense vocie model for better aed if kwargs.get("merge_vad", False): for i in range(len(res)): - res[i]['value'] = merge_vad(res[i]['value'], kwargs.get("merge_length", 15000)) + res[i]["value"] = merge_vad(res[i]["value"], kwargs.get("merge_length", 15000)) # step.2 compute asr model model = self.model deep_update(kwargs, cfg) - batch_size = max(int(kwargs.get("batch_size_s", 300))*1000, 1) - batch_size_threshold_ms = int(kwargs.get("batch_size_threshold_s", 60))*1000 + batch_size = max(int(kwargs.get("batch_size_s", 300)) * 1000, 1) + batch_size_threshold_ms = int(kwargs.get("batch_size_threshold_s", 60)) * 1000 kwargs["batch_size"] = batch_size - key_list, data_list = prepare_data_iterator(input, input_len=input_len, data_type=kwargs.get("data_type", None)) + key_list, data_list = prepare_data_iterator( + input, input_len=input_len, data_type=kwargs.get("data_type", None) + ) results_ret_list = [] time_speech_total_all_samples = 1e-6 beg_total = time.time() - pbar_total = tqdm(colour="red", total=len(res), dynamic_ncols=True) if not kwargs.get("disable_pbar", False) else None + pbar_total = ( + tqdm(colour="red", total=len(res), dynamic_ncols=True) + if not kwargs.get("disable_pbar", False) + else None + ) for i in range(len(res)): key = res[i]["key"] vadsegments = res[i]["value"] @@ -341,7 +364,7 @@ class AutoModel: batch_size_ms_cum = 0 beg_idx = 0 beg_asr_total = time.time() - time_speech_total_per_sample = speech_lengths/16000 + time_speech_total_per_sample = speech_lengths / 16000 time_speech_total_all_samples += time_speech_total_per_sample # pbar_sample = tqdm(colour="blue", total=n, dynamic_ncols=True) @@ -349,26 +372,40 @@ class AutoModel: all_segments = [] for j, _ in enumerate(range(0, n)): # pbar_sample.update(1) - batch_size_ms_cum += (sorted_data[j][0][1] - sorted_data[j][0][0]) - if j < n - 1 and ( - batch_size_ms_cum + sorted_data[j + 1][0][1] - sorted_data[j + 1][0][0]) < batch_size and ( - sorted_data[j + 1][0][1] - sorted_data[j + 1][0][0]) < batch_size_threshold_ms: + batch_size_ms_cum += sorted_data[j][0][1] - sorted_data[j][0][0] + if ( + j < n - 1 + and (batch_size_ms_cum + sorted_data[j + 1][0][1] - sorted_data[j + 1][0][0]) + < batch_size + and (sorted_data[j + 1][0][1] - sorted_data[j + 1][0][0]) + < batch_size_threshold_ms + ): continue batch_size_ms_cum = 0 end_idx = j + 1 - speech_j, speech_lengths_j = slice_padding_audio_samples(speech, speech_lengths, sorted_data[beg_idx:end_idx]) - results = self.inference(speech_j, input_len=None, model=model, kwargs=kwargs, **cfg) + speech_j, speech_lengths_j = slice_padding_audio_samples( + speech, speech_lengths, sorted_data[beg_idx:end_idx] + ) + results = self.inference( + speech_j, input_len=None, model=model, kwargs=kwargs, **cfg + ) if self.spk_model is not None: # compose vad segments: [[start_time_sec, end_time_sec, speech], [...]] for _b in range(len(speech_j)): - vad_segments = [[sorted_data[beg_idx:end_idx][_b][0][0]/1000.0, - sorted_data[beg_idx:end_idx][_b][0][1]/1000.0, - np.array(speech_j[_b])]] + vad_segments = [ + [ + sorted_data[beg_idx:end_idx][_b][0][0] / 1000.0, + sorted_data[beg_idx:end_idx][_b][0][1] / 1000.0, + np.array(speech_j[_b]), + ] + ] segments = sv_chunk(vad_segments) all_segments.extend(segments) speech_b = [i[2] for i in segments] - spk_res = self.inference(speech_b, input_len=None, model=self.spk_model, kwargs=kwargs, **cfg) - results[_b]['spk_embedding'] = spk_res[0]['spk_embedding'] + spk_res = self.inference( + speech_b, input_len=None, model=self.spk_model, kwargs=kwargs, **cfg + ) + results[_b]["spk_embedding"] = spk_res[0]["spk_embedding"] beg_idx = end_idx if len(results) < 1: continue @@ -398,12 +435,12 @@ class AutoModel: t[0] += vadsegments[j][0] t[1] += vadsegments[j][0] result[k].extend(restored_data[j][k]) - elif k == 'spk_embedding': + elif k == "spk_embedding": if k not in result: result[k] = restored_data[j][k] else: result[k] = torch.cat([result[k], restored_data[j][k]], dim=0) - elif 'text' in k: + elif "text" in k: if k not in result: result[k] = restored_data[j][k] else: @@ -414,62 +451,80 @@ class AutoModel: else: result[k] += restored_data[j][k] - return_raw_text = kwargs.get('return_raw_text', False) + return_raw_text = kwargs.get("return_raw_text", False) # step.3 compute punc model if self.punc_model is not None: if not len(result["text"].strip()): if return_raw_text: - result['raw_text'] = '' + result["raw_text"] = "" else: deep_update(self.punc_kwargs, cfg) - punc_res = self.inference(result["text"], model=self.punc_model, kwargs=self.punc_kwargs, **cfg) + punc_res = self.inference( + result["text"], model=self.punc_model, kwargs=self.punc_kwargs, **cfg + ) raw_text = copy.copy(result["text"]) - if return_raw_text: result['raw_text'] = raw_text + if return_raw_text: + result["raw_text"] = raw_text result["text"] = punc_res[0]["text"] else: raw_text = None # speaker embedding cluster after resorted - if self.spk_model is not None and kwargs.get('return_spk_res', True): + if self.spk_model is not None and kwargs.get("return_spk_res", True): if raw_text is None: logging.error("Missing punc_model, which is required by spk_model.") all_segments = sorted(all_segments, key=lambda x: x[0]) - spk_embedding = result['spk_embedding'] - labels = self.cb_model(spk_embedding.cpu(), oracle_num=kwargs.get('preset_spk_num', None)) + spk_embedding = result["spk_embedding"] + labels = self.cb_model( + spk_embedding.cpu(), oracle_num=kwargs.get("preset_spk_num", None) + ) # del result['spk_embedding'] sv_output = postprocess(all_segments, None, labels, spk_embedding.cpu()) - if self.spk_mode == 'vad_segment': # recover sentence_list + if self.spk_mode == "vad_segment": # recover sentence_list sentence_list = [] for res, vadsegment in zip(restored_data, vadsegments): - if 'timestamp' not in res: - logging.error("Only 'iic/speech_paraformer-large-vad-punc_asr_nat-zh-cn-16k-common-vocab8404-pytorch' \ + if "timestamp" not in res: + logging.error( + "Only 'iic/speech_paraformer-large-vad-punc_asr_nat-zh-cn-16k-common-vocab8404-pytorch' \ and 'iic/speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch'\ - can predict timestamp, and speaker diarization relies on timestamps.") - sentence_list.append({"start": vadsegment[0], - "end": vadsegment[1], - "sentence": res['text'], - "timestamp": res['timestamp']}) - elif self.spk_mode == 'punc_segment': - if 'timestamp' not in result: - logging.error("Only 'iic/speech_paraformer-large-vad-punc_asr_nat-zh-cn-16k-common-vocab8404-pytorch' \ + can predict timestamp, and speaker diarization relies on timestamps." + ) + sentence_list.append( + { + "start": vadsegment[0], + "end": vadsegment[1], + "sentence": res["text"], + "timestamp": res["timestamp"], + } + ) + elif self.spk_mode == "punc_segment": + if "timestamp" not in result: + logging.error( + "Only 'iic/speech_paraformer-large-vad-punc_asr_nat-zh-cn-16k-common-vocab8404-pytorch' \ and 'iic/speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch'\ - can predict timestamp, and speaker diarization relies on timestamps.") - sentence_list = timestamp_sentence(punc_res[0]['punc_array'], - result['timestamp'], - raw_text, - return_raw_text=return_raw_text) + can predict timestamp, and speaker diarization relies on timestamps." + ) + sentence_list = timestamp_sentence( + punc_res[0]["punc_array"], + result["timestamp"], + raw_text, + return_raw_text=return_raw_text, + ) distribute_spk(sentence_list, sv_output) - result['sentence_info'] = sentence_list + result["sentence_info"] = sentence_list elif kwargs.get("sentence_timestamp", False): - if not len(result['text']): + if not len(result["text"]): sentence_list = [] else: - sentence_list = timestamp_sentence(punc_res[0]['punc_array'], - result['timestamp'], - raw_text, - return_raw_text=return_raw_text) - result['sentence_info'] = sentence_list - if "spk_embedding" in result: del result['spk_embedding'] + sentence_list = timestamp_sentence( + punc_res[0]["punc_array"], + result["timestamp"], + raw_text, + return_raw_text=return_raw_text, + ) + result["sentence_info"] = sentence_list + if "spk_embedding" in result: + del result["spk_embedding"] result["key"] = key results_ret_list.append(result) @@ -477,10 +532,11 @@ class AutoModel: time_escape_total_per_sample = end_asr_total - beg_asr_total if pbar_total: pbar_total.update(1) - pbar_total.set_description(f"rtf_avg: {time_escape_total_per_sample / time_speech_total_per_sample:0.3f}, " - f"time_speech: {time_speech_total_per_sample: 0.3f}, " - f"time_escape: {time_escape_total_per_sample:0.3f}") - + pbar_total.set_description( + f"rtf_avg: {time_escape_total_per_sample / time_speech_total_per_sample:0.3f}, " + f"time_speech: {time_speech_total_per_sample: 0.3f}, " + f"time_escape: {time_escape_total_per_sample:0.3f}" + ) # end_total = time.time() # time_escape_total_all_samples = end_total - beg_total @@ -490,9 +546,8 @@ class AutoModel: return results_ret_list def export(self, input=None, **cfg): - """ - + :param input: :param type: :param quantize: @@ -502,7 +557,7 @@ class AutoModel: :param cfg: :return: """ - + device = cfg.get("device", "cpu") model = self.model.to(device=device) kwargs = self.kwargs @@ -513,19 +568,17 @@ class AutoModel: type = kwargs.get("type", "onnx") - key_list, data_list = prepare_data_iterator(input, input_len=None, data_type=kwargs.get("data_type", None), key=None) + key_list, data_list = prepare_data_iterator( + input, input_len=None, data_type=kwargs.get("data_type", None), key=None + ) with torch.no_grad(): - + if type == "onnx": - export_dir = export_utils.export_onnx( - model=model, - data_in=data_list, - **kwargs) + export_dir = export_utils.export_onnx(model=model, data_in=data_list, **kwargs) else: export_dir = export_utils.export_torchscripts( - model=model, - data_in=data_list, - **kwargs) + model=model, data_in=data_list, **kwargs + ) - return export_dir \ No newline at end of file + return export_dir diff --git a/funasr/auto/auto_tokenizer.py b/funasr/auto/auto_tokenizer.py index d5082e2e8..59ff8893d 100644 --- a/funasr/auto/auto_tokenizer.py +++ b/funasr/auto/auto_tokenizer.py @@ -1,8 +1,7 @@ - - class AutoTokenizer: - """ - Undo - """ - def __init__(self): - pass \ No newline at end of file + """ + Undo + """ + + def __init__(self): + pass diff --git a/funasr/bin/compute_audio_cmvn.py b/funasr/bin/compute_audio_cmvn.py index cd64329c3..91bb2ac84 100644 --- a/funasr/bin/compute_audio_cmvn.py +++ b/funasr/bin/compute_audio_cmvn.py @@ -14,13 +14,14 @@ from funasr.train_utils.set_all_random_seed import set_all_random_seed @hydra.main(config_name=None, version_base=None) def main_hydra(kwargs: DictConfig): if kwargs.get("debug", False): - import pdb; pdb.set_trace() + import pdb + + pdb.set_trace() assert "model" in kwargs if "model_conf" not in kwargs: logging.info("download models from model hub: {}".format(kwargs.get("hub", "ms"))) kwargs = download_model(is_training=kwargs.get("is_training", True), **kwargs) - main(**kwargs) @@ -33,12 +34,9 @@ def main(**kwargs): torch.backends.cudnn.enabled = kwargs.get("cudnn_enabled", torch.backends.cudnn.enabled) torch.backends.cudnn.benchmark = kwargs.get("cudnn_benchmark", torch.backends.cudnn.benchmark) torch.backends.cudnn.deterministic = kwargs.get("cudnn_deterministic", True) - - - tokenizer = kwargs.get("tokenizer", None) - + # build frontend if frontend is none None frontend = kwargs.get("frontend", None) if frontend is not None: @@ -47,11 +45,15 @@ def main(**kwargs): kwargs["frontend"] = frontend kwargs["input_size"] = frontend.output_size() - - # dataset dataset_class = tables.dataset_classes.get(kwargs.get("dataset", "AudioDataset")) - dataset_train = dataset_class(kwargs.get("train_data_set_list"), frontend=frontend, tokenizer=None, is_training=False, **kwargs.get("dataset_conf")) + dataset_train = dataset_class( + kwargs.get("train_data_set_list"), + frontend=frontend, + tokenizer=None, + is_training=False, + **kwargs.get("dataset_conf") + ) # dataloader batch_sampler = kwargs["dataset_conf"].get("batch_sampler", "BatchSampler") @@ -62,9 +64,11 @@ def main(**kwargs): dataset_conf["num_workers"] = os.cpu_count() or 32 batch_sampler_train = batch_sampler_class(dataset_train, is_training=False, **dataset_conf) - dataloader_train = torch.utils.data.DataLoader(dataset_train, collate_fn=dataset_train.collator, **batch_sampler_train) - - iter_stop = int(kwargs.get("scale", 1.0)*len(dataloader_train)) + dataloader_train = torch.utils.data.DataLoader( + dataset_train, collate_fn=dataset_train.collator, **batch_sampler_train + ) + + iter_stop = int(kwargs.get("scale", 1.0) * len(dataloader_train)) total_frames = 0 for batch_idx, batch in enumerate(dataloader_train): @@ -79,33 +83,46 @@ def main(**kwargs): mean_stats += np.sum(fbank, axis=0) var_stats += np.sum(np.square(fbank), axis=0) total_frames += fbank.shape[0] - - + cmvn_info = { - 'mean_stats': list(mean_stats.tolist()), - 'var_stats': list(var_stats.tolist()), - 'total_frames': total_frames + "mean_stats": list(mean_stats.tolist()), + "var_stats": list(var_stats.tolist()), + "total_frames": total_frames, } cmvn_file = kwargs.get("cmvn_file", "cmvn.json") # import pdb;pdb.set_trace() - with open(cmvn_file, 'w') as fout: + with open(cmvn_file, "w") as fout: fout.write(json.dumps(cmvn_info)) - + mean = -1.0 * mean_stats / total_frames var = 1.0 / np.sqrt(var_stats / total_frames - mean * mean) dims = mean.shape[0] am_mvn = os.path.dirname(cmvn_file) + "/am.mvn" - with open(am_mvn, 'w') as fout: - fout.write("" + "\n" + " " + str(dims) + " " + str(dims) + '\n' + "[ 0 ]" + "\n" + " " + str(dims) + " " + str(dims) + "\n") - mean_str = str(list(mean)).replace(',', '').replace('[', '[ ').replace(']', ' ]') - fout.write(" 0 " + mean_str + '\n') - fout.write(" " + str(dims) + " " + str(dims) + '\n') - var_str = str(list(var)).replace(',', '').replace('[', '[ ').replace(']', ' ]') - fout.write(" 0 " + var_str + '\n') - fout.write("" + '\n') - - - + with open(am_mvn, "w") as fout: + fout.write( + "" + + "\n" + + " " + + str(dims) + + " " + + str(dims) + + "\n" + + "[ 0 ]" + + "\n" + + " " + + str(dims) + + " " + + str(dims) + + "\n" + ) + mean_str = str(list(mean)).replace(",", "").replace("[", "[ ").replace("]", " ]") + fout.write(" 0 " + mean_str + "\n") + fout.write(" " + str(dims) + " " + str(dims) + "\n") + var_str = str(list(var)).replace(",", "").replace("[", "[ ").replace("]", " ]") + fout.write(" 0 " + var_str + "\n") + fout.write("" + "\n") + + """ python funasr/bin/compute_audio_cmvn.py \ --config-path "/Users/zhifu/funasr1.0/examples/aishell/paraformer/conf" \ diff --git a/funasr/bin/export.py b/funasr/bin/export.py index cb160e9c5..6c9b49fd8 100644 --- a/funasr/bin/export.py +++ b/funasr/bin/export.py @@ -15,28 +15,31 @@ def main_hydra(cfg: DictConfig): return {k: to_plain_list(v) for k, v in cfg_item.items()} else: return cfg_item - + kwargs = to_plain_list(cfg) log_level = getattr(logging, kwargs.get("log_level", "INFO").upper()) logging.basicConfig(level=log_level) if kwargs.get("debug", False): - import pdb; pdb.set_trace() + import pdb + + pdb.set_trace() if "device" not in kwargs: kwargs["device"] = "cpu" model = AutoModel(**kwargs) - - res = model.export(input=kwargs.get("input", None), - type=kwargs.get("type", "onnx"), - quantize=kwargs.get("quantize", False), - fallback_num=kwargs.get("fallback-num", 5), - calib_num=kwargs.get("calib_num", 100), - opset_version=kwargs.get("opset_version", 14), - ) + + res = model.export( + input=kwargs.get("input", None), + type=kwargs.get("type", "onnx"), + quantize=kwargs.get("quantize", False), + fallback_num=kwargs.get("fallback-num", 5), + calib_num=kwargs.get("calib_num", 100), + opset_version=kwargs.get("opset_version", 14), + ) print(res) -if __name__ == '__main__': - main_hydra() \ No newline at end of file +if __name__ == "__main__": + main_hydra() diff --git a/funasr/bin/inference.py b/funasr/bin/inference.py index d2f0c149d..2a1b6aaaa 100644 --- a/funasr/bin/inference.py +++ b/funasr/bin/inference.py @@ -14,18 +14,20 @@ def main_hydra(cfg: DictConfig): return {k: to_plain_list(v) for k, v in cfg_item.items()} else: return cfg_item - + kwargs = to_plain_list(cfg) log_level = getattr(logging, kwargs.get("log_level", "INFO").upper()) logging.basicConfig(level=log_level) if kwargs.get("debug", False): - import pdb; pdb.set_trace() + import pdb + + pdb.set_trace() model = AutoModel(**kwargs) res = model.generate(input=kwargs["input"]) print(res) -if __name__ == '__main__': - main_hydra() \ No newline at end of file +if __name__ == "__main__": + main_hydra() diff --git a/funasr/bin/tokenize_text.py b/funasr/bin/tokenize_text.py index 0ecf2f619..8c2aae050 100755 --- a/funasr/bin/tokenize_text.py +++ b/funasr/bin/tokenize_text.py @@ -133,7 +133,7 @@ def tokenize( if not write_vocabulary: return - + ## FIXME ## del duplicate add_symbols in counter for symbol_and_id in add_symbol: @@ -197,12 +197,8 @@ def get_parser() -> argparse.ArgumentParser: help="The verbose level of logging", ) - parser.add_argument( - "--input", "-i", required=True, help="Input text. - indicates sys.stdin" - ) - parser.add_argument( - "--output", "-o", required=True, help="Output text. - indicates sys.stdout" - ) + parser.add_argument("--input", "-i", required=True, help="Input text. - indicates sys.stdin") + parser.add_argument("--output", "-o", required=True, help="Output text. - indicates sys.stdout") parser.add_argument( "--field", "-f", diff --git a/funasr/bin/train.py b/funasr/bin/train.py index c02a66ff7..05942cdb1 100644 --- a/funasr/bin/train.py +++ b/funasr/bin/train.py @@ -34,22 +34,24 @@ from funasr.train_utils.load_pretrained_model import load_pretrained_model from funasr.utils.misc import prepare_model_dir from funasr import AutoModel + @hydra.main(config_name=None, version_base=None) def main_hydra(kwargs: DictConfig): if kwargs.get("debug", False): - import pdb; pdb.set_trace() + import pdb + + pdb.set_trace() assert "model" in kwargs if "model_conf" not in kwargs: logging.info("download models from model hub: {}".format(kwargs.get("hub", "ms"))) kwargs = download_model(is_training=kwargs.get("is_training", True), **kwargs) - main(**kwargs) def main(**kwargs): - + # set random seed set_all_random_seed(kwargs.get("seed", 0)) torch.backends.cudnn.enabled = kwargs.get("cudnn_enabled", torch.backends.cudnn.enabled) @@ -57,28 +59,32 @@ def main(**kwargs): torch.backends.cudnn.deterministic = kwargs.get("cudnn_deterministic", True) # open tf32 torch.backends.cuda.matmul.allow_tf32 = kwargs.get("enable_tf32", True) - - local_rank = int(os.environ.get('LOCAL_RANK', 0)) + + local_rank = int(os.environ.get("LOCAL_RANK", 0)) if local_rank == 0: tables.print() # Check if we are using DDP or FSDP - use_ddp = 'WORLD_SIZE' in os.environ and int(os.environ["WORLD_SIZE"]) > 1 + use_ddp = "WORLD_SIZE" in os.environ and int(os.environ["WORLD_SIZE"]) > 1 use_fsdp = kwargs.get("use_fsdp", False) # use_ddp = False if use_fsdp else use_fsdp if use_ddp or use_fsdp: - dist.init_process_group(backend=kwargs.get("backend", "nccl"), init_method='env://') + dist.init_process_group(backend=kwargs.get("backend", "nccl"), init_method="env://") torch.cuda.set_device(local_rank) logging.info("Build model, frontend, tokenizer") device = kwargs.get("device", "cuda") kwargs["device"] = "cpu" model = AutoModel(**kwargs) - - + # save config.yaml - if (use_ddp or use_fsdp) and dist.get_rank() == 0 or not (use_ddp or use_fsdp) and local_rank == 0: + if ( + (use_ddp or use_fsdp) + and dist.get_rank() == 0 + or not (use_ddp or use_fsdp) + and local_rank == 0 + ): prepare_model_dir(**kwargs) - + # parse kwargs kwargs = model.kwargs kwargs["device"] = device @@ -100,12 +106,16 @@ def main(**kwargs): if k.startswith(t + ".") or k == t: logging.info(f"Setting {k}.requires_grad = False") p.requires_grad = False - if use_ddp: model = model.cuda(local_rank) - model = DDP(model, device_ids=[local_rank], - find_unused_parameters=kwargs.get("train_conf", {}).get("find_unused_parameters", False)) + model = DDP( + model, + device_ids=[local_rank], + find_unused_parameters=kwargs.get("train_conf", {}).get( + "find_unused_parameters", False + ), + ) elif use_fsdp: # model = FSDP(model).cuda(local_rank) @@ -124,24 +134,26 @@ def main(**kwargs): # Configure a custom `min_num_params` my_auto_wrap_policy = functools.partial(custom_auto_wrap_policy, min_num_params=int(1e5)) torch.cuda.set_device(local_rank) - model = FSDP(model, - auto_wrap_policy=custom_auto_wrap_policy, - mixed_precision=None, - device_id=torch.cuda.current_device()) + model = FSDP( + model, + auto_wrap_policy=custom_auto_wrap_policy, + mixed_precision=None, + device_id=torch.cuda.current_device(), + ) else: model = model.to(device=kwargs.get("device", "cuda")) if local_rank == 0: logging.info(f"{model}") kwargs["device"] = next(model.parameters()).device - + # optim logging.info("Build optim") optim = kwargs.get("optim", "adam") assert optim in optim_classes optim_class = optim_classes.get(optim) optim = optim_class(model.parameters(), **kwargs.get("optim_conf")) - + # scheduler logging.info("Build scheduler") scheduler = kwargs.get("scheduler", "warmuplr") @@ -149,19 +161,21 @@ def main(**kwargs): scheduler_class = scheduler_classes.get(scheduler) scheduler = scheduler_class(optim, **kwargs.get("scheduler_conf")) - # dataset logging.info("Build dataloader") - dataloader_class = tables.dataloader_classes.get(kwargs["dataset_conf"].get("dataloader", "DataloaderMapStyle")) + dataloader_class = tables.dataloader_classes.get( + kwargs["dataset_conf"].get("dataloader", "DataloaderMapStyle") + ) dataloader = dataloader_class(**kwargs) # dataloader_tr, dataloader_val = dataloader_class(**kwargs) - trainer = Trainer(local_rank=local_rank, - use_ddp=use_ddp, - use_fsdp=use_fsdp, - device=kwargs["device"], - output_dir=kwargs.get("output_dir", "./exp"), - **kwargs.get("train_conf"), - ) + trainer = Trainer( + local_rank=local_rank, + use_ddp=use_ddp, + use_fsdp=use_fsdp, + device=kwargs["device"], + output_dir=kwargs.get("output_dir", "./exp"), + **kwargs.get("train_conf"), + ) scaler = GradScaler(enabled=trainer.use_fp16) if trainer.use_fp16 else None scaler = ShardedGradScaler(enabled=trainer.use_fp16) if trainer.use_fsdp else scaler @@ -172,38 +186,35 @@ def main(**kwargs): os.makedirs(tensorboard_dir, exist_ok=True) try: from tensorboardX import SummaryWriter + writer = SummaryWriter(tensorboard_dir) if trainer.rank == 0 else None except: writer = None - + dataloader_tr, dataloader_val = None, None for epoch in range(trainer.start_epoch, trainer.max_epoch + 1): time1 = time.perf_counter() - + for data_split_i in range(dataloader.data_split_num): dataloader_tr, dataloader_val = dataloader.build_iter(epoch, data_split_i=data_split_i) trainer.train_epoch( - model=model, - optim=optim, - scheduler=scheduler, - scaler=scaler, - dataloader_train=dataloader_tr, - dataloader_val=dataloader_val, - epoch=epoch, - writer=writer, - data_split_i=data_split_i, - data_split_num=dataloader.data_split_num, - ) - + model=model, + optim=optim, + scheduler=scheduler, + scaler=scaler, + dataloader_train=dataloader_tr, + dataloader_val=dataloader_val, + epoch=epoch, + writer=writer, + data_split_i=data_split_i, + data_split_num=dataloader.data_split_num, + ) + trainer.validate_epoch( - model=model, - dataloader_val=dataloader_val, - epoch=epoch, - writer=writer + model=model, dataloader_val=dataloader_val, epoch=epoch, writer=writer ) scheduler.step() - trainer.save_checkpoint(epoch, model=model, optim=optim, scheduler=scheduler, scaler=scaler) time2 = time.perf_counter() @@ -212,8 +223,8 @@ def main(**kwargs): f"rank: {local_rank}, " f"time_escaped_epoch: {time_escaped:.3f} hours, " f"estimated to finish {trainer.max_epoch} " - f"epoch: {(trainer.max_epoch - epoch) * time_escaped:.3f} hours\n") - + f"epoch: {(trainer.max_epoch - epoch) * time_escaped:.3f} hours\n" + ) if trainer.rank == 0: average_checkpoints(trainer.output_dir, trainer.avg_nbest_model) @@ -221,7 +232,5 @@ def main(**kwargs): trainer.close() - - if __name__ == "__main__": - main_hydra() \ No newline at end of file + main_hydra() diff --git a/funasr/datasets/audio_datasets/datasets.py b/funasr/datasets/audio_datasets/datasets.py index 9825c86e8..2aafde3f2 100644 --- a/funasr/datasets/audio_datasets/datasets.py +++ b/funasr/datasets/audio_datasets/datasets.py @@ -10,28 +10,33 @@ class AudioDataset(torch.utils.data.Dataset): """ AudioDataset """ - def __init__(self, - path, - index_ds: str = None, - frontend=None, - tokenizer=None, - int_pad_value: int = -1, - float_pad_value: float = 0.0, - **kwargs): + + def __init__( + self, + path, + index_ds: str = None, + frontend=None, + tokenizer=None, + int_pad_value: int = -1, + float_pad_value: float = 0.0, + **kwargs, + ): super().__init__() index_ds_class = tables.index_ds_classes.get(index_ds) self.index_ds = index_ds_class(path, **kwargs) preprocessor_speech = kwargs.get("preprocessor_speech", None) if preprocessor_speech: preprocessor_speech_class = tables.preprocessor_classes.get(preprocessor_speech) - preprocessor_speech = preprocessor_speech_class(**kwargs.get("preprocessor_speech_conf")) + preprocessor_speech = preprocessor_speech_class( + **kwargs.get("preprocessor_speech_conf") + ) self.preprocessor_speech = preprocessor_speech preprocessor_text = kwargs.get("preprocessor_text", None) if preprocessor_text: preprocessor_text_class = tables.preprocessor_classes.get(preprocessor_text) preprocessor_text = preprocessor_text_class(**kwargs.get("preprocessor_text_conf")) self.preprocessor_text = preprocessor_text - + self.frontend = frontend self.fs = 16000 if frontend is None else frontend.fs self.data_type = "sound" @@ -39,18 +44,18 @@ class AudioDataset(torch.utils.data.Dataset): self.int_pad_value = int_pad_value self.float_pad_value = float_pad_value - + def get_source_len(self, index): item = self.index_ds[index] return self.index_ds.get_source_len(item) - + def get_target_len(self, index): item = self.index_ds[index] return self.index_ds.get_target_len(item) - + def __len__(self): return len(self.index_ds) - + def __getitem__(self, index): item = self.index_ds[index] # import pdb; @@ -59,7 +64,9 @@ class AudioDataset(torch.utils.data.Dataset): data_src = load_audio_text_image_video(source, fs=self.fs) if self.preprocessor_speech: data_src = self.preprocessor_speech(data_src, fs=self.fs) - speech, speech_lengths = extract_fbank(data_src, data_type=self.data_type, frontend=self.frontend, is_final=True) # speech: [b, T, d] + speech, speech_lengths = extract_fbank( + data_src, data_type=self.data_type, frontend=self.frontend, is_final=True + ) # speech: [b, T, d] target = item["target"] if self.preprocessor_text: @@ -73,14 +80,14 @@ class AudioDataset(torch.utils.data.Dataset): ids_lengths = len(ids) text_lengths = torch.tensor([ids_lengths], dtype=torch.int32) - return {"speech": speech[0, :, :], - "speech_lengths": speech_lengths, - "text": text, - "text_lengths": text_lengths, - } - - - def collator(self, samples: list=None): + return { + "speech": speech[0, :, :], + "speech_lengths": speech_lengths, + "text": text, + "text_lengths": text_lengths, + } + + def collator(self, samples: list = None): outputs = {} for sample in samples: for key in sample.keys(): @@ -91,12 +98,14 @@ class AudioDataset(torch.utils.data.Dataset): for key, data_list in outputs.items(): if isinstance(data_list[0], torch.Tensor): if data_list[0].dtype == torch.int64 or data_list[0].dtype == torch.int32: - + pad_value = self.int_pad_value else: pad_value = self.float_pad_value - - outputs[key] = torch.nn.utils.rnn.pad_sequence(data_list, batch_first=True, padding_value=pad_value) + + outputs[key] = torch.nn.utils.rnn.pad_sequence( + data_list, batch_first=True, padding_value=pad_value + ) return outputs @@ -111,7 +120,7 @@ class AudioDatasetHotword(AudioDataset): ): super().__init__(*args, **kwargs) self.seaco_id = seaco_id - + def __getitem__(self, index): item = self.index_ds[index] # import pdb; @@ -120,7 +129,9 @@ class AudioDatasetHotword(AudioDataset): data_src = load_audio_text_image_video(source, fs=self.fs) if self.preprocessor_speech: data_src = self.preprocessor_speech(data_src, fs=self.fs) - speech, speech_lengths = extract_fbank(data_src, data_type=self.data_type, frontend=self.frontend, is_final=True) # speech: [b, T, d] + speech, speech_lengths = extract_fbank( + data_src, data_type=self.data_type, frontend=self.frontend, is_final=True + ) # speech: [b, T, d] target = item["target"] if self.preprocessor_text: @@ -133,57 +144,71 @@ class AudioDatasetHotword(AudioDataset): text = ids ids_lengths = len(ids) text_lengths = torch.tensor([ids_lengths], dtype=torch.int32) - - def generate_index(length, - hotword_min_length=2, - hotword_max_length=8, - sample_rate=0.75, - double_rate=0.1, - pre_prob=0.0, - pre_index=None, - pre_hwlist=None): + + def generate_index( + length, + hotword_min_length=2, + hotword_max_length=8, + sample_rate=0.75, + double_rate=0.1, + pre_prob=0.0, + pre_index=None, + pre_hwlist=None, + ): if length < hotword_min_length: return [-1] if random.random() < sample_rate: if pre_prob > 0 and random.random() < pre_prob and pre_index is not None: return pre_index if length == hotword_min_length: - return [0, length-1] - elif random.random() < double_rate and length > hotword_max_length + hotword_min_length + 2: + return [0, length - 1] + elif ( + random.random() < double_rate + and length > hotword_max_length + hotword_min_length + 2 + ): # sample two hotwords in a sentence _max_hw_length = min(hotword_max_length, length // 2) # first hotword start1 = random.randint(0, length // 3) - end1 = random.randint(start1 + hotword_min_length - 1, start1 + _max_hw_length - 1) + end1 = random.randint( + start1 + hotword_min_length - 1, start1 + _max_hw_length - 1 + ) # second hotword start2 = random.randint(end1 + 1, length - hotword_min_length) - end2 = random.randint(min(length-1, start2+hotword_min_length-1), min(length-1, start2+hotword_max_length-1)) + end2 = random.randint( + min(length - 1, start2 + hotword_min_length - 1), + min(length - 1, start2 + hotword_max_length - 1), + ) return [start1, end1, start2, end2] else: # single hotword start = random.randint(0, length - hotword_min_length) - end = random.randint(min(length-1, start+hotword_min_length-1), min(length-1, start+hotword_max_length-1)) + end = random.randint( + min(length - 1, start + hotword_min_length - 1), + min(length - 1, start + hotword_max_length - 1), + ) return [start, end] else: return [-1] - + hotword_indx = generate_index(text_lengths[0]) - return {"speech": speech[0, :, :], - "speech_lengths": speech_lengths, - "text": text, - "text_lengths": text_lengths, - "hotword_indx": hotword_indx, - "seaco_id": self.seaco_id, - } - - def collator(self, samples: list=None): + return { + "speech": speech[0, :, :], + "speech_lengths": speech_lengths, + "text": text, + "text_lengths": text_lengths, + "hotword_indx": hotword_indx, + "seaco_id": self.seaco_id, + } + + def collator(self, samples: list = None): outputs = {} hotword_indxs = [] - seaco_id = samples[0]['seaco_id'] + seaco_id = samples[0]["seaco_id"] for sample in samples: for key in sample.keys(): - if key == 'seaco_id': + if key == "seaco_id": continue - elif key == 'hotword_indx': + elif key == "hotword_indx": hotword_indxs.append(sample[key]) else: if key not in outputs: @@ -196,32 +221,40 @@ class AudioDatasetHotword(AudioDataset): pad_value = self.int_pad_value else: pad_value = self.float_pad_value - outputs[key] = torch.nn.utils.rnn.pad_sequence(data_list, batch_first=True, padding_value=pad_value) - + outputs[key] = torch.nn.utils.rnn.pad_sequence( + data_list, batch_first=True, padding_value=pad_value + ) + hotword_list, hotword_lengths = [], [] - text = outputs['text'] + text = outputs["text"] seaco_label_pad = torch.ones_like(text) * -1 if seaco_id else None - for b, (hotword_indx, one_text, length) in enumerate(zip(hotword_indxs, text, outputs['text_lengths'])): + for b, (hotword_indx, one_text, length) in enumerate( + zip(hotword_indxs, text, outputs["text_lengths"]) + ): length = length[0] - if seaco_label_pad is not None: seaco_label_pad[b][:length] = seaco_id + if seaco_label_pad is not None: + seaco_label_pad[b][:length] = seaco_id if hotword_indx[0] != -1: start, end = int(hotword_indx[0]), int(hotword_indx[1]) - hotword = one_text[start: end+1] + hotword = one_text[start : end + 1] hotword_list.append(hotword) - hotword_lengths.append(end-start+1) - if seaco_label_pad is not None: seaco_label_pad[b][start: end+1] = one_text[start: end+1] + hotword_lengths.append(end - start + 1) + if seaco_label_pad is not None: + seaco_label_pad[b][start : end + 1] = one_text[start : end + 1] if len(hotword_indx) == 4 and hotword_indx[2] != -1: # the second hotword if exist start, end = int(hotword_indx[2]), int(hotword_indx[3]) - hotword_list.append(one_text[start: end+1]) - hotword_lengths.append(end-start+1) - if seaco_label_pad is not None: seaco_label_pad[b][start: end+1] = one_text[start: end+1] + hotword_list.append(one_text[start : end + 1]) + hotword_lengths.append(end - start + 1) + if seaco_label_pad is not None: + seaco_label_pad[b][start : end + 1] = one_text[start : end + 1] hotword_list.append(torch.tensor([1])) hotword_lengths.append(1) - hotword_pad = torch.nn.utils.rnn.pad_sequence(hotword_list, - batch_first=True, - padding_value=0) + hotword_pad = torch.nn.utils.rnn.pad_sequence( + hotword_list, batch_first=True, padding_value=0 + ) outputs["hotword_pad"] = hotword_pad outputs["hotword_lengths"] = torch.tensor(hotword_lengths, dtype=torch.int32) - if seaco_label_pad is not None: outputs['seaco_label_pad'] = seaco_label_pad - return outputs \ No newline at end of file + if seaco_label_pad is not None: + outputs["seaco_label_pad"] = seaco_label_pad + return outputs diff --git a/funasr/datasets/audio_datasets/espnet_samplers.py b/funasr/datasets/audio_datasets/espnet_samplers.py index 6b38bc217..e6efe0aae 100644 --- a/funasr/datasets/audio_datasets/espnet_samplers.py +++ b/funasr/datasets/audio_datasets/espnet_samplers.py @@ -18,7 +18,7 @@ def EspnetStyleBatchSampler_fn(dataset, **kwargs): dataloader_args["batch_sampler"] = batch_sampler dataloader_args["num_workers"] = kwargs.get("num_workers", 4) dataloader_args["pin_memory"] = kwargs.get("pin_memory", True) - + return dataloader_args @@ -29,18 +29,20 @@ import random class EspnetStyleBatchSampler(DistributedSampler): - def __init__(self, dataset, - batch_size, - batch_type="token", - rank=None, - num_replicas=None, - rank_split=False, - shuffle=True, - drop_last=False, - is_training: bool = True, - sort_size: int = 1024, - **kwargs, - ): + def __init__( + self, + dataset, + batch_size, + batch_type="token", + rank=None, + num_replicas=None, + rank_split=False, + shuffle=True, + drop_last=False, + is_training: bool = True, + sort_size: int = 1024, + **kwargs, + ): try: rank = dist.get_rank() @@ -69,9 +71,9 @@ class EspnetStyleBatchSampler(DistributedSampler): self.min_token_length = kwargs.get("min_token_length", 0) self.length_scale_source = kwargs.get("length_scale_source", 1.0) - # super().__init__(dataset, num_replicas=num_replicas, rank=rank, # shuffle=shuffle, drop_last=drop_last) + def __iter__(self): if self.shuffle: g = torch.Generator() @@ -80,18 +82,21 @@ class EspnetStyleBatchSampler(DistributedSampler): indices = torch.randperm(len(self.dataset), generator=g).tolist() else: indices = list(range(len(self.dataset))) - + # Sort indices by sample length sorted_indices = sorted(indices, key=lambda idx: self.dataset.get_source_len(idx)) - + # Organize batches based on 'length' or 'example' buffer_batches = [] batch = [] max_len_in_batch = 0 # Tracks the max sample length within the current batch - + for idx in sorted_indices: original_sample_length = self.dataset.get_source_len(idx) - if original_sample_length < self.min_token_length or original_sample_length > self.max_token_length: # Skip samples that exceed the max length + if ( + original_sample_length < self.min_token_length + or original_sample_length > self.max_token_length + ): # Skip samples that exceed the max length continue # Set sample_length based on the batch type sample_length = 1 if self.batch_type == "example" else original_sample_length @@ -106,37 +111,35 @@ class EspnetStyleBatchSampler(DistributedSampler): buffer_batches.append(batch) batch = [idx] max_len_in_batch = sample_length - + # Add the last batch if it shouldn't be dropped if batch and (not self.drop_last or len(batch) * max_len_in_batch == self.batch_size): buffer_batches.append(batch) - + # Shuffle the list of batches if self.shuffle: random.seed(self.epoch) random.shuffle(buffer_batches) - + # Ensure each rank gets the same number of batches batches_per_rank = int(math.ceil(len(buffer_batches) / self.num_replicas)) total_batches_needed = batches_per_rank * self.num_replicas extra_batches = total_batches_needed - len(buffer_batches) # Add extra batches by random selection, if needed buffer_batches += random.choices(buffer_batches, k=extra_batches) - + # Allocate the batches to the current rank start_idx = self.rank * batches_per_rank end_idx = start_idx + batches_per_rank rank_batches = buffer_batches[start_idx:end_idx] - + # Return an iterator over the batches for the current rank return iter(rank_batches) - + def __len__(self): # Calculate the number of batches per epoch for the current rank return 1 - + def set_epoch(self, epoch): # Set the epoch for shuffling self.epoch = epoch - - diff --git a/funasr/datasets/audio_datasets/index_ds.py b/funasr/datasets/audio_datasets/index_ds.py index 2677d3348..d26124bed 100644 --- a/funasr/datasets/audio_datasets/index_ds.py +++ b/funasr/datasets/audio_datasets/index_ds.py @@ -14,7 +14,7 @@ from funasr.register import tables @tables.register("index_ds_classes", "IndexDSJsonlRankFull") @tables.register("index_ds_classes", "IndexDSJsonlRankSplit") class IndexDSJsonlRankFull(torch.utils.data.Dataset): - + def __init__(self, path: str, **kwargs): super().__init__() self.max_source_length = kwargs.get("max_source_length", 2048) @@ -27,21 +27,23 @@ class IndexDSJsonlRankFull(torch.utils.data.Dataset): # jsonl list file data_split_num = kwargs.get("data_split_num", 1) data_split_i = kwargs.get("data_split_i", 0) - + if not is_training: data_split_num = 1 data_split_i = 0 - with open(path, encoding='utf-8') as fin: + with open(path, encoding="utf-8") as fin: file_list_all = fin.readlines() - - num_per_slice = (len(file_list_all)-1) // data_split_num + 1 - file_list = file_list_all[data_split_i * num_per_slice:(data_split_i + 1) * num_per_slice] + + num_per_slice = (len(file_list_all) - 1) // data_split_num + 1 + file_list = file_list_all[ + data_split_i * num_per_slice : (data_split_i + 1) * num_per_slice + ] logging.info( - f"is_training: {is_training}, data_split_num: {data_split_num}, data_split_i: {data_split_i}, \nfile_list: {file_list}, \nfile_list_all: {file_list_all}") - + f"is_training: {is_training}, data_split_num: {data_split_num}, data_split_i: {data_split_i}, \nfile_list: {file_list}, \nfile_list_all: {file_list_all}" + ) + else: file_list = [path] - # total_num = len(file_list) # try: @@ -76,29 +78,38 @@ class IndexDSJsonlRankFull(torch.utils.data.Dataset): # for file_json in file_list_rank: contents = [] for file_json in file_list: - with open(file_json.strip(), encoding='utf-8') as fin: + with open(file_json.strip(), encoding="utf-8") as fin: for line in fin: data = json.loads(line.strip()) if "text" in data: # for sft - contents.append(data['text']) + contents.append(data["text"]) if "source" in data: # for speech lab pretrain prompt = data.get("prompt", "") - source = data["source"].replace("/cpfs01", "/cpfs_speech/data") # only use in alibaba gpu group: .replace("/cpfs01", "/cpfs_speech/data") + source = data["source"].replace( + "/cpfs01", "/cpfs_speech/data" + ) # only use in alibaba gpu group: .replace("/cpfs01", "/cpfs_speech/data") target = data["target"] source_len = data.get("source_len", 1) target_len = data.get("target_len", 0) if "aishell" in source: target = target.replace(" ", "") - if source_len < self.min_source_length or source_len > self.max_source_length: + if ( + source_len < self.min_source_length + or source_len > self.max_source_length + ): continue - if target_len < self.min_target_length or target_len > self.max_target_length: + if ( + target_len < self.min_target_length + or target_len > self.max_target_length + ): continue - contents_i = {"source": source, - "prompt": prompt, - "target": target, - "source_len": source_len, - "target_len": target_len, - } + contents_i = { + "source": source, + "prompt": prompt, + "target": target, + "source_len": source_len, + "target_len": target_len, + } text_language = data.get("text_language", None) if text_language is not None: contents_i["text_language"] = text_language @@ -108,23 +119,21 @@ class IndexDSJsonlRankFull(torch.utils.data.Dataset): contents.append(contents_i) self.contents = contents - - logging.info( - "total_num of samplers: {}, {}".format(len(self.contents), path)) - + + logging.info("total_num of samplers: {}, {}".format(len(self.contents), path)) + def __len__(self): return len(self.contents) - + def __getitem__(self, index): - + data = self.contents[index] return data - + def get_source_len(self, data_dict): return data_dict.get("source_len", 1) - - def get_target_len(self, data_dict): - - return data_dict.get("target_len", 0) + def get_target_len(self, data_dict): + + return data_dict.get("target_len", 0) diff --git a/funasr/datasets/audio_datasets/jsonl2scp.py b/funasr/datasets/audio_datasets/jsonl2scp.py index b2d1703d0..3ce96ca78 100644 --- a/funasr/datasets/audio_datasets/jsonl2scp.py +++ b/funasr/datasets/audio_datasets/jsonl2scp.py @@ -9,15 +9,14 @@ import librosa import torch.distributed as dist - def gen_scp_from_jsonl(jsonl_file, data_type_list, wav_scp_file, text_file): wav_f = open(wav_scp_file, "w") text_f = open(text_file, "w") - with open(jsonl_file, encoding='utf-8') as fin: + with open(jsonl_file, encoding="utf-8") as fin: for line in fin: data = json.loads(line.strip()) - + prompt = data.get("prompt", "") source = data[data_type_list[0]] target = data[data_type_list[1]] @@ -33,22 +32,26 @@ def gen_scp_from_jsonl(jsonl_file, data_type_list, wav_scp_file, text_file): wav_f.close() text_f.close() - - - + + @hydra.main(config_name=None, version_base=None) def main_hydra(cfg: DictConfig): - + kwargs = OmegaConf.to_container(cfg, resolve=True) print(kwargs) - scp_file_list = kwargs.get("scp_file_list", ("/Users/zhifu/funasr1.0/test_local/wav.scp", "/Users/zhifu/funasr1.0/test_local/text.txt")) + scp_file_list = kwargs.get( + "scp_file_list", + ("/Users/zhifu/funasr1.0/test_local/wav.scp", "/Users/zhifu/funasr1.0/test_local/text.txt"), + ) if isinstance(scp_file_list, str): scp_file_list = eval(scp_file_list) data_type_list = kwargs.get("data_type_list", ("source", "target")) - jsonl_file = kwargs.get("jsonl_file_in", "/Users/zhifu/funasr1.0/test_local/audio_datasets.jsonl") + jsonl_file = kwargs.get( + "jsonl_file_in", "/Users/zhifu/funasr1.0/test_local/audio_datasets.jsonl" + ) gen_scp_from_jsonl(jsonl_file, data_type_list, *scp_file_list) - + """ python -m funasr.datasets.audio_datasets.json2scp \ @@ -59,5 +62,3 @@ python -m funasr.datasets.audio_datasets.json2scp \ if __name__ == "__main__": main_hydra() - - \ No newline at end of file diff --git a/funasr/datasets/audio_datasets/preprocessor.py b/funasr/datasets/audio_datasets/preprocessor.py index ab751401b..b3f942156 100644 --- a/funasr/datasets/audio_datasets/preprocessor.py +++ b/funasr/datasets/audio_datasets/preprocessor.py @@ -17,35 +17,39 @@ from funasr.register import tables @tables.register("preprocessor_classes", "SpeechPreprocessSpeedPerturb") class SpeechPreprocessSpeedPerturb(nn.Module): - def __init__(self, speed_perturb: list=None, **kwargs): - super().__init__() - self.speed_perturb = speed_perturb - - def forward(self, waveform, fs, **kwargs): - if self.speed_perturb is None: - return waveform - speed = random.choice(self.speed_perturb) - if speed != 1.0: - if not isinstance(waveform, torch.Tensor): - waveform = torch.tensor(waveform) - waveform, _ = torchaudio.sox_effects.apply_effects_tensor( - waveform.view(1, -1), fs, [['speed', str(speed)], ['rate', str(fs)]]) - waveform = waveform.view(-1) - - return waveform + def __init__(self, speed_perturb: list = None, **kwargs): + super().__init__() + self.speed_perturb = speed_perturb + + def forward(self, waveform, fs, **kwargs): + if self.speed_perturb is None: + return waveform + speed = random.choice(self.speed_perturb) + if speed != 1.0: + if not isinstance(waveform, torch.Tensor): + waveform = torch.tensor(waveform) + waveform, _ = torchaudio.sox_effects.apply_effects_tensor( + waveform.view(1, -1), fs, [["speed", str(speed)], ["rate", str(fs)]] + ) + waveform = waveform.view(-1) + + return waveform @tables.register("preprocessor_classes", "TextPreprocessSegDict") class TextPreprocessSegDict(nn.Module): - def __init__(self, seg_dict: str = None, - text_cleaner: Collection[str] = None, - split_with_space: bool = False, - **kwargs): - super().__init__() - - self.text_cleaner = TextCleaner(text_cleaner) - - def forward(self, text, **kwargs): - text = self.text_cleaner(text) - - return text + def __init__( + self, + seg_dict: str = None, + text_cleaner: Collection[str] = None, + split_with_space: bool = False, + **kwargs + ): + super().__init__() + + self.text_cleaner = TextCleaner(text_cleaner) + + def forward(self, text, **kwargs): + text = self.text_cleaner(text) + + return text diff --git a/funasr/datasets/audio_datasets/samplers.py b/funasr/datasets/audio_datasets/samplers.py index 1394f7e8c..18f8f91b4 100644 --- a/funasr/datasets/audio_datasets/samplers.py +++ b/funasr/datasets/audio_datasets/samplers.py @@ -22,30 +22,33 @@ def CustomDistributedBatchSampler_fn(dataset, **kwargs): batch_type = kwargs.get("batch_type", "example") if batch_type == "example": batch_sampler = CustomDistributedBatchSampler(dataset, **kwargs) - + else: if kwargs.get("sort_size", -1) > 0: batch_sampler = CustomDistributedBufferDynamicBatchSampler(dataset, **kwargs) else: batch_sampler = CustomDistributedDynamicBatchSampler(dataset, **kwargs) # batch_sampler = CustomDistributedDynamicBatchSampler(dataset, **kwargs) - + dataloader_args["batch_sampler"] = batch_sampler dataloader_args["num_workers"] = kwargs.get("num_workers", 4) dataloader_args["pin_memory"] = kwargs.get("pin_memory", True) - + return dataloader_args + class CustomDistributedBatchSampler(Sampler): - def __init__(self, dataset, - batch_size, - num_replicas=None, - rank=None, - shuffle=True, - drop_last=False, - is_training: bool = True, - **kwargs, - ): + def __init__( + self, + dataset, + batch_size, + num_replicas=None, + rank=None, + shuffle=True, + drop_last=False, + is_training: bool = True, + **kwargs, + ): try: rank = dist.get_rank() @@ -62,9 +65,13 @@ class CustomDistributedBatchSampler(Sampler): self.drop_last = drop_last # self.total_size = len(dataset) if self.drop_last: - self.total_size = (len(self.dataset) // (batch_size * num_replicas)) * (batch_size * num_replicas) + self.total_size = (len(self.dataset) // (batch_size * num_replicas)) * ( + batch_size * num_replicas + ) else: - self.total_size = math.ceil(len(self.dataset) / (batch_size * num_replicas)) * (batch_size * num_replicas) + self.total_size = math.ceil(len(self.dataset) / (batch_size * num_replicas)) * ( + batch_size * num_replicas + ) self.num_samples = int(self.total_size // self.num_replicas) self.epoch = 0 self.max_token_length = kwargs.get("max_token_length", None) @@ -84,12 +91,14 @@ class CustomDistributedBatchSampler(Sampler): if padding_size <= len(indices): indices += indices[:padding_size] else: - indices += (indices * (padding_size // len(indices)) + indices[:padding_size % len(indices)]) + indices += ( + indices * (padding_size // len(indices)) + indices[: padding_size % len(indices)] + ) assert len(indices) == self.total_size # Subsample - indices = indices[self.rank:self.total_size:self.num_replicas] + indices = indices[self.rank : self.total_size : self.num_replicas] assert len(indices) == self.num_samples # Filter out indices with length greater than the max length, if provided @@ -102,7 +111,9 @@ class CustomDistributedBatchSampler(Sampler): indices = filtered_indices # Now that we have only the indices for this replica, chunk them into batches - batches = [indices[i:i + self.batch_size] for i in range(0, len(indices), self.batch_size)] + batches = [ + indices[i : i + self.batch_size] for i in range(0, len(indices), self.batch_size) + ] # Drop the last batch if it's not full and drop_last is True if self.drop_last and len(batches[-1]) != self.batch_size: @@ -117,18 +128,21 @@ class CustomDistributedBatchSampler(Sampler): def set_epoch(self, epoch): self.epoch = epoch + class CustomDistributedBufferBatchSampler(Sampler): - def __init__(self, dataset, - batch_size, - num_replicas=None, - rank=None, - shuffle=True, - drop_last=False, - is_training: bool = True, - sort_size: int = 1024, - **kwargs, - ): - + def __init__( + self, + dataset, + batch_size, + num_replicas=None, + rank=None, + shuffle=True, + drop_last=False, + is_training: bool = True, + sort_size: int = 1024, + **kwargs, + ): + try: rank = dist.get_rank() num_replicas = dist.get_world_size() @@ -144,15 +158,19 @@ class CustomDistributedBufferBatchSampler(Sampler): self.drop_last = drop_last # self.total_size = len(dataset) if self.drop_last: - self.total_size = (len(self.dataset) // (batch_size * num_replicas)) * (batch_size * num_replicas) + self.total_size = (len(self.dataset) // (batch_size * num_replicas)) * ( + batch_size * num_replicas + ) else: - self.total_size = math.ceil(len(self.dataset) / (batch_size * num_replicas)) * (batch_size * num_replicas) + self.total_size = math.ceil(len(self.dataset) / (batch_size * num_replicas)) * ( + batch_size * num_replicas + ) self.num_samples = int(self.total_size // self.num_replicas) self.epoch = 0 self.max_token_length = kwargs.get("max_token_length", None) self.length_scale_source = kwargs.get("length_scale_source", 1.0) self.sort_size = sort_size - + def __iter__(self): # Generate a list of indices if self.shuffle: @@ -161,20 +179,22 @@ class CustomDistributedBufferBatchSampler(Sampler): indices = torch.randperm(len(self.dataset), generator=g).tolist() else: indices = list(range(len(self.dataset))) - + # Add extra samples to make it evenly divisible padding_size = self.total_size - len(indices) if padding_size <= len(indices): indices += indices[:padding_size] else: - indices += (indices * (padding_size // len(indices)) + indices[:padding_size % len(indices)]) - + indices += ( + indices * (padding_size // len(indices)) + indices[: padding_size % len(indices)] + ) + assert len(indices) == self.total_size - + # Subsample - indices = indices[self.rank:self.total_size:self.num_replicas] + indices = indices[self.rank : self.total_size : self.num_replicas] assert len(indices) == self.num_samples - + # Filter out indices with length greater than the max length, if provided if self.max_token_length is not None: filtered_indices = [] @@ -205,29 +225,34 @@ class CustomDistributedBufferBatchSampler(Sampler): def _create_batches_from_buffer(self, buffer): # Function to convert the sorted buffer into batches - batched_buffer = [buffer[i:i + self.batch_size] for i in range(0, len(buffer), self.batch_size)] + batched_buffer = [ + buffer[i : i + self.batch_size] for i in range(0, len(buffer), self.batch_size) + ] if self.drop_last and len(batched_buffer[-1]) != self.batch_size: batched_buffer = batched_buffer[:-1] return batched_buffer - + def __len__(self): - + return self.num_samples // self.batch_size - + def set_epoch(self, epoch): self.epoch = epoch + class CustomDistributedDynamicBatchSampler(DistributedSampler): - def __init__(self, dataset, - batch_size, - num_replicas=None, - rank=None, - shuffle=True, - drop_last=False, - is_training: bool = True, - **kwargs, - ): - + def __init__( + self, + dataset, + batch_size, + num_replicas=None, + rank=None, + shuffle=True, + drop_last=False, + is_training: bool = True, + **kwargs, + ): + try: rank = dist.get_rank() num_replicas = dist.get_world_size() @@ -241,13 +266,13 @@ class CustomDistributedDynamicBatchSampler(DistributedSampler): self.is_training = is_training self.shuffle = shuffle and is_training self.drop_last = drop_last - + self.total_size = len(self.dataset) # self.num_samples = int(math.ceil(self.total_size / self.num_replicas)) self.epoch = 0 self.max_token_length = kwargs.get("max_token_length", 2048) self.length_scale_source = kwargs.get("length_scale_source", 1.0) - + def __iter__(self): if self.shuffle: g = torch.Generator() @@ -255,21 +280,22 @@ class CustomDistributedDynamicBatchSampler(DistributedSampler): indices = torch.randperm(len(self.dataset), generator=g).tolist() else: indices = list(range(len(self.dataset))) - - indices = indices[self.rank:self.total_size:self.num_replicas] - + + indices = indices[self.rank : self.total_size : self.num_replicas] + batches = [] batch = [] max_len_in_batch = 0 current_batch_length = 0 - + for idx in indices: sample_length = self.dataset.get_source_len(idx) if sample_length > self.max_token_length: continue - potential_batch_length = (max_len_in_batch if sample_length < max_len_in_batch else sample_length) * ( - len(batch) + 1) - + potential_batch_length = ( + max_len_in_batch if sample_length < max_len_in_batch else sample_length + ) * (len(batch) + 1) + if potential_batch_length <= self.batch_size: batch.append(idx) if sample_length > max_len_in_batch: @@ -280,35 +306,37 @@ class CustomDistributedDynamicBatchSampler(DistributedSampler): batch = [idx] max_len_in_batch = sample_length # current_batch_length = max_len_in_batch - + # Add the last batch if it's not empty and we're not dropping it if batch and (not self.drop_last or len(batch) * max_len_in_batch == self.batch_size): batches.append(batch) - + return iter(batches) - + def __len__(self): - + return 1 - + def set_epoch(self, epoch): self.epoch = epoch class CustomDistributedBufferDynamicBatchSampler(DistributedSampler): - def __init__(self, dataset, - batch_size, - batch_type="token", - num_replicas=None, - rank=None, - rank_split=False, - shuffle=True, - drop_last=False, - is_training: bool = True, - sort_size: int = 1024, - **kwargs, - ): - + def __init__( + self, + dataset, + batch_size, + batch_type="token", + num_replicas=None, + rank=None, + rank_split=False, + shuffle=True, + drop_last=False, + is_training: bool = True, + sort_size: int = 1024, + **kwargs, + ): + try: rank = dist.get_rank() num_replicas = dist.get_world_size() @@ -320,7 +348,7 @@ class CustomDistributedBufferDynamicBatchSampler(DistributedSampler): # logging.info(f"Warning, rank_split: {rank_split}, batch and shuffle data in local rank") # rank = 0 # num_replicas = 1 - + self.rank = rank self.num_replicas = num_replicas self.dataset = dataset @@ -329,22 +357,23 @@ class CustomDistributedBufferDynamicBatchSampler(DistributedSampler): self.is_training = is_training self.shuffle = shuffle and is_training self.drop_last = drop_last - + self.total_size = len(self.dataset) self.num_samples = int(math.ceil(self.total_size / self.num_replicas)) self.epoch = 0 self.sort_size = sort_size * num_replicas self.max_token_length = kwargs.get("max_token_length", 2048) self.length_scale_source = kwargs.get("length_scale_source", 1.0) - super().__init__(dataset, num_replicas=num_replicas, rank=rank, - shuffle=shuffle, drop_last=drop_last) + super().__init__( + dataset, num_replicas=num_replicas, rank=rank, shuffle=shuffle, drop_last=drop_last + ) def __iter__(self): if self.shuffle: g = torch.Generator() g.manual_seed(self.epoch) random.seed(self.epoch) - + indices = torch.randperm(len(self.dataset), generator=g).tolist() else: indices = list(range(len(self.dataset))) @@ -352,7 +381,9 @@ class CustomDistributedBufferDynamicBatchSampler(DistributedSampler): # Create sorted buffers and form batches buffer_batches = [] for i in range(0, len(indices), self.sort_size): - buffer = sorted(indices[i:i + self.sort_size], key=lambda idx: self.dataset.get_source_len(idx)) + buffer = sorted( + indices[i : i + self.sort_size], key=lambda idx: self.dataset.get_source_len(idx) + ) batch = [] max_len_in_batch = 0 for idx in buffer: @@ -374,10 +405,10 @@ class CustomDistributedBufferDynamicBatchSampler(DistributedSampler): # Ensure each rank gets the same number of batches, duplicate data if needed batches_per_rank = math.ceil(len(buffer_batches) / self.num_replicas) total_batches_needed = batches_per_rank * self.num_replicas - + extra_batches = total_batches_needed - len(buffer_batches) buffer_batches += random.choices(buffer_batches, k=extra_batches) - + # Evenly distribute batches from buffer_batches to each rank rank_batches = [[] for _ in range(self.num_replicas)] for i, batch in enumerate(buffer_batches): @@ -388,17 +419,18 @@ class CustomDistributedBufferDynamicBatchSampler(DistributedSampler): return iter(final_batches) - def __len__(self): - + return 1 - + def set_epoch(self, epoch): self.epoch = epoch class DistributedSamplerWarp(BatchSampler): - def __init__(self, dataset, batch_size, num_replicas=None, rank=None, shuffle=True, drop_last=False): + def __init__( + self, dataset, batch_size, num_replicas=None, rank=None, shuffle=True, drop_last=False + ): if num_replicas is None: if not torch.distributed.is_available(): raise RuntimeError("Requires distributed package to be available") @@ -407,32 +439,29 @@ class DistributedSamplerWarp(BatchSampler): if not torch.distributed.is_available(): raise RuntimeError("Requires distributed package to be available") rank = torch.distributed.get_rank() - + self.dataset = dataset self.batch_size = batch_size self.num_replicas = num_replicas self.rank = rank self.shuffle = shuffle self.drop_last = drop_last - + # Create an instance of the DistributedSampler self.sampler = DistributedSampler( - self.dataset, - num_replicas=self.num_replicas, - rank=self.rank, - shuffle=self.shuffle + self.dataset, num_replicas=self.num_replicas, rank=self.rank, shuffle=self.shuffle ) - + # Call BatchSampler's constructor super().__init__(self.sampler, batch_size, drop_last) - + def __iter__(self): # If we shuffle, we need to call the set_epoch method if self.shuffle: self.sampler.set_epoch(self.epoch) - + # Generate batch indices using the parent class return super().__iter__() - + def set_epoch(self, epoch): self.epoch = epoch diff --git a/funasr/datasets/audio_datasets/scp2jsonl.py b/funasr/datasets/audio_datasets/scp2jsonl.py index 0764298f4..f6ceb6977 100644 --- a/funasr/datasets/audio_datasets/scp2jsonl.py +++ b/funasr/datasets/audio_datasets/scp2jsonl.py @@ -9,8 +9,9 @@ import librosa import torch.distributed as dist - -def gen_jsonl_from_wav_text_list(path, data_type_list=("source", "target"), jsonl_file_out:str=None, **kwargs): +def gen_jsonl_from_wav_text_list( + path, data_type_list=("source", "target"), jsonl_file_out: str = None, **kwargs +): try: rank = dist.get_rank() world_size = dist.get_world_size() @@ -25,18 +26,27 @@ def gen_jsonl_from_wav_text_list(path, data_type_list=("source", "target"), json for data_type, data_file in zip(data_type_list, path): json_dict[data_type] = {} with open(data_file, "r") as f: - + data_file_lists = f.readlines() - lines_for_each_th = (len(data_file_lists)-1)//cpu_cores + 1 + lines_for_each_th = (len(data_file_lists) - 1) // cpu_cores + 1 task_num = cpu_cores if len(data_file_lists) > cpu_cores else 1 # import pdb;pdb.set_trace() if task_num > 1: with concurrent.futures.ThreadPoolExecutor(max_workers=cpu_cores) as executor: - - futures = [executor.submit(parse_context_length, data_file_lists[i*lines_for_each_th:(i+1)*lines_for_each_th], data_type) for i in range(task_num)] - + + futures = [ + executor.submit( + parse_context_length, + data_file_lists[ + i * lines_for_each_th : (i + 1) * lines_for_each_th + ], + data_type, + ) + for i in range(task_num) + ] + for future in concurrent.futures.as_completed(futures): - + json_dict[data_type].update(future.result()) else: res = parse_context_length(data_file_lists, data_type) @@ -48,19 +58,19 @@ def gen_jsonl_from_wav_text_list(path, data_type_list=("source", "target"), json for data_file in data_type_list: jsonl_line.update(json_dict[data_file][key]) jsonl_line = json.dumps(jsonl_line, ensure_ascii=False) - f.write(jsonl_line+"\n") + f.write(jsonl_line + "\n") f.flush() print(f"processed {len(json_dict[data_type_list[0]])} samples") - + else: pass - + if world_size > 1: dist.barrier() - - + + def parse_context_length(data_list: list, data_type: str): - + res = {} for i, line in enumerate(data_list): key, line = line.strip().split(maxsplit=1) @@ -68,26 +78,33 @@ def parse_context_length(data_list: list, data_type: str): if os.path.exists(line): waveform, _ = librosa.load(line, sr=16000) sample_num = len(waveform) - context_len = int(sample_num/16000*1000/10) + context_len = int(sample_num / 16000 * 1000 / 10) else: context_len = len(line.split()) if " " in line else len(line) res[key] = {data_type: line, f"{data_type}_len": context_len} return res - + @hydra.main(config_name=None, version_base=None) def main_hydra(cfg: DictConfig): - + kwargs = OmegaConf.to_container(cfg, resolve=True) print(kwargs) - scp_file_list = kwargs.get("scp_file_list", ("/Users/zhifu/funasr1.0/test_local/wav.scp", "/Users/zhifu/funasr1.0/test_local/text.txt")) + scp_file_list = kwargs.get( + "scp_file_list", + ("/Users/zhifu/funasr1.0/test_local/wav.scp", "/Users/zhifu/funasr1.0/test_local/text.txt"), + ) if isinstance(scp_file_list, str): scp_file_list = eval(scp_file_list) data_type_list = kwargs.get("data_type_list", ("source", "target")) - jsonl_file_out = kwargs.get("jsonl_file_out", "/Users/zhifu/funasr1.0/test_local/audio_datasets.jsonl") - gen_jsonl_from_wav_text_list(scp_file_list, data_type_list=data_type_list, jsonl_file_out=jsonl_file_out) - + jsonl_file_out = kwargs.get( + "jsonl_file_out", "/Users/zhifu/funasr1.0/test_local/audio_datasets.jsonl" + ) + gen_jsonl_from_wav_text_list( + scp_file_list, data_type_list=data_type_list, jsonl_file_out=jsonl_file_out + ) + """ python -m funasr.datasets.audio_datasets.scp2jsonl \ @@ -98,5 +115,3 @@ python -m funasr.datasets.audio_datasets.scp2jsonl \ if __name__ == "__main__": main_hydra() - - \ No newline at end of file diff --git a/funasr/datasets/dataloader_entry.py b/funasr/datasets/dataloader_entry.py index 70da722aa..95958057e 100644 --- a/funasr/datasets/dataloader_entry.py +++ b/funasr/datasets/dataloader_entry.py @@ -1,83 +1,134 @@ - import logging import torch from funasr.register import tables + # @tables.register("dataloader_classes", "DataloaderMapStyle") def DataloaderMapStyle(frontend=None, tokenizer=None, **kwargs): - # dataset - logging.info("Build dataloader") - dataset_class = tables.dataset_classes.get(kwargs.get("dataset", "AudioDataset")) - dataset_tr = dataset_class(kwargs.get("train_data_set_list"), frontend=frontend, tokenizer=tokenizer, is_training=True, **kwargs.get("dataset_conf")) - dataset_val = dataset_class(kwargs.get("valid_data_set_list"), frontend=frontend, tokenizer=tokenizer, is_training=False, **kwargs.get("dataset_conf")) - - # dataloader - batch_sampler = kwargs["dataset_conf"].get("batch_sampler", "BatchSampler") - batch_sampler_val = None - if batch_sampler is not None: - batch_sampler_class = tables.batch_sampler_classes.get(batch_sampler) - batch_sampler = batch_sampler_class(dataset_tr, **kwargs.get("dataset_conf")) - batch_sampler_val = batch_sampler_class(dataset_val, is_training=False, **kwargs.get("dataset_conf")) - - dataloader_tr = torch.utils.data.DataLoader(dataset_tr, collate_fn=dataset_tr.collator, **batch_sampler) - dataloader_val = torch.utils.data.DataLoader(dataset_val, collate_fn=dataset_val.collator, **batch_sampler_val) - - return dataloader_tr, dataloader_val + # dataset + logging.info("Build dataloader") + dataset_class = tables.dataset_classes.get(kwargs.get("dataset", "AudioDataset")) + dataset_tr = dataset_class( + kwargs.get("train_data_set_list"), + frontend=frontend, + tokenizer=tokenizer, + is_training=True, + **kwargs.get("dataset_conf") + ) + dataset_val = dataset_class( + kwargs.get("valid_data_set_list"), + frontend=frontend, + tokenizer=tokenizer, + is_training=False, + **kwargs.get("dataset_conf") + ) + + # dataloader + batch_sampler = kwargs["dataset_conf"].get("batch_sampler", "BatchSampler") + batch_sampler_val = None + if batch_sampler is not None: + batch_sampler_class = tables.batch_sampler_classes.get(batch_sampler) + batch_sampler = batch_sampler_class(dataset_tr, **kwargs.get("dataset_conf")) + batch_sampler_val = batch_sampler_class( + dataset_val, is_training=False, **kwargs.get("dataset_conf") + ) + + dataloader_tr = torch.utils.data.DataLoader( + dataset_tr, collate_fn=dataset_tr.collator, **batch_sampler + ) + dataloader_val = torch.utils.data.DataLoader( + dataset_val, collate_fn=dataset_val.collator, **batch_sampler_val + ) + + return dataloader_tr, dataloader_val + @tables.register("dataloader_classes", "DataloaderMapStyle") class DataloaderMapStyle: - def __init__(self, frontend=None, tokenizer=None, **kwargs): - # dataset - logging.info("Build dataloader") - dataset_class = tables.dataset_classes.get(kwargs.get("dataset", "AudioDataset")) - dataset_tr = dataset_class(kwargs.get("train_data_set_list"), frontend=frontend, tokenizer=tokenizer, - is_training=True, **kwargs.get("dataset_conf")) - dataset_val = dataset_class(kwargs.get("valid_data_set_list"), frontend=frontend, tokenizer=tokenizer, - is_training=False, **kwargs.get("dataset_conf")) + def __init__(self, frontend=None, tokenizer=None, **kwargs): + # dataset + logging.info("Build dataloader") + dataset_class = tables.dataset_classes.get(kwargs.get("dataset", "AudioDataset")) + dataset_tr = dataset_class( + kwargs.get("train_data_set_list"), + frontend=frontend, + tokenizer=tokenizer, + is_training=True, + **kwargs.get("dataset_conf") + ) + dataset_val = dataset_class( + kwargs.get("valid_data_set_list"), + frontend=frontend, + tokenizer=tokenizer, + is_training=False, + **kwargs.get("dataset_conf") + ) + + self.dataset_tr = dataset_tr + self.dataset_val = dataset_val + self.kwargs = kwargs + + # split dataset + self.data_split_num = kwargs["dataset_conf"].get("data_split_num", 1) + self.dataset_class = dataset_class + self.frontend = frontend + self.tokenizer = tokenizer + self.kwargs = kwargs + + def build_iter(self, epoch=0, data_split_i=0, **kwargs): + + # reload dataset slice + if self.data_split_num > 1: + del self.dataset_tr + self.dataset_tr = self.dataset_class( + self.kwargs.get("train_data_set_list"), + frontend=self.frontend, + tokenizer=self.tokenizer, + is_training=True, + **self.kwargs.get("dataset_conf"), + data_split_i=data_split_i + ) + + # dataloader + batch_sampler = self.kwargs["dataset_conf"].get("batch_sampler", "BatchSampler") + batch_sampler_val = None + if batch_sampler is not None: + batch_sampler_class = tables.batch_sampler_classes.get(batch_sampler) + batch_sampler = batch_sampler_class(self.dataset_tr, **self.kwargs.get("dataset_conf")) + batch_sampler_val = batch_sampler_class( + self.dataset_val, is_training=False, **self.kwargs.get("dataset_conf") + ) + + batch_sampler["batch_sampler"].set_epoch(epoch) + batch_sampler_val["batch_sampler"].set_epoch(epoch) + dataloader_tr = torch.utils.data.DataLoader( + self.dataset_tr, collate_fn=self.dataset_tr.collator, **batch_sampler + ) + dataloader_val = torch.utils.data.DataLoader( + self.dataset_val, collate_fn=self.dataset_val.collator, **batch_sampler_val + ) + + return dataloader_tr, dataloader_val - self.dataset_tr = dataset_tr - self.dataset_val = dataset_val - self.kwargs = kwargs - - # split dataset - self.data_split_num = kwargs["dataset_conf"].get("data_split_num", 1) - self.dataset_class = dataset_class - self.frontend = frontend - self.tokenizer = tokenizer - self.kwargs = kwargs - - def build_iter(self, epoch=0, data_split_i=0, **kwargs): - - # reload dataset slice - if self.data_split_num > 1: - del self.dataset_tr - self.dataset_tr = self.dataset_class(self.kwargs.get("train_data_set_list"), frontend=self.frontend, tokenizer=self.tokenizer, - is_training=True, **self.kwargs.get("dataset_conf"), data_split_i=data_split_i) - - # dataloader - batch_sampler = self.kwargs["dataset_conf"].get("batch_sampler", "BatchSampler") - batch_sampler_val = None - if batch_sampler is not None: - batch_sampler_class = tables.batch_sampler_classes.get(batch_sampler) - batch_sampler = batch_sampler_class(self.dataset_tr, **self.kwargs.get("dataset_conf")) - batch_sampler_val = batch_sampler_class(self.dataset_val, is_training=False, **self.kwargs.get("dataset_conf")) - - batch_sampler["batch_sampler"].set_epoch(epoch) - batch_sampler_val["batch_sampler"].set_epoch(epoch) - dataloader_tr = torch.utils.data.DataLoader(self.dataset_tr, collate_fn=self.dataset_tr.collator, **batch_sampler) - dataloader_val = torch.utils.data.DataLoader(self.dataset_val, collate_fn=self.dataset_val.collator, **batch_sampler_val) - - return dataloader_tr, dataloader_val - @tables.register("dataloader_classes", "DataloaderIterable") def DataloaderIterable(frontend=None, tokenizer=None, **kwargs): - logging.info("Build dataloader") - dataset_class = tables.dataset_classes.get(kwargs.get("dataset", "LargeDataset")) - dataset_tr = dataset_class(kwargs.get("train_data_set_list"), frontend=frontend, tokenizer=tokenizer, - is_training=True, **kwargs.get("dataset_conf")) - dataset_val = dataset_class(kwargs.get("valid_data_set_list"), frontend=frontend, tokenizer=tokenizer, - is_training=False, **kwargs.get("dataset_conf")) - - return dataset_tr, dataset_val \ No newline at end of file + logging.info("Build dataloader") + dataset_class = tables.dataset_classes.get(kwargs.get("dataset", "LargeDataset")) + dataset_tr = dataset_class( + kwargs.get("train_data_set_list"), + frontend=frontend, + tokenizer=tokenizer, + is_training=True, + **kwargs.get("dataset_conf") + ) + dataset_val = dataset_class( + kwargs.get("valid_data_set_list"), + frontend=frontend, + tokenizer=tokenizer, + is_training=False, + **kwargs.get("dataset_conf") + ) + + return dataset_tr, dataset_val diff --git a/funasr/datasets/large_datasets/build_dataloader.py b/funasr/datasets/large_datasets/build_dataloader.py index 8a255f99e..da0471727 100644 --- a/funasr/datasets/large_datasets/build_dataloader.py +++ b/funasr/datasets/large_datasets/build_dataloader.py @@ -13,6 +13,7 @@ from funasr.tokenizer.abs_tokenizer import AbsTokenizer from funasr.register import tables + def read_symbol_table(symbol_table_file): if isinstance(symbol_table_file, str): symbol_table = {} @@ -62,6 +63,7 @@ class SentencepiecesTokenizer(AbsTokenizer): self._build_sentence_piece_processor() return self.sp.DecodePieces(list(tokens)) + @tables.register("dataset_classes", "LargeDataset") class LargeDataLoader(AbsIterFactory): def __init__(self, args, mode="train"): @@ -76,22 +78,32 @@ class LargeDataLoader(AbsIterFactory): bpe_tokenizer = SentencepiecesTokenizer(args.bpemodel) self.dataset_conf = args.dataset_conf if "frontend_conf" not in args: - self.frontend_conf = None + self.frontend_conf = None else: self.frontend_conf = args.frontend_conf - self.speed_perturb = args.speed_perturb if hasattr(args, "speed_perturb") else None + self.speed_perturb = args.speed_perturb if hasattr(args, "speed_perturb") else None logging.info("dataloader config: {}".format(self.dataset_conf)) batch_mode = self.dataset_conf.get("batch_mode", "padding") data_list = args.train_data_file if mode == "train" else args.valid_data_file - self.dataset = Dataset(data_list, symbol_table, seg_dict, punc_dict, bpe_tokenizer, - self.dataset_conf, self.frontend_conf, - speed_perturb=self.speed_perturb if mode == "train" else None, - mode=mode, batch_mode=batch_mode) + self.dataset = Dataset( + data_list, + symbol_table, + seg_dict, + punc_dict, + bpe_tokenizer, + self.dataset_conf, + self.frontend_conf, + speed_perturb=self.speed_perturb if mode == "train" else None, + mode=mode, + batch_mode=batch_mode, + ) def build_iter(self, epoch, shuffle=True): self.dataset.set_epoch(epoch) - data_loader = DataLoader(self.dataset, - batch_size=None, - pin_memory=True, - num_workers=self.dataset_conf.get("num_workers", 8)) + data_loader = DataLoader( + self.dataset, + batch_size=None, + pin_memory=True, + num_workers=self.dataset_conf.get("num_workers", 8), + ) return data_loader diff --git a/funasr/datasets/large_datasets/collate_fn.py b/funasr/datasets/large_datasets/collate_fn.py index aff25a8e2..4648d8735 100644 --- a/funasr/datasets/large_datasets/collate_fn.py +++ b/funasr/datasets/large_datasets/collate_fn.py @@ -13,11 +13,11 @@ class CommonCollateFn: """Functor class of common_collate_fn()""" def __init__( - self, - float_pad_value: Union[float, int] = 0.0, - int_pad_value: int = -32768, - not_sequence: Collection[str] = (), - max_sample_size=None + self, + float_pad_value: Union[float, int] = 0.0, + int_pad_value: int = -32768, + not_sequence: Collection[str] = (), + max_sample_size=None, ): self.float_pad_value = float_pad_value self.int_pad_value = int_pad_value @@ -31,7 +31,7 @@ class CommonCollateFn: ) def __call__( - self, data: Collection[Tuple[str, Dict[str, np.ndarray]]] + self, data: Collection[Tuple[str, Dict[str, np.ndarray]]] ) -> Tuple[List[str], Dict[str, torch.Tensor]]: return common_collate_fn( data, @@ -42,13 +42,12 @@ class CommonCollateFn: def common_collate_fn( - data: Collection[Tuple[str, Dict[str, np.ndarray]]], - float_pad_value: Union[float, int] = 0.0, - int_pad_value: int = -32768, - not_sequence: Collection[str] = (), + data: Collection[Tuple[str, Dict[str, np.ndarray]]], + float_pad_value: Union[float, int] = 0.0, + int_pad_value: int = -32768, + not_sequence: Collection[str] = (), ) -> Tuple[List[str], Dict[str, torch.Tensor]]: - """Concatenate ndarray-list to an array and convert to torch.Tensor. - """ + """Concatenate ndarray-list to an array and convert to torch.Tensor.""" uttids = [u for u, _ in data] data = [d for _, d in data] @@ -81,11 +80,11 @@ class DiarCollateFn: """Functor class of common_collate_fn()""" def __init__( - self, - float_pad_value: Union[float, int] = 0.0, - int_pad_value: int = -32768, - not_sequence: Collection[str] = (), - max_sample_size=None + self, + float_pad_value: Union[float, int] = 0.0, + int_pad_value: int = -32768, + not_sequence: Collection[str] = (), + max_sample_size=None, ): self.float_pad_value = float_pad_value self.int_pad_value = int_pad_value @@ -99,7 +98,7 @@ class DiarCollateFn: ) def __call__( - self, data: Collection[Tuple[str, Dict[str, np.ndarray]]] + self, data: Collection[Tuple[str, Dict[str, np.ndarray]]] ) -> Tuple[List[str], Dict[str, torch.Tensor]]: return diar_collate_fn( data, @@ -110,13 +109,12 @@ class DiarCollateFn: def diar_collate_fn( - data: Collection[Tuple[str, Dict[str, np.ndarray]]], - float_pad_value: Union[float, int] = 0.0, - int_pad_value: int = -32768, - not_sequence: Collection[str] = (), + data: Collection[Tuple[str, Dict[str, np.ndarray]]], + float_pad_value: Union[float, int] = 0.0, + int_pad_value: int = -32768, + not_sequence: Collection[str] = (), ) -> Tuple[List[str], Dict[str, torch.Tensor]]: - """Concatenate ndarray-list to an array and convert to torch.Tensor. - """ + """Concatenate ndarray-list to an array and convert to torch.Tensor.""" uttids = [u for u, _ in data] data = [d for _, d in data] @@ -157,9 +155,9 @@ def crop_to_max_size(feature, target_size): def clipping_collate_fn( - data: Collection[Tuple[str, Dict[str, np.ndarray]]], - max_sample_size=None, - not_sequence: Collection[str] = (), + data: Collection[Tuple[str, Dict[str, np.ndarray]]], + max_sample_size=None, + not_sequence: Collection[str] = (), ) -> Tuple[List[str], Dict[str, torch.Tensor]]: # mainly for pre-training uttids = [u for u, _ in data] @@ -193,4 +191,4 @@ def clipping_collate_fn( output[key + "_lengths"] = lens output = (uttids, output) - return output \ No newline at end of file + return output diff --git a/funasr/datasets/large_datasets/datapipes/batch.py b/funasr/datasets/large_datasets/datapipes/batch.py index 35e5dbaeb..aeeb45146 100644 --- a/funasr/datasets/large_datasets/datapipes/batch.py +++ b/funasr/datasets/large_datasets/datapipes/batch.py @@ -19,13 +19,13 @@ def _token_len_fn(token, len_fn): class MaxTokenBucketizerIterDataPipe(IterableDataset): def __init__( - self, - datapipe, - batch_size=8000, - len_fn=_default_len_fn, - buffer_size=10240, - sort_size=500, - batch_mode="padding", + self, + datapipe, + batch_size=8000, + len_fn=_default_len_fn, + buffer_size=10240, + sort_size=500, + batch_mode="padding", ): assert batch_size > 0, "Batch size is required to be larger than 0!" assert buffer_size >= -1, "Buffer size is required to be larger than -1!" diff --git a/funasr/datasets/large_datasets/datapipes/filter.py b/funasr/datasets/large_datasets/datapipes/filter.py index 6fe71538d..c4f045dc1 100644 --- a/funasr/datasets/large_datasets/datapipes/filter.py +++ b/funasr/datasets/large_datasets/datapipes/filter.py @@ -1,14 +1,13 @@ from torch.utils.data import IterableDataset + def default_fn(data): return data class FilterIterDataPipe(IterableDataset): - def __init__(self, - datapipe, - fn=default_fn): + def __init__(self, datapipe, fn=default_fn): self.datapipe = datapipe self.fn = fn diff --git a/funasr/datasets/large_datasets/datapipes/map.py b/funasr/datasets/large_datasets/datapipes/map.py index dfcd6a045..f7211f902 100644 --- a/funasr/datasets/large_datasets/datapipes/map.py +++ b/funasr/datasets/large_datasets/datapipes/map.py @@ -7,9 +7,7 @@ def default_fn(data): class MapperIterDataPipe(IterableDataset): - def __init__(self, - datapipe, - fn=default_fn): + def __init__(self, datapipe, fn=default_fn): self.datapipe = datapipe self.fn = fn diff --git a/funasr/datasets/large_datasets/dataset.py b/funasr/datasets/large_datasets/dataset.py index d3489c199..1e5b6c12b 100644 --- a/funasr/datasets/large_datasets/dataset.py +++ b/funasr/datasets/large_datasets/dataset.py @@ -7,6 +7,7 @@ import torch import torch.distributed as dist import torchaudio import numpy as np + # import librosa import librosa from kaldiio import ReadHelper @@ -23,7 +24,7 @@ from funasr.datasets.large_datasets.utils.tokenize import tokenize def read_lists(list_file): lists = [] - with open(list_file, 'r', encoding='utf8') as fin: + with open(list_file, "r", encoding="utf8") as fin: for line in fin: parts = line.strip() lists.append(parts) @@ -31,8 +32,16 @@ def read_lists(list_file): class AudioDataset(IterableDataset): - def __init__(self, scp_lists, data_names, data_types, frontend_conf=None, shuffle=True, speed_perturb=None, - mode="train"): + def __init__( + self, + scp_lists, + data_names, + data_types, + frontend_conf=None, + shuffle=True, + speed_perturb=None, + mode="train", + ): self.scp_lists = scp_lists self.data_names = data_names self.data_types = data_types @@ -64,7 +73,7 @@ class AudioDataset(IterableDataset): if self.shuffle: random.seed(self.epoch) random.shuffle(data_index) - return data_index[self.rank::self.world_size] + return data_index[self.rank :: self.world_size] return data_index @@ -77,7 +86,7 @@ class AudioDataset(IterableDataset): self.worker_id = worker_info.id self.num_workers = worker_info.num_workers - return rank_data_index[self.worker_id::self.num_workers] + return rank_data_index[self.worker_id :: self.num_workers] def close_reader(self, reader_list): for reader in reader_list: @@ -91,8 +100,8 @@ class AudioDataset(IterableDataset): for index in worker_data_index: data = dict(scp=self.scp_lists[index]) - assert 'scp' in data - scp = data['scp'] + assert "scp" in data + scp = data["scp"] data_file_list = scp.strip().split() data_name_list = self.data_names.split(",") data_type_list = self.data_types.split(",") @@ -100,15 +109,16 @@ class AudioDataset(IterableDataset): for file in data_file_list: assert os.path.exists(file), "{} not exists".format(file) - assert len(data_file_list) == len(data_name_list) == len(data_type_list), \ - "The item number of data, data_names, data_types must be the same " + assert ( + len(data_file_list) == len(data_name_list) == len(data_type_list) + ), "The item number of data, data_names, data_types must be the same " reader_list = [] for data_file, data_type in zip(data_file_list, data_type_list): if data_type == "kaldi_ark": - ark_reader = ReadHelper('ark:{}'.format(data_file)) + ark_reader = ReadHelper("ark:{}".format(data_file)) reader_list.append(ark_reader) - elif data_type == "text" or data_type == "sound" or data_type == 'text_hotword': + elif data_type == "text" or data_type == "sound" or data_type == "text_hotword": text_reader = open(data_file, "r", encoding="utf-8") reader_list.append(text_reader) elif data_type == "none": @@ -130,15 +140,16 @@ class AudioDataset(IterableDataset): waveform, sampling_rate = torchaudio.load(path) except: # waveform, sampling_rate = librosa.load(path, dtype='float32') - waveform, sampling_rate = librosa.load(path, dtype='float32') + waveform, sampling_rate = librosa.load(path, dtype="float32") if waveform.ndim == 2: waveform = waveform[:, 0] waveform = np.expand_dims(waveform, axis=0) waveform = torch.tensor(waveform) if self.frontend_conf is not None: if sampling_rate != self.frontend_conf["fs"]: - waveform = torchaudio.transforms.Resample(orig_freq=sampling_rate, - new_freq=self.frontend_conf["fs"])(waveform) + waveform = torchaudio.transforms.Resample( + orig_freq=sampling_rate, new_freq=self.frontend_conf["fs"] + )(waveform) sampling_rate = self.frontend_conf["fs"] waveform = waveform.numpy() mat = waveform[0] @@ -146,7 +157,10 @@ class AudioDataset(IterableDataset): speed = random.choice(self.speed_perturb) if speed != 1.0: mat, _ = torchaudio.sox_effects.apply_effects_tensor( - torch.tensor(mat).view(1, -1), sampling_rate, [['speed', str(speed)], ['rate', str(sampling_rate)]]) + torch.tensor(mat).view(1, -1), + sampling_rate, + [["speed", str(speed)], ["rate", str(sampling_rate)]], + ) mat = mat.view(-1).numpy() sample_dict[data_name] = mat sample_dict["sampling_rate"] = sampling_rate @@ -158,7 +172,7 @@ class AudioDataset(IterableDataset): sample_dict[data_name] = segs[1:] if "key" not in sample_dict: sample_dict["key"] = segs[0] - sample_dict['hw_tag'] = 1 + sample_dict["hw_tag"] = 1 elif data_type == "text_nospace": text = item segs = text.strip().split(maxsplit=1) @@ -183,23 +197,25 @@ def len_fn_example(data): def len_fn_token(data): assert "speech" in data if "sampling_rate" in data: - return (data["speech"].shape[0] / data["sampling_rate"]) * 1000. + return (data["speech"].shape[0] / data["sampling_rate"]) * 1000.0 else: return data["speech"].shape[0] -def Dataset(data_list_file, - dict, - seg_dict, - punc_dict, - bpe_tokenizer, - conf, - frontend_conf, - speed_perturb=None, - mode="train", - batch_mode="padding"): +def Dataset( + data_list_file, + dict, + seg_dict, + punc_dict, + bpe_tokenizer, + conf, + frontend_conf, + speed_perturb=None, + mode="train", + batch_mode="padding", +): scp_lists = read_lists(data_list_file) - shuffle = conf.get('shuffle', True) + shuffle = conf.get("shuffle", True) data_names = conf.get("data_names", "speech,text") data_types = conf.get("data_types", "kaldi_ark,text") @@ -207,63 +223,72 @@ def Dataset(data_list_file, # pre_prob = conf.get("pre_prob", 0) # unused yet if pre_hwfile is not None: pre_hwlist = [] - with open(pre_hwfile, 'r', encoding="utf-8") as fin: + with open(pre_hwfile, "r", encoding="utf-8") as fin: for line in fin.readlines(): pre_hwlist.append(line.strip()) else: pre_hwlist = None - hw_config = {"sample_rate": conf.get("sample_rate", 0.6), - "double_rate": conf.get("double_rate", 0.1), - "hotword_min_length": conf.get("hotword_min_length", 2), - "hotword_max_length": conf.get("hotword_max_length", 8), - "pre_prob": conf.get("pre_prob", 0.0), - "pre_hwlist": pre_hwlist} + hw_config = { + "sample_rate": conf.get("sample_rate", 0.6), + "double_rate": conf.get("double_rate", 0.1), + "hotword_min_length": conf.get("hotword_min_length", 2), + "hotword_max_length": conf.get("hotword_max_length", 8), + "pre_prob": conf.get("pre_prob", 0.0), + "pre_hwlist": pre_hwlist, + } - - - dataset = AudioDataset(scp_lists, - data_names, - data_types, - frontend_conf=frontend_conf, - shuffle=shuffle, - speed_perturb=speed_perturb, - mode=mode, - ) + dataset = AudioDataset( + scp_lists, + data_names, + data_types, + frontend_conf=frontend_conf, + shuffle=shuffle, + speed_perturb=speed_perturb, + mode=mode, + ) if "text" in data_names: - vocab = {'vocab': dict, 'seg_dict': seg_dict, 'punc_dict': punc_dict, 'bpe_tokenizer': bpe_tokenizer, 'hw_config': hw_config} + vocab = { + "vocab": dict, + "seg_dict": seg_dict, + "punc_dict": punc_dict, + "bpe_tokenizer": bpe_tokenizer, + "hw_config": hw_config, + } tokenize_fn = partial(tokenize, **vocab) dataset = MapperIterDataPipe(dataset, fn=tokenize_fn) - filter_conf = conf.get('filter_conf', {}) + filter_conf = conf.get("filter_conf", {}) filter_fn = partial(filter, **filter_conf) dataset = FilterIterDataPipe(dataset, fn=filter_fn) if shuffle: - buffer_conf = conf.get('shuffle_conf', {}) - buffer_size = buffer_conf['shuffle_size'] - sort_size = buffer_conf['sort_size'] + buffer_conf = conf.get("shuffle_conf", {}) + buffer_size = buffer_conf["shuffle_size"] + sort_size = buffer_conf["sort_size"] else: buffer_size = 0 sort_size = 1 - batch_conf = conf.get('batch_conf', {}) - batch_size = batch_conf['batch_size'] - batch_type = batch_conf['batch_type'] + batch_conf = conf.get("batch_conf", {}) + batch_size = batch_conf["batch_size"] + batch_type = batch_conf["batch_type"] assert batch_type in ["example", "token"] - if batch_type == 'example': + if batch_type == "example": len_fn = len_fn_example else: len_fn = len_fn_token - dataset = MaxTokenBucketizerIterDataPipe(dataset, - batch_size=batch_size, - len_fn=len_fn, - buffer_size=buffer_size, - sort_size=sort_size, - batch_mode=batch_mode) + dataset = MaxTokenBucketizerIterDataPipe( + dataset, + batch_size=batch_size, + len_fn=len_fn, + buffer_size=buffer_size, + sort_size=sort_size, + batch_mode=batch_mode, + ) int_pad_value = conf.get("int_pad_value", -1) float_pad_value = conf.get("float_pad_value", 0.0) diff --git a/funasr/datasets/large_datasets/utils/clipping.py b/funasr/datasets/large_datasets/utils/clipping.py index 2554abac8..92f7d70ef 100644 --- a/funasr/datasets/large_datasets/utils/clipping.py +++ b/funasr/datasets/large_datasets/utils/clipping.py @@ -25,7 +25,9 @@ def clipping(data): tensor_lengths = torch.tensor([len(d[data_name]) for d in data], dtype=torch.int32) length_clip = min(tensor_lengths) - tensor_clip = tensor_list[0].new_zeros(len(tensor_list), length_clip, tensor_list[0].shape[1]) + tensor_clip = tensor_list[0].new_zeros( + len(tensor_list), length_clip, tensor_list[0].shape[1] + ) for i, (tensor, length) in enumerate(zip(tensor_list, tensor_lengths)): diff = length - length_clip assert diff >= 0 @@ -35,6 +37,8 @@ def clipping(data): tensor_clip[i] = crop_to_max_size(tensor, length_clip) batch[data_name] = tensor_clip - batch[data_name + "_lengths"] = torch.tensor([tensor.shape[0] for tensor in tensor_clip], dtype=torch.long) + batch[data_name + "_lengths"] = torch.tensor( + [tensor.shape[0] for tensor in tensor_clip], dtype=torch.long + ) return keys, batch diff --git a/funasr/datasets/large_datasets/utils/filter.py b/funasr/datasets/large_datasets/utils/filter.py index 1260a47c4..adc8fa038 100644 --- a/funasr/datasets/large_datasets/utils/filter.py +++ b/funasr/datasets/large_datasets/utils/filter.py @@ -1,26 +1,27 @@ #!/usr/bin/env python -def filter(data, - speech_length_min=100, - speech_length_max=15000, - token_length_min=0, - token_length_max=200): +def filter( + data, speech_length_min=100, speech_length_max=15000, token_length_min=0, token_length_max=200 +): assert "speech" in data or "text" in data if "speech" in data and "text" in data: if "sampling_rate" in data: - speech_length = (data["speech"].shape[0] / data["sampling_rate"]) * 1000. + speech_length = (data["speech"].shape[0] / data["sampling_rate"]) * 1000.0 else: speech_length = data["speech"].shape[0] - num_tokens = len(data['text']) - return speech_length_min < speech_length < speech_length_max and token_length_min < num_tokens < token_length_max + num_tokens = len(data["text"]) + return ( + speech_length_min < speech_length < speech_length_max + and token_length_min < num_tokens < token_length_max + ) elif "speech" in data: if "sampling_rate" in data: - speech_length = (data["speech"].shape[0] / data["sampling_rate"]) * 1000. + speech_length = (data["speech"].shape[0] / data["sampling_rate"]) * 1000.0 else: speech_length = data["speech"].shape[0] return speech_length_min < speech_length < speech_length_max else: - num_tokens = len(data['text']) + num_tokens = len(data["text"]) return token_length_min < num_tokens < token_length_max diff --git a/funasr/datasets/large_datasets/utils/hotword_utils.py b/funasr/datasets/large_datasets/utils/hotword_utils.py index 73f8bdd90..66c131ed3 100644 --- a/funasr/datasets/large_datasets/utils/hotword_utils.py +++ b/funasr/datasets/large_datasets/utils/hotword_utils.py @@ -1,33 +1,42 @@ import random -def sample_hotword(length, - hotword_min_length, - hotword_max_length, - sample_rate, - double_rate, - pre_prob, - pre_index=None, - pre_hwlist=None): - if length < hotword_min_length: - return [-1] - if random.random() < sample_rate: - if pre_prob > 0 and random.random() < pre_prob and pre_index is not None: - return pre_index - if length == hotword_min_length: - return [0, length-1] - elif random.random() < double_rate and length > hotword_max_length + hotword_min_length + 2: - # sample two hotwords in a sentence - _max_hw_length = min(hotword_max_length, length // 2) - # first hotword - start1 = random.randint(0, length // 3) - end1 = random.randint(start1 + hotword_min_length - 1, start1 + _max_hw_length - 1) - # second hotword - start2 = random.randint(end1 + 1, length - hotword_min_length) - end2 = random.randint(min(length-1, start2+hotword_min_length-1), min(length-1, start2+hotword_max_length-1)) - return [start1, end1, start2, end2] - else: # single hotword - start = random.randint(0, length - hotword_min_length) - end = random.randint(min(length-1, start+hotword_min_length-1), min(length-1, start+hotword_max_length-1)) - return [start, end] - else: - return [-1] \ No newline at end of file + +def sample_hotword( + length, + hotword_min_length, + hotword_max_length, + sample_rate, + double_rate, + pre_prob, + pre_index=None, + pre_hwlist=None, +): + if length < hotword_min_length: + return [-1] + if random.random() < sample_rate: + if pre_prob > 0 and random.random() < pre_prob and pre_index is not None: + return pre_index + if length == hotword_min_length: + return [0, length - 1] + elif random.random() < double_rate and length > hotword_max_length + hotword_min_length + 2: + # sample two hotwords in a sentence + _max_hw_length = min(hotword_max_length, length // 2) + # first hotword + start1 = random.randint(0, length // 3) + end1 = random.randint(start1 + hotword_min_length - 1, start1 + _max_hw_length - 1) + # second hotword + start2 = random.randint(end1 + 1, length - hotword_min_length) + end2 = random.randint( + min(length - 1, start2 + hotword_min_length - 1), + min(length - 1, start2 + hotword_max_length - 1), + ) + return [start1, end1, start2, end2] + else: # single hotword + start = random.randint(0, length - hotword_min_length) + end = random.randint( + min(length - 1, start + hotword_min_length - 1), + min(length - 1, start + hotword_max_length - 1), + ) + return [start, end] + else: + return [-1] diff --git a/funasr/datasets/large_datasets/utils/low_frame_rate.py b/funasr/datasets/large_datasets/utils/low_frame_rate.py index 76eb2da93..87718e938 100644 --- a/funasr/datasets/large_datasets/utils/low_frame_rate.py +++ b/funasr/datasets/large_datasets/utils/low_frame_rate.py @@ -20,10 +20,10 @@ def build_LFR_features(data, m, n): T_lfr = int(np.ceil(T / n)) for i in range(T_lfr): if m <= T - i * n: - LFR_inputs.append(np.hstack(data[i*n:i*n+m])) + LFR_inputs.append(np.hstack(data[i * n : i * n + m])) else: num_padding = m - (T - i * n) - frame = np.hstack(data[i*n:]) + frame = np.hstack(data[i * n :]) for _ in range(num_padding): frame = np.hstack((frame, data[-1])) LFR_inputs.append(frame) diff --git a/funasr/datasets/large_datasets/utils/padding.py b/funasr/datasets/large_datasets/utils/padding.py index 26c6e84ba..cb43a2726 100644 --- a/funasr/datasets/large_datasets/utils/padding.py +++ b/funasr/datasets/large_datasets/utils/padding.py @@ -7,7 +7,7 @@ def padding(data, float_pad_value=0.0, int_pad_value=-1): assert isinstance(data, list) assert "key" in data[0] assert "speech" in data[0] or "text" in data[0] - + keys = [x["key"] for x in data] batch = {} @@ -16,7 +16,7 @@ def padding(data, float_pad_value=0.0, int_pad_value=-1): if data_name == "key" or data_name == "sampling_rate": continue else: - if data_name != 'hotword_indxs': + if data_name != "hotword_indxs": if data[0][data_name].dtype.kind == "i": pad_value = int_pad_value tensor_type = torch.int64 @@ -26,9 +26,7 @@ def padding(data, float_pad_value=0.0, int_pad_value=-1): tensor_list = [torch.tensor(np.copy(d[data_name]), dtype=tensor_type) for d in data] tensor_lengths = torch.tensor([len(d[data_name]) for d in data], dtype=torch.int32) - tensor_pad = pad_sequence(tensor_list, - batch_first=True, - padding_value=pad_value) + tensor_pad = pad_sequence(tensor_list, batch_first=True, padding_value=pad_value) batch[data_name] = tensor_pad batch[data_name + "_lengths"] = tensor_lengths @@ -38,37 +36,37 @@ def padding(data, float_pad_value=0.0, int_pad_value=-1): # use it to slice hotwords out hotword_list = [] hotword_lengths = [] - text = batch['text'] - text_lengths = batch['text_lengths'] - hotword_indxs = batch['hotword_indxs'] + text = batch["text"] + text_lengths = batch["text_lengths"] + hotword_indxs = batch["hotword_indxs"] dha_pad = torch.ones_like(text) * -1 _, t1 = text.shape t1 += 1 # TODO: as parameter which is same as predictor_bias nth_hw = 0 - for b, (hotword_indx, one_text, length) in enumerate(zip(hotword_indxs, text, text_lengths)): + for b, (hotword_indx, one_text, length) in enumerate( + zip(hotword_indxs, text, text_lengths) + ): dha_pad[b][:length] = 8405 if hotword_indx[0] != -1: start, end = int(hotword_indx[0]), int(hotword_indx[1]) - hotword = one_text[start: end+1] + hotword = one_text[start : end + 1] hotword_list.append(hotword) - hotword_lengths.append(end-start+1) - dha_pad[b][start: end+1] = one_text[start: end+1] + hotword_lengths.append(end - start + 1) + dha_pad[b][start : end + 1] = one_text[start : end + 1] nth_hw += 1 if len(hotword_indx) == 4 and hotword_indx[2] != -1: # the second hotword if exist start, end = int(hotword_indx[2]), int(hotword_indx[3]) - hotword_list.append(one_text[start: end+1]) - hotword_lengths.append(end-start+1) - dha_pad[b][start: end+1] = one_text[start: end+1] + hotword_list.append(one_text[start : end + 1]) + hotword_lengths.append(end - start + 1) + dha_pad[b][start : end + 1] = one_text[start : end + 1] nth_hw += 1 hotword_list.append(torch.tensor([1])) hotword_lengths.append(1) - hotword_pad = pad_sequence(hotword_list, - batch_first=True, - padding_value=0) + hotword_pad = pad_sequence(hotword_list, batch_first=True, padding_value=0) batch["hotword_pad"] = hotword_pad batch["hotword_lengths"] = torch.tensor(hotword_lengths, dtype=torch.int32) - batch['dha_pad'] = dha_pad - del batch['hotword_indxs'] - del batch['hotword_indxs_lengths'] + batch["dha_pad"] = dha_pad + del batch["hotword_indxs"] + del batch["hotword_indxs_lengths"] return keys, batch diff --git a/funasr/datasets/large_datasets/utils/tokenize.py b/funasr/datasets/large_datasets/utils/tokenize.py index 34a97c176..5a1ddd276 100644 --- a/funasr/datasets/large_datasets/utils/tokenize.py +++ b/funasr/datasets/large_datasets/utils/tokenize.py @@ -3,6 +3,7 @@ import re import numpy as np from funasr.datasets.large_datasets.utils.hotword_utils import sample_hotword + def forward_segment(text, seg_dict): word_list = [] i = 0 @@ -17,8 +18,9 @@ def forward_segment(text, seg_dict): i += len(longest_word) return word_list + def seg_tokenize(txt, seg_dict): - pattern = re.compile(r'^[\u4E00-\u9FA50-9]+$') + pattern = re.compile(r"^[\u4E00-\u9FA50-9]+$") out_txt = "" for word in txt: word = word.lower() @@ -35,12 +37,8 @@ def seg_tokenize(txt, seg_dict): out_txt += "" + " " return out_txt.strip().split() -def tokenize(data, - vocab=None, - seg_dict=None, - punc_dict=None, - bpe_tokenizer=None, - hw_config=None): + +def tokenize(data, vocab=None, seg_dict=None, punc_dict=None, bpe_tokenizer=None, hw_config=None): assert "text" in data assert isinstance(vocab, dict) text = data["text"] @@ -53,11 +51,11 @@ def tokenize(data, text = seg_tokenize(text, seg_dict) length = len(text) - if 'hw_tag' in data: + if "hw_tag" in data: pre_index = None - if hw_config['pre_hwlist'] is not None and hw_config['pre_prob'] > 0: + if hw_config["pre_hwlist"] is not None and hw_config["pre_prob"] > 0: # enable preset hotword detect in sampling - for hw in hw_config['pre_hwlist']: + for hw in hw_config["pre_hwlist"]: hw = " ".join(seg_tokenize(hw, seg_dict)) _find = " ".join(text).find(hw) if _find != -1: @@ -65,11 +63,11 @@ def tokenize(data, pre_index = [_find, _find + max(hw.count(" "), 1)] break hotword_indxs = sample_hotword(length, **hw_config, pre_index=pre_index) - data['hotword_indxs'] = hotword_indxs - del data['hw_tag'] + data["hotword_indxs"] = hotword_indxs + del data["hw_tag"] for i in range(length): x = text[i] - if i == length-1 and "punc" in data and x.startswith("vad:"): + if i == length - 1 and "punc" in data and x.startswith("vad:"): vad = x[4:] if len(vad) == 0: vad = -1 @@ -78,7 +76,7 @@ def tokenize(data, elif x in vocab: token.append(vocab[x]) else: - token.append(vocab['']) + token.append(vocab[""]) if "punc" in data and punc_dict is not None: punc_token = [] @@ -87,9 +85,9 @@ def tokenize(data, punc_token.append(punc_dict[punc]) else: punc_token.append(punc_dict["_"]) - data["punc"] = np.array(punc_token) + data["punc"] = np.array(punc_token) data["text"] = np.array(token) if vad is not -2: - data["vad_indexes"]=np.array([vad], dtype=np.int64) + data["vad_indexes"] = np.array([vad], dtype=np.int64) return data diff --git a/funasr/datasets/llm_datasets/datasets.py b/funasr/datasets/llm_datasets/datasets.py index a49122da9..b66055412 100644 --- a/funasr/datasets/llm_datasets/datasets.py +++ b/funasr/datasets/llm_datasets/datasets.py @@ -10,28 +10,33 @@ class AudioLLMNARDataset(torch.utils.data.Dataset): """ AudioLLMDataset """ - def __init__(self, - path, - index_ds: str = None, - frontend=None, - tokenizer=None, - int_pad_value: int = -1, - float_pad_value: float = 0.0, - **kwargs): + + def __init__( + self, + path, + index_ds: str = None, + frontend=None, + tokenizer=None, + int_pad_value: int = -1, + float_pad_value: float = 0.0, + **kwargs + ): super().__init__() index_ds_class = tables.index_ds_classes.get(index_ds) self.index_ds = index_ds_class(path, **kwargs) preprocessor_speech = kwargs.get("preprocessor_speech", None) if preprocessor_speech: preprocessor_speech_class = tables.preprocessor_classes.get(preprocessor_speech) - preprocessor_speech = preprocessor_speech_class(**kwargs.get("preprocessor_speech_conf", {})) + preprocessor_speech = preprocessor_speech_class( + **kwargs.get("preprocessor_speech_conf", {}) + ) self.preprocessor_speech = preprocessor_speech preprocessor_text = kwargs.get("preprocessor_text", None) if preprocessor_text: preprocessor_text_class = tables.preprocessor_classes.get(preprocessor_text) preprocessor_text = preprocessor_text_class(**kwargs.get("preprocessor_text_conf", {})) self.preprocessor_text = preprocessor_text - + self.frontend = frontend self.fs = 16000 if frontend is None else frontend.fs self.data_type = "sound" @@ -39,22 +44,24 @@ class AudioLLMNARDataset(torch.utils.data.Dataset): self.float_pad_value = float_pad_value self.prompt = kwargs.get("prompt", "Please copy the following text.") - self.prompt_pre = "USER: \nINSTRUCTION: {}\nINPUT: ".format(self.prompt) # "USER: \nINSTRUCTION: {}\nINPUT: {}\nASSISTANT: " + self.prompt_pre = "USER: \nINSTRUCTION: {}\nINPUT: ".format( + self.prompt + ) # "USER: \nINSTRUCTION: {}\nINPUT: {}\nASSISTANT: " self.prompt_af = "" self.IGNORE_INDEX = kwargs.get("IGNORE_INDEX", -100) self.int_pad_value = self.IGNORE_INDEX - + def get_source_len(self, index): item = self.index_ds[index] return self.index_ds.get_source_len(item) - + def get_target_len(self, index): item = self.index_ds[index] return self.index_ds.get_target_len(item) - + def __len__(self): return len(self.index_ds) - + def __getitem__(self, index): item = self.index_ds[index] # import pdb; @@ -63,17 +70,18 @@ class AudioLLMNARDataset(torch.utils.data.Dataset): data_src = load_audio_text_image_video(source, fs=self.fs) if self.preprocessor_speech: data_src = self.preprocessor_speech(data_src, fs=self.fs) - speech, speech_lengths = extract_fbank(data_src, data_type=self.data_type, frontend=self.frontend, is_final=True) # speech: [b, T, d] + speech, speech_lengths = extract_fbank( + data_src, data_type=self.data_type, frontend=self.frontend, is_final=True + ) # speech: [b, T, d] speech = speech.squeeze(0) target = item["target"] if self.preprocessor_text: target = self.preprocessor_text(target) - - - prompt_ids_pre = self.tokenizer.encode(self.prompt_pre) # [bos,prompt] + + prompt_ids_pre = self.tokenizer.encode(self.prompt_pre) # [bos,prompt] prompt_ids_length = len(prompt_ids_pre) - + # bos prompt audio bos target # prompt_input = "{}{}".format(self.prompt_pre, target) # prompt_input_ids = self.tokenizer.encode(prompt_input) #[bos, prompt, input] @@ -83,11 +91,17 @@ class AudioLLMNARDataset(torch.utils.data.Dataset): target_ids = target_ids[1:] target_ids_length = len(target_ids) audio_length = target_ids_length - input_ids = prompt_ids_pre + target_ids + [self.tokenizer.pad_token_id] + target_ids #[bos, prompt, input, pad, target] - input_ids = torch.tensor(copy.deepcopy(input_ids), dtype=torch.int64) #[bos, prompt, input, pad, target] - input_ids[prompt_ids_length:prompt_ids_length+audio_length] = -1 # [bos, prompt,-1, pad, target] # it is no need, only for check - attention_mask = input_ids.ge(-1) # [true, true, true, true, true], length mask - + input_ids = ( + prompt_ids_pre + target_ids + [self.tokenizer.pad_token_id] + target_ids + ) # [bos, prompt, input, pad, target] + input_ids = torch.tensor( + copy.deepcopy(input_ids), dtype=torch.int64 + ) # [bos, prompt, input, pad, target] + input_ids[prompt_ids_length : prompt_ids_length + audio_length] = ( + -1 + ) # [bos, prompt,-1, pad, target] # it is no need, only for check + attention_mask = input_ids.ge(-1) # [true, true, true, true, true], length mask + # bos prompt audio target eos # prompt_answer = "{}{}".format(self.prompt_pre, target) # prompt_answer_ids = self.tokenizer.encode(prompt_answer) #[bos, prompt, input] @@ -96,35 +110,41 @@ class AudioLLMNARDataset(torch.utils.data.Dataset): if target_ids[0] == self.tokenizer.bos_token_id: target_ids = target_ids[1:] # target_ids_length = len(target_ids) - labels_ids = prompt_ids_pre + target_ids + target_ids + [self.tokenizer.eos_token_id] # [bos, prompt, input, target, eos] - labels_ids = torch.tensor(copy.deepcopy(labels_ids), dtype=torch.int64) # [bos, prompt, input, target, eos] + labels_ids = ( + prompt_ids_pre + target_ids + target_ids + [self.tokenizer.eos_token_id] + ) # [bos, prompt, input, target, eos] + labels_ids = torch.tensor( + copy.deepcopy(labels_ids), dtype=torch.int64 + ) # [bos, prompt, input, target, eos] labels_ids[:prompt_ids_length] = -1 # [-1, -1, input, target, eos] label_mask = labels_ids.ge(0) # [false, false, true, true, true], length mask labels_ids[~label_mask] = self.IGNORE_INDEX # [-1, -1, input, target, eos] - - audio_mask = [0] * prompt_ids_length + [1] * audio_length + [0] * target_ids_length + [0] # [0, 0, 1, 0, 0] + + audio_mask = ( + [0] * prompt_ids_length + [1] * audio_length + [0] * target_ids_length + [0] + ) # [0, 0, 1, 0, 0] audio_mask = torch.tensor(audio_mask, dtype=torch.float32) - - ids = target_ids #self.tokenizer.encode(target) # token ids is different from labels_ids + + ids = target_ids # self.tokenizer.encode(target) # token ids is different from labels_ids text = torch.tensor(ids, dtype=torch.int64) text_lengths = torch.tensor([len(ids)], dtype=torch.int32) prompt_bos_length = torch.tensor([len(prompt_ids_pre)], dtype=torch.int32) - - return {"speech": speech, - "speech_lengths": speech_lengths, - "text": text, - "text_lengths": text_lengths, - "input_ids": input_ids, - "attention_mask": attention_mask, - "labels_ids": labels_ids, - "label_mask": label_mask, - "audio_mask": audio_mask, - "prompt_bos_length": prompt_bos_length, - } - - - def collator(self, samples: list=None): + + return { + "speech": speech, + "speech_lengths": speech_lengths, + "text": text, + "text_lengths": text_lengths, + "input_ids": input_ids, + "attention_mask": attention_mask, + "labels_ids": labels_ids, + "label_mask": label_mask, + "audio_mask": audio_mask, + "prompt_bos_length": prompt_bos_length, + } + + def collator(self, samples: list = None): outputs = {} for sample in samples: for key in sample.keys(): @@ -135,12 +155,14 @@ class AudioLLMNARDataset(torch.utils.data.Dataset): for key, data_list in outputs.items(): if isinstance(data_list[0], torch.Tensor): if data_list[0].dtype == torch.int64 or data_list[0].dtype == torch.int32: - + pad_value = self.int_pad_value else: pad_value = self.float_pad_value - - outputs[key] = torch.nn.utils.rnn.pad_sequence(data_list, batch_first=True, padding_value=pad_value) + + outputs[key] = torch.nn.utils.rnn.pad_sequence( + data_list, batch_first=True, padding_value=pad_value + ) return outputs @@ -149,53 +171,58 @@ class AudioLLMDataset(torch.utils.data.Dataset): """ AudioLLMDataset """ - - def __init__(self, - path, - index_ds: str = None, - frontend=None, - tokenizer=None, - int_pad_value: int = -1, - float_pad_value: float = 0.0, - **kwargs): + + def __init__( + self, + path, + index_ds: str = None, + frontend=None, + tokenizer=None, + int_pad_value: int = -1, + float_pad_value: float = 0.0, + **kwargs + ): super().__init__() index_ds_class = tables.index_ds_classes.get(index_ds) self.index_ds = index_ds_class(path, **kwargs) preprocessor_speech = kwargs.get("preprocessor_speech", None) if preprocessor_speech: preprocessor_speech_class = tables.preprocessor_classes.get(preprocessor_speech) - preprocessor_speech = preprocessor_speech_class(**kwargs.get("preprocessor_speech_conf", {})) + preprocessor_speech = preprocessor_speech_class( + **kwargs.get("preprocessor_speech_conf", {}) + ) self.preprocessor_speech = preprocessor_speech preprocessor_text = kwargs.get("preprocessor_text", None) if preprocessor_text: preprocessor_text_class = tables.preprocessor_classes.get(preprocessor_text) preprocessor_text = preprocessor_text_class(**kwargs.get("preprocessor_text_conf", {})) self.preprocessor_text = preprocessor_text - + self.frontend = frontend self.fs = 16000 if frontend is None else frontend.fs self.data_type = "sound" self.tokenizer = tokenizer - + self.float_pad_value = float_pad_value self.prompt = kwargs.get("prompt", "Transcribe speech to text.") self.prompt_pre = "USER: \nINSTRUCTION: {}\nINPUT: ".format( - self.prompt) # "USER: \nINSTRUCTION: {}\nnINPUT: {}\nASSISTANT: " + self.prompt + ) # "USER: \nINSTRUCTION: {}\nnINPUT: {}\nASSISTANT: " self.prompt_af = "" self.IGNORE_INDEX = kwargs.get("IGNORE_INDEX", -100) self.int_pad_value = self.IGNORE_INDEX - + def get_source_len(self, index): item = self.index_ds[index] return self.index_ds.get_source_len(item) - + def get_target_len(self, index): item = self.index_ds[index] return self.index_ds.get_target_len(item) - + def __len__(self): return len(self.index_ds) - + def __getitem__(self, index): item = self.index_ds[index] # import pdb; @@ -204,17 +231,18 @@ class AudioLLMDataset(torch.utils.data.Dataset): data_src = load_audio_text_image_video(source, fs=self.fs) if self.preprocessor_speech: data_src = self.preprocessor_speech(data_src, fs=self.fs) - speech, speech_lengths = extract_fbank(data_src, data_type=self.data_type, frontend=self.frontend, - is_final=True) # speech: [b, T, d] + speech, speech_lengths = extract_fbank( + data_src, data_type=self.data_type, frontend=self.frontend, is_final=True + ) # speech: [b, T, d] speech = speech.squeeze(0) - + target = item["target"] if self.preprocessor_text: target = self.preprocessor_text(target) - + prompt_ids_pre = self.tokenizer.encode(self.prompt_pre) # [bos,prompt] prompt_ids_length = len(prompt_ids_pre) - + prompt_input = "{}{}".format(self.prompt_pre, target) prompt_input_ids = self.tokenizer.encode(prompt_input) audio_length = len(prompt_input_ids) - prompt_ids_length @@ -222,7 +250,7 @@ class AudioLLMDataset(torch.utils.data.Dataset): input_ids = torch.tensor(input_ids, dtype=torch.int64) # [bos, prompt, input, pad] input_ids[prompt_ids_length:] = -1 # [bos, prompt,-1,-1] attention_mask = input_ids.ge(-1) # [true, true, true, true], length mask - + prompt_answer = "{}{}".format(self.prompt_pre, target) prompt_answer_ids = self.tokenizer.encode(prompt_answer) answer_length = len(prompt_answer_ids) - prompt_ids_length @@ -231,25 +259,26 @@ class AudioLLMDataset(torch.utils.data.Dataset): labels_ids[:prompt_ids_length] = -1 # [-1, -1, input, eos] label_mask = labels_ids.ge(0) # [False,False,True,True] labels_ids[~label_mask] = self.IGNORE_INDEX # [-100,-100,input,eos] - + audio_mask = [0] * prompt_ids_length + [1] * audio_length + [0] audio_mask = torch.tensor(audio_mask, dtype=torch.float32) - + ids = self.tokenizer.encode(target) # token ids is different from labels_ids text = torch.tensor(ids, dtype=torch.int64) text_lengths = torch.tensor([len(ids)], dtype=torch.int32) - - return {"speech": speech, - "speech_lengths": speech_lengths, - "text": text, - "text_lengths": text_lengths, - "input_ids": input_ids, - "attention_mask": attention_mask, - "labels_ids": labels_ids, - "label_mask": label_mask, - "audio_mask": audio_mask, - } - + + return { + "speech": speech, + "speech_lengths": speech_lengths, + "text": text, + "text_lengths": text_lengths, + "input_ids": input_ids, + "attention_mask": attention_mask, + "labels_ids": labels_ids, + "label_mask": label_mask, + "audio_mask": audio_mask, + } + def collator(self, samples: list = None): outputs = {} for sample in samples: @@ -257,16 +286,18 @@ class AudioLLMDataset(torch.utils.data.Dataset): if key not in outputs: outputs[key] = [] outputs[key].append(sample[key]) - + for key, data_list in outputs.items(): if isinstance(data_list[0], torch.Tensor): if data_list[0].dtype == torch.int64 or data_list[0].dtype == torch.int32: - + pad_value = self.int_pad_value else: pad_value = self.float_pad_value - - outputs[key] = torch.nn.utils.rnn.pad_sequence(data_list, batch_first=True, padding_value=pad_value) + + outputs[key] = torch.nn.utils.rnn.pad_sequence( + data_list, batch_first=True, padding_value=pad_value + ) return outputs @@ -275,53 +306,58 @@ class AudioLLMARDataset(torch.utils.data.Dataset): """ AudioLLMDataset """ - - def __init__(self, - path, - index_ds: str = None, - frontend=None, - tokenizer=None, - int_pad_value: int = -1, - float_pad_value: float = 0.0, - **kwargs): + + def __init__( + self, + path, + index_ds: str = None, + frontend=None, + tokenizer=None, + int_pad_value: int = -1, + float_pad_value: float = 0.0, + **kwargs + ): super().__init__() index_ds_class = tables.index_ds_classes.get(index_ds) self.index_ds = index_ds_class(path, **kwargs) preprocessor_speech = kwargs.get("preprocessor_speech", None) if preprocessor_speech: preprocessor_speech_class = tables.preprocessor_classes.get(preprocessor_speech) - preprocessor_speech = preprocessor_speech_class(**kwargs.get("preprocessor_speech_conf", {})) + preprocessor_speech = preprocessor_speech_class( + **kwargs.get("preprocessor_speech_conf", {}) + ) self.preprocessor_speech = preprocessor_speech preprocessor_text = kwargs.get("preprocessor_text", None) if preprocessor_text: preprocessor_text_class = tables.preprocessor_classes.get(preprocessor_text) preprocessor_text = preprocessor_text_class(**kwargs.get("preprocessor_text_conf", {})) self.preprocessor_text = preprocessor_text - + self.frontend = frontend self.fs = 16000 if frontend is None else frontend.fs self.data_type = "sound" self.tokenizer = tokenizer - + self.float_pad_value = float_pad_value self.prompt = kwargs.get("prompt", "Transcribe speech to text.") self.prompt_pre = "USER: \nINSTRUCTION: {}\nINPUT: ".format( - self.prompt) # "USER: \nINSTRUCTION: {}\nnINPUT: {}\nASSISTANT: " + self.prompt + ) # "USER: \nINSTRUCTION: {}\nnINPUT: {}\nASSISTANT: " self.prompt_af = "" self.IGNORE_INDEX = kwargs.get("IGNORE_INDEX", -100) self.int_pad_value = self.IGNORE_INDEX - + def get_source_len(self, index): item = self.index_ds[index] return self.index_ds.get_source_len(item) - + def get_target_len(self, index): item = self.index_ds[index] return self.index_ds.get_target_len(item) - + def __len__(self): return len(self.index_ds) - + def __getitem__(self, index): item = self.index_ds[index] # import pdb; @@ -330,17 +366,18 @@ class AudioLLMARDataset(torch.utils.data.Dataset): data_src = load_audio_text_image_video(source, fs=self.fs) if self.preprocessor_speech: data_src = self.preprocessor_speech(data_src, fs=self.fs) - speech, speech_lengths = extract_fbank(data_src, data_type=self.data_type, frontend=self.frontend, - is_final=True) # speech: [b, T, d] + speech, speech_lengths = extract_fbank( + data_src, data_type=self.data_type, frontend=self.frontend, is_final=True + ) # speech: [b, T, d] speech = speech.squeeze(0) - + target = item["target"] if self.preprocessor_text: target = self.preprocessor_text(target) - + prompt_ids_pre = self.tokenizer.encode(self.prompt_pre) # [bos,prompt] prompt_ids_length = len(prompt_ids_pre) - + prompt_input = "{}{}".format(self.prompt_pre, target) prompt_input_ids = self.tokenizer.encode(prompt_input) audio_length = len(prompt_input_ids) - prompt_ids_length @@ -348,7 +385,7 @@ class AudioLLMARDataset(torch.utils.data.Dataset): input_ids = torch.tensor(input_ids, dtype=torch.int64) # [bos, prompt, input, pad] input_ids[prompt_ids_length:] = -1 # [bos, prompt,-1,-1] attention_mask = input_ids.ge(-1) # [true, true, true, true], length mask - + prompt_answer = "{}{}".format(self.prompt_pre, target) prompt_answer_ids = self.tokenizer.encode(prompt_answer) answer_length = len(prompt_answer_ids) - prompt_ids_length @@ -357,25 +394,26 @@ class AudioLLMARDataset(torch.utils.data.Dataset): labels_ids[:prompt_ids_length] = -1 # [-1, -1, input, eos] label_mask = labels_ids.ge(0) # [False,False,True,True] labels_ids[~label_mask] = self.IGNORE_INDEX # [-100,-100,input,eos] - + audio_mask = [0] * prompt_ids_length + [1] * audio_length + [0] audio_mask = torch.tensor(audio_mask, dtype=torch.float32) - + ids = self.tokenizer.encode(target) # token ids is different from labels_ids text = torch.tensor(ids, dtype=torch.int64) text_lengths = torch.tensor([len(ids)], dtype=torch.int32) - - return {"speech": speech, - "speech_lengths": speech_lengths, - "text": text, - "text_lengths": text_lengths, - "input_ids": input_ids, - "attention_mask": attention_mask, - "labels_ids": labels_ids, - "label_mask": label_mask, - "audio_mask": audio_mask, - } - + + return { + "speech": speech, + "speech_lengths": speech_lengths, + "text": text, + "text_lengths": text_lengths, + "input_ids": input_ids, + "attention_mask": attention_mask, + "labels_ids": labels_ids, + "label_mask": label_mask, + "audio_mask": audio_mask, + } + def collator(self, samples: list = None): outputs = {} for sample in samples: @@ -383,14 +421,16 @@ class AudioLLMARDataset(torch.utils.data.Dataset): if key not in outputs: outputs[key] = [] outputs[key].append(sample[key]) - + for key, data_list in outputs.items(): if isinstance(data_list[0], torch.Tensor): if data_list[0].dtype == torch.int64 or data_list[0].dtype == torch.int32: - + pad_value = self.int_pad_value else: pad_value = self.float_pad_value - - outputs[key] = torch.nn.utils.rnn.pad_sequence(data_list, batch_first=True, padding_value=pad_value) + + outputs[key] = torch.nn.utils.rnn.pad_sequence( + data_list, batch_first=True, padding_value=pad_value + ) return outputs diff --git a/funasr/datasets/llm_datasets/preprocessor.py b/funasr/datasets/llm_datasets/preprocessor.py index 4cbef3318..b99255e38 100644 --- a/funasr/datasets/llm_datasets/preprocessor.py +++ b/funasr/datasets/llm_datasets/preprocessor.py @@ -16,22 +16,19 @@ from funasr.tokenizer.cleaner import TextCleaner from funasr.register import tables - @tables.register("preprocessor_classes", "TextPreprocessRemovePunctuation") class TextPreprocessRemovePunctuation(nn.Module): - def __init__(self, - **kwargs): - super().__init__() - - - def forward(self, text, **kwargs): - # 定义英文标点符号 - en_punct = string.punctuation - # 定义中文标点符号（部分常用的） - cn_punct = '。？！，、；：“”‘’（）《》【】…—～·' - # 合并英文和中文标点符号 - all_punct = en_punct + cn_punct - # 创建正则表达式模式，匹配任何在all_punct中的字符 - punct_pattern = re.compile('[{}]'.format(re.escape(all_punct))) - # 使用正则表达式的sub方法替换掉这些字符 - return punct_pattern.sub('', text) + def __init__(self, **kwargs): + super().__init__() + + def forward(self, text, **kwargs): + # 定义英文标点符号 + en_punct = string.punctuation + # 定义中文标点符号（部分常用的） + cn_punct = "。？！，、；：“”‘’（）《》【】…—～·" + # 合并英文和中文标点符号 + all_punct = en_punct + cn_punct + # 创建正则表达式模式，匹配任何在all_punct中的字符 + punct_pattern = re.compile("[{}]".format(re.escape(all_punct))) + # 使用正则表达式的sub方法替换掉这些字符 + return punct_pattern.sub("", text) diff --git a/funasr/datasets/llm_datasets_qwenaudio/datasets.py b/funasr/datasets/llm_datasets_qwenaudio/datasets.py index 7a2ce2288..b56e624ac 100644 --- a/funasr/datasets/llm_datasets_qwenaudio/datasets.py +++ b/funasr/datasets/llm_datasets_qwenaudio/datasets.py @@ -10,34 +10,38 @@ class AudioLLMQwenAudioDataset(torch.utils.data.Dataset): """ AudioLLMDataset """ - - def __init__(self, - path, - index_ds: str = None, - frontend=None, - tokenizer=None, - int_pad_value: int = -1, - float_pad_value: float = 0.0, - **kwargs): + + def __init__( + self, + path, + index_ds: str = None, + frontend=None, + tokenizer=None, + int_pad_value: int = -1, + float_pad_value: float = 0.0, + **kwargs + ): super().__init__() index_ds_class = tables.index_ds_classes.get(index_ds) self.index_ds = index_ds_class(path, **kwargs) preprocessor_speech = kwargs.get("preprocessor_speech", None) if preprocessor_speech: preprocessor_speech_class = tables.preprocessor_classes.get(preprocessor_speech) - preprocessor_speech = preprocessor_speech_class(**kwargs.get("preprocessor_speech_conf", {})) + preprocessor_speech = preprocessor_speech_class( + **kwargs.get("preprocessor_speech_conf", {}) + ) self.preprocessor_speech = preprocessor_speech preprocessor_text = kwargs.get("preprocessor_text", None) if preprocessor_text: preprocessor_text_class = tables.preprocessor_classes.get(preprocessor_text) preprocessor_text = preprocessor_text_class(**kwargs.get("preprocessor_text_conf", {})) self.preprocessor_text = preprocessor_text - + self.frontend = frontend self.fs = 16000 if frontend is None else frontend.fs self.data_type = "sound" self.tokenizer = tokenizer - + self.float_pad_value = float_pad_value self.prompt = kwargs.get("prompt", "Transcribe speech to text.") # self.prompt_pre = "USER: \nINSTRUCTION: {}\nINPUT: ".format(self.prompt) # "USER: \nINSTRUCTION: {}\nnINPUT: {}\nASSISTANT: " @@ -48,19 +52,18 @@ class AudioLLMQwenAudioDataset(torch.utils.data.Dataset): self.audio_encoder_downsample_rate = kwargs.get("audio_encoder_downsample_rate", 2) self.prompt_template = "{}" self.answer_template = "{}" - - + def get_source_len(self, index): item = self.index_ds[index] return self.index_ds.get_source_len(item) - + def get_target_len(self, index): item = self.index_ds[index] return self.index_ds.get_target_len(item) - + def __len__(self): return len(self.index_ds) - + def __getitem__(self, index): item = self.index_ds[index] # import pdb; @@ -69,13 +72,18 @@ class AudioLLMQwenAudioDataset(torch.utils.data.Dataset): data_src = load_audio_text_image_video(source, fs=self.fs) if self.preprocessor_speech: data_src = self.preprocessor_speech(data_src, fs=self.fs) - speech, speech_lengths = extract_fbank(data_src, data_type=self.data_type, frontend=self.frontend, - is_final=True) # speech: [b, T, d] + speech, speech_lengths = extract_fbank( + data_src, data_type=self.data_type, frontend=self.frontend, is_final=True + ) # speech: [b, T, d] speech = speech.squeeze(0) - audio_pseudo_length = (speech.shape[0] + 1) // self.audio_adaptor_downsample_rate // self.audio_encoder_downsample_rate - audio_pseudo = torch.full((audio_pseudo_length,), -1) # placeholder - + audio_pseudo_length = ( + (speech.shape[0] + 1) + // self.audio_adaptor_downsample_rate + // self.audio_encoder_downsample_rate + ) + audio_pseudo = torch.full((audio_pseudo_length,), -1) # placeholder + target = item["target"] if self.preprocessor_text: target = self.preprocessor_text(target) @@ -83,21 +91,21 @@ class AudioLLMQwenAudioDataset(torch.utils.data.Dataset): self.prompt_pre = self.prompt_template.format(self.prompt) prompt_ids_pre = self.tokenizer.encode(self.prompt_pre) # [bos,prompt] prompt_pre_length = len(prompt_ids_pre) - + # input input = self.answer_template.format(target.lower()) prompt_input = "{}{}".format(self.prompt_pre, input) - prompt_input_ids = self.tokenizer.encode(prompt_input) # [bos, prompt, input] + prompt_input_ids = self.tokenizer.encode(prompt_input) # [bos, prompt, input] # audio_length = len(prompt_input_ids) - prompt_pre_length - input_ids = prompt_input_ids + [self.tokenizer.pad_token_id] # [bos, prompt, input, pad] + input_ids = prompt_input_ids + [self.tokenizer.pad_token_id] # [bos, prompt, input, pad] input_ids_length = len(input_ids) input_ids = torch.tensor(input_ids, dtype=torch.int64) # [bos, prompt, input, pad] - input_ids = torch.cat((audio_pseudo, input_ids)) # [audio, bos, prompt, input, pad] + input_ids = torch.cat((audio_pseudo, input_ids)) # [audio, bos, prompt, input, pad] # input_ids[:audio_pseudo_length] = -1 # [-1, bos, prompt, input, pad] attention_mask = input_ids.ge(-1) # [true, true, true, true, true], length mask # input_ids[prompt_pre_length:] = -1 # [bos, prompt,-1,-1] # attention_mask = input_ids.ge(-1) # [true, true, true, true], length mask - + # label answer = self.answer_template.format(target.lower()) prompt_answer = "{}{}".format(self.prompt_pre, answer) @@ -106,30 +114,31 @@ class AudioLLMQwenAudioDataset(torch.utils.data.Dataset): labels_ids = copy.deepcopy(prompt_answer_ids) + [self.tokenizer.eos_token_id] labels_ids = torch.tensor(labels_ids, dtype=torch.int64) # [bos, prompt, answer, eos] labels_ids = torch.cat((audio_pseudo, labels_ids)) # [audio, bos, prompt, answer, eos] - labels_ids[:audio_pseudo_length+prompt_pre_length] = -1 # [-1, -1, -1, answer, eos] + labels_ids[: audio_pseudo_length + prompt_pre_length] = -1 # [-1, -1, -1, answer, eos] # labels_ids[:prompt_pre_length] = -1 # [-1, -1, input, eos] label_mask = labels_ids.ge(0) # [false, false, false, true, true] labels_ids[~label_mask] = self.IGNORE_INDEX # [-100, -100, -100, answer, eos] - + # audio_mask for input_ids - audio_mask = [1]*audio_pseudo_length + [0]*input_ids_length + audio_mask = [1] * audio_pseudo_length + [0] * input_ids_length audio_mask = torch.tensor(audio_mask, dtype=torch.float32) - + ids = self.tokenizer.encode(target) # token ids is different from labels_ids text = torch.tensor(ids, dtype=torch.int64) text_lengths = torch.tensor([len(ids)], dtype=torch.int32) - - return {"speech": speech, - "speech_lengths": speech_lengths, - "text": text, - "text_lengths": text_lengths, - "input_ids": input_ids, - "attention_mask": attention_mask, - "labels_ids": labels_ids, - "label_mask": label_mask, - "audio_mask": audio_mask, - } - + + return { + "speech": speech, + "speech_lengths": speech_lengths, + "text": text, + "text_lengths": text_lengths, + "input_ids": input_ids, + "attention_mask": attention_mask, + "labels_ids": labels_ids, + "label_mask": label_mask, + "audio_mask": audio_mask, + } + def collator(self, samples: list = None): outputs = {} for sample in samples: @@ -137,14 +146,16 @@ class AudioLLMQwenAudioDataset(torch.utils.data.Dataset): if key not in outputs: outputs[key] = [] outputs[key].append(sample[key]) - + for key, data_list in outputs.items(): if isinstance(data_list[0], torch.Tensor): if data_list[0].dtype == torch.int64 or data_list[0].dtype == torch.int32: - + pad_value = self.int_pad_value else: pad_value = self.float_pad_value - - outputs[key] = torch.nn.utils.rnn.pad_sequence(data_list, batch_first=True, padding_value=pad_value) + + outputs[key] = torch.nn.utils.rnn.pad_sequence( + data_list, batch_first=True, padding_value=pad_value + ) return outputs diff --git a/funasr/datasets/llm_datasets_vicuna/datasets.py b/funasr/datasets/llm_datasets_vicuna/datasets.py index b5fc5bc71..04fa514a1 100644 --- a/funasr/datasets/llm_datasets_vicuna/datasets.py +++ b/funasr/datasets/llm_datasets_vicuna/datasets.py @@ -10,34 +10,38 @@ class AudioLLMVicunaDataset(torch.utils.data.Dataset): """ AudioLLMDataset """ - - def __init__(self, - path, - index_ds: str = None, - frontend=None, - tokenizer=None, - int_pad_value: int = -1, - float_pad_value: float = 0.0, - **kwargs): + + def __init__( + self, + path, + index_ds: str = None, + frontend=None, + tokenizer=None, + int_pad_value: int = -1, + float_pad_value: float = 0.0, + **kwargs + ): super().__init__() index_ds_class = tables.index_ds_classes.get(index_ds) self.index_ds = index_ds_class(path, **kwargs) preprocessor_speech = kwargs.get("preprocessor_speech", None) if preprocessor_speech: preprocessor_speech_class = tables.preprocessor_classes.get(preprocessor_speech) - preprocessor_speech = preprocessor_speech_class(**kwargs.get("preprocessor_speech_conf", {})) + preprocessor_speech = preprocessor_speech_class( + **kwargs.get("preprocessor_speech_conf", {}) + ) self.preprocessor_speech = preprocessor_speech preprocessor_text = kwargs.get("preprocessor_text", None) if preprocessor_text: preprocessor_text_class = tables.preprocessor_classes.get(preprocessor_text) preprocessor_text = preprocessor_text_class(**kwargs.get("preprocessor_text_conf", {})) self.preprocessor_text = preprocessor_text - + self.frontend = frontend self.fs = 16000 if frontend is None else frontend.fs self.data_type = "sound" self.tokenizer = tokenizer - + self.float_pad_value = float_pad_value self.prompt = kwargs.get("prompt", "Transcribe speech to text.") # self.prompt_pre = "USER: \nINSTRUCTION: {}\nINPUT: ".format(self.prompt) # "USER: \nINSTRUCTION: {}\nnINPUT: {}\nASSISTANT: " @@ -48,19 +52,18 @@ class AudioLLMVicunaDataset(torch.utils.data.Dataset): self.audio_encoder_downsample_rate = kwargs.get("audio_encoder_downsample_rate", 2) self.prompt_template = "USER: {}\n ASSISTANT:" self.answer_template = "{}" - - + def get_source_len(self, index): item = self.index_ds[index] return self.index_ds.get_source_len(item) - + def get_target_len(self, index): item = self.index_ds[index] return self.index_ds.get_target_len(item) - + def __len__(self): return len(self.index_ds) - + def __getitem__(self, index): item = self.index_ds[index] # import pdb; @@ -69,13 +72,18 @@ class AudioLLMVicunaDataset(torch.utils.data.Dataset): data_src = load_audio_text_image_video(source, fs=self.fs) if self.preprocessor_speech: data_src = self.preprocessor_speech(data_src, fs=self.fs) - speech, speech_lengths = extract_fbank(data_src, data_type=self.data_type, frontend=self.frontend, - is_final=True) # speech: [b, T, d] + speech, speech_lengths = extract_fbank( + data_src, data_type=self.data_type, frontend=self.frontend, is_final=True + ) # speech: [b, T, d] speech = speech.squeeze(0) - audio_pseudo_length = (speech.shape[0] + 1) // self.audio_adaptor_downsample_rate // self.audio_encoder_downsample_rate - audio_pseudo = torch.full((audio_pseudo_length,), -1) # placeholder - + audio_pseudo_length = ( + (speech.shape[0] + 1) + // self.audio_adaptor_downsample_rate + // self.audio_encoder_downsample_rate + ) + audio_pseudo = torch.full((audio_pseudo_length,), -1) # placeholder + target = item["target"] if self.preprocessor_text: target = self.preprocessor_text(target) @@ -83,21 +91,21 @@ class AudioLLMVicunaDataset(torch.utils.data.Dataset): self.prompt_pre = self.prompt_template.format(self.prompt) prompt_ids_pre = self.tokenizer.encode(self.prompt_pre) # [bos,prompt] prompt_pre_length = len(prompt_ids_pre) - + # input input = self.answer_template.format(target.lower()) prompt_input = "{}{}".format(self.prompt_pre, input) - prompt_input_ids = self.tokenizer.encode(prompt_input) # [bos, prompt, input] + prompt_input_ids = self.tokenizer.encode(prompt_input) # [bos, prompt, input] # audio_length = len(prompt_input_ids) - prompt_pre_length - input_ids = prompt_input_ids + [self.tokenizer.pad_token_id] # [bos, prompt, input, pad] + input_ids = prompt_input_ids + [self.tokenizer.pad_token_id] # [bos, prompt, input, pad] input_ids_length = len(input_ids) input_ids = torch.tensor(input_ids, dtype=torch.int64) # [bos, prompt, input, pad] - input_ids = torch.cat((audio_pseudo, input_ids)) # [audio, bos, prompt, input, pad] + input_ids = torch.cat((audio_pseudo, input_ids)) # [audio, bos, prompt, input, pad] # input_ids[:audio_pseudo_length] = -1 # [-1, bos, prompt, input, pad] attention_mask = input_ids.ge(-1) # [true, true, true, true, true], length mask # input_ids[prompt_pre_length:] = -1 # [bos, prompt,-1,-1] # attention_mask = input_ids.ge(-1) # [true, true, true, true], length mask - + # label answer = self.answer_template.format(target.lower()) prompt_answer = "{}{}".format(self.prompt_pre, answer) @@ -106,30 +114,31 @@ class AudioLLMVicunaDataset(torch.utils.data.Dataset): labels_ids = copy.deepcopy(prompt_answer_ids) + [self.tokenizer.eos_token_id] labels_ids = torch.tensor(labels_ids, dtype=torch.int64) # [bos, prompt, answer, eos] labels_ids = torch.cat((audio_pseudo, labels_ids)) # [audio, bos, prompt, answer, eos] - labels_ids[:audio_pseudo_length+prompt_pre_length] = -1 # [-1, -1, -1, answer, eos] + labels_ids[: audio_pseudo_length + prompt_pre_length] = -1 # [-1, -1, -1, answer, eos] # labels_ids[:prompt_pre_length] = -1 # [-1, -1, input, eos] label_mask = labels_ids.ge(0) # [false, false, false, true, true] labels_ids[~label_mask] = self.IGNORE_INDEX # [-100, -100, -100, answer, eos] - + # audio_mask for input_ids - audio_mask = [1]*audio_pseudo_length + [0]*input_ids_length + audio_mask = [1] * audio_pseudo_length + [0] * input_ids_length audio_mask = torch.tensor(audio_mask, dtype=torch.float32) - + ids = self.tokenizer.encode(target) # token ids is different from labels_ids text = torch.tensor(ids, dtype=torch.int64) text_lengths = torch.tensor([len(ids)], dtype=torch.int32) - - return {"speech": speech, - "speech_lengths": speech_lengths, - "text": text, - "text_lengths": text_lengths, - "input_ids": input_ids, - "attention_mask": attention_mask, - "labels_ids": labels_ids, - "label_mask": label_mask, - "audio_mask": audio_mask, - } - + + return { + "speech": speech, + "speech_lengths": speech_lengths, + "text": text, + "text_lengths": text_lengths, + "input_ids": input_ids, + "attention_mask": attention_mask, + "labels_ids": labels_ids, + "label_mask": label_mask, + "audio_mask": audio_mask, + } + def collator(self, samples: list = None): outputs = {} for sample in samples: @@ -137,14 +146,16 @@ class AudioLLMVicunaDataset(torch.utils.data.Dataset): if key not in outputs: outputs[key] = [] outputs[key].append(sample[key]) - + for key, data_list in outputs.items(): if isinstance(data_list[0], torch.Tensor): if data_list[0].dtype == torch.int64 or data_list[0].dtype == torch.int32: - + pad_value = self.int_pad_value else: pad_value = self.float_pad_value - - outputs[key] = torch.nn.utils.rnn.pad_sequence(data_list, batch_first=True, padding_value=pad_value) + + outputs[key] = torch.nn.utils.rnn.pad_sequence( + data_list, batch_first=True, padding_value=pad_value + ) return outputs diff --git a/funasr/datasets/sense_voice_datasets/datasets.py b/funasr/datasets/sense_voice_datasets/datasets.py index 956cf79ab..5468ea66e 100644 --- a/funasr/datasets/sense_voice_datasets/datasets.py +++ b/funasr/datasets/sense_voice_datasets/datasets.py @@ -10,28 +10,33 @@ class SenseVoiceDataset(torch.utils.data.Dataset): """ SenseVoiceDataset """ - def __init__(self, - path, - index_ds: str = None, - frontend=None, - tokenizer=None, - int_pad_value: int = -1, - float_pad_value: float = 0.0, - **kwargs): + + def __init__( + self, + path, + index_ds: str = None, + frontend=None, + tokenizer=None, + int_pad_value: int = -1, + float_pad_value: float = 0.0, + **kwargs, + ): super().__init__() index_ds_class = tables.index_ds_classes.get(index_ds) self.index_ds = index_ds_class(path, **kwargs) preprocessor_speech = kwargs.get("preprocessor_speech", None) if preprocessor_speech: preprocessor_speech_class = tables.preprocessor_classes.get(preprocessor_speech) - preprocessor_speech = preprocessor_speech_class(**kwargs.get("preprocessor_speech_conf")) + preprocessor_speech = preprocessor_speech_class( + **kwargs.get("preprocessor_speech_conf") + ) self.preprocessor_speech = preprocessor_speech preprocessor_text = kwargs.get("preprocessor_text", None) if preprocessor_text: preprocessor_text_class = tables.preprocessor_classes.get(preprocessor_text) preprocessor_text = preprocessor_text_class(**kwargs.get("preprocessor_text_conf")) self.preprocessor_text = preprocessor_text - + self.frontend = frontend self.fs = 16000 if frontend is None else frontend.fs self.data_type = "sound" @@ -41,18 +46,18 @@ class SenseVoiceDataset(torch.utils.data.Dataset): self.float_pad_value = float_pad_value self.sos = kwargs.get("sos", "<|startoftranscript|>") self.eos = kwargs.get("eos", "<|endoftext|>") - + def get_source_len(self, index): item = self.index_ds[index] return self.index_ds.get_source_len(item) - + def get_target_len(self, index): item = self.index_ds[index] return self.index_ds.get_target_len(item) - + def __len__(self): return len(self.index_ds) - + def __getitem__(self, index): item = self.index_ds[index] # import pdb; @@ -61,42 +66,46 @@ class SenseVoiceDataset(torch.utils.data.Dataset): data_src = load_audio_text_image_video(source, fs=self.fs) if self.preprocessor_speech: data_src = self.preprocessor_speech(data_src, fs=self.fs) - speech, speech_lengths = extract_fbank(data_src, data_type=self.data_type, frontend=self.frontend, is_final=True) # speech: [b, T, d] + speech, speech_lengths = extract_fbank( + data_src, data_type=self.data_type, frontend=self.frontend, is_final=True + ) # speech: [b, T, d] speech = speech.permute(0, 2, 1) target = item["target"] if self.preprocessor_text: target = self.preprocessor_text(target) - + task = item.get("prompt", "<|ASR|>") text_language = item.get("text_language", "<|zh|>") prompt = f"{self.sos}{task}{text_language}" prompt_ids = self.tokenizer.encode(prompt, allowed_special="all") - prompt_ids_len = len(prompt_ids) - 1 # [sos, task] + prompt_ids_len = len(prompt_ids) - 1 # [sos, task] target_ids = self.tokenizer.encode(target, allowed_special="all") - target_ids_len = len(target_ids) + 1 # [lid, text] - - eos = self.tokenizer.encode(self.eos, allowed_special="all") # [eos] - + target_ids_len = len(target_ids) + 1 # [lid, text] + + eos = self.tokenizer.encode(self.eos, allowed_special="all") # [eos] + ids = prompt_ids + target_ids + eos ids_lengths = len(ids) - + text = torch.tensor(ids, dtype=torch.int64) text_lengths = torch.tensor([ids_lengths], dtype=torch.int32) - target_mask = [0] * (prompt_ids_len) + [1] * (target_ids_len) + [1] # [sos, task, lid, text, eos]: [0, 0, 1, 1, 1] + target_mask = ( + [0] * (prompt_ids_len) + [1] * (target_ids_len) + [1] + ) # [sos, task, lid, text, eos]: [0, 0, 1, 1, 1] target_mask = torch.tensor(target_mask, dtype=torch.float32) - return {"speech": speech[0, :, :], - "speech_lengths": speech_lengths, - "text": text, - "text_lengths": text_lengths, - "target_mask": target_mask, - } - - - def collator(self, samples: list=None): + return { + "speech": speech[0, :, :], + "speech_lengths": speech_lengths, + "text": text, + "text_lengths": text_lengths, + "target_mask": target_mask, + } + + def collator(self, samples: list = None): outputs = {} for sample in samples: for key in sample.keys(): @@ -107,12 +116,12 @@ class SenseVoiceDataset(torch.utils.data.Dataset): for key, data_list in outputs.items(): if isinstance(data_list[0], torch.Tensor): if data_list[0].dtype == torch.int64 or data_list[0].dtype == torch.int32: - + pad_value = self.int_pad_value else: pad_value = self.float_pad_value - - outputs[key] = torch.nn.utils.rnn.pad_sequence(data_list, batch_first=True, padding_value=pad_value) + + outputs[key] = torch.nn.utils.rnn.pad_sequence( + data_list, batch_first=True, padding_value=pad_value + ) return outputs - - diff --git a/funasr/download/download_dataset_from_hub.py b/funasr/download/download_dataset_from_hub.py index d06d8213a..2419ce869 100644 --- a/funasr/download/download_dataset_from_hub.py +++ b/funasr/download/download_dataset_from_hub.py @@ -1,11 +1,14 @@ - def download_dataset(): - pass + pass + def download_dataset_from_ms(**kwargs): - from modelscope.msdatasets import MsDataset - dataset_name = kwargs.get("dataset_name", 'speech_asr/speech_asr_aishell1_trainsets') - subset_name = kwargs.get("subset_name", 'default') - split = kwargs.get("split", 'train') - data_dump_dir = kwargs.get("data_dump_dir", None) - ds = MsDataset.load(dataset_name=dataset_name, subset_name=subset_name, split=split, cache_dir=data_dump_dir) \ No newline at end of file + from modelscope.msdatasets import MsDataset + + dataset_name = kwargs.get("dataset_name", "speech_asr/speech_asr_aishell1_trainsets") + subset_name = kwargs.get("subset_name", "default") + split = kwargs.get("split", "train") + data_dump_dir = kwargs.get("data_dump_dir", None) + ds = MsDataset.load( + dataset_name=dataset_name, subset_name=subset_name, split=split, cache_dir=data_dump_dir + ) diff --git a/funasr/download/download_from_hub.py b/funasr/download/download_from_hub.py index 8255947bf..43f5b67e6 100644 --- a/funasr/download/download_from_hub.py +++ b/funasr/download/download_from_hub.py @@ -22,9 +22,10 @@ def download_model(**kwargs): if model_or_path in name_maps_openai: model_or_path = name_maps_openai[model_or_path] kwargs["model_path"] = model_or_path - + return kwargs + def download_from_ms(**kwargs): model_or_path = kwargs.get("model") if model_or_path in name_maps_ms: @@ -32,18 +33,21 @@ def download_from_ms(**kwargs): model_revision = kwargs.get("model_revision", "master") if not os.path.exists(model_or_path) and "model_path" not in kwargs: try: - model_or_path = get_or_download_model_dir(model_or_path, model_revision, - is_training=kwargs.get("is_training"), - check_latest=kwargs.get("check_latest", True)) + model_or_path = get_or_download_model_dir( + model_or_path, + model_revision, + is_training=kwargs.get("is_training"), + check_latest=kwargs.get("check_latest", True), + ) except Exception as e: print(f"Download: {model_or_path} failed!: {e}") - + kwargs["model_path"] = model_or_path if "model_path" not in kwargs else kwargs["model_path"] - + if os.path.exists(os.path.join(model_or_path, "configuration.json")): - with open(os.path.join(model_or_path, "configuration.json"), 'r', encoding='utf-8') as f: + with open(os.path.join(model_or_path, "configuration.json"), "r", encoding="utf-8") as f: conf_json = json.load(f) - + cfg = {} if "file_path_metas" in conf_json: add_file_root_path(model_or_path, conf_json["file_path_metas"], cfg) @@ -52,7 +56,9 @@ def download_from_ms(**kwargs): config = OmegaConf.load(cfg["config"]) kwargs = OmegaConf.merge(config, cfg) kwargs["model"] = config["model"] - elif os.path.exists(os.path.join(model_or_path, "config.yaml")) and os.path.exists(os.path.join(model_or_path, "model.pt")): + elif os.path.exists(os.path.join(model_or_path, "config.yaml")) and os.path.exists( + os.path.join(model_or_path, "model.pt") + ): config = OmegaConf.load(os.path.join(model_or_path, "config.yaml")) kwargs = OmegaConf.merge(config, kwargs) init_param = os.path.join(model_or_path, "model.pb") @@ -76,11 +82,13 @@ def download_from_ms(**kwargs): requirements = os.path.join(model_or_path, "requirements.txt") print(f"Detect model requirements, begin to install it: {requirements}") from funasr.utils.install_model_requirements import install_requirements + install_requirements(requirements) return kwargs -def add_file_root_path(model_or_path: str, file_path_metas: dict, cfg = {}): - + +def add_file_root_path(model_or_path: str, file_path_metas: dict, cfg={}): + if isinstance(file_path_metas, dict): for k, v in file_path_metas.items(): if isinstance(v, str): @@ -91,17 +99,17 @@ def add_file_root_path(model_or_path: str, file_path_metas: dict, cfg = {}): if k not in cfg: cfg[k] = {} add_file_root_path(model_or_path, v, cfg[k]) - + return cfg def get_or_download_model_dir( - model, - model_revision=None, - is_training=False, - check_latest=True, - ): - """ Get local model directory or download model if necessary. + model, + model_revision=None, + is_training=False, + check_latest=True, +): + """Get local model directory or download model if necessary. Args: model (str): model id or path to local model directory. @@ -112,27 +120,19 @@ def get_or_download_model_dir( from modelscope.hub.snapshot_download import snapshot_download from modelscope.utils.constant import Invoke, ThirdParty - + key = Invoke.LOCAL_TRAINER if is_training else Invoke.PIPELINE - + if os.path.exists(model) and check_latest: - model_cache_dir = model if os.path.isdir( - model) else os.path.dirname(model) + model_cache_dir = model if os.path.isdir(model) else os.path.dirname(model) try: check_local_model_is_latest( - model_cache_dir, - user_agent={ - Invoke.KEY: key, - ThirdParty.KEY: "funasr" - }) + model_cache_dir, user_agent={Invoke.KEY: key, ThirdParty.KEY: "funasr"} + ) except: print("could not check the latest version") else: model_cache_dir = snapshot_download( - model, - revision=model_revision, - user_agent={ - Invoke.KEY: key, - ThirdParty.KEY: "funasr" - }) - return model_cache_dir \ No newline at end of file + model, revision=model_revision, user_agent={Invoke.KEY: key, ThirdParty.KEY: "funasr"} + ) + return model_cache_dir diff --git a/funasr/download/file.py b/funasr/download/file.py index da4958aae..a8a13f874 100644 --- a/funasr/download/file.py +++ b/funasr/download/file.py @@ -10,6 +10,7 @@ from typing import Generator, Union import requests from urllib.parse import urlparse + def download_from_url(url): result = urlparse(url) file_path = None @@ -21,11 +22,12 @@ def download_from_url(url): if not os.path.exists(work_dir): os.makedirs(work_dir) file_path = os.path.join(work_dir, os.path.basename(url)) - with open(file_path, 'wb') as fb: + with open(file_path, "wb") as fb: fb.write(data) assert file_path is not None, f"failed to download: {url}" return file_path + class Storage(metaclass=ABCMeta): """Abstract class of storage. @@ -47,10 +49,7 @@ class Storage(metaclass=ABCMeta): pass @abstractmethod - def write_text(self, - obj: str, - filepath: Union[str, Path], - encoding: str = 'utf-8') -> None: + def write_text(self, obj: str, filepath: Union[str, Path], encoding: str = "utf-8") -> None: pass @@ -66,13 +65,11 @@ class LocalStorage(Storage): Returns: bytes: Expected bytes object. """ - with open(filepath, 'rb') as f: + with open(filepath, "rb") as f: content = f.read() return content - def read_text(self, - filepath: Union[str, Path], - encoding: str = 'utf-8') -> str: + def read_text(self, filepath: Union[str, Path], encoding: str = "utf-8") -> str: """Read data from a given ``filepath`` with 'r' mode. Args: @@ -83,7 +80,7 @@ class LocalStorage(Storage): Returns: str: Expected text reading from ``filepath``. """ - with open(filepath, 'r', encoding=encoding) as f: + with open(filepath, "r", encoding=encoding) as f: value_buf = f.read() return value_buf @@ -102,13 +99,10 @@ class LocalStorage(Storage): if dirname and not os.path.exists(dirname): os.makedirs(dirname, exist_ok=True) - with open(filepath, 'wb') as f: + with open(filepath, "wb") as f: f.write(obj) - def write_text(self, - obj: str, - filepath: Union[str, Path], - encoding: str = 'utf-8') -> None: + def write_text(self, obj: str, filepath: Union[str, Path], encoding: str = "utf-8") -> None: """Write data to a given ``filepath`` with 'w' mode. Note: @@ -125,14 +119,11 @@ class LocalStorage(Storage): if dirname and not os.path.exists(dirname): os.makedirs(dirname, exist_ok=True) - with open(filepath, 'w', encoding=encoding) as f: + with open(filepath, "w", encoding=encoding) as f: f.write(obj) @contextlib.contextmanager - def as_local_path( - self, - filepath: Union[str, - Path]) -> Generator[Union[str, Path], None, None]: + def as_local_path(self, filepath: Union[str, Path]) -> Generator[Union[str, Path], None, None]: """Only for unified API and do nothing.""" yield filepath @@ -152,8 +143,7 @@ class HTTPStorage(Storage): return r.text @contextlib.contextmanager - def as_local_path( - self, filepath: str) -> Generator[Union[str, Path], None, None]: + def as_local_path(self, filepath: str) -> Generator[Union[str, Path], None, None]: """Download a file from ``filepath``. ``as_local_path`` is decorated by :meth:`contextlib.contextmanager`. It @@ -179,14 +169,10 @@ class HTTPStorage(Storage): os.remove(f.name) def write(self, obj: bytes, url: Union[str, Path]) -> None: - raise NotImplementedError('write is not supported by HTTP Storage') + raise NotImplementedError("write is not supported by HTTP Storage") - def write_text(self, - obj: str, - url: Union[str, Path], - encoding: str = 'utf-8') -> None: - raise NotImplementedError( - 'write_text is not supported by HTTP Storage') + def write_text(self, obj: str, url: Union[str, Path], encoding: str = "utf-8") -> None: + raise NotImplementedError("write_text is not supported by HTTP Storage") class OSSStorage(Storage): @@ -194,20 +180,16 @@ class OSSStorage(Storage): def __init__(self, oss_config_file=None): # read from config file or env var - raise NotImplementedError( - 'OSSStorage.__init__ to be implemented in the future') + raise NotImplementedError("OSSStorage.__init__ to be implemented in the future") def read(self, filepath): - raise NotImplementedError( - 'OSSStorage.read to be implemented in the future') + raise NotImplementedError("OSSStorage.read to be implemented in the future") - def read_text(self, filepath, encoding='utf-8'): - raise NotImplementedError( - 'OSSStorage.read_text to be implemented in the future') + def read_text(self, filepath, encoding="utf-8"): + raise NotImplementedError("OSSStorage.read_text to be implemented in the future") @contextlib.contextmanager - def as_local_path( - self, filepath: str) -> Generator[Union[str, Path], None, None]: + def as_local_path(self, filepath: str) -> Generator[Union[str, Path], None, None]: """Download a file from ``filepath``. ``as_local_path`` is decorated by :meth:`contextlib.contextmanager`. It @@ -233,15 +215,10 @@ class OSSStorage(Storage): os.remove(f.name) def write(self, obj: bytes, filepath: Union[str, Path]) -> None: - raise NotImplementedError( - 'OSSStorage.write to be implemented in the future') + raise NotImplementedError("OSSStorage.write to be implemented in the future") - def write_text(self, - obj: str, - filepath: Union[str, Path], - encoding: str = 'utf-8') -> None: - raise NotImplementedError( - 'OSSStorage.write_text to be implemented in the future') + def write_text(self, obj: str, filepath: Union[str, Path], encoding: str = "utf-8") -> None: + raise NotImplementedError("OSSStorage.write_text to be implemented in the future") G_STORAGES = {} @@ -249,27 +226,26 @@ G_STORAGES = {} class File(object): _prefix_to_storage: dict = { - 'oss': OSSStorage, - 'http': HTTPStorage, - 'https': HTTPStorage, - 'local': LocalStorage, + "oss": OSSStorage, + "http": HTTPStorage, + "https": HTTPStorage, + "local": LocalStorage, } @staticmethod def _get_storage(uri): - assert isinstance(uri, - str), f'uri should be str type, but got {type(uri)}' + assert isinstance(uri, str), f"uri should be str type, but got {type(uri)}" - if '://' not in uri: + if "://" not in uri: # local path - storage_type = 'local' + storage_type = "local" else: - prefix, _ = uri.split('://') + prefix, _ = uri.split("://") storage_type = prefix - assert storage_type in File._prefix_to_storage, \ - f'Unsupported uri {uri}, valid prefixs: '\ - f'{list(File._prefix_to_storage.keys())}' + assert storage_type in File._prefix_to_storage, ( + f"Unsupported uri {uri}, valid prefixs: " f"{list(File._prefix_to_storage.keys())}" + ) if storage_type not in G_STORAGES: G_STORAGES[storage_type] = File._prefix_to_storage[storage_type]() @@ -290,7 +266,7 @@ class File(object): return storage.read(uri) @staticmethod - def read_text(uri: Union[str, Path], encoding: str = 'utf-8') -> str: + def read_text(uri: Union[str, Path], encoding: str = "utf-8") -> str: """Read data from a given ``filepath`` with 'r' mode. Args: @@ -320,7 +296,7 @@ class File(object): return storage.write(obj, uri) @staticmethod - def write_text(obj: str, uri: str, encoding: str = 'utf-8') -> None: + def write_text(obj: str, uri: str, encoding: str = "utf-8") -> None: """Write data to a given ``filepath`` with 'w' mode. Note: diff --git a/funasr/download/name_maps_from_hub.py b/funasr/download/name_maps_from_hub.py index fc0084336..87a89fcc6 100644 --- a/funasr/download/name_maps_from_hub.py +++ b/funasr/download/name_maps_from_hub.py @@ -14,9 +14,7 @@ name_maps_ms = { "Qwen-Audio": "Qwen/Qwen-Audio", } -name_maps_hf = { - -} +name_maps_hf = {} name_maps_openai = { "Whisper-tiny.en": "tiny.en", @@ -31,4 +29,4 @@ name_maps_openai = { "Whisper-large-v2": "large-v2", "Whisper-large-v3": "large-v3", "Whisper-large": "large", -} \ No newline at end of file +} diff --git a/funasr/download/runtime_sdk_download_tool.py b/funasr/download/runtime_sdk_download_tool.py index 1981347da..7776a7116 100644 --- a/funasr/download/runtime_sdk_download_tool.py +++ b/funasr/download/runtime_sdk_download_tool.py @@ -7,33 +7,38 @@ from funasr.utils.types import str2bool def main(): parser = argparse.ArgumentParser() - parser.add_argument('--model-name', type=str, required=True) - parser.add_argument('--export-dir', type=str, required=True) - parser.add_argument('--export', type=str2bool, default=True, help='whether to export model') - parser.add_argument('--type', type=str, default='onnx', help='["onnx", "torch"]') - parser.add_argument('--device', type=str, default='cpu', help='["cpu", "cuda"]') - parser.add_argument('--quantize', type=str2bool, default=False, help='export quantized model') - parser.add_argument('--fallback-num', type=int, default=0, help='amp fallback number') - parser.add_argument('--audio_in', type=str, default=None, help='["wav", "wav.scp"]') - parser.add_argument('--model_revision', type=str, default=None, help='model_revision') - parser.add_argument('--calib_num', type=int, default=200, help='calib max num') + parser.add_argument("--model-name", type=str, required=True) + parser.add_argument("--export-dir", type=str, required=True) + parser.add_argument("--export", type=str2bool, default=True, help="whether to export model") + parser.add_argument("--type", type=str, default="onnx", help='["onnx", "torch"]') + parser.add_argument("--device", type=str, default="cpu", help='["cpu", "cuda"]') + parser.add_argument("--quantize", type=str2bool, default=False, help="export quantized model") + parser.add_argument("--fallback-num", type=int, default=0, help="amp fallback number") + parser.add_argument("--audio_in", type=str, default=None, help='["wav", "wav.scp"]') + parser.add_argument("--model_revision", type=str, default=None, help="model_revision") + parser.add_argument("--calib_num", type=int, default=200, help="calib max num") args = parser.parse_args() - + model_dir = args.model_name if not Path(args.model_name).exists(): from modelscope.hub.snapshot_download import snapshot_download + try: - model_dir = snapshot_download(args.model_name, cache_dir=args.export_dir, revision=args.model_revision) + model_dir = snapshot_download( + args.model_name, cache_dir=args.export_dir, revision=args.model_revision + ) except: - raise "model_dir must be model_name in modelscope or local path downloaded from modelscope, but is {}".format \ - (model_dir) + raise "model_dir must be model_name in modelscope or local path downloaded from modelscope, but is {}".format( + model_dir + ) if args.export: - model_file = os.path.join(model_dir, 'model.onnx') + model_file = os.path.join(model_dir, "model.onnx") if args.quantize: - model_file = os.path.join(model_dir, 'model_quant.onnx') + model_file = os.path.join(model_dir, "model_quant.onnx") if not os.path.exists(model_file): print(".onnx is not exist, begin to export onnx") from funasr.bin.export_model import ModelExport + export_model = ModelExport( cache_dir=args.export_dir, onnx=True, @@ -44,4 +49,4 @@ def main(): if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/funasr/frontends/default.py b/funasr/frontends/default.py index c4bdbd774..462da8633 100644 --- a/funasr/frontends/default.py +++ b/funasr/frontends/default.py @@ -6,6 +6,7 @@ import logging import numpy as np import torch import torch.nn as nn + try: from torch_complex.tensor import ComplexTensor except: @@ -25,23 +26,23 @@ class DefaultFrontend(nn.Module): """ def __init__( - self, - fs: int = 16000, - n_fft: int = 512, - win_length: int = None, - hop_length: int = 128, - window: Optional[str] = "hann", - center: bool = True, - normalized: bool = False, - onesided: bool = True, - n_mels: int = 80, - fmin: int = None, - fmax: int = None, - htk: bool = False, - frontend_conf: Optional[dict] = None, - apply_stft: bool = True, - use_channel: int = None, - **kwargs, + self, + fs: int = 16000, + n_fft: int = 512, + win_length: int = None, + hop_length: int = 128, + window: Optional[str] = "hann", + center: bool = True, + normalized: bool = False, + onesided: bool = True, + n_mels: int = 80, + fmin: int = None, + fmax: int = None, + htk: bool = False, + frontend_conf: Optional[dict] = None, + apply_stft: bool = True, + use_channel: int = None, + **kwargs, ): super().__init__() @@ -85,11 +86,11 @@ class DefaultFrontend(nn.Module): return self.n_mels def forward( - self, input: torch.Tensor, input_lengths: Union[torch.Tensor, list] + self, input: torch.Tensor, input_lengths: Union[torch.Tensor, list] ) -> Tuple[torch.Tensor, torch.Tensor]: if isinstance(input_lengths, list): input_lengths = torch.tensor(input_lengths) - if input.dtype == torch.float64: + if input.dtype == torch.float64: input = input.float() # 1. Domain-conversion: e.g. Stft: time -> time-freq if self.stft is not None: @@ -119,7 +120,7 @@ class DefaultFrontend(nn.Module): # 4. STFT -> Power spectrum # h: ComplexTensor(B, T, F) -> torch.Tensor(B, T, F) - input_power = input_stft.real ** 2 + input_stft.imag ** 2 + input_power = input_stft.real**2 + input_stft.imag**2 # 5. Feature transform e.g. Stft -> Log-Mel-Fbank # input_power: (Batch, [Channel,] Length, Freq) @@ -128,9 +129,7 @@ class DefaultFrontend(nn.Module): return input_feats, feats_lens - def _compute_stft( - self, input: torch.Tensor, input_lengths: torch.Tensor - ) -> torch.Tensor: + def _compute_stft(self, input: torch.Tensor, input_lengths: torch.Tensor) -> torch.Tensor: input_stft, feats_lens = self.stft(input, input_lengths) assert input_stft.dim() >= 4, input_stft.shape @@ -149,28 +148,28 @@ class MultiChannelFrontend(nn.Module): """ def __init__( - self, - fs: int = 16000, - n_fft: int = 512, - win_length: int = None, - hop_length: int = None, - frame_length: int = None, - frame_shift: int = None, - window: Optional[str] = "hann", - center: bool = True, - normalized: bool = False, - onesided: bool = True, - n_mels: int = 80, - fmin: int = None, - fmax: int = None, - htk: bool = False, - frontend_conf: Optional[dict] = None, - apply_stft: bool = True, - use_channel: int = None, - lfr_m: int = 1, - lfr_n: int = 1, - cmvn_file: str = None, - mc: bool = True + self, + fs: int = 16000, + n_fft: int = 512, + win_length: int = None, + hop_length: int = None, + frame_length: int = None, + frame_shift: int = None, + window: Optional[str] = "hann", + center: bool = True, + normalized: bool = False, + onesided: bool = True, + n_mels: int = 80, + fmin: int = None, + fmax: int = None, + htk: bool = False, + frontend_conf: Optional[dict] = None, + apply_stft: bool = True, + use_channel: int = None, + lfr_m: int = 1, + lfr_n: int = 1, + cmvn_file: str = None, + mc: bool = True, ): super().__init__() # Deepcopy (In general, dict shouldn't be used as default arg) @@ -183,8 +182,7 @@ class MultiChannelFrontend(nn.Module): self.hop_length = self.hop_length else: logging.error( - "Only one of (win_length, hop_length) and (frame_length, frame_shift)" - "can be set." + "Only one of (win_length, hop_length) and (frame_length, frame_shift)" "can be set." ) exit(1) @@ -234,10 +232,10 @@ class MultiChannelFrontend(nn.Module): return self.n_mels def forward( - self, input: torch.Tensor, input_lengths: torch.Tensor + self, input: torch.Tensor, input_lengths: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: # 1. Domain-conversion: e.g. Stft: time -> time-freq - #import pdb;pdb.set_trace() + # import pdb;pdb.set_trace() if self.stft is not None: input_stft, feats_lens = self._compute_stft(input, input_lengths) else: @@ -255,7 +253,7 @@ class MultiChannelFrontend(nn.Module): if self.training: if self.use_channel is not None: input_stft = input_stft[:, :, self.use_channel, :] - + else: # Select 1ch randomly ch = np.random.randint(input_stft.size(2)) @@ -266,7 +264,7 @@ class MultiChannelFrontend(nn.Module): # 4. STFT -> Power spectrum # h: ComplexTensor(B, T, F) -> torch.Tensor(B, T, F) - input_power = input_stft.real ** 2 + input_stft.imag ** 2 + input_power = input_stft.real**2 + input_stft.imag**2 # 5. Feature transform e.g. Stft -> Log-Mel-Fbank # input_power: (Batch, [Channel,] Length, Freq) @@ -274,11 +272,13 @@ class MultiChannelFrontend(nn.Module): input_feats, _ = self.logmel(input_power, feats_lens) if self.mc: # MFCCA - if input_feats.dim() ==4: + if input_feats.dim() == 4: bt = input_feats.size(0) channel_size = input_feats.size(2) - input_feats = input_feats.transpose(1,2).reshape(bt*channel_size,-1,80).contiguous() - feats_lens = feats_lens.repeat(1,channel_size).squeeze() + input_feats = ( + input_feats.transpose(1, 2).reshape(bt * channel_size, -1, 80).contiguous() + ) + feats_lens = feats_lens.repeat(1, channel_size).squeeze() else: channel_size = 1 return input_feats, feats_lens, channel_size @@ -304,9 +304,7 @@ class MultiChannelFrontend(nn.Module): return input_feats, feats_lens - def _compute_stft( - self, input: torch.Tensor, input_lengths: torch.Tensor - ) -> torch.Tensor: + def _compute_stft(self, input: torch.Tensor, input_lengths: torch.Tensor) -> torch.Tensor: input_stft, feats_lens = self.stft(input, input_lengths) assert input_stft.dim() >= 4, input_stft.shape @@ -319,22 +317,22 @@ class MultiChannelFrontend(nn.Module): return input_stft, feats_lens def _load_cmvn(self, cmvn_file): - with open(cmvn_file, 'r', encoding='utf-8') as f: + with open(cmvn_file, "r", encoding="utf-8") as f: lines = f.readlines() means_list = [] vars_list = [] for i in range(len(lines)): line_item = lines[i].split() - if line_item[0] == '': + if line_item[0] == "": line_item = lines[i + 1].split() - if line_item[0] == '': - add_shift_line = line_item[3:(len(line_item) - 1)] + if line_item[0] == "": + add_shift_line = line_item[3 : (len(line_item) - 1)] means_list = list(add_shift_line) continue - elif line_item[0] == '': + elif line_item[0] == "": line_item = lines[i + 1].split() - if line_item[0] == '': - rescale_line = line_item[3:(len(line_item) - 1)] + if line_item[0] == "": + rescale_line = line_item[3 : (len(line_item) - 1)] vars_list = list(rescale_line) continue means = np.array(means_list).astype(np.float) diff --git a/funasr/frontends/eend_ola_feature.py b/funasr/frontends/eend_ola_feature.py index e15b71c25..0c79a3bb8 100644 --- a/funasr/frontends/eend_ola_feature.py +++ b/funasr/frontends/eend_ola_feature.py @@ -13,7 +13,7 @@ def transform(Y, dtype=np.float32): sr = 8000 n_mels = 23 mel_basis = librosa.filters.mel(sr, n_fft, n_mels) - Y = np.dot(Y ** 2, mel_basis.T) + Y = np.dot(Y**2, mel_basis.T) Y = np.log10(np.maximum(Y, 1e-10)) mean = np.mean(Y, axis=0) Y = Y - mean @@ -27,25 +27,21 @@ def subsample(Y, T, subsampling=1): def splice(Y, context_size=0): - Y_pad = np.pad( - Y, - [(context_size, context_size), (0, 0)], - 'constant') + Y_pad = np.pad(Y, [(context_size, context_size), (0, 0)], "constant") Y_spliced = np.lib.stride_tricks.as_strided( np.ascontiguousarray(Y_pad), (Y.shape[0], Y.shape[1] * (2 * context_size + 1)), - (Y.itemsize * Y.shape[1], Y.itemsize), writeable=False) + (Y.itemsize * Y.shape[1], Y.itemsize), + writeable=False, + ) return Y_spliced -def stft( - data, - frame_size=1024, - frame_shift=256): +def stft(data, frame_size=1024, frame_shift=256): fft_size = 1 << (frame_size - 1).bit_length() if len(data) % frame_shift == 0: - return librosa.stft(data, n_fft=fft_size, win_length=frame_size, - hop_length=frame_shift).T[:-1] + return librosa.stft(data, n_fft=fft_size, win_length=frame_size, hop_length=frame_shift).T[ + :-1 + ] else: - return librosa.stft(data, n_fft=fft_size, win_length=frame_size, - hop_length=frame_shift).T \ No newline at end of file + return librosa.stft(data, n_fft=fft_size, win_length=frame_size, hop_length=frame_shift).T diff --git a/funasr/frontends/fused.py b/funasr/frontends/fused.py index 24f73f4c2..0fa763900 100644 --- a/funasr/frontends/fused.py +++ b/funasr/frontends/fused.py @@ -7,14 +7,10 @@ from typing import Tuple class FusedFrontends(nn.Module): - def __init__( - self, frontends=None, align_method="linear_projection", proj_dim=100, fs=16000 - ): + def __init__(self, frontends=None, align_method="linear_projection", proj_dim=100, fs=16000): super().__init__() - self.align_method = ( - align_method # fusing method : linear_projection only for now - ) + self.align_method = align_method # fusing method : linear_projection only for now self.proj_dim = proj_dim # dim of the projection done on each frontend self.frontends = [] # list of the frontends to combine @@ -109,9 +105,7 @@ class FusedFrontends(nn.Module): input_feats, feats_lens = frontend.forward(input, input_lengths) self.feats.append([input_feats, feats_lens]) - if ( - self.align_method == "linear_projection" - ): # TODO(Dan): to add other align methods + if self.align_method == "linear_projection": # TODO(Dan): to add other align methods # first step : projections self.feats_proj = [] @@ -141,4 +135,4 @@ class FusedFrontends(nn.Module): else: raise NotImplementedError - return input_feats, feats_lens \ No newline at end of file + return input_feats, feats_lens diff --git a/funasr/frontends/s3prl.py b/funasr/frontends/s3prl.py index ff60592ca..5c358cca6 100644 --- a/funasr/frontends/s3prl.py +++ b/funasr/frontends/s3prl.py @@ -29,11 +29,11 @@ class S3prlFrontend(nn.Module): """Speech Pretrained Representation frontend structure for ASR.""" def __init__( - self, - fs: Union[int, str] = 16000, - frontend_conf: Optional[dict] = None, - download_dir: str = None, - multilayer_feature: bool = False, + self, + fs: Union[int, str] = 16000, + frontend_conf: Optional[dict] = None, + download_dir: str = None, + multilayer_feature: bool = False, ): super().__init__() if isinstance(fs, str): @@ -74,7 +74,7 @@ class S3prlFrontend(nn.Module): ).to("cpu") if getattr( - s3prl_upstream, "model", None + s3prl_upstream, "model", None ) is not None and s3prl_upstream.model.__class__.__name__ in [ "Wav2Vec2Model", "HubertModel", @@ -102,9 +102,9 @@ class S3prlFrontend(nn.Module): Output - sequence of tiled representations shape: (batch_size, seq_len * factor, feature_dim) """ - assert ( - len(feature.shape) == 3 - ), "Input argument `feature` has invalid shape: {}".format(feature.shape) + assert len(feature.shape) == 3, "Input argument `feature` has invalid shape: {}".format( + feature.shape + ) tiled_feature = feature.repeat(1, 1, self.args.tile_factor) tiled_feature = tiled_feature.reshape( feature.size(0), feature.size(1) * self.args.tile_factor, feature.size(2) @@ -115,7 +115,7 @@ class S3prlFrontend(nn.Module): return self.output_dim def forward( - self, input: torch.Tensor, input_lengths: torch.Tensor + self, input: torch.Tensor, input_lengths: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: wavs = [wav[: input_lengths[i]] for i, wav in enumerate(input)] self.upstream.eval() diff --git a/funasr/frontends/utils/beamformer.py b/funasr/frontends/utils/beamformer.py index f3eccee4c..d19abe0ab 100644 --- a/funasr/frontends/utils/beamformer.py +++ b/funasr/frontends/utils/beamformer.py @@ -76,9 +76,7 @@ def get_mvdr_vector( return beamform_vector -def apply_beamforming_vector( - beamform_vector: ComplexTensor, mix: ComplexTensor -) -> ComplexTensor: +def apply_beamforming_vector(beamform_vector: ComplexTensor, mix: ComplexTensor) -> ComplexTensor: # (..., C) x (..., C, T) -> (..., T) es = FC.einsum("...c,...ct->...t", [beamform_vector.conj(), mix]) return es diff --git a/funasr/frontends/utils/complex_utils.py b/funasr/frontends/utils/complex_utils.py index 5d313c6ac..06c69363b 100644 --- a/funasr/frontends/utils/complex_utils.py +++ b/funasr/frontends/utils/complex_utils.py @@ -1,10 +1,12 @@ """Beamformer module.""" + from distutils.version import LooseVersion from typing import Sequence from typing import Tuple from typing import Union import torch + try: from torch_complex import functional as FC from torch_complex.tensor import ComplexTensor @@ -12,7 +14,6 @@ except: print("Please install torch_complex firstly") - EPS = torch.finfo(torch.double).eps is_torch_1_8_plus = LooseVersion(torch.__version__) >= LooseVersion("1.8.0") is_torch_1_9_plus = LooseVersion(torch.__version__) >= LooseVersion("1.9.0") @@ -27,15 +28,11 @@ def new_complex_like( elif is_torch_complex_tensor(ref): return torch.complex(*real_imag) else: - raise ValueError( - "Please update your PyTorch version to 1.9+ for complex support." - ) + raise ValueError("Please update your PyTorch version to 1.9+ for complex support.") def is_torch_complex_tensor(c): - return ( - not isinstance(c, ComplexTensor) and is_torch_1_9_plus and torch.is_complex(c) - ) + return not isinstance(c, ComplexTensor) and is_torch_1_9_plus and torch.is_complex(c) def is_complex(c): @@ -59,8 +56,7 @@ def to_float(c): def cat(seq: Sequence[Union[ComplexTensor, torch.Tensor]], *args, **kwargs): if not isinstance(seq, (list, tuple)): raise TypeError( - "cat(): argument 'tensors' (position 1) must be tuple of Tensors, " - "not Tensor" + "cat(): argument 'tensors' (position 1) must be tuple of Tensors, " "not Tensor" ) if isinstance(seq[0], ComplexTensor): return FC.cat(seq, *args, **kwargs) @@ -68,17 +64,13 @@ def cat(seq: Sequence[Union[ComplexTensor, torch.Tensor]], *args, **kwargs): return torch.cat(seq, *args, **kwargs) -def complex_norm( - c: Union[torch.Tensor, ComplexTensor], dim=-1, keepdim=False -) -> torch.Tensor: +def complex_norm(c: Union[torch.Tensor, ComplexTensor], dim=-1, keepdim=False) -> torch.Tensor: if not is_complex(c): raise TypeError("Input is not a complex tensor.") if is_torch_complex_tensor(c): return torch.norm(c, dim=dim, keepdim=keepdim) else: - return torch.sqrt( - (c.real**2 + c.imag**2).sum(dim=dim, keepdim=keepdim) + EPS - ) + return torch.sqrt((c.real**2 + c.imag**2).sum(dim=dim, keepdim=keepdim) + EPS) def einsum(equation, *operands): @@ -116,9 +108,7 @@ def einsum(equation, *operands): return torch.einsum(equation, a, b) -def inverse( - c: Union[torch.Tensor, ComplexTensor] -) -> Union[torch.Tensor, ComplexTensor]: +def inverse(c: Union[torch.Tensor, ComplexTensor]) -> Union[torch.Tensor, ComplexTensor]: if isinstance(c, ComplexTensor): return c.inverse2() else: @@ -186,8 +176,7 @@ def solve(b: Union[torch.Tensor, ComplexTensor], a: Union[torch.Tensor, ComplexT def stack(seq: Sequence[Union[ComplexTensor, torch.Tensor]], *args, **kwargs): if not isinstance(seq, (list, tuple)): raise TypeError( - "stack(): argument 'tensors' (position 1) must be tuple of Tensors, " - "not Tensor" + "stack(): argument 'tensors' (position 1) must be tuple of Tensors, " "not Tensor" ) if isinstance(seq[0], ComplexTensor): return FC.stack(seq, *args, **kwargs) diff --git a/funasr/frontends/utils/dnn_beamformer.py b/funasr/frontends/utils/dnn_beamformer.py index 75637d287..73b6bea42 100644 --- a/funasr/frontends/utils/dnn_beamformer.py +++ b/funasr/frontends/utils/dnn_beamformer.py @@ -1,4 +1,5 @@ """DNN beamformer module.""" + from typing import Tuple import torch @@ -36,18 +37,14 @@ class DNN_Beamformer(torch.nn.Module): beamformer_type="mvdr", ): super().__init__() - self.mask = MaskEstimator( - btype, bidim, blayers, bunits, bprojs, dropout_rate, nmask=bnmask - ) + self.mask = MaskEstimator(btype, bidim, blayers, bunits, bprojs, dropout_rate, nmask=bnmask) self.ref = AttentionReference(bidim, badim) self.ref_channel = ref_channel self.nmask = bnmask if beamformer_type != "mvdr": - raise ValueError( - "Not supporting beamformer_type={}".format(beamformer_type) - ) + raise ValueError("Not supporting beamformer_type={}".format(beamformer_type)) self.beamformer_type = beamformer_type def forward( @@ -76,9 +73,7 @@ class DNN_Beamformer(torch.nn.Module): u, _ = self.ref(psd_speech, ilens) else: # (optional) Create onehot vector for fixed reference microphone - u = torch.zeros( - *(data.size()[:-3] + (data.size(-2),)), device=data.device - ) + u = torch.zeros(*(data.size()[:-3] + (data.size(-2),)), device=data.device) u[..., self.ref_channel].fill_(1) ws = get_mvdr_vector(psd_speech, psd_noise, u) @@ -108,9 +103,7 @@ class DNN_Beamformer(torch.nn.Module): mask_speech = list(masks[:-1]) mask_noise = masks[-1] - psd_speeches = [ - get_power_spectral_density_matrix(data, mask) for mask in mask_speech - ] + psd_speeches = [get_power_spectral_density_matrix(data, mask) for mask in mask_speech] psd_noise = get_power_spectral_density_matrix(data, mask_noise) enhanced = [] @@ -118,9 +111,7 @@ class DNN_Beamformer(torch.nn.Module): for i in range(self.nmask - 1): psd_speech = psd_speeches.pop(i) # treat all other speakers' psd_speech as noises - enh, w = apply_beamforming( - data, ilens, psd_speech, sum(psd_speeches) + psd_noise - ) + enh, w = apply_beamforming(data, ilens, psd_speech, sum(psd_speeches) + psd_noise) psd_speeches.insert(i, psd_speech) # (..., F, T) -> (..., T, F) @@ -155,9 +146,7 @@ class AttentionReference(torch.nn.Module): B, _, C = psd_in.size()[:3] assert psd_in.size(2) == psd_in.size(3), psd_in.size() # psd_in: (B, F, C, C) - psd = psd_in.masked_fill( - torch.eye(C, dtype=torch.bool, device=psd_in.device), 0 - ) + psd = psd_in.masked_fill(torch.eye(C, dtype=torch.bool, device=psd_in.device), 0) # psd: (B, F, C, C) -> (B, C, F) psd = (psd.sum(dim=-1) / (C - 1)).transpose(-1, -2) diff --git a/funasr/frontends/utils/feature_transform.py b/funasr/frontends/utils/feature_transform.py index a3ee26020..ab1486be9 100644 --- a/funasr/frontends/utils/feature_transform.py +++ b/funasr/frontends/utils/feature_transform.py @@ -36,9 +36,7 @@ class FeatureTransform(torch.nn.Module): self.global_mvn = None if self.apply_uttmvn is not None: - self.uttmvn = UtteranceMVN( - norm_means=uttmvn_norm_means, norm_vars=uttmvn_norm_vars - ) + self.uttmvn = UtteranceMVN(norm_means=uttmvn_norm_means, norm_vars=uttmvn_norm_vars) else: self.uttmvn = None @@ -191,9 +189,7 @@ class GlobalMVN(torch.nn.Module): class UtteranceMVN(torch.nn.Module): - def __init__( - self, norm_means: bool = True, norm_vars: bool = False, eps: float = 1.0e-20 - ): + def __init__(self, norm_means: bool = True, norm_vars: bool = False, eps: float = 1.0e-20): super().__init__() self.norm_means = norm_means self.norm_vars = norm_vars diff --git a/funasr/frontends/utils/log_mel.py b/funasr/frontends/utils/log_mel.py index 4983a1d0a..611cf1eb3 100644 --- a/funasr/frontends/utils/log_mel.py +++ b/funasr/frontends/utils/log_mel.py @@ -73,11 +73,7 @@ class LogMel(torch.nn.Module): # Zero padding if ilens is not None: - logmel_feat = logmel_feat.masked_fill( - make_pad_mask(ilens, logmel_feat, 1), 0.0 - ) + logmel_feat = logmel_feat.masked_fill(make_pad_mask(ilens, logmel_feat, 1), 0.0) else: - ilens = feat.new_full( - [feat.size(0)], fill_value=feat.size(1), dtype=torch.long - ) + ilens = feat.new_full([feat.size(0)], fill_value=feat.size(1), dtype=torch.long) return logmel_feat, ilens diff --git a/funasr/frontends/utils/mask_estimator.py b/funasr/frontends/utils/mask_estimator.py index 002e474cc..29a21d97a 100644 --- a/funasr/frontends/utils/mask_estimator.py +++ b/funasr/frontends/utils/mask_estimator.py @@ -23,9 +23,7 @@ class MaskEstimator(torch.nn.Module): self.type = type self.nmask = nmask - self.linears = torch.nn.ModuleList( - [torch.nn.Linear(projs, idim) for _ in range(nmask)] - ) + self.linears = torch.nn.ModuleList([torch.nn.Linear(projs, idim) for _ in range(nmask)]) def forward( self, xs: ComplexTensor, ilens: torch.LongTensor diff --git a/funasr/frontends/utils/stft.py b/funasr/frontends/utils/stft.py index 00d9ec5ad..381f1e986 100644 --- a/funasr/frontends/utils/stft.py +++ b/funasr/frontends/utils/stft.py @@ -86,13 +86,12 @@ class Stft(torch.nn.Module): # or (Batch, Channel, Freq, Frames, 2=real_imag) if self.window is not None: if self.window.lower() == "povey": - window = torch.hann_window(self.win_length, periodic=False, - device=input.device, dtype=input.dtype).pow(0.85) + window = torch.hann_window( + self.win_length, periodic=False, device=input.device, dtype=input.dtype + ).pow(0.85) else: window_func = getattr(torch, f"{self.window}_window") - window = window_func( - self.win_length, dtype=input.dtype, device=input.device - ) + window = window_func(self.win_length, dtype=input.dtype, device=input.device) else: window = None @@ -135,9 +134,7 @@ class Stft(torch.nn.Module): [torch.zeros(n_pad_left), window, torch.zeros(n_pad_right)], 0 ).numpy() else: - win_length = ( - self.win_length if self.win_length is not None else self.n_fft - ) + win_length = self.win_length if self.win_length is not None else self.n_fft stft_kwargs["window"] = torch.ones(win_length) output = [] @@ -160,9 +157,7 @@ class Stft(torch.nn.Module): if multi_channel: # output: (Batch * Channel, Frames, Freq, 2=real_imag) # -> (Batch, Frame, Channel, Freq, 2=real_imag) - output = output.view(bs, -1, output.size(1), output.size(2), 2).transpose( - 1, 2 - ) + output = output.view(bs, -1, output.size(1), output.size(2), 2).transpose(1, 2) if ilens is not None: if self.center: @@ -194,14 +189,10 @@ class Stft(torch.nn.Module): try: import torchaudio except ImportError: - raise ImportError( - "Please install torchaudio>=0.3.0 or use torch>=1.6.0" - ) + raise ImportError("Please install torchaudio>=0.3.0 or use torch>=1.6.0") if not hasattr(torchaudio.functional, "istft"): - raise ImportError( - "Please install torchaudio>=0.3.0 or use torch>=1.6.0" - ) + raise ImportError("Please install torchaudio>=0.3.0 or use torch>=1.6.0") istft = torchaudio.functional.istft if self.window is not None: diff --git a/funasr/frontends/wav_frontend.py b/funasr/frontends/wav_frontend.py index afa7421ca..a4002df5b 100644 --- a/funasr/frontends/wav_frontend.py +++ b/funasr/frontends/wav_frontend.py @@ -12,24 +12,23 @@ import funasr.frontends.eend_ola_feature as eend_ola_feature from funasr.register import tables - def load_cmvn(cmvn_file): - with open(cmvn_file, 'r', encoding='utf-8') as f: + with open(cmvn_file, "r", encoding="utf-8") as f: lines = f.readlines() means_list = [] vars_list = [] for i in range(len(lines)): line_item = lines[i].split() - if line_item[0] == '': + if line_item[0] == "": line_item = lines[i + 1].split() - if line_item[0] == '': - add_shift_line = line_item[3:(len(line_item) - 1)] + if line_item[0] == "": + add_shift_line = line_item[3 : (len(line_item) - 1)] means_list = list(add_shift_line) continue - elif line_item[0] == '': + elif line_item[0] == "": line_item = lines[i + 1].split() - if line_item[0] == '': - rescale_line = line_item[3:(len(line_item) - 1)] + if line_item[0] == "": + rescale_line = line_item[3 : (len(line_item) - 1)] vars_list = list(rescale_line) continue means = np.array(means_list).astype(np.float32) @@ -65,38 +64,38 @@ def apply_lfr(inputs, lfr_m, lfr_n): T = T + (lfr_m - 1) // 2 for i in range(T_lfr): if lfr_m <= T - i * lfr_n: - LFR_inputs.append((inputs[i * lfr_n:i * lfr_n + lfr_m]).view(1, -1)) + LFR_inputs.append((inputs[i * lfr_n : i * lfr_n + lfr_m]).view(1, -1)) else: # process last LFR frame num_padding = lfr_m - (T - i * lfr_n) - frame = (inputs[i * lfr_n:]).view(-1) + frame = (inputs[i * lfr_n :]).view(-1) for _ in range(num_padding): frame = torch.hstack((frame, inputs[-1])) LFR_inputs.append(frame) LFR_outputs = torch.vstack(LFR_inputs) return LFR_outputs.type(torch.float32) + @tables.register("frontend_classes", "wav_frontend") @tables.register("frontend_classes", "WavFrontend") class WavFrontend(nn.Module): - """Conventional frontend structure for ASR. - """ + """Conventional frontend structure for ASR.""" def __init__( - self, - cmvn_file: str = None, - fs: int = 16000, - window: str = 'hamming', - n_mels: int = 80, - frame_length: int = 25, - frame_shift: int = 10, - filter_length_min: int = -1, - filter_length_max: int = -1, - lfr_m: int = 1, - lfr_n: int = 1, - dither: float = 1.0, - snip_edges: bool = True, - upsacle_samples: bool = True, - **kwargs, + self, + cmvn_file: str = None, + fs: int = 16000, + window: str = "hamming", + n_mels: int = 80, + frame_length: int = 25, + frame_shift: int = 10, + filter_length_min: int = -1, + filter_length_max: int = -1, + lfr_m: int = 1, + lfr_n: int = 1, + dither: float = 1.0, + snip_edges: bool = True, + upsacle_samples: bool = True, + **kwargs, ): super().__init__() self.fs = fs @@ -118,10 +117,10 @@ class WavFrontend(nn.Module): return self.n_mels * self.lfr_m def forward( - self, - input: torch.Tensor, - input_lengths, - **kwargs, + self, + input: torch.Tensor, + input_lengths, + **kwargs, ) -> Tuple[torch.Tensor, torch.Tensor]: batch_size = input.size(0) feats = [] @@ -132,15 +131,17 @@ class WavFrontend(nn.Module): if self.upsacle_samples: waveform = waveform * (1 << 15) waveform = waveform.unsqueeze(0) - mat = kaldi.fbank(waveform, - num_mel_bins=self.n_mels, - frame_length=self.frame_length, - frame_shift=self.frame_shift, - dither=self.dither, - energy_floor=0.0, - window_type=self.window, - sample_frequency=self.fs, - snip_edges=self.snip_edges) + mat = kaldi.fbank( + waveform, + num_mel_bins=self.n_mels, + frame_length=self.frame_length, + frame_shift=self.frame_shift, + dither=self.dither, + energy_floor=0.0, + window_type=self.window, + sample_frequency=self.fs, + snip_edges=self.snip_edges, + ) if self.lfr_m != 1 or self.lfr_n != 1: mat = apply_lfr(mat, self.lfr_m, self.lfr_n) @@ -154,15 +155,12 @@ class WavFrontend(nn.Module): if batch_size == 1: feats_pad = feats[0][None, :, :] else: - feats_pad = pad_sequence(feats, - batch_first=True, - padding_value=0.0) + feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0) return feats_pad, feats_lens def forward_fbank( - self, - input: torch.Tensor, - input_lengths: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + self, input: torch.Tensor, input_lengths: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: batch_size = input.size(0) feats = [] feats_lens = [] @@ -171,34 +169,33 @@ class WavFrontend(nn.Module): waveform = input[i][:waveform_length] waveform = waveform * (1 << 15) waveform = waveform.unsqueeze(0) - mat = kaldi.fbank(waveform, - num_mel_bins=self.n_mels, - frame_length=self.frame_length, - frame_shift=self.frame_shift, - dither=self.dither, - energy_floor=0.0, - window_type=self.window, - sample_frequency=self.fs) + mat = kaldi.fbank( + waveform, + num_mel_bins=self.n_mels, + frame_length=self.frame_length, + frame_shift=self.frame_shift, + dither=self.dither, + energy_floor=0.0, + window_type=self.window, + sample_frequency=self.fs, + ) feat_length = mat.size(0) feats.append(mat) feats_lens.append(feat_length) feats_lens = torch.as_tensor(feats_lens) - feats_pad = pad_sequence(feats, - batch_first=True, - padding_value=0.0) + feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0) return feats_pad, feats_lens def forward_lfr_cmvn( - self, - input: torch.Tensor, - input_lengths: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + self, input: torch.Tensor, input_lengths: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: batch_size = input.size(0) feats = [] feats_lens = [] for i in range(batch_size): - mat = input[i, :input_lengths[i], :] + mat = input[i, : input_lengths[i], :] if self.lfr_m != 1 or self.lfr_n != 1: mat = apply_lfr(mat, self.lfr_m, self.lfr_n) if self.cmvn is not None: @@ -208,33 +205,30 @@ class WavFrontend(nn.Module): feats_lens.append(feat_length) feats_lens = torch.as_tensor(feats_lens) - feats_pad = pad_sequence(feats, - batch_first=True, - padding_value=0.0) + feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0) return feats_pad, feats_lens @tables.register("frontend_classes", "WavFrontendOnline") class WavFrontendOnline(nn.Module): - """Conventional frontend structure for streaming ASR/VAD. - """ + """Conventional frontend structure for streaming ASR/VAD.""" def __init__( - self, - cmvn_file: str = None, - fs: int = 16000, - window: str = 'hamming', - n_mels: int = 80, - frame_length: int = 25, - frame_shift: int = 10, - filter_length_min: int = -1, - filter_length_max: int = -1, - lfr_m: int = 1, - lfr_n: int = 1, - dither: float = 1.0, - snip_edges: bool = True, - upsacle_samples: bool = True, - **kwargs, + self, + cmvn_file: str = None, + fs: int = 16000, + window: str = "hamming", + n_mels: int = 80, + frame_length: int = 25, + frame_shift: int = 10, + filter_length_min: int = -1, + filter_length_max: int = -1, + lfr_m: int = 1, + lfr_n: int = 1, + dither: float = 1.0, + snip_edges: bool = True, + upsacle_samples: bool = True, + **kwargs, ): super().__init__() self.fs = fs @@ -281,8 +275,9 @@ class WavFrontendOnline(nn.Module): return inputs.type(torch.float32) @staticmethod - def apply_lfr(inputs: torch.Tensor, lfr_m: int, lfr_n: int, is_final: bool = False) -> Tuple[ - torch.Tensor, torch.Tensor, int]: + def apply_lfr( + inputs: torch.Tensor, lfr_m: int, lfr_n: int, is_final: bool = False + ) -> Tuple[torch.Tensor, torch.Tensor, int]: """ Apply lfr with data """ @@ -290,15 +285,17 @@ class WavFrontendOnline(nn.Module): LFR_inputs = [] # inputs = torch.vstack((inputs_lfr_cache, inputs)) T = inputs.shape[0] # include the right context - T_lfr = int(np.ceil((T - (lfr_m - 1) // 2) / lfr_n)) # minus the right context: (lfr_m - 1) // 2 + T_lfr = int( + np.ceil((T - (lfr_m - 1) // 2) / lfr_n) + ) # minus the right context: (lfr_m - 1) // 2 splice_idx = T_lfr for i in range(T_lfr): if lfr_m <= T - i * lfr_n: - LFR_inputs.append((inputs[i * lfr_n:i * lfr_n + lfr_m]).view(1, -1)) + LFR_inputs.append((inputs[i * lfr_n : i * lfr_n + lfr_m]).view(1, -1)) else: # process last LFR frame if is_final: num_padding = lfr_m - (T - i * lfr_n) - frame = (inputs[i * lfr_n:]).view(-1) + frame = (inputs[i * lfr_n :]).view(-1) for _ in range(num_padding): frame = torch.hstack((frame, inputs[-1])) LFR_inputs.append(frame) @@ -312,23 +309,29 @@ class WavFrontendOnline(nn.Module): return LFR_outputs.type(torch.float32), lfr_splice_cache, splice_idx @staticmethod - def compute_frame_num(sample_length: int, frame_sample_length: int, frame_shift_sample_length: int) -> int: + def compute_frame_num( + sample_length: int, frame_sample_length: int, frame_shift_sample_length: int + ) -> int: frame_num = int((sample_length - frame_sample_length) / frame_shift_sample_length + 1) return frame_num if frame_num >= 1 and sample_length >= frame_sample_length else 0 def forward_fbank( - self, - input: torch.Tensor, - input_lengths: torch.Tensor, - cache: dict = {}, - **kwargs, + self, + input: torch.Tensor, + input_lengths: torch.Tensor, + cache: dict = {}, + **kwargs, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: batch_size = input.size(0) input = torch.cat((cache["input_cache"], input), dim=1) - frame_num = self.compute_frame_num(input.shape[-1], self.frame_sample_length, self.frame_shift_sample_length) + frame_num = self.compute_frame_num( + input.shape[-1], self.frame_sample_length, self.frame_shift_sample_length + ) # update self.in_cache - cache["input_cache"] = input[:, -(input.shape[-1] - frame_num * self.frame_shift_sample_length):] + cache["input_cache"] = input[ + :, -(input.shape[-1] - frame_num * self.frame_shift_sample_length) : + ] waveforms = torch.empty(0) feats_pad = torch.empty(0) feats_lens = torch.empty(0) @@ -340,17 +343,25 @@ class WavFrontendOnline(nn.Module): waveform = input[i] # we need accurate wave samples that used for fbank extracting waveforms.append( - waveform[:((frame_num - 1) * self.frame_shift_sample_length + self.frame_sample_length)]) + waveform[ + : ( + (frame_num - 1) * self.frame_shift_sample_length + + self.frame_sample_length + ) + ] + ) waveform = waveform * (1 << 15) waveform = waveform.unsqueeze(0) - mat = kaldi.fbank(waveform, - num_mel_bins=self.n_mels, - frame_length=self.frame_length, - frame_shift=self.frame_shift, - dither=self.dither, - energy_floor=0.0, - window_type=self.window, - sample_frequency=self.fs) + mat = kaldi.fbank( + waveform, + num_mel_bins=self.n_mels, + frame_length=self.frame_length, + frame_shift=self.frame_shift, + dither=self.dither, + energy_floor=0.0, + window_type=self.window, + sample_frequency=self.fs, + ) feat_length = mat.size(0) feats.append(mat) @@ -358,33 +369,31 @@ class WavFrontendOnline(nn.Module): waveforms = torch.stack(waveforms) feats_lens = torch.as_tensor(feats_lens) - feats_pad = pad_sequence(feats, - batch_first=True, - padding_value=0.0) + feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0) cache["fbanks"] = feats_pad - cache["fbanks_lens"]= copy.deepcopy(feats_lens) + cache["fbanks_lens"] = copy.deepcopy(feats_lens) return waveforms, feats_pad, feats_lens - def forward_lfr_cmvn( - self, - input: torch.Tensor, - input_lengths: torch.Tensor, - is_final: bool = False, - cache: dict = {}, - **kwargs, + self, + input: torch.Tensor, + input_lengths: torch.Tensor, + is_final: bool = False, + cache: dict = {}, + **kwargs, ): batch_size = input.size(0) feats = [] feats_lens = [] lfr_splice_frame_idxs = [] for i in range(batch_size): - mat = input[i, :input_lengths[i], :] + mat = input[i, : input_lengths[i], :] if self.lfr_m != 1 or self.lfr_n != 1: # update self.lfr_splice_cache in self.apply_lfr # mat, self.lfr_splice_cache[i], lfr_splice_frame_idx = self.apply_lfr(mat, self.lfr_m, self.lfr_n, self.lfr_splice_cache[i], - mat, cache["lfr_splice_cache"][i], lfr_splice_frame_idx = self.apply_lfr(mat, self.lfr_m, self.lfr_n, - is_final) + mat, cache["lfr_splice_cache"][i], lfr_splice_frame_idx = self.apply_lfr( + mat, self.lfr_m, self.lfr_n, is_final + ) if self.cmvn_file is not None: mat = self.apply_cmvn(mat, self.cmvn) feat_length = mat.size(0) @@ -393,68 +402,93 @@ class WavFrontendOnline(nn.Module): lfr_splice_frame_idxs.append(lfr_splice_frame_idx) feats_lens = torch.as_tensor(feats_lens) - feats_pad = pad_sequence(feats, - batch_first=True, - padding_value=0.0) + feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0) lfr_splice_frame_idxs = torch.as_tensor(lfr_splice_frame_idxs) return feats_pad, feats_lens, lfr_splice_frame_idxs - def forward( - self, input: torch.Tensor, input_lengths: torch.Tensor, **kwargs - ): + def forward(self, input: torch.Tensor, input_lengths: torch.Tensor, **kwargs): is_final = kwargs.get("is_final", False) cache = kwargs.get("cache", {}) if len(cache) == 0: self.init_cache(cache) - + batch_size = input.shape[0] - assert batch_size == 1, 'we support to extract feature online only when the batch size is equal to 1 now' - - waveforms, feats, feats_lengths = self.forward_fbank(input, input_lengths, cache=cache) # input shape: B T D - + assert ( + batch_size == 1 + ), "we support to extract feature online only when the batch size is equal to 1 now" + + waveforms, feats, feats_lengths = self.forward_fbank( + input, input_lengths, cache=cache + ) # input shape: B T D + if feats.shape[0]: cache["waveforms"] = torch.cat((cache["reserve_waveforms"], waveforms), dim=1) - + if not cache["lfr_splice_cache"]: # 初始化splice_cache for i in range(batch_size): - cache["lfr_splice_cache"].append(feats[i][0, :].unsqueeze(dim=0).repeat((self.lfr_m - 1) // 2, 1)) + cache["lfr_splice_cache"].append( + feats[i][0, :].unsqueeze(dim=0).repeat((self.lfr_m - 1) // 2, 1) + ) # need the number of the input frames + self.lfr_splice_cache[0].shape[0] is greater than self.lfr_m if feats_lengths[0] + cache["lfr_splice_cache"][0].shape[0] >= self.lfr_m: lfr_splice_cache_tensor = torch.stack(cache["lfr_splice_cache"]) # B T D feats = torch.cat((lfr_splice_cache_tensor, feats), dim=1) feats_lengths += lfr_splice_cache_tensor[0].shape[0] frame_from_waveforms = int( - (cache["waveforms"].shape[1] - self.frame_sample_length) / self.frame_shift_sample_length + 1) - minus_frame = (self.lfr_m - 1) // 2 if cache["reserve_waveforms"].numel() == 0 else 0 - feats, feats_lengths, lfr_splice_frame_idxs = self.forward_lfr_cmvn(feats, feats_lengths, is_final, cache=cache) + (cache["waveforms"].shape[1] - self.frame_sample_length) + / self.frame_shift_sample_length + + 1 + ) + minus_frame = ( + (self.lfr_m - 1) // 2 if cache["reserve_waveforms"].numel() == 0 else 0 + ) + feats, feats_lengths, lfr_splice_frame_idxs = self.forward_lfr_cmvn( + feats, feats_lengths, is_final, cache=cache + ) if self.lfr_m == 1: cache["reserve_waveforms"] = torch.empty(0) else: reserve_frame_idx = lfr_splice_frame_idxs[0] - minus_frame # print('reserve_frame_idx: ' + str(reserve_frame_idx)) # print('frame_frame: ' + str(frame_from_waveforms)) - cache["reserve_waveforms"] = cache["waveforms"][:, reserve_frame_idx * self.frame_shift_sample_length:frame_from_waveforms * self.frame_shift_sample_length] - sample_length = (frame_from_waveforms - 1) * self.frame_shift_sample_length + self.frame_sample_length + cache["reserve_waveforms"] = cache["waveforms"][ + :, + reserve_frame_idx + * self.frame_shift_sample_length : frame_from_waveforms + * self.frame_shift_sample_length, + ] + sample_length = ( + frame_from_waveforms - 1 + ) * self.frame_shift_sample_length + self.frame_sample_length cache["waveforms"] = cache["waveforms"][:, :sample_length] else: # update self.reserve_waveforms and self.lfr_splice_cache - cache["reserve_waveforms"] = cache["waveforms"][:, :-(self.frame_sample_length - self.frame_shift_sample_length)] + cache["reserve_waveforms"] = cache["waveforms"][ + :, : -(self.frame_sample_length - self.frame_shift_sample_length) + ] for i in range(batch_size): - cache["lfr_splice_cache"][i] = torch.cat((cache["lfr_splice_cache"][i], feats[i]), dim=0) + cache["lfr_splice_cache"][i] = torch.cat( + (cache["lfr_splice_cache"][i], feats[i]), dim=0 + ) return torch.empty(0), feats_lengths else: if is_final: - cache["waveforms"] = waveforms if cache["reserve_waveforms"].numel() == 0 else cache["reserve_waveforms"] + cache["waveforms"] = ( + waveforms + if cache["reserve_waveforms"].numel() == 0 + else cache["reserve_waveforms"] + ) feats = torch.stack(cache["lfr_splice_cache"]) feats_lengths = torch.zeros(batch_size, dtype=torch.int) + feats.shape[1] - feats, feats_lengths, _ = self.forward_lfr_cmvn(feats, feats_lengths, is_final, cache=cache) + feats, feats_lengths, _ = self.forward_lfr_cmvn( + feats, feats_lengths, is_final, cache=cache + ) # if is_final: # self.init_cache(cache) return feats, feats_lengths - - def init_cache(self, cache: dict = {}): + def init_cache(self, cache: dict = {}): cache["reserve_waveforms"] = torch.empty(0) cache["input_cache"] = torch.empty(0) cache["lfr_splice_cache"] = [] @@ -465,17 +499,16 @@ class WavFrontendOnline(nn.Module): class WavFrontendMel23(nn.Module): - """Conventional frontend structure for ASR. - """ + """Conventional frontend structure for ASR.""" def __init__( - self, - fs: int = 16000, - frame_length: int = 25, - frame_shift: int = 10, - lfr_m: int = 1, - lfr_n: int = 1, - **kwargs, + self, + fs: int = 16000, + frame_length: int = 25, + frame_shift: int = 10, + lfr_m: int = 1, + lfr_n: int = 1, + **kwargs, ): super().__init__() self.fs = fs @@ -489,9 +522,8 @@ class WavFrontendMel23(nn.Module): return self.n_mels * (2 * self.lfr_m + 1) def forward( - self, - input: torch.Tensor, - input_lengths: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + self, input: torch.Tensor, input_lengths: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: batch_size = input.size(0) feats = [] feats_lens = [] @@ -502,14 +534,12 @@ class WavFrontendMel23(nn.Module): mat = eend_ola_feature.stft(waveform, self.frame_length, self.frame_shift) mat = eend_ola_feature.transform(mat) mat = eend_ola_feature.splice(mat, context_size=self.lfr_m) - mat = mat[::self.lfr_n] + mat = mat[:: self.lfr_n] mat = torch.from_numpy(mat) feat_length = mat.size(0) feats.append(mat) feats_lens.append(feat_length) feats_lens = torch.as_tensor(feats_lens) - feats_pad = pad_sequence(feats, - batch_first=True, - padding_value=0.0) + feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0) return feats_pad, feats_lens diff --git a/funasr/frontends/whisper_frontend.py b/funasr/frontends/whisper_frontend.py index acc99af01..acf32dbc7 100644 --- a/funasr/frontends/whisper_frontend.py +++ b/funasr/frontends/whisper_frontend.py @@ -15,19 +15,20 @@ class WhisperFrontend(nn.Module): """ def __init__( - self, - fs: int = 16000, - whisper_model: str = None, - do_pad_trim: bool = True, - n_mels: int = 80, - permute: bool = False, - **kwargs, + self, + fs: int = 16000, + whisper_model: str = None, + do_pad_trim: bool = True, + n_mels: int = 80, + permute: bool = False, + **kwargs, ): super().__init__() assert fs == 16000 self.fs = fs import whisper from whisper.audio import HOP_LENGTH, N_FFT, N_SAMPLES + self.n_fft = N_FFT self.win_length = N_FFT self.hop_length = HOP_LENGTH @@ -42,6 +43,7 @@ class WhisperFrontend(nn.Module): self.filters_path = filters_path if filters_path is not None: from funasr.models.sense_voice.whisper_lib.audio import mel_filters + self.mel_filters = mel_filters else: self.mel_filters = whisper.audio.mel_filters @@ -56,14 +58,12 @@ class WhisperFrontend(nn.Module): return self.n_mels def log_mel_spectrogram( - self, - audio: torch.Tensor, - ilens: torch.Tensor = None, + self, + audio: torch.Tensor, + ilens: torch.Tensor = None, ) -> torch.Tensor: window = torch.hann_window(self.win_length).to(audio.device) - stft = torch.stft( - audio, self.n_fft, self.hop_length, window=window, return_complex=True - ) + stft = torch.stft(audio, self.n_fft, self.hop_length, window=window, return_complex=True) # whisper deletes the last frame by default (Shih-Lun) magnitudes = stft[..., :-1].abs() ** 2 @@ -89,7 +89,10 @@ class WhisperFrontend(nn.Module): return log_spec, olens def forward( - self, input: torch.Tensor, input_lengths: torch.Tensor, **kwargs, + self, + input: torch.Tensor, + input_lengths: torch.Tensor, + **kwargs, ) -> Tuple[torch.Tensor, torch.Tensor]: batch_size = input.size(0) feats = [] @@ -108,9 +111,7 @@ class WhisperFrontend(nn.Module): if batch_size == 1: feats_pad = feats[0][None, :, :] else: - feats_pad = pad_sequence(feats, - batch_first=True, - padding_value=0.0) + feats_pad = pad_sequence(feats, batch_first=True, padding_value=0.0) if self.permute: feats_pad = feats_pad.permute(0, 2, 1) - return feats_pad, feats_lens \ No newline at end of file + return feats_pad, feats_lens diff --git a/funasr/frontends/windowing.py b/funasr/frontends/windowing.py index 3250550ef..84ec8e200 100644 --- a/funasr/frontends/windowing.py +++ b/funasr/frontends/windowing.py @@ -71,4 +71,4 @@ class SlidingWindow(nn.Module): def output_size(self) -> int: """Return output length of feature dimension D, i.e. the window length.""" - return self.win_length \ No newline at end of file + return self.win_length diff --git a/funasr/losses/label_smoothing_loss.py b/funasr/losses/label_smoothing_loss.py index 385025dab..ffc38dae7 100644 --- a/funasr/losses/label_smoothing_loss.py +++ b/funasr/losses/label_smoothing_loss.py @@ -65,11 +65,7 @@ class LabelSmoothingLoss(nn.Module): class SequenceBinaryCrossEntropy(nn.Module): - def __init__( - self, - normalize_length=False, - criterion=nn.BCEWithLogitsLoss(reduction="none") - ): + def __init__(self, normalize_length=False, criterion=nn.BCEWithLogitsLoss(reduction="none")): super().__init__() self.normalize_length = normalize_length self.criterion = criterion @@ -95,7 +91,7 @@ class NllLoss(nn.Module): size, padding_idx, normalize_length=False, - criterion=nn.NLLLoss(reduction='none'), + criterion=nn.NLLLoss(reduction="none"), ): """Construct an NllLoss object.""" super(NllLoss, self).__init__() @@ -122,6 +118,6 @@ class NllLoss(nn.Module): ignore = target == self.padding_idx # (B,) total = len(target) - ignore.sum().item() target = target.masked_fill(ignore, 0) # avoid -1 index - kl = self.criterion(x , target) + kl = self.criterion(x, target) denom = total if self.normalize_length else batch_size return kl.masked_fill(ignore, 0).sum() / denom diff --git a/funasr/metrics/common.py b/funasr/metrics/common.py index 92f90796a..2443e0dc5 100644 --- a/funasr/metrics/common.py +++ b/funasr/metrics/common.py @@ -36,9 +36,9 @@ def end_detect(ended_hyps, i, M=3, D_end=np.log(1 * np.exp(-10))): hyp_length = i - m hyps_same_length = [x for x in ended_hyps if len(x["yseq"]) == hyp_length] if len(hyps_same_length) > 0: - best_hyp_same_length = sorted( - hyps_same_length, key=lambda x: x["score"], reverse=True - )[0] + best_hyp_same_length = sorted(hyps_same_length, key=lambda x: x["score"], reverse=True)[ + 0 + ] if best_hyp_same_length["score"] - best_hyp["score"] < D_end: count += 1 @@ -63,9 +63,7 @@ def label_smoothing_dist(odim, lsm_type, transcript=None, blank=0): trans_json = json.load(f)["utts"] if lsm_type == "unigram": - assert transcript is not None, ( - "transcript is required for %s label smoothing" % lsm_type - ) + assert transcript is not None, "transcript is required for %s label smoothing" % lsm_type labelcount = np.zeros(odim) for k, v in trans_json.items(): ids = np.array([int(n) for n in v["output"][0]["tokenid"].split()]) @@ -108,9 +106,7 @@ class ErrorCalculator(object): :return: """ - def __init__( - self, char_list, sym_space, sym_blank, report_cer=False, report_wer=False - ): + def __init__(self, char_list, sym_space, sym_blank, report_cer=False, report_wer=False): """Construct an ErrorCalculator object.""" super(ErrorCalculator, self).__init__() diff --git a/funasr/metrics/compute_acc.py b/funasr/metrics/compute_acc.py index ec8067f3c..834586687 100644 --- a/funasr/metrics/compute_acc.py +++ b/funasr/metrics/compute_acc.py @@ -1,5 +1,6 @@ import torch + def th_accuracy(pad_outputs, pad_targets, ignore_label): """Calculate accuracy. @@ -16,12 +17,11 @@ def th_accuracy(pad_outputs, pad_targets, ignore_label): pad_targets.size(0), pad_targets.size(1), pad_outputs.size(1) ).argmax(2) mask = pad_targets != ignore_label - numerator = torch.sum( - pad_pred.masked_select(mask) == pad_targets.masked_select(mask) - ) + numerator = torch.sum(pad_pred.masked_select(mask) == pad_targets.masked_select(mask)) denominator = torch.sum(mask) return float(numerator) / float(denominator) + def compute_accuracy(pad_outputs, pad_targets, ignore_label): """Calculate accuracy. @@ -37,4 +37,4 @@ def compute_accuracy(pad_outputs, pad_targets, ignore_label): mask = pad_targets != ignore_label numerator = torch.sum(pad_outputs.masked_select(mask) == pad_targets.masked_select(mask)) denominator = torch.sum(mask) - return numerator.float() / denominator.float() #(FIX:MZY):return torch.Tensor type \ No newline at end of file + return numerator.float() / denominator.float() # (FIX:MZY):return torch.Tensor type diff --git a/funasr/metrics/compute_eer.py b/funasr/metrics/compute_eer.py index e2ea0d8e3..5180b1669 100644 --- a/funasr/metrics/compute_eer.py +++ b/funasr/metrics/compute_eer.py @@ -55,5 +55,5 @@ def main(): print("EER is {:.4f} at threshold {:.4f}".format(eer * 100.0, threshold)) -if __name__ == '__main__': - main() \ No newline at end of file +if __name__ == "__main__": + main() diff --git a/funasr/metrics/compute_min_dcf.py b/funasr/metrics/compute_min_dcf.py index 610113a8f..472f7b3e2 100644 --- a/funasr/metrics/compute_min_dcf.py +++ b/funasr/metrics/compute_min_dcf.py @@ -16,27 +16,45 @@ import sys, argparse, os def GetArgs(): - parser = argparse.ArgumentParser(description="Compute the minimum " - "detection cost function along with the threshold at which it occurs. " - "Usage: sid/compute_min_dcf.py [options...] " - " " - "E.g., sid/compute_min_dcf.py --p-target 0.01 --c-miss 1 --c-fa 1 " - "exp/scores/trials data/test/trials", - formatter_class=argparse.ArgumentDefaultsHelpFormatter) - parser.add_argument('--p-target', type=float, dest="p_target", - default=0.01, - help='The prior probability of the target speaker in a trial.') - parser.add_argument('--c-miss', type=float, dest="c_miss", default=1, - help='Cost of a missed detection. This is usually not changed.') - parser.add_argument('--c-fa', type=float, dest="c_fa", default=1, - help='Cost of a spurious detection. This is usually not changed.') - parser.add_argument("scores_filename", - help="Input scores file, with columns of the form " - " ") - parser.add_argument("trials_filename", - help="Input trials file, with columns of the form " - " ") - sys.stderr.write(' '.join(sys.argv) + "\n") + parser = argparse.ArgumentParser( + description="Compute the minimum " + "detection cost function along with the threshold at which it occurs. " + "Usage: sid/compute_min_dcf.py [options...] " + " " + "E.g., sid/compute_min_dcf.py --p-target 0.01 --c-miss 1 --c-fa 1 " + "exp/scores/trials data/test/trials", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "--p-target", + type=float, + dest="p_target", + default=0.01, + help="The prior probability of the target speaker in a trial.", + ) + parser.add_argument( + "--c-miss", + type=float, + dest="c_miss", + default=1, + help="Cost of a missed detection. This is usually not changed.", + ) + parser.add_argument( + "--c-fa", + type=float, + dest="c_fa", + default=1, + help="Cost of a spurious detection. This is usually not changed.", + ) + parser.add_argument( + "scores_filename", + help="Input scores file, with columns of the form " " ", + ) + parser.add_argument( + "trials_filename", + help="Input trials file, with columns of the form " " ", + ) + sys.stderr.write(" ".join(sys.argv) + "\n") args = parser.parse_args() args = CheckArgs(args) return args @@ -44,11 +62,11 @@ def GetArgs(): def CheckArgs(args): if args.c_fa <= 0: - raise Exception("--c-fa must be greater than 0") + raise Exception("--c-fa must be greater than 0") if args.c_miss <= 0: - raise Exception("--c-miss must be greater than 0") + raise Exception("--c-miss must be greater than 0") if args.p_target <= 0 or args.p_target >= 1: - raise Exception("--p-target must be greater than 0 and less than 1") + raise Exception("--p-target must be greater than 0 and less than 1") return args @@ -59,9 +77,9 @@ def ComputeErrorRates(scores, labels): # Sort the scores from smallest to largest, and also get the corresponding # indexes of the sorted scores. We will treat the sorted scores as the # thresholds at which the the error-rates are evaluated. - sorted_indexes, thresholds = zip(*sorted( - [(index, threshold) for index, threshold in enumerate(scores)], - key=itemgetter(1))) + sorted_indexes, thresholds = zip( + *sorted([(index, threshold) for index, threshold in enumerate(scores)], key=itemgetter(1)) + ) labels = [labels[i] for i in sorted_indexes] fns = [] tns = [] @@ -75,18 +93,18 @@ def ComputeErrorRates(scores, labels): fns.append(labels[i]) tns.append(1 - labels[i]) else: - fns.append(fns[i-1] + labels[i]) - tns.append(tns[i-1] + 1 - labels[i]) + fns.append(fns[i - 1] + labels[i]) + tns.append(tns[i - 1] + 1 - labels[i]) positives = sum(labels) negatives = len(labels) - positives - # Now divide the false negatives by the total number of + # Now divide the false negatives by the total number of # positives to obtain the false negative rates across # all thresholds fnrs = [fn / float(positives) for fn in fns] - # Divide the true negatives by the total number of - # negatives to get the true negative rate. Subtract these + # Divide the true negatives by the total number of + # negatives to get the true negative rate. Subtract these # quantities from 1 to get the false positive rates. fprs = [1 - tn / float(negatives) for tn in tns] return fnrs, fprs, thresholds @@ -111,8 +129,8 @@ def ComputeMinDcf(fnrs, fprs, thresholds, p_target, c_miss, c_fa): def compute_min_dcf(scores_filename, trials_filename, c_miss=1, c_fa=1, p_target=0.01): - scores_file = open(scores_filename, 'r').readlines() - trials_file = open(trials_filename, 'r').readlines() + scores_file = open(scores_filename, "r").readlines() + trials_file = open(trials_filename, "r").readlines() c_miss = c_miss c_fa = c_fa p_target = p_target @@ -136,23 +154,22 @@ def compute_min_dcf(scores_filename, trials_filename, c_miss=1, c_fa=1, p_target else: labels.append(0) else: - raise Exception("Missing entry for " + utt1 + " and " + utt2 - + " " + scores_filename) + raise Exception("Missing entry for " + utt1 + " and " + utt2 + " " + scores_filename) fnrs, fprs, thresholds = ComputeErrorRates(scores, labels) - mindcf, threshold = ComputeMinDcf(fnrs, fprs, thresholds, p_target, - c_miss, c_fa) + mindcf, threshold = ComputeMinDcf(fnrs, fprs, thresholds, p_target, c_miss, c_fa) return mindcf, threshold def main(): args = GetArgs() mindcf, threshold = compute_min_dcf( - args.scores_filename, args.trials_filename, - args.c_miss, args.c_fa, args.p_target + args.scores_filename, args.trials_filename, args.c_miss, args.c_fa, args.p_target + ) + sys.stdout.write( + "minDCF is {0:.4f} at threshold {1:.4f} (p-target={2}, c-miss={3}, " + "c-fa={4})\n".format(mindcf, threshold, args.p_target, args.c_miss, args.c_fa) ) - sys.stdout.write("minDCF is {0:.4f} at threshold {1:.4f} (p-target={2}, c-miss={3}, " - "c-fa={4})\n".format(mindcf, threshold, args.p_target, args.c_miss, args.c_fa)) if __name__ == "__main__": diff --git a/funasr/metrics/wer.py b/funasr/metrics/wer.py index b58daabc2..ab07bc29f 100755 --- a/funasr/metrics/wer.py +++ b/funasr/metrics/wer.py @@ -5,186 +5,221 @@ import hydra from omegaconf import DictConfig, OmegaConf, ListConfig -def compute_wer(ref_file, - hyp_file, - cer_file, - cn_postprocess=False, - ): - rst = { - 'Wrd': 0, - 'Corr': 0, - 'Ins': 0, - 'Del': 0, - 'Sub': 0, - 'Snt': 0, - 'Err': 0.0, - 'S.Err': 0.0, - 'wrong_words': 0, - 'wrong_sentences': 0 - } - - hyp_dict = {} - ref_dict = {} - with open(hyp_file, 'r') as hyp_reader: - for line in hyp_reader: - key = line.strip().split()[0] - value = line.strip().split()[1:] - if cn_postprocess: - value = " ".join(value) - value = value.replace(" ", "") - if value[0] == "请": - value = value[1:] - value = [x for x in value] - hyp_dict[key] = value - with open(ref_file, 'r') as ref_reader: - for line in ref_reader: - key = line.strip().split()[0] - value = line.strip().split()[1:] - if cn_postprocess: - value = " ".join(value) - value = value.replace(" ", "") - value = [x for x in value] - ref_dict[key] = value - - cer_detail_writer = open(cer_file, 'w') - for hyp_key in hyp_dict: - if hyp_key in ref_dict: - out_item = compute_wer_by_line(hyp_dict[hyp_key], ref_dict[hyp_key]) - rst['Wrd'] += out_item['nwords'] - rst['Corr'] += out_item['cor'] - rst['wrong_words'] += out_item['wrong'] - rst['Ins'] += out_item['ins'] - rst['Del'] += out_item['del'] - rst['Sub'] += out_item['sub'] - rst['Snt'] += 1 - if out_item['wrong'] > 0: - rst['wrong_sentences'] += 1 - cer_detail_writer.write(hyp_key + print_cer_detail(out_item) + '\n') - cer_detail_writer.write("ref:" + '\t' + " ".join(list(map(lambda x: x.lower(), ref_dict[hyp_key]))) + '\n') - cer_detail_writer.write("hyp:" + '\t' + " ".join(list(map(lambda x: x.lower(), hyp_dict[hyp_key]))) + '\n') - cer_detail_writer.flush() - - if rst['Wrd'] > 0: - rst['Err'] = round(rst['wrong_words'] * 100 / rst['Wrd'], 2) - if rst['Snt'] > 0: - rst['S.Err'] = round(rst['wrong_sentences'] * 100 / rst['Snt'], 2) - - cer_detail_writer.write('\n') - cer_detail_writer.write("%WER " + str(rst['Err']) + " [ " + str(rst['wrong_words']) + " / " + str(rst['Wrd']) + - ", " + str(rst['Ins']) + " ins, " + str(rst['Del']) + " del, " + str( - rst['Sub']) + " sub ]" + '\n') - cer_detail_writer.write( - "%SER " + str(rst['S.Err']) + " [ " + str(rst['wrong_sentences']) + " / " + str(rst['Snt']) + " ]" + '\n') - cer_detail_writer.write("Scored " + str(len(hyp_dict)) + " sentences, " + str( - len(hyp_dict) - rst['Snt']) + " not present in hyp." + '\n') - - cer_detail_writer.close() +def compute_wer( + ref_file, + hyp_file, + cer_file, + cn_postprocess=False, +): + rst = { + "Wrd": 0, + "Corr": 0, + "Ins": 0, + "Del": 0, + "Sub": 0, + "Snt": 0, + "Err": 0.0, + "S.Err": 0.0, + "wrong_words": 0, + "wrong_sentences": 0, + } + + hyp_dict = {} + ref_dict = {} + with open(hyp_file, "r") as hyp_reader: + for line in hyp_reader: + key = line.strip().split()[0] + value = line.strip().split()[1:] + if cn_postprocess: + value = " ".join(value) + value = value.replace(" ", "") + if value[0] == "请": + value = value[1:] + value = [x for x in value] + hyp_dict[key] = value + with open(ref_file, "r") as ref_reader: + for line in ref_reader: + key = line.strip().split()[0] + value = line.strip().split()[1:] + if cn_postprocess: + value = " ".join(value) + value = value.replace(" ", "") + value = [x for x in value] + ref_dict[key] = value + + cer_detail_writer = open(cer_file, "w") + for hyp_key in hyp_dict: + if hyp_key in ref_dict: + out_item = compute_wer_by_line(hyp_dict[hyp_key], ref_dict[hyp_key]) + rst["Wrd"] += out_item["nwords"] + rst["Corr"] += out_item["cor"] + rst["wrong_words"] += out_item["wrong"] + rst["Ins"] += out_item["ins"] + rst["Del"] += out_item["del"] + rst["Sub"] += out_item["sub"] + rst["Snt"] += 1 + if out_item["wrong"] > 0: + rst["wrong_sentences"] += 1 + cer_detail_writer.write(hyp_key + print_cer_detail(out_item) + "\n") + cer_detail_writer.write( + "ref:" + "\t" + " ".join(list(map(lambda x: x.lower(), ref_dict[hyp_key]))) + "\n" + ) + cer_detail_writer.write( + "hyp:" + "\t" + " ".join(list(map(lambda x: x.lower(), hyp_dict[hyp_key]))) + "\n" + ) + cer_detail_writer.flush() + + if rst["Wrd"] > 0: + rst["Err"] = round(rst["wrong_words"] * 100 / rst["Wrd"], 2) + if rst["Snt"] > 0: + rst["S.Err"] = round(rst["wrong_sentences"] * 100 / rst["Snt"], 2) + + cer_detail_writer.write("\n") + cer_detail_writer.write( + "%WER " + + str(rst["Err"]) + + " [ " + + str(rst["wrong_words"]) + + " / " + + str(rst["Wrd"]) + + ", " + + str(rst["Ins"]) + + " ins, " + + str(rst["Del"]) + + " del, " + + str(rst["Sub"]) + + " sub ]" + + "\n" + ) + cer_detail_writer.write( + "%SER " + + str(rst["S.Err"]) + + " [ " + + str(rst["wrong_sentences"]) + + " / " + + str(rst["Snt"]) + + " ]" + + "\n" + ) + cer_detail_writer.write( + "Scored " + + str(len(hyp_dict)) + + " sentences, " + + str(len(hyp_dict) - rst["Snt"]) + + " not present in hyp." + + "\n" + ) + + cer_detail_writer.close() -def compute_wer_by_line(hyp, - ref): - hyp = list(map(lambda x: x.lower(), hyp)) - ref = list(map(lambda x: x.lower(), ref)) - - len_hyp = len(hyp) - len_ref = len(ref) - - cost_matrix = np.zeros((len_hyp + 1, len_ref + 1), dtype=np.int16) - - ops_matrix = np.zeros((len_hyp + 1, len_ref + 1), dtype=np.int8) - - for i in range(len_hyp + 1): - cost_matrix[i][0] = i - for j in range(len_ref + 1): - cost_matrix[0][j] = j - - for i in range(1, len_hyp + 1): - for j in range(1, len_ref + 1): - if hyp[i - 1] == ref[j - 1]: - cost_matrix[i][j] = cost_matrix[i - 1][j - 1] - else: - substitution = cost_matrix[i - 1][j - 1] + 1 - insertion = cost_matrix[i - 1][j] + 1 - deletion = cost_matrix[i][j - 1] + 1 - - compare_val = [substitution, insertion, deletion] - - min_val = min(compare_val) - operation_idx = compare_val.index(min_val) + 1 - cost_matrix[i][j] = min_val - ops_matrix[i][j] = operation_idx - - match_idx = [] - i = len_hyp - j = len_ref - rst = { - 'nwords': len_ref, - 'cor': 0, - 'wrong': 0, - 'ins': 0, - 'del': 0, - 'sub': 0 - } - while i >= 0 or j >= 0: - i_idx = max(0, i) - j_idx = max(0, j) - - if ops_matrix[i_idx][j_idx] == 0: # correct - if i - 1 >= 0 and j - 1 >= 0: - match_idx.append((j - 1, i - 1)) - rst['cor'] += 1 - - i -= 1 - j -= 1 - - elif ops_matrix[i_idx][j_idx] == 2: # insert - i -= 1 - rst['ins'] += 1 - - elif ops_matrix[i_idx][j_idx] == 3: # delete - j -= 1 - rst['del'] += 1 - - elif ops_matrix[i_idx][j_idx] == 1: # substitute - i -= 1 - j -= 1 - rst['sub'] += 1 - - if i < 0 and j >= 0: - rst['del'] += 1 - elif j < 0 and i >= 0: - rst['ins'] += 1 - - match_idx.reverse() - wrong_cnt = cost_matrix[len_hyp][len_ref] - rst['wrong'] = wrong_cnt - - return rst +def compute_wer_by_line(hyp, ref): + hyp = list(map(lambda x: x.lower(), hyp)) + ref = list(map(lambda x: x.lower(), ref)) + + len_hyp = len(hyp) + len_ref = len(ref) + + cost_matrix = np.zeros((len_hyp + 1, len_ref + 1), dtype=np.int16) + + ops_matrix = np.zeros((len_hyp + 1, len_ref + 1), dtype=np.int8) + + for i in range(len_hyp + 1): + cost_matrix[i][0] = i + for j in range(len_ref + 1): + cost_matrix[0][j] = j + + for i in range(1, len_hyp + 1): + for j in range(1, len_ref + 1): + if hyp[i - 1] == ref[j - 1]: + cost_matrix[i][j] = cost_matrix[i - 1][j - 1] + else: + substitution = cost_matrix[i - 1][j - 1] + 1 + insertion = cost_matrix[i - 1][j] + 1 + deletion = cost_matrix[i][j - 1] + 1 + + compare_val = [substitution, insertion, deletion] + + min_val = min(compare_val) + operation_idx = compare_val.index(min_val) + 1 + cost_matrix[i][j] = min_val + ops_matrix[i][j] = operation_idx + + match_idx = [] + i = len_hyp + j = len_ref + rst = {"nwords": len_ref, "cor": 0, "wrong": 0, "ins": 0, "del": 0, "sub": 0} + while i >= 0 or j >= 0: + i_idx = max(0, i) + j_idx = max(0, j) + + if ops_matrix[i_idx][j_idx] == 0: # correct + if i - 1 >= 0 and j - 1 >= 0: + match_idx.append((j - 1, i - 1)) + rst["cor"] += 1 + + i -= 1 + j -= 1 + + elif ops_matrix[i_idx][j_idx] == 2: # insert + i -= 1 + rst["ins"] += 1 + + elif ops_matrix[i_idx][j_idx] == 3: # delete + j -= 1 + rst["del"] += 1 + + elif ops_matrix[i_idx][j_idx] == 1: # substitute + i -= 1 + j -= 1 + rst["sub"] += 1 + + if i < 0 and j >= 0: + rst["del"] += 1 + elif j < 0 and i >= 0: + rst["ins"] += 1 + + match_idx.reverse() + wrong_cnt = cost_matrix[len_hyp][len_ref] + rst["wrong"] = wrong_cnt + + return rst def print_cer_detail(rst): - return ("(" + "nwords=" + str(rst['nwords']) + ",cor=" + str(rst['cor']) - + ",ins=" + str(rst['ins']) + ",del=" + str(rst['del']) + ",sub=" - + str(rst['sub']) + ") corr:" + '{:.2%}'.format(rst['cor'] / rst['nwords']) - + ",cer:" + '{:.2%}'.format(rst['wrong'] / rst['nwords'])) + return ( + "(" + + "nwords=" + + str(rst["nwords"]) + + ",cor=" + + str(rst["cor"]) + + ",ins=" + + str(rst["ins"]) + + ",del=" + + str(rst["del"]) + + ",sub=" + + str(rst["sub"]) + + ") corr:" + + "{:.2%}".format(rst["cor"] / rst["nwords"]) + + ",cer:" + + "{:.2%}".format(rst["wrong"] / rst["nwords"]) + ) @hydra.main(config_name=None, version_base=None) def main_hydra(cfg: DictConfig): - ref_file = cfg.get("ref_file", None) - hyp_file = cfg.get("hyp_file", None) - cer_file = cfg.get("cer_file", None) - cn_postprocess = cfg.get("cn_postprocess", False) - if ref_file is None or hyp_file is None or cer_file is None: - print( - "usage : python -m funasr.metrics.wer ++ref_file=test.ref ++hyp_file=test.hyp ++cer_file=test.wer ++cn_postprocess=false") - sys.exit(0) - - compute_wer(ref_file, hyp_file, cer_file, cn_postprocess) + ref_file = cfg.get("ref_file", None) + hyp_file = cfg.get("hyp_file", None) + cer_file = cfg.get("cer_file", None) + cn_postprocess = cfg.get("cn_postprocess", False) + if ref_file is None or hyp_file is None or cer_file is None: + print( + "usage : python -m funasr.metrics.wer ++ref_file=test.ref ++hyp_file=test.hyp ++cer_file=test.wer ++cn_postprocess=false" + ) + sys.exit(0) - -if __name__ == '__main__': - main_hydra() - + compute_wer(ref_file, hyp_file, cer_file, cn_postprocess) +if __name__ == "__main__": + main_hydra() diff --git a/funasr/models/bat/model.py b/funasr/models/bat/model.py index bdfcba6e1..ec2244337 100644 --- a/funasr/models/bat/model.py +++ b/funasr/models/bat/model.py @@ -34,4 +34,4 @@ else: @tables.register("model_classes", "BAT") # TODO: BAT training class BAT(Transducer): - pass \ No newline at end of file + pass diff --git a/funasr/models/bicif_paraformer/cif_predictor.py b/funasr/models/bicif_paraformer/cif_predictor.py index 2cdbc16a8..3739c76f4 100644 --- a/funasr/models/bicif_paraformer/cif_predictor.py +++ b/funasr/models/bicif_paraformer/cif_predictor.py @@ -14,7 +14,7 @@ class mae_loss(torch.nn.Module): def __init__(self, normalize_length=False): super(mae_loss, self).__init__() self.normalize_length = normalize_length - self.criterion = torch.nn.L1Loss(reduction='sum') + self.criterion = torch.nn.L1Loss(reduction="sum") def forward(self, token_length, pre_token_length): loss_token_normalizer = token_length.size(0) @@ -43,19 +43,17 @@ def cif(hidden, alphas, threshold): list_fires.append(integrate) fire_place = integrate >= threshold - integrate = torch.where(fire_place, - integrate - torch.ones([batch_size], device=hidden.device), - integrate) - cur = torch.where(fire_place, - distribution_completion, - alpha) + integrate = torch.where( + fire_place, integrate - torch.ones([batch_size], device=hidden.device), integrate + ) + cur = torch.where(fire_place, distribution_completion, alpha) remainds = alpha - cur frame += cur[:, None] * hidden[:, t, :] list_frames.append(frame) - frame = torch.where(fire_place[:, None].repeat(1, hidden_size), - remainds[:, None] * hidden[:, t, :], - frame) + frame = torch.where( + fire_place[:, None].repeat(1, hidden_size), remainds[:, None] * hidden[:, t, :], frame + ) fires = torch.stack(list_fires, 1) frames = torch.stack(list_frames, 1) @@ -85,33 +83,37 @@ def cif_wo_hidden(alphas, threshold): list_fires.append(integrate) fire_place = integrate >= threshold - integrate = torch.where(fire_place, - integrate - torch.ones([batch_size], device=alphas.device)*threshold, - integrate) + integrate = torch.where( + fire_place, + integrate - torch.ones([batch_size], device=alphas.device) * threshold, + integrate, + ) fires = torch.stack(list_fires, 1) return fires + @tables.register("predictor_classes", "CifPredictorV3") class CifPredictorV3(torch.nn.Module): - def __init__(self, - idim, - l_order, - r_order, - threshold=1.0, - dropout=0.1, - smooth_factor=1.0, - noise_threshold=0, - tail_threshold=0.0, - tf2torch_tensor_name_prefix_torch="predictor", - tf2torch_tensor_name_prefix_tf="seq2seq/cif", - smooth_factor2=1.0, - noise_threshold2=0, - upsample_times=5, - upsample_type="cnn", - use_cif1_cnn=True, - tail_mask=True, - ): + def __init__( + self, + idim, + l_order, + r_order, + threshold=1.0, + dropout=0.1, + smooth_factor=1.0, + noise_threshold=0, + tail_threshold=0.0, + tf2torch_tensor_name_prefix_torch="predictor", + tf2torch_tensor_name_prefix_tf="seq2seq/cif", + smooth_factor2=1.0, + noise_threshold2=0, + upsample_times=5, + upsample_type="cnn", + use_cif1_cnn=True, + tail_mask=True, + ): super(CifPredictorV3, self).__init__() self.pad = torch.nn.ConstantPad1d((l_order, r_order), 0) @@ -128,39 +130,53 @@ class CifPredictorV3(torch.nn.Module): self.upsample_times = upsample_times self.upsample_type = upsample_type self.use_cif1_cnn = use_cif1_cnn - if self.upsample_type == 'cnn': - self.upsample_cnn = torch.nn.ConvTranspose1d(idim, idim, self.upsample_times, self.upsample_times) + if self.upsample_type == "cnn": + self.upsample_cnn = torch.nn.ConvTranspose1d( + idim, idim, self.upsample_times, self.upsample_times + ) self.cif_output2 = torch.nn.Linear(idim, 1) - elif self.upsample_type == 'cnn_blstm': - self.upsample_cnn = torch.nn.ConvTranspose1d(idim, idim, self.upsample_times, self.upsample_times) - self.blstm = torch.nn.LSTM(idim, idim, 1, bias=True, batch_first=True, dropout=0.0, bidirectional=True) - self.cif_output2 = torch.nn.Linear(idim*2, 1) - elif self.upsample_type == 'cnn_attn': - self.upsample_cnn = torch.nn.ConvTranspose1d(idim, idim, self.upsample_times, self.upsample_times) + elif self.upsample_type == "cnn_blstm": + self.upsample_cnn = torch.nn.ConvTranspose1d( + idim, idim, self.upsample_times, self.upsample_times + ) + self.blstm = torch.nn.LSTM( + idim, idim, 1, bias=True, batch_first=True, dropout=0.0, bidirectional=True + ) + self.cif_output2 = torch.nn.Linear(idim * 2, 1) + elif self.upsample_type == "cnn_attn": + self.upsample_cnn = torch.nn.ConvTranspose1d( + idim, idim, self.upsample_times, self.upsample_times + ) from funasr.models.transformer.encoder import EncoderLayer as TransformerEncoderLayer from funasr.models.transformer.attention import MultiHeadedAttention from funasr.models.transformer.positionwise_feed_forward import PositionwiseFeedForward + positionwise_layer_args = ( idim, - idim*2, + idim * 2, 0.1, ) self.self_attn = TransformerEncoderLayer( idim, - MultiHeadedAttention( - 4, idim, 0.1 - ), + MultiHeadedAttention(4, idim, 0.1), PositionwiseFeedForward(*positionwise_layer_args), 0.1, - True, #normalize_before, - False, #concat_after, + True, # normalize_before, + False, # concat_after, ) self.cif_output2 = torch.nn.Linear(idim, 1) self.smooth_factor2 = smooth_factor2 self.noise_threshold2 = noise_threshold2 - def forward(self, hidden, target_label=None, mask=None, ignore_id=-1, mask_chunk_predictor=None, - target_label_length=None): + def forward( + self, + hidden, + target_label=None, + mask=None, + ignore_id=-1, + mask_chunk_predictor=None, + target_label_length=None, + ): h = hidden context = h.transpose(1, 2) queries = self.pad(context) @@ -171,23 +187,27 @@ class CifPredictorV3(torch.nn.Module): _output = context else: _output = output - if self.upsample_type == 'cnn': + if self.upsample_type == "cnn": output2 = self.upsample_cnn(_output) - output2 = output2.transpose(1,2) - elif self.upsample_type == 'cnn_blstm': + output2 = output2.transpose(1, 2) + elif self.upsample_type == "cnn_blstm": output2 = self.upsample_cnn(_output) - output2 = output2.transpose(1,2) + output2 = output2.transpose(1, 2) output2, (_, _) = self.blstm(output2) - elif self.upsample_type == 'cnn_attn': + elif self.upsample_type == "cnn_attn": output2 = self.upsample_cnn(_output) - output2 = output2.transpose(1,2) + output2 = output2.transpose(1, 2) output2, _ = self.self_attn(output2, mask) # import pdb; pdb.set_trace() alphas2 = torch.sigmoid(self.cif_output2(output2)) alphas2 = torch.nn.functional.relu(alphas2 * self.smooth_factor2 - self.noise_threshold2) # repeat the mask in T demension to match the upsampled length if mask is not None: - mask2 = mask.repeat(1, self.upsample_times, 1).transpose(-1, -2).reshape(alphas2.shape[0], -1) + mask2 = ( + mask.repeat(1, self.upsample_times, 1) + .transpose(-1, -2) + .reshape(alphas2.shape[0], -1) + ) mask2 = mask2.unsqueeze(-1) alphas2 = alphas2 * mask2 alphas2 = alphas2.squeeze(-1) @@ -236,22 +256,26 @@ class CifPredictorV3(torch.nn.Module): _output = context else: _output = output - if self.upsample_type == 'cnn': + if self.upsample_type == "cnn": output2 = self.upsample_cnn(_output) - output2 = output2.transpose(1,2) - elif self.upsample_type == 'cnn_blstm': + output2 = output2.transpose(1, 2) + elif self.upsample_type == "cnn_blstm": output2 = self.upsample_cnn(_output) - output2 = output2.transpose(1,2) + output2 = output2.transpose(1, 2) output2, (_, _) = self.blstm(output2) - elif self.upsample_type == 'cnn_attn': + elif self.upsample_type == "cnn_attn": output2 = self.upsample_cnn(_output) - output2 = output2.transpose(1,2) + output2 = output2.transpose(1, 2) output2, _ = self.self_attn(output2, mask) alphas2 = torch.sigmoid(self.cif_output2(output2)) alphas2 = torch.nn.functional.relu(alphas2 * self.smooth_factor2 - self.noise_threshold2) # repeat the mask in T demension to match the upsampled length if mask is not None: - mask2 = mask.repeat(1, self.upsample_times, 1).transpose(-1, -2).reshape(alphas2.shape[0], -1) + mask2 = ( + mask.repeat(1, self.upsample_times, 1) + .transpose(-1, -2) + .reshape(alphas2.shape[0], -1) + ) mask2 = mask2.unsqueeze(-1) alphas2 = alphas2 * mask2 alphas2 = alphas2.squeeze(-1) @@ -289,9 +313,9 @@ class CifPredictorV3(torch.nn.Module): return hidden, alphas, token_num_floor - def gen_frame_alignments(self, - alphas: torch.Tensor = None, - encoder_sequence_length: torch.Tensor = None): + def gen_frame_alignments( + self, alphas: torch.Tensor = None, encoder_sequence_length: torch.Tensor = None + ): batch_size, maximum_length = alphas.size() int_type = torch.int32 @@ -314,11 +338,15 @@ class CifPredictorV3(torch.nn.Module): index_div = torch.floor(torch.true_divide(alphas_cumsum, index)).type(int_type) index_div_bool_zeros = index_div.eq(0) index_div_bool_zeros_count = torch.sum(index_div_bool_zeros, dim=-1) + 1 - index_div_bool_zeros_count = torch.clamp(index_div_bool_zeros_count, 0, encoder_sequence_length.max()) + index_div_bool_zeros_count = torch.clamp( + index_div_bool_zeros_count, 0, encoder_sequence_length.max() + ) token_num_mask = (~make_pad_mask(token_num, maxlen=max_token_num)).to(token_num.device) index_div_bool_zeros_count *= token_num_mask - index_div_bool_zeros_count_tile = index_div_bool_zeros_count[:, :, None].repeat(1, 1, maximum_length) + index_div_bool_zeros_count_tile = index_div_bool_zeros_count[:, :, None].repeat( + 1, 1, maximum_length + ) ones = torch.ones_like(index_div_bool_zeros_count_tile) zeros = torch.zeros_like(index_div_bool_zeros_count_tile) ones = torch.cumsum(ones, dim=2) @@ -329,19 +357,25 @@ class CifPredictorV3(torch.nn.Module): index_div_bool_zeros_count_tile = 1 - index_div_bool_zeros_count_tile_bool.type(int_type) index_div_bool_zeros_count_tile_out = torch.sum(index_div_bool_zeros_count_tile, dim=1) index_div_bool_zeros_count_tile_out = index_div_bool_zeros_count_tile_out.type(int_type) - predictor_mask = (~make_pad_mask(encoder_sequence_length, maxlen=encoder_sequence_length.max())).type( - int_type).to(encoder_sequence_length.device) + predictor_mask = ( + (~make_pad_mask(encoder_sequence_length, maxlen=encoder_sequence_length.max())) + .type(int_type) + .to(encoder_sequence_length.device) + ) index_div_bool_zeros_count_tile_out = index_div_bool_zeros_count_tile_out * predictor_mask predictor_alignments = index_div_bool_zeros_count_tile_out - predictor_alignments_length = predictor_alignments.sum(-1).type(encoder_sequence_length.dtype) + predictor_alignments_length = predictor_alignments.sum(-1).type( + encoder_sequence_length.dtype + ) return predictor_alignments.detach(), predictor_alignments_length.detach() + @tables.register("predictor_classes", "CifPredictorV3Export") class CifPredictorV3Export(torch.nn.Module): def __init__(self, model, **kwargs): super().__init__() - + self.pad = model.pad self.cif_conv1d = model.cif_conv1d self.cif_output = model.cif_output @@ -349,23 +383,25 @@ class CifPredictorV3Export(torch.nn.Module): self.smooth_factor = model.smooth_factor self.noise_threshold = model.noise_threshold self.tail_threshold = model.tail_threshold - + self.upsample_times = model.upsample_times self.upsample_cnn = model.upsample_cnn self.blstm = model.blstm self.cif_output2 = model.cif_output2 self.smooth_factor2 = model.smooth_factor2 self.noise_threshold2 = model.noise_threshold2 - - def forward(self, hidden: torch.Tensor, - mask: torch.Tensor, - ): + + def forward( + self, + hidden: torch.Tensor, + mask: torch.Tensor, + ): h = hidden context = h.transpose(1, 2) queries = self.pad(context) output = torch.relu(self.cif_conv1d(queries)) output = output.transpose(1, 2) - + output = self.cif_output(output) alphas = torch.sigmoid(output) alphas = torch.nn.functional.relu(alphas * self.smooth_factor - self.noise_threshold) @@ -373,18 +409,18 @@ class CifPredictorV3Export(torch.nn.Module): alphas = alphas * mask alphas = alphas.squeeze(-1) token_num = alphas.sum(-1) - + mask = mask.squeeze(-1) hidden, alphas, token_num = self.tail_process_fn(hidden, alphas, mask=mask) acoustic_embeds, cif_peak = cif_export(hidden, alphas, self.threshold) - + return acoustic_embeds, token_num, alphas, cif_peak - + def get_upsample_timestmap(self, hidden, mask=None, token_num=None): h = hidden b = hidden.shape[0] context = h.transpose(1, 2) - + # generate alphas2 _output = context output2 = self.upsample_cnn(_output) @@ -392,8 +428,10 @@ class CifPredictorV3Export(torch.nn.Module): output2, (_, _) = self.blstm(output2) alphas2 = torch.sigmoid(self.cif_output2(output2)) alphas2 = torch.nn.functional.relu(alphas2 * self.smooth_factor2 - self.noise_threshold2) - - mask = mask.repeat(1, self.upsample_times, 1).transpose(-1, -2).reshape(alphas2.shape[0], -1) + + mask = ( + mask.repeat(1, self.upsample_times, 1).transpose(-1, -2).reshape(alphas2.shape[0], -1) + ) mask = mask.unsqueeze(-1) alphas2 = alphas2 * mask alphas2 = alphas2.squeeze(-1) @@ -403,26 +441,26 @@ class CifPredictorV3Export(torch.nn.Module): us_alphas = alphas2 us_cif_peak = cif_wo_hidden_export(us_alphas, self.threshold - 1e-4) return us_alphas, us_cif_peak - + def tail_process_fn(self, hidden, alphas, token_num=None, mask=None): b, t, d = hidden.size() tail_threshold = self.tail_threshold - + zeros_t = torch.zeros((b, 1), dtype=torch.float32, device=alphas.device) ones_t = torch.ones_like(zeros_t) - + mask_1 = torch.cat([mask, zeros_t], dim=1) mask_2 = torch.cat([ones_t, mask], dim=1) mask = mask_2 - mask_1 tail_threshold = mask * tail_threshold alphas = torch.cat([alphas, zeros_t], dim=1) alphas = torch.add(alphas, tail_threshold) - + zeros = torch.zeros((b, 1, d), dtype=hidden.dtype).to(hidden.device) hidden = torch.cat([hidden, zeros], dim=1) token_num = alphas.sum(dim=-1) token_num_floor = torch.floor(token_num) - + return hidden, alphas, token_num_floor @@ -430,39 +468,41 @@ class CifPredictorV3Export(torch.nn.Module): def cif_export(hidden, alphas, threshold: float): batch_size, len_time, hidden_size = hidden.size() threshold = torch.tensor([threshold], dtype=alphas.dtype).to(alphas.device) - + # loop varss integrate = torch.zeros([batch_size], dtype=alphas.dtype, device=hidden.device) frame = torch.zeros([batch_size, hidden_size], dtype=hidden.dtype, device=hidden.device) # intermediate vars along time list_fires = [] list_frames = [] - + for t in range(len_time): alpha = alphas[:, t] - distribution_completion = torch.ones([batch_size], dtype=alphas.dtype, device=hidden.device) - integrate - + distribution_completion = ( + torch.ones([batch_size], dtype=alphas.dtype, device=hidden.device) - integrate + ) + integrate += alpha list_fires.append(integrate) - + fire_place = integrate >= threshold - integrate = torch.where(fire_place, - integrate - torch.ones([batch_size], dtype=alphas.dtype, device=hidden.device), - integrate) - cur = torch.where(fire_place, - distribution_completion, - alpha) + integrate = torch.where( + fire_place, + integrate - torch.ones([batch_size], dtype=alphas.dtype, device=hidden.device), + integrate, + ) + cur = torch.where(fire_place, distribution_completion, alpha) remainds = alpha - cur - + frame += cur[:, None] * hidden[:, t, :] list_frames.append(frame) - frame = torch.where(fire_place[:, None].repeat(1, hidden_size), - remainds[:, None] * hidden[:, t, :], - frame) - + frame = torch.where( + fire_place[:, None].repeat(1, hidden_size), remainds[:, None] * hidden[:, t, :], frame + ) + fires = torch.stack(list_fires, 1) frames = torch.stack(list_frames, 1) - + fire_idxs = fires >= threshold frame_fires = torch.zeros_like(hidden) max_label_len = frames[0, fire_idxs[0]].size(0) @@ -470,7 +510,7 @@ def cif_export(hidden, alphas, threshold: float): frame_fire = frames[b, fire_idxs[b]] frame_len = frame_fire.size(0) frame_fires[b, :frame_len, :] = frame_fire - + if frame_len >= max_label_len: max_label_len = frame_len frame_fires = frame_fires[:, :max_label_len, :] @@ -480,22 +520,24 @@ def cif_export(hidden, alphas, threshold: float): @torch.jit.script def cif_wo_hidden_export(alphas, threshold: float): batch_size, len_time = alphas.size() - + # loop varss integrate = torch.zeros([batch_size], dtype=alphas.dtype, device=alphas.device) # intermediate vars along time list_fires = [] - + for t in range(len_time): alpha = alphas[:, t] - + integrate += alpha list_fires.append(integrate) - + fire_place = integrate >= threshold - integrate = torch.where(fire_place, - integrate - torch.ones([batch_size], device=alphas.device) * threshold, - integrate) - + integrate = torch.where( + fire_place, + integrate - torch.ones([batch_size], device=alphas.device) * threshold, + integrate, + ) + fires = torch.stack(list_fires, 1) - return fires \ No newline at end of file + return fires diff --git a/funasr/models/bicif_paraformer/export_meta.py b/funasr/models/bicif_paraformer/export_meta.py index 7ae800ea7..e9d0a252b 100644 --- a/funasr/models/bicif_paraformer/export_meta.py +++ b/funasr/models/bicif_paraformer/export_meta.py @@ -8,6 +8,7 @@ import types from funasr.register import tables + def export_rebuild_model(model, **kwargs): is_onnx = kwargs.get("type", "onnx") == "onnx" encoder_class = tables.encoder_classes.get(kwargs["encoder"] + "Export") @@ -21,7 +22,7 @@ def export_rebuild_model(model, **kwargs): from funasr.utils.torch_function import sequence_mask - model.make_pad_mask = sequence_mask(kwargs['max_seq_len'], flip=False) + model.make_pad_mask = sequence_mask(kwargs["max_seq_len"], flip=False) model.forward = types.MethodType(export_forward, model) model.export_dummy_inputs = types.MethodType(export_dummy_inputs, model) @@ -32,6 +33,7 @@ def export_rebuild_model(model, **kwargs): return model + def export_forward( self, speech: torch.Tensor, @@ -53,39 +55,32 @@ def export_forward( return decoder_out, pre_token_length, us_alphas, us_cif_peak + def export_dummy_inputs(self): speech = torch.randn(2, 30, 560) speech_lengths = torch.tensor([6, 30], dtype=torch.int32) return (speech, speech_lengths) + def export_input_names(self): - return ['speech', 'speech_lengths'] + return ["speech", "speech_lengths"] + def export_output_names(self): - return ['logits', 'token_num', 'us_alphas', 'us_cif_peak'] + return ["logits", "token_num", "us_alphas", "us_cif_peak"] + def export_dynamic_axes(self): return { - 'speech': { - 0: 'batch_size', - 1: 'feats_length' - }, - 'speech_lengths': { - 0: 'batch_size', - }, - 'logits': { - 0: 'batch_size', - 1: 'logits_length' - }, - 'us_alphas': { - 0: 'batch_size', - 1: 'alphas_length' - }, - 'us_cif_peak': { - 0: 'batch_size', - 1: 'alphas_length' + "speech": {0: "batch_size", 1: "feats_length"}, + "speech_lengths": { + 0: "batch_size", }, + "logits": {0: "batch_size", 1: "logits_length"}, + "us_alphas": {0: "batch_size", 1: "alphas_length"}, + "us_cif_peak": {0: "batch_size", 1: "alphas_length"}, } + def export_name(self): - return "model.onnx" \ No newline at end of file + return "model.onnx" diff --git a/funasr/models/bicif_paraformer/model.py b/funasr/models/bicif_paraformer/model.py index 6f37dd4e0..4db9c76cd 100644 --- a/funasr/models/bicif_paraformer/model.py +++ b/funasr/models/bicif_paraformer/model.py @@ -43,7 +43,7 @@ class BiCifParaformer(Paraformer): Paper2: Achieving timestamp prediction while recognizing with non-autoregressive end-to-end ASR model https://arxiv.org/abs/2301.12343 """ - + def __init__( self, *args, @@ -51,7 +51,6 @@ class BiCifParaformer(Paraformer): ): super().__init__(*args, **kwargs) - def _calc_pre2_loss( self, encoder_out: torch.Tensor, @@ -59,19 +58,21 @@ class BiCifParaformer(Paraformer): ys_pad: torch.Tensor, ys_pad_lens: torch.Tensor, ): - encoder_out_mask = (~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :]).to( - encoder_out.device) + encoder_out_mask = ( + ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :] + ).to(encoder_out.device) if self.predictor_bias == 1: _, ys_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id) ys_pad_lens = ys_pad_lens + self.predictor_bias - _, _, _, _, pre_token_length2 = self.predictor(encoder_out, ys_pad, encoder_out_mask, ignore_id=self.ignore_id) - + _, _, _, _, pre_token_length2 = self.predictor( + encoder_out, ys_pad, encoder_out_mask, ignore_id=self.ignore_id + ) + # loss_pre = self.criterion_pre(ys_pad_lens.type_as(pre_token_length), pre_token_length) loss_pre2 = self.criterion_pre(ys_pad_lens.type_as(pre_token_length2), pre_token_length2) - + return loss_pre2 - - + def _calc_att_loss( self, encoder_out: torch.Tensor, @@ -79,29 +80,29 @@ class BiCifParaformer(Paraformer): ys_pad: torch.Tensor, ys_pad_lens: torch.Tensor, ): - encoder_out_mask = (~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :]).to( - encoder_out.device) + encoder_out_mask = ( + ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :] + ).to(encoder_out.device) if self.predictor_bias == 1: _, ys_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id) ys_pad_lens = ys_pad_lens + self.predictor_bias - pre_acoustic_embeds, pre_token_length, _, pre_peak_index, _ = self.predictor(encoder_out, ys_pad, - encoder_out_mask, - ignore_id=self.ignore_id) - + pre_acoustic_embeds, pre_token_length, _, pre_peak_index, _ = self.predictor( + encoder_out, ys_pad, encoder_out_mask, ignore_id=self.ignore_id + ) + # 0. sampler decoder_out_1st = None if self.sampling_ratio > 0.0: - sematic_embeds, decoder_out_1st = self.sampler(encoder_out, encoder_out_lens, ys_pad, ys_pad_lens, - pre_acoustic_embeds) + sematic_embeds, decoder_out_1st = self.sampler( + encoder_out, encoder_out_lens, ys_pad, ys_pad_lens, pre_acoustic_embeds + ) else: sematic_embeds = pre_acoustic_embeds - + # 1. Forward decoder - decoder_outs = self.decoder( - encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens - ) + decoder_outs = self.decoder(encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens) decoder_out, _ = decoder_outs[0], decoder_outs[1] - + if decoder_out_1st is None: decoder_out_1st = decoder_out # 2. Compute attention loss @@ -112,36 +113,34 @@ class BiCifParaformer(Paraformer): ignore_label=self.ignore_id, ) loss_pre = self.criterion_pre(ys_pad_lens.type_as(pre_token_length), pre_token_length) - + # Compute cer/wer using attention-decoder if self.training or self.error_calculator is None: cer_att, wer_att = None, None else: ys_hat = decoder_out_1st.argmax(dim=-1) cer_att, wer_att = self.error_calculator(ys_hat.cpu(), ys_pad.cpu()) - + return loss_att, acc_att, cer_att, wer_att, loss_pre - def calc_predictor(self, encoder_out, encoder_out_lens): - encoder_out_mask = (~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :]).to( - encoder_out.device) - pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index, pre_token_length2 = self.predictor(encoder_out, - None, - encoder_out_mask, - ignore_id=self.ignore_id) + encoder_out_mask = ( + ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :] + ).to(encoder_out.device) + pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index, pre_token_length2 = ( + self.predictor(encoder_out, None, encoder_out_mask, ignore_id=self.ignore_id) + ) return pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index - def calc_predictor_timestamp(self, encoder_out, encoder_out_lens, token_num): - encoder_out_mask = (~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :]).to( - encoder_out.device) - ds_alphas, ds_cif_peak, us_alphas, us_peaks = self.predictor.get_upsample_timestamp(encoder_out, - encoder_out_mask, - token_num) + encoder_out_mask = ( + ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :] + ).to(encoder_out.device) + ds_alphas, ds_cif_peak, us_alphas, us_peaks = self.predictor.get_upsample_timestamp( + encoder_out, encoder_out_mask, token_num + ) return ds_alphas, ds_cif_peak, us_alphas, us_peaks - - + def forward( self, speech: torch.Tensor, @@ -161,44 +160,48 @@ class BiCifParaformer(Paraformer): text_lengths = text_lengths[:, 0] if len(speech_lengths.size()) > 1: speech_lengths = speech_lengths[:, 0] - + batch_size = speech.shape[0] - + # Encoder encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) - loss_ctc, cer_ctc = None, None loss_pre = None stats = dict() - + # decoder: CTC branch if self.ctc_weight != 0.0: loss_ctc, cer_ctc = self._calc_ctc_loss( encoder_out, encoder_out_lens, text, text_lengths ) - + # Collect CTC branch stats stats["loss_ctc"] = loss_ctc.detach() if loss_ctc is not None else None stats["cer_ctc"] = cer_ctc - # decoder: Attention decoder branch loss_att, acc_att, cer_att, wer_att, loss_pre = self._calc_att_loss( encoder_out, encoder_out_lens, text, text_lengths ) - - loss_pre2 = self._calc_pre2_loss( - encoder_out, encoder_out_lens, text, text_lengths - ) - + + loss_pre2 = self._calc_pre2_loss(encoder_out, encoder_out_lens, text, text_lengths) + # 3. CTC-Att loss definition if self.ctc_weight == 0.0: - loss = loss_att + loss_pre * self.predictor_weight + loss_pre2 * self.predictor_weight * 0.5 + loss = ( + loss_att + + loss_pre * self.predictor_weight + + loss_pre2 * self.predictor_weight * 0.5 + ) else: - loss = self.ctc_weight * loss_ctc + ( - 1 - self.ctc_weight) * loss_att + loss_pre * self.predictor_weight + loss_pre2 * self.predictor_weight * 0.5 - + loss = ( + self.ctc_weight * loss_ctc + + (1 - self.ctc_weight) * loss_att + + loss_pre * self.predictor_weight + + loss_pre2 * self.predictor_weight * 0.5 + ) + # Collect Attn branch stats stats["loss_att"] = loss_att.detach() if loss_att is not None else None stats["acc"] = acc_att @@ -206,34 +209,36 @@ class BiCifParaformer(Paraformer): stats["wer"] = wer_att stats["loss_pre"] = loss_pre.detach().cpu() if loss_pre is not None else None stats["loss_pre2"] = loss_pre2.detach().cpu() - + stats["loss"] = torch.clone(loss.detach()) - + # force_gatherable: to-device and to-tensor if scalar for DataParallel if self.length_normalized_loss: batch_size = int((text_lengths + self.predictor_bias).sum()) - + loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) return loss, stats, weight + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): - def inference(self, - data_in, - data_lengths=None, - key: list = None, - tokenizer=None, - frontend=None, - **kwargs, - ): - # init beamsearch is_use_ctc = kwargs.get("decoding_ctc_weight", 0.0) > 0.00001 and self.ctc != None - is_use_lm = kwargs.get("lm_weight", 0.0) > 0.00001 and kwargs.get("lm_file", None) is not None + is_use_lm = ( + kwargs.get("lm_weight", 0.0) > 0.00001 and kwargs.get("lm_file", None) is not None + ) if self.beam_search is None and (is_use_lm or is_use_ctc): logging.info("enable beam_search") self.init_beam_search(**kwargs) self.nbest = kwargs.get("nbest", 1) - + meta_data = {} # if isinstance(data_in, torch.Tensor): # fbank # speech, speech_lengths = data_in, data_lengths @@ -244,59 +249,70 @@ class BiCifParaformer(Paraformer): # else: # extract fbank feats time1 = time.perf_counter() - audio_sample_list = load_audio_text_image_video(data_in, fs=frontend.fs, audio_fs=kwargs.get("fs", 16000)) + audio_sample_list = load_audio_text_image_video( + data_in, fs=frontend.fs, audio_fs=kwargs.get("fs", 16000) + ) time2 = time.perf_counter() meta_data["load_data"] = f"{time2 - time1:0.3f}" - speech, speech_lengths = extract_fbank(audio_sample_list, data_type=kwargs.get("data_type", "sound"), - frontend=frontend) + speech, speech_lengths = extract_fbank( + audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend + ) time3 = time.perf_counter() meta_data["extract_feat"] = f"{time3 - time2:0.3f}" - meta_data["batch_data_time"] = speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 - + meta_data["batch_data_time"] = ( + speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 + ) + speech = speech.to(device=kwargs["device"]) speech_lengths = speech_lengths.to(device=kwargs["device"]) - + # Encoder encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) if isinstance(encoder_out, tuple): encoder_out = encoder_out[0] - + # predictor predictor_outs = self.calc_predictor(encoder_out, encoder_out_lens) - pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = predictor_outs[0], predictor_outs[1], \ - predictor_outs[2], predictor_outs[3] + pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = ( + predictor_outs[0], + predictor_outs[1], + predictor_outs[2], + predictor_outs[3], + ) pre_token_length = pre_token_length.round().long() if torch.max(pre_token_length) < 1: return [] - decoder_outs = self.cal_decoder_with_predictor(encoder_out, encoder_out_lens, pre_acoustic_embeds, - pre_token_length) + decoder_outs = self.cal_decoder_with_predictor( + encoder_out, encoder_out_lens, pre_acoustic_embeds, pre_token_length + ) decoder_out, ys_pad_lens = decoder_outs[0], decoder_outs[1] - + # BiCifParaformer, test no bias cif2 - _, _, us_alphas, us_peaks = self.calc_predictor_timestamp(encoder_out, encoder_out_lens, - pre_token_length) - + _, _, us_alphas, us_peaks = self.calc_predictor_timestamp( + encoder_out, encoder_out_lens, pre_token_length + ) + results = [] b, n, d = decoder_out.size() for i in range(b): - x = encoder_out[i, :encoder_out_lens[i], :] - am_scores = decoder_out[i, :pre_token_length[i], :] + x = encoder_out[i, : encoder_out_lens[i], :] + am_scores = decoder_out[i, : pre_token_length[i], :] if self.beam_search is not None: nbest_hyps = self.beam_search( - x=x, am_scores=am_scores, maxlenratio=kwargs.get("maxlenratio", 0.0), - minlenratio=kwargs.get("minlenratio", 0.0) + x=x, + am_scores=am_scores, + maxlenratio=kwargs.get("maxlenratio", 0.0), + minlenratio=kwargs.get("minlenratio", 0.0), ) - + nbest_hyps = nbest_hyps[: self.nbest] else: - + yseq = am_scores.argmax(dim=-1) score = am_scores.max(dim=-1)[0] score = torch.sum(score, dim=-1) # pad with mask tokens to ensure compatibility with sos/eos tokens - yseq = torch.tensor( - [self.sos] + yseq.tolist() + [self.eos], device=yseq.device - ) + yseq = torch.tensor([self.sos] + yseq.tolist() + [self.eos], device=yseq.device) nbest_hyps = [Hypothesis(yseq=yseq, score=score)] for nbest_idx, hyp in enumerate(nbest_hyps): ibest_writer = None @@ -304,34 +320,43 @@ class BiCifParaformer(Paraformer): if not hasattr(self, "writer"): self.writer = DatadirWriter(kwargs.get("output_dir")) ibest_writer = self.writer[f"{nbest_idx+1}best_recog"] - + # remove sos/eos and get results last_pos = -1 if isinstance(hyp.yseq, list): token_int = hyp.yseq[1:last_pos] else: token_int = hyp.yseq[1:last_pos].tolist() - + # remove blank symbol id, which is assumed to be 0 - token_int = list(filter(lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int)) - + token_int = list( + filter( + lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int + ) + ) + if tokenizer is not None: # Change integer-ids to tokens token = tokenizer.ids2tokens(token_int) text = tokenizer.tokens2text(token) - - _, timestamp = ts_prediction_lfr6_standard(us_alphas[i][:encoder_out_lens[i] * 3], - us_peaks[i][:encoder_out_lens[i] * 3], - copy.copy(token), - vad_offset=kwargs.get("begin_time", 0)) - - text_postprocessed, time_stamp_postprocessed, word_lists = postprocess_utils.sentence_postprocess( - token, timestamp) - result_i = {"key": key[i], "text": text_postprocessed, - "timestamp": time_stamp_postprocessed, - } - + _, timestamp = ts_prediction_lfr6_standard( + us_alphas[i][: encoder_out_lens[i] * 3], + us_peaks[i][: encoder_out_lens[i] * 3], + copy.copy(token), + vad_offset=kwargs.get("begin_time", 0), + ) + + text_postprocessed, time_stamp_postprocessed, word_lists = ( + postprocess_utils.sentence_postprocess(token, timestamp) + ) + + result_i = { + "key": key[i], + "text": text_postprocessed, + "timestamp": time_stamp_postprocessed, + } + if ibest_writer is not None: ibest_writer["token"][key[i]] = " ".join(token) # ibest_writer["text"][key[i]] = text @@ -340,12 +365,13 @@ class BiCifParaformer(Paraformer): else: result_i = {"key": key[i], "token_int": token_int} results.append(result_i) - + return results, meta_data def export(self, **kwargs): from .export_meta import export_rebuild_model - if 'max_seq_len' not in kwargs: - kwargs['max_seq_len'] = 512 + + if "max_seq_len" not in kwargs: + kwargs["max_seq_len"] = 512 models = export_rebuild_model(model=self, **kwargs) return models diff --git a/funasr/models/branchformer/encoder.py b/funasr/models/branchformer/encoder.py index 4b5b2377d..a15a3f282 100644 --- a/funasr/models/branchformer/encoder.py +++ b/funasr/models/branchformer/encoder.py @@ -45,6 +45,7 @@ from funasr.models.transformer.utils.subsampling import ( from funasr.register import tables + class BranchformerEncoderLayer(torch.nn.Module): """Branchformer encoder layer module. @@ -73,9 +74,7 @@ class BranchformerEncoderLayer(torch.nn.Module): stochastic_depth_rate: float = 0.0, ): super().__init__() - assert (attn is not None) or ( - cgmlp is not None - ), "At least one branch should be valid" + assert (attn is not None) or (cgmlp is not None), "At least one branch should be valid" self.size = size self.attn = attn @@ -111,9 +110,7 @@ class BranchformerEncoderLayer(torch.nn.Module): self.merge_proj = torch.nn.Linear(size, size) elif merge_method == "fixed_ave": - assert ( - 0.0 <= cgmlp_weight <= 1.0 - ), "cgmlp weight should be between 0.0 and 1.0" + assert 0.0 <= cgmlp_weight <= 1.0, "cgmlp weight should be between 0.0 and 1.0" # remove the other branch if only one branch is used if cgmlp_weight == 0.0: @@ -223,14 +220,10 @@ class BranchformerEncoderLayer(torch.nn.Module): ) # (batch, 1, time) if mask is not None: min_value = float( - numpy.finfo( - torch.tensor(0, dtype=score1.dtype).numpy().dtype - ).min + numpy.finfo(torch.tensor(0, dtype=score1.dtype).numpy().dtype).min ) score1 = score1.masked_fill(mask.eq(0), min_value) - score1 = torch.softmax(score1, dim=-1).masked_fill( - mask.eq(0), 0.0 - ) + score1 = torch.softmax(score1, dim=-1).masked_fill(mask.eq(0), 0.0) else: score1 = torch.softmax(score1, dim=-1) pooled1 = torch.matmul(score1, x1).squeeze(1) # (batch, size) @@ -242,14 +235,10 @@ class BranchformerEncoderLayer(torch.nn.Module): ) # (batch, 1, time) if mask is not None: min_value = float( - numpy.finfo( - torch.tensor(0, dtype=score2.dtype).numpy().dtype - ).min + numpy.finfo(torch.tensor(0, dtype=score2.dtype).numpy().dtype).min ) score2 = score2.masked_fill(mask.eq(0), min_value) - score2 = torch.softmax(score2, dim=-1).masked_fill( - mask.eq(0), 0.0 - ) + score2 = torch.softmax(score2, dim=-1).masked_fill(mask.eq(0), 0.0) else: score2 = torch.softmax(score2, dim=-1) pooled2 = torch.matmul(score2, x2).squeeze(1) # (batch, size) @@ -259,19 +248,13 @@ class BranchformerEncoderLayer(torch.nn.Module): merge_weights = torch.softmax( torch.cat([weight1, weight2], dim=-1), dim=-1 ) # (batch, 2) - merge_weights = merge_weights.unsqueeze(-1).unsqueeze( - -1 - ) # (batch, 2, 1, 1) + merge_weights = merge_weights.unsqueeze(-1).unsqueeze(-1) # (batch, 2, 1, 1) w1, w2 = merge_weights[:, 0], merge_weights[:, 1] # (batch, 1, 1) - x = x + stoch_layer_coeff * self.dropout( - self.merge_proj(w1 * x1 + w2 * x2) - ) + x = x + stoch_layer_coeff * self.dropout(self.merge_proj(w1 * x1 + w2 * x2)) elif self.merge_method == "fixed_ave": x = x + stoch_layer_coeff * self.dropout( - self.merge_proj( - (1.0 - self.cgmlp_weight) * x1 + self.cgmlp_weight * x2 - ) + self.merge_proj((1.0 - self.cgmlp_weight) * x1 + self.cgmlp_weight * x2) ) else: raise RuntimeError(f"unknown merge method: {self.merge_method}") @@ -291,6 +274,7 @@ class BranchformerEncoderLayer(torch.nn.Module): return x, mask + @tables.register("encoder_classes", "BranchformerEncoder") class BranchformerEncoder(nn.Module): """Branchformer encoder module.""" @@ -345,9 +329,7 @@ class BranchformerEncoder(nn.Module): elif pos_enc_layer_type == "legacy_rel_pos": assert attention_layer_type == "legacy_rel_selfattn" pos_enc_class = LegacyRelPositionalEncoding - logging.warning( - "Using legacy_rel_pos and it will be deprecated in the future." - ) + logging.warning("Using legacy_rel_pos and it will be deprecated in the future.") else: raise ValueError("unknown pos_enc_layer: " + pos_enc_layer_type) @@ -419,9 +401,7 @@ class BranchformerEncoder(nn.Module): output_size, attention_dropout_rate, ) - logging.warning( - "Using legacy_rel_selfattn and it will be deprecated in the future." - ) + logging.warning("Using legacy_rel_selfattn and it will be deprecated in the future.") elif attention_layer_type == "rel_selfattn": assert pos_enc_layer_type == "rel_pos" encoder_selfattn_layer = RelPositionMultiHeadedAttention @@ -480,9 +460,7 @@ class BranchformerEncoder(nn.Module): num_blocks, lambda lnum: BranchformerEncoderLayer( output_size, - encoder_selfattn_layer(*encoder_selfattn_layer_args) - if use_attn - else None, + encoder_selfattn_layer(*encoder_selfattn_layer_args) if use_attn else None, cgmlp_layer(*cgmlp_layer_args) if use_cgmlp else None, dropout_rate, merge_method, diff --git a/funasr/models/branchformer/fastformer.py b/funasr/models/branchformer/fastformer.py index 24ca9475c..ede94dc30 100644 --- a/funasr/models/branchformer/fastformer.py +++ b/funasr/models/branchformer/fastformer.py @@ -81,14 +81,11 @@ class FastSelfAttention(torch.nn.Module): # (batch, n_heads, time) query_for_score = ( - self.query_att(mixed_query_layer).transpose(1, 2) - / self.attention_head_size**0.5 + self.query_att(mixed_query_layer).transpose(1, 2) / self.attention_head_size**0.5 ) if mask is not None: min_value = float( - numpy.finfo( - torch.tensor(0, dtype=query_for_score.dtype).numpy().dtype - ).min + numpy.finfo(torch.tensor(0, dtype=query_for_score.dtype).numpy().dtype).min ) query_for_score = query_for_score.masked_fill(mask, min_value) query_weight = torch.softmax(query_for_score, dim=-1).masked_fill(mask, 0.0) @@ -108,9 +105,7 @@ class FastSelfAttention(torch.nn.Module): pooled_query = self.dropout(pooled_query) pooled_query_repeat = pooled_query.repeat(1, seq_len, 1) # (batch, time, size) - mixed_query_key_layer = ( - mixed_key_layer * pooled_query_repeat - ) # (batch, time, size) + mixed_query_key_layer = mixed_key_layer * pooled_query_repeat # (batch, time, size) # (batch, n_heads, time) query_key_score = ( @@ -118,14 +113,10 @@ class FastSelfAttention(torch.nn.Module): ).transpose(1, 2) if mask is not None: min_value = float( - numpy.finfo( - torch.tensor(0, dtype=query_key_score.dtype).numpy().dtype - ).min + numpy.finfo(torch.tensor(0, dtype=query_key_score.dtype).numpy().dtype).min ) query_key_score = query_key_score.masked_fill(mask, min_value) - query_key_weight = torch.softmax(query_key_score, dim=-1).masked_fill( - mask, 0.0 - ) + query_key_weight = torch.softmax(query_key_score, dim=-1).masked_fill(mask, 0.0) else: query_key_weight = torch.softmax(query_key_score, dim=-1) @@ -133,9 +124,7 @@ class FastSelfAttention(torch.nn.Module): key_layer = self.transpose_for_scores( mixed_query_key_layer ) # (batch, n_heads, time, attn_dim) - pooled_key = torch.matmul( - query_key_weight, key_layer - ) # (batch, n_heads, 1, attn_dim) + pooled_key = torch.matmul(query_key_weight, key_layer) # (batch, n_heads, 1, attn_dim) pooled_key = self.dropout(pooled_key) # NOTE: value = query, due to param sharing @@ -143,11 +132,8 @@ class FastSelfAttention(torch.nn.Module): 1, 2 ) # (batch, time, n_heads, attn_dim) weighted_value = weighted_value.reshape( - weighted_value.shape[:-2] - + (self.num_attention_heads * self.attention_head_size,) + weighted_value.shape[:-2] + (self.num_attention_heads * self.attention_head_size,) ) # (batch, time, size) - weighted_value = ( - self.dropout(self.transform(weighted_value)) + mixed_query_layer - ) + weighted_value = self.dropout(self.transform(weighted_value)) + mixed_query_layer return weighted_value diff --git a/funasr/models/branchformer/model.py b/funasr/models/branchformer/model.py index 7fa99b3a2..f7e4b3623 100644 --- a/funasr/models/branchformer/model.py +++ b/funasr/models/branchformer/model.py @@ -3,6 +3,7 @@ import logging from funasr.models.transformer.model import Transformer from funasr.register import tables + @tables.register("model_classes", "Branchformer") class Branchformer(Transformer): """CTC-attention hybrid Encoder-Decoder model""" diff --git a/funasr/models/campplus/cluster_backend.py b/funasr/models/campplus/cluster_backend.py index 14fbbe180..b98721eb8 100644 --- a/funasr/models/campplus/cluster_backend.py +++ b/funasr/models/campplus/cluster_backend.py @@ -51,7 +51,7 @@ class SpectralCluster: def p_pruning(self, A): if A.shape[0] * self.pval < 6: - pval = 6. / A.shape[0] + pval = 6.0 / A.shape[0] else: pval = self.pval @@ -80,7 +80,8 @@ class SpectralCluster: num_of_spk = k_oracle else: lambda_gap_list = self.getEigenGaps( - lambdas[self.min_num_spks - 1:self.max_num_spks + 1]) + lambdas[self.min_num_spks - 1 : self.max_num_spks + 1] + ) num_of_spk = np.argmax(lambda_gap_list) + self.min_num_spks emb = eig_vecs[:, :num_of_spk] @@ -105,12 +106,9 @@ class UmapHdbscan: Emphasis On Topological Structure. ICASSP2022 """ - def __init__(self, - n_neighbors=20, - n_components=60, - min_samples=10, - min_cluster_size=10, - metric='cosine'): + def __init__( + self, n_neighbors=20, n_components=60, min_samples=10, min_cluster_size=10, metric="cosine" + ): self.n_neighbors = n_neighbors self.n_components = n_components self.min_samples = min_samples @@ -119,6 +117,7 @@ class UmapHdbscan: def __call__(self, X): import umap.umap_ as umap + umap_X = umap.UMAP( n_neighbors=self.n_neighbors, min_dist=0.0, @@ -128,7 +127,8 @@ class UmapHdbscan: labels = HDBSCAN( min_samples=self.min_samples, min_cluster_size=self.min_cluster_size, - allow_single_cluster=True).fit_predict(umap_X) + allow_single_cluster=True, + ).fit_predict(umap_X) return labels @@ -141,7 +141,7 @@ class ClusterBackend(torch.nn.Module): def __init__(self): super().__init__() - self.model_config = {'merge_thr':0.78} + self.model_config = {"merge_thr": 0.78} # self.other_config = kwargs self.spectral_cluster = SpectralCluster() @@ -149,21 +149,18 @@ class ClusterBackend(torch.nn.Module): def forward(self, X, **params): # clustering and return the labels - k = params['oracle_num'] if 'oracle_num' in params else None - assert len( - X.shape - ) == 2, 'modelscope error: the shape of input should be [N, C]' + k = params["oracle_num"] if "oracle_num" in params else None + assert len(X.shape) == 2, "modelscope error: the shape of input should be [N, C]" if X.shape[0] < 20: - return np.zeros(X.shape[0], dtype='int') + return np.zeros(X.shape[0], dtype="int") if X.shape[0] < 2048 or k is not None: # unexpected corner case labels = self.spectral_cluster(X, k) else: labels = self.umap_hdbscan_cluster(X) - if k is None and 'merge_thr' in self.model_config: - labels = self.merge_by_cos(labels, X, - self.model_config['merge_thr']) + if k is None and "merge_thr" in self.model_config: + labels = self.merge_by_cos(labels, X, self.model_config["merge_thr"]) return labels @@ -180,8 +177,7 @@ class ClusterBackend(torch.nn.Module): spk_center.append(spk_emb) assert len(spk_center) > 0 spk_center = np.stack(spk_center, axis=0) - norm_spk_center = spk_center / np.linalg.norm( - spk_center, axis=1, keepdims=True) + norm_spk_center = spk_center / np.linalg.norm(spk_center, axis=1, keepdims=True) affinity = np.matmul(norm_spk_center, norm_spk_center.T) affinity = np.triu(affinity, 1) spks = np.unravel_index(np.argmax(affinity), affinity.shape) diff --git a/funasr/models/campplus/components.py b/funasr/models/campplus/components.py index 8db9aef75..263b94334 100644 --- a/funasr/models/campplus/components.py +++ b/funasr/models/campplus/components.py @@ -14,30 +14,25 @@ class BasicResBlock(torch.nn.Module): def __init__(self, in_planes, planes, stride=1): super(BasicResBlock, self).__init__() - self.conv1 = torch.nn.Conv2d(in_planes, - planes, - kernel_size=3, - stride=(stride, 1), - padding=1, - bias=False) + self.conv1 = torch.nn.Conv2d( + in_planes, planes, kernel_size=3, stride=(stride, 1), padding=1, bias=False + ) self.bn1 = torch.nn.BatchNorm2d(planes) - self.conv2 = torch.nn.Conv2d(planes, - planes, - kernel_size=3, - stride=1, - padding=1, - bias=False) + self.conv2 = torch.nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False) self.bn2 = torch.nn.BatchNorm2d(planes) self.shortcut = torch.nn.Sequential() if stride != 1 or in_planes != self.expansion * planes: self.shortcut = torch.nn.Sequential( - torch.nn.Conv2d(in_planes, - self.expansion * planes, - kernel_size=1, - stride=(stride, 1), - bias=False), - torch.nn.BatchNorm2d(self.expansion * planes)) + torch.nn.Conv2d( + in_planes, + self.expansion * planes, + kernel_size=1, + stride=(stride, 1), + bias=False, + ), + torch.nn.BatchNorm2d(self.expansion * planes), + ) def forward(self, x): out = F.relu(self.bn1(self.conv1(x))) @@ -48,11 +43,7 @@ class BasicResBlock(torch.nn.Module): class FCM(torch.nn.Module): - def __init__(self, - block=BasicResBlock, - num_blocks=[2, 2], - m_channels=32, - feat_dim=80): + def __init__(self, block=BasicResBlock, num_blocks=[2, 2], m_channels=32, feat_dim=80): super(FCM, self).__init__() self.in_planes = m_channels self.conv1 = torch.nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False) @@ -61,7 +52,9 @@ class FCM(torch.nn.Module): self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=2) self.layer2 = self._make_layer(block, m_channels, num_blocks[0], stride=2) - self.conv2 = torch.nn.Conv2d(m_channels, m_channels, kernel_size=3, stride=(2, 1), padding=1, bias=False) + self.conv2 = torch.nn.Conv2d( + m_channels, m_channels, kernel_size=3, stride=(2, 1), padding=1, bias=False + ) self.bn2 = torch.nn.BatchNorm2d(m_channels) self.out_channels = m_channels * (feat_dim // 8) @@ -87,18 +80,17 @@ class FCM(torch.nn.Module): def get_nonlinear(config_str, channels): nonlinear = torch.nn.Sequential() - for name in config_str.split('-'): - if name == 'relu': - nonlinear.add_module('relu', torch.nn.ReLU(inplace=True)) - elif name == 'prelu': - nonlinear.add_module('prelu', torch.nn.PReLU(channels)) - elif name == 'batchnorm': - nonlinear.add_module('batchnorm', torch.nn.BatchNorm1d(channels)) - elif name == 'batchnorm_': - nonlinear.add_module('batchnorm', - torch.nn.BatchNorm1d(channels, affine=False)) + for name in config_str.split("-"): + if name == "relu": + nonlinear.add_module("relu", torch.nn.ReLU(inplace=True)) + elif name == "prelu": + nonlinear.add_module("prelu", torch.nn.PReLU(channels)) + elif name == "batchnorm": + nonlinear.add_module("batchnorm", torch.nn.BatchNorm1d(channels)) + elif name == "batchnorm_": + nonlinear.add_module("batchnorm", torch.nn.BatchNorm1d(channels, affine=False)) else: - raise ValueError('Unexpected module ({}).'.format(name)) + raise ValueError("Unexpected module ({}).".format(name)) return nonlinear @@ -117,27 +109,32 @@ class StatsPool(torch.nn.Module): class TDNNLayer(torch.nn.Module): - def __init__(self, - in_channels, - out_channels, - kernel_size, - stride=1, - padding=0, - dilation=1, - bias=False, - config_str='batchnorm-relu'): + def __init__( + self, + in_channels, + out_channels, + kernel_size, + stride=1, + padding=0, + dilation=1, + bias=False, + config_str="batchnorm-relu", + ): super(TDNNLayer, self).__init__() if padding < 0: - assert kernel_size % 2 == 1, 'Expect equal paddings, but got even kernel size ({})'.format( - kernel_size) + assert ( + kernel_size % 2 == 1 + ), "Expect equal paddings, but got even kernel size ({})".format(kernel_size) padding = (kernel_size - 1) // 2 * dilation - self.linear = torch.nn.Conv1d(in_channels, - out_channels, - kernel_size, - stride=stride, - padding=padding, - dilation=dilation, - bias=bias) + self.linear = torch.nn.Conv1d( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=bias, + ) self.nonlinear = get_nonlinear(config_str, out_channels) def forward(self, x): @@ -147,23 +144,19 @@ class TDNNLayer(torch.nn.Module): class CAMLayer(torch.nn.Module): - def __init__(self, - bn_channels, - out_channels, - kernel_size, - stride, - padding, - dilation, - bias, - reduction=2): + def __init__( + self, bn_channels, out_channels, kernel_size, stride, padding, dilation, bias, reduction=2 + ): super(CAMLayer, self).__init__() - self.linear_local = torch.nn.Conv1d(bn_channels, - out_channels, - kernel_size, - stride=stride, - padding=padding, - dilation=dilation, - bias=bias) + self.linear_local = torch.nn.Conv1d( + bn_channels, + out_channels, + kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=bias, + ) self.linear1 = torch.nn.Conv1d(bn_channels, bn_channels // reduction, 1) self.relu = torch.nn.ReLU(inplace=True) self.linear2 = torch.nn.Conv1d(bn_channels // reduction, out_channels, 1) @@ -176,45 +169,50 @@ class CAMLayer(torch.nn.Module): m = self.sigmoid(self.linear2(context)) return y * m - def seg_pooling(self, x, seg_len=100, stype='avg'): - if stype == 'avg': + def seg_pooling(self, x, seg_len=100, stype="avg"): + if stype == "avg": seg = F.avg_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True) - elif stype == 'max': + elif stype == "max": seg = F.max_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True) else: - raise ValueError('Wrong segment pooling type.') + raise ValueError("Wrong segment pooling type.") shape = seg.shape seg = seg.unsqueeze(-1).expand(*shape, seg_len).reshape(*shape[:-1], -1) - seg = seg[..., :x.shape[-1]] + seg = seg[..., : x.shape[-1]] return seg class CAMDenseTDNNLayer(torch.nn.Module): - def __init__(self, - in_channels, - out_channels, - bn_channels, - kernel_size, - stride=1, - dilation=1, - bias=False, - config_str='batchnorm-relu', - memory_efficient=False): + def __init__( + self, + in_channels, + out_channels, + bn_channels, + kernel_size, + stride=1, + dilation=1, + bias=False, + config_str="batchnorm-relu", + memory_efficient=False, + ): super(CAMDenseTDNNLayer, self).__init__() - assert kernel_size % 2 == 1, 'Expect equal paddings, but got even kernel size ({})'.format( - kernel_size) + assert kernel_size % 2 == 1, "Expect equal paddings, but got even kernel size ({})".format( + kernel_size + ) padding = (kernel_size - 1) // 2 * dilation self.memory_efficient = memory_efficient self.nonlinear1 = get_nonlinear(config_str, in_channels) self.linear1 = torch.nn.Conv1d(in_channels, bn_channels, 1, bias=False) self.nonlinear2 = get_nonlinear(config_str, bn_channels) - self.cam_layer = CAMLayer(bn_channels, - out_channels, - kernel_size, - stride=stride, - padding=padding, - dilation=dilation, - bias=bias) + self.cam_layer = CAMLayer( + bn_channels, + out_channels, + kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=bias, + ) def bn_function(self, x): return self.linear1(self.nonlinear1(x)) @@ -229,29 +227,33 @@ class CAMDenseTDNNLayer(torch.nn.Module): class CAMDenseTDNNBlock(torch.nn.ModuleList): - def __init__(self, - num_layers, - in_channels, - out_channels, - bn_channels, - kernel_size, - stride=1, - dilation=1, - bias=False, - config_str='batchnorm-relu', - memory_efficient=False): + def __init__( + self, + num_layers, + in_channels, + out_channels, + bn_channels, + kernel_size, + stride=1, + dilation=1, + bias=False, + config_str="batchnorm-relu", + memory_efficient=False, + ): super(CAMDenseTDNNBlock, self).__init__() for i in range(num_layers): - layer = CAMDenseTDNNLayer(in_channels=in_channels + i * out_channels, - out_channels=out_channels, - bn_channels=bn_channels, - kernel_size=kernel_size, - stride=stride, - dilation=dilation, - bias=bias, - config_str=config_str, - memory_efficient=memory_efficient) - self.add_module('tdnnd%d' % (i + 1), layer) + layer = CAMDenseTDNNLayer( + in_channels=in_channels + i * out_channels, + out_channels=out_channels, + bn_channels=bn_channels, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + bias=bias, + config_str=config_str, + memory_efficient=memory_efficient, + ) + self.add_module("tdnnd%d" % (i + 1), layer) def forward(self, x): for layer in self: @@ -260,11 +262,7 @@ class CAMDenseTDNNBlock(torch.nn.ModuleList): class TransitLayer(torch.nn.Module): - def __init__(self, - in_channels, - out_channels, - bias=True, - config_str='batchnorm-relu'): + def __init__(self, in_channels, out_channels, bias=True, config_str="batchnorm-relu"): super(TransitLayer, self).__init__() self.nonlinear = get_nonlinear(config_str, in_channels) self.linear = torch.nn.Conv1d(in_channels, out_channels, 1, bias=bias) @@ -276,11 +274,7 @@ class TransitLayer(torch.nn.Module): class DenseLayer(torch.nn.Module): - def __init__(self, - in_channels, - out_channels, - bias=False, - config_str='batchnorm-relu'): + def __init__(self, in_channels, out_channels, bias=False, config_str="batchnorm-relu"): super(DenseLayer, self).__init__() self.linear = torch.nn.Conv1d(in_channels, out_channels, 1, bias=bias) self.nonlinear = get_nonlinear(config_str, out_channels) @@ -292,5 +286,3 @@ class DenseLayer(torch.nn.Module): x = self.linear(x) x = self.nonlinear(x) return x - - diff --git a/funasr/models/campplus/model.py b/funasr/models/campplus/model.py index b2d65f1b2..e3a829b1f 100644 --- a/funasr/models/campplus/model.py +++ b/funasr/models/campplus/model.py @@ -14,8 +14,15 @@ from distutils.version import LooseVersion from funasr.register import tables from funasr.models.campplus.utils import extract_feature from funasr.utils.load_utils import load_audio_text_image_video -from funasr.models.campplus.components import DenseLayer, StatsPool, \ - TDNNLayer, CAMDenseTDNNBlock, TransitLayer, get_nonlinear, FCM +from funasr.models.campplus.components import ( + DenseLayer, + StatsPool, + TDNNLayer, + CAMDenseTDNNBlock, + TransitLayer, + get_nonlinear, + FCM, +) if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): @@ -29,16 +36,18 @@ else: @tables.register("model_classes", "CAMPPlus") class CAMPPlus(torch.nn.Module): - def __init__(self, - feat_dim=80, - embedding_size=192, - growth_rate=32, - bn_size=4, - init_channels=128, - config_str='batchnorm-relu', - memory_efficient=True, - output_level='segment', - **kwargs,): + def __init__( + self, + feat_dim=80, + embedding_size=192, + growth_rate=32, + bn_size=4, + init_channels=128, + config_str="batchnorm-relu", + memory_efficient=True, + output_level="segment", + **kwargs, + ): super().__init__() self.head = FCM(feat_dim=feat_dim) @@ -46,49 +55,56 @@ class CAMPPlus(torch.nn.Module): self.output_level = output_level self.xvector = torch.nn.Sequential( - OrderedDict([ - - ('tdnn', - TDNNLayer(channels, - init_channels, - 5, - stride=2, - dilation=1, - padding=-1, - config_str=config_str)), - ])) + OrderedDict( + [ + ( + "tdnn", + TDNNLayer( + channels, + init_channels, + 5, + stride=2, + dilation=1, + padding=-1, + config_str=config_str, + ), + ), + ] + ) + ) channels = init_channels - for i, (num_layers, kernel_size, - dilation) in enumerate(zip((12, 24, 16), (3, 3, 3), (1, 2, 2))): - block = CAMDenseTDNNBlock(num_layers=num_layers, - in_channels=channels, - out_channels=growth_rate, - bn_channels=bn_size * growth_rate, - kernel_size=kernel_size, - dilation=dilation, - config_str=config_str, - memory_efficient=memory_efficient) - self.xvector.add_module('block%d' % (i + 1), block) + for i, (num_layers, kernel_size, dilation) in enumerate( + zip((12, 24, 16), (3, 3, 3), (1, 2, 2)) + ): + block = CAMDenseTDNNBlock( + num_layers=num_layers, + in_channels=channels, + out_channels=growth_rate, + bn_channels=bn_size * growth_rate, + kernel_size=kernel_size, + dilation=dilation, + config_str=config_str, + memory_efficient=memory_efficient, + ) + self.xvector.add_module("block%d" % (i + 1), block) channels = channels + num_layers * growth_rate self.xvector.add_module( - 'transit%d' % (i + 1), - TransitLayer(channels, - channels // 2, - bias=False, - config_str=config_str)) + "transit%d" % (i + 1), + TransitLayer(channels, channels // 2, bias=False, config_str=config_str), + ) channels //= 2 - self.xvector.add_module( - 'out_nonlinear', get_nonlinear(config_str, channels)) + self.xvector.add_module("out_nonlinear", get_nonlinear(config_str, channels)) - if self.output_level == 'segment': - self.xvector.add_module('stats', StatsPool()) + if self.output_level == "segment": + self.xvector.add_module("stats", StatsPool()) self.xvector.add_module( - 'dense', - DenseLayer( - channels * 2, embedding_size, config_str='batchnorm_')) + "dense", DenseLayer(channels * 2, embedding_size, config_str="batchnorm_") + ) else: - assert self.output_level == 'frame', '`output_level` should be set to \'segment\' or \'frame\'. ' + assert ( + self.output_level == "frame" + ), "`output_level` should be set to 'segment' or 'frame'. " for m in self.modules(): if isinstance(m, (torch.nn.Conv1d, torch.nn.Linear)): @@ -100,22 +116,25 @@ class CAMPPlus(torch.nn.Module): x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T) x = self.head(x) x = self.xvector(x) - if self.output_level == 'frame': + if self.output_level == "frame": x = x.transpose(1, 2) return x - def inference(self, - data_in, - data_lengths=None, - key: list=None, - tokenizer=None, - frontend=None, - **kwargs, - ): + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): # extract fbank feats meta_data = {} time1 = time.perf_counter() - audio_sample_list = load_audio_text_image_video(data_in, fs=16000, audio_fs=kwargs.get("fs", 16000), data_type="sound") + audio_sample_list = load_audio_text_image_video( + data_in, fs=16000, audio_fs=kwargs.get("fs", 16000), data_type="sound" + ) time2 = time.perf_counter() meta_data["load_data"] = f"{time2 - time1:0.3f}" speech, speech_lengths, speech_times = extract_feature(audio_sample_list) @@ -124,4 +143,4 @@ class CAMPPlus(torch.nn.Module): meta_data["extract_feat"] = f"{time3 - time2:0.3f}" meta_data["batch_data_time"] = np.array(speech_times).sum().item() / 16000.0 results = [{"spk_embedding": self.forward(speech.to(torch.float32))}] - return results, meta_data \ No newline at end of file + return results, meta_data diff --git a/funasr/models/campplus/utils.py b/funasr/models/campplus/utils.py index 399be4b0a..cb8d4b31a 100644 --- a/funasr/models/campplus/utils.py +++ b/funasr/models/campplus/utils.py @@ -25,15 +25,13 @@ def check_audio_list(audio: list): audio_dur = 0 for i in range(len(audio)): seg = audio[i] - assert seg[1] >= seg[0], 'modelscope error: Wrong time stamps.' - assert isinstance(seg[2], np.ndarray), 'modelscope error: Wrong data type.' - assert int(seg[1] * 16000) - int( - seg[0] * 16000 - ) == seg[2].shape[ - 0], 'modelscope error: audio data in list is inconsistent with time length.' + assert seg[1] >= seg[0], "modelscope error: Wrong time stamps." + assert isinstance(seg[2], np.ndarray), "modelscope error: Wrong data type." + assert ( + int(seg[1] * 16000) - int(seg[0] * 16000) == seg[2].shape[0] + ), "modelscope error: audio data in list is inconsistent with time length." if i > 0: - assert seg[0] >= audio[ - i - 1][1], 'modelscope error: Wrong time stamps.' + assert seg[0] >= audio[i - 1][1], "modelscope error: Wrong time stamps." audio_dur += seg[1] - seg[0] return audio_dur # assert audio_dur > 5, 'modelscope error: The effective audio duration is too short.' @@ -44,39 +42,38 @@ def sv_preprocess(inputs: Union[np.ndarray, list]): for i in range(len(inputs)): if isinstance(inputs[i], str): file_bytes = File.read(inputs[i]) - data, fs = sf.load(io.BytesIO(file_bytes), dtype='float32') + data, fs = sf.load(io.BytesIO(file_bytes), dtype="float32") if len(data.shape) == 2: data = data[:, 0] data = torch.from_numpy(data).unsqueeze(0) data = data.squeeze(0) elif isinstance(inputs[i], np.ndarray): - assert len( - inputs[i].shape - ) == 1, 'modelscope error: Input array should be [N, T]' + assert len(inputs[i].shape) == 1, "modelscope error: Input array should be [N, T]" data = inputs[i] - if data.dtype in ['int16', 'int32', 'int64']: - data = (data / (1 << 15)).astype('float32') + if data.dtype in ["int16", "int32", "int64"]: + data = (data / (1 << 15)).astype("float32") else: - data = data.astype('float32') + data = data.astype("float32") data = torch.from_numpy(data) else: raise ValueError( - 'modelscope error: The input type is restricted to audio address and nump array.' + "modelscope error: The input type is restricted to audio address and nump array." ) output.append(data) return output -def sv_chunk(vad_segments: list, fs = 16000) -> list: +def sv_chunk(vad_segments: list, fs=16000) -> list: config = { - 'seg_dur': 1.5, - 'seg_shift': 0.75, - } + "seg_dur": 1.5, + "seg_shift": 0.75, + } + def seg_chunk(seg_data): seg_st = seg_data[0] data = seg_data[2] - chunk_len = int(config['seg_dur'] * fs) - chunk_shift = int(config['seg_shift'] * fs) + chunk_len = int(config["seg_dur"] * fs) + chunk_shift = int(config["seg_shift"] * fs) last_chunk_ed = 0 seg_res = [] for chunk_st in range(0, data.shape[0], chunk_shift): @@ -87,13 +84,8 @@ def sv_chunk(vad_segments: list, fs = 16000) -> list: chunk_st = max(0, chunk_ed - chunk_len) chunk_data = data[chunk_st:chunk_ed] if chunk_data.shape[0] < chunk_len: - chunk_data = np.pad(chunk_data, - (0, chunk_len - chunk_data.shape[0]), - 'constant') - seg_res.append([ - chunk_st / fs + seg_st, chunk_ed / fs + seg_st, - chunk_data - ]) + chunk_data = np.pad(chunk_data, (0, chunk_len - chunk_data.shape[0]), "constant") + seg_res.append([chunk_st / fs + seg_st, chunk_ed / fs + seg_st, chunk_data]) return seg_res segs = [] @@ -108,8 +100,7 @@ def extract_feature(audio): feature_times = [] feature_lengths = [] for au in audio: - feature = Kaldi.fbank( - au.unsqueeze(0), num_mel_bins=80) + feature = Kaldi.fbank(au.unsqueeze(0), num_mel_bins=80) feature = feature - feature.mean(dim=0, keepdim=True) features.append(feature) feature_times.append(au.shape[0]) @@ -120,8 +111,9 @@ def extract_feature(audio): return features_padded, feature_lengths, feature_times -def postprocess(segments: list, vad_segments: list, - labels: np.ndarray, embeddings: np.ndarray) -> list: +def postprocess( + segments: list, vad_segments: list, labels: np.ndarray, embeddings: np.ndarray +) -> list: assert len(segments) == len(labels) labels = correct_labels(labels) distribute_res = [] @@ -166,16 +158,17 @@ def correct_labels(labels): new_labels.append(id2id[i]) return np.array(new_labels) + def merge_seque(distribute_res): res = [distribute_res[0]] for i in range(1, len(distribute_res)): - if distribute_res[i][2] != res[-1][2] or distribute_res[i][ - 0] > res[-1][1]: + if distribute_res[i][2] != res[-1][2] or distribute_res[i][0] > res[-1][1]: res.append(distribute_res[i]) else: res[-1][1] = distribute_res[i][1] return res + def smooth(res, mindur=1): # if only one segment, return directly if len(res) < 2: @@ -202,20 +195,19 @@ def smooth(res, mindur=1): def distribute_spk(sentence_list, sd_time_list): sd_sentence_list = [] for d in sentence_list: - sentence_start = d['start'] - sentence_end = d['end'] + sentence_start = d["start"] + sentence_end = d["end"] sentence_spk = 0 max_overlap = 0 for sd_time in sd_time_list: spk_st, spk_ed, spk = sd_time - spk_st = spk_st*1000 - spk_ed = spk_ed*1000 - overlap = max( - min(sentence_end, spk_ed) - max(sentence_start, spk_st), 0) + spk_st = spk_st * 1000 + spk_ed = spk_ed * 1000 + overlap = max(min(sentence_end, spk_ed) - max(sentence_start, spk_st), 0) if overlap > max_overlap: max_overlap = overlap sentence_spk = spk - d['spk'] = int(sentence_spk) + d["spk"] = int(sentence_spk) sd_sentence_list.append(d) return sd_sentence_list @@ -241,10 +233,7 @@ class Storage(metaclass=ABCMeta): pass @abstractmethod - def write_text(self, - obj: str, - filepath: Union[str, Path], - encoding: str = 'utf-8') -> None: + def write_text(self, obj: str, filepath: Union[str, Path], encoding: str = "utf-8") -> None: pass @@ -260,13 +249,11 @@ class LocalStorage(Storage): Returns: bytes: Expected bytes object. """ - with open(filepath, 'rb') as f: + with open(filepath, "rb") as f: content = f.read() return content - def read_text(self, - filepath: Union[str, Path], - encoding: str = 'utf-8') -> str: + def read_text(self, filepath: Union[str, Path], encoding: str = "utf-8") -> str: """Read data from a given ``filepath`` with 'r' mode. Args: @@ -277,7 +264,7 @@ class LocalStorage(Storage): Returns: str: Expected text reading from ``filepath``. """ - with open(filepath, 'r', encoding=encoding) as f: + with open(filepath, "r", encoding=encoding) as f: value_buf = f.read() return value_buf @@ -296,13 +283,10 @@ class LocalStorage(Storage): if dirname and not os.path.exists(dirname): os.makedirs(dirname, exist_ok=True) - with open(filepath, 'wb') as f: + with open(filepath, "wb") as f: f.write(obj) - def write_text(self, - obj: str, - filepath: Union[str, Path], - encoding: str = 'utf-8') -> None: + def write_text(self, obj: str, filepath: Union[str, Path], encoding: str = "utf-8") -> None: """Write data to a given ``filepath`` with 'w' mode. Note: @@ -319,14 +303,11 @@ class LocalStorage(Storage): if dirname and not os.path.exists(dirname): os.makedirs(dirname, exist_ok=True) - with open(filepath, 'w', encoding=encoding) as f: + with open(filepath, "w", encoding=encoding) as f: f.write(obj) @contextlib.contextmanager - def as_local_path( - self, - filepath: Union[str, - Path]) -> Generator[Union[str, Path], None, None]: + def as_local_path(self, filepath: Union[str, Path]) -> Generator[Union[str, Path], None, None]: """Only for unified API and do nothing.""" yield filepath @@ -346,8 +327,7 @@ class HTTPStorage(Storage): return r.text @contextlib.contextmanager - def as_local_path( - self, filepath: str) -> Generator[Union[str, Path], None, None]: + def as_local_path(self, filepath: str) -> Generator[Union[str, Path], None, None]: """Download a file from ``filepath``. ``as_local_path`` is decorated by :meth:`contextlib.contextmanager`. It @@ -373,14 +353,10 @@ class HTTPStorage(Storage): os.remove(f.name) def write(self, obj: bytes, url: Union[str, Path]) -> None: - raise NotImplementedError('write is not supported by HTTP Storage') + raise NotImplementedError("write is not supported by HTTP Storage") - def write_text(self, - obj: str, - url: Union[str, Path], - encoding: str = 'utf-8') -> None: - raise NotImplementedError( - 'write_text is not supported by HTTP Storage') + def write_text(self, obj: str, url: Union[str, Path], encoding: str = "utf-8") -> None: + raise NotImplementedError("write_text is not supported by HTTP Storage") class OSSStorage(Storage): @@ -388,20 +364,16 @@ class OSSStorage(Storage): def __init__(self, oss_config_file=None): # read from config file or env var - raise NotImplementedError( - 'OSSStorage.__init__ to be implemented in the future') + raise NotImplementedError("OSSStorage.__init__ to be implemented in the future") def read(self, filepath): - raise NotImplementedError( - 'OSSStorage.read to be implemented in the future') + raise NotImplementedError("OSSStorage.read to be implemented in the future") - def read_text(self, filepath, encoding='utf-8'): - raise NotImplementedError( - 'OSSStorage.read_text to be implemented in the future') + def read_text(self, filepath, encoding="utf-8"): + raise NotImplementedError("OSSStorage.read_text to be implemented in the future") @contextlib.contextmanager - def as_local_path( - self, filepath: str) -> Generator[Union[str, Path], None, None]: + def as_local_path(self, filepath: str) -> Generator[Union[str, Path], None, None]: """Download a file from ``filepath``. ``as_local_path`` is decorated by :meth:`contextlib.contextmanager`. It @@ -427,15 +399,10 @@ class OSSStorage(Storage): os.remove(f.name) def write(self, obj: bytes, filepath: Union[str, Path]) -> None: - raise NotImplementedError( - 'OSSStorage.write to be implemented in the future') + raise NotImplementedError("OSSStorage.write to be implemented in the future") - def write_text(self, - obj: str, - filepath: Union[str, Path], - encoding: str = 'utf-8') -> None: - raise NotImplementedError( - 'OSSStorage.write_text to be implemented in the future') + def write_text(self, obj: str, filepath: Union[str, Path], encoding: str = "utf-8") -> None: + raise NotImplementedError("OSSStorage.write_text to be implemented in the future") G_STORAGES = {} @@ -443,27 +410,26 @@ G_STORAGES = {} class File(object): _prefix_to_storage: dict = { - 'oss': OSSStorage, - 'http': HTTPStorage, - 'https': HTTPStorage, - 'local': LocalStorage, + "oss": OSSStorage, + "http": HTTPStorage, + "https": HTTPStorage, + "local": LocalStorage, } @staticmethod def _get_storage(uri): - assert isinstance(uri, - str), f'uri should be str type, but got {type(uri)}' + assert isinstance(uri, str), f"uri should be str type, but got {type(uri)}" - if '://' not in uri: + if "://" not in uri: # local path - storage_type = 'local' + storage_type = "local" else: - prefix, _ = uri.split('://') + prefix, _ = uri.split("://") storage_type = prefix - assert storage_type in File._prefix_to_storage, \ - f'Unsupported uri {uri}, valid prefixs: '\ - f'{list(File._prefix_to_storage.keys())}' + assert storage_type in File._prefix_to_storage, ( + f"Unsupported uri {uri}, valid prefixs: " f"{list(File._prefix_to_storage.keys())}" + ) if storage_type not in G_STORAGES: G_STORAGES[storage_type] = File._prefix_to_storage[storage_type]() @@ -484,7 +450,7 @@ class File(object): return storage.read(uri) @staticmethod - def read_text(uri: Union[str, Path], encoding: str = 'utf-8') -> str: + def read_text(uri: Union[str, Path], encoding: str = "utf-8") -> str: """Read data from a given ``filepath`` with 'r' mode. Args: @@ -514,7 +480,7 @@ class File(object): return storage.write(obj, uri) @staticmethod - def write_text(obj: str, uri: str, encoding: str = 'utf-8') -> None: + def write_text(obj: str, uri: str, encoding: str = "utf-8") -> None: """Write data to a given ``filepath`` with 'w' mode. Note: diff --git a/funasr/models/conformer/encoder.py b/funasr/models/conformer/encoder.py index be973c641..7c939b47c 100644 --- a/funasr/models/conformer/encoder.py +++ b/funasr/models/conformer/encoder.py @@ -49,6 +49,7 @@ from funasr.models.transformer.utils.subsampling import StreamingConvInput from funasr.register import tables import pdb + class ConvolutionModule(nn.Module): """ConvolutionModule in Conformer model. @@ -146,16 +147,16 @@ class EncoderLayer(nn.Module): """ def __init__( - self, - size, - self_attn, - feed_forward, - feed_forward_macaron, - conv_module, - dropout_rate, - normalize_before=True, - concat_after=False, - stochastic_depth_rate=0.0, + self, + size, + self_attn, + feed_forward, + feed_forward_macaron, + conv_module, + dropout_rate, + normalize_before=True, + concat_after=False, + stochastic_depth_rate=0.0, ): """Construct an EncoderLayer object.""" super(EncoderLayer, self).__init__() @@ -266,9 +267,7 @@ class EncoderLayer(nn.Module): residual = x if self.normalize_before: x = self.norm_ff(x) - x = residual + stoch_layer_coeff * self.ff_scale * self.dropout( - self.feed_forward(x) - ) + x = residual + stoch_layer_coeff * self.ff_scale * self.dropout(self.feed_forward(x)) if not self.normalize_before: x = self.norm_ff(x) @@ -321,32 +320,32 @@ class ConformerEncoder(nn.Module): """ def __init__( - self, - input_size: int, - output_size: int = 256, - attention_heads: int = 4, - linear_units: int = 2048, - num_blocks: int = 6, - dropout_rate: float = 0.1, - positional_dropout_rate: float = 0.1, - attention_dropout_rate: float = 0.0, - input_layer: str = "conv2d", - normalize_before: bool = True, - concat_after: bool = False, - positionwise_layer_type: str = "linear", - positionwise_conv_kernel_size: int = 3, - macaron_style: bool = False, - rel_pos_type: str = "legacy", - pos_enc_layer_type: str = "rel_pos", - selfattention_layer_type: str = "rel_selfattn", - activation_type: str = "swish", - use_cnn_module: bool = True, - zero_triu: bool = False, - cnn_module_kernel: int = 31, - padding_idx: int = -1, - interctc_layer_idx: List[int] = [], - interctc_use_conditioning: bool = False, - stochastic_depth_rate: Union[float, List[float]] = 0.0, + self, + input_size: int, + output_size: int = 256, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + attention_dropout_rate: float = 0.0, + input_layer: str = "conv2d", + normalize_before: bool = True, + concat_after: bool = False, + positionwise_layer_type: str = "linear", + positionwise_conv_kernel_size: int = 3, + macaron_style: bool = False, + rel_pos_type: str = "legacy", + pos_enc_layer_type: str = "rel_pos", + selfattention_layer_type: str = "rel_selfattn", + activation_type: str = "swish", + use_cnn_module: bool = True, + zero_triu: bool = False, + cnn_module_kernel: int = 31, + padding_idx: int = -1, + interctc_layer_idx: List[int] = [], + interctc_use_conditioning: bool = False, + stochastic_depth_rate: Union[float, List[float]] = 0.0, ): super().__init__() self._output_size = output_size @@ -373,9 +372,7 @@ class ConformerEncoder(nn.Module): elif pos_enc_layer_type == "legacy_rel_pos": assert selfattention_layer_type == "legacy_rel_selfattn" pos_enc_class = LegacyRelPositionalEncoding - logging.warning( - "Using legacy_rel_pos and it will be deprecated in the future." - ) + logging.warning("Using legacy_rel_pos and it will be deprecated in the future.") else: raise ValueError("unknown pos_enc_layer: " + pos_enc_layer_type) @@ -432,9 +429,7 @@ class ConformerEncoder(nn.Module): pos_enc_class(output_size, positional_dropout_rate), ) elif input_layer is None: - self.embed = torch.nn.Sequential( - pos_enc_class(output_size, positional_dropout_rate) - ) + self.embed = torch.nn.Sequential(pos_enc_class(output_size, positional_dropout_rate)) else: raise ValueError("unknown input_layer: " + input_layer) self.normalize_before = normalize_before @@ -480,9 +475,7 @@ class ConformerEncoder(nn.Module): output_size, attention_dropout_rate, ) - logging.warning( - "Using legacy_rel_selfattn and it will be deprecated in the future." - ) + logging.warning("Using legacy_rel_selfattn and it will be deprecated in the future.") elif selfattention_layer_type == "rel_selfattn": assert pos_enc_layer_type == "rel_pos" encoder_selfattn_layer = RelPositionMultiHeadedAttention @@ -534,11 +527,11 @@ class ConformerEncoder(nn.Module): return self._output_size def forward( - self, - xs_pad: torch.Tensor, - ilens: torch.Tensor, - prev_states: torch.Tensor = None, - ctc: CTC = None, + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, + prev_states: torch.Tensor = None, + ctc: CTC = None, ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: """Calculate forward propagation. @@ -556,11 +549,11 @@ class ConformerEncoder(nn.Module): masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device) if ( - isinstance(self.embed, Conv2dSubsampling) - or isinstance(self.embed, Conv2dSubsampling2) - or isinstance(self.embed, Conv2dSubsampling6) - or isinstance(self.embed, Conv2dSubsampling8) - or isinstance(self.embed, Conv2dSubsamplingPad) + isinstance(self.embed, Conv2dSubsampling) + or isinstance(self.embed, Conv2dSubsampling2) + or isinstance(self.embed, Conv2dSubsampling6) + or isinstance(self.embed, Conv2dSubsampling8) + or isinstance(self.embed, Conv2dSubsamplingPad) ): short_status, limit_size = check_short_utt(self.embed, xs_pad.size(1)) if short_status: @@ -611,7 +604,7 @@ class ConformerEncoder(nn.Module): if len(intermediate_outs) > 0: return (xs_pad, intermediate_outs), olens, None return xs_pad, olens, None - + class CausalConvolution(torch.nn.Module): """ConformerConvolution module definition. @@ -708,6 +701,7 @@ class CausalConvolution(torch.nn.Module): return x, cache + class ChunkEncoderLayer(torch.nn.Module): """Chunk Conformer module definition. Args: @@ -797,9 +791,7 @@ class ChunkEncoderLayer(torch.nn.Module): residual = x x = self.norm_macaron(x) - x = residual + self.feed_forward_scale * self.dropout( - self.feed_forward_macaron(x) - ) + x = residual + self.feed_forward_scale * self.dropout(self.feed_forward_macaron(x)) residual = x x = self.norm_self_att(x) @@ -876,9 +868,7 @@ class ChunkEncoderLayer(torch.nn.Module): residual = x x = self.norm_conv(x) - x, conv_cache = self.conv_mod( - x, cache=self.cache[1], right_context=right_context - ) + x, conv_cache = self.conv_mod(x, cache=self.cache[1], right_context=right_context) x = residual + x residual = x @@ -890,6 +880,7 @@ class ChunkEncoderLayer(torch.nn.Module): return x, pos_enc + @tables.register("encoder_classes", "ChunkConformerEncoder") class ConformerChunkEncoder(torch.nn.Module): """Encoder module definition. @@ -940,7 +931,6 @@ class ConformerChunkEncoder(torch.nn.Module): """Construct an Encoder object.""" super().__init__() - self.embed = StreamingConvInput( input_size=input_size, conv_size=output_size, @@ -954,9 +944,7 @@ class ConformerChunkEncoder(torch.nn.Module): positional_dropout_rate, ) - activation = get_activation( - activation_type - ) + activation = get_activation(activation_type) pos_wise_args = ( output_size, @@ -985,7 +973,6 @@ class ConformerChunkEncoder(torch.nn.Module): simplified_att_score, ) - fn_modules = [] for _ in range(num_blocks): module = lambda: ChunkEncoderLayer( @@ -996,7 +983,7 @@ class ConformerChunkEncoder(torch.nn.Module): CausalConvolution(*conv_mod_args), dropout_rate=dropout_rate, ) - fn_modules.append(module) + fn_modules.append(module) self.encoders = MultiBlocks( [fn() for fn in fn_modules], @@ -1040,7 +1027,6 @@ class ConformerChunkEncoder(torch.nn.Module): """ return self.embed.get_size_before_subsampling(size) - def reset_streaming_cache(self, left_context: int, device: torch.device) -> None: """Initialize/Reset encoder streaming cache. Args: @@ -1062,9 +1048,7 @@ class ConformerChunkEncoder(torch.nn.Module): x: Encoder outputs. (B, T_out, D_enc) x_len: Encoder outputs lenghts. (B,) """ - short_status, limit_size = check_short_utt( - self.embed.subsampling_factor, x.size(1) - ) + short_status, limit_size = check_short_utt(self.embed.subsampling_factor, x.size(1)) if short_status: raise TooShortUttError( @@ -1078,7 +1062,10 @@ class ConformerChunkEncoder(torch.nn.Module): if self.unified_model_training: if self.training: - chunk_size = self.default_chunk_size + torch.randint(-self.jitter_range, self.jitter_range+1, (1,)).item() + chunk_size = ( + self.default_chunk_size + + torch.randint(-self.jitter_range, self.jitter_range + 1, (1,)).item() + ) else: chunk_size = self.default_chunk_size x, mask = self.embed(x, mask, chunk_size) @@ -1104,9 +1091,9 @@ class ConformerChunkEncoder(torch.nn.Module): olens = mask.eq(0).sum(1) if self.time_reduction_factor > 1: - x_utt = x_utt[:,::self.time_reduction_factor,:] - x_chunk = x_chunk[:,::self.time_reduction_factor,:] - olens = torch.floor_divide(olens-1, self.time_reduction_factor) + 1 + x_utt = x_utt[:, :: self.time_reduction_factor, :] + x_chunk = x_chunk[:, :: self.time_reduction_factor, :] + olens = torch.floor_divide(olens - 1, self.time_reduction_factor) + 1 return x_utt, x_chunk, olens @@ -1144,8 +1131,8 @@ class ConformerChunkEncoder(torch.nn.Module): olens = mask.eq(0).sum(1) if self.time_reduction_factor > 1: - x = x[:,::self.time_reduction_factor,:] - olens = torch.floor_divide(olens-1, self.time_reduction_factor) + 1 + x = x[:, :: self.time_reduction_factor, :] + olens = torch.floor_divide(olens - 1, self.time_reduction_factor) + 1 return x, olens, None @@ -1162,9 +1149,7 @@ class ConformerChunkEncoder(torch.nn.Module): x: Encoder outputs. (B, T_out, D_enc) x_len: Encoder outputs lenghts. (B,) """ - short_status, limit_size = check_short_utt( - self.embed.subsampling_factor, x.size(1) - ) + short_status, limit_size = check_short_utt(self.embed.subsampling_factor, x.size(1)) if short_status: raise TooShortUttError( @@ -1185,7 +1170,7 @@ class ConformerChunkEncoder(torch.nn.Module): ) if self.time_reduction_factor > 1: - x_utt = x_utt[:,::self.time_reduction_factor,:] + x_utt = x_utt[:, :: self.time_reduction_factor, :] return x_utt def simu_chunk_forward( @@ -1196,9 +1181,7 @@ class ConformerChunkEncoder(torch.nn.Module): left_context: int = 32, right_context: int = 0, ) -> torch.Tensor: - short_status, limit_size = check_short_utt( - self.embed.subsampling_factor, x.size(1) - ) + short_status, limit_size = check_short_utt(self.embed.subsampling_factor, x.size(1)) if short_status: raise TooShortUttError( @@ -1227,7 +1210,7 @@ class ConformerChunkEncoder(torch.nn.Module): ) olens = mask.eq(0).sum(1) if self.time_reduction_factor > 1: - x = x[:,::self.time_reduction_factor,:] + x = x[:, :: self.time_reduction_factor, :] return x @@ -1255,9 +1238,7 @@ class ConformerChunkEncoder(torch.nn.Module): if left_context > 0: processed_mask = ( - torch.arange(left_context, device=x.device) - .view(1, left_context) - .flip(1) + torch.arange(left_context, device=x.device).view(1, left_context).flip(1) ) processed_mask = processed_mask >= processed_frames mask = torch.cat([processed_mask, mask], dim=1) @@ -1275,5 +1256,5 @@ class ConformerChunkEncoder(torch.nn.Module): x = x[:, 0:-right_context, :] if self.time_reduction_factor > 1: - x = x[:,::self.time_reduction_factor,:] + x = x[:, :: self.time_reduction_factor, :] return x diff --git a/funasr/models/conformer/model.py b/funasr/models/conformer/model.py index 171014b80..d7e77b155 100644 --- a/funasr/models/conformer/model.py +++ b/funasr/models/conformer/model.py @@ -5,11 +5,11 @@ import torch from funasr.models.transformer.model import Transformer from funasr.register import tables + @tables.register("model_classes", "Conformer") class Conformer(Transformer): """CTC-attention hybrid Encoder-Decoder model""" - def __init__( self, *args, diff --git a/funasr/models/conformer_rwkv/decoder.py b/funasr/models/conformer_rwkv/decoder.py index d7f113dc8..90e56e59c 100644 --- a/funasr/models/conformer_rwkv/decoder.py +++ b/funasr/models/conformer_rwkv/decoder.py @@ -28,6 +28,7 @@ from funasr.models.transformer.scorers.scorer_interface import BatchScorerInterf from omegaconf import OmegaConf from funasr.register import tables + class DecoderLayer(nn.Module): """Single decoder layer module. @@ -51,16 +52,16 @@ class DecoderLayer(nn.Module): """ def __init__( - self, - size, - self_attn, - src_attn, - feed_forward, - dropout_rate, - normalize_before=True, - concat_after=False, - layer_id=None, - args={}, + self, + size, + self_attn, + src_attn, + feed_forward, + dropout_rate, + normalize_before=True, + concat_after=False, + layer_id=None, + args={}, ): """Construct an DecoderLayer object.""" super(DecoderLayer, self).__init__() @@ -86,7 +87,7 @@ class DecoderLayer(nn.Module): layer_id = 0 scale = ((1 + layer_id) / args.get("n_layer")) ** 0.7 nn.init.constant_(self.ln0.weight, scale) - + # init if args.get("init_rwkv", True): print("init_rwkv") @@ -115,28 +116,24 @@ class DecoderLayer(nn.Module): if self.layer_id == 0 and self.ln0 is not None: tgt = self.ln0(tgt) - + residual = tgt - - + tgt = self.norm1(tgt) if cache is None: - + x = residual + self.dropout(self.self_attn(tgt, mask=tgt_mask)) else: - + # tgt_q = tgt[:, -1:, :] # residual_q = residual[:, -1:, :] tgt_q_mask = None - + x = residual + self.dropout(self.self_attn(tgt, mask=tgt_q_mask)) x = x[:, -1, :] - - # x = residual + self.dropout(self.self_attn(tgt_q, tgt, tgt, tgt_q_mask)) - residual = x x = self.norm2(x) @@ -145,14 +142,12 @@ class DecoderLayer(nn.Module): x = self.norm3(x) x = residual + self.dropout(self.feed_forward(x)) - if cache is not None: x = torch.cat([cache, x], dim=1) return x, tgt_mask, memory, memory_mask - class BaseTransformerDecoder(nn.Module, BatchScorerInterface): """Base class of Transfomer decoder module. @@ -176,15 +171,15 @@ class BaseTransformerDecoder(nn.Module, BatchScorerInterface): """ def __init__( - self, - vocab_size: int, - encoder_output_size: int, - dropout_rate: float = 0.1, - positional_dropout_rate: float = 0.1, - input_layer: str = "embed", - use_output_layer: bool = True, - pos_enc_class=PositionalEncoding, - normalize_before: bool = True, + self, + vocab_size: int, + encoder_output_size: int, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, ): super().__init__() attention_dim = encoder_output_size @@ -217,11 +212,11 @@ class BaseTransformerDecoder(nn.Module, BatchScorerInterface): self.decoders = None def forward( - self, - hs_pad: torch.Tensor, - hlens: torch.Tensor, - ys_in_pad: torch.Tensor, - ys_in_lens: torch.Tensor, + self, + hs_pad: torch.Tensor, + hlens: torch.Tensor, + ys_in_pad: torch.Tensor, + ys_in_lens: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: """Forward decoder. @@ -249,20 +244,14 @@ class BaseTransformerDecoder(nn.Module, BatchScorerInterface): tgt_mask = tgt_mask & m memory = hs_pad - memory_mask = (~make_pad_mask(hlens, maxlen=memory.size(1)))[:, None, :].to( - memory.device - ) + memory_mask = (~make_pad_mask(hlens, maxlen=memory.size(1)))[:, None, :].to(memory.device) # Padding for Longformer if memory_mask.shape[-1] != memory.shape[1]: padlen = memory.shape[1] - memory_mask.shape[-1] - memory_mask = torch.nn.functional.pad( - memory_mask, (0, padlen), "constant", False - ) + memory_mask = torch.nn.functional.pad(memory_mask, (0, padlen), "constant", False) x = self.embed(tgt) - x, tgt_mask, memory, memory_mask = self.decoders( - x, tgt_mask, memory, memory_mask - ) + x, tgt_mask, memory, memory_mask = self.decoders(x, tgt_mask, memory, memory_mask) if self.normalize_before: x = self.after_norm(x) if self.output_layer is not None: @@ -272,11 +261,11 @@ class BaseTransformerDecoder(nn.Module, BatchScorerInterface): return x, olens def forward_one_step( - self, - tgt: torch.Tensor, - tgt_mask: torch.Tensor, - memory: torch.Tensor, - cache: List[torch.Tensor] = None, + self, + tgt: torch.Tensor, + tgt_mask: torch.Tensor, + memory: torch.Tensor, + cache: List[torch.Tensor] = None, ) -> Tuple[torch.Tensor, List[torch.Tensor]]: """Forward one step. @@ -296,9 +285,7 @@ class BaseTransformerDecoder(nn.Module, BatchScorerInterface): cache = [None] * len(self.decoders) new_cache = [] for c, decoder in zip(cache, self.decoders): - x, tgt_mask, memory, memory_mask = decoder( - x, tgt_mask, memory, None, cache=c - ) + x, tgt_mask, memory, memory_mask = decoder(x, tgt_mask, memory, None, cache=c) new_cache.append(x) if self.normalize_before: @@ -313,13 +300,11 @@ class BaseTransformerDecoder(nn.Module, BatchScorerInterface): def score(self, ys, state, x): """Score.""" ys_mask = subsequent_mask(len(ys), device=x.device).unsqueeze(0) - logp, state = self.forward_one_step( - ys.unsqueeze(0), ys_mask, x.unsqueeze(0), cache=state - ) + logp, state = self.forward_one_step(ys.unsqueeze(0), ys_mask, x.unsqueeze(0), cache=state) return logp.squeeze(0), state def batch_score( - self, ys: torch.Tensor, states: List[Any], xs: torch.Tensor + self, ys: torch.Tensor, states: List[Any], xs: torch.Tensor ) -> Tuple[torch.Tensor, List[Any]]: """Score new token batch. @@ -343,8 +328,7 @@ class BaseTransformerDecoder(nn.Module, BatchScorerInterface): else: # transpose state of [batch, layer] into [layer, batch] batch_state = [ - torch.stack([states[b][i] for b in range(n_batch)]) - for i in range(n_layers) + torch.stack([states[b][i] for b in range(n_batch)]) for i in range(n_layers) ] # batch decoding @@ -355,25 +339,26 @@ class BaseTransformerDecoder(nn.Module, BatchScorerInterface): state_list = [[states[i][b] for i in range(n_layers)] for b in range(n_batch)] return logp, state_list + @tables.register("decoder_classes", "TransformerRWKVDecoder") class TransformerRWKVDecoder(BaseTransformerDecoder): def __init__( - self, - vocab_size: int, - encoder_output_size: int, - attention_heads: int = 4, - linear_units: int = 2048, - num_blocks: int = 6, - dropout_rate: float = 0.1, - positional_dropout_rate: float = 0.1, - self_attention_dropout_rate: float = 0.0, - src_attention_dropout_rate: float = 0.0, - input_layer: str = "embed", - use_output_layer: bool = True, - pos_enc_class=PositionalEncoding, - normalize_before: bool = True, - concat_after: bool = False, - **kwargs, + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + **kwargs, ): super().__init__( vocab_size=vocab_size, @@ -386,6 +371,7 @@ class TransformerRWKVDecoder(BaseTransformerDecoder): normalize_before=normalize_before, ) from funasr.models.sense_voice.rwkv_v6 import RWKVLayer + rwkv_cfg = kwargs.get("rwkv_cfg", {}) args = OmegaConf.create(rwkv_cfg) # self.attn = RWKVLayer(args=args, layer_id=layer_id) @@ -395,9 +381,7 @@ class TransformerRWKVDecoder(BaseTransformerDecoder): lambda lnum: DecoderLayer( attention_dim, RWKVLayer(args=args, layer_id=lnum), - MultiHeadedAttention( - attention_heads, attention_dim, src_attention_dropout_rate - ), + MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate), PositionwiseFeedForward(attention_dim, linear_units, dropout_rate), dropout_rate, normalize_before, @@ -406,18 +390,18 @@ class TransformerRWKVDecoder(BaseTransformerDecoder): args=args, ), ) - + # init if args.get("init_rwkv", True): print("init_rwkv") nn.init.uniform_(self.embed[0].weight, a=-1e-4, b=1e-4) def forward( - self, - hs_pad: torch.Tensor, - hlens: torch.Tensor, - ys_in_pad: torch.Tensor, - ys_in_lens: torch.Tensor, + self, + hs_pad: torch.Tensor, + hlens: torch.Tensor, + ys_in_pad: torch.Tensor, + ys_in_lens: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: """Forward decoder. @@ -445,20 +429,14 @@ class TransformerRWKVDecoder(BaseTransformerDecoder): tgt_mask = tgt_mask & m memory = hs_pad - memory_mask = (~make_pad_mask(hlens, maxlen=memory.size(1)))[:, None, :].to( - memory.device - ) + memory_mask = (~make_pad_mask(hlens, maxlen=memory.size(1)))[:, None, :].to(memory.device) # Padding for Longformer if memory_mask.shape[-1] != memory.shape[1]: padlen = memory.shape[1] - memory_mask.shape[-1] - memory_mask = torch.nn.functional.pad( - memory_mask, (0, padlen), "constant", False - ) + memory_mask = torch.nn.functional.pad(memory_mask, (0, padlen), "constant", False) x = self.embed(tgt) - x, tgt_mask, memory, memory_mask = self.decoders( - x, tgt_mask, memory, memory_mask - ) + x, tgt_mask, memory, memory_mask = self.decoders(x, tgt_mask, memory, memory_mask) if self.normalize_before: x = self.after_norm(x) if self.output_layer is not None: @@ -468,11 +446,11 @@ class TransformerRWKVDecoder(BaseTransformerDecoder): return x, olens def forward_one_step( - self, - tgt: torch.Tensor, - tgt_mask: torch.Tensor, - memory: torch.Tensor, - cache: List[torch.Tensor] = None, + self, + tgt: torch.Tensor, + tgt_mask: torch.Tensor, + memory: torch.Tensor, + cache: List[torch.Tensor] = None, ) -> Tuple[torch.Tensor, List[torch.Tensor]]: """Forward one step. @@ -492,9 +470,7 @@ class TransformerRWKVDecoder(BaseTransformerDecoder): cache = [None] * len(self.decoders) new_cache = [] for c, decoder in zip(cache, self.decoders): - x, tgt_mask, memory, memory_mask = decoder( - x, tgt_mask, memory, None, cache=c - ) + x, tgt_mask, memory, memory_mask = decoder(x, tgt_mask, memory, None, cache=c) new_cache.append(x) if self.normalize_before: @@ -504,4 +480,4 @@ class TransformerRWKVDecoder(BaseTransformerDecoder): if self.output_layer is not None: y = torch.log_softmax(self.output_layer(y), dim=-1) - return y, new_cache \ No newline at end of file + return y, new_cache diff --git a/funasr/models/conformer_rwkv/model.py b/funasr/models/conformer_rwkv/model.py index 171014b80..d7e77b155 100644 --- a/funasr/models/conformer_rwkv/model.py +++ b/funasr/models/conformer_rwkv/model.py @@ -5,11 +5,11 @@ import torch from funasr.models.transformer.model import Transformer from funasr.register import tables + @tables.register("model_classes", "Conformer") class Conformer(Transformer): """CTC-attention hybrid Encoder-Decoder model""" - def __init__( self, *args, diff --git a/funasr/models/contextual_paraformer/decoder.py b/funasr/models/contextual_paraformer/decoder.py index 1116f8440..0b30c990d 100644 --- a/funasr/models/contextual_paraformer/decoder.py +++ b/funasr/models/contextual_paraformer/decoder.py @@ -15,7 +15,10 @@ from funasr.models.transformer.layer_norm import LayerNorm from funasr.models.transformer.embedding import PositionalEncoding from funasr.models.paraformer.decoder import DecoderLayerSANM, ParaformerSANMDecoder from funasr.models.sanm.positionwise_feed_forward import PositionwiseFeedForwardDecoderSANM -from funasr.models.sanm.attention import MultiHeadedAttentionSANMDecoder, MultiHeadedAttentionCrossAtt +from funasr.models.sanm.attention import ( + MultiHeadedAttentionSANMDecoder, + MultiHeadedAttentionCrossAtt, +) class ContextualDecoderLayer(torch.nn.Module): @@ -47,7 +50,14 @@ class ContextualDecoderLayer(torch.nn.Module): self.concat_linear1 = torch.nn.Linear(size + size, size) self.concat_linear2 = torch.nn.Linear(size + size, size) - def forward(self, tgt, tgt_mask, memory, memory_mask, cache=None,): + def forward( + self, + tgt, + tgt_mask, + memory, + memory_mask, + cache=None, + ): # tgt = self.dropout(tgt) if isinstance(tgt, Tuple): tgt, _ = tgt @@ -97,9 +107,10 @@ class ContextualBiasDecoder(torch.nn.Module): if self.src_attn is not None: if self.normalize_before: x = self.norm3(x) - x = self.dropout(self.src_attn(x, memory, memory_mask)) + x = self.dropout(self.src_attn(x, memory, memory_mask)) return x, tgt_mask, memory, memory_mask, cache + @tables.register("decoder_classes", "ContextualParaformerDecoder") class ContextualParaformerDecoder(ParaformerSANMDecoder): """ @@ -107,6 +118,7 @@ class ContextualParaformerDecoder(ParaformerSANMDecoder): Paraformer: Fast and Accurate Parallel Transformer for Non-autoregressive End-to-End Speech Recognition https://arxiv.org/abs/2006.01713 """ + def __init__( self, vocab_size: int, @@ -139,7 +151,7 @@ class ContextualParaformerDecoder(ParaformerSANMDecoder): ) attention_dim = encoder_output_size - if input_layer == 'none': + if input_layer == "none": self.embed = None if input_layer == "embed": self.embed = torch.nn.Sequential( @@ -194,20 +206,20 @@ class ContextualParaformerDecoder(ParaformerSANMDecoder): dropout_rate=dropout_rate, normalize_before=True, ) - self.bias_output = torch.nn.Conv1d(attention_dim*2, attention_dim, 1, bias=False) + self.bias_output = torch.nn.Conv1d(attention_dim * 2, attention_dim, 1, bias=False) self.last_decoder = ContextualDecoderLayer( - attention_dim, - MultiHeadedAttentionSANMDecoder( - attention_dim, self_attention_dropout_rate, kernel_size, sanm_shfit=sanm_shfit - ), - MultiHeadedAttentionCrossAtt( - attention_heads, attention_dim, src_attention_dropout_rate - ), - PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate), - dropout_rate, - normalize_before, - concat_after, - ) + attention_dim, + MultiHeadedAttentionSANMDecoder( + attention_dim, self_attention_dropout_rate, kernel_size, sanm_shfit=sanm_shfit + ), + MultiHeadedAttentionCrossAtt( + attention_heads, attention_dim, src_attention_dropout_rate + ), + PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate), + dropout_rate, + normalize_before, + concat_after, + ) if num_blocks - att_layer_num <= 0: self.decoders2 = None else: @@ -273,31 +285,25 @@ class ContextualParaformerDecoder(ParaformerSANMDecoder): memory_mask = myutils.sequence_mask(hlens, device=memory.device)[:, None, :] x = tgt - x, tgt_mask, memory, memory_mask, _ = self.decoders( - x, tgt_mask, memory, memory_mask - ) - _, _, x_self_attn, x_src_attn = self.last_decoder( - x, tgt_mask, memory, memory_mask - ) + x, tgt_mask, memory, memory_mask, _ = self.decoders(x, tgt_mask, memory, memory_mask) + _, _, x_self_attn, x_src_attn = self.last_decoder(x, tgt_mask, memory, memory_mask) # contextual paraformer related contextual_length = torch.Tensor([contextual_info.shape[1]]).int().repeat(hs_pad.shape[0]) contextual_mask = myutils.sequence_mask(contextual_length, device=memory.device)[:, None, :] - cx, tgt_mask, _, _, _ = self.bias_decoder(x_self_attn, tgt_mask, contextual_info, memory_mask=contextual_mask) + cx, tgt_mask, _, _, _ = self.bias_decoder( + x_self_attn, tgt_mask, contextual_info, memory_mask=contextual_mask + ) if self.bias_output is not None: - x = torch.cat([x_src_attn, cx*clas_scale], dim=2) + x = torch.cat([x_src_attn, cx * clas_scale], dim=2) x = self.bias_output(x.transpose(1, 2)).transpose(1, 2) # 2D -> D x = x_self_attn + self.dropout(x) if self.decoders2 is not None: - x, tgt_mask, memory, memory_mask, _ = self.decoders2( - x, tgt_mask, memory, memory_mask - ) + x, tgt_mask, memory, memory_mask, _ = self.decoders2(x, tgt_mask, memory, memory_mask) - x, tgt_mask, memory, memory_mask, _ = self.decoders3( - x, tgt_mask, memory, memory_mask - ) + x, tgt_mask, memory, memory_mask, _ = self.decoders3(x, tgt_mask, memory, memory_mask) if self.normalize_before: x = self.after_norm(x) olens = tgt_mask.sum(1) @@ -308,20 +314,26 @@ class ContextualParaformerDecoder(ParaformerSANMDecoder): @tables.register("decoder_classes", "ContextualParaformerDecoderExport") class ContextualParaformerDecoderExport(torch.nn.Module): - def __init__(self, model, - max_seq_len=512, - model_name='decoder', - onnx: bool = True, - **kwargs,): + def __init__( + self, + model, + max_seq_len=512, + model_name="decoder", + onnx: bool = True, + **kwargs, + ): super().__init__() from funasr.utils.torch_function import sequence_mask + self.model = model self.make_pad_mask = sequence_mask(max_seq_len, flip=False) - + from funasr.models.sanm.attention import MultiHeadedAttentionSANMDecoderExport from funasr.models.sanm.attention import MultiHeadedAttentionCrossAttExport from funasr.models.paraformer.decoder import DecoderLayerSANMExport - from funasr.models.transformer.positionwise_feed_forward import PositionwiseFeedForwardDecoderSANMExport + from funasr.models.transformer.positionwise_feed_forward import ( + PositionwiseFeedForwardDecoderSANMExport, + ) for i, d in enumerate(self.model.decoders): if isinstance(d.feed_forward, PositionwiseFeedForwardDecoderSANM): @@ -344,27 +356,34 @@ class ContextualParaformerDecoderExport(torch.nn.Module): if isinstance(d.feed_forward, PositionwiseFeedForwardDecoderSANM): d.feed_forward = PositionwiseFeedForwardDecoderSANMExport(d.feed_forward) self.model.decoders3[i] = DecoderLayerSANMExport(d) - + self.output_layer = model.output_layer self.after_norm = model.after_norm self.model_name = model_name # bias decoder if isinstance(self.model.bias_decoder.src_attn, MultiHeadedAttentionCrossAtt): - self.model.bias_decoder.src_attn = MultiHeadedAttentionCrossAttExport(self.model.bias_decoder.src_attn) + self.model.bias_decoder.src_attn = MultiHeadedAttentionCrossAttExport( + self.model.bias_decoder.src_attn + ) self.bias_decoder = self.model.bias_decoder - + # last decoder if isinstance(self.model.last_decoder.src_attn, MultiHeadedAttentionCrossAtt): - self.model.last_decoder.src_attn = MultiHeadedAttentionCrossAttExport(self.model.last_decoder.src_attn) + self.model.last_decoder.src_attn = MultiHeadedAttentionCrossAttExport( + self.model.last_decoder.src_attn + ) if isinstance(self.model.last_decoder.self_attn, MultiHeadedAttentionSANMDecoder): - self.model.last_decoder.self_attn = MultiHeadedAttentionSANMDecoderExport(self.model.last_decoder.self_attn) + self.model.last_decoder.self_attn = MultiHeadedAttentionSANMDecoderExport( + self.model.last_decoder.self_attn + ) if isinstance(self.model.last_decoder.feed_forward, PositionwiseFeedForwardDecoderSANM): - self.model.last_decoder.feed_forward = PositionwiseFeedForwardDecoderSANMExport(self.model.last_decoder.feed_forward) + self.model.last_decoder.feed_forward = PositionwiseFeedForwardDecoderSANMExport( + self.model.last_decoder.feed_forward + ) self.last_decoder = self.model.last_decoder self.bias_output = self.model.bias_output self.dropout = self.model.dropout - def prepare_mask(self, mask): mask_3d_btd = mask[:, :, None] @@ -373,7 +392,7 @@ class ContextualParaformerDecoderExport(torch.nn.Module): elif len(mask.shape) == 3: mask_4d_bhlt = 1 - mask[:, None, :] mask_4d_bhlt = mask_4d_bhlt * -10000.0 - + return mask_3d_btd, mask_4d_bhlt def forward( @@ -396,13 +415,9 @@ class ContextualParaformerDecoderExport(torch.nn.Module): # memory_mask = myutils.sequence_mask(hlens, device=memory.device)[:, None, :] x = tgt - x, tgt_mask, memory, memory_mask, _ = self.model.decoders( - x, tgt_mask, memory, memory_mask - ) + x, tgt_mask, memory, memory_mask, _ = self.model.decoders(x, tgt_mask, memory, memory_mask) - _, _, x_self_attn, x_src_attn = self.last_decoder( - x, tgt_mask, memory, memory_mask - ) + _, _, x_self_attn, x_src_attn = self.last_decoder(x, tgt_mask, memory, memory_mask) # contextual paraformer related contextual_length = torch.Tensor([bias_embed.shape[1]]).int().repeat(hs_pad.shape[0]) @@ -411,7 +426,9 @@ class ContextualParaformerDecoderExport(torch.nn.Module): contextual_mask, _ = self.prepare_mask(contextual_mask) # import pdb; pdb.set_trace() contextual_mask = contextual_mask.transpose(2, 1).unsqueeze(1) - cx, tgt_mask, _, _, _ = self.bias_decoder(x_self_attn, tgt_mask, bias_embed, memory_mask=contextual_mask) + cx, tgt_mask, _, _, _ = self.bias_decoder( + x_self_attn, tgt_mask, bias_embed, memory_mask=contextual_mask + ) if self.bias_output is not None: x = torch.cat([x_src_attn, cx], dim=2) @@ -422,11 +439,8 @@ class ContextualParaformerDecoderExport(torch.nn.Module): x, tgt_mask, memory, memory_mask, _ = self.model.decoders2( x, tgt_mask, memory, memory_mask ) - x, tgt_mask, memory, memory_mask, _ = self.model.decoders3( - x, tgt_mask, memory, memory_mask - ) + x, tgt_mask, memory, memory_mask, _ = self.model.decoders3(x, tgt_mask, memory, memory_mask) x = self.after_norm(x) x = self.output_layer(x) return x, ys_in_lens - diff --git a/funasr/models/contextual_paraformer/export_meta.py b/funasr/models/contextual_paraformer/export_meta.py index 6e1067a60..602057f51 100644 --- a/funasr/models/contextual_paraformer/export_meta.py +++ b/funasr/models/contextual_paraformer/export_meta.py @@ -11,9 +11,7 @@ from funasr.models.seaco_paraformer.export_meta import ContextualEmbedderExport class ContextualEmbedderExport2(ContextualEmbedderExport): - def __init__(self, - model, - **kwargs): + def __init__(self, model, **kwargs): super().__init__(model) self.embedding = model.bias_embed model.bias_encoder.batch_first = False @@ -22,45 +20,56 @@ class ContextualEmbedderExport2(ContextualEmbedderExport): def export_rebuild_model(model, **kwargs): is_onnx = kwargs.get("type", "onnx") == "onnx" - + encoder_class = tables.encoder_classes.get(kwargs["encoder"] + "Export") model.encoder = encoder_class(model.encoder, onnx=is_onnx) predictor_class = tables.predictor_classes.get(kwargs["predictor"] + "Export") model.predictor = predictor_class(model.predictor, onnx=is_onnx) - + # little difference with bias encoder with seaco paraformer embedder_class = ContextualEmbedderExport2 embedder_model = embedder_class(model, onnx=is_onnx) - + if kwargs["decoder"] == "ParaformerSANMDecoder": kwargs["decoder"] = "ParaformerSANMDecoderOnline" decoder_class = tables.decoder_classes.get(kwargs["decoder"] + "Export") model.decoder = decoder_class(model.decoder, onnx=is_onnx) from funasr.utils.torch_function import sequence_mask + model.make_pad_mask = sequence_mask(kwargs["max_seq_len"], flip=False) model.feats_dim = 560 import copy + backbone_model = copy.copy(model) # backbone backbone_model.forward = types.MethodType(export_backbone_forward, backbone_model) - backbone_model.export_dummy_inputs = types.MethodType(export_backbone_dummy_inputs, backbone_model) - backbone_model.export_input_names = types.MethodType(export_backbone_input_names, backbone_model) - backbone_model.export_output_names = types.MethodType(export_backbone_output_names, backbone_model) - backbone_model.export_dynamic_axes = types.MethodType(export_backbone_dynamic_axes, backbone_model) + backbone_model.export_dummy_inputs = types.MethodType( + export_backbone_dummy_inputs, backbone_model + ) + backbone_model.export_input_names = types.MethodType( + export_backbone_input_names, backbone_model + ) + backbone_model.export_output_names = types.MethodType( + export_backbone_output_names, backbone_model + ) + backbone_model.export_dynamic_axes = types.MethodType( + export_backbone_dynamic_axes, backbone_model + ) backbone_model.export_name = types.MethodType(export_backbone_name, backbone_model) - + return backbone_model, embedder_model + def export_backbone_forward( - self, - speech: torch.Tensor, - speech_lengths: torch.Tensor, - bias_embed: torch.Tensor, - ): + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + bias_embed: torch.Tensor, +): batch = {"speech": speech, "speech_lengths": speech_lengths} enc, enc_len = self.encoder(**batch) @@ -73,36 +82,32 @@ def export_backbone_forward( return decoder_out, pre_token_length + def export_backbone_dummy_inputs(self): speech = torch.randn(2, 30, self.feats_dim) speech_lengths = torch.tensor([6, 30], dtype=torch.int32) bias_embed = torch.randn(2, 1, 512) return (speech, speech_lengths, bias_embed) + def export_backbone_input_names(self): - return ['speech', 'speech_lengths', 'bias_embed'] + return ["speech", "speech_lengths", "bias_embed"] + def export_backbone_output_names(self): - return ['logits', 'token_num'] + return ["logits", "token_num"] + def export_backbone_dynamic_axes(self): return { - 'speech': { - 0: 'batch_size', - 1: 'feats_length' - }, - 'speech_lengths': { - 0: 'batch_size', - }, - 'bias_embed': { - 0: 'batch_size', - 1: 'num_hotwords' - }, - 'logits': { - 0: 'batch_size', - 1: 'logits_length' + "speech": {0: "batch_size", 1: "feats_length"}, + "speech_lengths": { + 0: "batch_size", }, + "bias_embed": {0: "batch_size", 1: "num_hotwords"}, + "logits": {0: "batch_size", 1: "logits_length"}, } - + + def export_backbone_name(self): - return 'model.onnx' \ No newline at end of file + return "model.onnx" diff --git a/funasr/models/contextual_paraformer/model.py b/funasr/models/contextual_paraformer/model.py index b9fd3c463..fd8822025 100644 --- a/funasr/models/contextual_paraformer/model.py +++ b/funasr/models/contextual_paraformer/model.py @@ -44,14 +44,14 @@ class ContextualParaformer(Paraformer): FunASR: A Fundamental End-to-End Speech Recognition Toolkit https://arxiv.org/abs/2305.11013 """ - + def __init__( self, *args, **kwargs, ): super().__init__(*args, **kwargs) - + self.target_buffer_length = kwargs.get("target_buffer_length", -1) inner_dim = kwargs.get("inner_dim", 256) bias_encoder_type = kwargs.get("bias_encoder_type", "lstm") @@ -60,14 +60,16 @@ class ContextualParaformer(Paraformer): crit_attn_smooth = kwargs.get("crit_attn_smooth", 0.0) bias_encoder_dropout_rate = kwargs.get("bias_encoder_dropout_rate", 0.0) - if bias_encoder_type == 'lstm': - self.bias_encoder = torch.nn.LSTM(inner_dim, inner_dim, 1, batch_first=True, dropout=bias_encoder_dropout_rate) + if bias_encoder_type == "lstm": + self.bias_encoder = torch.nn.LSTM( + inner_dim, inner_dim, 1, batch_first=True, dropout=bias_encoder_dropout_rate + ) self.bias_embed = torch.nn.Embedding(self.vocab_size, inner_dim) - elif bias_encoder_type == 'mean': + elif bias_encoder_type == "mean": self.bias_embed = torch.nn.Embedding(self.vocab_size, inner_dim) else: logging.error("Unsupport bias encoder type: {}".format(bias_encoder_type)) - + if self.target_buffer_length > 0: self.hotword_buffer = None self.length_record = [] @@ -87,7 +89,7 @@ class ContextualParaformer(Paraformer): **kwargs, ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]: """Frontend + Encoder + Decoder + Calc loss - + Args: speech: (Batch, Length, ...) speech_lengths: (Batch, ) @@ -107,19 +109,19 @@ class ContextualParaformer(Paraformer): encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) loss_ctc, cer_ctc = None, None - + stats = dict() - + # 1. CTC branch if self.ctc_weight != 0.0: loss_ctc, cer_ctc = self._calc_ctc_loss( encoder_out, encoder_out_lens, text, text_lengths ) - + # Collect CTC branch stats stats["loss_ctc"] = loss_ctc.detach() if loss_ctc is not None else None stats["cer_ctc"] = cer_ctc - + # 2b. Attention decoder branch loss_att, acc_att, cer_att, wer_att, loss_pre, loss_ideal = self._calc_att_clas_loss( encoder_out, encoder_out_lens, text, text_lengths, hotword_pad, hotword_lengths @@ -129,27 +131,31 @@ class ContextualParaformer(Paraformer): if self.ctc_weight == 0.0: loss = loss_att + loss_pre * self.predictor_weight else: - loss = self.ctc_weight * loss_ctc + (1 - self.ctc_weight) * loss_att + loss_pre * self.predictor_weight - + loss = ( + self.ctc_weight * loss_ctc + + (1 - self.ctc_weight) * loss_att + + loss_pre * self.predictor_weight + ) + if loss_ideal is not None: loss = loss + loss_ideal * self.crit_attn_weight stats["loss_ideal"] = loss_ideal.detach().cpu() - + # Collect Attn branch stats stats["loss_att"] = loss_att.detach() if loss_att is not None else None stats["acc"] = acc_att stats["cer"] = cer_att stats["wer"] = wer_att stats["loss_pre"] = loss_pre.detach().cpu() if loss_pre is not None else None - + stats["loss"] = torch.clone(loss.detach()) # force_gatherable: to-device and to-tensor if scalar for DataParallel if self.length_normalized_loss: batch_size = int((text_lengths + self.predictor_bias).sum()) - + loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) return loss, stats, weight - + def _calc_att_clas_loss( self, encoder_out: torch.Tensor, @@ -159,15 +165,17 @@ class ContextualParaformer(Paraformer): hotword_pad: torch.Tensor, hotword_lengths: torch.Tensor, ): - encoder_out_mask = (~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :]).to( - encoder_out.device) + encoder_out_mask = ( + ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :] + ).to(encoder_out.device) if self.predictor_bias == 1: _, ys_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id) ys_pad_lens = ys_pad_lens + self.predictor_bias - pre_acoustic_embeds, pre_token_length, _, _ = self.predictor(encoder_out, ys_pad, encoder_out_mask, - ignore_id=self.ignore_id) + pre_acoustic_embeds, pre_token_length, _, _ = self.predictor( + encoder_out, ys_pad, encoder_out_mask, ignore_id=self.ignore_id + ) # -1. bias encoder if self.use_decoder_embedding: hw_embed = self.decoder.embed(hotword_pad) @@ -183,17 +191,27 @@ class ContextualParaformer(Paraformer): decoder_out_1st = None if self.sampling_ratio > 0.0: - sematic_embeds, decoder_out_1st = self.sampler(encoder_out, encoder_out_lens, ys_pad, ys_pad_lens, - pre_acoustic_embeds, contextual_info) + sematic_embeds, decoder_out_1st = self.sampler( + encoder_out, + encoder_out_lens, + ys_pad, + ys_pad_lens, + pre_acoustic_embeds, + contextual_info, + ) else: sematic_embeds = pre_acoustic_embeds # 1. Forward decoder decoder_outs = self.decoder( - encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens, contextual_info=contextual_info + encoder_out, + encoder_out_lens, + sematic_embeds, + ys_pad_lens, + contextual_info=contextual_info, ) decoder_out, _ = decoder_outs[0], decoder_outs[1] - ''' + """ if self.crit_attn_weight > 0 and attn.shape[-1] > 1: ideal_attn = ideal_attn + self.crit_attn_smooth / (self.crit_attn_smooth + 1.0) attn_non_blank = attn[:,:,:,:-1] @@ -201,7 +219,7 @@ class ContextualParaformer(Paraformer): loss_ideal = self.attn_loss(attn_non_blank.max(1)[0], ideal_attn_non_blank.to(attn.device)) else: loss_ideal = None - ''' + """ loss_ideal = None if decoder_out_1st is None: @@ -214,18 +232,28 @@ class ContextualParaformer(Paraformer): ignore_label=self.ignore_id, ) loss_pre = self.criterion_pre(ys_pad_lens.type_as(pre_token_length), pre_token_length) - + # Compute cer/wer using attention-decoder if self.training or self.error_calculator is None: cer_att, wer_att = None, None else: ys_hat = decoder_out_1st.argmax(dim=-1) cer_att, wer_att = self.error_calculator(ys_hat.cpu(), ys_pad.cpu()) - + return loss_att, acc_att, cer_att, wer_att, loss_pre, loss_ideal - - def sampler(self, encoder_out, encoder_out_lens, ys_pad, ys_pad_lens, pre_acoustic_embeds, contextual_info): - tgt_mask = (~make_pad_mask(ys_pad_lens, maxlen=ys_pad_lens.max())[:, :, None]).to(ys_pad.device) + + def sampler( + self, + encoder_out, + encoder_out_lens, + ys_pad, + ys_pad_lens, + pre_acoustic_embeds, + contextual_info, + ): + tgt_mask = (~make_pad_mask(ys_pad_lens, maxlen=ys_pad_lens.max())[:, :, None]).to( + ys_pad.device + ) ys_pad = ys_pad * tgt_mask[:, :, 0] if self.share_embedding: ys_pad_embed = self.decoder.output_layer.weight[ys_pad] @@ -233,7 +261,11 @@ class ContextualParaformer(Paraformer): ys_pad_embed = self.decoder.embed(ys_pad) with torch.no_grad(): decoder_outs = self.decoder( - encoder_out, encoder_out_lens, pre_acoustic_embeds, ys_pad_lens, contextual_info=contextual_info + encoder_out, + encoder_out_lens, + pre_acoustic_embeds, + ys_pad_lens, + contextual_info=contextual_info, ) decoder_out, _ = decoder_outs[0], decoder_outs[1] pred_tokens = decoder_out.argmax(-1) @@ -243,21 +275,35 @@ class ContextualParaformer(Paraformer): input_mask = torch.ones_like(nonpad_positions) bsz, seq_len = ys_pad.size() for li in range(bsz): - target_num = (((seq_lens[li] - same_num[li].sum()).float()) * self.sampling_ratio).long() + target_num = ( + ((seq_lens[li] - same_num[li].sum()).float()) * self.sampling_ratio + ).long() if target_num > 0: - input_mask[li].scatter_(dim=0, - index=torch.randperm(seq_lens[li])[:target_num].to(pre_acoustic_embeds.device), - value=0) + input_mask[li].scatter_( + dim=0, + index=torch.randperm(seq_lens[li])[:target_num].to( + pre_acoustic_embeds.device + ), + value=0, + ) input_mask = input_mask.eq(1) input_mask = input_mask.masked_fill(~nonpad_positions, False) input_mask_expand_dim = input_mask.unsqueeze(2).to(pre_acoustic_embeds.device) - - sematic_embeds = pre_acoustic_embeds.masked_fill(~input_mask_expand_dim, 0) + ys_pad_embed.masked_fill( - input_mask_expand_dim, 0) + + sematic_embeds = pre_acoustic_embeds.masked_fill( + ~input_mask_expand_dim, 0 + ) + ys_pad_embed.masked_fill(input_mask_expand_dim, 0) return sematic_embeds * tgt_mask, decoder_out * tgt_mask - - def cal_decoder_with_predictor(self, encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens, hw_list=None, - clas_scale=1.0): + + def cal_decoder_with_predictor( + self, + encoder_out, + encoder_out_lens, + sematic_embeds, + ys_pad_lens, + hw_list=None, + clas_scale=1.0, + ): if hw_list is None: hw_list = [torch.Tensor([1]).long().to(encoder_out.device)] # empty hotword list hw_list_pad = pad_list(hw_list, 0) @@ -269,63 +315,80 @@ class ContextualParaformer(Paraformer): hw_embed = h_n.repeat(encoder_out.shape[0], 1, 1) else: hw_lengths = [len(i) for i in hw_list] - hw_list_pad = pad_list([torch.Tensor(i).long() for i in hw_list], 0).to(encoder_out.device) + hw_list_pad = pad_list([torch.Tensor(i).long() for i in hw_list], 0).to( + encoder_out.device + ) if self.use_decoder_embedding: hw_embed = self.decoder.embed(hw_list_pad) else: hw_embed = self.bias_embed(hw_list_pad) - hw_embed = torch.nn.utils.rnn.pack_padded_sequence(hw_embed, hw_lengths, batch_first=True, - enforce_sorted=False) + hw_embed = torch.nn.utils.rnn.pack_padded_sequence( + hw_embed, hw_lengths, batch_first=True, enforce_sorted=False + ) _, (h_n, _) = self.bias_encoder(hw_embed) hw_embed = h_n.repeat(encoder_out.shape[0], 1, 1) decoder_outs = self.decoder( - encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens, contextual_info=hw_embed, clas_scale=clas_scale + encoder_out, + encoder_out_lens, + sematic_embeds, + ys_pad_lens, + contextual_info=hw_embed, + clas_scale=clas_scale, ) decoder_out = decoder_outs[0] decoder_out = torch.log_softmax(decoder_out, dim=-1) return decoder_out, ys_pad_lens - - def inference(self, - data_in, - data_lengths=None, - key: list = None, - tokenizer=None, - frontend=None, - **kwargs, - ): + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): # init beamsearch is_use_ctc = kwargs.get("decoding_ctc_weight", 0.0) > 0.00001 and self.ctc != None - is_use_lm = kwargs.get("lm_weight", 0.0) > 0.00001 and kwargs.get("lm_file", None) is not None + is_use_lm = ( + kwargs.get("lm_weight", 0.0) > 0.00001 and kwargs.get("lm_file", None) is not None + ) if self.beam_search is None and (is_use_lm or is_use_ctc): logging.info("enable beam_search") self.init_beam_search(**kwargs) self.nbest = kwargs.get("nbest", 1) meta_data = {} - + # extract fbank feats time1 = time.perf_counter() - audio_sample_list = load_audio_text_image_video(data_in, fs=frontend.fs, audio_fs=kwargs.get("fs", 16000)) + audio_sample_list = load_audio_text_image_video( + data_in, fs=frontend.fs, audio_fs=kwargs.get("fs", 16000) + ) time2 = time.perf_counter() meta_data["load_data"] = f"{time2 - time1:0.3f}" - speech, speech_lengths = extract_fbank(audio_sample_list, data_type=kwargs.get("data_type", "sound"), - frontend=frontend) + speech, speech_lengths = extract_fbank( + audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend + ) time3 = time.perf_counter() meta_data["extract_feat"] = f"{time3 - time2:0.3f}" - meta_data[ - "batch_data_time"] = speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 + meta_data["batch_data_time"] = ( + speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 + ) speech = speech.to(device=kwargs["device"]) speech_lengths = speech_lengths.to(device=kwargs["device"]) # hotword - self.hotword_list = self.generate_hotwords_list(kwargs.get("hotword", None), tokenizer=tokenizer, frontend=frontend) + self.hotword_list = self.generate_hotwords_list( + kwargs.get("hotword", None), tokenizer=tokenizer, frontend=frontend + ) # Encoder encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) @@ -334,40 +397,47 @@ class ContextualParaformer(Paraformer): # predictor predictor_outs = self.calc_predictor(encoder_out, encoder_out_lens) - pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = predictor_outs[0], predictor_outs[1], \ - predictor_outs[2], predictor_outs[3] + pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = ( + predictor_outs[0], + predictor_outs[1], + predictor_outs[2], + predictor_outs[3], + ) pre_token_length = pre_token_length.round().long() if torch.max(pre_token_length) < 1: return [] - - decoder_outs = self.cal_decoder_with_predictor(encoder_out, encoder_out_lens, - pre_acoustic_embeds, - pre_token_length, - hw_list=self.hotword_list, - clas_scale=kwargs.get("clas_scale", 1.0)) + + decoder_outs = self.cal_decoder_with_predictor( + encoder_out, + encoder_out_lens, + pre_acoustic_embeds, + pre_token_length, + hw_list=self.hotword_list, + clas_scale=kwargs.get("clas_scale", 1.0), + ) decoder_out, ys_pad_lens = decoder_outs[0], decoder_outs[1] - + results = [] b, n, d = decoder_out.size() for i in range(b): - x = encoder_out[i, :encoder_out_lens[i], :] - am_scores = decoder_out[i, :pre_token_length[i], :] + x = encoder_out[i, : encoder_out_lens[i], :] + am_scores = decoder_out[i, : pre_token_length[i], :] if self.beam_search is not None: nbest_hyps = self.beam_search( - x=x, am_scores=am_scores, maxlenratio=kwargs.get("maxlenratio", 0.0), - minlenratio=kwargs.get("minlenratio", 0.0) + x=x, + am_scores=am_scores, + maxlenratio=kwargs.get("maxlenratio", 0.0), + minlenratio=kwargs.get("minlenratio", 0.0), ) - + nbest_hyps = nbest_hyps[: self.nbest] else: - + yseq = am_scores.argmax(dim=-1) score = am_scores.max(dim=-1)[0] score = torch.sum(score, dim=-1) # pad with mask tokens to ensure compatibility with sos/eos tokens - yseq = torch.tensor( - [self.sos] + yseq.tolist() + [self.eos], device=yseq.device - ) + yseq = torch.tensor([self.sos] + yseq.tolist() + [self.eos], device=yseq.device) nbest_hyps = [Hypothesis(yseq=yseq, score=score)] for nbest_idx, hyp in enumerate(nbest_hyps): ibest_writer = None @@ -375,26 +445,29 @@ class ContextualParaformer(Paraformer): if not hasattr(self, "writer"): self.writer = DatadirWriter(kwargs.get("output_dir")) ibest_writer = self.writer[f"{nbest_idx + 1}best_recog"] - + # remove sos/eos and get results last_pos = -1 if isinstance(hyp.yseq, list): token_int = hyp.yseq[1:last_pos] else: token_int = hyp.yseq[1:last_pos].tolist() - + # remove blank symbol id, which is assumed to be 0 token_int = list( - filter(lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int)) - + filter( + lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int + ) + ) + if tokenizer is not None: # Change integer-ids to tokens token = tokenizer.ids2tokens(token_int) text = tokenizer.tokens2text(token) - + text_postprocessed, _ = postprocess_utils.sentence_postprocess(token) result_i = {"key": key[i], "text": text_postprocessed} - + if ibest_writer is not None: ibest_writer["token"][key[i]] = " ".join(token) ibest_writer["text"][key[i]] = text @@ -402,7 +475,7 @@ class ContextualParaformer(Paraformer): else: result_i = {"key": key[i], "token_int": token_int} results.append(result_i) - + return results, meta_data def generate_hotwords_list(self, hotword_list_or_file, tokenizer=None, frontend=None): @@ -417,9 +490,9 @@ class ContextualParaformer(Paraformer): value = s[1:] seg_dict[key] = " ".join(value) return seg_dict - + def seg_tokenize(txt, seg_dict): - pattern = re.compile(r'^[\u4E00-\u9FA50-9]+$') + pattern = re.compile(r"^[\u4E00-\u9FA50-9]+$") out_txt = "" for word in txt: word = word.lower() @@ -435,11 +508,11 @@ class ContextualParaformer(Paraformer): else: out_txt += "" + " " return out_txt.strip().split() - + seg_dict = None if frontend.cmvn_file is not None: model_dir = os.path.dirname(frontend.cmvn_file) - seg_dict_file = os.path.join(model_dir, 'seg_dict') + seg_dict_file = os.path.join(model_dir, "seg_dict") if os.path.exists(seg_dict_file): seg_dict = load_seg_dict(seg_dict_file) else: @@ -448,11 +521,11 @@ class ContextualParaformer(Paraformer): if hotword_list_or_file is None: hotword_list = None # for local txt inputs - elif os.path.exists(hotword_list_or_file) and hotword_list_or_file.endswith('.txt'): + elif os.path.exists(hotword_list_or_file) and hotword_list_or_file.endswith(".txt"): logging.info("Attempting to parse hotwords from local txt...") hotword_list = [] hotword_str_list = [] - with codecs.open(hotword_list_or_file, 'r') as fin: + with codecs.open(hotword_list_or_file, "r") as fin: for line in fin.readlines(): hw = line.strip() hw_list = hw.split() @@ -461,11 +534,14 @@ class ContextualParaformer(Paraformer): hotword_str_list.append(hw) hotword_list.append(tokenizer.tokens2ids(hw_list)) hotword_list.append([self.sos]) - hotword_str_list.append('') - logging.info("Initialized hotword list from file: {}, hotword list: {}." - .format(hotword_list_or_file, hotword_str_list)) + hotword_str_list.append("") + logging.info( + "Initialized hotword list from file: {}, hotword list: {}.".format( + hotword_list_or_file, hotword_str_list + ) + ) # for url, download and generate txt - elif hotword_list_or_file.startswith('http'): + elif hotword_list_or_file.startswith("http"): logging.info("Attempting to parse hotwords from url...") work_dir = tempfile.TemporaryDirectory().name if not os.path.exists(work_dir): @@ -476,7 +552,7 @@ class ContextualParaformer(Paraformer): hotword_list_or_file = text_file_path hotword_list = [] hotword_str_list = [] - with codecs.open(hotword_list_or_file, 'r') as fin: + with codecs.open(hotword_list_or_file, "r") as fin: for line in fin.readlines(): hw = line.strip() hw_list = hw.split() @@ -485,11 +561,14 @@ class ContextualParaformer(Paraformer): hotword_str_list.append(hw) hotword_list.append(tokenizer.tokens2ids(hw_list)) hotword_list.append([self.sos]) - hotword_str_list.append('') - logging.info("Initialized hotword list from file: {}, hotword list: {}." - .format(hotword_list_or_file, hotword_str_list)) + hotword_str_list.append("") + logging.info( + "Initialized hotword list from file: {}, hotword list: {}.".format( + hotword_list_or_file, hotword_str_list + ) + ) # for text str input - elif not hotword_list_or_file.endswith('.txt'): + elif not hotword_list_or_file.endswith(".txt"): logging.info("Attempting to parse hotwords as str...") hotword_list = [] hotword_str_list = [] @@ -500,7 +579,7 @@ class ContextualParaformer(Paraformer): hw_list = seg_tokenize(hw_list, seg_dict) hotword_list.append(tokenizer.tokens2ids(hw_list)) hotword_list.append([self.sos]) - hotword_str_list.append('') + hotword_str_list.append("") logging.info("Hotword list: {}.".format(hotword_str_list)) else: hotword_list = None @@ -510,8 +589,9 @@ class ContextualParaformer(Paraformer): self, **kwargs, ): - if 'max_seq_len' not in kwargs: - kwargs['max_seq_len'] = 512 + if "max_seq_len" not in kwargs: + kwargs["max_seq_len"] = 512 from .export_meta import export_rebuild_model + models = export_rebuild_model(model=self, **kwargs) return models diff --git a/funasr/models/ct_transformer/export_meta.py b/funasr/models/ct_transformer/export_meta.py index 691f21124..c222008ce 100644 --- a/funasr/models/ct_transformer/export_meta.py +++ b/funasr/models/ct_transformer/export_meta.py @@ -9,18 +9,18 @@ from funasr.register import tables def export_rebuild_model(model, **kwargs): - + is_onnx = kwargs.get("type", "onnx") == "onnx" encoder_class = tables.encoder_classes.get(kwargs["encoder"] + "Export") model.encoder = encoder_class(model.encoder, onnx=is_onnx) - + model.forward = types.MethodType(export_forward, model) model.export_dummy_inputs = types.MethodType(export_dummy_inputs, model) model.export_input_names = types.MethodType(export_input_names, model) model.export_output_names = types.MethodType(export_output_names, model) model.export_dynamic_axes = types.MethodType(export_dynamic_axes, model) model.export_name = types.MethodType(export_name, model) - + return model @@ -37,31 +37,31 @@ def export_forward(self, inputs: torch.Tensor, text_lengths: torch.Tensor): y = self.decoder(h) return y + def export_dummy_inputs(self): length = 120 text_indexes = torch.randint(0, self.embed.num_embeddings, (2, length)).type(torch.int32) - text_lengths = torch.tensor([length-20, length], dtype=torch.int32) + text_lengths = torch.tensor([length - 20, length], dtype=torch.int32) return (text_indexes, text_lengths) + def export_input_names(self): - return ['inputs', 'text_lengths'] + return ["inputs", "text_lengths"] + def export_output_names(self): - return ['logits'] + return ["logits"] + def export_dynamic_axes(self): return { - 'inputs': { - 0: 'batch_size', - 1: 'feats_length' - }, - 'text_lengths': { - 0: 'batch_size', - }, - 'logits': { - 0: 'batch_size', - 1: 'logits_length' + "inputs": {0: "batch_size", 1: "feats_length"}, + "text_lengths": { + 0: "batch_size", }, + "logits": {0: "batch_size", 1: "logits_length"}, } + + def export_name(self): - return "model.onnx" \ No newline at end of file + return "model.onnx" diff --git a/funasr/models/ct_transformer/model.py b/funasr/models/ct_transformer/model.py index edb398ff3..534c7ab94 100644 --- a/funasr/models/ct_transformer/model.py +++ b/funasr/models/ct_transformer/model.py @@ -17,6 +17,7 @@ from funasr.train_utils.device_funcs import force_gatherable from funasr.utils.load_utils import load_audio_text_image_video from funasr.models.transformer.utils.nets_utils import make_pad_mask from funasr.models.ct_transformer.utils import split_to_mini_sentence, split_words + try: import jieba except: @@ -37,6 +38,7 @@ class CTTransformer(torch.nn.Module): CT-Transformer: Controllable time-delay transformer for real-time punctuation prediction and disfluency detection https://arxiv.org/pdf/2003.01309.pdf """ + def __init__( self, encoder: str = None, @@ -58,8 +60,7 @@ class CTTransformer(torch.nn.Module): punc_size = len(punc_list) if punc_weight is None: punc_weight = [1] * punc_size - - + self.embed = torch.nn.Embedding(vocab_size, embed_unit) encoder_class = tables.encoder_classes.get(encoder) encoder = encoder_class(**encoder_conf) @@ -76,8 +77,6 @@ class CTTransformer(torch.nn.Module): if kwargs.get("jieba_usr_dict", None) is not None: jieba.load_userdict(kwargs["jieba_usr_dict"]) self.jieba_usr_dict = jieba - - def punc_forward(self, text: torch.Tensor, text_lengths: torch.Tensor, **kwargs): """Compute loss value from buffer sequences. @@ -111,12 +110,16 @@ class CTTransformer(torch.nn.Module): """ y = y.unsqueeze(0) - h, _, cache = self.encoder.forward_one_step(self.embed(y), self._target_mask(y), cache=state) + h, _, cache = self.encoder.forward_one_step( + self.embed(y), self._target_mask(y), cache=state + ) h = self.decoder(h[:, -1]) logp = h.log_softmax(dim=-1).squeeze(0) return logp, cache - def batch_score(self, ys: torch.Tensor, states: List[Any], xs: torch.Tensor) -> Tuple[torch.Tensor, List[Any]]: + def batch_score( + self, ys: torch.Tensor, states: List[Any], xs: torch.Tensor + ) -> Tuple[torch.Tensor, List[Any]]: """Score new token batch. Args: @@ -138,10 +141,14 @@ class CTTransformer(torch.nn.Module): batch_state = None else: # transpose state of [batch, layer] into [layer, batch] - batch_state = [torch.stack([states[b][i] for b in range(n_batch)]) for i in range(n_layers)] + batch_state = [ + torch.stack([states[b][i] for b in range(n_batch)]) for i in range(n_layers) + ] # batch decoding - h, _, states = self.encoder.forward_one_step(self.embed(ys), self._target_mask(ys), cache=batch_state) + h, _, states = self.encoder.forward_one_step( + self.embed(ys), self._target_mask(ys), cache=batch_state + ) h = self.decoder(h[:, -1]) logp = h.log_softmax(dim=-1) @@ -171,12 +178,12 @@ class CTTransformer(torch.nn.Module): batch_size = text.size(0) # For data parallel if max_length is None: - text = text[:, :text_lengths.max()] - punc = punc[:, :text_lengths.max()] + text = text[:, : text_lengths.max()] + punc = punc[:, : text_lengths.max()] else: text = text[:, :max_length] punc = punc[:, :max_length] - + if self.with_vad(): # Should be VadRealtimeTransformer assert vad_indexes is not None @@ -184,21 +191,29 @@ class CTTransformer(torch.nn.Module): else: # Should be TargetDelayTransformer, y, _ = self.punc_forward(text, text_lengths) - + # Calc negative log likelihood # nll: (BxL,) if self.training == False: _, indices = y.view(-1, y.shape[-1]).topk(1, dim=1) from sklearn.metrics import f1_score - f1_score = f1_score(punc.view(-1).detach().cpu().numpy(), - indices.squeeze(-1).detach().cpu().numpy(), - average='micro') + + f1_score = f1_score( + punc.view(-1).detach().cpu().numpy(), + indices.squeeze(-1).detach().cpu().numpy(), + average="micro", + ) nll = torch.Tensor([f1_score]).repeat(text_lengths.sum()) return nll, text_lengths else: self.punc_weight = self.punc_weight.to(punc.device) - nll = F.cross_entropy(y.view(-1, y.shape[-1]), punc.view(-1), self.punc_weight, reduction="none", - ignore_index=self.ignore_id) + nll = F.cross_entropy( + y.view(-1, y.shape[-1]), + punc.view(-1), + self.punc_weight, + reduction="none", + ignore_index=self.ignore_id, + ) # nll: (BxL,) -> (BxL,) if max_length is None: nll.masked_fill_(make_pad_mask(text_lengths).to(nll.device).view(-1), 0.0) @@ -211,7 +226,6 @@ class CTTransformer(torch.nn.Module): nll = nll.view(batch_size, -1) return nll, text_lengths - def forward( self, text: torch.Tensor, @@ -225,26 +239,27 @@ class CTTransformer(torch.nn.Module): ntokens = y_lengths.sum() loss = nll.sum() / ntokens stats = dict(loss=loss.detach()) - + # force_gatherable: to-device and to-tensor if scalar for DataParallel loss, stats, weight = force_gatherable((loss, stats, ntokens), loss.device) return loss, stats, weight - - def inference(self, - data_in, - data_lengths=None, - key: list = None, - tokenizer=None, - frontend=None, - **kwargs, - ): + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): assert len(data_in) == 1 text = load_audio_text_image_video(data_in, data_type=kwargs.get("kwargs", "text"))[0] vad_indexes = kwargs.get("vad_indexes", None) # text = data_in[0] # text_lengths = data_lengths[0] if data_lengths is not None else None split_size = kwargs.get("split_size", 20) - + tokens = split_words(text, jieba_usr_dict=self.jieba_usr_dict) tokens_int = tokenizer.encode(tokens) @@ -252,7 +267,7 @@ class CTTransformer(torch.nn.Module): mini_sentences_id = split_to_mini_sentence(tokens_int, split_size) assert len(mini_sentences) == len(mini_sentences_id) cache_sent = [] - cache_sent_id = torch.from_numpy(np.array([], dtype='int32')) + cache_sent_id = torch.from_numpy(np.array([], dtype="int32")) new_mini_sentence = "" new_mini_sentence_punc = [] cache_pop_trigger_limit = 200 @@ -266,7 +281,7 @@ class CTTransformer(torch.nn.Module): mini_sentence_id = np.concatenate((cache_sent_id, mini_sentence_id), axis=0) data = { "text": torch.unsqueeze(torch.from_numpy(mini_sentence_id), 0), - "text_lengths": torch.from_numpy(np.array([len(mini_sentence_id)], dtype='int32')), + "text_lengths": torch.from_numpy(np.array([len(mini_sentence_id)], dtype="int32")), } data = to_device(data, kwargs["device"]) # y, _ = self.wrapped_model(**data) @@ -282,20 +297,27 @@ class CTTransformer(torch.nn.Module): sentenceEnd = -1 last_comma_index = -1 for i in range(len(punctuations) - 2, 1, -1): - if self.punc_list[punctuations[i]] == "。" or self.punc_list[punctuations[i]] == "？": + if ( + self.punc_list[punctuations[i]] == "。" + or self.punc_list[punctuations[i]] == "？" + ): sentenceEnd = i break if last_comma_index < 0 and self.punc_list[punctuations[i]] == "，": last_comma_index = i - if sentenceEnd < 0 and len(mini_sentence) > cache_pop_trigger_limit and last_comma_index >= 0: + if ( + sentenceEnd < 0 + and len(mini_sentence) > cache_pop_trigger_limit + and last_comma_index >= 0 + ): # The sentence it too long, cut off at a comma. sentenceEnd = last_comma_index punctuations[sentenceEnd] = self.sentence_end_id - cache_sent = mini_sentence[sentenceEnd + 1:] - cache_sent_id = mini_sentence_id[sentenceEnd + 1:] - mini_sentence = mini_sentence[0:sentenceEnd + 1] - punctuations = punctuations[0:sentenceEnd + 1] + cache_sent = mini_sentence[sentenceEnd + 1 :] + cache_sent_id = mini_sentence_id[sentenceEnd + 1 :] + mini_sentence = mini_sentence[0 : sentenceEnd + 1] + punctuations = punctuations[0 : sentenceEnd + 1] # if len(punctuations) == 0: # continue @@ -304,13 +326,20 @@ class CTTransformer(torch.nn.Module): new_mini_sentence_punc += [int(x) for x in punctuations_np] words_with_punc = [] for i in range(len(mini_sentence)): - if (i==0 or self.punc_list[punctuations[i-1]] == "。" or self.punc_list[punctuations[i-1]] == "？") and len(mini_sentence[i][0].encode()) == 1: + if ( + i == 0 + or self.punc_list[punctuations[i - 1]] == "。" + or self.punc_list[punctuations[i - 1]] == "？" + ) and len(mini_sentence[i][0].encode()) == 1: mini_sentence[i] = mini_sentence[i].capitalize() if i == 0: if len(mini_sentence[i][0].encode()) == 1: mini_sentence[i] = " " + mini_sentence[i] if i > 0: - if len(mini_sentence[i][0].encode()) == 1 and len(mini_sentence[i - 1][0].encode()) == 1: + if ( + len(mini_sentence[i][0].encode()) == 1 + and len(mini_sentence[i - 1][0].encode()) == 1 + ): mini_sentence[i] = " " + mini_sentence[i] words_with_punc.append(mini_sentence[i]) if self.punc_list[punctuations[i]] != "_": @@ -330,19 +359,37 @@ class CTTransformer(torch.nn.Module): if mini_sentence_i == len(mini_sentences) - 1: if new_mini_sentence[-1] == "，" or new_mini_sentence[-1] == "、": new_mini_sentence_out = new_mini_sentence[:-1] + "。" - new_mini_sentence_punc_out = new_mini_sentence_punc[:-1] + [self.sentence_end_id] + new_mini_sentence_punc_out = new_mini_sentence_punc[:-1] + [ + self.sentence_end_id + ] elif new_mini_sentence[-1] == ",": new_mini_sentence_out = new_mini_sentence[:-1] + "." - new_mini_sentence_punc_out = new_mini_sentence_punc[:-1] + [self.sentence_end_id] - elif new_mini_sentence[-1] != "。" and new_mini_sentence[-1] != "？" and len(new_mini_sentence[-1].encode())!=1: + new_mini_sentence_punc_out = new_mini_sentence_punc[:-1] + [ + self.sentence_end_id + ] + elif ( + new_mini_sentence[-1] != "。" + and new_mini_sentence[-1] != "？" + and len(new_mini_sentence[-1].encode()) != 1 + ): new_mini_sentence_out = new_mini_sentence + "。" - new_mini_sentence_punc_out = new_mini_sentence_punc[:-1] + [self.sentence_end_id] - if len(punctuations): punctuations[-1] = 2 - elif new_mini_sentence[-1] != "." and new_mini_sentence[-1] != "?" and len(new_mini_sentence[-1].encode())==1: + new_mini_sentence_punc_out = new_mini_sentence_punc[:-1] + [ + self.sentence_end_id + ] + if len(punctuations): + punctuations[-1] = 2 + elif ( + new_mini_sentence[-1] != "." + and new_mini_sentence[-1] != "?" + and len(new_mini_sentence[-1].encode()) == 1 + ): new_mini_sentence_out = new_mini_sentence + "." - new_mini_sentence_punc_out = new_mini_sentence_punc[:-1] + [self.sentence_end_id] - if len(punctuations): punctuations[-1] = 2 - # keep a punctuations array for punc segment + new_mini_sentence_punc_out = new_mini_sentence_punc[:-1] + [ + self.sentence_end_id + ] + if len(punctuations): + punctuations[-1] = 2 + # keep a punctuations array for punc segment if punc_array is None: punc_array = punctuations else: @@ -355,22 +402,21 @@ class CTTransformer(torch.nn.Module): new_punc_array = copy.copy(punc_array).tolist() # for i, (token, punc_id) in enumerate(zip(tokens[::-1], punc_array.tolist()[::-1])): for i, token in enumerate(tokens[::-1]): - if '\u0e00' <= token[0] <= '\u9fa5': # ignore en words + if "\u0e00" <= token[0] <= "\u9fa5": # ignore en words if len(token) > 1: num_append = len(token) - 1 ind_append = len_tokens - i - 1 for _ in range(num_append): new_punc_array.insert(ind_append, 1) punc_array = torch.tensor(new_punc_array) - + result_i = {"key": key[0], "text": new_mini_sentence_out, "punc_array": punc_array} results.append(result_i) return results, meta_data def export(self, **kwargs): - + from .export_meta import export_rebuild_model + models = export_rebuild_model(model=self, **kwargs) return models - - diff --git a/funasr/models/ct_transformer/utils.py b/funasr/models/ct_transformer/utils.py index 01b1850d7..b6e11e7af 100644 --- a/funasr/models/ct_transformer/utils.py +++ b/funasr/models/ct_transformer/utils.py @@ -14,11 +14,12 @@ def split_to_mini_sentence(words: list, word_limit: int = 20): length = len(words) sentence_len = length // word_limit for i in range(sentence_len): - sentences.append(words[i * word_limit:(i + 1) * word_limit]) + sentences.append(words[i * word_limit : (i + 1) * word_limit]) if length % word_limit > 0: - sentences.append(words[sentence_len * word_limit:]) + sentences.append(words[sentence_len * word_limit :]) return sentences + def split_words(text: str, jieba_usr_dict=None, **kwargs): if jieba_usr_dict: input_list = text.split() @@ -27,21 +28,21 @@ def split_words(text: str, jieba_usr_dict=None, **kwargs): token_list_tmp = [] language_flag = None for token in input_list: - if isEnglish(token) and language_flag == 'Chinese': + if isEnglish(token) and language_flag == "Chinese": token_list_all.append(token_list_tmp) - langauge_list.append('Chinese') + langauge_list.append("Chinese") token_list_tmp = [] - elif not isEnglish(token) and language_flag == 'English': + elif not isEnglish(token) and language_flag == "English": token_list_all.append(token_list_tmp) - langauge_list.append('English') + langauge_list.append("English") token_list_tmp = [] token_list_tmp.append(token) if isEnglish(token): - language_flag = 'English' + language_flag = "English" else: - language_flag = 'Chinese' + language_flag = "Chinese" if token_list_tmp: token_list_all.append(token_list_tmp) @@ -49,7 +50,7 @@ def split_words(text: str, jieba_usr_dict=None, **kwargs): result_list = [] for token_list_tmp, language_flag in zip(token_list_all, langauge_list): - if language_flag == 'English': + if language_flag == "English": result_list.extend(token_list_tmp) else: seg_list = jieba_usr_dict.cut(join_chinese_and_english(token_list_tmp), HMM=False) @@ -77,17 +78,19 @@ def split_words(text: str, jieba_usr_dict=None, **kwargs): words.append(current_word) return words -def isEnglish(text:str): - if re.search('^[a-zA-Z\']+$', text): + +def isEnglish(text: str): + if re.search("^[a-zA-Z']+$", text): return True else: return False + def join_chinese_and_english(input_list): - line = '' + line = "" for token in input_list: if isEnglish(token): - line = line + ' ' + token + line = line + " " + token else: line = line + token diff --git a/funasr/models/ct_transformer_streaming/attention.py b/funasr/models/ct_transformer_streaming/attention.py index 3177ecae8..97e770bbb 100644 --- a/funasr/models/ct_transformer_streaming/attention.py +++ b/funasr/models/ct_transformer_streaming/attention.py @@ -18,5 +18,3 @@ class MultiHeadedAttentionSANMwithMask(MultiHeadedAttentionSANM): scores = torch.matmul(q_h, k_h.transpose(-2, -1)) att_outs = self.forward_attention(v_h, scores, mask[1], mask_att_chunk_encoder) return att_outs + fsmn_memory - - diff --git a/funasr/models/ct_transformer_streaming/encoder.py b/funasr/models/ct_transformer_streaming/encoder.py index bf0b8b277..a61319aa4 100644 --- a/funasr/models/ct_transformer_streaming/encoder.py +++ b/funasr/models/ct_transformer_streaming/encoder.py @@ -20,7 +20,12 @@ from funasr.models.transformer.utils.mask import subsequent_mask, vad_mask from funasr.models.transformer.utils.multi_layer_conv import MultiLayeredConv1d from funasr.models.transformer.positionwise_feed_forward import PositionwiseFeedForward from funasr.models.ct_transformer_streaming.attention import MultiHeadedAttentionSANMwithMask -from funasr.models.transformer.utils.subsampling import Conv2dSubsampling, Conv2dSubsampling2, Conv2dSubsampling6, Conv2dSubsampling8 +from funasr.models.transformer.utils.subsampling import ( + Conv2dSubsampling, + Conv2dSubsampling2, + Conv2dSubsampling6, + Conv2dSubsampling8, +) class EncoderLayerSANM(torch.nn.Module): @@ -82,7 +87,18 @@ class EncoderLayerSANM(torch.nn.Module): x = self.norm1(x) if self.concat_after: - x_concat = torch.cat((x, self.self_attn(x, mask, mask_shfit_chunk=mask_shfit_chunk, mask_att_chunk_encoder=mask_att_chunk_encoder)), dim=-1) + x_concat = torch.cat( + ( + x, + self.self_attn( + x, + mask, + mask_shfit_chunk=mask_shfit_chunk, + mask_att_chunk_encoder=mask_att_chunk_encoder, + ), + ), + dim=-1, + ) if self.in_size == self.size: x = residual + stoch_layer_coeff * self.concat_linear(x_concat) else: @@ -90,11 +106,21 @@ class EncoderLayerSANM(torch.nn.Module): else: if self.in_size == self.size: x = residual + stoch_layer_coeff * self.dropout( - self.self_attn(x, mask, mask_shfit_chunk=mask_shfit_chunk, mask_att_chunk_encoder=mask_att_chunk_encoder) + self.self_attn( + x, + mask, + mask_shfit_chunk=mask_shfit_chunk, + mask_att_chunk_encoder=mask_att_chunk_encoder, + ) ) else: x = stoch_layer_coeff * self.dropout( - self.self_attn(x, mask, mask_shfit_chunk=mask_shfit_chunk, mask_att_chunk_encoder=mask_att_chunk_encoder) + self.self_attn( + x, + mask, + mask_shfit_chunk=mask_shfit_chunk, + mask_att_chunk_encoder=mask_att_chunk_encoder, + ) ) if not self.normalize_before: x = self.norm1(x) @@ -171,8 +197,8 @@ class SANMVadEncoder(torch.nn.Module): padding_idx: int = -1, interctc_layer_idx: List[int] = [], interctc_use_conditioning: bool = False, - kernel_size : int = 11, - sanm_shfit : int = 0, + kernel_size: int = 11, + sanm_shfit: int = 0, selfattention_layer_type: str = "sanm", ): super().__init__() @@ -277,7 +303,7 @@ class SANMVadEncoder(torch.nn.Module): ) self.encoders = repeat( - num_blocks-1, + num_blocks - 1, lambda lnum: EncoderLayerSANM( output_size, output_size, @@ -321,16 +347,20 @@ class SANMVadEncoder(torch.nn.Module): masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device) sub_masks = subsequent_mask(masks.size(-1), device=xs_pad.device).unsqueeze(0) no_future_masks = masks & sub_masks - xs_pad *= self.output_size()**0.5 + xs_pad *= self.output_size() ** 0.5 if self.embed is None: xs_pad = xs_pad - elif (isinstance(self.embed, Conv2dSubsampling) or isinstance(self.embed, Conv2dSubsampling2) - or isinstance(self.embed, Conv2dSubsampling6) or isinstance(self.embed, Conv2dSubsampling8)): + elif ( + isinstance(self.embed, Conv2dSubsampling) + or isinstance(self.embed, Conv2dSubsampling2) + or isinstance(self.embed, Conv2dSubsampling6) + or isinstance(self.embed, Conv2dSubsampling8) + ): short_status, limit_size = check_short_utt(self.embed, xs_pad.size(1)) if short_status: raise TooShortUttError( - f"has {xs_pad.size(1)} frames and is too short for subsampling " + - f"(it needs more than {limit_size} frames), return empty results", + f"has {xs_pad.size(1)} frames and is too short for subsampling " + + f"(it needs more than {limit_size} frames), return empty results", xs_pad.size(1), limit_size, ) @@ -344,25 +374,26 @@ class SANMVadEncoder(torch.nn.Module): xs_pad, _ = encoder_outs[0], encoder_outs[1] intermediate_outs = [] - for layer_idx, encoder_layer in enumerate(self.encoders): - if layer_idx + 1 == len(self.encoders): - # This is last layer. - coner_mask = torch.ones(masks.size(0), - masks.size(-1), - masks.size(-1), - device=xs_pad.device, - dtype=torch.bool) - for word_index, length in enumerate(ilens): - coner_mask[word_index, :, :] = vad_mask(masks.size(-1), - vad_indexes[word_index], - device=xs_pad.device) - layer_mask = masks & coner_mask - else: - layer_mask = no_future_masks - mask_tup1 = [masks, layer_mask] - encoder_outs = encoder_layer(xs_pad, mask_tup1) - xs_pad, layer_mask = encoder_outs[0], encoder_outs[1] + if layer_idx + 1 == len(self.encoders): + # This is last layer. + coner_mask = torch.ones( + masks.size(0), + masks.size(-1), + masks.size(-1), + device=xs_pad.device, + dtype=torch.bool, + ) + for word_index, length in enumerate(ilens): + coner_mask[word_index, :, :] = vad_mask( + masks.size(-1), vad_indexes[word_index], device=xs_pad.device + ) + layer_mask = masks & coner_mask + else: + layer_mask = no_future_masks + mask_tup1 = [masks, layer_mask] + encoder_outs = encoder_layer(xs_pad, mask_tup1) + xs_pad, layer_mask = encoder_outs[0], encoder_outs[1] if self.normalize_before: xs_pad = self.after_norm(xs_pad) @@ -401,6 +432,7 @@ class EncoderLayerSANMExport(torch.nn.Module): return x, mask + @tables.register("encoder_classes", "SANMVadEncoderExport") class SANMVadEncoderExport(torch.nn.Module): def __init__( @@ -408,68 +440,67 @@ class SANMVadEncoderExport(torch.nn.Module): model, max_seq_len=512, feats_dim=560, - model_name='encoder', + model_name="encoder", onnx: bool = True, ): super().__init__() self.embed = model.embed self.model = model self._output_size = model._output_size - + from funasr.utils.torch_function import sequence_mask - self.make_pad_mask = sequence_mask(max_seq_len, flip=False) - + from funasr.models.sanm.attention import MultiHeadedAttentionSANMExport - - if hasattr(model, 'encoders0'): + + if hasattr(model, "encoders0"): for i, d in enumerate(self.model.encoders0): if isinstance(d.self_attn, MultiHeadedAttentionSANMwithMask): d.self_attn = MultiHeadedAttentionSANMExport(d.self_attn) self.model.encoders0[i] = EncoderLayerSANMExport(d) - + for i, d in enumerate(self.model.encoders): if isinstance(d.self_attn, MultiHeadedAttentionSANMwithMask): d.self_attn = MultiHeadedAttentionSANMExport(d.self_attn) self.model.encoders[i] = EncoderLayerSANMExport(d) - - + def prepare_mask(self, mask, sub_masks): mask_3d_btd = mask[:, :, None] mask_4d_bhlt = (1 - sub_masks) * -10000.0 - + return mask_3d_btd, mask_4d_bhlt - - def forward(self, - speech: torch.Tensor, - speech_lengths: torch.Tensor, - vad_masks: torch.Tensor, - sub_masks: torch.Tensor, - ): - speech = speech * self._output_size ** 0.5 + + def forward( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + vad_masks: torch.Tensor, + sub_masks: torch.Tensor, + ): + speech = speech * self._output_size**0.5 mask = self.make_pad_mask(speech_lengths) vad_masks = self.prepare_mask(mask, vad_masks) mask = self.prepare_mask(mask, sub_masks) - + if self.embed is None: xs_pad = speech else: xs_pad = self.embed(speech) - + encoder_outs = self.model.encoders0(xs_pad, mask) xs_pad, masks = encoder_outs[0], encoder_outs[1] - + # encoder_outs = self.model.encoders(xs_pad, mask) for layer_idx, encoder_layer in enumerate(self.model.encoders): if layer_idx == len(self.model.encoders) - 1: mask = vad_masks encoder_outs = encoder_layer(xs_pad, mask) xs_pad, masks = encoder_outs[0], encoder_outs[1] - + xs_pad = self.model.after_norm(xs_pad) - + return xs_pad, speech_lengths - + def get_output_size(self): - return self.model.encoders[0].size \ No newline at end of file + return self.model.encoders[0].size diff --git a/funasr/models/ct_transformer_streaming/export_meta.py b/funasr/models/ct_transformer_streaming/export_meta.py index e4745d6fb..ba0283e5d 100644 --- a/funasr/models/ct_transformer_streaming/export_meta.py +++ b/funasr/models/ct_transformer_streaming/export_meta.py @@ -9,25 +9,28 @@ from funasr.register import tables def export_rebuild_model(model, **kwargs): - + is_onnx = kwargs.get("type", "onnx") == "onnx" encoder_class = tables.encoder_classes.get(kwargs["encoder"] + "Export") model.encoder = encoder_class(model.encoder, onnx=is_onnx) - + model.forward = types.MethodType(export_forward, model) model.export_dummy_inputs = types.MethodType(export_dummy_inputs, model) model.export_input_names = types.MethodType(export_input_names, model) model.export_output_names = types.MethodType(export_output_names, model) model.export_dynamic_axes = types.MethodType(export_dynamic_axes, model) model.export_name = types.MethodType(export_name, model) - + return model -def export_forward(self, inputs: torch.Tensor, - text_lengths: torch.Tensor, - vad_indexes: torch.Tensor, - sub_masks: torch.Tensor, - ): + +def export_forward( + self, + inputs: torch.Tensor, + text_lengths: torch.Tensor, + vad_indexes: torch.Tensor, + sub_masks: torch.Tensor, +): """Compute loss value from buffer sequences. Args: @@ -41,6 +44,7 @@ def export_forward(self, inputs: torch.Tensor, y = self.decoder(h) return y + def export_dummy_inputs(self): length = 120 text_indexes = torch.randint(0, self.embed.num_embeddings, (1, length)).type(torch.int32) @@ -50,28 +54,23 @@ def export_dummy_inputs(self): sub_masks = torch.tril(sub_masks).type(torch.float32) return (text_indexes, text_lengths, vad_mask, sub_masks[None, None, :, :]) + def export_input_names(self): - return ['inputs', 'text_lengths', 'vad_masks', 'sub_masks'] + return ["inputs", "text_lengths", "vad_masks", "sub_masks"] + def export_output_names(self): - return ['logits'] + return ["logits"] + def export_dynamic_axes(self): return { - 'inputs': { - 1: 'feats_length' - }, - 'vad_masks': { - 2: 'feats_length1', - 3: 'feats_length2' - }, - 'sub_masks': { - 2: 'feats_length1', - 3: 'feats_length2' - }, - 'logits': { - 1: 'logits_length' - }, + "inputs": {1: "feats_length"}, + "vad_masks": {2: "feats_length1", 3: "feats_length2"}, + "sub_masks": {2: "feats_length1", 3: "feats_length2"}, + "logits": {1: "logits_length"}, } + + def export_name(self): return "model.onnx" diff --git a/funasr/models/ct_transformer_streaming/model.py b/funasr/models/ct_transformer_streaming/model.py index e6977ad19..a1afc7aea 100644 --- a/funasr/models/ct_transformer_streaming/model.py +++ b/funasr/models/ct_transformer_streaming/model.py @@ -31,6 +31,7 @@ class CTTransformerStreaming(CTTransformer): CT-Transformer: Controllable time-delay transformer for real-time punctuation prediction and disfluency detection https://arxiv.org/pdf/2003.01309.pdf """ + def __init__( self, *args, @@ -38,8 +39,9 @@ class CTTransformerStreaming(CTTransformer): ): super().__init__(*args, **kwargs) - - def punc_forward(self, text: torch.Tensor, text_lengths: torch.Tensor, vad_indexes: torch.Tensor, **kwargs): + def punc_forward( + self, text: torch.Tensor, text_lengths: torch.Tensor, vad_indexes: torch.Tensor, **kwargs + ): """Compute loss value from buffer sequences. Args: @@ -55,24 +57,24 @@ class CTTransformerStreaming(CTTransformer): def with_vad(self): return True - - def inference(self, - data_in, - data_lengths=None, - key: list = None, - tokenizer=None, - frontend=None, - cache: dict = {}, - **kwargs, - ): + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + cache: dict = {}, + **kwargs, + ): assert len(data_in) == 1 - + if len(cache) == 0: cache["pre_text"] = [] text = load_audio_text_image_video(data_in, data_type=kwargs.get("kwargs", "text"))[0] text = "".join(cache["pre_text"]) + " " + text - split_size = kwargs.get("split_size", 20) tokens = split_words(text) @@ -82,7 +84,7 @@ class CTTransformerStreaming(CTTransformer): mini_sentences_id = split_to_mini_sentence(tokens_int, split_size) assert len(mini_sentences) == len(mini_sentences_id) cache_sent = [] - cache_sent_id = torch.from_numpy(np.array([], dtype='int32')) + cache_sent_id = torch.from_numpy(np.array([], dtype="int32")) skip_num = 0 sentence_punc_list = [] sentence_words_list = [] @@ -97,8 +99,8 @@ class CTTransformerStreaming(CTTransformer): mini_sentence_id = np.concatenate((cache_sent_id, mini_sentence_id), axis=0) data = { "text": torch.unsqueeze(torch.from_numpy(mini_sentence_id), 0), - "text_lengths": torch.from_numpy(np.array([len(mini_sentence_id)], dtype='int32')), - "vad_indexes": torch.from_numpy(np.array([len(cache["pre_text"])], dtype='int32')), + "text_lengths": torch.from_numpy(np.array([len(mini_sentence_id)], dtype="int32")), + "vad_indexes": torch.from_numpy(np.array([len(cache["pre_text"])], dtype="int32")), } data = to_device(data, kwargs["device"]) # y, _ = self.wrapped_model(**data) @@ -114,20 +116,27 @@ class CTTransformerStreaming(CTTransformer): sentenceEnd = -1 last_comma_index = -1 for i in range(len(punctuations) - 2, 1, -1): - if self.punc_list[punctuations[i]] == "。" or self.punc_list[punctuations[i]] == "？": + if ( + self.punc_list[punctuations[i]] == "。" + or self.punc_list[punctuations[i]] == "？" + ): sentenceEnd = i break if last_comma_index < 0 and self.punc_list[punctuations[i]] == "，": last_comma_index = i - if sentenceEnd < 0 and len(mini_sentence) > cache_pop_trigger_limit and last_comma_index >= 0: + if ( + sentenceEnd < 0 + and len(mini_sentence) > cache_pop_trigger_limit + and last_comma_index >= 0 + ): # The sentence it too long, cut off at a comma. sentenceEnd = last_comma_index punctuations[sentenceEnd] = self.sentence_end_id - cache_sent = mini_sentence[sentenceEnd + 1:] - cache_sent_id = mini_sentence_id[sentenceEnd + 1:] - mini_sentence = mini_sentence[0:sentenceEnd + 1] - punctuations = punctuations[0:sentenceEnd + 1] + cache_sent = mini_sentence[sentenceEnd + 1 :] + cache_sent_id = mini_sentence_id[sentenceEnd + 1 :] + mini_sentence = mini_sentence[0 : sentenceEnd + 1] + punctuations = punctuations[0 : sentenceEnd + 1] # if len(punctuations) == 0: # continue @@ -141,7 +150,10 @@ class CTTransformerStreaming(CTTransformer): sentence_punc_list_out = [] for i in range(0, len(sentence_words_list)): if i > 0: - if len(sentence_words_list[i][0].encode()) == 1 and len(sentence_words_list[i - 1][-1].encode()) == 1: + if ( + len(sentence_words_list[i][0].encode()) == 1 + and len(sentence_words_list[i - 1][-1].encode()) == 1 + ): sentence_words_list[i] = " " + sentence_words_list[i] if skip_num < len(cache["pre_text"]): skip_num += 1 @@ -158,7 +170,7 @@ class CTTransformerStreaming(CTTransformer): if sentence_punc_list[i] == "。" or sentence_punc_list[i] == "？": sentenceEnd = i break - cache["pre_text"] = sentence_words_list[sentenceEnd + 1:] + cache["pre_text"] = sentence_words_list[sentenceEnd + 1 :] if sentence_out[-1] in self.punc_list: sentence_out = sentence_out[:-1] sentence_punc_list_out[-1] = "_" @@ -167,15 +179,15 @@ class CTTransformerStreaming(CTTransformer): punc_array = punctuations else: punc_array = torch.cat([punc_array, punctuations], dim=0) - + result_i = {"key": key[0], "text": sentence_out, "punc_array": punc_array} results.append(result_i) - + return results, meta_data def export(self, **kwargs): - + from .export_meta import export_rebuild_model + models = export_rebuild_model(model=self, **kwargs) return models - diff --git a/funasr/models/ctc/ctc.py b/funasr/models/ctc/ctc.py index d3c10fa40..bdfb3a68a 100644 --- a/funasr/models/ctc/ctc.py +++ b/funasr/models/ctc/ctc.py @@ -40,9 +40,7 @@ class CTC(torch.nn.Module): logging.warning("ignore_nan_grad option is not supported for warp_ctc") self.ctc_loss = warp_ctc.CTCLoss(size_average=True, reduce=reduce) else: - raise ValueError( - f'ctc_type must be "builtin" or "warpctc": {self.ctc_type}' - ) + raise ValueError(f'ctc_type must be "builtin" or "warpctc": {self.ctc_type}') self.reduce = reduce diff --git a/funasr/models/data2vec/data2vec.py b/funasr/models/data2vec/data2vec.py index c77cedf4b..ed58d1c7c 100644 --- a/funasr/models/data2vec/data2vec.py +++ b/funasr/models/data2vec/data2vec.py @@ -16,6 +16,7 @@ import torch.nn as nn # from funasr.models.base_model import FunASRModel # from funasr.models.encoder.abs_encoder import AbsEncoder from funasr.frontends.abs_frontend import AbsFrontend + # from funasr.models.preencoder.abs_preencoder import AbsPreEncoder # from funasr.models.specaug.abs_specaug import AbsSpecAug from funasr.train_utils.device_funcs import force_gatherable @@ -33,12 +34,12 @@ class Data2VecPretrainModel(nn.Module): """Data2Vec Pretrain model""" def __init__( - self, - frontend = None, - specaug = None, - normalize = None, - encoder = None, - preencoder = None, + self, + frontend=None, + specaug=None, + normalize=None, + encoder=None, + preencoder=None, ): super().__init__() @@ -51,9 +52,9 @@ class Data2VecPretrainModel(nn.Module): self.num_updates = 0 def forward( - self, - speech: torch.Tensor, - speech_lengths: torch.Tensor, + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]: """Frontend + Encoder + Calc loss Args: @@ -61,10 +62,7 @@ class Data2VecPretrainModel(nn.Module): speech_lengths: (Batch, ) """ # Check that batch_size is unified - assert ( - speech.shape[0] - == speech_lengths.shape[0] - ), (speech.shape, speech_lengths.shape) + assert speech.shape[0] == speech_lengths.shape[0], (speech.shape, speech_lengths.shape) self.encoder.set_num_updates(self.num_updates) @@ -91,17 +89,15 @@ class Data2VecPretrainModel(nn.Module): return loss, stats, weight def collect_feats( - self, - speech: torch.Tensor, - speech_lengths: torch.Tensor + self, speech: torch.Tensor, speech_lengths: torch.Tensor ) -> Dict[str, torch.Tensor]: feats, feats_lengths = self._extract_feats(speech, speech_lengths) return {"feats": feats, "feats_lengths": feats_lengths} def encode( - self, - speech: torch.Tensor, - speech_lengths: torch.Tensor, + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, ): """Frontend + Encoder. Args: @@ -132,7 +128,7 @@ class Data2VecPretrainModel(nn.Module): return encoder_out def _extract_feats( - self, speech: torch.Tensor, speech_lengths: torch.Tensor + self, speech: torch.Tensor, speech_lengths: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: assert speech_lengths.dim() == 1, speech_lengths.shape diff --git a/funasr/models/data2vec/data2vec_encoder.py b/funasr/models/data2vec/data2vec_encoder.py index 4689e200c..f591dd6a3 100644 --- a/funasr/models/data2vec/data2vec_encoder.py +++ b/funasr/models/data2vec/data2vec_encoder.py @@ -29,71 +29,71 @@ def get_annealed_rate(start, end, curr_step, total_steps): class Data2VecEncoder(nn.Module): def __init__( - self, - # for ConvFeatureExtractionModel - input_size: int = None, - extractor_mode: str = None, - conv_feature_layers: str = "[(512,2,2)] + [(512,2,2)]", - # for Transformer Encoder - ## model architecture - layer_type: str = "transformer", - layer_norm_first: bool = False, - encoder_layers: int = 12, - encoder_embed_dim: int = 768, - encoder_ffn_embed_dim: int = 3072, - encoder_attention_heads: int = 12, - activation_fn: str = "gelu", - ## dropouts - dropout: float = 0.1, - attention_dropout: float = 0.1, - activation_dropout: float = 0.0, - encoder_layerdrop: float = 0.0, - dropout_input: float = 0.0, - dropout_features: float = 0.0, - ## grad settings - feature_grad_mult: float = 1.0, - ## masking - mask_prob: float = 0.65, - mask_length: int = 10, - mask_selection: str = "static", - mask_other: int = 0, - no_mask_overlap: bool = False, - mask_min_space: int = 1, - require_same_masks: bool = True, # if set as True, collate_fn should be clipping - mask_dropout: float = 0.0, - ## channel masking - mask_channel_length: int = 10, - mask_channel_prob: float = 0.0, - mask_channel_before: bool = False, - mask_channel_selection: str = "static", - mask_channel_other: int = 0, - no_mask_channel_overlap: bool = False, - mask_channel_min_space: int = 1, - ## positional embeddings - conv_pos: int = 128, - conv_pos_groups: int = 16, - pos_conv_depth: int = 1, - max_positions: int = 100000, - # EMA module - average_top_k_layers: int = 8, - layer_norm_target_layer: bool = False, - instance_norm_target_layer: bool = False, - instance_norm_targets: bool = False, - layer_norm_targets: bool = False, - batch_norm_target_layer: bool = False, - group_norm_target_layer: bool = False, - ema_decay: float = 0.999, - ema_end_decay: float = 0.9999, - ema_anneal_end_step: int = 100000, - ema_transformer_only: bool = True, - ema_layers_only: bool = True, - min_target_var: float = 0.1, - min_pred_var: float = 0.01, - # Loss - loss_beta: float = 0.0, - loss_scale: float = None, - # FP16 optimization - required_seq_len_multiple: int = 2, + self, + # for ConvFeatureExtractionModel + input_size: int = None, + extractor_mode: str = None, + conv_feature_layers: str = "[(512,2,2)] + [(512,2,2)]", + # for Transformer Encoder + ## model architecture + layer_type: str = "transformer", + layer_norm_first: bool = False, + encoder_layers: int = 12, + encoder_embed_dim: int = 768, + encoder_ffn_embed_dim: int = 3072, + encoder_attention_heads: int = 12, + activation_fn: str = "gelu", + ## dropouts + dropout: float = 0.1, + attention_dropout: float = 0.1, + activation_dropout: float = 0.0, + encoder_layerdrop: float = 0.0, + dropout_input: float = 0.0, + dropout_features: float = 0.0, + ## grad settings + feature_grad_mult: float = 1.0, + ## masking + mask_prob: float = 0.65, + mask_length: int = 10, + mask_selection: str = "static", + mask_other: int = 0, + no_mask_overlap: bool = False, + mask_min_space: int = 1, + require_same_masks: bool = True, # if set as True, collate_fn should be clipping + mask_dropout: float = 0.0, + ## channel masking + mask_channel_length: int = 10, + mask_channel_prob: float = 0.0, + mask_channel_before: bool = False, + mask_channel_selection: str = "static", + mask_channel_other: int = 0, + no_mask_channel_overlap: bool = False, + mask_channel_min_space: int = 1, + ## positional embeddings + conv_pos: int = 128, + conv_pos_groups: int = 16, + pos_conv_depth: int = 1, + max_positions: int = 100000, + # EMA module + average_top_k_layers: int = 8, + layer_norm_target_layer: bool = False, + instance_norm_target_layer: bool = False, + instance_norm_targets: bool = False, + layer_norm_targets: bool = False, + batch_norm_target_layer: bool = False, + group_norm_target_layer: bool = False, + ema_decay: float = 0.999, + ema_end_decay: float = 0.9999, + ema_anneal_end_step: int = 100000, + ema_transformer_only: bool = True, + ema_layers_only: bool = True, + min_target_var: float = 0.1, + min_pred_var: float = 0.01, + # Loss + loss_beta: float = 0.0, + loss_scale: float = None, + # FP16 optimization + required_seq_len_multiple: int = 2, ): super().__init__() @@ -133,7 +133,9 @@ class Data2VecEncoder(nn.Module): self.mask_other = mask_other self.no_mask_overlap = no_mask_overlap self.mask_min_space = mask_min_space - self.require_same_masks = require_same_masks # if set as True, collate_fn should be clipping + self.require_same_masks = ( + require_same_masks # if set as True, collate_fn should be clipping + ) self.mask_dropout = mask_dropout ## channel masking self.mask_channel_length = mask_channel_length @@ -239,11 +241,11 @@ class Data2VecEncoder(nn.Module): self.num_updates = num_updates def apply_mask( - self, - x, - padding_mask, - mask_indices=None, - mask_channel_indices=None, + self, + x, + padding_mask, + mask_indices=None, + mask_channel_indices=None, ): B, T, C = x.shape @@ -259,10 +261,7 @@ class Data2VecEncoder(nn.Module): min_space=self.mask_channel_min_space, ) mask_channel_indices = ( - torch.from_numpy(mask_channel_indices) - .to(x.device) - .unsqueeze(1) - .expand(-1, T, -1) + torch.from_numpy(mask_channel_indices).to(x.device).unsqueeze(1).expand(-1, T, -1) ) x[mask_channel_indices] = 0 @@ -300,9 +299,9 @@ class Data2VecEncoder(nn.Module): ) mask_channel_indices = ( torch.from_numpy(mask_channel_indices) - .to(x.device) - .unsqueeze(1) - .expand(-1, T, -1) + .to(x.device) + .unsqueeze(1) + .expand(-1, T, -1) ) x[mask_channel_indices] = 0 @@ -326,15 +325,15 @@ class Data2VecEncoder(nn.Module): return input_lengths.to(torch.long) def forward( - self, - xs_pad, - ilens=None, - mask=False, - features_only=True, - layer=None, - mask_indices=None, - mask_channel_indices=None, - padding_count=None, + self, + xs_pad, + ilens=None, + mask=False, + features_only=True, + layer=None, + mask_indices=None, + mask_channel_indices=None, + padding_count=None, ): # create padding_mask by ilens if ilens is not None: @@ -446,16 +445,12 @@ class Data2VecEncoder(nn.Module): if self.batch_norm_target_layer: target_layer_results = [ - F.batch_norm( - tl.float(), running_mean=None, running_var=None, training=True - ) + F.batch_norm(tl.float(), running_mean=None, running_var=None, training=True) for tl in target_layer_results ] if self.instance_norm_target_layer: - target_layer_results = [ - F.instance_norm(tl.float()) for tl in target_layer_results - ] + target_layer_results = [F.instance_norm(tl.float()) for tl in target_layer_results] if permuted: target_layer_results = [ @@ -464,14 +459,12 @@ class Data2VecEncoder(nn.Module): if self.group_norm_target_layer: target_layer_results = [ - F.layer_norm(tl.float(), tl.shape[-2:]) - for tl in target_layer_results + F.layer_norm(tl.float(), tl.shape[-2:]) for tl in target_layer_results ] if self.layer_norm_target_layer: target_layer_results = [ - F.layer_norm(tl.float(), tl.shape[-1:]) - for tl in target_layer_results + F.layer_norm(tl.float(), tl.shape[-1:]) for tl in target_layer_results ] y = sum(target_layer_results) / len(target_layer_results) @@ -521,9 +514,7 @@ class Data2VecEncoder(nn.Module): f"target var is {result['target_var'].item()} < {self.min_target_var}, exiting" ) if self.num_updates > 5000 and result["pred_var"] < self.min_pred_var: - logging.error( - f"pred var is {result['pred_var'].item()} < {self.min_pred_var}, exiting" - ) + logging.error(f"pred var is {result['pred_var'].item()} < {self.min_pred_var}, exiting") raise Exception( f"pred var is {result['pred_var'].item()} < {self.min_pred_var}, exiting" ) @@ -539,20 +530,18 @@ class Data2VecEncoder(nn.Module): if dist.is_initialized(): zc = torch.tensor(y.size(0)).cuda() zs = y.sum(dim=0) - zss = (y ** 2).sum(dim=0) + zss = (y**2).sum(dim=0) dist.all_reduce(zc) dist.all_reduce(zs) dist.all_reduce(zss) - var = zss / (zc - 1) - (zs ** 2) / (zc * (zc - 1)) + var = zss / (zc - 1) - (zs**2) / (zc * (zc - 1)) return torch.sqrt(var + 1e-6).mean() else: return torch.sqrt(y.var(dim=0) + 1e-6).mean() - def extract_features( - self, xs_pad, ilens, mask=False, layer=None - ): + def extract_features(self, xs_pad, ilens, mask=False, layer=None): res = self.forward( xs_pad, ilens, @@ -571,4 +560,4 @@ class Data2VecEncoder(nn.Module): ) def output_size(self) -> int: - return self.encoder_embed_dim \ No newline at end of file + return self.encoder_embed_dim diff --git a/funasr/models/data2vec/data_utils.py b/funasr/models/data2vec/data_utils.py index abd6d516d..f486edf5a 100644 --- a/funasr/models/data2vec/data_utils.py +++ b/funasr/models/data2vec/data_utils.py @@ -11,17 +11,17 @@ import torch def compute_mask_indices( - shape: Tuple[int, int], - padding_mask: Optional[torch.Tensor], - mask_prob: float, - mask_length: int, - mask_type: str = "static", - mask_other: float = 0.0, - min_masks: int = 0, - no_overlap: bool = False, - min_space: int = 0, - require_same_masks: bool = True, - mask_dropout: float = 0.0, + shape: Tuple[int, int], + padding_mask: Optional[torch.Tensor], + mask_prob: float, + mask_length: int, + mask_type: str = "static", + mask_other: float = 0.0, + min_masks: int = 0, + no_overlap: bool = False, + min_space: int = 0, + require_same_masks: bool = True, + mask_dropout: float = 0.0, ) -> np.ndarray: """ Computes random mask spans for a given shape @@ -123,11 +123,7 @@ def compute_mask_indices( mask_idc = np.random.choice(sz - min_len, num_mask, replace=False) mask_idc = np.asarray( - [ - mask_idc[j] + offset - for j in range(len(mask_idc)) - for offset in range(lengths[j]) - ] + [mask_idc[j] + offset for j in range(len(mask_idc)) for offset in range(lengths[j])] ) mask_idcs.append(np.unique(mask_idc[mask_idc < sz])) @@ -138,9 +134,7 @@ def compute_mask_indices( mask_idc = np.random.choice(mask_idc, min_len, replace=False) if mask_dropout > 0: num_holes = np.rint(len(mask_idc) * mask_dropout).astype(int) - mask_idc = np.random.choice( - mask_idc, len(mask_idc) - num_holes, replace=False - ) + mask_idc = np.random.choice(mask_idc, len(mask_idc) - num_holes, replace=False) mask[i, mask_idc] = True diff --git a/funasr/models/data2vec/ema_module.py b/funasr/models/data2vec/ema_module.py index 4e46f502e..3d7c21e98 100644 --- a/funasr/models/data2vec/ema_module.py +++ b/funasr/models/data2vec/ema_module.py @@ -84,18 +84,14 @@ class EMAModule: decay = self.decay ema_state_dict = {} - ema_params = ( - self.fp32_params if self.ema_fp32 else self.model.state_dict() - ) + ema_params = self.fp32_params if self.ema_fp32 else self.model.state_dict() for key, param in new_model.state_dict().items(): if isinstance(param, dict): continue try: ema_param = ema_params[key] except KeyError: - ema_param = ( - param.float().clone() if param.ndim == 1 else copy.deepcopy(param) - ) + ema_param = param.float().clone() if param.ndim == 1 else copy.deepcopy(param) if param.shape != ema_param.shape: raise ValueError( @@ -107,7 +103,9 @@ class EMAModule: # Do not decay a model.version pytorch param continue - if key in self.skip_keys or ("num_batches_tracked" in key and ema_param.dtype == torch.int64): + if key in self.skip_keys or ( + "num_batches_tracked" in key and ema_param.dtype == torch.int64 + ): ema_param = param.to(dtype=ema_param.dtype).clone() ema_params[key].copy_(ema_param) else: diff --git a/funasr/models/data2vec/multihead_attention.py b/funasr/models/data2vec/multihead_attention.py index d0dbe359a..85bd0387d 100644 --- a/funasr/models/data2vec/multihead_attention.py +++ b/funasr/models/data2vec/multihead_attention.py @@ -29,11 +29,11 @@ class FairseqDropout(nn.Module): return x def make_generation_fast_( - self, - name: str, - retain_dropout: bool = False, - retain_dropout_modules: Optional[List[str]] = None, - **kwargs + self, + name: str, + retain_dropout: bool = False, + retain_dropout_modules: Optional[List[str]] = None, + **kwargs, ): if retain_dropout: if retain_dropout_modules is not None and self.module_name is None: @@ -42,12 +42,10 @@ class FairseqDropout(nn.Module): "because module_name was not set".format(name) ) elif ( - retain_dropout_modules is None # if None, apply to all modules - or self.module_name in retain_dropout_modules + retain_dropout_modules is None # if None, apply to all modules + or self.module_name in retain_dropout_modules ): - logging.info( - "Enabling dropout during inference for module: {}".format(name) - ) + logging.info("Enabling dropout during inference for module: {}".format(name)) self.apply_during_inference = True else: logging.info("Disabling dropout for module: {}".format(name)) @@ -60,19 +58,19 @@ class MultiheadAttention(nn.Module): """ def __init__( - self, - embed_dim, - num_heads, - kdim=None, - vdim=None, - dropout=0.0, - bias=True, - add_bias_kv=False, - add_zero_attn=False, - self_attention=False, - encoder_decoder_attention=False, - q_noise=0.0, - qn_block_size=8, + self, + embed_dim, + num_heads, + kdim=None, + vdim=None, + dropout=0.0, + bias=True, + add_bias_kv=False, + add_zero_attn=False, + self_attention=False, + encoder_decoder_attention=False, + q_noise=0.0, + qn_block_size=8, ): super().__init__() self.embed_dim = embed_dim @@ -81,15 +79,13 @@ class MultiheadAttention(nn.Module): self.qkv_same_dim = self.kdim == embed_dim and self.vdim == embed_dim self.num_heads = num_heads - self.dropout_module = FairseqDropout( - dropout, module_name=self.__class__.__name__ - ) + self.dropout_module = FairseqDropout(dropout, module_name=self.__class__.__name__) self.head_dim = embed_dim // num_heads assert ( - self.head_dim * num_heads == self.embed_dim + self.head_dim * num_heads == self.embed_dim ), "embed_dim must be divisible by num_heads" - self.scaling = self.head_dim ** -0.5 + self.scaling = self.head_dim**-0.5 self.self_attention = self_attention self.encoder_decoder_attention = encoder_decoder_attention @@ -157,45 +153,23 @@ class MultiheadAttention(nn.Module): start_idx = i * self.head_dim end_idx = (i + 1) * self.head_dim k_proj_heads_norm.append( - torch.sum( - torch.abs( - self.k_proj.weight[ - start_idx:end_idx, - ] - ) - ).tolist() + torch.sum(torch.abs(self.k_proj.weight[start_idx:end_idx,])).tolist() + torch.sum(torch.abs(self.k_proj.bias[start_idx:end_idx])).tolist() ) q_proj_heads_norm.append( - torch.sum( - torch.abs( - self.q_proj.weight[ - start_idx:end_idx, - ] - ) - ).tolist() + torch.sum(torch.abs(self.q_proj.weight[start_idx:end_idx,])).tolist() + torch.sum(torch.abs(self.q_proj.bias[start_idx:end_idx])).tolist() ) v_proj_heads_norm.append( - torch.sum( - torch.abs( - self.v_proj.weight[ - start_idx:end_idx, - ] - ) - ).tolist() + torch.sum(torch.abs(self.v_proj.weight[start_idx:end_idx,])).tolist() + torch.sum(torch.abs(self.v_proj.bias[start_idx:end_idx])).tolist() ) heads_norm = [] for i in range(self.num_heads): - heads_norm.append( - k_proj_heads_norm[i] + q_proj_heads_norm[i] + v_proj_heads_norm[i] - ) + heads_norm.append(k_proj_heads_norm[i] + q_proj_heads_norm[i] + v_proj_heads_norm[i]) - sorted_head_index = sorted( - range(self.num_heads), key=lambda k: heads_norm[k], reverse=True - ) + sorted_head_index = sorted(range(self.num_heads), key=lambda k: heads_norm[k], reverse=True) reserve_head_index = [] for i in range(num_heads_to_keep): start = sorted_head_index[i] * self.head_dim @@ -214,26 +188,14 @@ class MultiheadAttention(nn.Module): for ele in reserve_head_index: start_idx, end_idx = ele - new_q_weight.append( - self.q_proj.weight[ - start_idx:end_idx, - ] - ) + new_q_weight.append(self.q_proj.weight[start_idx:end_idx,]) new_q_bias.append(self.q_proj.bias[start_idx:end_idx]) - new_k_weight.append( - self.k_proj.weight[ - start_idx:end_idx, - ] - ) + new_k_weight.append(self.k_proj.weight[start_idx:end_idx,]) new_k_bias.append(self.k_proj.bias[start_idx:end_idx]) - new_v_weight.append( - self.v_proj.weight[ - start_idx:end_idx, - ] - ) + new_v_weight.append(self.v_proj.weight[start_idx:end_idx,]) new_v_bias.append(self.v_proj.bias[start_idx:end_idx]) new_out_proj_weight.append(self.out_proj.weight[:, start_idx:end_idx]) @@ -277,17 +239,17 @@ class MultiheadAttention(nn.Module): self.skip_embed_dim_check = True def forward( - self, - query, - key: Optional[Tensor], - value: Optional[Tensor], - key_padding_mask: Optional[Tensor] = None, - incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None, - need_weights: bool = True, - static_kv: bool = False, - attn_mask: Optional[Tensor] = None, - before_softmax: bool = False, - need_head_weights: bool = False, + self, + query, + key: Optional[Tensor], + value: Optional[Tensor], + key_padding_mask: Optional[Tensor] = None, + incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None, + need_weights: bool = True, + static_kv: bool = False, + attn_mask: Optional[Tensor] = None, + before_softmax: bool = False, + need_head_weights: bool = False, ) -> Tuple[Tensor, Optional[Tensor]]: """Input shape: Time x Batch x Channel @@ -314,9 +276,7 @@ class MultiheadAttention(nn.Module): tgt_len, bsz, embed_dim = query.size() src_len = tgt_len if not self.skip_embed_dim_check: - assert ( - embed_dim == self.embed_dim - ), f"query dim {embed_dim} != {self.embed_dim}" + assert embed_dim == self.embed_dim, f"query dim {embed_dim} != {self.embed_dim}" assert list(query.size()) == [tgt_len, bsz, embed_dim] if key is not None: src_len, key_bsz, _ = key.size() @@ -326,18 +286,18 @@ class MultiheadAttention(nn.Module): assert src_len, bsz == value.shape[:2] if ( - not self.onnx_trace - and not is_tpu # don't use PyTorch version on TPUs - and incremental_state is None - and not static_kv - # A workaround for quantization to work. Otherwise JIT compilation - # treats bias in linear module as method. - and not torch.jit.is_scripting() - # The Multihead attention implemented in pytorch forces strong dimension check - # for input embedding dimention and K,Q,V projection dimension. - # Since pruning will break the dimension check and it is not easy to modify the pytorch API, - # it is preferred to bypass the pytorch MHA when we need to skip embed_dim_check - and not self.skip_embed_dim_check + not self.onnx_trace + and not is_tpu # don't use PyTorch version on TPUs + and incremental_state is None + and not static_kv + # A workaround for quantization to work. Otherwise JIT compilation + # treats bias in linear module as method. + and not torch.jit.is_scripting() + # The Multihead attention implemented in pytorch forces strong dimension check + # for input embedding dimention and K,Q,V projection dimension. + # Since pruning will break the dimension check and it is not easy to modify the pytorch API, + # it is preferred to bypass the pytorch MHA when we need to skip embed_dim_check + and not self.skip_embed_dim_check ): assert key is not None and value is not None return F.multi_head_attention_forward( @@ -401,9 +361,7 @@ class MultiheadAttention(nn.Module): k = torch.cat([k, self.bias_k.repeat(1, bsz, 1)]) v = torch.cat([v, self.bias_v.repeat(1, bsz, 1)]) if attn_mask is not None: - attn_mask = torch.cat( - [attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1 - ) + attn_mask = torch.cat([attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1) if key_padding_mask is not None: key_padding_mask = torch.cat( [ @@ -413,23 +371,11 @@ class MultiheadAttention(nn.Module): dim=1, ) - q = ( - q.contiguous() - .view(tgt_len, bsz * self.num_heads, self.head_dim) - .transpose(0, 1) - ) + q = q.contiguous().view(tgt_len, bsz * self.num_heads, self.head_dim).transpose(0, 1) if k is not None: - k = ( - k.contiguous() - .view(-1, bsz * self.num_heads, self.head_dim) - .transpose(0, 1) - ) + k = k.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1) if v is not None: - v = ( - v.contiguous() - .view(-1, bsz * self.num_heads, self.head_dim) - .transpose(0, 1) - ) + v = v.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1) if saved_state is not None: # saved states are stored with shape (bsz, num_heads, seq_len, head_dim) @@ -488,16 +434,12 @@ class MultiheadAttention(nn.Module): k = torch.cat([k, k.new_zeros((k.size(0), 1) + k.size()[2:])], dim=1) v = torch.cat([v, v.new_zeros((v.size(0), 1) + v.size()[2:])], dim=1) if attn_mask is not None: - attn_mask = torch.cat( - [attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1 - ) + attn_mask = torch.cat([attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1) if key_padding_mask is not None: key_padding_mask = torch.cat( [ key_padding_mask, - torch.zeros(key_padding_mask.size(0), 1).type_as( - key_padding_mask - ), + torch.zeros(key_padding_mask.size(0), 1).type_as(key_padding_mask), ], dim=1, ) @@ -546,9 +488,9 @@ class MultiheadAttention(nn.Module): attn = self.out_proj(attn) attn_weights: Optional[Tensor] = None if need_weights: - attn_weights = attn_weights_float.view( - bsz, self.num_heads, tgt_len, src_len - ).transpose(1, 0) + attn_weights = attn_weights_float.view(bsz, self.num_heads, tgt_len, src_len).transpose( + 1, 0 + ) if not need_head_weights: # average attention weights over heads attn_weights = attn_weights.mean(dim=0) @@ -557,11 +499,11 @@ class MultiheadAttention(nn.Module): @staticmethod def _append_prev_key_padding_mask( - key_padding_mask: Optional[Tensor], - prev_key_padding_mask: Optional[Tensor], - batch_size: int, - src_len: int, - static_kv: bool, + key_padding_mask: Optional[Tensor], + prev_key_padding_mask: Optional[Tensor], + batch_size: int, + src_len: int, + static_kv: bool, ) -> Optional[Tensor]: # saved key padding masks have shape (bsz, seq_len) if prev_key_padding_mask is not None and static_kv: @@ -590,9 +532,7 @@ class MultiheadAttention(nn.Module): (batch_size, src_len - key_padding_mask.size(1)), device=key_padding_mask.device, ) - new_key_padding_mask = torch.cat( - [filler.float(), key_padding_mask.float()], dim=1 - ) + new_key_padding_mask = torch.cat([filler.float(), key_padding_mask.float()], dim=1) else: new_key_padding_mask = key_padding_mask.float() else: @@ -601,9 +541,9 @@ class MultiheadAttention(nn.Module): @torch.jit.export def reorder_incremental_state( - self, - incremental_state: Dict[str, Dict[str, Optional[Tensor]]], - new_order: Tensor, + self, + incremental_state: Dict[str, Dict[str, Optional[Tensor]]], + new_order: Tensor, ): """Reorder buffered internal state (for incremental generation).""" input_buffer = self._get_input_buffer(incremental_state) @@ -611,16 +551,16 @@ class MultiheadAttention(nn.Module): for k in input_buffer.keys(): input_buffer_k = input_buffer[k] if input_buffer_k is not None: - if self.encoder_decoder_attention and input_buffer_k.size( - 0 - ) == new_order.size(0): + if self.encoder_decoder_attention and input_buffer_k.size(0) == new_order.size( + 0 + ): break input_buffer[k] = input_buffer_k.index_select(0, new_order) incremental_state = self._set_input_buffer(incremental_state, input_buffer) return incremental_state def _get_input_buffer( - self, incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] + self, incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] ) -> Dict[str, Optional[Tensor]]: result = self.get_incremental_state(incremental_state, "attn_state") if result is not None: @@ -630,9 +570,9 @@ class MultiheadAttention(nn.Module): return empty_result def _set_input_buffer( - self, - incremental_state: Dict[str, Dict[str, Optional[Tensor]]], - buffer: Dict[str, Optional[Tensor]], + self, + incremental_state: Dict[str, Dict[str, Optional[Tensor]]], + buffer: Dict[str, Optional[Tensor]], ): return self.set_incremental_state(incremental_state, "attn_state", buffer) @@ -648,8 +588,8 @@ class MultiheadAttention(nn.Module): # in_proj_weight used to be q + k + v with same dimensions dim = int(state_dict[k].shape[0] / 3) items_to_add[prefix + "q_proj.weight"] = state_dict[k][:dim] - items_to_add[prefix + "k_proj.weight"] = state_dict[k][dim: 2 * dim] - items_to_add[prefix + "v_proj.weight"] = state_dict[k][2 * dim:] + items_to_add[prefix + "k_proj.weight"] = state_dict[k][dim : 2 * dim] + items_to_add[prefix + "v_proj.weight"] = state_dict[k][2 * dim :] keys_to_remove.append(k) @@ -657,10 +597,8 @@ class MultiheadAttention(nn.Module): if k_bias in state_dict.keys(): dim = int(state_dict[k].shape[0] / 3) items_to_add[prefix + "q_proj.bias"] = state_dict[k_bias][:dim] - items_to_add[prefix + "k_proj.bias"] = state_dict[k_bias][ - dim: 2 * dim - ] - items_to_add[prefix + "v_proj.bias"] = state_dict[k_bias][2 * dim:] + items_to_add[prefix + "k_proj.bias"] = state_dict[k_bias][dim : 2 * dim] + items_to_add[prefix + "v_proj.bias"] = state_dict[k_bias][2 * dim :] keys_to_remove.append(prefix + "in_proj_bias") diff --git a/funasr/models/data2vec/quant_noise.py b/funasr/models/data2vec/quant_noise.py index 11a82b649..50b2c1528 100644 --- a/funasr/models/data2vec/quant_noise.py +++ b/funasr/models/data2vec/quant_noise.py @@ -40,7 +40,7 @@ def quant_noise(module, p, block_size): # 2D matrix if not is_conv: assert ( - module.weight.size(1) % block_size == 0 + module.weight.size(1) % block_size == 0 ), "Input features must be a multiple of block sizes" # 4D matrix @@ -48,7 +48,7 @@ def quant_noise(module, p, block_size): # 1x1 convolutions if module.kernel_size == (1, 1): assert ( - module.in_channels % block_size == 0 + module.in_channels % block_size == 0 ), "Input channels must be a multiple of block sizes" # regular convolutions else: @@ -65,9 +65,7 @@ def quant_noise(module, p, block_size): out_features = weight.size(0) # split weight matrix into blocks and randomly drop selected blocks - mask = torch.zeros( - in_features // block_size * out_features, device=weight.device - ) + mask = torch.zeros(in_features // block_size * out_features, device=weight.device) mask.bernoulli_(p) mask = mask.repeat_interleave(block_size, -1).view(-1, in_features) @@ -86,20 +84,16 @@ def quant_noise(module, p, block_size): mask.bernoulli_(p) mask = mask.repeat_interleave(block_size, -1).view(-1, in_channels) else: - mask = torch.zeros( - weight.size(0), weight.size(1), device=weight.device - ) + mask = torch.zeros(weight.size(0), weight.size(1), device=weight.device) mask.bernoulli_(p) mask = ( mask.unsqueeze(2) - .unsqueeze(3) - .repeat(1, 1, mod.kernel_size[0], mod.kernel_size[1]) + .unsqueeze(3) + .repeat(1, 1, mod.kernel_size[0], mod.kernel_size[1]) ) # scale weights and apply mask - mask = mask.to( - torch.bool - ) # x.bool() is not currently supported in TorchScript + mask = mask.to(torch.bool) # x.bool() is not currently supported in TorchScript s = 1 / (1 - p) mod.weight.data = s * weight.masked_fill(mask, 0) diff --git a/funasr/models/data2vec/utils.py b/funasr/models/data2vec/utils.py index 70889dc91..b4c160358 100644 --- a/funasr/models/data2vec/utils.py +++ b/funasr/models/data2vec/utils.py @@ -84,9 +84,7 @@ def pad_to_multiple(x, multiple, dim=-1, value=0): def gelu_accurate(x): if not hasattr(gelu_accurate, "_a"): gelu_accurate._a = math.sqrt(2 / math.pi) - return ( - 0.5 * x * (1 + torch.tanh(gelu_accurate._a * (x + 0.044715 * torch.pow(x, 3)))) - ) + return 0.5 * x * (1 + torch.tanh(gelu_accurate._a * (x + 0.044715 * torch.pow(x, 3)))) def gelu(x: torch.Tensor) -> torch.Tensor: diff --git a/funasr/models/data2vec/wav2vec2.py b/funasr/models/data2vec/wav2vec2.py index 6631c57fe..a3e837ad3 100644 --- a/funasr/models/data2vec/wav2vec2.py +++ b/funasr/models/data2vec/wav2vec2.py @@ -18,25 +18,25 @@ from funasr.models.data2vec.multihead_attention import MultiheadAttention class ConvFeatureExtractionModel(nn.Module): def __init__( - self, - conv_layers: List[Tuple[int, int, int]], - dropout: float = 0.0, - mode: str = "default", - conv_bias: bool = False, - in_d: int = 1 + self, + conv_layers: List[Tuple[int, int, int]], + dropout: float = 0.0, + mode: str = "default", + conv_bias: bool = False, + in_d: int = 1, ): super().__init__() assert mode in {"default", "layer_norm"} def block( - n_in, - n_out, - k, - stride, - is_layer_norm=False, - is_group_norm=False, - conv_bias=False, + n_in, + n_out, + k, + stride, + is_layer_norm=False, + is_group_norm=False, + conv_bias=False, ): def make_conv(): conv = nn.Conv1d(n_in, n_out, k, stride=stride, bias=conv_bias) @@ -44,8 +44,8 @@ class ConvFeatureExtractionModel(nn.Module): return conv assert ( - is_layer_norm and is_group_norm - ) == False, "layer norm and group norm are exclusive" + is_layer_norm and is_group_norm + ) == False, "layer norm and group norm are exclusive" if is_layer_norm: return nn.Sequential( @@ -134,25 +134,25 @@ class TransformerEncoder(nn.Module): return layer def __init__( - self, - # position - dropout, - encoder_embed_dim, - required_seq_len_multiple, - pos_conv_depth, - conv_pos, - conv_pos_groups, - # transformer layers - layer_type, - encoder_layers, - encoder_ffn_embed_dim, - encoder_attention_heads, - attention_dropout, - activation_dropout, - activation_fn, - layer_norm_first, - encoder_layerdrop, - max_positions, + self, + # position + dropout, + encoder_embed_dim, + required_seq_len_multiple, + pos_conv_depth, + conv_pos, + conv_pos_groups, + # transformer layers + layer_type, + encoder_layers, + encoder_ffn_embed_dim, + encoder_attention_heads, + attention_dropout, + activation_dropout, + activation_fn, + layer_norm_first, + encoder_layerdrop, + max_positions, ): super().__init__() @@ -185,9 +185,7 @@ class TransformerEncoder(nn.Module): ] ) - self.pos_conv = make_conv_block( - self.embedding_dim, k, conv_pos_groups, num_layers - ) + self.pos_conv = make_conv_block(self.embedding_dim, k, conv_pos_groups, num_layers) else: self.pos_conv = make_conv_pos( @@ -206,9 +204,7 @@ class TransformerEncoder(nn.Module): self.layer_norm_first = layer_norm_first self.layerdrop = encoder_layerdrop self.max_positions = max_positions - self.layers = nn.ModuleList( - [self.build_encoder_layer() for _ in range(encoder_layers)] - ) + self.layers = nn.ModuleList([self.build_encoder_layer() for _ in range(encoder_layers)]) self.layer_norm = torch.nn.LayerNorm(self.embedding_dim) self.apply(utils.init_bert_params) @@ -222,11 +218,11 @@ class TransformerEncoder(nn.Module): return x, layer_results def extract_features( - self, - x, - padding_mask=None, - tgt_layer=None, - min_layer=0, + self, + x, + padding_mask=None, + tgt_layer=None, + min_layer=0, ): if padding_mask is not None: @@ -240,9 +236,7 @@ class TransformerEncoder(nn.Module): x = self.layer_norm(x) # pad to the sequence length dimension - x, pad_length = utils.pad_to_multiple( - x, self.required_seq_len_multiple, dim=-2, value=0 - ) + x, pad_length = utils.pad_to_multiple(x, self.required_seq_len_multiple, dim=-2, value=0) if pad_length > 0 and padding_mask is None: padding_mask = x.new_zeros((x.size(0), x.size(1)), dtype=torch.bool) padding_mask[:, -pad_length:] = True @@ -304,15 +298,15 @@ class TransformerSentenceEncoderLayer(nn.Module): """ def __init__( - self, - embedding_dim: int = 768, - ffn_embedding_dim: int = 3072, - num_attention_heads: int = 8, - dropout: float = 0.1, - attention_dropout: float = 0.1, - activation_dropout: float = 0.1, - activation_fn: str = "relu", - layer_norm_first: bool = False, + self, + embedding_dim: int = 768, + ffn_embedding_dim: int = 3072, + num_attention_heads: int = 8, + dropout: float = 0.1, + attention_dropout: float = 0.1, + activation_dropout: float = 0.1, + activation_fn: str = "relu", + layer_norm_first: bool = False, ) -> None: super().__init__() @@ -345,10 +339,10 @@ class TransformerSentenceEncoderLayer(nn.Module): self.final_layer_norm = torch.nn.LayerNorm(self.embedding_dim) def forward( - self, - x: torch.Tensor, # (T, B, C) - self_attn_mask: torch.Tensor = None, - self_attn_padding_mask: torch.Tensor = None, + self, + x: torch.Tensor, # (T, B, C) + self_attn_mask: torch.Tensor = None, + self_attn_padding_mask: torch.Tensor = None, ): """ LayerNorm is applied either before or after the self-attention/ffn diff --git a/funasr/models/e_branchformer/encoder.py b/funasr/models/e_branchformer/encoder.py index 4084e2116..adc6642f4 100644 --- a/funasr/models/e_branchformer/encoder.py +++ b/funasr/models/e_branchformer/encoder.py @@ -44,6 +44,7 @@ from funasr.models.transformer.utils.subsampling import ( ) from funasr.register import tables + class EBranchformerEncoderLayer(torch.nn.Module): """E-Branchformer encoder layer module. @@ -174,6 +175,7 @@ class EBranchformerEncoderLayer(torch.nn.Module): return x, mask + @tables.register("encoder_classes", "EBranchformerEncoder") class EBranchformerEncoder(nn.Module): """E-Branchformer encoder module.""" @@ -232,9 +234,7 @@ class EBranchformerEncoder(nn.Module): elif pos_enc_layer_type == "legacy_rel_pos": assert attention_layer_type == "legacy_rel_selfattn" pos_enc_class = LegacyRelPositionalEncoding - logging.warning( - "Using legacy_rel_pos and it will be deprecated in the future." - ) + logging.warning("Using legacy_rel_pos and it will be deprecated in the future.") else: raise ValueError("unknown pos_enc_layer: " + pos_enc_layer_type) @@ -320,9 +320,7 @@ class EBranchformerEncoder(nn.Module): output_size, attention_dropout_rate, ) - logging.warning( - "Using legacy_rel_selfattn and it will be deprecated in the future." - ) + logging.warning("Using legacy_rel_selfattn and it will be deprecated in the future.") elif attention_layer_type == "rel_selfattn": assert pos_enc_layer_type == "rel_pos" encoder_selfattn_layer = RelPositionMultiHeadedAttention @@ -360,9 +358,7 @@ class EBranchformerEncoder(nn.Module): encoder_selfattn_layer(*encoder_selfattn_layer_args), cgmlp_layer(*cgmlp_layer_args), positionwise_layer(*positionwise_layer_args) if use_ffn else None, - positionwise_layer(*positionwise_layer_args) - if use_ffn and macaron_ffn - else None, + positionwise_layer(*positionwise_layer_args) if use_ffn and macaron_ffn else None, dropout_rate, merge_conv_kernel, ), diff --git a/funasr/models/e_branchformer/model.py b/funasr/models/e_branchformer/model.py index 14c8c4d4d..9853f7352 100644 --- a/funasr/models/e_branchformer/model.py +++ b/funasr/models/e_branchformer/model.py @@ -3,6 +3,7 @@ import logging from funasr.models.transformer.model import Transformer from funasr.register import tables + @tables.register("model_classes", "EBranchformer") class EBranchformer(Transformer): """CTC-attention hybrid Encoder-Decoder model""" diff --git a/funasr/models/eend/e2e_diar_eend_ola.py b/funasr/models/eend/e2e_diar_eend_ola.py index ee8083631..cae5d1f23 100644 --- a/funasr/models/eend/e2e_diar_eend_ola.py +++ b/funasr/models/eend/e2e_diar_eend_ola.py @@ -4,13 +4,17 @@ from typing import Dict, List, Tuple, Optional import numpy as np import torch -import torch.nn as nn +import torch.nn as nn import torch.nn.functional as F from funasr.frontends.wav_frontend import WavFrontendMel23 from funasr.models.eend.encoder import EENDOLATransformerEncoder from funasr.models.eend.encoder_decoder_attractor import EncoderDecoderAttractor -from funasr.models.eend.utils.losses import standard_loss, cal_power_loss, fast_batch_pit_n_speaker_loss +from funasr.models.eend.utils.losses import ( + standard_loss, + cal_power_loss, + fast_batch_pit_n_speaker_loss, +) from funasr.models.eend.utils.power import create_powerlabel from funasr.models.eend.utils.power import generate_mapping_dict from funasr.train_utils.device_funcs import force_gatherable @@ -27,19 +31,16 @@ else: def pad_attractor(att, max_n_speakers): C, D = att.shape if C < max_n_speakers: - att = torch.cat([att, torch.zeros(max_n_speakers - C, D).to(torch.float32).to(att.device)], dim=0) + att = torch.cat( + [att, torch.zeros(max_n_speakers - C, D).to(torch.float32).to(att.device)], dim=0 + ) return att def pad_labels(ts, out_size): for i, t in enumerate(ts): if t.shape[1] < out_size: - ts[i] = F.pad( - t, - (0, out_size - t.shape[1], 0, 0), - mode='constant', - value=0. - ) + ts[i] = F.pad(t, (0, out_size - t.shape[1], 0, 0), mode="constant", value=0.0) return ts @@ -48,7 +49,16 @@ def pad_results(ys, out_size): for i, y in enumerate(ys): if y.shape[1] < out_size: ys_padded.append( - torch.cat([y, torch.zeros(y.shape[0], out_size - y.shape[1]).to(torch.float32).to(y.device)], dim=1)) + torch.cat( + [ + y, + torch.zeros(y.shape[0], out_size - y.shape[1]) + .to(torch.float32) + .to(y.device), + ], + dim=1, + ) + ) else: ys_padded.append(y) return ys_padded @@ -58,15 +68,15 @@ class DiarEENDOLAModel(nn.Module): """EEND-OLA diarization model""" def __init__( - self, - frontend: Optional[WavFrontendMel23], - encoder: EENDOLATransformerEncoder, - encoder_decoder_attractor: EncoderDecoderAttractor, - n_units: int = 256, - max_n_speaker: int = 8, - attractor_loss_weight: float = 1.0, - mapping_dict=None, - **kwargs, + self, + frontend: Optional[WavFrontendMel23], + encoder: EENDOLATransformerEncoder, + encoder_decoder_attractor: EncoderDecoderAttractor, + n_units: int = 256, + max_n_speaker: int = 8, + attractor_loss_weight: float = 1.0, + mapping_dict=None, + **kwargs, ): super().__init__() self.frontend = frontend @@ -79,13 +89,15 @@ class DiarEENDOLAModel(nn.Module): self.mapping_dict = mapping_dict # PostNet self.postnet = nn.LSTM(self.max_n_speaker, n_units, 1, batch_first=True) - self.output_layer = nn.Linear(n_units, mapping_dict['oov'] + 1) + self.output_layer = nn.Linear(n_units, mapping_dict["oov"] + 1) def forward_encoder(self, xs, ilens): xs = nn.utils.rnn.pad_sequence(xs, batch_first=True, padding_value=-1) pad_shape = xs.shape xs_mask = [torch.ones(ilen).to(xs.device) for ilen in ilens] - xs_mask = torch.nn.utils.rnn.pad_sequence(xs_mask, batch_first=True, padding_value=0).unsqueeze(-2) + xs_mask = torch.nn.utils.rnn.pad_sequence( + xs_mask, batch_first=True, padding_value=0 + ).unsqueeze(-2) emb = self.enc(xs, xs_mask) emb = torch.split(emb.view(pad_shape[0], pad_shape[1], -1), 1, dim=0) emb = [e[0][:ilen] for e, ilen in zip(emb, ilens)] @@ -94,34 +106,42 @@ class DiarEENDOLAModel(nn.Module): def forward_post_net(self, logits, ilens): maxlen = torch.max(ilens).to(torch.int).item() logits = nn.utils.rnn.pad_sequence(logits, batch_first=True, padding_value=-1) - logits = nn.utils.rnn.pack_padded_sequence(logits, ilens.cpu().to(torch.int64), batch_first=True, - enforce_sorted=False) + logits = nn.utils.rnn.pack_padded_sequence( + logits, ilens.cpu().to(torch.int64), batch_first=True, enforce_sorted=False + ) outputs, (_, _) = self.postnet(logits) - outputs = nn.utils.rnn.pad_packed_sequence(outputs, batch_first=True, padding_value=-1, total_length=maxlen)[0] - outputs = [output[:ilens[i].to(torch.int).item()] for i, output in enumerate(outputs)] + outputs = nn.utils.rnn.pad_packed_sequence( + outputs, batch_first=True, padding_value=-1, total_length=maxlen + )[0] + outputs = [output[: ilens[i].to(torch.int).item()] for i, output in enumerate(outputs)] outputs = [self.output_layer(output) for output in outputs] return outputs def forward( - self, - speech: List[torch.Tensor], - speaker_labels: List[torch.Tensor], - orders: torch.Tensor, + self, + speech: List[torch.Tensor], + speaker_labels: List[torch.Tensor], + orders: torch.Tensor, ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]: # Check that batch_size is unified - assert (len(speech) == len(speaker_labels)), (len(speech), len(speaker_labels)) + assert len(speech) == len(speaker_labels), (len(speech), len(speaker_labels)) speech_lengths = torch.tensor([len(sph) for sph in speech]).to(torch.int64) - speaker_labels_lengths = torch.tensor([spk.shape[-1] for spk in speaker_labels]).to(torch.int64) + speaker_labels_lengths = torch.tensor([spk.shape[-1] for spk in speaker_labels]).to( + torch.int64 + ) batch_size = len(speech) # Encoder encoder_out = self.forward_encoder(speech, speech_lengths) # Encoder-decoder attractor - attractor_loss, attractors = self.encoder_decoder_attractor([e[order] for e, order in zip(encoder_out, orders)], - speaker_labels_lengths) - speaker_logits = [torch.matmul(e, att.permute(1, 0)) for e, att in zip(encoder_out, attractors)] + attractor_loss, attractors = self.encoder_decoder_attractor( + [e[order] for e, order in zip(encoder_out, orders)], speaker_labels_lengths + ) + speaker_logits = [ + torch.matmul(e, att.permute(1, 0)) for e, att in zip(encoder_out, attractors) + ] # pit loss pit_speaker_labels = fast_batch_pit_n_speaker_loss(speaker_logits, speaker_labels) @@ -129,10 +149,17 @@ class DiarEENDOLAModel(nn.Module): # pse loss with torch.no_grad(): - power_ts = [create_powerlabel(label.cpu().numpy(), self.mapping_dict, self.max_n_speaker). - to(encoder_out[0].device, non_blocking=True) for label in pit_speaker_labels] + power_ts = [ + create_powerlabel(label.cpu().numpy(), self.mapping_dict, self.max_n_speaker).to( + encoder_out[0].device, non_blocking=True + ) + for label in pit_speaker_labels + ] pad_attractors = [pad_attractor(att, self.max_n_speaker) for att in attractors] - pse_speaker_logits = [torch.matmul(e, pad_att.permute(1, 0)) for e, pad_att in zip(encoder_out, pad_attractors)] + pse_speaker_logits = [ + torch.matmul(e, pad_att.permute(1, 0)) + for e, pad_att in zip(encoder_out, pad_attractors) + ] pse_speaker_logits = self.forward_post_net(pse_speaker_logits, speech_lengths) pse_loss = cal_power_loss(pse_speaker_logits, power_ts) @@ -151,12 +178,14 @@ class DiarEENDOLAModel(nn.Module): loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) return loss, stats, weight - def estimate_sequential(self, - speech: torch.Tensor, - n_speakers: int = None, - shuffle: bool = True, - threshold: float = 0.5, - **kwargs): + def estimate_sequential( + self, + speech: torch.Tensor, + n_speakers: int = None, + shuffle: bool = True, + threshold: float = 0.5, + **kwargs, + ): speech_lengths = torch.tensor([len(sph) for sph in speech]).to(torch.int64) emb = self.forward_encoder(speech, speech_lengths) if shuffle: @@ -164,35 +193,46 @@ class DiarEENDOLAModel(nn.Module): for order in orders: np.random.shuffle(order) attractors, probs = self.encoder_decoder_attractor.estimate( - [e[torch.from_numpy(order).to(torch.long).to(speech[0].device)] for e, order in zip(emb, orders)]) + [ + e[torch.from_numpy(order).to(torch.long).to(speech[0].device)] + for e, order in zip(emb, orders) + ] + ) else: attractors, probs = self.encoder_decoder_attractor.estimate(emb) attractors_active = [] for p, att, e in zip(probs, attractors, emb): if n_speakers and n_speakers >= 0: - att = att[:n_speakers, ] + att = att[:n_speakers,] attractors_active.append(att) elif threshold is not None: silence = torch.nonzero(p < threshold)[0] n_spk = silence[0] if silence.size else None - att = att[:n_spk, ] + att = att[:n_spk,] attractors_active.append(att) else: - NotImplementedError('n_speakers or threshold has to be given.') + NotImplementedError("n_speakers or threshold has to be given.") raw_n_speakers = [att.shape[0] for att in attractors_active] attractors = [ - pad_attractor(att, self.max_n_speaker) if att.shape[0] <= self.max_n_speaker else att[:self.max_n_speaker] - for att in attractors_active] + ( + pad_attractor(att, self.max_n_speaker) + if att.shape[0] <= self.max_n_speaker + else att[: self.max_n_speaker] + ) + for att in attractors_active + ] ys = [torch.matmul(e, att.permute(1, 0)) for e, att in zip(emb, attractors)] logits = self.forward_post_net(ys, speech_lengths) - ys = [self.recover_y_from_powerlabel(logit, raw_n_speaker) for logit, raw_n_speaker in - zip(logits, raw_n_speakers)] + ys = [ + self.recover_y_from_powerlabel(logit, raw_n_speaker) + for logit, raw_n_speaker in zip(logits, raw_n_speakers) + ] return ys, emb, attractors, raw_n_speakers def recover_y_from_powerlabel(self, logit, n_speaker): pred = torch.argmax(torch.softmax(logit, dim=-1), dim=-1) - oov_index = torch.where(pred == self.mapping_dict['oov'])[0] + oov_index = torch.where(pred == self.mapping_dict["oov"])[0] for i in oov_index: if i > 0: pred[i] = pred[i - 1] @@ -200,9 +240,13 @@ class DiarEENDOLAModel(nn.Module): pred[i] = 0 pred = [self.inv_mapping_func(i) for i in pred] decisions = [bin(num)[2:].zfill(self.max_n_speaker)[::-1] for num in pred] - decisions = torch.from_numpy( - np.stack([np.array([int(i) for i in dec]) for dec in decisions], axis=0)).to(logit.device).to( - torch.float32) + decisions = ( + torch.from_numpy( + np.stack([np.array([int(i) for i in dec]) for dec in decisions], axis=0) + ) + .to(logit.device) + .to(torch.float32) + ) decisions = decisions[:, :n_speaker] return decisions @@ -210,11 +254,11 @@ class DiarEENDOLAModel(nn.Module): if not isinstance(label, int): label = int(label) - if label in self.mapping_dict['label2dec'].keys(): - num = self.mapping_dict['label2dec'][label] + if label in self.mapping_dict["label2dec"].keys(): + num = self.mapping_dict["label2dec"][label] else: num = -1 return num def collect_feats(self, **batch: torch.Tensor) -> Dict[str, torch.Tensor]: - pass \ No newline at end of file + pass diff --git a/funasr/models/eend/eend_ola_dataloader.py b/funasr/models/eend/eend_ola_dataloader.py index 2ee9272f5..43f00fb65 100644 --- a/funasr/models/eend/eend_ola_dataloader.py +++ b/funasr/models/eend/eend_ola_dataloader.py @@ -12,17 +12,15 @@ def custom_collate(batch): speech = [torch.from_numpy(np.copy(sph)).to(torch.float32) for sph in speech] speaker_labels = [torch.from_numpy(np.copy(spk)).to(torch.float32) for spk in speaker_labels] orders = [torch.from_numpy(np.copy(o)).to(torch.int64) for o in orders] - batch = dict(speech=speech, - speaker_labels=speaker_labels, - orders=orders) + batch = dict(speech=speech, speaker_labels=speaker_labels, orders=orders) return keys, batch class EENDOLADataset(Dataset): def __init__( - self, - data_file, + self, + data_file, ): self.data_file = data_file with open(data_file) as f: @@ -44,14 +42,16 @@ class EENDOLADataset(Dataset): return key, speech, speaker_label, order -class EENDOLADataLoader(): +class EENDOLADataLoader: def __init__(self, data_file, batch_size, shuffle=True, num_workers=8): dataset = EENDOLADataset(data_file) - self.data_loader = DataLoader(dataset, - batch_size=batch_size, - collate_fn=custom_collate, - shuffle=shuffle, - num_workers=num_workers) + self.data_loader = DataLoader( + dataset, + batch_size=batch_size, + collate_fn=custom_collate, + shuffle=shuffle, + num_workers=num_workers, + ) def build_iter(self, epoch): - return self.data_loader \ No newline at end of file + return self.data_loader diff --git a/funasr/models/eend/encoder.py b/funasr/models/eend/encoder.py index 430f0d02c..0dbd98fdf 100644 --- a/funasr/models/eend/encoder.py +++ b/funasr/models/eend/encoder.py @@ -20,8 +20,7 @@ class MultiHeadSelfAttention(nn.Module): q = self.linearQ(x).view(batch_size, -1, self.h, self.d_k) k = self.linearK(x).view(batch_size, -1, self.h, self.d_k) v = self.linearV(x).view(batch_size, -1, self.h, self.d_k) - scores = torch.matmul( - q.permute(0, 2, 1, 3), k.permute(0, 2, 3, 1)) / math.sqrt(self.d_k) + scores = torch.matmul(q.permute(0, 2, 1, 3), k.permute(0, 2, 3, 1)) / math.sqrt(self.d_k) if x_mask is not None: x_mask = x_mask.unsqueeze(1) scores = scores.masked_fill(x_mask == 0, -1e9) @@ -61,9 +60,7 @@ class PositionalEncoding(torch.nn.Module): return pe = torch.zeros(x.size(1), self.d_model) if self.reverse: - position = torch.arange( - x.size(1) - 1, -1, -1.0, dtype=torch.float32 - ).unsqueeze(1) + position = torch.arange(x.size(1) - 1, -1, -1.0, dtype=torch.float32).unsqueeze(1) else: position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1) div_term = torch.exp( @@ -82,38 +79,40 @@ class PositionalEncoding(torch.nn.Module): class EENDOLATransformerEncoder(nn.Module): - def __init__(self, - idim: int, - n_layers: int, - n_units: int, - e_units: int = 2048, - h: int = 4, - dropout_rate: float = 0.1, - use_pos_emb: bool = False): + def __init__( + self, + idim: int, + n_layers: int, + n_units: int, + e_units: int = 2048, + h: int = 4, + dropout_rate: float = 0.1, + use_pos_emb: bool = False, + ): super(EENDOLATransformerEncoder, self).__init__() self.linear_in = nn.Linear(idim, n_units) self.lnorm_in = nn.LayerNorm(n_units) self.n_layers = n_layers self.dropout = nn.Dropout(dropout_rate) for i in range(n_layers): - setattr(self, '{}{:d}'.format("lnorm1_", i), - nn.LayerNorm(n_units)) - setattr(self, '{}{:d}'.format("self_att_", i), - MultiHeadSelfAttention(n_units, h)) - setattr(self, '{}{:d}'.format("lnorm2_", i), - nn.LayerNorm(n_units)) - setattr(self, '{}{:d}'.format("ff_", i), - PositionwiseFeedForward(n_units, e_units, dropout_rate)) + setattr(self, "{}{:d}".format("lnorm1_", i), nn.LayerNorm(n_units)) + setattr(self, "{}{:d}".format("self_att_", i), MultiHeadSelfAttention(n_units, h)) + setattr(self, "{}{:d}".format("lnorm2_", i), nn.LayerNorm(n_units)) + setattr( + self, + "{}{:d}".format("ff_", i), + PositionwiseFeedForward(n_units, e_units, dropout_rate), + ) self.lnorm_out = nn.LayerNorm(n_units) def __call__(self, x, x_mask=None): BT_size = x.shape[0] * x.shape[1] e = self.linear_in(x.reshape(BT_size, -1)) for i in range(self.n_layers): - e = getattr(self, '{}{:d}'.format("lnorm1_", i))(e) - s = getattr(self, '{}{:d}'.format("self_att_", i))(e, x.shape[0], x_mask) + e = getattr(self, "{}{:d}".format("lnorm1_", i))(e) + s = getattr(self, "{}{:d}".format("self_att_", i))(e, x.shape[0], x_mask) e = e + self.dropout(s) - e = getattr(self, '{}{:d}'.format("lnorm2_", i))(e) - s = getattr(self, '{}{:d}'.format("ff_", i))(e) + e = getattr(self, "{}{:d}".format("lnorm2_", i))(e) + s = getattr(self, "{}{:d}".format("ff_", i))(e) e = e + self.dropout(s) - return self.lnorm_out(e) \ No newline at end of file + return self.lnorm_out(e) diff --git a/funasr/models/eend/encoder_decoder_attractor.py b/funasr/models/eend/encoder_decoder_attractor.py index 45ac98219..65007912d 100644 --- a/funasr/models/eend/encoder_decoder_attractor.py +++ b/funasr/models/eend/encoder_decoder_attractor.py @@ -25,26 +25,40 @@ class EncoderDecoderAttractor(nn.Module): max_zlen = torch.max(zlens).to(torch.int).item() zeros = [self.enc0_dropout(z) for z in zeros] zeros = nn.utils.rnn.pad_sequence(zeros, batch_first=True, padding_value=-1) - zeros = nn.utils.rnn.pack_padded_sequence(zeros, zlens, batch_first=True, enforce_sorted=False) + zeros = nn.utils.rnn.pack_padded_sequence( + zeros, zlens, batch_first=True, enforce_sorted=False + ) attractors, (_, _) = self.decoder(zeros, (hx, cx)) - attractors = nn.utils.rnn.pad_packed_sequence(attractors, batch_first=True, padding_value=-1, - total_length=max_zlen)[0] - attractors = [att[:zlens[i].to(torch.int).item()] for i, att in enumerate(attractors)] + attractors = nn.utils.rnn.pad_packed_sequence( + attractors, batch_first=True, padding_value=-1, total_length=max_zlen + )[0] + attractors = [att[: zlens[i].to(torch.int).item()] for i, att in enumerate(attractors)] return attractors def forward(self, xs, n_speakers): - zeros = [torch.zeros(n_spk + 1, self.n_units).to(torch.float32).to(xs[0].device) for n_spk in n_speakers] + zeros = [ + torch.zeros(n_spk + 1, self.n_units).to(torch.float32).to(xs[0].device) + for n_spk in n_speakers + ] attractors = self.forward_core(xs, zeros) - labels = torch.cat([torch.from_numpy(np.array([[1] * n_spk + [0]], np.float32)) for n_spk in n_speakers], dim=1) + labels = torch.cat( + [torch.from_numpy(np.array([[1] * n_spk + [0]], np.float32)) for n_spk in n_speakers], + dim=1, + ) labels = labels.to(xs[0].device) - logit = torch.cat([self.counter(att).view(-1, n_spk + 1) for att, n_spk in zip(attractors, n_speakers)], dim=1) + logit = torch.cat( + [self.counter(att).view(-1, n_spk + 1) for att, n_spk in zip(attractors, n_speakers)], + dim=1, + ) loss = F.binary_cross_entropy(torch.sigmoid(logit), labels) attractors = [att[slice(0, att.shape[0] - 1)] for att in attractors] return loss, attractors def estimate(self, xs, max_n_speakers=15): - zeros = [torch.zeros(max_n_speakers, self.n_units).to(torch.float32).to(xs[0].device) for _ in xs] + zeros = [ + torch.zeros(max_n_speakers, self.n_units).to(torch.float32).to(xs[0].device) for _ in xs + ] attractors = self.forward_core(xs, zeros) probs = [torch.sigmoid(torch.flatten(self.counter(att))) for att in attractors] return attractors, probs diff --git a/funasr/models/eend/utils/feature.py b/funasr/models/eend/utils/feature.py index 544a3521d..1dd55dc05 100644 --- a/funasr/models/eend/utils/feature.py +++ b/funasr/models/eend/utils/feature.py @@ -8,13 +8,13 @@ import librosa def get_input_dim( - frame_size, - context_size, - transform_type, + frame_size, + context_size, + transform_type, ): - if transform_type.startswith('logmel23'): + if transform_type.startswith("logmel23"): frame_size = 23 - elif transform_type.startswith('logmel'): + elif transform_type.startswith("logmel"): frame_size = 40 else: fft_size = 1 << (frame_size - 1).bit_length() @@ -23,11 +23,8 @@ def get_input_dim( return input_dim -def transform( - Y, - transform_type=None, - dtype=np.float32): - """ Transform STFT feature +def transform(Y, transform_type=None, dtype=np.float32): + """Transform STFT feature Args: Y: STFT @@ -42,37 +39,37 @@ def transform( Y = np.abs(Y) if not transform_type: pass - elif transform_type == 'log': + elif transform_type == "log": Y = np.log(np.maximum(Y, 1e-10)) - elif transform_type == 'logmel': + elif transform_type == "logmel": n_fft = 2 * (Y.shape[1] - 1) sr = 16000 n_mels = 40 mel_basis = librosa.filters.mel(sr, n_fft, n_mels) - Y = np.dot(Y ** 2, mel_basis.T) + Y = np.dot(Y**2, mel_basis.T) Y = np.log10(np.maximum(Y, 1e-10)) - elif transform_type == 'logmel23': + elif transform_type == "logmel23": n_fft = 2 * (Y.shape[1] - 1) sr = 8000 n_mels = 23 mel_basis = librosa.filters.mel(sr, n_fft, n_mels) - Y = np.dot(Y ** 2, mel_basis.T) + Y = np.dot(Y**2, mel_basis.T) Y = np.log10(np.maximum(Y, 1e-10)) - elif transform_type == 'logmel23_mn': + elif transform_type == "logmel23_mn": n_fft = 2 * (Y.shape[1] - 1) sr = 8000 n_mels = 23 mel_basis = librosa.filters.mel(sr, n_fft, n_mels) - Y = np.dot(Y ** 2, mel_basis.T) + Y = np.dot(Y**2, mel_basis.T) Y = np.log10(np.maximum(Y, 1e-10)) mean = np.mean(Y, axis=0) Y = Y - mean - elif transform_type == 'logmel23_swn': + elif transform_type == "logmel23_swn": n_fft = 2 * (Y.shape[1] - 1) sr = 8000 n_mels = 23 mel_basis = librosa.filters.mel(sr, n_fft, n_mels) - Y = np.dot(Y ** 2, mel_basis.T) + Y = np.dot(Y**2, mel_basis.T) Y = np.log10(np.maximum(Y, 1e-10)) # b = np.ones(300)/300 # mean = scipy.signal.convolve2d(Y, b[:, None], mode='same') @@ -84,32 +81,31 @@ def transform( th = (np.mean(powers[powers >= th]) + np.mean(powers[powers < th])) / 2 mean = np.mean(Y[powers > th, :], axis=0) Y = Y - mean - elif transform_type == 'logmel23_mvn': + elif transform_type == "logmel23_mvn": n_fft = 2 * (Y.shape[1] - 1) sr = 8000 n_mels = 23 mel_basis = librosa.filters.mel(sr, n_fft, n_mels) - Y = np.dot(Y ** 2, mel_basis.T) + Y = np.dot(Y**2, mel_basis.T) Y = np.log10(np.maximum(Y, 1e-10)) mean = np.mean(Y, axis=0) Y = Y - mean std = np.maximum(np.std(Y, axis=0), 1e-10) Y = Y / std else: - raise ValueError('Unknown transform_type: %s' % transform_type) + raise ValueError("Unknown transform_type: %s" % transform_type) return Y.astype(dtype) def subsample(Y, T, subsampling=1): - """ Frame subsampling - """ + """Frame subsampling""" Y_ss = Y[::subsampling] T_ss = T[::subsampling] return Y_ss, T_ss def splice(Y, context_size=0): - """ Frame splicing + """Frame splicing Args: Y: feature @@ -122,22 +118,18 @@ def splice(Y, context_size=0): Y_spliced: spliced feature (n_frames, n_featdim * (2 * context_size + 1))-shaped """ - Y_pad = np.pad( - Y, - [(context_size, context_size), (0, 0)], - 'constant') + Y_pad = np.pad(Y, [(context_size, context_size), (0, 0)], "constant") Y_spliced = np.lib.stride_tricks.as_strided( np.ascontiguousarray(Y_pad), (Y.shape[0], Y.shape[1] * (2 * context_size + 1)), - (Y.itemsize * Y.shape[1], Y.itemsize), writeable=False) + (Y.itemsize * Y.shape[1], Y.itemsize), + writeable=False, + ) return Y_spliced -def stft( - data, - frame_size=1024, - frame_shift=256): - """ Compute STFT features +def stft(data, frame_size=1024, frame_shift=256): + """Compute STFT features Args: data: audio signal @@ -154,11 +146,11 @@ def stft( # HACK: The last frame is ommited # as librosa.stft produces such an excessive frame if len(data) % frame_shift == 0: - return librosa.stft(data, n_fft=fft_size, win_length=frame_size, - hop_length=frame_shift).T[:-1] + return librosa.stft(data, n_fft=fft_size, win_length=frame_size, hop_length=frame_shift).T[ + :-1 + ] else: - return librosa.stft(data, n_fft=fft_size, win_length=frame_size, - hop_length=frame_shift).T + return librosa.stft(data, n_fft=fft_size, win_length=frame_size, hop_length=frame_shift).T def _count_frames(data_len, size, shift): @@ -170,14 +162,9 @@ def _count_frames(data_len, size, shift): def get_frame_labels( - kaldi_obj, - rec, - start=0, - end=None, - frame_size=1024, - frame_shift=256, - n_speakers=None): - """ Get frame-aligned labels of given recording + kaldi_obj, rec, start=0, end=None, frame_size=1024, frame_shift=256, n_speakers=None +): + """Get frame-aligned labels of given recording Args: kaldi_obj (KaldiData) rec (str): recording id @@ -192,10 +179,8 @@ def get_frame_labels( T: label (n_frames, n_speakers)-shaped np.int32 array """ - filtered_segments = kaldi_obj.segments[kaldi_obj.segments['rec'] == rec] - speakers = np.unique( - [kaldi_obj.utt2spk[seg['utt']] for seg - in filtered_segments]).tolist() + filtered_segments = kaldi_obj.segments[kaldi_obj.segments["rec"] == rec] + speakers = np.unique([kaldi_obj.utt2spk[seg["utt"]] for seg in filtered_segments]).tolist() if n_speakers is None: n_speakers = len(speakers) es = end * frame_shift if end is not None else None @@ -206,11 +191,9 @@ def get_frame_labels( end = n_frames for seg in filtered_segments: - speaker_index = speakers.index(kaldi_obj.utt2spk[seg['utt']]) - start_frame = np.rint( - seg['st'] * rate / frame_shift).astype(int) - end_frame = np.rint( - seg['et'] * rate / frame_shift).astype(int) + speaker_index = speakers.index(kaldi_obj.utt2spk[seg["utt"]]) + start_frame = np.rint(seg["st"] * rate / frame_shift).astype(int) + end_frame = np.rint(seg["et"] * rate / frame_shift).astype(int) rel_start = rel_end = None if start <= start_frame and start_frame < end: rel_start = start_frame - start @@ -222,11 +205,9 @@ def get_frame_labels( def get_labeledSTFT( - kaldi_obj, - rec, start, end, frame_size, frame_shift, - n_speakers=None, - use_speaker_id=False): - """ Extracts STFT and corresponding labels + kaldi_obj, rec, start, end, frame_size, frame_shift, n_speakers=None, use_speaker_id=False +): + """Extracts STFT and corresponding labels Extracts STFT and corresponding diarization labels for given recording id and start/end times @@ -246,14 +227,11 @@ def get_labeledSTFT( T: label (n_frmaes, n_speakers)-shaped np.int32 array. """ - data, rate = kaldi_obj.load_wav( - rec, start * frame_shift, end * frame_shift) + data, rate = kaldi_obj.load_wav(rec, start * frame_shift, end * frame_shift) Y = stft(data, frame_size, frame_shift) filtered_segments = kaldi_obj.segments[rec] # filtered_segments = kaldi_obj.segments[kaldi_obj.segments['rec'] == rec] - speakers = np.unique( - [kaldi_obj.utt2spk[seg['utt']] for seg - in filtered_segments]).tolist() + speakers = np.unique([kaldi_obj.utt2spk[seg["utt"]] for seg in filtered_segments]).tolist() if n_speakers is None: n_speakers = len(speakers) T = np.zeros((Y.shape[0], n_speakers), dtype=np.int32) @@ -263,13 +241,11 @@ def get_labeledSTFT( S = np.zeros((Y.shape[0], len(all_speakers)), dtype=np.int32) for seg in filtered_segments: - speaker_index = speakers.index(kaldi_obj.utt2spk[seg['utt']]) + speaker_index = speakers.index(kaldi_obj.utt2spk[seg["utt"]]) if use_speaker_id: - all_speaker_index = all_speakers.index(kaldi_obj.utt2spk[seg['utt']]) - start_frame = np.rint( - seg['st'] * rate / frame_shift).astype(int) - end_frame = np.rint( - seg['et'] * rate / frame_shift).astype(int) + all_speaker_index = all_speakers.index(kaldi_obj.utt2spk[seg["utt"]]) + start_frame = np.rint(seg["st"] * rate / frame_shift).astype(int) + end_frame = np.rint(seg["et"] * rate / frame_shift).astype(int) rel_start = rel_end = None if start <= start_frame and start_frame < end: rel_start = start_frame - start diff --git a/funasr/models/eend/utils/kaldi_data.py b/funasr/models/eend/utils/kaldi_data.py index 53f6230b7..59e7a16f8 100644 --- a/funasr/models/eend/utils/kaldi_data.py +++ b/funasr/models/eend/utils/kaldi_data.py @@ -15,16 +15,14 @@ from functools import lru_cache def load_segments(segments_file): - """ load segments file as array """ + """load segments file as array""" if not os.path.exists(segments_file): return None return np.loadtxt( - segments_file, - dtype=[('utt', 'object'), - ('rec', 'object'), - ('st', 'f'), - ('et', 'f')], - ndmin=1) + segments_file, + dtype=[("utt", "object"), ("rec", "object"), ("st", "f"), ("et", "f")], + ndmin=1, + ) def load_segments_hash(segments_file): @@ -45,35 +43,33 @@ def load_segments_rechash(segments_file): utt, rec, st, et = line.strip().split() if rec not in ret: ret[rec] = [] - ret[rec].append({'utt':utt, 'st':float(st), 'et':float(et)}) + ret[rec].append({"utt": utt, "st": float(st), "et": float(et)}) return ret def load_wav_scp(wav_scp_file): - """ return dictionary { rec: wav_rxfilename } """ + """return dictionary { rec: wav_rxfilename }""" lines = [line.strip().split(None, 1) for line in open(wav_scp_file)] return {x[0]: x[1] for x in lines} @lru_cache(maxsize=1) def load_wav(wav_rxfilename, start=0, end=None): - """ This function reads audio file and return data in numpy.float32 array. - "lru_cache" holds recently loaded audio so that can be called - many times on the same audio file. - OPTIMIZE: controls lru_cache size for random access, - considering memory size + """This function reads audio file and return data in numpy.float32 array. + "lru_cache" holds recently loaded audio so that can be called + many times on the same audio file. + OPTIMIZE: controls lru_cache size for random access, + considering memory size """ - if wav_rxfilename.endswith('|'): + if wav_rxfilename.endswith("|"): # input piped command - p = subprocess.Popen(wav_rxfilename[:-1], shell=True, - stdout=subprocess.PIPE) - data, samplerate = sf.load(io.BytesIO(p.stdout.read()), - dtype='float32') + p = subprocess.Popen(wav_rxfilename[:-1], shell=True, stdout=subprocess.PIPE) + data, samplerate = sf.load(io.BytesIO(p.stdout.read()), dtype="float32") # cannot seek data = data[start:end] - elif wav_rxfilename == '-': + elif wav_rxfilename == "-": # stdin - data, samplerate = sf.load(sys.stdin, dtype='float32') + data, samplerate = sf.load(sys.stdin, dtype="float32") # cannot seek data = data[start:end] else: @@ -83,13 +79,13 @@ def load_wav(wav_rxfilename, start=0, end=None): def load_utt2spk(utt2spk_file): - """ returns dictionary { uttid: spkid } """ + """returns dictionary { uttid: spkid }""" lines = [line.strip().split(None, 1) for line in open(utt2spk_file)] return {x[0]: x[1] for x in lines} def load_spk2utt(spk2utt_file): - """ returns dictionary { spkid: list of uttids } """ + """returns dictionary { spkid: list of uttids }""" if not os.path.exists(spk2utt_file): return None lines = [line.strip().split() for line in open(spk2utt_file)] @@ -97,7 +93,7 @@ def load_spk2utt(spk2utt_file): def load_reco2dur(reco2dur_file): - """ returns dictionary { recid: duration } """ + """returns dictionary { recid: duration }""" if not os.path.exists(reco2dur_file): return None lines = [line.strip().split(None, 1) for line in open(reco2dur_file)] @@ -105,7 +101,7 @@ def load_reco2dur(reco2dur_file): def process_wav(wav_rxfilename, process): - """ This function returns preprocessed wav_rxfilename + """This function returns preprocessed wav_rxfilename Args: wav_rxfilename: input process: command which can be connected via pipe, @@ -113,7 +109,7 @@ def process_wav(wav_rxfilename, process): Returns: wav_rxfilename: output piped command """ - if wav_rxfilename.endswith('|'): + if wav_rxfilename.endswith("|"): # input piped command return wav_rxfilename + process + "|" else: @@ -122,18 +118,18 @@ def process_wav(wav_rxfilename, process): def extract_segments(wavs, segments=None): - """ This function returns generator of segmented audio as - (utterance id, numpy.float32 array) - TODO?: sampling rate is not converted. + """This function returns generator of segmented audio as + (utterance id, numpy.float32 array) + TODO?: sampling rate is not converted. """ if segments is not None: # segments should be sorted by rec-id for seg in segments: - wav = wavs[seg['rec']] + wav = wavs[seg["rec"]] data, samplerate = load_wav(wav) - st_sample = np.rint(seg['st'] * samplerate).astype(int) - et_sample = np.rint(seg['et'] * samplerate).astype(int) - yield seg['utt'], data[st_sample:et_sample] + st_sample = np.rint(seg["st"] * samplerate).astype(int) + et_sample = np.rint(seg["et"] * samplerate).astype(int) + yield seg["utt"], data[st_sample:et_sample] else: # segments file not found, # wav.scp is used as segmented audio list @@ -145,18 +141,12 @@ def extract_segments(wavs, segments=None): class KaldiData: def __init__(self, data_dir): self.data_dir = data_dir - self.segments = load_segments_rechash( - os.path.join(self.data_dir, 'segments')) - self.utt2spk = load_utt2spk( - os.path.join(self.data_dir, 'utt2spk')) - self.wavs = load_wav_scp( - os.path.join(self.data_dir, 'wav.scp')) - self.reco2dur = load_reco2dur( - os.path.join(self.data_dir, 'reco2dur')) - self.spk2utt = load_spk2utt( - os.path.join(self.data_dir, 'spk2utt')) + self.segments = load_segments_rechash(os.path.join(self.data_dir, "segments")) + self.utt2spk = load_utt2spk(os.path.join(self.data_dir, "utt2spk")) + self.wavs = load_wav_scp(os.path.join(self.data_dir, "wav.scp")) + self.reco2dur = load_reco2dur(os.path.join(self.data_dir, "reco2dur")) + self.spk2utt = load_spk2utt(os.path.join(self.data_dir, "spk2utt")) def load_wav(self, recid, start=0, end=None): - data, rate = load_wav( - self.wavs[recid], start, end) + data, rate = load_wav(self.wavs[recid], start, end) return data, rate diff --git a/funasr/models/eend/utils/losses.py b/funasr/models/eend/utils/losses.py index 756952d03..82fb57358 100644 --- a/funasr/models/eend/utils/losses.py +++ b/funasr/models/eend/utils/losses.py @@ -7,7 +7,11 @@ from scipy.optimize import linear_sum_assignment def standard_loss(ys, ts): losses = [F.binary_cross_entropy(torch.sigmoid(y), t) * len(y) for y, t in zip(ys, ts)] loss = torch.sum(torch.stack(losses)) - n_frames = torch.from_numpy(np.array(np.sum([t.shape[0] for t in ts]))).to(torch.float32).to(ys[0].device) + n_frames = ( + torch.from_numpy(np.array(np.sum([t.shape[0] for t in ts]))) + .to(torch.float32) + .to(ys[0].device) + ) loss = loss / n_frames return loss @@ -31,10 +35,15 @@ def fast_batch_pit_n_speaker_loss(ys, ts): def cal_power_loss(logits, power_ts): - losses = [F.cross_entropy(input=logit, target=power_t.to(torch.long)) * len(logit) for logit, power_t in - zip(logits, power_ts)] + losses = [ + F.cross_entropy(input=logit, target=power_t.to(torch.long)) * len(logit) + for logit, power_t in zip(logits, power_ts) + ] loss = torch.sum(torch.stack(losses)) - n_frames = torch.from_numpy(np.array(np.sum([power_t.shape[0] for power_t in power_ts]))).to(torch.float32).to( - power_ts[0].device) + n_frames = ( + torch.from_numpy(np.array(np.sum([power_t.shape[0] for power_t in power_ts]))) + .to(torch.float32) + .to(power_ts[0].device) + ) loss = loss / n_frames return loss diff --git a/funasr/models/eend/utils/power.py b/funasr/models/eend/utils/power.py index 7144e24da..b6b32c556 100644 --- a/funasr/models/eend/utils/power.py +++ b/funasr/models/eend/utils/power.py @@ -20,14 +20,14 @@ def generate_mapping_dict(max_speaker_num=6, max_olp_speaker_num=3): all_kinds_order = sorted(all_kinds) mapping_dict = {} - mapping_dict['dec2label'] = {} - mapping_dict['label2dec'] = {} + mapping_dict["dec2label"] = {} + mapping_dict["label2dec"] = {} for i in range(len(all_kinds_order)): dec = all_kinds_order[i] - mapping_dict['dec2label'][dec] = i - mapping_dict['label2dec'][i] = dec + mapping_dict["dec2label"][dec] = i + mapping_dict["label2dec"][i] = dec oov_id = len(all_kinds_order) - mapping_dict['oov'] = oov_id + mapping_dict["oov"] = oov_id return mapping_dict @@ -45,10 +45,10 @@ def raw_dec_trans(x, max_speaker_num): def mapping_func(num, mapping_dict): - if num in mapping_dict['dec2label'].keys(): - label = mapping_dict['dec2label'][num] + if num in mapping_dict["dec2label"].keys(): + label = mapping_dict["dec2label"][num] else: - label = mapping_dict['oov'] + label = mapping_dict["oov"] return label @@ -79,16 +79,25 @@ def generate_perm_pse(label, n_speaker, mapping_dict, max_speaker_num, max_olp_s perms = np.array(list(permutations(range(n_speaker)))).astype(np.float32) perms = torch.from_numpy(perms).to(label.device).to(torch.int64) perm_labels = [label[:, perm] for perm in perms] - perm_pse_labels = [create_powerlabel(perm_label.cpu().numpy(), mapping_dict, max_speaker_num). - to(perm_label.device, non_blocking=True) for perm_label in perm_labels] + perm_pse_labels = [ + create_powerlabel(perm_label.cpu().numpy(), mapping_dict, max_speaker_num).to( + perm_label.device, non_blocking=True + ) + for perm_label in perm_labels + ] return perm_labels, perm_pse_labels -def generate_min_pse(label, n_speaker, mapping_dict, max_speaker_num, pse_logit, max_olp_speaker_num=3): - perm_labels, perm_pse_labels = generate_perm_pse(label, n_speaker, mapping_dict, max_speaker_num, - max_olp_speaker_num=max_olp_speaker_num) - losses = [F.cross_entropy(input=pse_logit, target=perm_pse_label.to(torch.long)) * len(pse_logit) - for perm_pse_label in perm_pse_labels] +def generate_min_pse( + label, n_speaker, mapping_dict, max_speaker_num, pse_logit, max_olp_speaker_num=3 +): + perm_labels, perm_pse_labels = generate_perm_pse( + label, n_speaker, mapping_dict, max_speaker_num, max_olp_speaker_num=max_olp_speaker_num + ) + losses = [ + F.cross_entropy(input=pse_logit, target=perm_pse_label.to(torch.long)) * len(pse_logit) + for perm_pse_label in perm_pse_labels + ] loss = torch.stack(losses) min_index = torch.argmin(loss) selected_perm_label, selected_pse_label = perm_labels[min_index], perm_pse_labels[min_index] diff --git a/funasr/models/eend/utils/report.py b/funasr/models/eend/utils/report.py index cd100f066..eb1bdacd1 100644 --- a/funasr/models/eend/utils/report.py +++ b/funasr/models/eend/utils/report.py @@ -6,13 +6,13 @@ from funasr.models.eend.utils.power import create_powerlabel from itertools import combinations metrics = [ - ('diarization_error', 'speaker_scored', 'DER'), - ('speech_miss', 'speech_scored', 'SAD_MR'), - ('speech_falarm', 'speech_scored', 'SAD_FR'), - ('speaker_miss', 'speaker_scored', 'MI'), - ('speaker_falarm', 'speaker_scored', 'FA'), - ('speaker_error', 'speaker_scored', 'CF'), - ('correct', 'frames', 'accuracy') + ("diarization_error", "speaker_scored", "DER"), + ("speech_miss", "speech_scored", "SAD_MR"), + ("speech_falarm", "speech_scored", "SAD_FR"), + ("speaker_miss", "speaker_scored", "MI"), + ("speaker_falarm", "speaker_scored", "FA"), + ("speaker_error", "speaker_scored", "CF"), + ("correct", "frames", "accuracy"), ] @@ -20,9 +20,9 @@ def recover_prediction(y, n_speaker): if n_speaker <= 1: return y elif n_speaker == 2: - com_index = torch.from_numpy( - np.array(list(combinations(np.arange(n_speaker), 2)))).to( - y.dtype) + com_index = torch.from_numpy(np.array(list(combinations(np.arange(n_speaker), 2)))).to( + y.dtype + ) num_coms = com_index.shape[0] y_single = y[:, :-num_coms] y_olp = y[:, -num_coms:] @@ -30,24 +30,29 @@ def recover_prediction(y, n_speaker): olp_map_index = torch.stack(olp_map_index, dim=1) com_map_index = com_index[olp_map_index[:, -1]] speaker_map_index = torch.from_numpy(np.array(com_map_index)).view(-1).to(torch.int64) - frame_map_index = olp_map_index[:, 0][:, None].repeat([1, 2]).view(-1).to( - torch.int64) + frame_map_index = olp_map_index[:, 0][:, None].repeat([1, 2]).view(-1).to(torch.int64) y_single[frame_map_index] = 0 y_single[frame_map_index, speaker_map_index] = 1 return y_single else: - olp2_com_index = torch.from_numpy(np.array(list(combinations(np.arange(n_speaker), 2)))).to(y.dtype) + olp2_com_index = torch.from_numpy(np.array(list(combinations(np.arange(n_speaker), 2)))).to( + y.dtype + ) olp2_num_coms = olp2_com_index.shape[0] - olp3_com_index = torch.from_numpy(np.array(list(combinations(np.arange(n_speaker), 3)))).to(y.dtype) + olp3_com_index = torch.from_numpy(np.array(list(combinations(np.arange(n_speaker), 3)))).to( + y.dtype + ) olp3_num_coms = olp3_com_index.shape[0] y_single = y[:, :n_speaker] - y_olp2 = y[:, n_speaker:n_speaker + olp2_num_coms] + y_olp2 = y[:, n_speaker : n_speaker + olp2_num_coms] y_olp3 = y[:, -olp3_num_coms:] olp3_map_index = torch.where(y_olp3 > 0.5) olp3_map_index = torch.stack(olp3_map_index, dim=1) olp3_com_map_index = olp3_com_index[olp3_map_index[:, -1]] - olp3_speaker_map_index = torch.from_numpy(np.array(olp3_com_map_index)).view(-1).to(torch.int64) + olp3_speaker_map_index = ( + torch.from_numpy(np.array(olp3_com_map_index)).view(-1).to(torch.int64) + ) olp3_frame_map_index = olp3_map_index[:, 0][:, None].repeat([1, 3]).view(-1).to(torch.int64) y_single[olp3_frame_map_index] = 0 y_single[olp3_frame_map_index, olp3_speaker_map_index] = 1 @@ -56,14 +61,16 @@ def recover_prediction(y, n_speaker): olp2_map_index = torch.where(y_olp2 > 0.5) olp2_map_index = torch.stack(olp2_map_index, dim=1) olp2_com_map_index = olp2_com_index[olp2_map_index[:, -1]] - olp2_speaker_map_index = torch.from_numpy(np.array(olp2_com_map_index)).view(-1).to(torch.int64) + olp2_speaker_map_index = ( + torch.from_numpy(np.array(olp2_com_map_index)).view(-1).to(torch.int64) + ) olp2_frame_map_index = olp2_map_index[:, 0][:, None].repeat([1, 2]).view(-1).to(torch.int64) y_single[olp2_frame_map_index] = 0 y_single[olp2_frame_map_index, olp2_speaker_map_index] = 1 return y_single -class PowerReporter(): +class PowerReporter: def __init__(self, valid_data_loader, mapping_dict, max_n_speaker): valid_data_loader_cp = copy.deepcopy(valid_data_loader) self.valid_data_loader = valid_data_loader_cp @@ -77,22 +84,31 @@ class PowerReporter(): def report_val(self, model, eidx, device): model.eval() ud_valid_start = time.time() - valid_res, valid_loss, stats_keys, vad_valid_accuracy = self.report_core(model, self.valid_data_loader, device) + valid_res, valid_loss, stats_keys, vad_valid_accuracy = self.report_core( + model, self.valid_data_loader, device + ) # Epoch Display - valid_der = valid_res['diarization_error'] / valid_res['speaker_scored'] - valid_accuracy = valid_res['correct'].to(torch.float32) / valid_res['frames'] * 100 + valid_der = valid_res["diarization_error"] / valid_res["speaker_scored"] + valid_accuracy = valid_res["correct"].to(torch.float32) / valid_res["frames"] * 100 vad_valid_accuracy = vad_valid_accuracy * 100 - print('Epoch ', eidx + 1, 'Valid Loss ', valid_loss, 'Valid_DER %.5f' % valid_der, - 'Valid_Accuracy %.5f%% ' % valid_accuracy, 'VAD_Valid_Accuracy %.5f%% ' % vad_valid_accuracy) - ud_valid = (time.time() - ud_valid_start) / 60. - print('Valid cost time ... ', ud_valid) + print( + "Epoch ", + eidx + 1, + "Valid Loss ", + valid_loss, + "Valid_DER %.5f" % valid_der, + "Valid_Accuracy %.5f%% " % valid_accuracy, + "VAD_Valid_Accuracy %.5f%% " % vad_valid_accuracy, + ) + ud_valid = (time.time() - ud_valid_start) / 60.0 + print("Valid cost time ... ", ud_valid) def inv_mapping_func(self, label, mapping_dict): if not isinstance(label, int): label = int(label) - if label in mapping_dict['label2dec'].keys(): - num = mapping_dict['label2dec'][label] + if label in mapping_dict["label2dec"].keys(): + num = mapping_dict["label2dec"][label] else: num = -1 return num @@ -100,22 +116,24 @@ class PowerReporter(): def report_core(self, model, data_loader, device): res = {} for item in metrics: - res[item[0]] = 0. - res[item[1]] = 0. + res[item[0]] = 0.0 + res[item[1]] = 0.0 with torch.no_grad(): - loss_s = 0. + loss_s = 0.0 uidx = 0 for xs, ts, orders in data_loader: xs = [x.to(device) for x in xs] ts = [t.to(device) for t in ts] orders = [o.to(device) for o in orders] - loss, pit_loss, mpit_loss, att_loss, ys, logits, labels, attractors = model(xs, ts, orders) + loss, pit_loss, mpit_loss, att_loss, ys, logits, labels, attractors = model( + xs, ts, orders + ) loss_s += loss.item() uidx += 1 for logit, t, att in zip(logits, labels, attractors): pred = torch.argmax(torch.softmax(logit, dim=-1), dim=-1) # (T, ) - oov_index = torch.where(pred == self.mapping_dict['oov'])[0] + oov_index = torch.where(pred == self.mapping_dict["oov"])[0] for i in oov_index: if i > 0: pred[i] = pred[i - 1] @@ -123,15 +141,19 @@ class PowerReporter(): pred[i] = 0 pred = [self.inv_mapping_func(i, self.mapping_dict) for i in pred] decisions = [bin(num)[2:].zfill(self.max_n_speaker)[::-1] for num in pred] - decisions = torch.from_numpy( - np.stack([np.array([int(i) for i in dec]) for dec in decisions], axis=0)).to(att.device).to( - torch.float32) - decisions = decisions[:, :att.shape[0]] + decisions = ( + torch.from_numpy( + np.stack([np.array([int(i) for i in dec]) for dec in decisions], axis=0) + ) + .to(att.device) + .to(torch.float32) + ) + decisions = decisions[:, : att.shape[0]] stats = self.calc_diarization_error(decisions, t) - res['speaker_scored'] += stats['speaker_scored'] - res['speech_scored'] += stats['speech_scored'] - res['frames'] += stats['frames'] + res["speaker_scored"] += stats["speaker_scored"] + res["speech_scored"] += stats["speech_scored"] + res["frames"] += stats["frames"] for item in metrics: res[item[0]] += stats[item[0]] loss_s /= uidx @@ -140,20 +162,21 @@ class PowerReporter(): return res, loss_s, stats.keys(), vad_acc def calc_diarization_error(self, decisions, label, label_delay=0): - label = label[:len(label) - label_delay, ...] + label = label[: len(label) - label_delay, ...] n_ref = torch.sum(label, dim=-1) n_sys = torch.sum(decisions, dim=-1) res = {} - res['speech_scored'] = torch.sum(n_ref > 0) - res['speech_miss'] = torch.sum((n_ref > 0) & (n_sys == 0)) - res['speech_falarm'] = torch.sum((n_ref == 0) & (n_sys > 0)) - res['speaker_scored'] = torch.sum(n_ref) - res['speaker_miss'] = torch.sum(torch.max(n_ref - n_sys, torch.zeros_like(n_ref))) - res['speaker_falarm'] = torch.sum(torch.max(n_sys - n_ref, torch.zeros_like(n_ref))) + res["speech_scored"] = torch.sum(n_ref > 0) + res["speech_miss"] = torch.sum((n_ref > 0) & (n_sys == 0)) + res["speech_falarm"] = torch.sum((n_ref == 0) & (n_sys > 0)) + res["speaker_scored"] = torch.sum(n_ref) + res["speaker_miss"] = torch.sum(torch.max(n_ref - n_sys, torch.zeros_like(n_ref))) + res["speaker_falarm"] = torch.sum(torch.max(n_sys - n_ref, torch.zeros_like(n_ref))) n_map = torch.sum(((label == 1) & (decisions == 1)), dim=-1).to(torch.float32) - res['speaker_error'] = torch.sum(torch.min(n_ref, n_sys) - n_map) - res['correct'] = torch.sum(label == decisions) / label.shape[1] - res['diarization_error'] = ( - res['speaker_miss'] + res['speaker_falarm'] + res['speaker_error']) - res['frames'] = len(label) + res["speaker_error"] = torch.sum(torch.min(n_ref, n_sys) - n_map) + res["correct"] = torch.sum(label == decisions) / label.shape[1] + res["diarization_error"] = ( + res["speaker_miss"] + res["speaker_falarm"] + res["speaker_error"] + ) + res["frames"] = len(label) return res diff --git a/funasr/models/emotion2vec/audio.py b/funasr/models/emotion2vec/audio.py index d21500b24..260957343 100644 --- a/funasr/models/emotion2vec/audio.py +++ b/funasr/models/emotion2vec/audio.py @@ -22,7 +22,6 @@ from funasr.models.emotion2vec.base import ModalitySpecificEncoder, get_alibi_bi class AudioEncoder(ModalitySpecificEncoder): - def __init__( self, modality_cfg, @@ -91,9 +90,7 @@ class AudioEncoder(ModalitySpecificEncoder): ) decoder = ( - Decoder1d(modality_cfg.decoder, embed_dim) - if modality_cfg.decoder is not None - else None + Decoder1d(modality_cfg.decoder, embed_dim) if modality_cfg.decoder is not None else None ) alibi_bias_fn = partial(get_alibi_bias, alibi_biases=alibi_biases) @@ -144,13 +141,9 @@ class AudioEncoder(ModalitySpecificEncoder): output_lengths - 1, ) ] = 1 - padding_mask = ( - 1 - padding_mask.flip([-1]).cumsum(-1).flip([-1]) - ).bool() + padding_mask = (1 - padding_mask.flip([-1]).cumsum(-1).flip([-1])).bool() else: - padding_mask = torch.zeros( - x.shape[:2], dtype=torch.bool, device=x.device - ) + padding_mask = torch.zeros(x.shape[:2], dtype=torch.bool, device=x.device) return padding_mask diff --git a/funasr/models/emotion2vec/base.py b/funasr/models/emotion2vec/base.py index cd87a99cc..1b4301ea3 100644 --- a/funasr/models/emotion2vec/base.py +++ b/funasr/models/emotion2vec/base.py @@ -22,7 +22,6 @@ from funasr.models.emotion2vec.fairseq_modules import index_put logger = logging.getLogger(__name__) - MaskSeed = namedtuple("MaskSeed", ["seed", "update", "ids"]) MaskInfo = namedtuple("MaskInfo", ["x_unmasked", "mask", "ids_restore", "ids_keep"]) @@ -69,13 +68,17 @@ class ModalitySpecificEncoder(nn.Module): self.alibi_scale = nn.Parameter( torch.full( ( - (modality_cfg.prenet_depth + modality_cfg.model_depth) - if modality_cfg.learned_alibi_scale_per_layer - else 1, + ( + (modality_cfg.prenet_depth + modality_cfg.model_depth) + if modality_cfg.learned_alibi_scale_per_layer + else 1 + ), 1, - self.modality_cfg.num_alibi_heads - if modality_cfg.learned_alibi_scale_per_head - else 1, + ( + self.modality_cfg.num_alibi_heads + if modality_cfg.learned_alibi_scale_per_head + else 1 + ), 1, 1, ), @@ -96,9 +99,7 @@ class ModalitySpecificEncoder(nn.Module): device="cpu", ) self.alibi_bias = nn.Parameter(alibi_bias) - self.get_alibi_bias = partial( - _learned_alibi_bias, alibi_bias=self.alibi_bias - ) + self.get_alibi_bias = partial(_learned_alibi_bias, alibi_bias=self.alibi_bias) def upgrade_state_dict_named(self, state_dict, name): k = f"{name}.alibi_scale" @@ -147,9 +148,7 @@ class ModalitySpecificEncoder(nn.Module): if self.local_grad_mult == 1.0: x = self.local_encoder(features) else: - x = GradMultiply.apply( - self.local_encoder(features), self.local_grad_mult - ) + x = GradMultiply.apply(self.local_encoder(features), self.local_grad_mult) else: with torch.no_grad(): x = self.local_encoder(features) @@ -188,8 +187,7 @@ class ModalitySpecificEncoder(nn.Module): x = x.repeat_interleave(clone_batch, 0) if mask_seeds is not None: clone_hash = [ - int(hash((mask_seeds.seed, ind)) % 1e10) - for ind in range(clone_batch - 1) + int(hash((mask_seeds.seed, ind)) % 1e10) for ind in range(clone_batch - 1) ] clone_hash = torch.tensor([0] + clone_hash).long().view(1, -1) @@ -197,9 +195,7 @@ class ModalitySpecificEncoder(nn.Module): id = id.repeat_interleave(clone_batch, 0) id = id.view(-1, clone_batch) + clone_hash.to(id) id = id.view(-1) - mask_seeds = MaskSeed( - seed=mask_seeds.seed, update=mask_seeds.update, ids=id - ) + mask_seeds = MaskSeed(seed=mask_seeds.seed, update=mask_seeds.update, ids=id) if padding_mask is not None: padding_mask = padding_mask.repeat_interleave(clone_batch, 0) @@ -268,9 +264,7 @@ class ModalitySpecificEncoder(nn.Module): x, masked_padding_mask, alibi_bias, - alibi_scale[: self.modality_cfg.prenet_depth] - if alibi_scale is not None - else None, + alibi_scale[: self.modality_cfg.prenet_depth] if alibi_scale is not None else None, ) return { @@ -278,9 +272,11 @@ class ModalitySpecificEncoder(nn.Module): "local_features": local_features, "padding_mask": masked_padding_mask, "alibi_bias": alibi_bias, - "alibi_scale": alibi_scale[self.modality_cfg.prenet_depth :] - if alibi_scale is not None and alibi_scale.size(0) > 1 - else alibi_scale, + "alibi_scale": ( + alibi_scale[self.modality_cfg.prenet_depth :] + if alibi_scale is not None and alibi_scale.size(0) > 1 + else alibi_scale + ), "encoder_mask": mask_info, } @@ -405,9 +401,7 @@ class ModalitySpecificEncoder(nn.Module): x = x * (1 - mask.type_as(x).unsqueeze(-1)) else: num_masks = mask.sum().item() - masks = x.new_empty(num_masks, x.size(-1)).normal_( - 0, cfg.mask_noise_std - ) + masks = x.new_empty(num_masks, x.size(-1)).normal_(0, cfg.mask_noise_std) x = index_put(x, mask, masks) if cfg.mask_channel_prob > 0: mask_channel = compute_mask_indices( @@ -417,10 +411,7 @@ class ModalitySpecificEncoder(nn.Module): cfg.mask_channel_length, ) mask_channel = ( - torch.from_numpy(mask_channel) - .to(x.device) - .unsqueeze(1) - .expand(-1, T, -1) + torch.from_numpy(mask_channel).to(x.device).unsqueeze(1).expand(-1, T, -1) ) x = index_put(x, mask_channel, 0) return x @@ -445,9 +436,7 @@ def random_masking(x, mask_ratio, mask_seed: Optional[MaskSeed]): generator = None if mask_seed is not None: - seed = int( - hash((mask_seed.seed, mask_seed.update, mask_seed.ids.sum().item())) % 1e6 - ) + seed = int(hash((mask_seed.seed, mask_seed.update, mask_seed.ids.sum().item())) % 1e6) generator = torch.Generator(device=x.device) generator.manual_seed(seed) @@ -470,9 +459,7 @@ def random_masking(x, mask_ratio, mask_seed: Optional[MaskSeed]): ids_restore = ids_restore.unsqueeze(-1).expand(-1, -1, D) - return MaskInfo( - x_unmasked=x_unmasked, mask=mask, ids_restore=ids_restore, ids_keep=ids_keep - ) + return MaskInfo(x_unmasked=x_unmasked, mask=mask, ids_restore=ids_restore, ids_keep=ids_keep) def gather_unmasked(x: torch.Tensor, mask_info: MaskInfo) -> torch.Tensor: @@ -525,10 +512,7 @@ def get_alibi( # autoregressive model so we want a symmetric mask with 0 on the # diagonal and other wise linear decreasing valuees pos_bias = ( - torch.abs( - torch.arange(maxpos).unsqueeze(0) - torch.arange(maxpos).unsqueeze(1) - ) - * -1 + torch.abs(torch.arange(maxpos).unsqueeze(0) - torch.arange(maxpos).unsqueeze(1)) * -1 ) elif dims == 2: if distance == "manhattan": @@ -553,9 +537,7 @@ def get_alibi( else: raise Exception(f"unsupported number of alibi dims: {dims}") - alibi_bias = slopes.unsqueeze(1).unsqueeze(1) * pos_bias.unsqueeze(0).expand( - attn_heads, -1, -1 - ) + alibi_bias = slopes.unsqueeze(1).unsqueeze(1) * pos_bias.unsqueeze(0).expand(attn_heads, -1, -1) return alibi_bias diff --git a/funasr/models/emotion2vec/fairseq_modules.py b/funasr/models/emotion2vec/fairseq_modules.py index fa9888513..209dedaa3 100644 --- a/funasr/models/emotion2vec/fairseq_modules.py +++ b/funasr/models/emotion2vec/fairseq_modules.py @@ -4,142 +4,146 @@ import torch.nn.functional as F from typing import Optional, Tuple, List import numpy as np + def LayerNorm(normalized_shape, eps=1e-5, elementwise_affine=True, export=False): - return torch.nn.LayerNorm(normalized_shape, eps, elementwise_affine) + return torch.nn.LayerNorm(normalized_shape, eps, elementwise_affine) + class SamePad(nn.Module): - def __init__(self, kernel_size, causal=False): - super().__init__() - if causal: - self.remove = kernel_size - 1 - else: - self.remove = 1 if kernel_size % 2 == 0 else 0 - - def forward(self, x): - if self.remove > 0: - x = x[:, :, : -self.remove] - return x + def __init__(self, kernel_size, causal=False): + super().__init__() + if causal: + self.remove = kernel_size - 1 + else: + self.remove = 1 if kernel_size % 2 == 0 else 0 + + def forward(self, x): + if self.remove > 0: + x = x[:, :, : -self.remove] + return x + class TransposeLast(nn.Module): - def __init__(self, deconstruct_idx=None): - super().__init__() - self.deconstruct_idx = deconstruct_idx - - def forward(self, x): - if self.deconstruct_idx is not None: - x = x[self.deconstruct_idx] - return x.transpose(-2, -1) + def __init__(self, deconstruct_idx=None): + super().__init__() + self.deconstruct_idx = deconstruct_idx + + def forward(self, x): + if self.deconstruct_idx is not None: + x = x[self.deconstruct_idx] + return x.transpose(-2, -1) class Fp32LayerNorm(nn.LayerNorm): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - def forward(self, input): - output = F.layer_norm( - input.float(), - self.normalized_shape, - self.weight.float() if self.weight is not None else None, - self.bias.float() if self.bias is not None else None, - self.eps, - ) - return output.type_as(input) + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def forward(self, input): + output = F.layer_norm( + input.float(), + self.normalized_shape, + self.weight.float() if self.weight is not None else None, + self.bias.float() if self.bias is not None else None, + self.eps, + ) + return output.type_as(input) class Fp32GroupNorm(nn.GroupNorm): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - def forward(self, input): - output = F.group_norm( - input.float(), - self.num_groups, - self.weight.float() if self.weight is not None else None, - self.bias.float() if self.bias is not None else None, - self.eps, - ) - return output.type_as(input) + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def forward(self, input): + output = F.group_norm( + input.float(), + self.num_groups, + self.weight.float() if self.weight is not None else None, + self.bias.float() if self.bias is not None else None, + self.eps, + ) + return output.type_as(input) class ConvFeatureExtractionModel(nn.Module): - def __init__( - self, - conv_layers: List[Tuple[int, int, int]], - dropout: float = 0.0, - mode: str = "default", - conv_bias: bool = False, - ): - super().__init__() - - assert mode in {"default", "layer_norm"} - - def block( - n_in, - n_out, - k, - stride, - is_layer_norm=False, - is_group_norm=False, - conv_bias=False, - ): - def make_conv(): - conv = nn.Conv1d(n_in, n_out, k, stride=stride, bias=conv_bias) - nn.init.kaiming_normal_(conv.weight) - return conv - - assert ( - is_layer_norm and is_group_norm - ) == False, "layer norm and group norm are exclusive" - - if is_layer_norm: - return nn.Sequential( - make_conv(), - nn.Dropout(p=dropout), - nn.Sequential( - TransposeLast(), - Fp32LayerNorm(dim, elementwise_affine=True), - TransposeLast(), - ), - nn.GELU(), - ) - elif is_group_norm: - return nn.Sequential( - make_conv(), - nn.Dropout(p=dropout), - Fp32GroupNorm(dim, dim, affine=True), - nn.GELU(), - ) - else: - return nn.Sequential(make_conv(), nn.Dropout(p=dropout), nn.GELU()) - - in_d = 1 - self.conv_layers = nn.ModuleList() - for i, cl in enumerate(conv_layers): - assert len(cl) == 3, "invalid conv definition: " + str(cl) - (dim, k, stride) = cl - - self.conv_layers.append( - block( - in_d, - dim, - k, - stride, - is_layer_norm=mode == "layer_norm", - is_group_norm=mode == "default" and i == 0, - conv_bias=conv_bias, - ) - ) - in_d = dim - - def forward(self, x): - - # BxT -> BxCxT - x = x.unsqueeze(1) - - for conv in self.conv_layers: - x = conv(x) - - return x + def __init__( + self, + conv_layers: List[Tuple[int, int, int]], + dropout: float = 0.0, + mode: str = "default", + conv_bias: bool = False, + ): + super().__init__() + + assert mode in {"default", "layer_norm"} + + def block( + n_in, + n_out, + k, + stride, + is_layer_norm=False, + is_group_norm=False, + conv_bias=False, + ): + def make_conv(): + conv = nn.Conv1d(n_in, n_out, k, stride=stride, bias=conv_bias) + nn.init.kaiming_normal_(conv.weight) + return conv + + assert ( + is_layer_norm and is_group_norm + ) == False, "layer norm and group norm are exclusive" + + if is_layer_norm: + return nn.Sequential( + make_conv(), + nn.Dropout(p=dropout), + nn.Sequential( + TransposeLast(), + Fp32LayerNorm(dim, elementwise_affine=True), + TransposeLast(), + ), + nn.GELU(), + ) + elif is_group_norm: + return nn.Sequential( + make_conv(), + nn.Dropout(p=dropout), + Fp32GroupNorm(dim, dim, affine=True), + nn.GELU(), + ) + else: + return nn.Sequential(make_conv(), nn.Dropout(p=dropout), nn.GELU()) + + in_d = 1 + self.conv_layers = nn.ModuleList() + for i, cl in enumerate(conv_layers): + assert len(cl) == 3, "invalid conv definition: " + str(cl) + (dim, k, stride) = cl + + self.conv_layers.append( + block( + in_d, + dim, + k, + stride, + is_layer_norm=mode == "layer_norm", + is_group_norm=mode == "default" and i == 0, + conv_bias=conv_bias, + ) + ) + in_d = dim + + def forward(self, x): + + # BxT -> BxCxT + x = x.unsqueeze(1) + + for conv in self.conv_layers: + x = conv(x) + + return x + def compute_mask_indices( shape: Tuple[int, int], @@ -254,11 +258,7 @@ def compute_mask_indices( mask_idc = np.random.choice(sz - min_len, num_mask, replace=False) mask_idc = np.asarray( - [ - mask_idc[j] + offset - for j in range(len(mask_idc)) - for offset in range(lengths[j]) - ] + [mask_idc[j] + offset for j in range(len(mask_idc)) for offset in range(lengths[j])] ) mask_idcs.append(np.unique(mask_idc[mask_idc < sz])) @@ -269,9 +269,7 @@ def compute_mask_indices( mask_idc = np.random.choice(mask_idc, min_len, replace=False) if mask_dropout > 0: num_holes = np.rint(len(mask_idc) * mask_dropout).astype(int) - mask_idc = np.random.choice( - mask_idc, len(mask_idc) - num_holes, replace=False - ) + mask_idc = np.random.choice(mask_idc, len(mask_idc) - num_holes, replace=False) mask[i, mask_idc] = True @@ -288,8 +286,8 @@ class GradMultiply(torch.autograd.Function): @staticmethod def backward(ctx, grad): return grad * ctx.scale, None - - + + def is_xla_tensor(tensor): return torch.is_tensor(tensor) and tensor.device.type == "xla" @@ -303,4 +301,4 @@ def index_put(tensor, indices, value): tensor = torch.mul(tensor, ~indices) + torch.mul(value, indices) else: tensor[indices] = value - return tensor \ No newline at end of file + return tensor diff --git a/funasr/models/emotion2vec/model.py b/funasr/models/emotion2vec/model.py index 58b9f39ff..48b87160f 100644 --- a/funasr/models/emotion2vec/model.py +++ b/funasr/models/emotion2vec/model.py @@ -38,6 +38,7 @@ class Emotion2vec(torch.nn.Module): emotion2vec: Self-Supervised Pre-Training for Speech Emotion Representation https://arxiv.org/abs/2312.15185 """ + def __init__(self, **kwargs): super().__init__() # import pdb; pdb.set_trace() @@ -75,7 +76,7 @@ class Emotion2vec(torch.nn.Module): cfg.get("layer_norm_first"), self.alibi_biases, ) - self.modality_encoders['AUDIO'] = enc + self.modality_encoders["AUDIO"] = enc self.ema = None @@ -85,7 +86,9 @@ class Emotion2vec(torch.nn.Module): self.dropout_input = torch.nn.Dropout(cfg.get("dropout_input")) - dpr = np.linspace(cfg.get("start_drop_path_rate"), cfg.get("end_drop_path_rate"), cfg.get("depth")) + dpr = np.linspace( + cfg.get("start_drop_path_rate"), cfg.get("end_drop_path_rate"), cfg.get("depth") + ) self.blocks = torch.nn.ModuleList([make_block(dpr[i]) for i in range(cfg.get("depth"))]) @@ -98,7 +101,6 @@ class Emotion2vec(torch.nn.Module): if vocab_size > 0: self.proj = torch.nn.Linear(cfg.get("embed_dim"), vocab_size) - def forward( self, source, @@ -114,7 +116,7 @@ class Emotion2vec(torch.nn.Module): **kwargs, ): - feature_extractor = self.modality_encoders['AUDIO'] + feature_extractor = self.modality_encoders["AUDIO"] mask_seeds = None @@ -146,11 +148,7 @@ class Emotion2vec(torch.nn.Module): ): ab = masked_alibi_bias if ab is not None and alibi_scale is not None: - scale = ( - alibi_scale[i] - if alibi_scale.size(0) > 1 - else alibi_scale.squeeze(0) - ) + scale = alibi_scale[i] if alibi_scale.size(0) > 1 else alibi_scale.squeeze(0) ab = ab * scale.type_as(ab) x, lr = blk( @@ -192,15 +190,16 @@ class Emotion2vec(torch.nn.Module): ) return res - def inference(self, - data_in, - data_lengths=None, - key: list = None, - tokenizer=None, - frontend=None, - **kwargs, - ): - + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): + # if source_file.endswith('.wav'): # wav, sr = sf.read(source_file) # channel = sf.info(source_file).channels @@ -213,12 +212,17 @@ class Emotion2vec(torch.nn.Module): meta_data = {} # extract fbank feats time1 = time.perf_counter() - audio_sample_list = load_audio_text_image_video(data_in, fs=16000, audio_fs=kwargs.get("fs", 16000), - data_type=kwargs.get("data_type", "sound"), tokenizer=tokenizer) + audio_sample_list = load_audio_text_image_video( + data_in, + fs=16000, + audio_fs=kwargs.get("fs", 16000), + data_type=kwargs.get("data_type", "sound"), + tokenizer=tokenizer, + ) time2 = time.perf_counter() meta_data["load_data"] = f"{time2 - time1:0.3f}" - meta_data["batch_data_time"] = len(audio_sample_list[0])/kwargs.get("fs", 16000) - + meta_data["batch_data_time"] = len(audio_sample_list[0]) / kwargs.get("fs", 16000) + results = [] output_dir = kwargs.get("output_dir") if output_dir: @@ -230,13 +234,13 @@ class Emotion2vec(torch.nn.Module): source = source.view(1, -1) feats = self.extract_features(source, padding_mask=None) - x = feats['x'] - feats = feats['x'].squeeze(0).cpu().numpy() - if granularity == 'frame': + x = feats["x"] + feats = feats["x"].squeeze(0).cpu().numpy() + if granularity == "frame": feats = feats - elif granularity == 'utterance': + elif granularity == "utterance": feats = np.mean(feats, axis=0) - + if output_dir and extract_embedding: np.save(os.path.join(output_dir, "{}.npy".format(key[i])), feats) @@ -253,5 +257,4 @@ class Emotion2vec(torch.nn.Module): result_i["feats"] = feats results.append(result_i) - - return results, meta_data \ No newline at end of file + return results, meta_data diff --git a/funasr/models/emotion2vec/modules.py b/funasr/models/emotion2vec/modules.py index fcf99ebb8..6e3640948 100644 --- a/funasr/models/emotion2vec/modules.py +++ b/funasr/models/emotion2vec/modules.py @@ -77,18 +77,10 @@ class BlockEncoder(nn.Module): x = self.dropout(x) for i, blk in enumerate(self.blocks): - if ( - not self.training - or self.layerdrop == 0 - or (np.random.random() > self.layerdrop) - ): + if not self.training or self.layerdrop == 0 or (np.random.random() > self.layerdrop): ab = alibi_bias if ab is not None and alibi_scale is not None: - scale = ( - alibi_scale[i] - if alibi_scale.size(0) > 1 - else alibi_scale.squeeze(0) - ) + scale = alibi_scale[i] if alibi_scale.size(0) > 1 else alibi_scale.squeeze(0) ab = ab * scale.type_as(ab) x, _ = blk(x, padding_mask, ab) @@ -264,7 +256,7 @@ class AltAttention(nn.Module): super().__init__() self.num_heads = num_heads head_dim = dim // num_heads - self.scale = qk_scale or head_dim ** -0.5 + self.scale = qk_scale or head_dim**-0.5 self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) self.attn_drop = nn.Dropout(attn_drop) diff --git a/funasr/models/emotion2vec/timm_modules.py b/funasr/models/emotion2vec/timm_modules.py index b26da5271..416d2cbaf 100644 --- a/funasr/models/emotion2vec/timm_modules.py +++ b/funasr/models/emotion2vec/timm_modules.py @@ -4,7 +4,7 @@ from itertools import repeat from functools import partial -def drop_path(x, drop_prob: float = 0., training: bool = False, scale_by_keep: bool = True): +def drop_path(x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True): """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). This is the same as the DropConnect impl I created for EfficientNet, etc networks, however, @@ -14,7 +14,7 @@ def drop_path(x, drop_prob: float = 0., training: bool = False, scale_by_keep: b 'survival rate' as the argument. """ - if drop_prob == 0. or not training: + if drop_prob == 0.0 or not training: return x keep_prob = 1 - drop_prob shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets @@ -23,10 +23,11 @@ def drop_path(x, drop_prob: float = 0., training: bool = False, scale_by_keep: b random_tensor.div_(keep_prob) return x * random_tensor + class DropPath(nn.Module): - """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). - """ - def __init__(self, drop_prob: float = 0., scale_by_keep: bool = True): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" + + def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True): super(DropPath, self).__init__() self.drop_prob = drop_prob self.scale_by_keep = scale_by_keep @@ -35,10 +36,7 @@ class DropPath(nn.Module): return drop_path(x, self.drop_prob, self.training, self.scale_by_keep) def extra_repr(self): - return f'drop_prob={round(self.drop_prob,3):0.3f}' - - - + return f"drop_prob={round(self.drop_prob,3):0.3f}" # From PyTorch internals @@ -47,6 +45,7 @@ def _ntuple(n): if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): return tuple(x) return tuple(repeat(x, n)) + return parse @@ -56,19 +55,20 @@ to_3tuple = _ntuple(3) to_4tuple = _ntuple(4) to_ntuple = _ntuple + class Mlp(nn.Module): - """ MLP as used in Vision Transformer, MLP-Mixer and related networks - """ + """MLP as used in Vision Transformer, MLP-Mixer and related networks""" + def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - act_layer=nn.GELU, - norm_layer=None, - bias=True, - drop=0., - use_conv=False, + self, + in_features, + hidden_features=None, + out_features=None, + act_layer=nn.GELU, + norm_layer=None, + bias=True, + drop=0.0, + use_conv=False, ): super().__init__() out_features = out_features or in_features @@ -92,5 +92,3 @@ class Mlp(nn.Module): x = self.fc2(x) x = self.drop2(x) return x - - diff --git a/funasr/models/eres2net/eres2net.py b/funasr/models/eres2net/eres2net.py index 3ea9fdfba..de6c2bbdd 100644 --- a/funasr/models/eres2net/eres2net.py +++ b/funasr/models/eres2net/eres2net.py @@ -26,21 +26,18 @@ class ReLU(nn.Hardtanh): super(ReLU, self).__init__(0, 20, inplace) def __repr__(self): - inplace_str = 'inplace' if self.inplace else '' - return self.__class__.__name__ + ' (' \ - + inplace_str + ')' + inplace_str = "inplace" if self.inplace else "" + return self.__class__.__name__ + " (" + inplace_str + ")" def conv1x1(in_planes, out_planes, stride=1): "1x1 convolution without padding" - return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, - padding=0, bias=False) + return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, padding=0, bias=False) def conv3x3(in_planes, out_planes, stride=1): "3x3 convolution with padding" - return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, - padding=1, bias=False) + return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False) class BasicBlockERes2Net(nn.Module): @@ -67,12 +64,11 @@ class BasicBlockERes2Net(nn.Module): self.shortcut = nn.Sequential() if stride != 1 or in_planes != self.expansion * planes: self.shortcut = nn.Sequential( - nn.Conv2d(in_planes, - self.expansion * planes, - kernel_size=1, - stride=stride, - bias=False), - nn.BatchNorm2d(self.expansion * planes)) + nn.Conv2d( + in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False + ), + nn.BatchNorm2d(self.expansion * planes), + ) self.stride = stride self.width = width self.scale = scale @@ -135,12 +131,11 @@ class BasicBlockERes2Net_diff_AFF(nn.Module): self.shortcut = nn.Sequential() if stride != 1 or in_planes != self.expansion * planes: self.shortcut = nn.Sequential( - nn.Conv2d(in_planes, - self.expansion * planes, - kernel_size=1, - stride=stride, - bias=False), - nn.BatchNorm2d(self.expansion * planes)) + nn.Conv2d( + in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False + ), + nn.BatchNorm2d(self.expansion * planes), + ) self.stride = stride self.width = width self.scale = scale @@ -176,15 +171,17 @@ class BasicBlockERes2Net_diff_AFF(nn.Module): class ERes2Net(nn.Module): - def __init__(self, - block=BasicBlockERes2Net, - block_fuse=BasicBlockERes2Net_diff_AFF, - num_blocks=[3, 4, 6, 3], - m_channels=32, - feat_dim=80, - embedding_size=192, - pooling_func='TSTP', - two_emb_layer=False): + def __init__( + self, + block=BasicBlockERes2Net, + block_fuse=BasicBlockERes2Net_diff_AFF, + num_blocks=[3, 4, 6, 3], + m_channels=32, + feat_dim=80, + embedding_size=192, + pooling_func="TSTP", + two_emb_layer=False, + ): super(ERes2Net, self).__init__() self.in_planes = m_channels self.feat_dim = feat_dim @@ -192,48 +189,32 @@ class ERes2Net(nn.Module): self.stats_dim = int(feat_dim / 8) * m_channels * 8 self.two_emb_layer = two_emb_layer - self.conv1 = nn.Conv2d(1, - m_channels, - kernel_size=3, - stride=1, - padding=1, - bias=False) + self.conv1 = nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False) self.bn1 = nn.BatchNorm2d(m_channels) - self.layer1 = self._make_layer(block, - m_channels, - num_blocks[0], - stride=1) - self.layer2 = self._make_layer(block, - m_channels * 2, - num_blocks[1], - stride=2) - self.layer3 = self._make_layer(block_fuse, - m_channels * 4, - num_blocks[2], - stride=2) - self.layer4 = self._make_layer(block_fuse, - m_channels * 8, - num_blocks[3], - stride=2) + self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=1) + self.layer2 = self._make_layer(block, m_channels * 2, num_blocks[1], stride=2) + self.layer3 = self._make_layer(block_fuse, m_channels * 4, num_blocks[2], stride=2) + self.layer4 = self._make_layer(block_fuse, m_channels * 8, num_blocks[3], stride=2) # Downsampling module for each layer - self.layer1_downsample = nn.Conv2d(m_channels * 2, m_channels * 4, kernel_size=3, stride=2, padding=1, - bias=False) - self.layer2_downsample = nn.Conv2d(m_channels * 4, m_channels * 8, kernel_size=3, padding=1, stride=2, - bias=False) - self.layer3_downsample = nn.Conv2d(m_channels * 8, m_channels * 16, kernel_size=3, padding=1, stride=2, - bias=False) + self.layer1_downsample = nn.Conv2d( + m_channels * 2, m_channels * 4, kernel_size=3, stride=2, padding=1, bias=False + ) + self.layer2_downsample = nn.Conv2d( + m_channels * 4, m_channels * 8, kernel_size=3, padding=1, stride=2, bias=False + ) + self.layer3_downsample = nn.Conv2d( + m_channels * 8, m_channels * 16, kernel_size=3, padding=1, stride=2, bias=False + ) # Bottom-up fusion module self.fuse_mode12 = AFF(channels=m_channels * 4) self.fuse_mode123 = AFF(channels=m_channels * 8) self.fuse_mode1234 = AFF(channels=m_channels * 16) - self.n_stats = 1 if pooling_func == 'TAP' or pooling_func == "TSDP" else 2 - self.pool = getattr(pooling_layers, pooling_func)( - in_dim=self.stats_dim * block.expansion) - self.seg_1 = nn.Linear(self.stats_dim * block.expansion * self.n_stats, - embedding_size) + self.n_stats = 1 if pooling_func == "TAP" or pooling_func == "TSDP" else 2 + self.pool = getattr(pooling_layers, pooling_func)(in_dim=self.stats_dim * block.expansion) + self.seg_1 = nn.Linear(self.stats_dim * block.expansion * self.n_stats, embedding_size) if self.two_emb_layer: self.seg_bn_1 = nn.BatchNorm1d(embedding_size, affine=False) self.seg_2 = nn.Linear(embedding_size, embedding_size) @@ -298,12 +279,11 @@ class BasicBlockRes2Net(nn.Module): self.shortcut = nn.Sequential() if stride != 1 or in_planes != self.expansion * planes: self.shortcut = nn.Sequential( - nn.Conv2d(in_planes, - self.expansion * planes, - kernel_size=1, - stride=stride, - bias=False), - nn.BatchNorm2d(self.expansion * planes)) + nn.Conv2d( + in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False + ), + nn.BatchNorm2d(self.expansion * planes), + ) self.stride = stride self.width = width self.scale = scale @@ -340,14 +320,16 @@ class BasicBlockRes2Net(nn.Module): class Res2Net(nn.Module): - def __init__(self, - block=BasicBlockRes2Net, - num_blocks=[3, 4, 6, 3], - m_channels=32, - feat_dim=80, - embedding_size=192, - pooling_func='TSTP', - two_emb_layer=False): + def __init__( + self, + block=BasicBlockRes2Net, + num_blocks=[3, 4, 6, 3], + m_channels=32, + feat_dim=80, + embedding_size=192, + pooling_func="TSTP", + two_emb_layer=False, + ): super(Res2Net, self).__init__() self.in_planes = m_channels self.feat_dim = feat_dim @@ -355,35 +337,16 @@ class Res2Net(nn.Module): self.stats_dim = int(feat_dim / 8) * m_channels * 8 self.two_emb_layer = two_emb_layer - self.conv1 = nn.Conv2d(1, - m_channels, - kernel_size=3, - stride=1, - padding=1, - bias=False) + self.conv1 = nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False) self.bn1 = nn.BatchNorm2d(m_channels) - self.layer1 = self._make_layer(block, - m_channels, - num_blocks[0], - stride=1) - self.layer2 = self._make_layer(block, - m_channels * 2, - num_blocks[1], - stride=2) - self.layer3 = self._make_layer(block, - m_channels * 4, - num_blocks[2], - stride=2) - self.layer4 = self._make_layer(block, - m_channels * 8, - num_blocks[3], - stride=2) + self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=1) + self.layer2 = self._make_layer(block, m_channels * 2, num_blocks[1], stride=2) + self.layer3 = self._make_layer(block, m_channels * 4, num_blocks[2], stride=2) + self.layer4 = self._make_layer(block, m_channels * 8, num_blocks[3], stride=2) - self.n_stats = 1 if pooling_func == 'TAP' or pooling_func == "TSDP" else 2 - self.pool = getattr(pooling_layers, pooling_func)( - in_dim=self.stats_dim * block.expansion) - self.seg_1 = nn.Linear(self.stats_dim * block.expansion * self.n_stats, - embedding_size) + self.n_stats = 1 if pooling_func == "TAP" or pooling_func == "TSDP" else 2 + self.pool = getattr(pooling_layers, pooling_func)(in_dim=self.stats_dim * block.expansion) + self.seg_1 = nn.Linear(self.stats_dim * block.expansion * self.n_stats, embedding_size) if self.two_emb_layer: self.seg_bn_1 = nn.BatchNorm1d(embedding_size, affine=False) self.seg_2 = nn.Linear(embedding_size, embedding_size) diff --git a/funasr/models/eres2net/eres2net_aug.py b/funasr/models/eres2net/eres2net_aug.py index 8e4fc211d..f700b9c98 100644 --- a/funasr/models/eres2net/eres2net_aug.py +++ b/funasr/models/eres2net/eres2net_aug.py @@ -26,21 +26,18 @@ class ReLU(nn.Hardtanh): super(ReLU, self).__init__(0, 20, inplace) def __repr__(self): - inplace_str = 'inplace' if self.inplace else '' - return self.__class__.__name__ + ' (' \ - + inplace_str + ')' + inplace_str = "inplace" if self.inplace else "" + return self.__class__.__name__ + " (" + inplace_str + ")" def conv1x1(in_planes, out_planes, stride=1): "1x1 convolution without padding" - return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, - padding=0, bias=False) + return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, padding=0, bias=False) def conv3x3(in_planes, out_planes, stride=1): "3x3 convolution with padding" - return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, - padding=1, bias=False) + return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False) class BasicBlockERes2Net(nn.Module): @@ -67,12 +64,11 @@ class BasicBlockERes2Net(nn.Module): self.shortcut = nn.Sequential() if stride != 1 or in_planes != self.expansion * planes: self.shortcut = nn.Sequential( - nn.Conv2d(in_planes, - self.expansion * planes, - kernel_size=1, - stride=stride, - bias=False), - nn.BatchNorm2d(self.expansion * planes)) + nn.Conv2d( + in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False + ), + nn.BatchNorm2d(self.expansion * planes), + ) self.stride = stride self.width = width self.scale = scale @@ -136,12 +132,11 @@ class BasicBlockERes2Net_diff_AFF(nn.Module): self.shortcut = nn.Sequential() if stride != 1 or in_planes != self.expansion * planes: self.shortcut = nn.Sequential( - nn.Conv2d(in_planes, - self.expansion * planes, - kernel_size=1, - stride=stride, - bias=False), - nn.BatchNorm2d(self.expansion * planes)) + nn.Conv2d( + in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False + ), + nn.BatchNorm2d(self.expansion * planes), + ) self.stride = stride self.width = width self.scale = scale @@ -177,15 +172,17 @@ class BasicBlockERes2Net_diff_AFF(nn.Module): class ERes2NetAug(nn.Module): - def __init__(self, - block=BasicBlockERes2Net, - block_fuse=BasicBlockERes2Net_diff_AFF, - num_blocks=[3, 4, 6, 3], - m_channels=64, - feat_dim=80, - embedding_size=192, - pooling_func='TSTP', - two_emb_layer=False): + def __init__( + self, + block=BasicBlockERes2Net, + block_fuse=BasicBlockERes2Net_diff_AFF, + num_blocks=[3, 4, 6, 3], + m_channels=64, + feat_dim=80, + embedding_size=192, + pooling_func="TSTP", + two_emb_layer=False, + ): super(ERes2NetAug, self).__init__() self.in_planes = m_channels self.feat_dim = feat_dim @@ -193,45 +190,29 @@ class ERes2NetAug(nn.Module): self.stats_dim = int(feat_dim / 8) * m_channels * 8 self.two_emb_layer = two_emb_layer - self.conv1 = nn.Conv2d(1, - m_channels, - kernel_size=3, - stride=1, - padding=1, - bias=False) + self.conv1 = nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False) self.bn1 = nn.BatchNorm2d(m_channels) - self.layer1 = self._make_layer(block, - m_channels, - num_blocks[0], - stride=1) - self.layer2 = self._make_layer(block, - m_channels * 2, - num_blocks[1], - stride=2) - self.layer3 = self._make_layer(block_fuse, - m_channels * 4, - num_blocks[2], - stride=2) - self.layer4 = self._make_layer(block_fuse, - m_channels * 8, - num_blocks[3], - stride=2) + self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=1) + self.layer2 = self._make_layer(block, m_channels * 2, num_blocks[1], stride=2) + self.layer3 = self._make_layer(block_fuse, m_channels * 4, num_blocks[2], stride=2) + self.layer4 = self._make_layer(block_fuse, m_channels * 8, num_blocks[3], stride=2) - self.layer1_downsample = nn.Conv2d(m_channels * 4, m_channels * 8, kernel_size=3, padding=1, stride=2, - bias=False) - self.layer2_downsample = nn.Conv2d(m_channels * 8, m_channels * 16, kernel_size=3, padding=1, stride=2, - bias=False) - self.layer3_downsample = nn.Conv2d(m_channels * 16, m_channels * 32, kernel_size=3, padding=1, stride=2, - bias=False) + self.layer1_downsample = nn.Conv2d( + m_channels * 4, m_channels * 8, kernel_size=3, padding=1, stride=2, bias=False + ) + self.layer2_downsample = nn.Conv2d( + m_channels * 8, m_channels * 16, kernel_size=3, padding=1, stride=2, bias=False + ) + self.layer3_downsample = nn.Conv2d( + m_channels * 16, m_channels * 32, kernel_size=3, padding=1, stride=2, bias=False + ) self.fuse_mode12 = AFF(channels=m_channels * 8) self.fuse_mode123 = AFF(channels=m_channels * 16) self.fuse_mode1234 = AFF(channels=m_channels * 32) - self.n_stats = 1 if pooling_func == 'TAP' or pooling_func == "TSDP" else 2 - self.pool = getattr(pooling_layers, pooling_func)( - in_dim=self.stats_dim * block.expansion) - self.seg_1 = nn.Linear(self.stats_dim * block.expansion * self.n_stats, - embedding_size) + self.n_stats = 1 if pooling_func == "TAP" or pooling_func == "TSDP" else 2 + self.pool = getattr(pooling_layers, pooling_func)(in_dim=self.stats_dim * block.expansion) + self.seg_1 = nn.Linear(self.stats_dim * block.expansion * self.n_stats, embedding_size) if self.two_emb_layer: self.seg_bn_1 = nn.BatchNorm1d(embedding_size, affine=False) self.seg_2 = nn.Linear(embedding_size, embedding_size) diff --git a/funasr/models/eres2net/fusion.py b/funasr/models/eres2net/fusion.py index 2aff7a721..fbe699efb 100644 --- a/funasr/models/eres2net/fusion.py +++ b/funasr/models/eres2net/fusion.py @@ -23,7 +23,6 @@ class AFF(nn.Module): xa = torch.cat((x, ds_y), dim=1) x_att = self.local_att(xa) x_att = 1.0 + torch.tanh(x_att) - xo = torch.mul(x, x_att) + torch.mul(ds_y, 2.0-x_att) + xo = torch.mul(x, x_att) + torch.mul(ds_y, 2.0 - x_att) return xo - diff --git a/funasr/models/fsmn_vad_streaming/encoder.py b/funasr/models/fsmn_vad_streaming/encoder.py index bc51a6fcd..6668c5d51 100755 --- a/funasr/models/fsmn_vad_streaming/encoder.py +++ b/funasr/models/fsmn_vad_streaming/encoder.py @@ -9,6 +9,7 @@ import torch.nn.functional as F from funasr.register import tables + class LinearTransform(nn.Module): def __init__(self, input_dim, output_dim): @@ -53,13 +54,13 @@ class RectifiedLinear(nn.Module): class FSMNBlock(nn.Module): def __init__( - self, - input_dim: int, - output_dim: int, - lorder=None, - rorder=None, - lstride=1, - rstride=1, + self, + input_dim: int, + output_dim: int, + lorder=None, + rorder=None, + lstride=1, + rstride=1, ): super(FSMNBlock, self).__init__() @@ -74,28 +75,30 @@ class FSMNBlock(nn.Module): self.rstride = rstride self.conv_left = nn.Conv2d( - self.dim, self.dim, [lorder, 1], dilation=[lstride, 1], groups=self.dim, bias=False) + self.dim, self.dim, [lorder, 1], dilation=[lstride, 1], groups=self.dim, bias=False + ) if self.rorder > 0: self.conv_right = nn.Conv2d( - self.dim, self.dim, [rorder, 1], dilation=[rstride, 1], groups=self.dim, bias=False) + self.dim, self.dim, [rorder, 1], dilation=[rstride, 1], groups=self.dim, bias=False + ) else: self.conv_right = None def forward(self, input: torch.Tensor, cache: torch.Tensor): x = torch.unsqueeze(input, 1) x_per = x.permute(0, 3, 2, 1) # B D T C - + cache = cache.to(x_per.device) y_left = torch.cat((cache, x_per), dim=2) - cache = y_left[:, :, -(self.lorder - 1) * self.lstride:, :] + cache = y_left[:, :, -(self.lorder - 1) * self.lstride :, :] y_left = self.conv_left(y_left) out = x_per + y_left if self.conv_right is not None: # maybe need to check y_right = F.pad(x_per, [0, 0, 0, self.rorder * self.rstride]) - y_right = y_right[:, :, self.rstride:, :] + y_right = y_right[:, :, self.rstride :, :] y_right = self.conv_right(y_right) out += y_right @@ -106,15 +109,16 @@ class FSMNBlock(nn.Module): class BasicBlock(nn.Module): - def __init__(self, - linear_dim: int, - proj_dim: int, - lorder: int, - rorder: int, - lstride: int, - rstride: int, - stack_layer: int - ): + def __init__( + self, + linear_dim: int, + proj_dim: int, + lorder: int, + rorder: int, + lstride: int, + rstride: int, + stack_layer: int, + ): super(BasicBlock, self).__init__() self.lorder = lorder self.rorder = rorder @@ -128,17 +132,22 @@ class BasicBlock(nn.Module): def forward(self, input: torch.Tensor, cache: Dict[str, torch.Tensor]): x1 = self.linear(input) # B T D - cache_layer_name = 'cache_layer_{}'.format(self.stack_layer) + cache_layer_name = "cache_layer_{}".format(self.stack_layer) if cache_layer_name not in cache: - cache[cache_layer_name] = torch.zeros(x1.shape[0], x1.shape[-1], (self.lorder - 1) * self.lstride, 1) + cache[cache_layer_name] = torch.zeros( + x1.shape[0], x1.shape[-1], (self.lorder - 1) * self.lstride, 1 + ) x2, cache[cache_layer_name] = self.fsmn_block(x1, cache[cache_layer_name]) x3 = self.affine(x2) x4 = self.relu(x3) return x4 + + class BasicBlock_export(nn.Module): - def __init__(self, - model, - ): + def __init__( + self, + model, + ): super(BasicBlock_export, self).__init__() self.linear = model.linear self.fsmn_block = model.fsmn_block @@ -167,7 +176,7 @@ class FsmnStack(nn.Sequential): return x -''' +""" FSMN net for keyword spotting input_dim: input dimension linear_dim: fsmn input dimensionll @@ -176,23 +185,24 @@ lorder: fsmn left order rorder: fsmn right order num_syn: output dimension fsmn_layers: no. of sequential fsmn layers -''' +""" + @tables.register("encoder_classes", "FSMN") class FSMN(nn.Module): def __init__( - self, - input_dim: int, - input_affine_dim: int, - fsmn_layers: int, - linear_dim: int, - proj_dim: int, - lorder: int, - rorder: int, - lstride: int, - rstride: int, - output_affine_dim: int, - output_dim: int + self, + input_dim: int, + input_affine_dim: int, + fsmn_layers: int, + linear_dim: int, + proj_dim: int, + lorder: int, + rorder: int, + lstride: int, + rstride: int, + output_affine_dim: int, + output_dim: int, ): super().__init__() @@ -207,19 +217,21 @@ class FSMN(nn.Module): self.in_linear1 = AffineTransform(input_dim, input_affine_dim) self.in_linear2 = AffineTransform(input_affine_dim, linear_dim) self.relu = RectifiedLinear(linear_dim, linear_dim) - self.fsmn = FsmnStack(*[BasicBlock(linear_dim, proj_dim, lorder, rorder, lstride, rstride, i) for i in - range(fsmn_layers)]) + self.fsmn = FsmnStack( + *[ + BasicBlock(linear_dim, proj_dim, lorder, rorder, lstride, rstride, i) + for i in range(fsmn_layers) + ] + ) self.out_linear1 = AffineTransform(linear_dim, output_affine_dim) self.out_linear2 = AffineTransform(output_affine_dim, output_dim) self.softmax = nn.Softmax(dim=-1) - + def fuse_modules(self): pass def forward( - self, - input: torch.Tensor, - cache: Dict[str, torch.Tensor] + self, input: torch.Tensor, cache: Dict[str, torch.Tensor] ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: """ Args: @@ -242,10 +254,12 @@ class FSMN(nn.Module): @tables.register("encoder_classes", "FSMNExport") class FSMNExport(nn.Module): def __init__( - self, model, **kwargs, + self, + model, + **kwargs, ): super().__init__() - + # self.input_dim = input_dim # self.input_affine_dim = input_affine_dim # self.fsmn_layers = fsmn_layers @@ -273,10 +287,10 @@ class FSMNExport(nn.Module): for i, d in enumerate(model.fsmn): if isinstance(d, BasicBlock): self.fsmn[i] = BasicBlock_export(d) - + def fuse_modules(self): pass - + def forward( self, input: torch.Tensor, @@ -288,7 +302,7 @@ class FSMNExport(nn.Module): in_cache: when in_cache is not None, the forward is in streaming. The type of in_cache is a dict, egs, {'cache_layer_1': torch.Tensor(B, T1, D)}, T1 is equal to self.lorder. It is {} for the 1st frame """ - + x = self.in_linear1(input) x = self.in_linear2(x) x = self.relu(x) @@ -301,11 +315,11 @@ class FSMNExport(nn.Module): x = self.out_linear1(x) x = self.out_linear2(x) x = self.softmax(x) - + return x, out_caches -''' +""" one deep fsmn layer dimproj: projection dimension, input and output dimension of memory blocks dimlinear: dimension of mapping layer @@ -313,7 +327,8 @@ lorder: left order rorder: right order lstride: left stride rstride: right stride -''' +""" + @tables.register("encoder_classes", "DFSMN") class DFSMN(nn.Module): @@ -330,11 +345,13 @@ class DFSMN(nn.Module): self.shrink = LinearTransform(dimlinear, dimproj) self.conv_left = nn.Conv2d( - dimproj, dimproj, [lorder, 1], dilation=[lstride, 1], groups=dimproj, bias=False) + dimproj, dimproj, [lorder, 1], dilation=[lstride, 1], groups=dimproj, bias=False + ) if rorder > 0: self.conv_right = nn.Conv2d( - dimproj, dimproj, [rorder, 1], dilation=[rstride, 1], groups=dimproj, bias=False) + dimproj, dimproj, [rorder, 1], dilation=[rstride, 1], groups=dimproj, bias=False + ) else: self.conv_right = None @@ -349,7 +366,7 @@ class DFSMN(nn.Module): if self.conv_right is not None: y_right = F.pad(x_per, [0, 0, 0, (self.rorder) * self.rstride]) - y_right = y_right[:, :, self.rstride:, :] + y_right = y_right[:, :, self.rstride :, :] out = x_per + self.conv_left(y_left) + self.conv_right(y_right) else: out = x_per + self.conv_left(y_left) @@ -360,30 +377,28 @@ class DFSMN(nn.Module): return output -''' +""" build stacked dfsmn layers -''' +""" def buildDFSMNRepeats(linear_dim=128, proj_dim=64, lorder=20, rorder=1, fsmn_layers=6): repeats = [ - nn.Sequential( - DFSMN(proj_dim, linear_dim, lorder, rorder, 1, 1)) - for i in range(fsmn_layers) + nn.Sequential(DFSMN(proj_dim, linear_dim, lorder, rorder, 1, 1)) for i in range(fsmn_layers) ] return nn.Sequential(*repeats) -if __name__ == '__main__': +if __name__ == "__main__": fsmn = FSMN(400, 140, 4, 250, 128, 10, 2, 1, 1, 140, 2599) print(fsmn) num_params = sum(p.numel() for p in fsmn.parameters()) - print('the number of model params: {}'.format(num_params)) + print("the number of model params: {}".format(num_params)) x = torch.zeros(128, 200, 400) # batch-size * time * dim y, _ = fsmn(x) # batch-size * time * dim - print('input shape: {}'.format(x.shape)) - print('output shape: {}'.format(y.shape)) + print("input shape: {}".format(x.shape)) + print("output shape: {}".format(y.shape)) print(fsmn.to_kaldi_net()) diff --git a/funasr/models/fsmn_vad_streaming/export_meta.py b/funasr/models/fsmn_vad_streaming/export_meta.py index 71830265c..f76016931 100644 --- a/funasr/models/fsmn_vad_streaming/export_meta.py +++ b/funasr/models/fsmn_vad_streaming/export_meta.py @@ -12,7 +12,7 @@ def export_rebuild_model(model, **kwargs): is_onnx = kwargs.get("type", "onnx") == "onnx" encoder_class = tables.encoder_classes.get(kwargs["encoder"] + "Export") model.encoder = encoder_class(model.encoder, onnx=is_onnx) - + model.forward = types.MethodType(export_forward, model) model.export_dummy_inputs = types.MethodType(export_dummy_inputs, model) model.export_input_names = types.MethodType(export_input_names, model) @@ -22,38 +22,44 @@ def export_rebuild_model(model, **kwargs): return model + def export_forward(self, feats: torch.Tensor, *args, **kwargs): - + scores, out_caches = self.encoder(feats, *args) - + return scores, out_caches + def export_dummy_inputs(self, data_in=None, frame=30): if data_in is None: speech = torch.randn(1, frame, self.encoder_conf.get("input_dim")) else: - speech = None # Undo - + speech = None # Undo + cache_frames = self.encoder_conf.get("lorder") + self.encoder_conf.get("rorder") - 1 in_cache0 = torch.randn(1, self.encoder_conf.get("proj_dim"), cache_frames, 1) in_cache1 = torch.randn(1, self.encoder_conf.get("proj_dim"), cache_frames, 1) in_cache2 = torch.randn(1, self.encoder_conf.get("proj_dim"), cache_frames, 1) in_cache3 = torch.randn(1, self.encoder_conf.get("proj_dim"), cache_frames, 1) - + return (speech, in_cache0, in_cache1, in_cache2, in_cache3) + def export_input_names(self): - return ['speech', 'in_cache0', 'in_cache1', 'in_cache2', 'in_cache3'] + return ["speech", "in_cache0", "in_cache1", "in_cache2", "in_cache3"] + def export_output_names(self): - return ['logits', 'out_cache0', 'out_cache1', 'out_cache2', 'out_cache3'] + return ["logits", "out_cache0", "out_cache1", "out_cache2", "out_cache3"] + def export_dynamic_axes(self): return { - 'speech': { - 1: 'feats_length' - }, + "speech": {1: "feats_length"}, } -def export_name(self, ): + +def export_name( + self, +): return "model.onnx" diff --git a/funasr/models/fsmn_vad_streaming/model.py b/funasr/models/fsmn_vad_streaming/model.py index f1d416104..04689bed9 100644 --- a/funasr/models/fsmn_vad_streaming/model.py +++ b/funasr/models/fsmn_vad_streaming/model.py @@ -19,773 +19,898 @@ from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank class VadStateMachine(Enum): - kVadInStateStartPointNotDetected = 1 - kVadInStateInSpeechSegment = 2 - kVadInStateEndPointDetected = 3 + kVadInStateStartPointNotDetected = 1 + kVadInStateInSpeechSegment = 2 + kVadInStateEndPointDetected = 3 class FrameState(Enum): - kFrameStateInvalid = -1 - kFrameStateSpeech = 1 - kFrameStateSil = 0 + kFrameStateInvalid = -1 + kFrameStateSpeech = 1 + kFrameStateSil = 0 # final voice/unvoice state per frame class AudioChangeState(Enum): - kChangeStateSpeech2Speech = 0 - kChangeStateSpeech2Sil = 1 - kChangeStateSil2Sil = 2 - kChangeStateSil2Speech = 3 - kChangeStateNoBegin = 4 - kChangeStateInvalid = 5 + kChangeStateSpeech2Speech = 0 + kChangeStateSpeech2Sil = 1 + kChangeStateSil2Sil = 2 + kChangeStateSil2Speech = 3 + kChangeStateNoBegin = 4 + kChangeStateInvalid = 5 class VadDetectMode(Enum): - kVadSingleUtteranceDetectMode = 0 - kVadMutipleUtteranceDetectMode = 1 + kVadSingleUtteranceDetectMode = 0 + kVadMutipleUtteranceDetectMode = 1 class VADXOptions: - """ - Author: Speech Lab of DAMO Academy, Alibaba Group - Deep-FSMN for Large Vocabulary Continuous Speech Recognition - https://arxiv.org/abs/1803.05030 - """ - - def __init__( - self, - sample_rate: int = 16000, - detect_mode: int = VadDetectMode.kVadMutipleUtteranceDetectMode.value, - snr_mode: int = 0, - max_end_silence_time: int = 800, - max_start_silence_time: int = 3000, - do_start_point_detection: bool = True, - do_end_point_detection: bool = True, - window_size_ms: int = 200, - sil_to_speech_time_thres: int = 150, - speech_to_sil_time_thres: int = 150, - speech_2_noise_ratio: float = 1.0, - do_extend: int = 1, - lookback_time_start_point: int = 200, - lookahead_time_end_point: int = 100, - max_single_segment_time: int = 60000, - nn_eval_block_size: int = 8, - dcd_block_size: int = 4, - snr_thres: int = -100.0, - noise_frame_num_used_for_snr: int = 100, - decibel_thres: int = -100.0, - speech_noise_thres: float = 0.6, - fe_prior_thres: float = 1e-4, - silence_pdf_num: int = 1, - sil_pdf_ids: List[int] = [0], - speech_noise_thresh_low: float = -0.1, - speech_noise_thresh_high: float = 0.3, - output_frame_probs: bool = False, - frame_in_ms: int = 10, - frame_length_ms: int = 25, - **kwargs, - ): - self.sample_rate = sample_rate - self.detect_mode = detect_mode - self.snr_mode = snr_mode - self.max_end_silence_time = max_end_silence_time - self.max_start_silence_time = max_start_silence_time - self.do_start_point_detection = do_start_point_detection - self.do_end_point_detection = do_end_point_detection - self.window_size_ms = window_size_ms - self.sil_to_speech_time_thres = sil_to_speech_time_thres - self.speech_to_sil_time_thres = speech_to_sil_time_thres - self.speech_2_noise_ratio = speech_2_noise_ratio - self.do_extend = do_extend - self.lookback_time_start_point = lookback_time_start_point - self.lookahead_time_end_point = lookahead_time_end_point - self.max_single_segment_time = max_single_segment_time - self.nn_eval_block_size = nn_eval_block_size - self.dcd_block_size = dcd_block_size - self.snr_thres = snr_thres - self.noise_frame_num_used_for_snr = noise_frame_num_used_for_snr - self.decibel_thres = decibel_thres - self.speech_noise_thres = speech_noise_thres - self.fe_prior_thres = fe_prior_thres - self.silence_pdf_num = silence_pdf_num - self.sil_pdf_ids = sil_pdf_ids - self.speech_noise_thresh_low = speech_noise_thresh_low - self.speech_noise_thresh_high = speech_noise_thresh_high - self.output_frame_probs = output_frame_probs - self.frame_in_ms = frame_in_ms - self.frame_length_ms = frame_length_ms + """ + Author: Speech Lab of DAMO Academy, Alibaba Group + Deep-FSMN for Large Vocabulary Continuous Speech Recognition + https://arxiv.org/abs/1803.05030 + """ + + def __init__( + self, + sample_rate: int = 16000, + detect_mode: int = VadDetectMode.kVadMutipleUtteranceDetectMode.value, + snr_mode: int = 0, + max_end_silence_time: int = 800, + max_start_silence_time: int = 3000, + do_start_point_detection: bool = True, + do_end_point_detection: bool = True, + window_size_ms: int = 200, + sil_to_speech_time_thres: int = 150, + speech_to_sil_time_thres: int = 150, + speech_2_noise_ratio: float = 1.0, + do_extend: int = 1, + lookback_time_start_point: int = 200, + lookahead_time_end_point: int = 100, + max_single_segment_time: int = 60000, + nn_eval_block_size: int = 8, + dcd_block_size: int = 4, + snr_thres: int = -100.0, + noise_frame_num_used_for_snr: int = 100, + decibel_thres: int = -100.0, + speech_noise_thres: float = 0.6, + fe_prior_thres: float = 1e-4, + silence_pdf_num: int = 1, + sil_pdf_ids: List[int] = [0], + speech_noise_thresh_low: float = -0.1, + speech_noise_thresh_high: float = 0.3, + output_frame_probs: bool = False, + frame_in_ms: int = 10, + frame_length_ms: int = 25, + **kwargs, + ): + self.sample_rate = sample_rate + self.detect_mode = detect_mode + self.snr_mode = snr_mode + self.max_end_silence_time = max_end_silence_time + self.max_start_silence_time = max_start_silence_time + self.do_start_point_detection = do_start_point_detection + self.do_end_point_detection = do_end_point_detection + self.window_size_ms = window_size_ms + self.sil_to_speech_time_thres = sil_to_speech_time_thres + self.speech_to_sil_time_thres = speech_to_sil_time_thres + self.speech_2_noise_ratio = speech_2_noise_ratio + self.do_extend = do_extend + self.lookback_time_start_point = lookback_time_start_point + self.lookahead_time_end_point = lookahead_time_end_point + self.max_single_segment_time = max_single_segment_time + self.nn_eval_block_size = nn_eval_block_size + self.dcd_block_size = dcd_block_size + self.snr_thres = snr_thres + self.noise_frame_num_used_for_snr = noise_frame_num_used_for_snr + self.decibel_thres = decibel_thres + self.speech_noise_thres = speech_noise_thres + self.fe_prior_thres = fe_prior_thres + self.silence_pdf_num = silence_pdf_num + self.sil_pdf_ids = sil_pdf_ids + self.speech_noise_thresh_low = speech_noise_thresh_low + self.speech_noise_thresh_high = speech_noise_thresh_high + self.output_frame_probs = output_frame_probs + self.frame_in_ms = frame_in_ms + self.frame_length_ms = frame_length_ms class E2EVadSpeechBufWithDoa(object): - """ - Author: Speech Lab of DAMO Academy, Alibaba Group - Deep-FSMN for Large Vocabulary Continuous Speech Recognition - https://arxiv.org/abs/1803.05030 - """ - - def __init__(self): - self.start_ms = 0 - self.end_ms = 0 - self.buffer = [] - self.contain_seg_start_point = False - self.contain_seg_end_point = False - self.doa = 0 - - def Reset(self): - self.start_ms = 0 - self.end_ms = 0 - self.buffer = [] - self.contain_seg_start_point = False - self.contain_seg_end_point = False - self.doa = 0 + """ + Author: Speech Lab of DAMO Academy, Alibaba Group + Deep-FSMN for Large Vocabulary Continuous Speech Recognition + https://arxiv.org/abs/1803.05030 + """ + + def __init__(self): + self.start_ms = 0 + self.end_ms = 0 + self.buffer = [] + self.contain_seg_start_point = False + self.contain_seg_end_point = False + self.doa = 0 + + def Reset(self): + self.start_ms = 0 + self.end_ms = 0 + self.buffer = [] + self.contain_seg_start_point = False + self.contain_seg_end_point = False + self.doa = 0 class E2EVadFrameProb(object): - """ - Author: Speech Lab of DAMO Academy, Alibaba Group - Deep-FSMN for Large Vocabulary Continuous Speech Recognition - https://arxiv.org/abs/1803.05030 - """ - - def __init__(self): - self.noise_prob = 0.0 - self.speech_prob = 0.0 - self.score = 0.0 - self.frame_id = 0 - self.frm_state = 0 + """ + Author: Speech Lab of DAMO Academy, Alibaba Group + Deep-FSMN for Large Vocabulary Continuous Speech Recognition + https://arxiv.org/abs/1803.05030 + """ + + def __init__(self): + self.noise_prob = 0.0 + self.speech_prob = 0.0 + self.score = 0.0 + self.frame_id = 0 + self.frm_state = 0 class WindowDetector(object): - """ - Author: Speech Lab of DAMO Academy, Alibaba Group - Deep-FSMN for Large Vocabulary Continuous Speech Recognition - https://arxiv.org/abs/1803.05030 - """ - - def __init__(self, window_size_ms: int, - sil_to_speech_time: int, - speech_to_sil_time: int, - frame_size_ms: int): - self.window_size_ms = window_size_ms - self.sil_to_speech_time = sil_to_speech_time - self.speech_to_sil_time = speech_to_sil_time - self.frame_size_ms = frame_size_ms - - self.win_size_frame = int(window_size_ms / frame_size_ms) - self.win_sum = 0 - self.win_state = [0] * self.win_size_frame # 初始化窗 - - self.cur_win_pos = 0 - self.pre_frame_state = FrameState.kFrameStateSil - self.cur_frame_state = FrameState.kFrameStateSil - self.sil_to_speech_frmcnt_thres = int(sil_to_speech_time / frame_size_ms) - self.speech_to_sil_frmcnt_thres = int(speech_to_sil_time / frame_size_ms) - - self.voice_last_frame_count = 0 - self.noise_last_frame_count = 0 - self.hydre_frame_count = 0 - - def Reset(self) -> None: - self.cur_win_pos = 0 - self.win_sum = 0 - self.win_state = [0] * self.win_size_frame - self.pre_frame_state = FrameState.kFrameStateSil - self.cur_frame_state = FrameState.kFrameStateSil - self.voice_last_frame_count = 0 - self.noise_last_frame_count = 0 - self.hydre_frame_count = 0 - - def GetWinSize(self) -> int: - return int(self.win_size_frame) - - def DetectOneFrame(self, frameState: FrameState, frame_count: int, cache: dict = {}) -> AudioChangeState: - cur_frame_state = FrameState.kFrameStateSil - if frameState == FrameState.kFrameStateSpeech: - cur_frame_state = 1 - elif frameState == FrameState.kFrameStateSil: - cur_frame_state = 0 - else: - return AudioChangeState.kChangeStateInvalid - self.win_sum -= self.win_state[self.cur_win_pos] - self.win_sum += cur_frame_state - self.win_state[self.cur_win_pos] = cur_frame_state - self.cur_win_pos = (self.cur_win_pos + 1) % self.win_size_frame - - if self.pre_frame_state == FrameState.kFrameStateSil and self.win_sum >= self.sil_to_speech_frmcnt_thres: - self.pre_frame_state = FrameState.kFrameStateSpeech - return AudioChangeState.kChangeStateSil2Speech - - if self.pre_frame_state == FrameState.kFrameStateSpeech and self.win_sum <= self.speech_to_sil_frmcnt_thres: - self.pre_frame_state = FrameState.kFrameStateSil - return AudioChangeState.kChangeStateSpeech2Sil - - if self.pre_frame_state == FrameState.kFrameStateSil: - return AudioChangeState.kChangeStateSil2Sil - if self.pre_frame_state == FrameState.kFrameStateSpeech: - return AudioChangeState.kChangeStateSpeech2Speech - return AudioChangeState.kChangeStateInvalid - - def FrameSizeMs(self) -> int: - return int(self.frame_size_ms) + """ + Author: Speech Lab of DAMO Academy, Alibaba Group + Deep-FSMN for Large Vocabulary Continuous Speech Recognition + https://arxiv.org/abs/1803.05030 + """ + + def __init__( + self, + window_size_ms: int, + sil_to_speech_time: int, + speech_to_sil_time: int, + frame_size_ms: int, + ): + self.window_size_ms = window_size_ms + self.sil_to_speech_time = sil_to_speech_time + self.speech_to_sil_time = speech_to_sil_time + self.frame_size_ms = frame_size_ms + + self.win_size_frame = int(window_size_ms / frame_size_ms) + self.win_sum = 0 + self.win_state = [0] * self.win_size_frame # 初始化窗 + + self.cur_win_pos = 0 + self.pre_frame_state = FrameState.kFrameStateSil + self.cur_frame_state = FrameState.kFrameStateSil + self.sil_to_speech_frmcnt_thres = int(sil_to_speech_time / frame_size_ms) + self.speech_to_sil_frmcnt_thres = int(speech_to_sil_time / frame_size_ms) + + self.voice_last_frame_count = 0 + self.noise_last_frame_count = 0 + self.hydre_frame_count = 0 + + def Reset(self) -> None: + self.cur_win_pos = 0 + self.win_sum = 0 + self.win_state = [0] * self.win_size_frame + self.pre_frame_state = FrameState.kFrameStateSil + self.cur_frame_state = FrameState.kFrameStateSil + self.voice_last_frame_count = 0 + self.noise_last_frame_count = 0 + self.hydre_frame_count = 0 + + def GetWinSize(self) -> int: + return int(self.win_size_frame) + + def DetectOneFrame( + self, frameState: FrameState, frame_count: int, cache: dict = {} + ) -> AudioChangeState: + cur_frame_state = FrameState.kFrameStateSil + if frameState == FrameState.kFrameStateSpeech: + cur_frame_state = 1 + elif frameState == FrameState.kFrameStateSil: + cur_frame_state = 0 + else: + return AudioChangeState.kChangeStateInvalid + self.win_sum -= self.win_state[self.cur_win_pos] + self.win_sum += cur_frame_state + self.win_state[self.cur_win_pos] = cur_frame_state + self.cur_win_pos = (self.cur_win_pos + 1) % self.win_size_frame + + if ( + self.pre_frame_state == FrameState.kFrameStateSil + and self.win_sum >= self.sil_to_speech_frmcnt_thres + ): + self.pre_frame_state = FrameState.kFrameStateSpeech + return AudioChangeState.kChangeStateSil2Speech + + if ( + self.pre_frame_state == FrameState.kFrameStateSpeech + and self.win_sum <= self.speech_to_sil_frmcnt_thres + ): + self.pre_frame_state = FrameState.kFrameStateSil + return AudioChangeState.kChangeStateSpeech2Sil + + if self.pre_frame_state == FrameState.kFrameStateSil: + return AudioChangeState.kChangeStateSil2Sil + if self.pre_frame_state == FrameState.kFrameStateSpeech: + return AudioChangeState.kChangeStateSpeech2Speech + return AudioChangeState.kChangeStateInvalid + + def FrameSizeMs(self) -> int: + return int(self.frame_size_ms) class Stats(object): - def __init__(self, - sil_pdf_ids, - max_end_sil_frame_cnt_thresh, - speech_noise_thres, - ): - self.data_buf_start_frame = 0 - self.frm_cnt = 0 - self.latest_confirmed_speech_frame = 0 - self.lastest_confirmed_silence_frame = -1 - self.continous_silence_frame_count = 0 - self.vad_state_machine = VadStateMachine.kVadInStateStartPointNotDetected - self.confirmed_start_frame = -1 - self.confirmed_end_frame = -1 - self.number_end_time_detected = 0 - self.sil_frame = 0 - self.sil_pdf_ids = sil_pdf_ids - self.noise_average_decibel = -100.0 - self.pre_end_silence_detected = False - self.next_seg = True - - self.output_data_buf = [] - self.output_data_buf_offset = 0 - self.frame_probs = [] - self.max_end_sil_frame_cnt_thresh = max_end_sil_frame_cnt_thresh - self.speech_noise_thres = speech_noise_thres - self.scores = None - self.max_time_out = False - self.decibel = [] - self.data_buf = None - self.data_buf_all = None - self.waveform = None - self.last_drop_frames = 0 + def __init__( + self, + sil_pdf_ids, + max_end_sil_frame_cnt_thresh, + speech_noise_thres, + ): + self.data_buf_start_frame = 0 + self.frm_cnt = 0 + self.latest_confirmed_speech_frame = 0 + self.lastest_confirmed_silence_frame = -1 + self.continous_silence_frame_count = 0 + self.vad_state_machine = VadStateMachine.kVadInStateStartPointNotDetected + self.confirmed_start_frame = -1 + self.confirmed_end_frame = -1 + self.number_end_time_detected = 0 + self.sil_frame = 0 + self.sil_pdf_ids = sil_pdf_ids + self.noise_average_decibel = -100.0 + self.pre_end_silence_detected = False + self.next_seg = True + + self.output_data_buf = [] + self.output_data_buf_offset = 0 + self.frame_probs = [] + self.max_end_sil_frame_cnt_thresh = max_end_sil_frame_cnt_thresh + self.speech_noise_thres = speech_noise_thres + self.scores = None + self.max_time_out = False + self.decibel = [] + self.data_buf = None + self.data_buf_all = None + self.waveform = None + self.last_drop_frames = 0 @tables.register("model_classes", "FsmnVADStreaming") class FsmnVADStreaming(nn.Module): - """ - Author: Speech Lab of DAMO Academy, Alibaba Group - Deep-FSMN for Large Vocabulary Continuous Speech Recognition - https://arxiv.org/abs/1803.05030 - """ - - def __init__(self, - encoder: str = None, - encoder_conf: Optional[Dict] = None, - vad_post_args: Dict[str, Any] = None, - **kwargs, - ): - super().__init__() - self.vad_opts = VADXOptions(**kwargs) - - encoder_class = tables.encoder_classes.get(encoder) - encoder = encoder_class(**encoder_conf) - self.encoder = encoder - self.encoder_conf = encoder_conf - - def ResetDetection(self, cache: dict = {}): - cache["stats"].continous_silence_frame_count = 0 - cache["stats"].latest_confirmed_speech_frame = 0 - cache["stats"].lastest_confirmed_silence_frame = -1 - cache["stats"].confirmed_start_frame = -1 - cache["stats"].confirmed_end_frame = -1 - cache["stats"].vad_state_machine = VadStateMachine.kVadInStateStartPointNotDetected - cache["windows_detector"].Reset() - cache["stats"].sil_frame = 0 - cache["stats"].frame_probs = [] - - if cache["stats"].output_data_buf: - assert cache["stats"].output_data_buf[-1].contain_seg_end_point == True - drop_frames = int(cache["stats"].output_data_buf[-1].end_ms / self.vad_opts.frame_in_ms) - real_drop_frames = drop_frames - cache["stats"].last_drop_frames - cache["stats"].last_drop_frames = drop_frames - cache["stats"].data_buf_all = cache["stats"].data_buf_all[real_drop_frames * int( - self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000):] - cache["stats"].decibel = cache["stats"].decibel[real_drop_frames:] - cache["stats"].scores = cache["stats"].scores[:, real_drop_frames:, :] - - def ComputeDecibel(self, cache: dict = {}) -> None: - frame_sample_length = int(self.vad_opts.frame_length_ms * self.vad_opts.sample_rate / 1000) - frame_shift_length = int(self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000) - if cache["stats"].data_buf_all is None: - cache["stats"].data_buf_all = cache["stats"].waveform[ - 0] # cache["stats"].data_buf is pointed to cache["stats"].waveform[0] - cache["stats"].data_buf = cache["stats"].data_buf_all - else: - cache["stats"].data_buf_all = torch.cat((cache["stats"].data_buf_all, cache["stats"].waveform[0])) - for offset in range(0, cache["stats"].waveform.shape[1] - frame_sample_length + 1, frame_shift_length): - cache["stats"].decibel.append( - 10 * math.log10((cache["stats"].waveform[0][offset: offset + frame_sample_length]).square().sum() + \ - 0.000001)) - - def ComputeScores(self, feats: torch.Tensor, cache: dict = {}) -> None: - scores = self.encoder(feats, cache=cache["encoder"]).to('cpu') # return B * T * D - assert scores.shape[1] == feats.shape[1], "The shape between feats and scores does not match" - self.vad_opts.nn_eval_block_size = scores.shape[1] - cache["stats"].frm_cnt += scores.shape[1] # count total frames - if cache["stats"].scores is None: - cache["stats"].scores = scores # the first calculation - else: - cache["stats"].scores = torch.cat((cache["stats"].scores, scores), dim=1) - - def PopDataBufTillFrame(self, frame_idx: int, cache: dict = {}) -> None: # need check again - while cache["stats"].data_buf_start_frame < frame_idx: - if len(cache["stats"].data_buf) >= int(self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000): - cache["stats"].data_buf_start_frame += 1 - cache["stats"].data_buf = cache["stats"].data_buf_all[ - (cache["stats"].data_buf_start_frame - cache["stats"].last_drop_frames) * int( - self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000):] - - def PopDataToOutputBuf(self, start_frm: int, frm_cnt: int, first_frm_is_start_point: bool, - last_frm_is_end_point: bool, end_point_is_sent_end: bool, cache: dict = {}) -> None: - self.PopDataBufTillFrame(start_frm, cache=cache) - expected_sample_number = int(frm_cnt * self.vad_opts.sample_rate * self.vad_opts.frame_in_ms / 1000) - if last_frm_is_end_point: - extra_sample = max(0, int(self.vad_opts.frame_length_ms * self.vad_opts.sample_rate / 1000 - \ - self.vad_opts.sample_rate * self.vad_opts.frame_in_ms / 1000)) - expected_sample_number += int(extra_sample) - if end_point_is_sent_end: - expected_sample_number = max(expected_sample_number, len(cache["stats"].data_buf)) - if len(cache["stats"].data_buf) < expected_sample_number: - print('error in calling pop data_buf\n') - - if len(cache["stats"].output_data_buf) == 0 or first_frm_is_start_point: - cache["stats"].output_data_buf.append(E2EVadSpeechBufWithDoa()) - cache["stats"].output_data_buf[-1].Reset() - cache["stats"].output_data_buf[-1].start_ms = start_frm * self.vad_opts.frame_in_ms - cache["stats"].output_data_buf[-1].end_ms = cache["stats"].output_data_buf[-1].start_ms - cache["stats"].output_data_buf[-1].doa = 0 - cur_seg = cache["stats"].output_data_buf[-1] - if cur_seg.end_ms != start_frm * self.vad_opts.frame_in_ms: - print('warning\n') - out_pos = len(cur_seg.buffer) # cur_seg.buff现在没做任何操作 - data_to_pop = 0 - if end_point_is_sent_end: - data_to_pop = expected_sample_number - else: - data_to_pop = int(frm_cnt * self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000) - if data_to_pop > len(cache["stats"].data_buf): - print('VAD data_to_pop is bigger than cache["stats"].data_buf.size()!!!\n') - data_to_pop = len(cache["stats"].data_buf) - expected_sample_number = len(cache["stats"].data_buf) - - cur_seg.doa = 0 - for sample_cpy_out in range(0, data_to_pop): - # cur_seg.buffer[out_pos ++] = data_buf_.back(); - out_pos += 1 - for sample_cpy_out in range(data_to_pop, expected_sample_number): - # cur_seg.buffer[out_pos++] = data_buf_.back() - out_pos += 1 - if cur_seg.end_ms != start_frm * self.vad_opts.frame_in_ms: - print('Something wrong with the VAD algorithm\n') - cache["stats"].data_buf_start_frame += frm_cnt - cur_seg.end_ms = (start_frm + frm_cnt) * self.vad_opts.frame_in_ms - if first_frm_is_start_point: - cur_seg.contain_seg_start_point = True - if last_frm_is_end_point: - cur_seg.contain_seg_end_point = True - - def OnSilenceDetected(self, valid_frame: int, cache: dict = {}): - cache["stats"].lastest_confirmed_silence_frame = valid_frame - if cache["stats"].vad_state_machine == VadStateMachine.kVadInStateStartPointNotDetected: - self.PopDataBufTillFrame(valid_frame, cache=cache) - - # silence_detected_callback_ - # pass - - def OnVoiceDetected(self, valid_frame: int, cache: dict = {}) -> None: - cache["stats"].latest_confirmed_speech_frame = valid_frame - self.PopDataToOutputBuf(valid_frame, 1, False, False, False, cache=cache) - - def OnVoiceStart(self, start_frame: int, fake_result: bool = False, cache: dict = {}) -> None: - if self.vad_opts.do_start_point_detection: - pass - if cache["stats"].confirmed_start_frame != -1: - print('not reset vad properly\n') - else: - cache["stats"].confirmed_start_frame = start_frame - - if not fake_result and cache["stats"].vad_state_machine == VadStateMachine.kVadInStateStartPointNotDetected: - self.PopDataToOutputBuf(cache["stats"].confirmed_start_frame, 1, True, False, False, cache=cache) - - def OnVoiceEnd(self, end_frame: int, fake_result: bool, is_last_frame: bool, cache: dict = {}) -> None: - for t in range(cache["stats"].latest_confirmed_speech_frame + 1, end_frame): - self.OnVoiceDetected(t, cache=cache) - if self.vad_opts.do_end_point_detection: - pass - if cache["stats"].confirmed_end_frame != -1: - print('not reset vad properly\n') - else: - cache["stats"].confirmed_end_frame = end_frame - if not fake_result: - cache["stats"].sil_frame = 0 - self.PopDataToOutputBuf(cache["stats"].confirmed_end_frame, 1, False, True, is_last_frame, cache=cache) - cache["stats"].number_end_time_detected += 1 - - def MaybeOnVoiceEndIfLastFrame(self, is_final_frame: bool, cur_frm_idx: int, cache: dict = {}) -> None: - if is_final_frame: - self.OnVoiceEnd(cur_frm_idx, False, True, cache=cache) - cache["stats"].vad_state_machine = VadStateMachine.kVadInStateEndPointDetected - - def GetLatency(self, cache: dict = {}) -> int: - return int(self.LatencyFrmNumAtStartPoint(cache=cache) * self.vad_opts.frame_in_ms) - - def LatencyFrmNumAtStartPoint(self, cache: dict = {}) -> int: - vad_latency = cache["windows_detector"].GetWinSize() - if self.vad_opts.do_extend: - vad_latency += int(self.vad_opts.lookback_time_start_point / self.vad_opts.frame_in_ms) - return vad_latency - - def GetFrameState(self, t: int, cache: dict = {}): - frame_state = FrameState.kFrameStateInvalid - cur_decibel = cache["stats"].decibel[t] - cur_snr = cur_decibel - cache["stats"].noise_average_decibel - # for each frame, calc log posterior probability of each state - if cur_decibel < self.vad_opts.decibel_thres: - frame_state = FrameState.kFrameStateSil - self.DetectOneFrame(frame_state, t, False, cache=cache) - return frame_state - - sum_score = 0.0 - noise_prob = 0.0 - assert len(cache["stats"].sil_pdf_ids) == self.vad_opts.silence_pdf_num - if len(cache["stats"].sil_pdf_ids) > 0: - assert len(cache["stats"].scores) == 1 # 只支持batch_size = 1的测试 - sil_pdf_scores = [cache["stats"].scores[0][t][sil_pdf_id] for sil_pdf_id in cache["stats"].sil_pdf_ids] - sum_score = sum(sil_pdf_scores) - noise_prob = math.log(sum_score) * self.vad_opts.speech_2_noise_ratio - total_score = 1.0 - sum_score = total_score - sum_score - speech_prob = math.log(sum_score) - if self.vad_opts.output_frame_probs: - frame_prob = E2EVadFrameProb() - frame_prob.noise_prob = noise_prob - frame_prob.speech_prob = speech_prob - frame_prob.score = sum_score - frame_prob.frame_id = t - cache["stats"].frame_probs.append(frame_prob) - if math.exp(speech_prob) >= math.exp(noise_prob) + cache["stats"].speech_noise_thres: - if cur_snr >= self.vad_opts.snr_thres and cur_decibel >= self.vad_opts.decibel_thres: - frame_state = FrameState.kFrameStateSpeech - else: - frame_state = FrameState.kFrameStateSil - else: - frame_state = FrameState.kFrameStateSil - if cache["stats"].noise_average_decibel < -99.9: - cache["stats"].noise_average_decibel = cur_decibel - else: - cache["stats"].noise_average_decibel = (cur_decibel + cache["stats"].noise_average_decibel * ( - self.vad_opts.noise_frame_num_used_for_snr - - 1)) / self.vad_opts.noise_frame_num_used_for_snr - - return frame_state - - def forward(self, feats: torch.Tensor, - waveform: torch.tensor, - cache: dict = {}, - is_final: bool = False, - **kwargs, - ): - # if len(cache) == 0: - # self.AllResetDetection() - # self.waveform = waveform # compute decibel for each frame - cache["stats"].waveform = waveform - is_streaming_input = kwargs.get("is_streaming_input", True) - self.ComputeDecibel(cache=cache) - self.ComputeScores(feats, cache=cache) - if not is_final: - self.DetectCommonFrames(cache=cache) - else: - self.DetectLastFrames(cache=cache) - segments = [] - for batch_num in range(0, feats.shape[0]): # only support batch_size = 1 now - segment_batch = [] - if len(cache["stats"].output_data_buf) > 0: - for i in range(cache["stats"].output_data_buf_offset, len(cache["stats"].output_data_buf)): - if is_streaming_input: # in this case, return [beg, -1], [], [-1, end], [beg, end] - if not cache["stats"].output_data_buf[i].contain_seg_start_point: - continue - if not cache["stats"].next_seg and not cache["stats"].output_data_buf[i].contain_seg_end_point: - continue - start_ms = cache["stats"].output_data_buf[i].start_ms if cache["stats"].next_seg else -1 - if cache["stats"].output_data_buf[i].contain_seg_end_point: - end_ms = cache["stats"].output_data_buf[i].end_ms - cache["stats"].next_seg = True - cache["stats"].output_data_buf_offset += 1 - else: - end_ms = -1 - cache["stats"].next_seg = False - segment = [start_ms, end_ms] - - else: # in this case, return [beg, end] - - if not is_final and (not cache["stats"].output_data_buf[i].contain_seg_start_point or not - cache["stats"].output_data_buf[ - i].contain_seg_end_point): - continue - segment = [cache["stats"].output_data_buf[i].start_ms, cache["stats"].output_data_buf[i].end_ms] - cache["stats"].output_data_buf_offset += 1 # need update this parameter - - segment_batch.append(segment) - - if segment_batch: - segments.append(segment_batch) - # if is_final: - # # reset class variables and clear the dict for the next query - # self.AllResetDetection() - return segments - - def init_cache(self, cache: dict = {}, **kwargs): - - cache["frontend"] = {} - cache["prev_samples"] = torch.empty(0) - cache["encoder"] = {} + """ + Author: Speech Lab of DAMO Academy, Alibaba Group + Deep-FSMN for Large Vocabulary Continuous Speech Recognition + https://arxiv.org/abs/1803.05030 + """ - if kwargs.get("max_end_silence_time") is not None: - # update the max_end_silence_time - self.vad_opts.max_end_silence_time = kwargs.get("max_end_silence_time") + def __init__( + self, + encoder: str = None, + encoder_conf: Optional[Dict] = None, + vad_post_args: Dict[str, Any] = None, + **kwargs, + ): + super().__init__() + self.vad_opts = VADXOptions(**kwargs) - windows_detector = WindowDetector(self.vad_opts.window_size_ms, - self.vad_opts.sil_to_speech_time_thres, - self.vad_opts.speech_to_sil_time_thres, - self.vad_opts.frame_in_ms) - windows_detector.Reset() - - stats = Stats(sil_pdf_ids=self.vad_opts.sil_pdf_ids, - max_end_sil_frame_cnt_thresh=self.vad_opts.max_end_silence_time - self.vad_opts.speech_to_sil_time_thres, - speech_noise_thres=self.vad_opts.speech_noise_thres - ) - cache["windows_detector"] = windows_detector - cache["stats"] = stats - return cache - - def inference(self, - data_in, - data_lengths=None, - key: list = None, - tokenizer=None, - frontend=None, - cache: dict = {}, - **kwargs, - ): - - if len(cache) == 0: - self.init_cache(cache, **kwargs) - - meta_data = {} - chunk_size = kwargs.get("chunk_size", 60000) # 50ms - chunk_stride_samples = int(chunk_size * frontend.fs / 1000) - - time1 = time.perf_counter() - is_streaming_input = kwargs.get("is_streaming_input", False) if chunk_size >= 15000 else kwargs.get("is_streaming_input", True) - is_final = kwargs.get("is_final", False) if is_streaming_input else kwargs.get("is_final", True) - cfg = {"is_final": is_final, "is_streaming_input": is_streaming_input} - audio_sample_list = load_audio_text_image_video(data_in, - fs=frontend.fs, - audio_fs=kwargs.get("fs", 16000), - data_type=kwargs.get("data_type", "sound"), - tokenizer=tokenizer, - cache=cfg, - ) - _is_final = cfg["is_final"] # if data_in is a file or url, set is_final=True - is_streaming_input = cfg["is_streaming_input"] - time2 = time.perf_counter() - meta_data["load_data"] = f"{time2 - time1:0.3f}" - assert len(audio_sample_list) == 1, "batch_size must be set 1" - - audio_sample = torch.cat((cache["prev_samples"], audio_sample_list[0])) - - n = int(len(audio_sample) // chunk_stride_samples + int(_is_final)) - m = int(len(audio_sample) % chunk_stride_samples * (1 - int(_is_final))) - segments = [] - for i in range(n): - kwargs["is_final"] = _is_final and i == n - 1 - audio_sample_i = audio_sample[i * chunk_stride_samples:(i + 1) * chunk_stride_samples] - - # extract fbank feats - speech, speech_lengths = extract_fbank([audio_sample_i], data_type=kwargs.get("data_type", "sound"), - frontend=frontend, cache=cache["frontend"], - is_final=kwargs["is_final"]) - time3 = time.perf_counter() - meta_data["extract_feat"] = f"{time3 - time2:0.3f}" - meta_data["batch_data_time"] = speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 - speech = speech.to(device=kwargs["device"]) - speech_lengths = speech_lengths.to(device=kwargs["device"]) - - batch = { - "feats": speech, - "waveform": cache["frontend"]["waveforms"], - "is_final": kwargs["is_final"], - "cache": cache, - "is_streaming_input": is_streaming_input - } - segments_i = self.forward(**batch) - if len(segments_i) > 0: - segments.extend(*segments_i) - - cache["prev_samples"] = audio_sample[:-m] - if _is_final: - self.init_cache(cache) - - ibest_writer = None - if kwargs.get("output_dir") is not None: - if not hasattr(self, "writer"): - self.writer = DatadirWriter(kwargs.get("output_dir")) - ibest_writer = self.writer[f"{1}best_recog"] - - results = [] - result_i = {"key": key[0], "value": segments} - # if "MODELSCOPE_ENVIRONMENT" in os.environ and os.environ["MODELSCOPE_ENVIRONMENT"] == "eas": - # result_i = json.dumps(result_i) - - results.append(result_i) - - if ibest_writer is not None: - ibest_writer["text"][key[0]] = segments - - return results, meta_data - - def export(self, **kwargs): + encoder_class = tables.encoder_classes.get(encoder) + encoder = encoder_class(**encoder_conf) + self.encoder = encoder + self.encoder_conf = encoder_conf - from .export_meta import export_rebuild_model - models = export_rebuild_model(model=self, **kwargs) - return models + def ResetDetection(self, cache: dict = {}): + cache["stats"].continous_silence_frame_count = 0 + cache["stats"].latest_confirmed_speech_frame = 0 + cache["stats"].lastest_confirmed_silence_frame = -1 + cache["stats"].confirmed_start_frame = -1 + cache["stats"].confirmed_end_frame = -1 + cache["stats"].vad_state_machine = VadStateMachine.kVadInStateStartPointNotDetected + cache["windows_detector"].Reset() + cache["stats"].sil_frame = 0 + cache["stats"].frame_probs = [] - def DetectCommonFrames(self, cache: dict = {}) -> int: - if cache["stats"].vad_state_machine == VadStateMachine.kVadInStateEndPointDetected: - return 0 - for i in range(self.vad_opts.nn_eval_block_size - 1, -1, -1): - frame_state = FrameState.kFrameStateInvalid - frame_state = self.GetFrameState(cache["stats"].frm_cnt - 1 - i - cache["stats"].last_drop_frames, - cache=cache) - self.DetectOneFrame(frame_state, cache["stats"].frm_cnt - 1 - i, False, cache=cache) - - return 0 - - def DetectLastFrames(self, cache: dict = {}) -> int: - if cache["stats"].vad_state_machine == VadStateMachine.kVadInStateEndPointDetected: - return 0 - for i in range(self.vad_opts.nn_eval_block_size - 1, -1, -1): - frame_state = FrameState.kFrameStateInvalid - frame_state = self.GetFrameState(cache["stats"].frm_cnt - 1 - i - cache["stats"].last_drop_frames, - cache=cache) - if i != 0: - self.DetectOneFrame(frame_state, cache["stats"].frm_cnt - 1 - i, False, cache=cache) - else: - self.DetectOneFrame(frame_state, cache["stats"].frm_cnt - 1, True, cache=cache) - - return 0 - - def DetectOneFrame(self, cur_frm_state: FrameState, cur_frm_idx: int, is_final_frame: bool, - cache: dict = {}) -> None: - tmp_cur_frm_state = FrameState.kFrameStateInvalid - if cur_frm_state == FrameState.kFrameStateSpeech: - if math.fabs(1.0) > self.vad_opts.fe_prior_thres: - tmp_cur_frm_state = FrameState.kFrameStateSpeech - else: - tmp_cur_frm_state = FrameState.kFrameStateSil - elif cur_frm_state == FrameState.kFrameStateSil: - tmp_cur_frm_state = FrameState.kFrameStateSil - state_change = cache["windows_detector"].DetectOneFrame(tmp_cur_frm_state, cur_frm_idx, cache=cache) - frm_shift_in_ms = self.vad_opts.frame_in_ms - if AudioChangeState.kChangeStateSil2Speech == state_change: - silence_frame_count = cache["stats"].continous_silence_frame_count - cache["stats"].continous_silence_frame_count = 0 - cache["stats"].pre_end_silence_detected = False - start_frame = 0 - if cache["stats"].vad_state_machine == VadStateMachine.kVadInStateStartPointNotDetected: - start_frame = max(cache["stats"].data_buf_start_frame, - cur_frm_idx - self.LatencyFrmNumAtStartPoint(cache=cache)) - self.OnVoiceStart(start_frame, cache=cache) - cache["stats"].vad_state_machine = VadStateMachine.kVadInStateInSpeechSegment - for t in range(start_frame + 1, cur_frm_idx + 1): - self.OnVoiceDetected(t, cache=cache) - elif cache["stats"].vad_state_machine == VadStateMachine.kVadInStateInSpeechSegment: - for t in range(cache["stats"].latest_confirmed_speech_frame + 1, cur_frm_idx): - self.OnVoiceDetected(t, cache=cache) - if cur_frm_idx - cache["stats"].confirmed_start_frame + 1 > \ - self.vad_opts.max_single_segment_time / frm_shift_in_ms: - self.OnVoiceEnd(cur_frm_idx, False, False, cache=cache) - cache["stats"].vad_state_machine = VadStateMachine.kVadInStateEndPointDetected - elif not is_final_frame: - self.OnVoiceDetected(cur_frm_idx, cache=cache) - else: - self.MaybeOnVoiceEndIfLastFrame(is_final_frame, cur_frm_idx, cache=cache) - else: - pass - elif AudioChangeState.kChangeStateSpeech2Sil == state_change: - cache["stats"].continous_silence_frame_count = 0 - if cache["stats"].vad_state_machine == VadStateMachine.kVadInStateStartPointNotDetected: - pass - elif cache["stats"].vad_state_machine == VadStateMachine.kVadInStateInSpeechSegment: - if cur_frm_idx - cache["stats"].confirmed_start_frame + 1 > \ - self.vad_opts.max_single_segment_time / frm_shift_in_ms: - self.OnVoiceEnd(cur_frm_idx, False, False, cache=cache) - cache["stats"].vad_state_machine = VadStateMachine.kVadInStateEndPointDetected - elif not is_final_frame: - self.OnVoiceDetected(cur_frm_idx, cache=cache) - else: - self.MaybeOnVoiceEndIfLastFrame(is_final_frame, cur_frm_idx, cache=cache) - else: - pass - elif AudioChangeState.kChangeStateSpeech2Speech == state_change: - cache["stats"].continous_silence_frame_count = 0 - if cache["stats"].vad_state_machine == VadStateMachine.kVadInStateInSpeechSegment: - if cur_frm_idx - cache["stats"].confirmed_start_frame + 1 > \ - self.vad_opts.max_single_segment_time / frm_shift_in_ms: - cache["stats"].max_time_out = True - self.OnVoiceEnd(cur_frm_idx, False, False, cache=cache) - cache["stats"].vad_state_machine = VadStateMachine.kVadInStateEndPointDetected - elif not is_final_frame: - self.OnVoiceDetected(cur_frm_idx, cache=cache) - else: - self.MaybeOnVoiceEndIfLastFrame(is_final_frame, cur_frm_idx, cache=cache) - else: - pass - elif AudioChangeState.kChangeStateSil2Sil == state_change: - cache["stats"].continous_silence_frame_count += 1 - if cache["stats"].vad_state_machine == VadStateMachine.kVadInStateStartPointNotDetected: - # silence timeout, return zero length decision - if ((self.vad_opts.detect_mode == VadDetectMode.kVadSingleUtteranceDetectMode.value) and ( - cache[ - "stats"].continous_silence_frame_count * frm_shift_in_ms > self.vad_opts.max_start_silence_time)) \ - or (is_final_frame and cache["stats"].number_end_time_detected == 0): - for t in range(cache["stats"].lastest_confirmed_silence_frame + 1, cur_frm_idx): - self.OnSilenceDetected(t, cache=cache) - self.OnVoiceStart(0, True, cache=cache) - self.OnVoiceEnd(0, True, False, cache=cache) - cache["stats"].vad_state_machine = VadStateMachine.kVadInStateEndPointDetected - else: - if cur_frm_idx >= self.LatencyFrmNumAtStartPoint(cache=cache): - self.OnSilenceDetected(cur_frm_idx - self.LatencyFrmNumAtStartPoint(cache=cache), cache=cache) - elif cache["stats"].vad_state_machine == VadStateMachine.kVadInStateInSpeechSegment: - if cache["stats"].continous_silence_frame_count * frm_shift_in_ms >= cache[ - "stats"].max_end_sil_frame_cnt_thresh: - lookback_frame = int(cache["stats"].max_end_sil_frame_cnt_thresh / frm_shift_in_ms) - if self.vad_opts.do_extend: - lookback_frame -= int(self.vad_opts.lookahead_time_end_point / frm_shift_in_ms) - lookback_frame -= 1 - lookback_frame = max(0, lookback_frame) - self.OnVoiceEnd(cur_frm_idx - lookback_frame, False, False, cache=cache) - cache["stats"].vad_state_machine = VadStateMachine.kVadInStateEndPointDetected - elif cur_frm_idx - cache["stats"].confirmed_start_frame + 1 > \ - self.vad_opts.max_single_segment_time / frm_shift_in_ms: - self.OnVoiceEnd(cur_frm_idx, False, False, cache=cache) - cache["stats"].vad_state_machine = VadStateMachine.kVadInStateEndPointDetected - elif self.vad_opts.do_extend and not is_final_frame: - if cache["stats"].continous_silence_frame_count <= int( - self.vad_opts.lookahead_time_end_point / frm_shift_in_ms): - self.OnVoiceDetected(cur_frm_idx, cache=cache) - else: - self.MaybeOnVoiceEndIfLastFrame(is_final_frame, cur_frm_idx, cache=cache) - else: - pass - - if cache["stats"].vad_state_machine == VadStateMachine.kVadInStateEndPointDetected and \ - self.vad_opts.detect_mode == VadDetectMode.kVadMutipleUtteranceDetectMode.value: - self.ResetDetection(cache=cache) + if cache["stats"].output_data_buf: + assert cache["stats"].output_data_buf[-1].contain_seg_end_point == True + drop_frames = int(cache["stats"].output_data_buf[-1].end_ms / self.vad_opts.frame_in_ms) + real_drop_frames = drop_frames - cache["stats"].last_drop_frames + cache["stats"].last_drop_frames = drop_frames + cache["stats"].data_buf_all = cache["stats"].data_buf_all[ + real_drop_frames + * int(self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000) : + ] + cache["stats"].decibel = cache["stats"].decibel[real_drop_frames:] + cache["stats"].scores = cache["stats"].scores[:, real_drop_frames:, :] + def ComputeDecibel(self, cache: dict = {}) -> None: + frame_sample_length = int(self.vad_opts.frame_length_ms * self.vad_opts.sample_rate / 1000) + frame_shift_length = int(self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000) + if cache["stats"].data_buf_all is None: + cache["stats"].data_buf_all = cache["stats"].waveform[ + 0 + ] # cache["stats"].data_buf is pointed to cache["stats"].waveform[0] + cache["stats"].data_buf = cache["stats"].data_buf_all + else: + cache["stats"].data_buf_all = torch.cat( + (cache["stats"].data_buf_all, cache["stats"].waveform[0]) + ) + for offset in range( + 0, cache["stats"].waveform.shape[1] - frame_sample_length + 1, frame_shift_length + ): + cache["stats"].decibel.append( + 10 + * math.log10( + (cache["stats"].waveform[0][offset : offset + frame_sample_length]) + .square() + .sum() + + 0.000001 + ) + ) + def ComputeScores(self, feats: torch.Tensor, cache: dict = {}) -> None: + scores = self.encoder(feats, cache=cache["encoder"]).to("cpu") # return B * T * D + assert ( + scores.shape[1] == feats.shape[1] + ), "The shape between feats and scores does not match" + self.vad_opts.nn_eval_block_size = scores.shape[1] + cache["stats"].frm_cnt += scores.shape[1] # count total frames + if cache["stats"].scores is None: + cache["stats"].scores = scores # the first calculation + else: + cache["stats"].scores = torch.cat((cache["stats"].scores, scores), dim=1) + + def PopDataBufTillFrame(self, frame_idx: int, cache: dict = {}) -> None: # need check again + while cache["stats"].data_buf_start_frame < frame_idx: + if len(cache["stats"].data_buf) >= int( + self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000 + ): + cache["stats"].data_buf_start_frame += 1 + cache["stats"].data_buf = cache["stats"].data_buf_all[ + (cache["stats"].data_buf_start_frame - cache["stats"].last_drop_frames) + * int(self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000) : + ] + + def PopDataToOutputBuf( + self, + start_frm: int, + frm_cnt: int, + first_frm_is_start_point: bool, + last_frm_is_end_point: bool, + end_point_is_sent_end: bool, + cache: dict = {}, + ) -> None: + self.PopDataBufTillFrame(start_frm, cache=cache) + expected_sample_number = int( + frm_cnt * self.vad_opts.sample_rate * self.vad_opts.frame_in_ms / 1000 + ) + if last_frm_is_end_point: + extra_sample = max( + 0, + int( + self.vad_opts.frame_length_ms * self.vad_opts.sample_rate / 1000 + - self.vad_opts.sample_rate * self.vad_opts.frame_in_ms / 1000 + ), + ) + expected_sample_number += int(extra_sample) + if end_point_is_sent_end: + expected_sample_number = max(expected_sample_number, len(cache["stats"].data_buf)) + if len(cache["stats"].data_buf) < expected_sample_number: + print("error in calling pop data_buf\n") + + if len(cache["stats"].output_data_buf) == 0 or first_frm_is_start_point: + cache["stats"].output_data_buf.append(E2EVadSpeechBufWithDoa()) + cache["stats"].output_data_buf[-1].Reset() + cache["stats"].output_data_buf[-1].start_ms = start_frm * self.vad_opts.frame_in_ms + cache["stats"].output_data_buf[-1].end_ms = cache["stats"].output_data_buf[-1].start_ms + cache["stats"].output_data_buf[-1].doa = 0 + cur_seg = cache["stats"].output_data_buf[-1] + if cur_seg.end_ms != start_frm * self.vad_opts.frame_in_ms: + print("warning\n") + out_pos = len(cur_seg.buffer) # cur_seg.buff现在没做任何操作 + data_to_pop = 0 + if end_point_is_sent_end: + data_to_pop = expected_sample_number + else: + data_to_pop = int( + frm_cnt * self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000 + ) + if data_to_pop > len(cache["stats"].data_buf): + print('VAD data_to_pop is bigger than cache["stats"].data_buf.size()!!!\n') + data_to_pop = len(cache["stats"].data_buf) + expected_sample_number = len(cache["stats"].data_buf) + + cur_seg.doa = 0 + for sample_cpy_out in range(0, data_to_pop): + # cur_seg.buffer[out_pos ++] = data_buf_.back(); + out_pos += 1 + for sample_cpy_out in range(data_to_pop, expected_sample_number): + # cur_seg.buffer[out_pos++] = data_buf_.back() + out_pos += 1 + if cur_seg.end_ms != start_frm * self.vad_opts.frame_in_ms: + print("Something wrong with the VAD algorithm\n") + cache["stats"].data_buf_start_frame += frm_cnt + cur_seg.end_ms = (start_frm + frm_cnt) * self.vad_opts.frame_in_ms + if first_frm_is_start_point: + cur_seg.contain_seg_start_point = True + if last_frm_is_end_point: + cur_seg.contain_seg_end_point = True + + def OnSilenceDetected(self, valid_frame: int, cache: dict = {}): + cache["stats"].lastest_confirmed_silence_frame = valid_frame + if cache["stats"].vad_state_machine == VadStateMachine.kVadInStateStartPointNotDetected: + self.PopDataBufTillFrame(valid_frame, cache=cache) + + # silence_detected_callback_ + # pass + + def OnVoiceDetected(self, valid_frame: int, cache: dict = {}) -> None: + cache["stats"].latest_confirmed_speech_frame = valid_frame + self.PopDataToOutputBuf(valid_frame, 1, False, False, False, cache=cache) + + def OnVoiceStart(self, start_frame: int, fake_result: bool = False, cache: dict = {}) -> None: + if self.vad_opts.do_start_point_detection: + pass + if cache["stats"].confirmed_start_frame != -1: + print("not reset vad properly\n") + else: + cache["stats"].confirmed_start_frame = start_frame + + if ( + not fake_result + and cache["stats"].vad_state_machine == VadStateMachine.kVadInStateStartPointNotDetected + ): + self.PopDataToOutputBuf( + cache["stats"].confirmed_start_frame, 1, True, False, False, cache=cache + ) + + def OnVoiceEnd( + self, end_frame: int, fake_result: bool, is_last_frame: bool, cache: dict = {} + ) -> None: + for t in range(cache["stats"].latest_confirmed_speech_frame + 1, end_frame): + self.OnVoiceDetected(t, cache=cache) + if self.vad_opts.do_end_point_detection: + pass + if cache["stats"].confirmed_end_frame != -1: + print("not reset vad properly\n") + else: + cache["stats"].confirmed_end_frame = end_frame + if not fake_result: + cache["stats"].sil_frame = 0 + self.PopDataToOutputBuf( + cache["stats"].confirmed_end_frame, 1, False, True, is_last_frame, cache=cache + ) + cache["stats"].number_end_time_detected += 1 + + def MaybeOnVoiceEndIfLastFrame( + self, is_final_frame: bool, cur_frm_idx: int, cache: dict = {} + ) -> None: + if is_final_frame: + self.OnVoiceEnd(cur_frm_idx, False, True, cache=cache) + cache["stats"].vad_state_machine = VadStateMachine.kVadInStateEndPointDetected + + def GetLatency(self, cache: dict = {}) -> int: + return int(self.LatencyFrmNumAtStartPoint(cache=cache) * self.vad_opts.frame_in_ms) + + def LatencyFrmNumAtStartPoint(self, cache: dict = {}) -> int: + vad_latency = cache["windows_detector"].GetWinSize() + if self.vad_opts.do_extend: + vad_latency += int(self.vad_opts.lookback_time_start_point / self.vad_opts.frame_in_ms) + return vad_latency + + def GetFrameState(self, t: int, cache: dict = {}): + frame_state = FrameState.kFrameStateInvalid + cur_decibel = cache["stats"].decibel[t] + cur_snr = cur_decibel - cache["stats"].noise_average_decibel + # for each frame, calc log posterior probability of each state + if cur_decibel < self.vad_opts.decibel_thres: + frame_state = FrameState.kFrameStateSil + self.DetectOneFrame(frame_state, t, False, cache=cache) + return frame_state + + sum_score = 0.0 + noise_prob = 0.0 + assert len(cache["stats"].sil_pdf_ids) == self.vad_opts.silence_pdf_num + if len(cache["stats"].sil_pdf_ids) > 0: + assert len(cache["stats"].scores) == 1 # 只支持batch_size = 1的测试 + sil_pdf_scores = [ + cache["stats"].scores[0][t][sil_pdf_id] for sil_pdf_id in cache["stats"].sil_pdf_ids + ] + sum_score = sum(sil_pdf_scores) + noise_prob = math.log(sum_score) * self.vad_opts.speech_2_noise_ratio + total_score = 1.0 + sum_score = total_score - sum_score + speech_prob = math.log(sum_score) + if self.vad_opts.output_frame_probs: + frame_prob = E2EVadFrameProb() + frame_prob.noise_prob = noise_prob + frame_prob.speech_prob = speech_prob + frame_prob.score = sum_score + frame_prob.frame_id = t + cache["stats"].frame_probs.append(frame_prob) + if math.exp(speech_prob) >= math.exp(noise_prob) + cache["stats"].speech_noise_thres: + if cur_snr >= self.vad_opts.snr_thres and cur_decibel >= self.vad_opts.decibel_thres: + frame_state = FrameState.kFrameStateSpeech + else: + frame_state = FrameState.kFrameStateSil + else: + frame_state = FrameState.kFrameStateSil + if cache["stats"].noise_average_decibel < -99.9: + cache["stats"].noise_average_decibel = cur_decibel + else: + cache["stats"].noise_average_decibel = ( + cur_decibel + + cache["stats"].noise_average_decibel + * (self.vad_opts.noise_frame_num_used_for_snr - 1) + ) / self.vad_opts.noise_frame_num_used_for_snr + + return frame_state + + def forward( + self, + feats: torch.Tensor, + waveform: torch.tensor, + cache: dict = {}, + is_final: bool = False, + **kwargs, + ): + # if len(cache) == 0: + # self.AllResetDetection() + # self.waveform = waveform # compute decibel for each frame + cache["stats"].waveform = waveform + is_streaming_input = kwargs.get("is_streaming_input", True) + self.ComputeDecibel(cache=cache) + self.ComputeScores(feats, cache=cache) + if not is_final: + self.DetectCommonFrames(cache=cache) + else: + self.DetectLastFrames(cache=cache) + segments = [] + for batch_num in range(0, feats.shape[0]): # only support batch_size = 1 now + segment_batch = [] + if len(cache["stats"].output_data_buf) > 0: + for i in range( + cache["stats"].output_data_buf_offset, len(cache["stats"].output_data_buf) + ): + if ( + is_streaming_input + ): # in this case, return [beg, -1], [], [-1, end], [beg, end] + if not cache["stats"].output_data_buf[i].contain_seg_start_point: + continue + if ( + not cache["stats"].next_seg + and not cache["stats"].output_data_buf[i].contain_seg_end_point + ): + continue + start_ms = ( + cache["stats"].output_data_buf[i].start_ms + if cache["stats"].next_seg + else -1 + ) + if cache["stats"].output_data_buf[i].contain_seg_end_point: + end_ms = cache["stats"].output_data_buf[i].end_ms + cache["stats"].next_seg = True + cache["stats"].output_data_buf_offset += 1 + else: + end_ms = -1 + cache["stats"].next_seg = False + segment = [start_ms, end_ms] + + else: # in this case, return [beg, end] + + if not is_final and ( + not cache["stats"].output_data_buf[i].contain_seg_start_point + or not cache["stats"].output_data_buf[i].contain_seg_end_point + ): + continue + segment = [ + cache["stats"].output_data_buf[i].start_ms, + cache["stats"].output_data_buf[i].end_ms, + ] + cache["stats"].output_data_buf_offset += 1 # need update this parameter + + segment_batch.append(segment) + + if segment_batch: + segments.append(segment_batch) + # if is_final: + # # reset class variables and clear the dict for the next query + # self.AllResetDetection() + return segments + + def init_cache(self, cache: dict = {}, **kwargs): + + cache["frontend"] = {} + cache["prev_samples"] = torch.empty(0) + cache["encoder"] = {} + + if kwargs.get("max_end_silence_time") is not None: + # update the max_end_silence_time + self.vad_opts.max_end_silence_time = kwargs.get("max_end_silence_time") + + windows_detector = WindowDetector( + self.vad_opts.window_size_ms, + self.vad_opts.sil_to_speech_time_thres, + self.vad_opts.speech_to_sil_time_thres, + self.vad_opts.frame_in_ms, + ) + windows_detector.Reset() + + stats = Stats( + sil_pdf_ids=self.vad_opts.sil_pdf_ids, + max_end_sil_frame_cnt_thresh=self.vad_opts.max_end_silence_time + - self.vad_opts.speech_to_sil_time_thres, + speech_noise_thres=self.vad_opts.speech_noise_thres, + ) + cache["windows_detector"] = windows_detector + cache["stats"] = stats + return cache + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + cache: dict = {}, + **kwargs, + ): + + if len(cache) == 0: + self.init_cache(cache, **kwargs) + + meta_data = {} + chunk_size = kwargs.get("chunk_size", 60000) # 50ms + chunk_stride_samples = int(chunk_size * frontend.fs / 1000) + + time1 = time.perf_counter() + is_streaming_input = ( + kwargs.get("is_streaming_input", False) + if chunk_size >= 15000 + else kwargs.get("is_streaming_input", True) + ) + is_final = ( + kwargs.get("is_final", False) if is_streaming_input else kwargs.get("is_final", True) + ) + cfg = {"is_final": is_final, "is_streaming_input": is_streaming_input} + audio_sample_list = load_audio_text_image_video( + data_in, + fs=frontend.fs, + audio_fs=kwargs.get("fs", 16000), + data_type=kwargs.get("data_type", "sound"), + tokenizer=tokenizer, + cache=cfg, + ) + _is_final = cfg["is_final"] # if data_in is a file or url, set is_final=True + is_streaming_input = cfg["is_streaming_input"] + time2 = time.perf_counter() + meta_data["load_data"] = f"{time2 - time1:0.3f}" + assert len(audio_sample_list) == 1, "batch_size must be set 1" + + audio_sample = torch.cat((cache["prev_samples"], audio_sample_list[0])) + + n = int(len(audio_sample) // chunk_stride_samples + int(_is_final)) + m = int(len(audio_sample) % chunk_stride_samples * (1 - int(_is_final))) + segments = [] + for i in range(n): + kwargs["is_final"] = _is_final and i == n - 1 + audio_sample_i = audio_sample[i * chunk_stride_samples : (i + 1) * chunk_stride_samples] + + # extract fbank feats + speech, speech_lengths = extract_fbank( + [audio_sample_i], + data_type=kwargs.get("data_type", "sound"), + frontend=frontend, + cache=cache["frontend"], + is_final=kwargs["is_final"], + ) + time3 = time.perf_counter() + meta_data["extract_feat"] = f"{time3 - time2:0.3f}" + meta_data["batch_data_time"] = ( + speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 + ) + speech = speech.to(device=kwargs["device"]) + speech_lengths = speech_lengths.to(device=kwargs["device"]) + + batch = { + "feats": speech, + "waveform": cache["frontend"]["waveforms"], + "is_final": kwargs["is_final"], + "cache": cache, + "is_streaming_input": is_streaming_input, + } + segments_i = self.forward(**batch) + if len(segments_i) > 0: + segments.extend(*segments_i) + + cache["prev_samples"] = audio_sample[:-m] + if _is_final: + self.init_cache(cache) + + ibest_writer = None + if kwargs.get("output_dir") is not None: + if not hasattr(self, "writer"): + self.writer = DatadirWriter(kwargs.get("output_dir")) + ibest_writer = self.writer[f"{1}best_recog"] + + results = [] + result_i = {"key": key[0], "value": segments} + # if "MODELSCOPE_ENVIRONMENT" in os.environ and os.environ["MODELSCOPE_ENVIRONMENT"] == "eas": + # result_i = json.dumps(result_i) + + results.append(result_i) + + if ibest_writer is not None: + ibest_writer["text"][key[0]] = segments + + return results, meta_data + + def export(self, **kwargs): + + from .export_meta import export_rebuild_model + + models = export_rebuild_model(model=self, **kwargs) + return models + + def DetectCommonFrames(self, cache: dict = {}) -> int: + if cache["stats"].vad_state_machine == VadStateMachine.kVadInStateEndPointDetected: + return 0 + for i in range(self.vad_opts.nn_eval_block_size - 1, -1, -1): + frame_state = FrameState.kFrameStateInvalid + frame_state = self.GetFrameState( + cache["stats"].frm_cnt - 1 - i - cache["stats"].last_drop_frames, cache=cache + ) + self.DetectOneFrame(frame_state, cache["stats"].frm_cnt - 1 - i, False, cache=cache) + + return 0 + + def DetectLastFrames(self, cache: dict = {}) -> int: + if cache["stats"].vad_state_machine == VadStateMachine.kVadInStateEndPointDetected: + return 0 + for i in range(self.vad_opts.nn_eval_block_size - 1, -1, -1): + frame_state = FrameState.kFrameStateInvalid + frame_state = self.GetFrameState( + cache["stats"].frm_cnt - 1 - i - cache["stats"].last_drop_frames, cache=cache + ) + if i != 0: + self.DetectOneFrame(frame_state, cache["stats"].frm_cnt - 1 - i, False, cache=cache) + else: + self.DetectOneFrame(frame_state, cache["stats"].frm_cnt - 1, True, cache=cache) + + return 0 + + def DetectOneFrame( + self, cur_frm_state: FrameState, cur_frm_idx: int, is_final_frame: bool, cache: dict = {} + ) -> None: + tmp_cur_frm_state = FrameState.kFrameStateInvalid + if cur_frm_state == FrameState.kFrameStateSpeech: + if math.fabs(1.0) > self.vad_opts.fe_prior_thres: + tmp_cur_frm_state = FrameState.kFrameStateSpeech + else: + tmp_cur_frm_state = FrameState.kFrameStateSil + elif cur_frm_state == FrameState.kFrameStateSil: + tmp_cur_frm_state = FrameState.kFrameStateSil + state_change = cache["windows_detector"].DetectOneFrame( + tmp_cur_frm_state, cur_frm_idx, cache=cache + ) + frm_shift_in_ms = self.vad_opts.frame_in_ms + if AudioChangeState.kChangeStateSil2Speech == state_change: + silence_frame_count = cache["stats"].continous_silence_frame_count + cache["stats"].continous_silence_frame_count = 0 + cache["stats"].pre_end_silence_detected = False + start_frame = 0 + if cache["stats"].vad_state_machine == VadStateMachine.kVadInStateStartPointNotDetected: + start_frame = max( + cache["stats"].data_buf_start_frame, + cur_frm_idx - self.LatencyFrmNumAtStartPoint(cache=cache), + ) + self.OnVoiceStart(start_frame, cache=cache) + cache["stats"].vad_state_machine = VadStateMachine.kVadInStateInSpeechSegment + for t in range(start_frame + 1, cur_frm_idx + 1): + self.OnVoiceDetected(t, cache=cache) + elif cache["stats"].vad_state_machine == VadStateMachine.kVadInStateInSpeechSegment: + for t in range(cache["stats"].latest_confirmed_speech_frame + 1, cur_frm_idx): + self.OnVoiceDetected(t, cache=cache) + if ( + cur_frm_idx - cache["stats"].confirmed_start_frame + 1 + > self.vad_opts.max_single_segment_time / frm_shift_in_ms + ): + self.OnVoiceEnd(cur_frm_idx, False, False, cache=cache) + cache["stats"].vad_state_machine = VadStateMachine.kVadInStateEndPointDetected + elif not is_final_frame: + self.OnVoiceDetected(cur_frm_idx, cache=cache) + else: + self.MaybeOnVoiceEndIfLastFrame(is_final_frame, cur_frm_idx, cache=cache) + else: + pass + elif AudioChangeState.kChangeStateSpeech2Sil == state_change: + cache["stats"].continous_silence_frame_count = 0 + if cache["stats"].vad_state_machine == VadStateMachine.kVadInStateStartPointNotDetected: + pass + elif cache["stats"].vad_state_machine == VadStateMachine.kVadInStateInSpeechSegment: + if ( + cur_frm_idx - cache["stats"].confirmed_start_frame + 1 + > self.vad_opts.max_single_segment_time / frm_shift_in_ms + ): + self.OnVoiceEnd(cur_frm_idx, False, False, cache=cache) + cache["stats"].vad_state_machine = VadStateMachine.kVadInStateEndPointDetected + elif not is_final_frame: + self.OnVoiceDetected(cur_frm_idx, cache=cache) + else: + self.MaybeOnVoiceEndIfLastFrame(is_final_frame, cur_frm_idx, cache=cache) + else: + pass + elif AudioChangeState.kChangeStateSpeech2Speech == state_change: + cache["stats"].continous_silence_frame_count = 0 + if cache["stats"].vad_state_machine == VadStateMachine.kVadInStateInSpeechSegment: + if ( + cur_frm_idx - cache["stats"].confirmed_start_frame + 1 + > self.vad_opts.max_single_segment_time / frm_shift_in_ms + ): + cache["stats"].max_time_out = True + self.OnVoiceEnd(cur_frm_idx, False, False, cache=cache) + cache["stats"].vad_state_machine = VadStateMachine.kVadInStateEndPointDetected + elif not is_final_frame: + self.OnVoiceDetected(cur_frm_idx, cache=cache) + else: + self.MaybeOnVoiceEndIfLastFrame(is_final_frame, cur_frm_idx, cache=cache) + else: + pass + elif AudioChangeState.kChangeStateSil2Sil == state_change: + cache["stats"].continous_silence_frame_count += 1 + if cache["stats"].vad_state_machine == VadStateMachine.kVadInStateStartPointNotDetected: + # silence timeout, return zero length decision + if ( + (self.vad_opts.detect_mode == VadDetectMode.kVadSingleUtteranceDetectMode.value) + and ( + cache["stats"].continous_silence_frame_count * frm_shift_in_ms + > self.vad_opts.max_start_silence_time + ) + ) or (is_final_frame and cache["stats"].number_end_time_detected == 0): + for t in range(cache["stats"].lastest_confirmed_silence_frame + 1, cur_frm_idx): + self.OnSilenceDetected(t, cache=cache) + self.OnVoiceStart(0, True, cache=cache) + self.OnVoiceEnd(0, True, False, cache=cache) + cache["stats"].vad_state_machine = VadStateMachine.kVadInStateEndPointDetected + else: + if cur_frm_idx >= self.LatencyFrmNumAtStartPoint(cache=cache): + self.OnSilenceDetected( + cur_frm_idx - self.LatencyFrmNumAtStartPoint(cache=cache), cache=cache + ) + elif cache["stats"].vad_state_machine == VadStateMachine.kVadInStateInSpeechSegment: + if ( + cache["stats"].continous_silence_frame_count * frm_shift_in_ms + >= cache["stats"].max_end_sil_frame_cnt_thresh + ): + lookback_frame = int( + cache["stats"].max_end_sil_frame_cnt_thresh / frm_shift_in_ms + ) + if self.vad_opts.do_extend: + lookback_frame -= int( + self.vad_opts.lookahead_time_end_point / frm_shift_in_ms + ) + lookback_frame -= 1 + lookback_frame = max(0, lookback_frame) + self.OnVoiceEnd(cur_frm_idx - lookback_frame, False, False, cache=cache) + cache["stats"].vad_state_machine = VadStateMachine.kVadInStateEndPointDetected + elif ( + cur_frm_idx - cache["stats"].confirmed_start_frame + 1 + > self.vad_opts.max_single_segment_time / frm_shift_in_ms + ): + self.OnVoiceEnd(cur_frm_idx, False, False, cache=cache) + cache["stats"].vad_state_machine = VadStateMachine.kVadInStateEndPointDetected + elif self.vad_opts.do_extend and not is_final_frame: + if cache["stats"].continous_silence_frame_count <= int( + self.vad_opts.lookahead_time_end_point / frm_shift_in_ms + ): + self.OnVoiceDetected(cur_frm_idx, cache=cache) + else: + self.MaybeOnVoiceEndIfLastFrame(is_final_frame, cur_frm_idx, cache=cache) + else: + pass + + if ( + cache["stats"].vad_state_machine == VadStateMachine.kVadInStateEndPointDetected + and self.vad_opts.detect_mode == VadDetectMode.kVadMutipleUtteranceDetectMode.value + ): + self.ResetDetection(cache=cache) diff --git a/funasr/models/language_model/rnn/argument.py b/funasr/models/language_model/rnn/argument.py index b4c89d25f..8935262cc 100644 --- a/funasr/models/language_model/rnn/argument.py +++ b/funasr/models/language_model/rnn/argument.py @@ -43,9 +43,7 @@ def add_arguments_rnn_encoder_common(group): type=int, help="Number of encoder hidden units", ) - group.add_argument( - "--eprojs", default=320, type=int, help="Number of encoder projection units" - ) + group.add_argument("--eprojs", default=320, type=int, help="Number of encoder projection units") group.add_argument( "--subsample", default="1", @@ -66,12 +64,8 @@ def add_arguments_rnn_decoder_common(group): choices=["lstm", "gru"], help="Type of decoder network architecture", ) - group.add_argument( - "--dlayers", default=1, type=int, help="Number of decoder layers" - ) - group.add_argument( - "--dunits", default=320, type=int, help="Number of decoder hidden units" - ) + group.add_argument("--dlayers", default=1, type=int, help="Number of decoder layers") + group.add_argument("--dunits", default=320, type=int, help="Number of decoder hidden units") group.add_argument( "--dropout-rate-decoder", default=0.0, @@ -124,9 +118,7 @@ def add_arguments_rnn_attention_common(group): type=int, help="Number of attention transformation dimensions", ) - group.add_argument( - "--awin", default=5, type=int, help="Window size for location2d attention" - ) + group.add_argument("--awin", default=5, type=int, help="Window size for location2d attention") group.add_argument( "--aheads", default=4, diff --git a/funasr/models/language_model/rnn/attentions.py b/funasr/models/language_model/rnn/attentions.py index 30cdff99d..f5c450c10 100644 --- a/funasr/models/language_model/rnn/attentions.py +++ b/funasr/models/language_model/rnn/attentions.py @@ -10,9 +10,7 @@ from funasr.models.transformer.utils.nets_utils import make_pad_mask from funasr.models.transformer.utils.nets_utils import to_device -def _apply_attention_constraint( - e, last_attended_idx, backward_window=1, forward_window=3 -): +def _apply_attention_constraint(e, last_attended_idx, backward_window=1, forward_window=3): """Apply monotonic attention constraint. This function apply the monotonic attention constraint @@ -84,9 +82,7 @@ class NoAtt(torch.nn.Module): mask = 1.0 - make_pad_mask(enc_hs_len).float() att_prev = mask / mask.new(enc_hs_len).unsqueeze(-1) att_prev = att_prev.to(self.enc_h) - self.c = torch.sum( - self.enc_h * att_prev.view(batch, self.h_length, 1), dim=1 - ) + self.c = torch.sum(self.enc_h * att_prev.view(batch, self.h_length, 1), dim=1) return self.c, att_prev @@ -150,8 +146,7 @@ class AttDot(torch.nn.Module): dec_z = dec_z.view(batch, self.dunits) e = torch.sum( - self.pre_compute_enc_h - * torch.tanh(self.mlp_dec(dec_z)).view(batch, 1, self.att_dim), + self.pre_compute_enc_h * torch.tanh(self.mlp_dec(dec_z)).view(batch, 1, self.att_dim), dim=2, ) # utt x frame @@ -262,9 +257,7 @@ class AttLoc(torch.nn.Module): and not store pre_compute_enc_h """ - def __init__( - self, eprojs, dunits, att_dim, aconv_chans, aconv_filts, han_mode=False - ): + def __init__(self, eprojs, dunits, att_dim, aconv_chans, aconv_filts, han_mode=False): super(AttLoc, self).__init__() self.mlp_enc = torch.nn.Linear(eprojs, att_dim) self.mlp_dec = torch.nn.Linear(dunits, att_dim, bias=False) @@ -338,9 +331,7 @@ class AttLoc(torch.nn.Module): # initialize attention weight with uniform dist. if att_prev is None: # if no bias, 0 0-pad goes 0 - att_prev = 1.0 - make_pad_mask(enc_hs_len).to( - device=dec_z.device, dtype=dec_z.dtype - ) + att_prev = 1.0 - make_pad_mask(enc_hs_len).to(device=dec_z.device, dtype=dec_z.dtype) att_prev = att_prev / att_prev.new(enc_hs_len).unsqueeze(-1) # att_prev: utt x frame -> utt x 1 x 1 x frame @@ -356,9 +347,7 @@ class AttLoc(torch.nn.Module): # dot with gvec # utt x frame x att_dim -> utt x frame - e = self.gvec( - torch.tanh(att_conv + self.pre_compute_enc_h + dec_z_tiled) - ).squeeze(2) + e = self.gvec(torch.tanh(att_conv + self.pre_compute_enc_h + dec_z_tiled)).squeeze(2) # NOTE: consider zero padding when compute w. if self.mask is None: @@ -367,9 +356,7 @@ class AttLoc(torch.nn.Module): # apply monotonic attention constraint (mainly for TTS) if last_attended_idx is not None: - e = _apply_attention_constraint( - e, last_attended_idx, backward_window, forward_window - ) + e = _apply_attention_constraint(e, last_attended_idx, backward_window, forward_window) w = F.softmax(scaling * e, dim=1) @@ -446,12 +433,8 @@ class AttCov(torch.nn.Module): # initialize attention weight with uniform dist. if att_prev_list is None: # if no bias, 0 0-pad goes 0 - att_prev_list = to_device( - enc_hs_pad, (1.0 - make_pad_mask(enc_hs_len).float()) - ) - att_prev_list = [ - att_prev_list / att_prev_list.new(enc_hs_len).unsqueeze(-1) - ] + att_prev_list = to_device(enc_hs_pad, (1.0 - make_pad_mask(enc_hs_len).float())) + att_prev_list = [att_prev_list / att_prev_list.new(enc_hs_len).unsqueeze(-1)] # att_prev_list: L' * [B x T] => cov_vec B x T cov_vec = sum(att_prev_list) @@ -463,9 +446,7 @@ class AttCov(torch.nn.Module): # dot with gvec # utt x frame x att_dim -> utt x frame - e = self.gvec( - torch.tanh(cov_vec + self.pre_compute_enc_h + dec_z_tiled) - ).squeeze(2) + e = self.gvec(torch.tanh(cov_vec + self.pre_compute_enc_h + dec_z_tiled)).squeeze(2) # NOTE consider zero padding when compute w. if self.mask is None: @@ -499,9 +480,7 @@ class AttLoc2D(torch.nn.Module): flag to swith on mode of hierarchical attention and not store pre_compute_enc_h """ - def __init__( - self, eprojs, dunits, att_dim, att_win, aconv_chans, aconv_filts, han_mode=False - ): + def __init__(self, eprojs, dunits, att_dim, att_win, aconv_chans, aconv_filts, han_mode=False): super(AttLoc2D, self).__init__() self.mlp_enc = torch.nn.Linear(eprojs, att_dim) self.mlp_dec = torch.nn.Linear(dunits, att_dim, bias=False) @@ -580,9 +559,7 @@ class AttLoc2D(torch.nn.Module): # dot with gvec # utt x frame x att_dim -> utt x frame - e = self.gvec( - torch.tanh(att_conv + self.pre_compute_enc_h + dec_z_tiled) - ).squeeze(2) + e = self.gvec(torch.tanh(att_conv + self.pre_compute_enc_h + dec_z_tiled)).squeeze(2) # NOTE consider zero padding when compute w. if self.mask is None: @@ -619,9 +596,7 @@ class AttLocRec(torch.nn.Module): flag to swith on mode of hierarchical attention and not store pre_compute_enc_h """ - def __init__( - self, eprojs, dunits, att_dim, aconv_chans, aconv_filts, han_mode=False - ): + def __init__(self, eprojs, dunits, att_dim, aconv_chans, aconv_filts, han_mode=False): super(AttLocRec, self).__init__() self.mlp_enc = torch.nn.Linear(eprojs, att_dim) self.mlp_dec = torch.nn.Linear(dunits, att_dim, bias=False) @@ -740,9 +715,7 @@ class AttCovLoc(torch.nn.Module): flag to swith on mode of hierarchical attention and not store pre_compute_enc_h """ - def __init__( - self, eprojs, dunits, att_dim, aconv_chans, aconv_filts, han_mode=False - ): + def __init__(self, eprojs, dunits, att_dim, aconv_chans, aconv_filts, han_mode=False): super(AttCovLoc, self).__init__() self.mlp_enc = torch.nn.Linear(eprojs, att_dim) self.mlp_dec = torch.nn.Linear(dunits, att_dim, bias=False) @@ -804,9 +777,7 @@ class AttCovLoc(torch.nn.Module): if att_prev_list is None: # if no bias, 0 0-pad goes 0 mask = 1.0 - make_pad_mask(enc_hs_len).float() - att_prev_list = [ - to_device(enc_hs_pad, mask / mask.new(enc_hs_len).unsqueeze(-1)) - ] + att_prev_list = [to_device(enc_hs_pad, mask / mask.new(enc_hs_len).unsqueeze(-1))] # att_prev_list: L' * [B x T] => cov_vec B x T cov_vec = sum(att_prev_list) @@ -823,9 +794,7 @@ class AttCovLoc(torch.nn.Module): # dot with gvec # utt x frame x att_dim -> utt x frame - e = self.gvec( - torch.tanh(att_conv + self.pre_compute_enc_h + dec_z_tiled) - ).squeeze(2) + e = self.gvec(torch.tanh(att_conv + self.pre_compute_enc_h + dec_z_tiled)).squeeze(2) # NOTE consider zero padding when compute w. if self.mask is None: @@ -908,17 +877,14 @@ class AttMultiHeadDot(torch.nn.Module): self.h_length = self.enc_h.size(1) # utt x frame x att_dim self.pre_compute_k = [ - torch.tanh(self.mlp_k[h](self.enc_h)) - for h in six.moves.range(self.aheads) + torch.tanh(self.mlp_k[h](self.enc_h)) for h in six.moves.range(self.aheads) ] if self.pre_compute_v is None or self.han_mode: self.enc_h = enc_hs_pad # utt x frame x hdim self.h_length = self.enc_h.size(1) # utt x frame x att_dim - self.pre_compute_v = [ - self.mlp_v[h](self.enc_h) for h in six.moves.range(self.aheads) - ] + self.pre_compute_v = [self.mlp_v[h](self.enc_h) for h in six.moves.range(self.aheads)] if dec_z is None: dec_z = enc_hs_pad.new_zeros(batch, self.dunits) @@ -943,11 +909,7 @@ class AttMultiHeadDot(torch.nn.Module): # weighted sum over flames # utt x hdim # NOTE use bmm instead of sum(*) - c += [ - torch.sum( - self.pre_compute_v[h] * w[h].view(batch, self.h_length, 1), dim=1 - ) - ] + c += [torch.sum(self.pre_compute_v[h] * w[h].view(batch, self.h_length, 1), dim=1)] # concat all of c c = self.mlp_o(torch.cat(c, dim=1)) @@ -1024,17 +986,13 @@ class AttMultiHeadAdd(torch.nn.Module): self.enc_h = enc_hs_pad # utt x frame x hdim self.h_length = self.enc_h.size(1) # utt x frame x att_dim - self.pre_compute_k = [ - self.mlp_k[h](self.enc_h) for h in six.moves.range(self.aheads) - ] + self.pre_compute_k = [self.mlp_k[h](self.enc_h) for h in six.moves.range(self.aheads)] if self.pre_compute_v is None or self.han_mode: self.enc_h = enc_hs_pad # utt x frame x hdim self.h_length = self.enc_h.size(1) # utt x frame x att_dim - self.pre_compute_v = [ - self.mlp_v[h](self.enc_h) for h in six.moves.range(self.aheads) - ] + self.pre_compute_v = [self.mlp_v[h](self.enc_h) for h in six.moves.range(self.aheads)] if dec_z is None: dec_z = enc_hs_pad.new_zeros(batch, self.dunits) @@ -1046,8 +1004,7 @@ class AttMultiHeadAdd(torch.nn.Module): for h in six.moves.range(self.aheads): e = self.gvec[h]( torch.tanh( - self.pre_compute_k[h] - + self.mlp_q[h](dec_z).view(batch, 1, self.att_dim_k) + self.pre_compute_k[h] + self.mlp_q[h](dec_z).view(batch, 1, self.att_dim_k) ) ).squeeze(2) @@ -1060,11 +1017,7 @@ class AttMultiHeadAdd(torch.nn.Module): # weighted sum over flames # utt x hdim # NOTE use bmm instead of sum(*) - c += [ - torch.sum( - self.pre_compute_v[h] * w[h].view(batch, self.h_length, 1), dim=1 - ) - ] + c += [torch.sum(self.pre_compute_v[h] * w[h].view(batch, self.h_length, 1), dim=1)] # concat all of c c = self.mlp_o(torch.cat(c, dim=1)) @@ -1167,17 +1120,13 @@ class AttMultiHeadLoc(torch.nn.Module): self.enc_h = enc_hs_pad # utt x frame x hdim self.h_length = self.enc_h.size(1) # utt x frame x att_dim - self.pre_compute_k = [ - self.mlp_k[h](self.enc_h) for h in six.moves.range(self.aheads) - ] + self.pre_compute_k = [self.mlp_k[h](self.enc_h) for h in six.moves.range(self.aheads)] if self.pre_compute_v is None or self.han_mode: self.enc_h = enc_hs_pad # utt x frame x hdim self.h_length = self.enc_h.size(1) # utt x frame x att_dim - self.pre_compute_v = [ - self.mlp_v[h](self.enc_h) for h in six.moves.range(self.aheads) - ] + self.pre_compute_v = [self.mlp_v[h](self.enc_h) for h in six.moves.range(self.aheads)] if dec_z is None: dec_z = enc_hs_pad.new_zeros(batch, self.dunits) @@ -1189,9 +1138,7 @@ class AttMultiHeadLoc(torch.nn.Module): for _ in six.moves.range(self.aheads): # if no bias, 0 0-pad goes 0 mask = 1.0 - make_pad_mask(enc_hs_len).float() - att_prev += [ - to_device(enc_hs_pad, mask / mask.new(enc_hs_len).unsqueeze(-1)) - ] + att_prev += [to_device(enc_hs_pad, mask / mask.new(enc_hs_len).unsqueeze(-1))] c = [] w = [] @@ -1217,11 +1164,7 @@ class AttMultiHeadLoc(torch.nn.Module): # weighted sum over flames # utt x hdim # NOTE use bmm instead of sum(*) - c += [ - torch.sum( - self.pre_compute_v[h] * w[h].view(batch, self.h_length, 1), dim=1 - ) - ] + c += [torch.sum(self.pre_compute_v[h] * w[h].view(batch, self.h_length, 1), dim=1)] # concat all of c c = self.mlp_o(torch.cat(c, dim=1)) @@ -1323,17 +1266,13 @@ class AttMultiHeadMultiResLoc(torch.nn.Module): self.enc_h = enc_hs_pad # utt x frame x hdim self.h_length = self.enc_h.size(1) # utt x frame x att_dim - self.pre_compute_k = [ - self.mlp_k[h](self.enc_h) for h in six.moves.range(self.aheads) - ] + self.pre_compute_k = [self.mlp_k[h](self.enc_h) for h in six.moves.range(self.aheads)] if self.pre_compute_v is None or self.han_mode: self.enc_h = enc_hs_pad # utt x frame x hdim self.h_length = self.enc_h.size(1) # utt x frame x att_dim - self.pre_compute_v = [ - self.mlp_v[h](self.enc_h) for h in six.moves.range(self.aheads) - ] + self.pre_compute_v = [self.mlp_v[h](self.enc_h) for h in six.moves.range(self.aheads)] if dec_z is None: dec_z = enc_hs_pad.new_zeros(batch, self.dunits) @@ -1345,9 +1284,7 @@ class AttMultiHeadMultiResLoc(torch.nn.Module): for _ in six.moves.range(self.aheads): # if no bias, 0 0-pad goes 0 mask = 1.0 - make_pad_mask(enc_hs_len).float() - att_prev += [ - to_device(enc_hs_pad, mask / mask.new(enc_hs_len).unsqueeze(-1)) - ] + att_prev += [to_device(enc_hs_pad, mask / mask.new(enc_hs_len).unsqueeze(-1))] c = [] w = [] @@ -1373,11 +1310,7 @@ class AttMultiHeadMultiResLoc(torch.nn.Module): # weighted sum over flames # utt x hdim # NOTE use bmm instead of sum(*) - c += [ - torch.sum( - self.pre_compute_v[h] * w[h].view(batch, self.h_length, 1), dim=1 - ) - ] + c += [torch.sum(self.pre_compute_v[h] * w[h].view(batch, self.h_length, 1), dim=1)] # concat all of c c = self.mlp_o(torch.cat(c, dim=1)) @@ -1484,9 +1417,7 @@ class AttForward(torch.nn.Module): # dot with gvec # utt x frame x att_dim -> utt x frame - e = self.gvec( - torch.tanh(self.pre_compute_enc_h + dec_z_tiled + att_conv) - ).squeeze(2) + e = self.gvec(torch.tanh(self.pre_compute_enc_h + dec_z_tiled + att_conv)).squeeze(2) # NOTE: consider zero padding when compute w. if self.mask is None: @@ -1495,9 +1426,7 @@ class AttForward(torch.nn.Module): # apply monotonic attention constraint (mainly for TTS) if last_attended_idx is not None: - e = _apply_attention_constraint( - e, last_attended_idx, backward_window, forward_window - ) + e = _apply_attention_constraint(e, last_attended_idx, backward_window, forward_window) w = F.softmax(scaling * e, dim=1) @@ -1619,9 +1548,7 @@ class AttForwardTA(torch.nn.Module): # dot with gvec # utt x frame x att_dim -> utt x frame - e = self.gvec( - torch.tanh(att_conv + self.pre_compute_enc_h + dec_z_tiled) - ).squeeze(2) + e = self.gvec(torch.tanh(att_conv + self.pre_compute_enc_h + dec_z_tiled)).squeeze(2) # NOTE consider zero padding when compute w. if self.mask is None: @@ -1630,18 +1557,13 @@ class AttForwardTA(torch.nn.Module): # apply monotonic attention constraint (mainly for TTS) if last_attended_idx is not None: - e = _apply_attention_constraint( - e, last_attended_idx, backward_window, forward_window - ) + e = _apply_attention_constraint(e, last_attended_idx, backward_window, forward_window) w = F.softmax(scaling * e, dim=1) # forward attention att_prev_shift = F.pad(att_prev, (1, 0))[:, :-1] - w = ( - self.trans_agent_prob * att_prev - + (1 - self.trans_agent_prob) * att_prev_shift - ) * w + w = (self.trans_agent_prob * att_prev + (1 - self.trans_agent_prob) * att_prev_shift) * w # NOTE: clamp is needed to avoid nan gradient w = F.normalize(torch.clamp(w, 1e-6), p=1, dim=1) @@ -1651,9 +1573,7 @@ class AttForwardTA(torch.nn.Module): c = torch.sum(self.enc_h * w.view(batch, self.h_length, 1), dim=1) # update transition agent prob - self.trans_agent_prob = torch.sigmoid( - self.mlp_ta(torch.cat([c, out_prev, dec_z], dim=1)) - ) + self.trans_agent_prob = torch.sigmoid(self.mlp_ta(torch.cat([c, out_prev, dec_z], dim=1))) return c, w @@ -1717,9 +1637,7 @@ def att_for(args, num_att=1, han_mode=False): ) att_list.append(att) else: - raise ValueError( - "Number of encoders needs to be more than one. {}".format(num_encs) - ) + raise ValueError("Number of encoders needs to be more than one. {}".format(num_encs)) return att_list @@ -1785,9 +1703,7 @@ def att_to_numpy(att_ws, att): att_ws = torch.stack([aw[:, -1] for aw in att_ws], dim=1).cpu().numpy() elif isinstance(att, (AttCov, AttCovLoc)): # att_ws => list of list of previous attentions - att_ws = ( - torch.stack([aw[idx] for idx, aw in enumerate(att_ws)], dim=1).cpu().numpy() - ) + att_ws = torch.stack([aw[idx] for idx, aw in enumerate(att_ws)], dim=1).cpu().numpy() elif isinstance(att, AttLocRec): # att_ws => list of tuple of attention and hidden states att_ws = torch.stack([aw[0] for aw in att_ws], dim=1).cpu().numpy() diff --git a/funasr/models/language_model/rnn/decoders.py b/funasr/models/language_model/rnn/decoders.py index fd8bc2902..e7d35e92f 100644 --- a/funasr/models/language_model/rnn/decoders.py +++ b/funasr/models/language_model/rnn/decoders.py @@ -1,4 +1,5 @@ """RNN decoder module.""" + import logging import math import random @@ -45,24 +46,24 @@ class Decoder(torch.nn.Module, ScorerInterface): """ def __init__( - self, - eprojs, - odim, - dtype, - dlayers, - dunits, - sos, - eos, - att, - verbose=0, - char_list=None, - labeldist=None, - lsm_weight=0.0, - sampling_probability=0.0, - dropout=0.0, - context_residual=False, - replace_sos=False, - num_encs=1, + self, + eprojs, + odim, + dtype, + dlayers, + dunits, + sos, + eos, + att, + verbose=0, + char_list=None, + labeldist=None, + lsm_weight=0.0, + sampling_probability=0.0, + dropout=0.0, + context_residual=False, + replace_sos=False, + num_encs=1, ): torch.nn.Module.__init__(self) @@ -76,16 +77,20 @@ class Decoder(torch.nn.Module, ScorerInterface): self.decoder = torch.nn.ModuleList() self.dropout_dec = torch.nn.ModuleList() self.decoder += [ - torch.nn.LSTMCell(dunits + eprojs, dunits) - if self.dtype == "lstm" - else torch.nn.GRUCell(dunits + eprojs, dunits) + ( + torch.nn.LSTMCell(dunits + eprojs, dunits) + if self.dtype == "lstm" + else torch.nn.GRUCell(dunits + eprojs, dunits) + ) ] self.dropout_dec += [torch.nn.Dropout(p=dropout)] for _ in six.moves.range(1, self.dlayers): self.decoder += [ - torch.nn.LSTMCell(dunits, dunits) - if self.dtype == "lstm" - else torch.nn.GRUCell(dunits, dunits) + ( + torch.nn.LSTMCell(dunits, dunits) + if self.dtype == "lstm" + else torch.nn.GRUCell(dunits, dunits) + ) ] self.dropout_dec += [torch.nn.Dropout(p=dropout)] # NOTE: dropout is applied only for the vertical connections @@ -131,9 +136,7 @@ class Decoder(torch.nn.Module, ScorerInterface): else: z_list[0] = self.decoder[0](ey, z_prev[0]) for i in six.moves.range(1, self.dlayers): - z_list[i] = self.decoder[i]( - self.dropout_dec[i - 1](z_list[i - 1]), z_prev[i] - ) + z_list[i] = self.decoder[i](self.dropout_dec[i - 1](z_list[i - 1]), z_prev[i]) return z_list, c_list def forward(self, hs_pad, hlens, ys_pad, strm_idx=0, lang_ids=None): @@ -197,9 +200,7 @@ class Decoder(torch.nn.Module, ScorerInterface): ) ) logging.info( - self.__class__.__name__ - + " output lengths: " - + str([y.size(0) for y in ys_out]) + self.__class__.__name__ + " output lengths: " + str([y.size(0) for y in ys_out]) ) # initialization @@ -280,7 +281,7 @@ class Decoder(torch.nn.Module, ScorerInterface): ys_hat = y_all.view(batch, olength, -1) ys_true = ys_out_pad for (i, y_hat), y_true in zip( - enumerate(ys_hat.detach().cpu().numpy()), ys_true.detach().cpu().numpy() + enumerate(ys_hat.detach().cpu().numpy()), ys_true.detach().cpu().numpy() ): if i == MAX_DECODER_OUTPUT: break @@ -362,9 +363,7 @@ class Decoder(torch.nn.Module, ScorerInterface): weights_ctc_dec = recog_args.weights_ctc_dec / np.sum( recog_args.weights_ctc_dec ) # normalize - logging.info( - "ctc weights (decoding): " + " ".join([str(x) for x in weights_ctc_dec]) - ) + logging.info("ctc weights (decoding): " + " ".join([str(x) for x in weights_ctc_dec])) else: weights_ctc_dec = [1.0] @@ -450,9 +449,7 @@ class Decoder(torch.nn.Module, ScorerInterface): hyp["a_prev"][self.num_encs], ) ey = torch.cat((ey, att_c), dim=1) # utt(1) x (zdim + hdim) - z_list, c_list = self.rnn_forward( - ey, z_list, c_list, hyp["z_prev"], hyp["c_prev"] - ) + z_list, c_list = self.rnn_forward(ey, z_list, c_list, hyp["z_prev"], hyp["c_prev"]) # get nbest local scores and their ids if self.context_residual: @@ -464,9 +461,7 @@ class Decoder(torch.nn.Module, ScorerInterface): local_att_scores = F.log_softmax(logits, dim=1) if rnnlm: rnnlm_state, local_lm_scores = rnnlm.predict(hyp["rnnlm_prev"], vy) - local_scores = ( - local_att_scores + recog_args.lm_weight * local_lm_scores - ) + local_scores = local_att_scores + recog_args.lm_weight * local_lm_scores else: local_scores = local_att_scores @@ -482,9 +477,7 @@ class Decoder(torch.nn.Module, ScorerInterface): ctc_scores[idx], ctc_states[idx] = ctc_prefix_score[idx]( hyp["yseq"], local_best_ids[0], hyp["ctc_state_prev"][idx] ) - local_scores = (1.0 - ctc_weight) * local_att_scores[ - :, local_best_ids[0] - ] + local_scores = (1.0 - ctc_weight) * local_att_scores[:, local_best_ids[0]] if self.num_encs == 1: local_scores += ctc_weight * torch.from_numpy( ctc_scores[0] - hyp["ctc_score_prev"][0] @@ -492,24 +485,16 @@ class Decoder(torch.nn.Module, ScorerInterface): else: for idx in range(self.num_encs): local_scores += ( - ctc_weight - * weights_ctc_dec[idx] - * torch.from_numpy( - ctc_scores[idx] - hyp["ctc_score_prev"][idx] - ) + ctc_weight + * weights_ctc_dec[idx] + * torch.from_numpy(ctc_scores[idx] - hyp["ctc_score_prev"][idx]) ) if rnnlm: - local_scores += ( - recog_args.lm_weight * local_lm_scores[:, local_best_ids[0]] - ) - local_best_scores, joint_best_ids = torch.topk( - local_scores, beam, dim=1 - ) + local_scores += recog_args.lm_weight * local_lm_scores[:, local_best_ids[0]] + local_best_scores, joint_best_ids = torch.topk(local_scores, beam, dim=1) local_best_ids = local_best_ids[:, joint_best_ids[0]] else: - local_best_scores, local_best_ids = torch.topk( - local_scores, beam, dim=1 - ) + local_best_scores, local_best_ids = torch.topk(local_scores, beam, dim=1) for j in six.moves.range(beam): new_hyp = {} @@ -519,9 +504,7 @@ class Decoder(torch.nn.Module, ScorerInterface): if self.num_encs == 1: new_hyp["a_prev"] = att_w[:] else: - new_hyp["a_prev"] = [ - att_w_list[idx][:] for idx in range(self.num_encs + 1) - ] + new_hyp["a_prev"] = [att_w_list[idx][:] for idx in range(self.num_encs + 1)] new_hyp["score"] = hyp["score"] + local_best_scores[0, j] new_hyp["yseq"] = [0] * (1 + len(hyp["yseq"])) new_hyp["yseq"][: len(hyp["yseq"])] = hyp["yseq"] @@ -530,27 +513,22 @@ class Decoder(torch.nn.Module, ScorerInterface): new_hyp["rnnlm_prev"] = rnnlm_state if lpz[0] is not None: new_hyp["ctc_state_prev"] = [ - ctc_states[idx][joint_best_ids[0, j]] - for idx in range(self.num_encs) + ctc_states[idx][joint_best_ids[0, j]] for idx in range(self.num_encs) ] new_hyp["ctc_score_prev"] = [ - ctc_scores[idx][joint_best_ids[0, j]] - for idx in range(self.num_encs) + ctc_scores[idx][joint_best_ids[0, j]] for idx in range(self.num_encs) ] # will be (2 x beam) hyps at most hyps_best_kept.append(new_hyp) - hyps_best_kept = sorted( - hyps_best_kept, key=lambda x: x["score"], reverse=True - )[:beam] + hyps_best_kept = sorted(hyps_best_kept, key=lambda x: x["score"], reverse=True)[ + :beam + ] # sort and get nbest hyps = hyps_best_kept logging.debug("number of pruned hypotheses: " + str(len(hyps))) - logging.debug( - "best hypo: " - + "".join([char_list[int(x)] for x in hyps[0]["yseq"][1:]]) - ) + logging.debug("best hypo: " + "".join([char_list[int(x)] for x in hyps[0]["yseq"][1:]])) # add eos in the final loop to avoid that there are no ended hyps if i == maxlen - 1: @@ -569,9 +547,7 @@ class Decoder(torch.nn.Module, ScorerInterface): if len(hyp["yseq"]) > minlen: hyp["score"] += (i + 1) * penalty if rnnlm: # Word LM needs to add final score - hyp["score"] += recog_args.lm_weight * rnnlm.final( - hyp["rnnlm_prev"] - ) + hyp["score"] += recog_args.lm_weight * rnnlm.final(hyp["rnnlm_prev"]) ended_hyps.append(hyp) else: remained_hyps.append(hyp) @@ -589,21 +565,18 @@ class Decoder(torch.nn.Module, ScorerInterface): break for hyp in hyps: - logging.debug( - "hypo: " + "".join([char_list[int(x)] for x in hyp["yseq"][1:]]) - ) + logging.debug("hypo: " + "".join([char_list[int(x)] for x in hyp["yseq"][1:]])) logging.debug("number of ended hypotheses: " + str(len(ended_hyps))) nbest_hyps = sorted(ended_hyps, key=lambda x: x["score"], reverse=True)[ - : min(len(ended_hyps), recog_args.nbest) - ] + : min(len(ended_hyps), recog_args.nbest) + ] # check number of hypotheses if len(nbest_hyps) == 0: logging.warning( - "there is no N-best results, " - "perform recognition again with smaller minlenratio." + "there is no N-best results, " "perform recognition again with smaller minlenratio." ) # should copy because Namespace will be overwritten globally recog_args = Namespace(**vars(recog_args)) @@ -623,16 +596,16 @@ class Decoder(torch.nn.Module, ScorerInterface): return nbest_hyps def recognize_beam_batch( - self, - h, - hlens, - lpz, - recog_args, - char_list, - rnnlm=None, - normalize_score=True, - strm_idx=0, - lang_ids=None, + self, + h, + hlens, + lpz, + recog_args, + char_list, + rnnlm=None, + normalize_score=True, + strm_idx=0, + lang_ids=None, ): # to support mutiple encoder asr mode, in single encoder mode, # convert torch.Tensor to List of torch.Tensor @@ -667,9 +640,7 @@ class Decoder(torch.nn.Module, ScorerInterface): weights_ctc_dec = recog_args.weights_ctc_dec / np.sum( recog_args.weights_ctc_dec ) # normalize - logging.info( - "ctc weights (decoding): " + " ".join([str(x) for x in weights_ctc_dec]) - ) + logging.info("ctc weights (decoding): " + " ".join([str(x) for x in weights_ctc_dec])) else: weights_ctc_dec = [1.0] @@ -686,18 +657,10 @@ class Decoder(torch.nn.Module, ScorerInterface): logging.info("min output length: " + str(minlen)) # initialization - c_prev = [ - to_device(h[0], torch.zeros(n_bb, self.dunits)) for _ in range(self.dlayers) - ] - z_prev = [ - to_device(h[0], torch.zeros(n_bb, self.dunits)) for _ in range(self.dlayers) - ] - c_list = [ - to_device(h[0], torch.zeros(n_bb, self.dunits)) for _ in range(self.dlayers) - ] - z_list = [ - to_device(h[0], torch.zeros(n_bb, self.dunits)) for _ in range(self.dlayers) - ] + c_prev = [to_device(h[0], torch.zeros(n_bb, self.dunits)) for _ in range(self.dlayers)] + z_prev = [to_device(h[0], torch.zeros(n_bb, self.dunits)) for _ in range(self.dlayers)] + c_list = [to_device(h[0], torch.zeros(n_bb, self.dunits)) for _ in range(self.dlayers)] + z_list = [to_device(h[0], torch.zeros(n_bb, self.dunits)) for _ in range(self.dlayers)] vscores = to_device(h[0], torch.zeros(batch, beam)) rnnlm_state = None @@ -716,14 +679,10 @@ class Decoder(torch.nn.Module, ScorerInterface): if self.replace_sos and recog_args.tgt_lang: logging.info(" index: " + str(char_list.index(recog_args.tgt_lang))) logging.info(" mark: " + recog_args.tgt_lang) - yseq = [ - [char_list.index(recog_args.tgt_lang)] for _ in six.moves.range(n_bb) - ] + yseq = [[char_list.index(recog_args.tgt_lang)] for _ in six.moves.range(n_bb)] elif lang_ids is not None: # NOTE: used for evaluation during training - yseq = [ - [lang_ids[b // recog_args.beam_size]] for b in six.moves.range(n_bb) - ] + yseq = [[lang_ids[b // recog_args.beam_size]] for b in six.moves.range(n_bb)] else: logging.info(" index: " + str(self.sos)) logging.info(" mark: " + char_list[self.sos]) @@ -740,8 +699,7 @@ class Decoder(torch.nn.Module, ScorerInterface): ] exp_hlens = [exp_hlens[idx].view(-1).tolist() for idx in range(self.num_encs)] exp_h = [ - h[idx].unsqueeze(1).repeat(1, beam, 1, 1).contiguous() - for idx in range(self.num_encs) + h[idx].unsqueeze(1).repeat(1, beam, 1, 1).contiguous() for idx in range(self.num_encs) ] exp_h = [ exp_h[idx].view(n_bb, h[idx].size()[1], h[idx].size()[2]) @@ -750,9 +708,7 @@ class Decoder(torch.nn.Module, ScorerInterface): if lpz[0] is not None: scoring_num = min( - int(beam * CTC_SCORING_RATIO) - if att_weight > 0.0 and not lpz[0].is_cuda - else 0, + int(beam * CTC_SCORING_RATIO) if att_weight > 0.0 and not lpz[0].is_cuda else 0, lpz[0].size(-1), ) ctc_scorer = [ @@ -796,9 +752,7 @@ class Decoder(torch.nn.Module, ScorerInterface): # attention decoder z_list, c_list = self.rnn_forward(ey, z_list, c_list, z_prev, c_prev) if self.context_residual: - logits = self.output( - torch.cat((self.dropout_dec[-1](z_list[-1]), att_c), dim=-1) - ) + logits = self.output(torch.cat((self.dropout_dec[-1](z_list[-1]), att_c), dim=-1)) else: logits = self.output(self.dropout_dec[-1](z_list[-1])) local_scores = att_weight * F.log_softmax(logits, dim=1) @@ -812,9 +766,7 @@ class Decoder(torch.nn.Module, ScorerInterface): if ctc_scorer[0]: local_scores[:, 0] = self.logzero # avoid choosing blank part_ids = ( - torch.topk(local_scores, scoring_num, dim=-1)[1] - if scoring_num > 0 - else None + torch.topk(local_scores, scoring_num, dim=-1)[1] if scoring_num > 0 else None ) for idx in range(self.num_encs): att_w = att_w_list[idx] @@ -823,8 +775,7 @@ class Decoder(torch.nn.Module, ScorerInterface): yseq, ctc_state[idx], part_ids, att_w_ ) local_scores = ( - local_scores - + ctc_weight * weights_ctc_dec[idx] * local_ctc_scores + local_scores + ctc_weight * weights_ctc_dec[idx] * local_ctc_scores ) local_scores = local_scores.view(batch, beam, self.odim) @@ -839,9 +790,7 @@ class Decoder(torch.nn.Module, ScorerInterface): # global pruning accum_best_scores, accum_best_ids = torch.topk(vscores, beam, 1) - accum_odim_ids = ( - torch.fmod(accum_best_ids, self.odim).view(-1).data.cpu().tolist() - ) + accum_odim_ids = torch.fmod(accum_best_ids, self.odim).view(-1).data.cpu().tolist() accum_padded_beam_ids = ( (accum_best_ids // self.odim + pad_b).view(-1).data.cpu().tolist() ) @@ -867,24 +816,16 @@ class Decoder(torch.nn.Module, ScorerInterface): ] else: # handle the case of location_recurrent when return is a tuple - _a_prev_ = torch.index_select( - att_w_list[idx][0].view(n_bb, -1), 0, vidx - ) - _h_prev_ = torch.index_select( - att_w_list[idx][1][0].view(n_bb, -1), 0, vidx - ) - _c_prev_ = torch.index_select( - att_w_list[idx][1][1].view(n_bb, -1), 0, vidx - ) + _a_prev_ = torch.index_select(att_w_list[idx][0].view(n_bb, -1), 0, vidx) + _h_prev_ = torch.index_select(att_w_list[idx][1][0].view(n_bb, -1), 0, vidx) + _c_prev_ = torch.index_select(att_w_list[idx][1][1].view(n_bb, -1), 0, vidx) _a_prev = (_a_prev_, (_h_prev_, _c_prev_)) a_prev.append(_a_prev) z_prev = [ - torch.index_select(z_list[li].view(n_bb, -1), 0, vidx) - for li in range(self.dlayers) + torch.index_select(z_list[li].view(n_bb, -1), 0, vidx) for li in range(self.dlayers) ] c_prev = [ - torch.index_select(c_list[li].view(n_bb, -1), 0, vidx) - for li in range(self.dlayers) + torch.index_select(c_list[li].view(n_bb, -1), 0, vidx) for li in range(self.dlayers) ] # pick ended hyps @@ -900,9 +841,7 @@ class Decoder(torch.nn.Module, ScorerInterface): _vscore = None if eos_vscores[samp_i, beam_j] > thr[samp_i]: yk = y_prev[k][:] - if len(yk) <= min( - hlens[idx][samp_i] for idx in range(self.num_encs) - ): + if len(yk) <= min(hlens[idx][samp_i] for idx in range(self.num_encs)): _vscore = eos_vscores[samp_i][beam_j] + penalty_i elif i == maxlen - 1: yk = yseq[k][:] @@ -910,9 +849,7 @@ class Decoder(torch.nn.Module, ScorerInterface): if _vscore: yk.append(self.eos) if rnnlm: - _vscore += recog_args.lm_weight * rnnlm.final( - rnnlm_state, index=k - ) + _vscore += recog_args.lm_weight * rnnlm.final(rnnlm_state, index=k) _score = _vscore.data.cpu().numpy() ended_hyps[samp_i].append( {"yseq": yk, "vscore": _vscore, "score": _score} @@ -938,9 +875,7 @@ class Decoder(torch.nn.Module, ScorerInterface): torch.cuda.empty_cache() - dummy_hyps = [ - {"yseq": [self.sos, self.eos], "score": np.array([-float("inf")])} - ] + dummy_hyps = [{"yseq": [self.sos, self.eos], "score": np.array([-float("inf")])}] ended_hyps = [ ended_hyps[samp_i] if len(ended_hyps[samp_i]) != 0 else dummy_hyps for samp_i in six.moves.range(batch) @@ -952,7 +887,7 @@ class Decoder(torch.nn.Module, ScorerInterface): nbest_hyps = [ sorted(ended_hyps[samp_i], key=lambda x: x["score"], reverse=True)[ - : min(len(ended_hyps[samp_i]), recog_args.nbest) + : min(len(ended_hyps[samp_i]), recog_args.nbest) ] for samp_i in six.moves.range(batch) ] @@ -1168,13 +1103,9 @@ class Decoder(torch.nn.Module, ScorerInterface): state["a_prev"][self.num_encs], ) ey = torch.cat((ey, att_c), dim=1) # utt(1) x (zdim + hdim) - z_list, c_list = self.rnn_forward( - ey, z_list, c_list, state["z_prev"], state["c_prev"] - ) + z_list, c_list = self.rnn_forward(ey, z_list, c_list, state["z_prev"], state["c_prev"]) if self.context_residual: - logits = self.output( - torch.cat((self.dropout_dec[-1](z_list[-1]), att_c), dim=-1) - ) + logits = self.output(torch.cat((self.dropout_dec[-1](z_list[-1]), att_c), dim=-1)) else: logits = self.output(self.dropout_dec[-1](z_list[-1])) logp = F.log_softmax(logits, dim=1).squeeze(0) diff --git a/funasr/models/language_model/rnn/encoders.py b/funasr/models/language_model/rnn/encoders.py index d047a8ed2..fff376650 100644 --- a/funasr/models/language_model/rnn/encoders.py +++ b/funasr/models/language_model/rnn/encoders.py @@ -34,9 +34,7 @@ class RNNP(torch.nn.Module): inputdim = hdim RNN = torch.nn.LSTM if "lstm" in typ else torch.nn.GRU - rnn = RNN( - inputdim, cdim, num_layers=1, bidirectional=bidir, batch_first=True - ) + rnn = RNN(inputdim, cdim, num_layers=1, bidirectional=bidir, batch_first=True) setattr(self, "%s%d" % ("birnn" if bidir else "rnn", i), rnn) @@ -72,9 +70,7 @@ class RNNP(torch.nn.Module): rnn.flatten_parameters() if prev_state is not None and rnn.bidirectional: prev_state = reset_backward_rnn_state(prev_state) - ys, states = rnn( - xs_pack, hx=None if prev_state is None else prev_state[layer] - ) + ys, states = rnn(xs_pack, hx=None if prev_state is None else prev_state[layer]) elayer_states.append(states) # ys: utt list of frame x cdim x 2 (2: means bidirectional) ys_pad, ilens = pad_packed_sequence(ys, batch_first=True) @@ -155,9 +151,7 @@ class RNN(torch.nn.Module): # ys: utt list of frame x cdim x 2 (2: means bidirectional) ys_pad, ilens = pad_packed_sequence(ys, batch_first=True) # (sum _utt frame_utt) x dim - projected = torch.tanh( - self.l_last(ys_pad.contiguous().view(-1, ys_pad.size(2))) - ) + projected = torch.tanh(self.l_last(ys_pad.contiguous().view(-1, ys_pad.size(2)))) xs_pad = projected.view(ys_pad.size(0), ys_pad.size(1), -1) return xs_pad, ilens, states # x: utt list of frame x dim @@ -225,9 +219,7 @@ class VGG2L(torch.nn.Module): else: ilens = np.array(ilens, dtype=np.float32) ilens = np.array(np.ceil(ilens / 2), dtype=np.int64) - ilens = np.array( - np.ceil(np.array(ilens, dtype=np.float32) / 2), dtype=np.int64 - ).tolist() + ilens = np.array(np.ceil(np.array(ilens, dtype=np.float32) / 2), dtype=np.int64).tolist() # x: utt_list of frame (remove zeropaded frames) x (input channel num x dim) xs_pad = xs_pad.transpose(1, 2) @@ -250,9 +242,7 @@ class Encoder(torch.nn.Module): :param int in_channel: number of input channels """ - def __init__( - self, etype, idim, elayers, eunits, eprojs, subsample, dropout, in_channel=1 - ): + def __init__(self, etype, idim, elayers, eunits, eprojs, subsample, dropout, in_channel=1): super(Encoder, self).__init__() typ = etype.lstrip("vgg").rstrip("p") if typ not in ["lstm", "gru", "blstm", "bgru"]: @@ -367,6 +357,4 @@ def encoder_for(args, idim, subsample): enc_list.append(enc) return enc_list else: - raise ValueError( - "Number of encoders needs to be more than one. {}".format(num_encs) - ) + raise ValueError("Number of encoders needs to be more than one. {}".format(num_encs)) diff --git a/funasr/models/language_model/seq_rnn_lm.py b/funasr/models/language_model/seq_rnn_lm.py index bef497477..baf63bc30 100644 --- a/funasr/models/language_model/seq_rnn_lm.py +++ b/funasr/models/language_model/seq_rnn_lm.py @@ -1,4 +1,5 @@ """Sequential implementation of Recurrent Neural Network Language Model.""" + from typing import Tuple from typing import Union @@ -37,9 +38,7 @@ class SequentialRNNLM(AbsLM): self.encoder = nn.Embedding(vocab_size, ninp, padding_idx=ignore_id) if rnn_type in ["LSTM", "GRU"]: rnn_class = getattr(nn, rnn_type) - self.rnn = rnn_class( - ninp, nhid, nlayers, dropout=dropout_rate, batch_first=True - ) + self.rnn = rnn_class(ninp, nhid, nlayers, dropout=dropout_rate, batch_first=True) else: try: nonlinearity = {"RNN_TANH": "tanh", "RNN_RELU": "relu"}[rnn_type] @@ -67,9 +66,7 @@ class SequentialRNNLM(AbsLM): # https://arxiv.org/abs/1611.01462 if tie_weights: if nhid != ninp: - raise ValueError( - "When using the tied flag, nhid must be equal to emsize" - ) + raise ValueError("When using the tied flag, nhid must be equal to emsize") self.decoder.weight = self.encoder.weight self.rnn_type = rnn_type diff --git a/funasr/models/language_model/transformer_encoder.py b/funasr/models/language_model/transformer_encoder.py index 21f354844..7c22550a8 100644 --- a/funasr/models/language_model/transformer_encoder.py +++ b/funasr/models/language_model/transformer_encoder.py @@ -56,14 +56,14 @@ class EncoderLayer(nn.Module): """ def __init__( - self, - size, - self_attn, - feed_forward, - dropout_rate, - normalize_before=True, - concat_after=False, - stochastic_depth_rate=0.0, + self, + size, + self_attn, + feed_forward, + dropout_rate, + normalize_before=True, + concat_after=False, + stochastic_depth_rate=0.0, ): """Construct an EncoderLayer object.""" super(EncoderLayer, self).__init__() @@ -121,9 +121,7 @@ class EncoderLayer(nn.Module): x_concat = torch.cat((x, self.self_attn(x_q, x, x, mask)), dim=-1) x = residual + stoch_layer_coeff * self.concat_linear(x_concat) else: - x = residual + stoch_layer_coeff * self.dropout( - self.self_attn(x_q, x, x, mask) - ) + x = residual + stoch_layer_coeff * self.dropout(self.self_attn(x_q, x, x, mask)) if not self.normalize_before: x = self.norm1(x) @@ -245,9 +243,7 @@ class TransformerEncoder_lm(nn.Module): pos_enc_class(attention_dim, positional_dropout_rate), ) elif input_layer is None: - self.embed = torch.nn.Sequential( - pos_enc_class(attention_dim, positional_dropout_rate) - ) + self.embed = torch.nn.Sequential(pos_enc_class(attention_dim, positional_dropout_rate)) else: raise ValueError("unknown input_layer: " + input_layer) self.normalize_before = normalize_before @@ -288,8 +284,7 @@ class TransformerEncoder_lm(nn.Module): ] elif selfattention_layer_type == "lightconv2d": logging.info( - "encoder self-attention layer " - "type = lightweight convolution 2-dimensional" + "encoder self-attention layer " "type = lightweight convolution 2-dimensional" ) encoder_selfattn_layer = LightweightConvolution2D encoder_selfattn_layer_args = [ @@ -318,9 +313,7 @@ class TransformerEncoder_lm(nn.Module): for lnum in range(num_blocks) ] elif selfattention_layer_type == "dynamicconv2d": - logging.info( - "encoder self-attention layer type = dynamic convolution 2-dimensional" - ) + logging.info("encoder self-attention layer type = dynamic convolution 2-dimensional") encoder_selfattn_layer = DynamicConvolution2D encoder_selfattn_layer_args = [ ( @@ -355,9 +348,7 @@ class TransformerEncoder_lm(nn.Module): self.use_conditioning = True if ctc_softmax is not None else False if self.use_conditioning: self.ctc_softmax = ctc_softmax - self.conditioning_layer = torch.nn.Linear( - conditioning_layer_dim, attention_dim - ) + self.conditioning_layer = torch.nn.Linear(conditioning_layer_dim, attention_dim) def get_positionwise_layer( self, @@ -466,4 +457,3 @@ class TransformerEncoder_lm(nn.Module): if self.normalize_before: xs = self.after_norm(xs) return xs, masks, new_cache - diff --git a/funasr/models/language_model/transformer_lm.py b/funasr/models/language_model/transformer_lm.py index cc71142bb..3b90e3dc8 100644 --- a/funasr/models/language_model/transformer_lm.py +++ b/funasr/models/language_model/transformer_lm.py @@ -66,9 +66,7 @@ class TransformerLM(AbsLM): y = self.decoder(h) return y, None - def score( - self, y: torch.Tensor, state: Any, x: torch.Tensor - ) -> Tuple[torch.Tensor, Any]: + def score(self, y: torch.Tensor, state: Any, x: torch.Tensor) -> Tuple[torch.Tensor, Any]: """Score new token. Args: @@ -115,8 +113,7 @@ class TransformerLM(AbsLM): else: # transpose state of [batch, layer] into [layer, batch] batch_state = [ - torch.stack([states[b][i] for b in range(n_batch)]) - for i in range(n_layers) + torch.stack([states[b][i] for b in range(n_batch)]) for i in range(n_layers) ] # batch decoding diff --git a/funasr/models/lcbnet/attention.py b/funasr/models/lcbnet/attention.py index 8e8c5943a..05a5041f2 100644 --- a/funasr/models/lcbnet/attention.py +++ b/funasr/models/lcbnet/attention.py @@ -11,6 +11,7 @@ import math import torch from torch import nn + class MultiHeadedAttentionReturnWeight(nn.Module): """Multi-Head Attention layer. @@ -108,5 +109,3 @@ class MultiHeadedAttentionReturnWeight(nn.Module): q, k, v = self.forward_qkv(query, key, value) scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) return self.forward_attention(v, scores, mask) - - diff --git a/funasr/models/lcbnet/encoder.py b/funasr/models/lcbnet/encoder.py index c65823cb0..f5f249776 100644 --- a/funasr/models/lcbnet/encoder.py +++ b/funasr/models/lcbnet/encoder.py @@ -21,6 +21,7 @@ from funasr.models.transformer.positionwise_feed_forward import PositionwiseFeed from funasr.models.transformer.utils.repeat import repeat from funasr.register import tables + class EncoderLayer(nn.Module): """Encoder layer module. @@ -44,14 +45,14 @@ class EncoderLayer(nn.Module): """ def __init__( - self, - size, - self_attn, - feed_forward, - dropout_rate, - normalize_before=True, - concat_after=False, - stochastic_depth_rate=0.0, + self, + size, + self_attn, + feed_forward, + dropout_rate, + normalize_before=True, + concat_after=False, + stochastic_depth_rate=0.0, ): """Construct an EncoderLayer object.""" super(EncoderLayer, self).__init__() @@ -109,9 +110,7 @@ class EncoderLayer(nn.Module): x_concat = torch.cat((x, self.self_attn(x_q, x, x, mask)), dim=-1) x = residual + stoch_layer_coeff * self.concat_linear(x_concat) else: - x = residual + stoch_layer_coeff * self.dropout( - self.self_attn(x_q, x, x, mask) - ) + x = residual + stoch_layer_coeff * self.dropout(self.self_attn(x_q, x, x, mask)) if not self.normalize_before: x = self.norm1(x) @@ -127,6 +126,7 @@ class EncoderLayer(nn.Module): return x, mask + @tables.register("encoder_classes", "TransformerTextEncoder") class TransformerTextEncoder(nn.Module): """Transformer text encoder module. @@ -154,18 +154,18 @@ class TransformerTextEncoder(nn.Module): """ def __init__( - self, - input_size: int, - output_size: int = 256, - attention_heads: int = 4, - linear_units: int = 2048, - num_blocks: int = 6, - dropout_rate: float = 0.1, - positional_dropout_rate: float = 0.1, - attention_dropout_rate: float = 0.0, - pos_enc_class=PositionalEncoding, - normalize_before: bool = True, - concat_after: bool = False, + self, + input_size: int, + output_size: int = 256, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + attention_dropout_rate: float = 0.0, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, ): super().__init__() self._output_size = output_size @@ -187,9 +187,7 @@ class TransformerTextEncoder(nn.Module): num_blocks, lambda lnum: EncoderLayer( output_size, - MultiHeadedAttention( - attention_heads, output_size, attention_dropout_rate - ), + MultiHeadedAttention(attention_heads, output_size, attention_dropout_rate), positionwise_layer(*positionwise_layer_args), dropout_rate, normalize_before, @@ -203,9 +201,9 @@ class TransformerTextEncoder(nn.Module): return self._output_size def forward( - self, - xs_pad: torch.Tensor, - ilens: torch.Tensor, + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: """Embed positions in tensor. @@ -227,8 +225,6 @@ class TransformerTextEncoder(nn.Module): return xs_pad, olens, None - - @tables.register("encoder_classes", "FusionSANEncoder") class SelfSrcAttention(nn.Module): """Single decoder layer module. @@ -251,25 +247,32 @@ class SelfSrcAttention(nn.Module): """ + def __init__( - self, - size, - attention_heads, - attention_dim, - linear_units, - self_attention_dropout_rate, - src_attention_dropout_rate, - positional_dropout_rate, - dropout_rate, - normalize_before=True, - concat_after=False, + self, + size, + attention_heads, + attention_dim, + linear_units, + self_attention_dropout_rate, + src_attention_dropout_rate, + positional_dropout_rate, + dropout_rate, + normalize_before=True, + concat_after=False, ): """Construct an SelfSrcAttention object.""" super(SelfSrcAttention, self).__init__() self.size = size - self.self_attn = MultiHeadedAttention(attention_heads, attention_dim, self_attention_dropout_rate) - self.src_attn = MultiHeadedAttentionReturnWeight(attention_heads, attention_dim, src_attention_dropout_rate) - self.feed_forward = PositionwiseFeedForward(attention_dim, linear_units, positional_dropout_rate) + self.self_attn = MultiHeadedAttention( + attention_heads, attention_dim, self_attention_dropout_rate + ) + self.src_attn = MultiHeadedAttentionReturnWeight( + attention_heads, attention_dim, src_attention_dropout_rate + ) + self.feed_forward = PositionwiseFeedForward( + attention_dim, linear_units, positional_dropout_rate + ) self.norm1 = LayerNorm(size) self.norm2 = LayerNorm(size) self.norm3 = LayerNorm(size) @@ -319,9 +322,7 @@ class SelfSrcAttention(nn.Module): tgt_q_mask = tgt_mask[:, -1:, :] if self.concat_after: - tgt_concat = torch.cat( - (tgt_q, self.self_attn(tgt_q, tgt, tgt, tgt_q_mask)), dim=-1 - ) + tgt_concat = torch.cat((tgt_q, self.self_attn(tgt_q, tgt, tgt, tgt_q_mask)), dim=-1) x = residual + self.concat_linear1(tgt_concat) else: x = residual + self.dropout(self.self_attn(tgt_q, tgt, tgt, tgt_q_mask)) @@ -332,9 +333,7 @@ class SelfSrcAttention(nn.Module): if self.normalize_before: x = self.norm2(x) if self.concat_after: - x_concat = torch.cat( - (x, self.src_attn(x, memory, memory, memory_mask)), dim=-1 - ) + x_concat = torch.cat((x, self.src_attn(x, memory, memory, memory_mask)), dim=-1) x = residual + self.concat_linear2(x_concat) else: x, score = self.src_attn(x, memory, memory, memory_mask) @@ -357,7 +356,15 @@ class SelfSrcAttention(nn.Module): @tables.register("encoder_classes", "ConvBiasPredictor") class ConvPredictor(nn.Module): - def __init__(self, size=256, l_order=3, r_order=3, attention_heads=4, attention_dropout_rate=0.1, linear_units=2048): + def __init__( + self, + size=256, + l_order=3, + r_order=3, + attention_heads=4, + attention_dropout_rate=0.1, + linear_units=2048, + ): super().__init__() self.atten = MultiHeadedAttention(attention_heads, size, attention_dropout_rate) self.norm1 = LayerNorm(size) @@ -367,17 +374,16 @@ class ConvPredictor(nn.Module): self.conv1d = nn.Conv1d(size, size, l_order + r_order + 1, groups=size) self.output_linear = nn.Linear(size, 1) - def forward(self, text_enc, asr_enc): # stage1 cross-attention residual = text_enc text_enc = residual + self.atten(text_enc, asr_enc, asr_enc, None) - + # stage2 FFN residual = text_enc text_enc = self.norm1(text_enc) text_enc = residual + self.feed_forward(text_enc) - + # stage Conv predictor text_enc = self.norm2(text_enc) context = text_enc.transpose(1, 2) @@ -387,6 +393,6 @@ class ConvPredictor(nn.Module): output = output.transpose(1, 2) output = torch.relu(output) output = self.output_linear(output) - if output.dim()==3: - output = output.squeeze(2) + if output.dim() == 3: + output = output.squeeze(2) return output diff --git a/funasr/models/lcbnet/model.py b/funasr/models/lcbnet/model.py index 3ac319c61..d3df25a91 100644 --- a/funasr/models/lcbnet/model.py +++ b/funasr/models/lcbnet/model.py @@ -15,6 +15,7 @@ from funasr.losses.label_smoothing_loss import LabelSmoothingLoss from funasr.models.ctc.ctc import CTC from funasr.models.transformer.utils.add_sos_eos import add_sos_eos from funasr.metrics.compute_acc import th_accuracy + # from funasr.models.e2e_asr_common import ErrorCalculator from funasr.train_utils.device_funcs import force_gatherable from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank @@ -23,6 +24,8 @@ from funasr.utils.datadir_writer import DatadirWriter from funasr.register import tables import pdb + + @tables.register("model_classes", "LCBNet") class LCBNet(nn.Module): """ @@ -30,7 +33,7 @@ class LCBNet(nn.Module): LCB-NET: LONG-CONTEXT BIASING FOR AUDIO-VISUAL SPEECH RECOGNITION https://arxiv.org/abs/2401.06390 """ - + def __init__( self, specaug: str = None, @@ -93,7 +96,6 @@ class LCBNet(nn.Module): bias_predictor_class = tables.encoder_classes.get(bias_predictor) bias_predictor = bias_predictor_class(**bias_predictor_conf) - if decoder is not None: decoder_class = tables.decoder_classes.get(decoder) decoder = decoder_class( @@ -102,14 +104,12 @@ class LCBNet(nn.Module): **decoder_conf, ) if ctc_weight > 0.0: - + if ctc_conf is None: ctc_conf = {} - - ctc = CTC( - odim=vocab_size, encoder_output_size=encoder_output_size, **ctc_conf - ) - + + ctc = CTC(odim=vocab_size, encoder_output_size=encoder_output_size, **ctc_conf) + self.blank_id = blank_id self.sos = vocab_size - 1 self.eos = vocab_size - 1 @@ -140,7 +140,7 @@ class LCBNet(nn.Module): self.decoder = None else: self.decoder = decoder - + self.criterion_att = LabelSmoothingLoss( size=vocab_size, padding_idx=ignore_id, @@ -158,14 +158,14 @@ class LCBNet(nn.Module): self.ctc = None else: self.ctc = ctc - + self.share_embedding = share_embedding if self.share_embedding: self.decoder.embed = None - + self.length_normalized_loss = length_normalized_loss self.beam_search = None - + def forward( self, speech: torch.Tensor, @@ -186,30 +186,30 @@ class LCBNet(nn.Module): text_lengths = text_lengths[:, 0] if len(speech_lengths.size()) > 1: speech_lengths = speech_lengths[:, 0] - + batch_size = speech.shape[0] - + # 1. Encoder encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) intermediate_outs = None if isinstance(encoder_out, tuple): intermediate_outs = encoder_out[1] encoder_out = encoder_out[0] - + loss_att, acc_att, cer_att, wer_att = None, None, None, None loss_ctc, cer_ctc = None, None stats = dict() - + # decoder: CTC branch if self.ctc_weight != 0.0: loss_ctc, cer_ctc = self._calc_ctc_loss( encoder_out, encoder_out_lens, text, text_lengths ) - + # Collect CTC branch stats stats["loss_ctc"] = loss_ctc.detach() if loss_ctc is not None else None stats["cer_ctc"] = cer_ctc - + # Intermediate CTC (optional) loss_interctc = 0.0 if self.interctc_weight != 0.0 and intermediate_outs is not None: @@ -220,25 +220,23 @@ class LCBNet(nn.Module): intermediate_out, encoder_out_lens, text, text_lengths ) loss_interctc = loss_interctc + loss_ic - + # Collect Intermedaite CTC stats stats["loss_interctc_layer{}".format(layer_idx)] = ( loss_ic.detach() if loss_ic is not None else None ) stats["cer_interctc_layer{}".format(layer_idx)] = cer_ic - + loss_interctc = loss_interctc / len(intermediate_outs) - + # calculate whole encoder loss - loss_ctc = ( - 1 - self.interctc_weight - ) * loss_ctc + self.interctc_weight * loss_interctc - + loss_ctc = (1 - self.interctc_weight) * loss_ctc + self.interctc_weight * loss_interctc + # decoder: Attention decoder branch loss_att, acc_att, cer_att, wer_att = self._calc_att_loss( encoder_out, encoder_out_lens, text, text_lengths ) - + # 3. CTC-Att loss definition if self.ctc_weight == 0.0: loss = loss_att @@ -246,25 +244,27 @@ class LCBNet(nn.Module): loss = loss_ctc else: loss = self.ctc_weight * loss_ctc + (1 - self.ctc_weight) * loss_att - + # Collect Attn branch stats stats["loss_att"] = loss_att.detach() if loss_att is not None else None stats["acc"] = acc_att stats["cer"] = cer_att stats["wer"] = wer_att - + # Collect total loss stats stats["loss"] = torch.clone(loss.detach()) - + # force_gatherable: to-device and to-tensor if scalar for DataParallel if self.length_normalized_loss: batch_size = int((text_lengths + 1).sum()) loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) return loss, stats, weight - def encode( - self, speech: torch.Tensor, speech_lengths: torch.Tensor, **kwargs, + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + **kwargs, ) -> Tuple[torch.Tensor, torch.Tensor]: """Frontend + Encoder. Note that this method is used by asr_inference.py Args: @@ -283,20 +283,18 @@ class LCBNet(nn.Module): # feats: (Batch, Length, Dim) # -> encoder_out: (Batch, Length2, Dim2) if self.encoder.interctc_use_conditioning: - encoder_out, encoder_out_lens, _ = self.encoder( - speech, speech_lengths, ctc=self.ctc - ) + encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths, ctc=self.ctc) else: encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths) intermediate_outs = None if isinstance(encoder_out, tuple): intermediate_outs = encoder_out[1] encoder_out = encoder_out[0] - + if intermediate_outs is not None: return (encoder_out, intermediate_outs), encoder_out_lens return encoder_out, encoder_out_lens - + def _calc_att_loss( self, encoder_out: torch.Tensor, @@ -306,12 +304,10 @@ class LCBNet(nn.Module): ): ys_in_pad, ys_out_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id) ys_in_lens = ys_pad_lens + 1 - + # 1. Forward decoder - decoder_out, _ = self.decoder( - encoder_out, encoder_out_lens, ys_in_pad, ys_in_lens - ) - + decoder_out, _ = self.decoder(encoder_out, encoder_out_lens, ys_in_pad, ys_in_lens) + # 2. Compute attention loss loss_att = self.criterion_att(decoder_out, ys_out_pad) acc_att = th_accuracy( @@ -319,16 +315,16 @@ class LCBNet(nn.Module): ys_out_pad, ignore_label=self.ignore_id, ) - + # Compute cer/wer using attention-decoder if self.training or self.error_calculator is None: cer_att, wer_att = None, None else: ys_hat = decoder_out.argmax(dim=-1) cer_att, wer_att = self.error_calculator(ys_hat.cpu(), ys_pad.cpu()) - + return loss_att, acc_att, cer_att, wer_att - + def _calc_ctc_loss( self, encoder_out: torch.Tensor, @@ -338,41 +334,39 @@ class LCBNet(nn.Module): ): # Calc CTC loss loss_ctc = self.ctc(encoder_out, encoder_out_lens, ys_pad, ys_pad_lens) - + # Calc CER using CTC cer_ctc = None if not self.training and self.error_calculator is not None: ys_hat = self.ctc.argmax(encoder_out).data cer_ctc = self.error_calculator(ys_hat.cpu(), ys_pad.cpu(), is_ctc=True) return loss_ctc, cer_ctc - - def init_beam_search(self, - **kwargs, - ): + + def init_beam_search( + self, + **kwargs, + ): from funasr.models.transformer.search import BeamSearch from funasr.models.transformer.scorers.ctc import CTCPrefixScorer from funasr.models.transformer.scorers.length_bonus import LengthBonus - + # 1. Build ASR model scorers = {} - + if self.ctc != None: ctc = CTCPrefixScorer(ctc=self.ctc, eos=self.eos) - scorers.update( - ctc=ctc - ) + scorers.update(ctc=ctc) token_list = kwargs.get("token_list") scorers.update( decoder=self.decoder, length_bonus=LengthBonus(len(token_list)), ) - # 3. Build ngram model # ngram is not supported now ngram = None scorers["ngram"] = ngram - + weights = dict( decoder=1.0 - kwargs.get("decoding_ctc_weight", 0.3), ctc=kwargs.get("decoding_ctc_weight", 0.3), @@ -392,19 +386,20 @@ class LCBNet(nn.Module): ) self.beam_search = beam_search - - def inference(self, - data_in, - data_lengths=None, - key: list=None, - tokenizer=None, - frontend=None, - **kwargs, - ): - + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): + if kwargs.get("batch_size", 1) > 1: raise NotImplementedError("batch decoding is not implemented") - + # init beamsearch if self.beam_search is None: logging.info("enable beam_search") @@ -412,7 +407,9 @@ class LCBNet(nn.Module): self.nbest = kwargs.get("nbest", 1) meta_data = {} - if isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank": # fbank + if ( + isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank" + ): # fbank speech, speech_lengths = data_in, data_lengths if len(speech.shape) < 3: speech = speech[None, :, :] @@ -421,21 +418,26 @@ class LCBNet(nn.Module): else: # extract fbank feats time1 = time.perf_counter() - sample_list = load_audio_text_image_video(data_in, fs=frontend.fs, audio_fs=kwargs.get("fs", 16000), - data_type=kwargs.get("data_type", "sound"), - tokenizer=tokenizer) + sample_list = load_audio_text_image_video( + data_in, + fs=frontend.fs, + audio_fs=kwargs.get("fs", 16000), + data_type=kwargs.get("data_type", "sound"), + tokenizer=tokenizer, + ) time2 = time.perf_counter() meta_data["load_data"] = f"{time2 - time1:0.3f}" audio_sample_list = sample_list[0] - if len(sample_list) >1: + if len(sample_list) > 1: ocr_sample_list = sample_list[1] else: ocr_sample_list = [[294, 0]] - speech, speech_lengths = extract_fbank(audio_sample_list, data_type=kwargs.get("data_type", "sound"), - frontend=frontend) + speech, speech_lengths = extract_fbank( + audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend + ) time3 = time.perf_counter() meta_data["extract_feat"] = f"{time3 - time2:0.3f}" - frame_shift = 10 + frame_shift = 10 meta_data["batch_data_time"] = speech_lengths.sum().item() * frame_shift / 1000 speech = speech.to(device=kwargs["device"]) @@ -447,15 +449,19 @@ class LCBNet(nn.Module): ocr_list_new = [[x + 1 if x != 0 else x for x in sublist] for sublist in ocr_sample_list] ocr = torch.tensor(ocr_list_new).to(device=kwargs["device"]) - ocr_lengths = ocr.new_full([1], dtype=torch.long, fill_value=ocr.size(1)).to(device=kwargs["device"]) + ocr_lengths = ocr.new_full([1], dtype=torch.long, fill_value=ocr.size(1)).to( + device=kwargs["device"] + ) ocr, ocr_lens, _ = self.text_encoder(ocr, ocr_lengths) - fusion_out, _, _, _ = self.fusion_encoder(encoder_out,None, ocr, None) + fusion_out, _, _, _ = self.fusion_encoder(encoder_out, None, ocr, None) encoder_out = encoder_out + fusion_out # c. Passed the encoder result and the beam search nbest_hyps = self.beam_search( - x=encoder_out[0], maxlenratio=kwargs.get("maxlenratio", 0.0), minlenratio=kwargs.get("minlenratio", 0.0) + x=encoder_out[0], + maxlenratio=kwargs.get("maxlenratio", 0.0), + minlenratio=kwargs.get("minlenratio", 0.0), ) - + nbest_hyps = nbest_hyps[: self.nbest] results = [] @@ -468,17 +474,21 @@ class LCBNet(nn.Module): if not hasattr(self, "writer"): self.writer = DatadirWriter(kwargs.get("output_dir")) ibest_writer = self.writer[f"{nbest_idx + 1}best_recog"] - + # remove sos/eos and get results last_pos = -1 if isinstance(hyp.yseq, list): token_int = hyp.yseq[1:last_pos] else: token_int = hyp.yseq[1:last_pos].tolist() - + # remove blank symbol id, which is assumed to be 0 - token_int = list(filter(lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int)) - + token_int = list( + filter( + lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int + ) + ) + # Change integer-ids to tokens token = tokenizer.ids2tokens(token_int) text = tokenizer.tokens2text(token) @@ -486,10 +496,9 @@ class LCBNet(nn.Module): text_postprocessed, _ = postprocess_utils.sentence_postprocess(token) result_i = {"key": key[i], "token": token, "text": text_postprocessed} results.append(result_i) - + if ibest_writer is not None: ibest_writer["token"][key[i]] = " ".join(token) ibest_writer["text"][key[i]] = text_postprocessed - - return results, meta_data + return results, meta_data diff --git a/funasr/models/llm_asr/adaptor.py b/funasr/models/llm_asr/adaptor.py index 20935883f..8c2a80434 100644 --- a/funasr/models/llm_asr/adaptor.py +++ b/funasr/models/llm_asr/adaptor.py @@ -3,6 +3,7 @@ import torch.nn as nn from funasr.register import tables + @tables.register("adaptor_classes", "Linear") class Linear(nn.Module): def __init__(self, downsample_rate, encoder_dim, llm_dim, ffn_dim: int = 2048, **kwargs): @@ -20,7 +21,7 @@ class Linear(nn.Module): if num_frames_to_discard > 0: x = x[:, :-num_frames_to_discard, :] seq_len = x.size(1) - + x = x.contiguous() x = x.view(batch_size, seq_len // self.k, dim * self.k) x = self.linear1(x) @@ -28,6 +29,7 @@ class Linear(nn.Module): x = self.linear2(x) return x + @tables.register("adaptor_classes", "QFormer") class EncoderProjectorQFormer(nn.Module): def __init__(self, downsample_rate, encoder_dim, llm_dim, ffn_dim: int = 2048, **kwargs): @@ -35,28 +37,29 @@ class EncoderProjectorQFormer(nn.Module): self.encoder_dim = encoder_dim self.llm_dim = llm_dim from transformers import Blip2QFormerConfig, Blip2QFormerModel + configuration = Blip2QFormerConfig() configuration.encoder_hidden_size = self.encoder_dim configuration.num_hidden_layers = 2 - + self.query_len = 64 self.query = nn.Parameter(torch.zeros(1, self.query_len, configuration.hidden_size)) self.query.data.normal_(mean=0.0, std=1.0) self.qformer = Blip2QFormerModel(configuration) - + self.linear = nn.Linear(configuration.hidden_size, self.llm_dim) self.norm = nn.LayerNorm(self.llm_dim, eps=1e-5) - + def forward(self, x, atts): query = self.query.expand(x.shape[0], -1, -1) - + query_output = self.qformer( query_embeds=query, encoder_hidden_states=x, encoder_attention_mask=atts, return_dict=True, ) - + query_proj = self.norm(self.linear(query_output.last_hidden_state)) - - return query_proj \ No newline at end of file + + return query_proj diff --git a/funasr/models/llm_asr/model.py b/funasr/models/llm_asr/model.py index 90cbd948a..4345f696d 100644 --- a/funasr/models/llm_asr/model.py +++ b/funasr/models/llm_asr/model.py @@ -23,7 +23,7 @@ from funasr.register import tables @tables.register("model_classes", "LLMASR") class LLMASR(nn.Module): """ """ - + def __init__( self, specaug: str = None, @@ -59,28 +59,29 @@ class LLMASR(nn.Module): # postencoder: Optional[AbsPostEncoder] = None, **kwargs, ): - + super().__init__() - + if specaug is not None: specaug_class = tables.specaug_classes.get(specaug) specaug = specaug_class(**specaug_conf) if normalize is not None: normalize_class = tables.normalize_classes.get(normalize) normalize = normalize_class(**normalize_conf) - + # audio encoder hub = audio_encoder_conf.get("hub", None) if hub == "ms": from funasr import AutoModel + model = AutoModel(model=audio_encoder, model_revision="master") # frontend = model.kwargs.get("frontend") audio_encoder_output_size = model.model.encoder_output_size audio_encoder = model.model.model.encoder - + # self.frontend = frontend - + elif hub == "hf": pass else: @@ -92,7 +93,7 @@ class LLMASR(nn.Module): for name, param in audio_encoder.named_parameters(): param.requires_grad = False audio_encoder.eval() - + self.audio_encoder = audio_encoder # llm @@ -102,7 +103,7 @@ class LLMASR(nn.Module): from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig init_param_path = llm_conf.get("init_param_path", "vicuna-7b-v1.5") - + model = AutoModelForCausalLM.from_pretrained( init_param_path, load_in_8bit=None, @@ -115,15 +116,14 @@ class LLMASR(nn.Module): param.requires_grad = False model.eval() self.llm = model - + # adaptor adaptor_class = tables.adaptor_classes.get(audio_adaptor) audio_adaptor_conf["encoder_dim"] = audio_encoder_output_size audio_adaptor = adaptor_class(**audio_adaptor_conf) - + self.audio_adaptor = audio_adaptor - - + self.blank_id = blank_id self.sos = sos if sos is not None else vocab_size - 1 self.eos = eos if eos is not None else vocab_size - 1 @@ -143,7 +143,7 @@ class LLMASR(nn.Module): self.length_normalized_loss = length_normalized_loss self.beam_search = None - + def forward( self, speech: torch.Tensor, @@ -151,7 +151,7 @@ class LLMASR(nn.Module): text: torch.Tensor, text_lengths: torch.Tensor, input_ids: torch.Tensor, - attention_mask:torch.Tensor, + attention_mask: torch.Tensor, labels_ids: torch.Tensor, label_mask: torch.Tensor, audio_mask: torch.Tensor, @@ -170,15 +170,15 @@ class LLMASR(nn.Module): text_lengths = text_lengths[:, 0] if len(speech_lengths.size()) > 1: speech_lengths = speech_lengths[:, 0] - + batch_size = speech.shape[0] - + # audio encoder encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) - + # audio_adaptor encoder_out = self.audio_adaptor(encoder_out) - + input_ids[input_ids == -1] = 0 input_ids[input_ids == -100] = 0 if hasattr(self.llm.model, "embed_tokens"): @@ -193,12 +193,15 @@ class LLMASR(nn.Module): _, l, _ = encoder_out.shape # [audio, bos, prompt, input, pad] encoder_outs_pad = F.pad(encoder_out, (0, 0, 0, token_num - l, 0, 0), value=0.0) - inputs_embeds = encoder_outs_pad * audio_mask[:, :, None] + inputs_embeds * (1.0-audio_mask[:, :, None]) + inputs_embeds = encoder_outs_pad * audio_mask[:, :, None] + inputs_embeds * ( + 1.0 - audio_mask[:, :, None] + ) - model_outputs = self.llm(inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=labels_ids) + model_outputs = self.llm( + inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=labels_ids + ) loss = model_outputs.loss - stats = {} with torch.no_grad(): preds = torch.argmax(model_outputs.logits, -1) @@ -214,7 +217,10 @@ class LLMASR(nn.Module): return loss, stats, weight def encode( - self, speech: torch.Tensor, speech_lengths: torch.Tensor, **kwargs, + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + **kwargs, ): speech = speech.permute(0, 2, 1) res = self.audio_encoder(speech) @@ -223,25 +229,26 @@ class LLMASR(nn.Module): else: encoder_out, encoder_out_lens = res, speech_lengths return encoder_out, encoder_out_lens - - def inference(self, - data_in, - data_lengths=None, - key: list = None, - tokenizer=None, - frontend=None, - **kwargs, - ): - + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): + prompt = kwargs.get("prompt", "Transcribe speech to text.") - + if kwargs.get("batch_size", 1) > 1: raise NotImplementedError("batch decoding is not implemented") - - meta_data = {} - if isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank": # fbank + if ( + isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank" + ): # fbank speech, speech_lengths = data_in, data_lengths if len(speech.shape) < 3: speech = speech[None, :, :] @@ -250,33 +257,38 @@ class LLMASR(nn.Module): else: # extract fbank feats time1 = time.perf_counter() - audio_sample_list = load_audio_text_image_video(data_in, fs=frontend.fs, audio_fs=kwargs.get("fs", 16000), - data_type=kwargs.get("data_type", "sound"), - tokenizer=tokenizer) + audio_sample_list = load_audio_text_image_video( + data_in, + fs=frontend.fs, + audio_fs=kwargs.get("fs", 16000), + data_type=kwargs.get("data_type", "sound"), + tokenizer=tokenizer, + ) time2 = time.perf_counter() meta_data["load_data"] = f"{time2 - time1:0.3f}" - speech, speech_lengths = extract_fbank(audio_sample_list, data_type=kwargs.get("data_type", "sound"), - frontend=frontend) + speech, speech_lengths = extract_fbank( + audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend + ) time3 = time.perf_counter() meta_data["extract_feat"] = f"{time3 - time2:0.3f}" - meta_data["batch_data_time"] = speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 - + meta_data["batch_data_time"] = ( + speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 + ) + speech = speech.to(device=kwargs["device"]) speech_lengths = speech_lengths.to(device=kwargs["device"]) - + # Encoder encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) # adaptor encoder_out = self.audio_adaptor(encoder_out) - - + prompt_pre = "USER: \nINSTRUCTION: {}\nINPUT: ".format(prompt) prompt_ids = tokenizer.encode(prompt_pre) prompt_length = len(prompt_ids) prompt_ids = torch.tensor(prompt_ids, dtype=torch.int64).to(kwargs["device"]) - if hasattr(self.llm.model, "embed_tokens"): inputs_embeds = self.llm.model.embed_tokens(prompt_ids) elif hasattr(self.llm.model.model, "embed_tokens"): @@ -284,9 +296,13 @@ class LLMASR(nn.Module): else: inputs_embeds = self.llm.model.model.model.embed_tokens(prompt_ids) - inputs_embeds = torch.cat((inputs_embeds[None, :, :], encoder_out), dim=1) # [prompt, audio] - attention_mask = torch.ones(inputs_embeds.size()[:-1], dtype=torch.long).to(kwargs["device"]) - + inputs_embeds = torch.cat( + (inputs_embeds[None, :, :], encoder_out), dim=1 + ) # [prompt, audio] + attention_mask = torch.ones(inputs_embeds.size()[:-1], dtype=torch.long).to( + kwargs["device"] + ) + preds = self.llm.generate( inputs_embeds=inputs_embeds, max_length=kwargs.get("max_length", 200), @@ -301,17 +317,16 @@ class LLMASR(nn.Module): attention_mask=attention_mask, bos_token_id=tokenizer.bos_token_id, eos_token_id=tokenizer.eos_token_id, - pad_token_id=tokenizer.pad_token_id + pad_token_id=tokenizer.pad_token_id, ) - text = tokenizer.batch_decode(preds, add_special_tokens=False, skip_special_tokens=True) - text = text[0].split(': ')[-1] + text = text[0].split(": ")[-1] text = text.strip() - + # preds = torch.argmax(model_outputs.logits, -1) - + ibest_writer = None if kwargs.get("output_dir") is not None: if not hasattr(self, "writer"): @@ -324,9 +339,5 @@ class LLMASR(nn.Module): if ibest_writer is not None: ibest_writer["text"][key[0]] = text - - - - - return results, meta_data + return results, meta_data diff --git a/funasr/models/llm_asr_nar/adaptor.py b/funasr/models/llm_asr_nar/adaptor.py index 0676e7deb..0d73a1a0e 100644 --- a/funasr/models/llm_asr_nar/adaptor.py +++ b/funasr/models/llm_asr_nar/adaptor.py @@ -3,6 +3,7 @@ import torch.nn as nn from funasr.register import tables + @tables.register("adaptor_classes", "Linear") class Linear(nn.Module): def __init__(self, downsample_rate, encoder_dim, llm_dim, ffn_dim: int = 2048, **kwargs): @@ -20,7 +21,7 @@ class Linear(nn.Module): if num_frames_to_discard > 0: x = x[:, :-num_frames_to_discard, :] seq_len = x.size(1) - + x = x.contiguous() x = x.view(batch_size, seq_len // self.k, dim * self.k) x = self.linear1(x) diff --git a/funasr/models/llm_asr_nar/model.py b/funasr/models/llm_asr_nar/model.py index 994259a8f..8c0c3ff18 100644 --- a/funasr/models/llm_asr_nar/model.py +++ b/funasr/models/llm_asr_nar/model.py @@ -12,6 +12,7 @@ from funasr.losses.label_smoothing_loss import LabelSmoothingLoss from funasr.models.ctc.ctc import CTC from funasr.models.transformer.utils.add_sos_eos import add_sos_eos from funasr.metrics.compute_acc import th_accuracy, compute_accuracy + # from funasr.models.e2e_asr_common import ErrorCalculator from funasr.train_utils.device_funcs import force_gatherable from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank @@ -24,7 +25,7 @@ from funasr.register import tables @tables.register("model_classes", "LLMASRNAR") class LLMASRNAR(nn.Module): """ """ - + def __init__( self, specaug: str = None, @@ -60,28 +61,32 @@ class LLMASRNAR(nn.Module): # postencoder: Optional[AbsPostEncoder] = None, **kwargs, ): - + super().__init__() - + if specaug is not None: specaug_class = tables.specaug_classes.get(specaug) specaug = specaug_class(**specaug_conf) if normalize is not None: normalize_class = tables.normalize_classes.get(normalize) normalize = normalize_class(**normalize_conf) - + # audio encoder hub = encoder_conf.get("hub", None) if hub == "funasr": from funasr import AutoModel - init_param_path = encoder_conf.get("init_param_path", "iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch") + + init_param_path = encoder_conf.get( + "init_param_path", + "iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch", + ) model = AutoModel(model=init_param_path, model_revision="master") # frontend = model.kwargs.get("frontend") model.model.decoder = None - + self.audio_encoder = model.model # self.frontend = frontend - + elif hub == "hf": pass else: @@ -108,14 +113,13 @@ class LLMASRNAR(nn.Module): param.requires_grad = False model.eval() self.llm = model - + # adaptor adaptor_class = tables.adaptor_classes.get(adaptor) adaptor = adaptor_class(**adaptor_conf) - + self.adaptor = adaptor - - + self.blank_id = blank_id self.sos = sos if sos is not None else vocab_size - 1 self.eos = eos if eos is not None else vocab_size - 1 @@ -125,7 +129,6 @@ class LLMASRNAR(nn.Module): self.normalize = normalize self.encoder = encoder - self.criterion_att = LabelSmoothingLoss( size=vocab_size, padding_idx=ignore_id, @@ -142,7 +145,7 @@ class LLMASRNAR(nn.Module): self.length_normalized_loss = length_normalized_loss self.beam_search = None - + def forward( self, speech: torch.Tensor, @@ -150,7 +153,7 @@ class LLMASRNAR(nn.Module): text: torch.Tensor, text_lengths: torch.Tensor, input_ids: torch.Tensor, - attention_mask:torch.Tensor, + attention_mask: torch.Tensor, labels_ids: torch.Tensor, label_mask: torch.Tensor, audio_mask: torch.Tensor, @@ -169,12 +172,12 @@ class LLMASRNAR(nn.Module): text_lengths = text_lengths[:, 0] if len(speech_lengths.size()) > 1: speech_lengths = speech_lengths[:, 0] - + batch_size = speech.shape[0] - + # audio encoder encoder_out, encoder_out_lens = self.encode(speech, speech_lengths, audio_mask=audio_mask) - + # adaptor encoder_out = self.adaptor(encoder_out) @@ -191,14 +194,17 @@ class LLMASRNAR(nn.Module): if audio_mask is not None: batch_size, token_num, dims = inputs_embeds.shape _, l, _ = encoder_out.shape - encoder_outs_pad = F.pad(encoder_out, (0, 0, token_num-l-1, 1, 0, 0), value=0.0) - inputs_embeds = encoder_outs_pad * audio_mask[:, :, None] + inputs_embeds * (1.0-audio_mask[:, :, None]) + encoder_outs_pad = F.pad(encoder_out, (0, 0, token_num - l - 1, 1, 0, 0), value=0.0) + inputs_embeds = encoder_outs_pad * audio_mask[:, :, None] + inputs_embeds * ( + 1.0 - audio_mask[:, :, None] + ) inputs_embeds = F.pad(inputs_embeds[:, 1:, :], (0, 0, 0, 1, 0, 0), value=0.0) - model_outputs = self.llm(inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=labels_ids) + model_outputs = self.llm( + inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=labels_ids + ) loss = model_outputs.loss - stats = {} with torch.no_grad(): preds = torch.argmax(model_outputs.logits, -1) @@ -212,11 +218,14 @@ class LLMASRNAR(nn.Module): batch_size = int((text_lengths + 1).sum()) loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) return loss, stats, weight - + def encode( - self, speech: torch.Tensor, speech_lengths: torch.Tensor, **kwargs, + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + **kwargs, ): - + audio_mask = kwargs.get("audio_mask", None) audio_token_lengths = audio_mask.sum(-1) if audio_mask is not None else None text_token_int = kwargs.get("text_token_int", None) @@ -227,32 +236,33 @@ class LLMASRNAR(nn.Module): enc, enc_lens = self.audio_encoder.encode(**batch) with autocast(False): enc_mask = sequence_mask(enc_lens, enc.size(1), device=enc.device)[:, None, :] - pre_acoustic_embeds, pre_token_length, _, _ = self.audio_encoder.predictor(enc, - mask=enc_mask, - target_label_length=audio_token_lengths, - ) + pre_acoustic_embeds, pre_token_length, _, _ = self.audio_encoder.predictor( + enc, + mask=enc_mask, + target_label_length=audio_token_lengths, + ) return pre_acoustic_embeds, pre_token_length + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): - def inference(self, - data_in, - data_lengths=None, - key: list = None, - tokenizer=None, - frontend=None, - **kwargs, - ): - prompt = kwargs.get("prompt", "Transcribe speech to text.") - + if kwargs.get("batch_size", 1) > 1: raise NotImplementedError("batch decoding is not implemented") - - meta_data = {} - if isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank": # fbank + if ( + isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank" + ): # fbank speech, speech_lengths = data_in, data_lengths if len(speech.shape) < 3: speech = speech[None, :, :] @@ -261,9 +271,13 @@ class LLMASRNAR(nn.Module): else: # extract fbank feats time1 = time.perf_counter() - audio_sample_list = load_audio_text_image_video(data_in, fs=frontend.fs, audio_fs=kwargs.get("fs", 16000), - data_type=kwargs.get("data_type", "sound"), - tokenizer=None) + audio_sample_list = load_audio_text_image_video( + data_in, + fs=frontend.fs, + audio_fs=kwargs.get("fs", 16000), + data_type=kwargs.get("data_type", "sound"), + tokenizer=None, + ) if len(kwargs.get("data_type", [])) > 1: audio_sample_list, text_token_int_list = audio_sample_list text_token_int = text_token_int_list[0].replace(" ", "") @@ -272,28 +286,31 @@ class LLMASRNAR(nn.Module): text_token_int = None time2 = time.perf_counter() meta_data["load_data"] = f"{time2 - time1:0.3f}" - speech, speech_lengths = extract_fbank(audio_sample_list, data_type=kwargs.get("data_type", "sound"), - frontend=frontend) + speech, speech_lengths = extract_fbank( + audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend + ) time3 = time.perf_counter() meta_data["extract_feat"] = f"{time3 - time2:0.3f}" - meta_data["batch_data_time"] = speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 - + meta_data["batch_data_time"] = ( + speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 + ) + speech = speech.to(device=kwargs["device"]) speech_lengths = speech_lengths.to(device=kwargs["device"]) - + # Encoder - encoder_out, encoder_out_lens = self.encode(speech, speech_lengths, text_token_int=text_token_int) + encoder_out, encoder_out_lens = self.encode( + speech, speech_lengths, text_token_int=text_token_int + ) # adaptor encoder_out = self.adaptor(encoder_out) - - + prompt_pre = "USER: \nINSTRUCTION: {}\nINPUT: ".format(prompt) prompt_ids = tokenizer.encode(prompt_pre) prompt_length = len(prompt_ids) prompt_ids = torch.tensor(prompt_ids, dtype=torch.int64).to(kwargs["device"]) - if hasattr(self.llm.model, "embed_tokens"): inputs_embeds = self.llm.model.embed_tokens(prompt_ids) elif hasattr(self.llm.model.model, "embed_tokens"): @@ -301,9 +318,13 @@ class LLMASRNAR(nn.Module): else: inputs_embeds = self.llm.model.model.model.embed_tokens(prompt_ids) - inputs_embeds = torch.cat((inputs_embeds[None, :, :], encoder_out), dim=1) # [prompt, audio] - attention_mask = torch.ones(inputs_embeds.size()[:-1], dtype=torch.long).to(kwargs["device"]) - + inputs_embeds = torch.cat( + (inputs_embeds[None, :, :], encoder_out), dim=1 + ) # [prompt, audio] + attention_mask = torch.ones(inputs_embeds.size()[:-1], dtype=torch.long).to( + kwargs["device"] + ) + # model_outputs = self.llm.generate( # inputs_embeds=inputs_embeds, # max_length=kwargs.get("max_length", 200), @@ -321,16 +342,17 @@ class LLMASRNAR(nn.Module): # pad_token_id=tokenizer.pad_token_id # ) - - model_outputs = self.llm(inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=None) + model_outputs = self.llm( + inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=None + ) preds = torch.argmax(model_outputs.logits, -1) text = tokenizer.batch_decode(preds, add_special_tokens=False, skip_special_tokens=True) - text = text[0].split(': ')[-1] + text = text[0].split(": ")[-1] text = text.strip() - + # preds = torch.argmax(model_outputs.logits, -1) - + ibest_writer = None if kwargs.get("output_dir") is not None: if not hasattr(self, "writer"): @@ -343,17 +365,14 @@ class LLMASRNAR(nn.Module): if ibest_writer is not None: ibest_writer["text"][key[0]] = text - - - - + return results, meta_data @tables.register("model_classes", "LLMASRNARPrompt") class LLMASRNARPrompt(nn.Module): """ """ - + def __init__( self, specaug: str = None, @@ -390,43 +409,46 @@ class LLMASRNARPrompt(nn.Module): # postencoder: Optional[AbsPostEncoder] = None, **kwargs, ): - + super().__init__() - + if specaug is not None: specaug_class = tables.specaug_classes.get(specaug) specaug = specaug_class(**specaug_conf) if normalize is not None: normalize_class = tables.normalize_classes.get(normalize) normalize = normalize_class(**normalize_conf) - + # audio encoder hub = encoder_conf.get("hub", None) if hub == "funasr": from funasr import AutoModel - init_param_path = encoder_conf.get("init_param_path", - "iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch") + + init_param_path = encoder_conf.get( + "init_param_path", + "iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch", + ) model = AutoModel(model=init_param_path, model_revision="master") # frontend = model.kwargs.get("frontend") model.model.decoder = None - + self.audio_encoder = model.model # self.frontend = frontend self.predictor_weight = predictor_weight - + elif hub == "hf": pass else: encoder_class = tables.encoder_classes.get(encoder) encoder = encoder_class(input_size=input_size, **encoder_conf) encoder_output_size = encoder.output_size() - + # llm hub = llm_conf.get("hub", "hf") self.llm = None if hub == "hf": from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig - + init_param_path = llm_conf.get("init_param_path", "vicuna-7b-v1.5") model = AutoModelForCausalLM.from_pretrained( init_param_path, @@ -440,13 +462,13 @@ class LLMASRNARPrompt(nn.Module): param.requires_grad = False model.eval() self.llm = model - + # adaptor adaptor_class = tables.adaptor_classes.get(adaptor) adaptor = adaptor_class(**adaptor_conf) - + self.adaptor = adaptor - + self.blank_id = blank_id self.sos = sos if sos is not None else vocab_size - 1 self.eos = eos if eos is not None else vocab_size - 1 @@ -455,7 +477,7 @@ class LLMASRNARPrompt(nn.Module): self.specaug = specaug self.normalize = normalize self.encoder = encoder - + self.criterion_att = LabelSmoothingLoss( size=vocab_size, padding_idx=ignore_id, @@ -470,19 +492,17 @@ class LLMASRNARPrompt(nn.Module): # ) # self.error_calculator = None - + self.length_normalized_loss = length_normalized_loss self.beam_search = None if ctc_weight > 0.0: if ctc_conf is None: ctc_conf = {} - - ctc = CTC( - odim=vocab_size, encoder_output_size=adaptor_conf["encoder_dim"], **ctc_conf - ) + + ctc = CTC(odim=vocab_size, encoder_output_size=adaptor_conf["encoder_dim"], **ctc_conf) self.ctc_weight = ctc_weight self.ctc = ctc - + def forward( self, speech: torch.Tensor, @@ -509,7 +529,7 @@ class LLMASRNARPrompt(nn.Module): text_lengths = text_lengths[:, 0] if len(speech_lengths.size()) > 1: speech_lengths = speech_lengths[:, 0] - + batch_size = speech.shape[0] stats = {} @@ -517,21 +537,18 @@ class LLMASRNARPrompt(nn.Module): outs = self.encode(speech, speech_lengths, audio_mask=audio_mask) enc, enc_lens = outs[0], outs[1] encoder_out, encoder_out_lens, loss_pre = outs[2], outs[3], outs[4] - # decoder: CTC branch - + if self.ctc_weight != 0.0: - loss_ctc, cer_ctc = self._calc_ctc_loss( - enc, enc_lens, text, text_lengths - ) - + loss_ctc, cer_ctc = self._calc_ctc_loss(enc, enc_lens, text, text_lengths) + # Collect CTC branch stats stats["loss_ctc"] = torch.clone(loss_ctc.detach()) if loss_ctc is not None else None - + # adaptor encoder_out = self.adaptor(encoder_out) - + if input_ids is not None: input_ids[input_ids == -1] = 0 input_ids[input_ids == -100] = 0 @@ -541,7 +558,7 @@ class LLMASRNARPrompt(nn.Module): inputs_embeds = self.llm.model.model.embed_tokens(input_ids) else: inputs_embeds = self.llm.model.model.model.embed_tokens(input_ids) - + if audio_mask is not None: # inputs_embeds： [bos, prompt, input, pad, target] prompt_bos_length = kwargs.get("prompt_bos_length", None) @@ -549,58 +566,74 @@ class LLMASRNARPrompt(nn.Module): prompt_bos_length = prompt_bos_length[0].item() batch_size, token_num, dims = inputs_embeds.shape _, l, _ = encoder_out.shape - encoder_outs_pad = F.pad(encoder_out, (0, 0, prompt_bos_length, token_num - prompt_bos_length - l, 0, 0), value=0.0) - inputs_embeds = encoder_outs_pad * audio_mask[:, :, None] + inputs_embeds * (1.0 - audio_mask[:, :, None]) - inputs_embeds = F.pad(inputs_embeds[:, 1:, :], (0, 0, 0, 1, 0, 0), value=0.0) # [prompt, input, pad, target, 0.0] - + encoder_outs_pad = F.pad( + encoder_out, + (0, 0, prompt_bos_length, token_num - prompt_bos_length - l, 0, 0), + value=0.0, + ) + inputs_embeds = encoder_outs_pad * audio_mask[:, :, None] + inputs_embeds * ( + 1.0 - audio_mask[:, :, None] + ) + inputs_embeds = F.pad( + inputs_embeds[:, 1:, :], (0, 0, 0, 1, 0, 0), value=0.0 + ) # [prompt, input, pad, target, 0.0] + # labels_ids: [bos, prompt, input, target, eos] -> [-1, -1, input, target, eos] # loss: # inputs_embeds[:-1] -> [prompt, input, pad, target] # labels_ids[1:] -> [prompt, input, target, eos] -> [-1, input, target, eos]; - model_outputs = self.llm(inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=labels_ids) + model_outputs = self.llm( + inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=labels_ids + ) loss_llm = model_outputs.loss stats["loss_llm"] = torch.clone(loss_llm.detach()) if self.ctc_weight > 0.0: loss_llm = self.ctc_weight * loss_ctc + loss_llm loss = loss_llm + loss_pre * self.predictor_weight - + with torch.no_grad(): preds = torch.argmax(model_outputs.logits, -1) acc_att = compute_accuracy(preds[:, :-1], labels_ids[:, 1:], ignore_label=-100) stats["acc"] = acc_att - + stats["loss_pre"] = torch.clone(loss_pre.detach()) stats["loss"] = torch.clone(loss.detach()) stats["batch_size"] = batch_size - + # force_gatherable: to-device and to-tensor if scalar for DataParallel if self.length_normalized_loss: batch_size = int((text_lengths + 1).sum()) loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) return loss, stats, weight - + def encode( - self, speech: torch.Tensor, speech_lengths: torch.Tensor, **kwargs, + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + **kwargs, ): - + audio_mask = kwargs.get("audio_mask", None) audio_token_lengths = audio_mask.sum(-1) if audio_mask is not None else None text_token_int = kwargs.get("text_token_int", None) if audio_token_lengths is None and text_token_int is not None: audio_token_lengths = torch.tensor([len(text_token_int)], dtype=torch.int64) - + batch = {"speech": speech, "speech_lengths": speech_lengths} enc, enc_lens = self.audio_encoder.encode(**batch) with autocast(False): enc_mask = sequence_mask(enc_lens, enc.size(1), device=enc.device)[:, None, :] - pre_acoustic_embeds, pre_token_length, _, _ = self.audio_encoder.predictor(enc, - mask=enc_mask, - target_label_length=audio_token_lengths, - ) + pre_acoustic_embeds, pre_token_length, _, _ = self.audio_encoder.predictor( + enc, + mask=enc_mask, + target_label_length=audio_token_lengths, + ) loss_pre = 0.0 if audio_token_lengths is not None: - loss_pre = self.criterion_pre(audio_token_lengths.type_as(pre_token_length), pre_token_length) - + loss_pre = self.criterion_pre( + audio_token_lengths.type_as(pre_token_length), pre_token_length + ) + return enc, enc_lens, pre_acoustic_embeds, pre_token_length, loss_pre def _calc_ctc_loss( @@ -612,30 +645,33 @@ class LLMASRNARPrompt(nn.Module): ): # Calc CTC loss loss_ctc = self.ctc(encoder_out, encoder_out_lens, ys_pad, ys_pad_lens) - + # Calc CER using CTC cer_ctc = None if not self.training and self.error_calculator is not None: ys_hat = self.ctc.argmax(encoder_out).data cer_ctc = self.error_calculator(ys_hat.cpu(), ys_pad.cpu(), is_ctc=True) return loss_ctc, cer_ctc - - def inference(self, - data_in, - data_lengths=None, - key: list = None, - tokenizer=None, - frontend=None, - **kwargs, - ): - + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): + prompt = kwargs.get("prompt", "Transcribe speech to text.") - + if kwargs.get("batch_size", 1) > 1: raise NotImplementedError("batch decoding is not implemented") - + meta_data = {} - if isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank": # fbank + if ( + isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank" + ): # fbank speech, speech_lengths = data_in, data_lengths if len(speech.shape) < 3: speech = speech[None, :, :] @@ -644,9 +680,13 @@ class LLMASRNARPrompt(nn.Module): else: # extract fbank feats time1 = time.perf_counter() - audio_sample_list = load_audio_text_image_video(data_in, fs=frontend.fs, audio_fs=kwargs.get("fs", 16000), - data_type=kwargs.get("data_type", "sound"), - tokenizer=None) + audio_sample_list = load_audio_text_image_video( + data_in, + fs=frontend.fs, + audio_fs=kwargs.get("fs", 16000), + data_type=kwargs.get("data_type", "sound"), + tokenizer=None, + ) if len(kwargs.get("data_type", [])) > 1: audio_sample_list, text_token_int_list = audio_sample_list text_token_int = text_token_int_list[0] @@ -657,22 +697,25 @@ class LLMASRNARPrompt(nn.Module): text_token_int = None time2 = time.perf_counter() meta_data["load_data"] = f"{time2 - time1:0.3f}" - speech, speech_lengths = extract_fbank(audio_sample_list, data_type=kwargs.get("data_type", "sound"), - frontend=frontend) + speech, speech_lengths = extract_fbank( + audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend + ) time3 = time.perf_counter() meta_data["extract_feat"] = f"{time3 - time2:0.3f}" - meta_data["batch_data_time"] = speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 - + meta_data["batch_data_time"] = ( + speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 + ) + speech = speech.to(device=kwargs["device"]) speech_lengths = speech_lengths.to(device=kwargs["device"]) - + # Encoder res = self.encode(speech, speech_lengths, text_token_int=text_token_int) encoder_out = res[0] - + # adaptor encoder_out = self.adaptor(encoder_out) - + prompt_pre = "USER: \nINSTRUCTION: {}\nINPUT: ".format(prompt) prompt_ids = tokenizer.encode(prompt_pre) if prompt_ids[0] == tokenizer.bos_token_id: @@ -681,7 +724,7 @@ class LLMASRNARPrompt(nn.Module): prompt_length = len(prompt_ids) prompt_ids = torch.tensor(prompt_ids, dtype=torch.int64).to(kwargs["device"]) pad = torch.tensor([tokenizer.pad_token_id], dtype=torch.int64).to(kwargs["device"]) - + if hasattr(self.llm.model, "embed_tokens"): inputs_embeds = self.llm.model.embed_tokens(prompt_ids) pad = self.llm.model.embed_tokens(pad) @@ -689,11 +732,15 @@ class LLMASRNARPrompt(nn.Module): inputs_embeds = self.llm.model.model.embed_tokens(prompt_ids) else: inputs_embeds = self.llm.model.model.model.embed_tokens(prompt_ids) - + # inputs_embeds = torch.cat((inputs_embeds[None, :, :], encoder_out, pad[None, :, :]), dim=1) # [prompt, audio, pad] - inputs_embeds = torch.cat((inputs_embeds[None, :, :], encoder_out), dim=1) # [prompt, audio] - attention_mask = torch.ones(inputs_embeds.size()[:-1], dtype=torch.long).to(kwargs["device"]) - + inputs_embeds = torch.cat( + (inputs_embeds[None, :, :], encoder_out), dim=1 + ) # [prompt, audio] + attention_mask = torch.ones(inputs_embeds.size()[:-1], dtype=torch.long).to( + kwargs["device"] + ) + # model_outputs = self.llm.generate( # inputs_embeds=inputs_embeds, # max_length=kwargs.get("max_length", 200), @@ -710,30 +757,32 @@ class LLMASRNARPrompt(nn.Module): # eos_token_id=tokenizer.eos_token_id, # pad_token_id=tokenizer.pad_token_id # ) - - model_outputs = self.llm(inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=None) + + model_outputs = self.llm( + inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=None + ) preds = torch.argmax(model_outputs.logits, -1) text = tokenizer.batch_decode(preds, add_special_tokens=False, skip_special_tokens=True) - - text = text[0].split(':')[-1] + + text = text[0].split(":")[-1] text = text.strip() if text.startswith("Please\n "): text = text.replace("Please\n ", "") text = text.strip() - + # preds = torch.argmax(model_outputs.logits, -1) - + ibest_writer = None if kwargs.get("output_dir") is not None: if not hasattr(self, "writer"): self.writer = DatadirWriter(kwargs.get("output_dir")) ibest_writer = self.writer[f"{0 + 1}best_recog"] - + results = [] result_i = {"key": key[0], "text": text} results.append(result_i) - + if ibest_writer is not None: ibest_writer["text"][key[0]] = text - + return results, meta_data diff --git a/funasr/models/lora/layers.py b/funasr/models/lora/layers.py index 76f046c26..b91253d56 100644 --- a/funasr/models/lora/layers.py +++ b/funasr/models/lora/layers.py @@ -9,7 +9,8 @@ import torch.nn.functional as F import math from typing import Optional, List -class LoRALayer(): + +class LoRALayer: def __init__( self, r: int, @@ -20,7 +21,7 @@ class LoRALayer(): self.r = r self.lora_alpha = lora_alpha # Optional dropout - if lora_dropout > 0.: + if lora_dropout > 0.0: self.lora_dropout = nn.Dropout(p=lora_dropout) else: self.lora_dropout = lambda x: x @@ -41,8 +42,9 @@ class Embedding(nn.Embedding, LoRALayer): **kwargs ): nn.Embedding.__init__(self, num_embeddings, embedding_dim, **kwargs) - LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=0, - merge_weights=merge_weights) + LoRALayer.__init__( + self, r=r, lora_alpha=lora_alpha, lora_dropout=0, merge_weights=merge_weights + ) # Actual trainable parameters if r > 0: self.lora_A = nn.Parameter(self.weight.new_zeros((r, num_embeddings))) @@ -54,7 +56,7 @@ class Embedding(nn.Embedding, LoRALayer): def reset_parameters(self): nn.Embedding.reset_parameters(self) - if hasattr(self, 'lora_A'): + if hasattr(self, "lora_A"): # initialize A the same way as the default for nn.Linear and B to zero nn.init.zeros_(self.lora_A) nn.init.normal_(self.lora_B) @@ -80,8 +82,13 @@ class Embedding(nn.Embedding, LoRALayer): result = nn.Embedding.forward(self, x) if self.r > 0: after_A = F.embedding( - x, self.lora_A.T, self.padding_idx, self.max_norm, - self.norm_type, self.scale_grad_by_freq, self.sparse + x, + self.lora_A.T, + self.padding_idx, + self.max_norm, + self.norm_type, + self.scale_grad_by_freq, + self.sparse, ) result += (after_A @ self.lora_B.T) * self.scaling return result @@ -97,14 +104,15 @@ class Linear(nn.Linear, LoRALayer): out_features: int, r: int = 0, lora_alpha: int = 1, - lora_dropout: float = 0., - fan_in_fan_out: bool = False, # Set this to True if the layer to replace stores weight like (fan_in, fan_out) + lora_dropout: float = 0.0, + fan_in_fan_out: bool = False, # Set this to True if the layer to replace stores weight like (fan_in, fan_out) merge_weights: bool = True, **kwargs ): nn.Linear.__init__(self, in_features, out_features, **kwargs) - LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, - merge_weights=merge_weights) + LoRALayer.__init__( + self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=merge_weights + ) self.fan_in_fan_out = fan_in_fan_out # Actual trainable parameters @@ -120,7 +128,7 @@ class Linear(nn.Linear, LoRALayer): def reset_parameters(self): nn.Linear.reset_parameters(self) - if hasattr(self, 'lora_A'): + if hasattr(self, "lora_A"): # initialize A the same way as the default for nn.Linear and B to zero nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) nn.init.zeros_(self.lora_B) @@ -128,6 +136,7 @@ class Linear(nn.Linear, LoRALayer): def train(self, mode: bool = True): def T(w): return w.T if self.fan_in_fan_out else w + nn.Linear.train(self, mode) if self.merge_weights and self.merged: # Make sure that the weights are not merged @@ -138,6 +147,7 @@ class Linear(nn.Linear, LoRALayer): def eval(self): def T(w): return w.T if self.fan_in_fan_out else w + nn.Linear.eval(self) if self.merge_weights and not self.merged: # Merge the weights and mark it @@ -148,6 +158,7 @@ class Linear(nn.Linear, LoRALayer): def forward(self, x: torch.Tensor): def T(w): return w.T if self.fan_in_fan_out else w + if self.r > 0 and not self.merged: result = F.linear(x, T(self.weight), bias=self.bias) if self.r > 0: @@ -165,33 +176,34 @@ class MergedLinear(nn.Linear, LoRALayer): out_features: int, r: int = 0, lora_alpha: int = 1, - lora_dropout: float = 0., + lora_dropout: float = 0.0, enable_lora: List[bool] = [False], fan_in_fan_out: bool = False, merge_weights: bool = True, **kwargs ): nn.Linear.__init__(self, in_features, out_features, **kwargs) - LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, - merge_weights=merge_weights) - assert out_features % len(enable_lora) == 0, \ - 'The length of enable_lora must divide out_features' + LoRALayer.__init__( + self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=merge_weights + ) + assert ( + out_features % len(enable_lora) == 0 + ), "The length of enable_lora must divide out_features" self.enable_lora = enable_lora self.fan_in_fan_out = fan_in_fan_out # Actual trainable parameters if r > 0 and any(enable_lora): - self.lora_A = nn.Parameter( - self.weight.new_zeros((r * sum(enable_lora), in_features))) + self.lora_A = nn.Parameter(self.weight.new_zeros((r * sum(enable_lora), in_features))) self.lora_B = nn.Parameter( self.weight.new_zeros((out_features // len(enable_lora) * sum(enable_lora), r)) - ) # weights for Conv1D with groups=sum(enable_lora) + ) # weights for Conv1D with groups=sum(enable_lora) self.scaling = self.lora_alpha / self.r # Freezing the pre-trained weight matrix self.weight.requires_grad = False # Compute the indices - self.lora_ind = self.weight.new_zeros( - (out_features, ), dtype=torch.bool - ).view(len(enable_lora), -1) + self.lora_ind = self.weight.new_zeros((out_features,), dtype=torch.bool).view( + len(enable_lora), -1 + ) self.lora_ind[enable_lora, :] = True self.lora_ind = self.lora_ind.view(-1) self.reset_parameters() @@ -200,7 +212,7 @@ class MergedLinear(nn.Linear, LoRALayer): def reset_parameters(self): nn.Linear.reset_parameters(self) - if hasattr(self, 'lora_A'): + if hasattr(self, "lora_A"): # initialize A the same way as the default for nn.Linear and B to zero nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) nn.init.zeros_(self.lora_B) @@ -216,6 +228,7 @@ class MergedLinear(nn.Linear, LoRALayer): def train(self, mode: bool = True): def T(w): return w.T if self.fan_in_fan_out else w + nn.Linear.train(self, mode) if self.merge_weights and self.merged: # Make sure that the weights are not merged @@ -223,7 +236,7 @@ class MergedLinear(nn.Linear, LoRALayer): delta_w = F.conv1d( self.lora_A.data.unsqueeze(0), self.lora_B.data.unsqueeze(-1), - groups=sum(self.enable_lora) + groups=sum(self.enable_lora), ).squeeze(0) self.weight.data -= self.zero_pad(T(delta_w * self.scaling)) self.merged = False @@ -231,6 +244,7 @@ class MergedLinear(nn.Linear, LoRALayer): def eval(self): def T(w): return w.T if self.fan_in_fan_out else w + nn.Linear.eval(self) if self.merge_weights and not self.merged: # Merge the weights and mark it @@ -238,7 +252,7 @@ class MergedLinear(nn.Linear, LoRALayer): delta_w = F.conv1d( self.lora_A.data.unsqueeze(0), self.lora_B.data.unsqueeze(-1), - groups=sum(self.enable_lora) + groups=sum(self.enable_lora), ).squeeze(0) self.weight.data += self.zero_pad(T(delta_w * self.scaling)) self.merged = True @@ -246,6 +260,7 @@ class MergedLinear(nn.Linear, LoRALayer): def forward(self, x: torch.Tensor): def T(w): return w.T if self.fan_in_fan_out else w + if self.merged: return F.linear(x, T(self.weight), bias=self.bias) else: @@ -255,7 +270,7 @@ class MergedLinear(nn.Linear, LoRALayer): after_B = F.conv1d( after_A.transpose(-2, -1), self.lora_B.unsqueeze(-1), - groups=sum(self.enable_lora) + groups=sum(self.enable_lora), ).transpose(-2, -1) result += self.zero_pad(after_B) * self.scaling return result @@ -270,21 +285,22 @@ class Conv2d(nn.Conv2d, LoRALayer): kernel_size: int, r: int = 0, lora_alpha: int = 1, - lora_dropout: float = 0., + lora_dropout: float = 0.0, merge_weights: bool = True, **kwargs ): nn.Conv2d.__init__(self, in_channels, out_channels, kernel_size, **kwargs) - LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, - merge_weights=merge_weights) + LoRALayer.__init__( + self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=merge_weights + ) assert type(kernel_size) is int # Actual trainable parameters if r > 0: self.lora_A = nn.Parameter( - self.weight.new_zeros((r*kernel_size, in_channels*kernel_size)) + self.weight.new_zeros((r * kernel_size, in_channels * kernel_size)) ) self.lora_B = nn.Parameter( - self.weight.new_zeros((out_channels*kernel_size, r*kernel_size)) + self.weight.new_zeros((out_channels * kernel_size, r * kernel_size)) ) self.scaling = self.lora_alpha / self.r # Freezing the pre-trained weight matrix @@ -293,7 +309,7 @@ class Conv2d(nn.Conv2d, LoRALayer): def reset_parameters(self): nn.Conv2d.reset_parameters(self) - if hasattr(self, 'lora_A'): + if hasattr(self, "lora_A"): # initialize A the same way as the default for nn.Linear and B to zero nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) nn.init.zeros_(self.lora_B) @@ -317,7 +333,10 @@ class Conv2d(nn.Conv2d, LoRALayer): return F.conv2d( x, self.weight + (self.lora_B @ self.lora_A).view(self.weight.shape) * self.scaling, - self.bias, self.stride, self.padding, self.dilation, self.groups + self.bias, + self.stride, + self.padding, + self.dilation, + self.groups, ) return nn.Conv2d.forward(self, x) - diff --git a/funasr/models/lora/utils.py b/funasr/models/lora/utils.py index e18bf44f6..670c1dcc0 100644 --- a/funasr/models/lora/utils.py +++ b/funasr/models/lora/utils.py @@ -10,41 +10,38 @@ from typing import Dict from .layers import LoRALayer -def mark_only_lora_as_trainable(model: nn.Module, bias: str = 'none') -> None: +def mark_only_lora_as_trainable(model: nn.Module, bias: str = "none") -> None: for n, p in model.named_parameters(): - if 'lora_' not in n and 'cif' not in n: + if "lora_" not in n and "cif" not in n: p.requires_grad = False - if bias == 'none': + if bias == "none": return - elif bias == 'all': + elif bias == "all": for n, p in model.named_parameters(): - if 'bias' in n: + if "bias" in n: p.requires_grad = True - elif bias == 'lora_only': + elif bias == "lora_only": for m in model.modules(): - if isinstance(m, LoRALayer) and \ - hasattr(m, 'bias') and \ - m.bias is not None: - m.bias.requires_grad = True + if isinstance(m, LoRALayer) and hasattr(m, "bias") and m.bias is not None: + m.bias.requires_grad = True else: raise NotImplementedError -def lora_state_dict(model: nn.Module, bias: str = 'none') -> Dict[str, torch.Tensor]: +def lora_state_dict(model: nn.Module, bias: str = "none") -> Dict[str, torch.Tensor]: my_state_dict = model.state_dict() - if bias == 'none': - return {k: my_state_dict[k] for k in my_state_dict if 'lora_' in k} - elif bias == 'all': - return {k: my_state_dict[k] for k in my_state_dict if 'lora_' in k or 'bias' in k} - elif bias == 'lora_only': + if bias == "none": + return {k: my_state_dict[k] for k in my_state_dict if "lora_" in k} + elif bias == "all": + return {k: my_state_dict[k] for k in my_state_dict if "lora_" in k or "bias" in k} + elif bias == "lora_only": to_return = {} for k in my_state_dict: - if 'lora_' in k: + if "lora_" in k: to_return[k] = my_state_dict[k] - bias_name = k.split('lora_')[0]+'bias' + bias_name = k.split("lora_")[0] + "bias" if bias_name in my_state_dict: to_return[bias_name] = my_state_dict[bias_name] return to_return else: raise NotImplementedError - diff --git a/funasr/models/mfcca/e2e_asr_mfcca.py b/funasr/models/mfcca/e2e_asr_mfcca.py index 48534ddb4..681d1f61f 100644 --- a/funasr/models/mfcca/e2e_asr_mfcca.py +++ b/funasr/models/mfcca/e2e_asr_mfcca.py @@ -31,10 +31,13 @@ else: @contextmanager def autocast(enabled=True): yield + + import pdb import random import math + class MFCCA(FunASRModel): """ Author: Audio, Speech and Language Processing Group (ASLP@NPU), Northwestern Polytechnical University @@ -43,26 +46,26 @@ class MFCCA(FunASRModel): """ def __init__( - self, - vocab_size: int, - token_list: Union[Tuple[str, ...], List[str]], - frontend: Optional[AbsFrontend], - specaug: Optional[AbsSpecAug], - normalize: Optional[AbsNormalize], - encoder: AbsEncoder, - decoder: AbsDecoder, - ctc: CTC, - rnnt_decoder: None = None, - ctc_weight: float = 0.5, - ignore_id: int = -1, - lsm_weight: float = 0.0, - mask_ratio: float = 0.0, - length_normalized_loss: bool = False, - report_cer: bool = True, - report_wer: bool = True, - sym_space: str = "", - sym_blank: str = "", - preencoder: Optional[AbsPreEncoder] = None, + self, + vocab_size: int, + token_list: Union[Tuple[str, ...], List[str]], + frontend: Optional[AbsFrontend], + specaug: Optional[AbsSpecAug], + normalize: Optional[AbsNormalize], + encoder: AbsEncoder, + decoder: AbsDecoder, + ctc: CTC, + rnnt_decoder: None = None, + ctc_weight: float = 0.5, + ignore_id: int = -1, + lsm_weight: float = 0.0, + mask_ratio: float = 0.0, + length_normalized_loss: bool = False, + report_cer: bool = True, + report_wer: bool = True, + sym_space: str = "", + sym_blank: str = "", + preencoder: Optional[AbsPreEncoder] = None, ): assert 0.0 <= ctc_weight <= 1.0, ctc_weight assert rnnt_decoder is None, "Not implemented" @@ -111,11 +114,11 @@ class MFCCA(FunASRModel): self.error_calculator = None def forward( - self, - speech: torch.Tensor, - speech_lengths: torch.Tensor, - text: torch.Tensor, - text_lengths: torch.Tensor, + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + text: torch.Tensor, + text_lengths: torch.Tensor, ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]: """Frontend + Encoder + Decoder + Calc loss Args: @@ -127,18 +130,15 @@ class MFCCA(FunASRModel): assert text_lengths.dim() == 1, text_lengths.shape # Check that batch_size is unified assert ( - speech.shape[0] - == speech_lengths.shape[0] - == text.shape[0] - == text_lengths.shape[0] + speech.shape[0] == speech_lengths.shape[0] == text.shape[0] == text_lengths.shape[0] ), (speech.shape, speech_lengths.shape, text.shape, text_lengths.shape) # pdb.set_trace() - if (speech.dim() == 3 and speech.size(2) == 8 and self.mask_ratio != 0): + if speech.dim() == 3 and speech.size(2) == 8 and self.mask_ratio != 0: rate_num = random.random() # rate_num = 0.1 - if (rate_num <= self.mask_ratio): + if rate_num <= self.mask_ratio: retain_channel = math.ceil(random.random() * 8) - if (retain_channel > 1): + if retain_channel > 1: speech = speech[:, :, torch.randperm(8)[0:retain_channel].sort().values] else: speech = speech[:, :, torch.randperm(8)[0]] @@ -192,17 +192,17 @@ class MFCCA(FunASRModel): return loss, stats, weight def collect_feats( - self, - speech: torch.Tensor, - speech_lengths: torch.Tensor, - text: torch.Tensor, - text_lengths: torch.Tensor, + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + text: torch.Tensor, + text_lengths: torch.Tensor, ) -> Dict[str, torch.Tensor]: feats, feats_lengths, channel_size = self._extract_feats(speech, speech_lengths) return {"feats": feats, "feats_lengths": feats_lengths} def encode( - self, speech: torch.Tensor, speech_lengths: torch.Tensor + self, speech: torch.Tensor, speech_lengths: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: """Frontend + Encoder. Note that this method is used by asr_inference.py Args: @@ -230,7 +230,7 @@ class MFCCA(FunASRModel): encoder_out.size(), speech.size(0), ) - if (encoder_out.dim() == 4): + if encoder_out.dim() == 4: assert encoder_out.size(2) <= encoder_out_lens.max(), ( encoder_out.size(), encoder_out_lens.max(), @@ -244,7 +244,7 @@ class MFCCA(FunASRModel): return encoder_out, encoder_out_lens def _extract_feats( - self, speech: torch.Tensor, speech_lengths: torch.Tensor + self, speech: torch.Tensor, speech_lengths: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: assert speech_lengths.dim() == 1, speech_lengths.shape # for data-parallel @@ -262,19 +262,17 @@ class MFCCA(FunASRModel): return feats, feats_lengths, channel_size def _calc_att_loss( - self, - encoder_out: torch.Tensor, - encoder_out_lens: torch.Tensor, - ys_pad: torch.Tensor, - ys_pad_lens: torch.Tensor, + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_pad_lens: torch.Tensor, ): ys_in_pad, ys_out_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id) ys_in_lens = ys_pad_lens + 1 # 1. Forward decoder - decoder_out, _ = self.decoder( - encoder_out, encoder_out_lens, ys_in_pad, ys_in_lens - ) + decoder_out, _ = self.decoder(encoder_out, encoder_out_lens, ys_in_pad, ys_in_lens) # 2. Compute attention loss loss_att = self.criterion_att(decoder_out, ys_out_pad) @@ -294,14 +292,14 @@ class MFCCA(FunASRModel): return loss_att, acc_att, cer_att, wer_att def _calc_ctc_loss( - self, - encoder_out: torch.Tensor, - encoder_out_lens: torch.Tensor, - ys_pad: torch.Tensor, - ys_pad_lens: torch.Tensor, + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_pad_lens: torch.Tensor, ): # Calc CTC loss - if (encoder_out.dim() == 4): + if encoder_out.dim() == 4: encoder_out = encoder_out.mean(1) loss_ctc = self.ctc(encoder_out, encoder_out_lens, ys_pad, ys_pad_lens) @@ -313,10 +311,10 @@ class MFCCA(FunASRModel): return loss_ctc, cer_ctc def _calc_rnnt_loss( - self, - encoder_out: torch.Tensor, - encoder_out_lens: torch.Tensor, - ys_pad: torch.Tensor, - ys_pad_lens: torch.Tensor, + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_pad_lens: torch.Tensor, ): - raise NotImplementedError \ No newline at end of file + raise NotImplementedError diff --git a/funasr/models/mfcca/encoder_layer_mfcca.py b/funasr/models/mfcca/encoder_layer_mfcca.py index cc86f89a7..f16a4e644 100644 --- a/funasr/models/mfcca/encoder_layer_mfcca.py +++ b/funasr/models/mfcca/encoder_layer_mfcca.py @@ -15,7 +15,6 @@ from funasr.models.transformer.layer_norm import LayerNorm from torch.autograd import Variable - class Encoder_Conformer_Layer(nn.Module): """Encoder layer module. @@ -111,7 +110,6 @@ class Encoder_Conformer_Layer(nn.Module): if self.normalize_before: x = self.norm_mha(x) - if cache is None: x_q = x else: @@ -120,12 +118,12 @@ class Encoder_Conformer_Layer(nn.Module): residual = residual[:, -1:, :] mask = None if mask is None else mask[:, -1:, :] - if self.cca_pos<2: + if self.cca_pos < 2: if pos_emb is not None: x_att = self.self_attn(x_q, x, x, pos_emb, mask) else: x_att = self.self_attn(x_q, x, x, mask) - else: + else: x_att = self.self_attn(x_q, x, x, mask) if self.concat_after: @@ -165,8 +163,6 @@ class Encoder_Conformer_Layer(nn.Module): return x, mask - - class EncoderLayer(nn.Module): """Encoder layer module. @@ -209,19 +205,19 @@ class EncoderLayer(nn.Module): self.encoder_cros_channel_atten = self_attn_cros_channel self.encoder_csa = Encoder_Conformer_Layer( - size, - self_attn_conformer, - feed_forward_csa, - feed_forward_macaron_csa, - conv_module_csa, - dropout_rate, - normalize_before, - concat_after, - cca_pos=0) + size, + self_attn_conformer, + feed_forward_csa, + feed_forward_macaron_csa, + conv_module_csa, + dropout_rate, + normalize_before, + concat_after, + cca_pos=0, + ) self.norm_mha = LayerNorm(size) # for the MHA module self.dropout = nn.Dropout(dropout_rate) - def forward(self, x_input, mask, channel_size, cache=None): """Compute encoded features. @@ -245,26 +241,33 @@ class EncoderLayer(nn.Module): x = self.norm_mha(x) t_leng = x.size(1) d_dim = x.size(2) - x_new = x.reshape(-1,channel_size,t_leng,d_dim).transpose(1,2) # x_new B*T * C * D - x_k_v = x_new.new(x_new.size(0),x_new.size(1),5,x_new.size(2),x_new.size(3)) - pad_before = Variable(torch.zeros(x_new.size(0),2,x_new.size(2),x_new.size(3))).type(x_new.type()) - pad_after = Variable(torch.zeros(x_new.size(0),2,x_new.size(2),x_new.size(3))).type(x_new.type()) - x_pad = torch.cat([pad_before,x_new, pad_after], 1) - x_k_v[:,:,0,:,:]=x_pad[:,0:-4,:,:] - x_k_v[:,:,1,:,:]=x_pad[:,1:-3,:,:] - x_k_v[:,:,2,:,:]=x_pad[:,2:-2,:,:] - x_k_v[:,:,3,:,:]=x_pad[:,3:-1,:,:] - x_k_v[:,:,4,:,:]=x_pad[:,4:,:,:] - x_new = x_new.reshape(-1,channel_size,d_dim) - x_k_v = x_k_v.reshape(-1,5*channel_size,d_dim) + x_new = x.reshape(-1, channel_size, t_leng, d_dim).transpose(1, 2) # x_new B*T * C * D + x_k_v = x_new.new(x_new.size(0), x_new.size(1), 5, x_new.size(2), x_new.size(3)) + pad_before = Variable(torch.zeros(x_new.size(0), 2, x_new.size(2), x_new.size(3))).type( + x_new.type() + ) + pad_after = Variable(torch.zeros(x_new.size(0), 2, x_new.size(2), x_new.size(3))).type( + x_new.type() + ) + x_pad = torch.cat([pad_before, x_new, pad_after], 1) + x_k_v[:, :, 0, :, :] = x_pad[:, 0:-4, :, :] + x_k_v[:, :, 1, :, :] = x_pad[:, 1:-3, :, :] + x_k_v[:, :, 2, :, :] = x_pad[:, 2:-2, :, :] + x_k_v[:, :, 3, :, :] = x_pad[:, 3:-1, :, :] + x_k_v[:, :, 4, :, :] = x_pad[:, 4:, :, :] + x_new = x_new.reshape(-1, channel_size, d_dim) + x_k_v = x_k_v.reshape(-1, 5 * channel_size, d_dim) x_att = self.encoder_cros_channel_atten(x_new, x_k_v, x_k_v, None) - x_att = x_att.reshape(-1,t_leng,channel_size,d_dim).transpose(1,2).reshape(-1,t_leng,d_dim) + x_att = ( + x_att.reshape(-1, t_leng, channel_size, d_dim) + .transpose(1, 2) + .reshape(-1, t_leng, d_dim) + ) x = residual + self.dropout(x_att) if pos_emb is not None: - x_input = (x, pos_emb) + x_input = (x, pos_emb) else: x_input = x x_input, mask = self.encoder_csa(x_input, mask) - - return x_input, mask , channel_size + return x_input, mask, channel_size diff --git a/funasr/models/mfcca/mfcca_encoder.py b/funasr/models/mfcca/mfcca_encoder.py index 52da33d79..19a6df92a 100644 --- a/funasr/models/mfcca/mfcca_encoder.py +++ b/funasr/models/mfcca/mfcca_encoder.py @@ -136,29 +136,29 @@ class MFCCAEncoder(AbsEncoder): """ def __init__( - self, - input_size: int, - output_size: int = 256, - attention_heads: int = 4, - linear_units: int = 2048, - num_blocks: int = 6, - dropout_rate: float = 0.1, - positional_dropout_rate: float = 0.1, - attention_dropout_rate: float = 0.0, - input_layer: str = "conv2d", - normalize_before: bool = True, - concat_after: bool = False, - positionwise_layer_type: str = "linear", - positionwise_conv_kernel_size: int = 3, - macaron_style: bool = False, - rel_pos_type: str = "legacy", - pos_enc_layer_type: str = "rel_pos", - selfattention_layer_type: str = "rel_selfattn", - activation_type: str = "swish", - use_cnn_module: bool = True, - zero_triu: bool = False, - cnn_module_kernel: int = 31, - padding_idx: int = -1, + self, + input_size: int, + output_size: int = 256, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + attention_dropout_rate: float = 0.0, + input_layer: str = "conv2d", + normalize_before: bool = True, + concat_after: bool = False, + positionwise_layer_type: str = "linear", + positionwise_conv_kernel_size: int = 3, + macaron_style: bool = False, + rel_pos_type: str = "legacy", + pos_enc_layer_type: str = "rel_pos", + selfattention_layer_type: str = "rel_selfattn", + activation_type: str = "swish", + use_cnn_module: bool = True, + zero_triu: bool = False, + cnn_module_kernel: int = 31, + padding_idx: int = -1, ): super().__init__() self._output_size = output_size @@ -185,9 +185,7 @@ class MFCCAEncoder(AbsEncoder): elif pos_enc_layer_type == "legacy_rel_pos": assert selfattention_layer_type == "legacy_rel_selfattn" pos_enc_class = LegacyRelPositionalEncoding - logging.warning( - "Using legacy_rel_pos and it will be deprecated in the future." - ) + logging.warning("Using legacy_rel_pos and it will be deprecated in the future.") else: raise ValueError("unknown pos_enc_layer: " + pos_enc_layer_type) @@ -230,9 +228,7 @@ class MFCCAEncoder(AbsEncoder): pos_enc_class(output_size, positional_dropout_rate), ) elif input_layer is None: - self.embed = torch.nn.Sequential( - pos_enc_class(output_size, positional_dropout_rate) - ) + self.embed = torch.nn.Sequential(pos_enc_class(output_size, positional_dropout_rate)) else: raise ValueError("unknown input_layer: " + input_layer) self.normalize_before = normalize_before @@ -278,9 +274,7 @@ class MFCCAEncoder(AbsEncoder): output_size, attention_dropout_rate, ) - logging.warning( - "Using legacy_rel_selfattn and it will be deprecated in the future." - ) + logging.warning("Using legacy_rel_selfattn and it will be deprecated in the future.") elif selfattention_layer_type == "rel_selfattn": assert pos_enc_layer_type == "rel_pos" encoder_selfattn_layer = RelPositionMultiHeadedAttention @@ -329,11 +323,11 @@ class MFCCAEncoder(AbsEncoder): return self._output_size def forward( - self, - xs_pad: torch.Tensor, - ilens: torch.Tensor, - channel_size: torch.Tensor, - prev_states: torch.Tensor = None, + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, + channel_size: torch.Tensor, + prev_states: torch.Tensor = None, ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: """Calculate forward propagation. Args: @@ -347,9 +341,9 @@ class MFCCAEncoder(AbsEncoder): """ masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device) if ( - isinstance(self.embed, Conv2dSubsampling) - or isinstance(self.embed, Conv2dSubsampling6) - or isinstance(self.embed, Conv2dSubsampling8) + isinstance(self.embed, Conv2dSubsampling) + or isinstance(self.embed, Conv2dSubsampling6) + or isinstance(self.embed, Conv2dSubsampling8) ): short_status, limit_size = check_short_utt(self.embed, xs_pad.size(1)) if short_status: @@ -370,7 +364,7 @@ class MFCCAEncoder(AbsEncoder): d_dim = xs_pad.size(2) xs_pad = xs_pad.reshape(-1, channel_size, t_leng, d_dim) # pdb.set_trace() - if (channel_size < 8): + if channel_size < 8: repeat_num = math.ceil(8 / channel_size) xs_pad = xs_pad.repeat(1, repeat_num, 1, 1)[:, 0:8, :, :] xs_pad = self.conv1(xs_pad) @@ -388,10 +382,10 @@ class MFCCAEncoder(AbsEncoder): return xs_pad, olens, None def forward_hidden( - self, - xs_pad: torch.Tensor, - ilens: torch.Tensor, - prev_states: torch.Tensor = None, + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, + prev_states: torch.Tensor = None, ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: """Calculate forward propagation. Args: @@ -405,9 +399,9 @@ class MFCCAEncoder(AbsEncoder): """ masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device) if ( - isinstance(self.embed, Conv2dSubsampling) - or isinstance(self.embed, Conv2dSubsampling6) - or isinstance(self.embed, Conv2dSubsampling8) + isinstance(self.embed, Conv2dSubsampling) + or isinstance(self.embed, Conv2dSubsampling6) + or isinstance(self.embed, Conv2dSubsampling8) ): short_status, limit_size = check_short_utt(self.embed, xs_pad.size(1)) if short_status: @@ -433,4 +427,4 @@ class MFCCAEncoder(AbsEncoder): self.hidden_feature = self.after_norm(hidden_feature) olens = masks.squeeze(1).sum(1) - return xs_pad, olens, None \ No newline at end of file + return xs_pad, olens, None diff --git a/funasr/models/monotonic_aligner/model.py b/funasr/models/monotonic_aligner/model.py index 718923bd6..4754d1fb3 100644 --- a/funasr/models/monotonic_aligner/model.py +++ b/funasr/models/monotonic_aligner/model.py @@ -28,6 +28,7 @@ class MonotonicAligner(torch.nn.Module): Achieving timestamp prediction while recognizing with non-autoregressive end-to-end ASR model https://arxiv.org/abs/2301.12343 """ + def __init__( self, input_size: int = 80, @@ -64,11 +65,11 @@ class MonotonicAligner(torch.nn.Module): self.predictor_bias = predictor_bias def forward( - self, - speech: torch.Tensor, - speech_lengths: torch.Tensor, - text: torch.Tensor, - text_lengths: torch.Tensor, + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + text: torch.Tensor, + text_lengths: torch.Tensor, ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]: """Frontend + Encoder + Decoder + Calc loss Args: @@ -80,25 +81,25 @@ class MonotonicAligner(torch.nn.Module): assert text_lengths.dim() == 1, text_lengths.shape # Check that batch_size is unified assert ( - speech.shape[0] - == speech_lengths.shape[0] - == text.shape[0] - == text_lengths.shape[0] + speech.shape[0] == speech_lengths.shape[0] == text.shape[0] == text_lengths.shape[0] ), (speech.shape, speech_lengths.shape, text.shape, text_lengths.shape) batch_size = speech.shape[0] # for data-parallel text = text[:, : text_lengths.max()] - speech = speech[:, :speech_lengths.max()] + speech = speech[:, : speech_lengths.max()] # 1. Encoder encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) - encoder_out_mask = (~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :]).to( - encoder_out.device) + encoder_out_mask = ( + ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :] + ).to(encoder_out.device) if self.predictor_bias == 1: _, text = add_sos_eos(text, 1, 2, -1) text_lengths = text_lengths + self.predictor_bias - _, _, _, _, pre_token_length2 = self.predictor(encoder_out, text, encoder_out_mask, ignore_id=-1) + _, _, _, _, pre_token_length2 = self.predictor( + encoder_out, text, encoder_out_mask, ignore_id=-1 + ) # loss_pre = self.criterion_pre(ys_pad_lens.type_as(pre_token_length), pre_token_length) loss_pre = self.criterion_pre(text_lengths.type_as(pre_token_length2), pre_token_length2) @@ -115,15 +116,19 @@ class MonotonicAligner(torch.nn.Module): return loss, stats, weight def calc_predictor_timestamp(self, encoder_out, encoder_out_lens, token_num): - encoder_out_mask = (~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :]).to( - encoder_out.device) - ds_alphas, ds_cif_peak, us_alphas, us_peaks = self.predictor.get_upsample_timestamp(encoder_out, - encoder_out_mask, - token_num) + encoder_out_mask = ( + ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :] + ).to(encoder_out.device) + ds_alphas, ds_cif_peak, us_alphas, us_peaks = self.predictor.get_upsample_timestamp( + encoder_out, encoder_out_mask, token_num + ) return ds_alphas, ds_cif_peak, us_alphas, us_peaks def encode( - self, speech: torch.Tensor, speech_lengths: torch.Tensor, **kwargs, + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + **kwargs, ) -> Tuple[torch.Tensor, torch.Tensor]: """Encoder. Note that this method is used by asr_inference.py Args: @@ -136,51 +141,62 @@ class MonotonicAligner(torch.nn.Module): # Data augmentation if self.specaug is not None and self.training: speech, speech_lengths = self.specaug(speech, speech_lengths) - + # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN if self.normalize is not None: speech, speech_lengths = self.normalize(speech, speech_lengths) - + # Forward encoder encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths) if isinstance(encoder_out, tuple): encoder_out = encoder_out[0] return encoder_out, encoder_out_lens - - def inference(self, - data_in, - data_lengths=None, - key: list=None, - tokenizer=None, - frontend=None, - **kwargs, - ): + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): meta_data = {} # extract fbank feats time1 = time.perf_counter() - audio_list, text_token_int_list = load_audio_text_image_video(data_in, - fs=frontend.fs, - audio_fs=kwargs.get("fs", 16000), - data_type=kwargs.get("data_type", "sound"), - tokenizer=tokenizer) + audio_list, text_token_int_list = load_audio_text_image_video( + data_in, + fs=frontend.fs, + audio_fs=kwargs.get("fs", 16000), + data_type=kwargs.get("data_type", "sound"), + tokenizer=tokenizer, + ) time2 = time.perf_counter() meta_data["load_data"] = f"{time2 - time1:0.3f}" - speech, speech_lengths = extract_fbank(audio_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend) + speech, speech_lengths = extract_fbank( + audio_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend + ) time3 = time.perf_counter() meta_data["extract_feat"] = f"{time3 - time2:0.3f}" - meta_data["batch_data_time"] = speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 - + meta_data["batch_data_time"] = ( + speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 + ) + speech = speech.to(device=kwargs["device"]) speech_lengths = speech_lengths.to(device=kwargs["device"]) # Encoder encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) - + # predictor - text_lengths = torch.tensor([len(i)+1 for i in text_token_int_list]).to(encoder_out.device) - _, _, us_alphas, us_peaks = self.calc_predictor_timestamp(encoder_out, encoder_out_lens, token_num=text_lengths) - + text_lengths = torch.tensor([len(i) + 1 for i in text_token_int_list]).to( + encoder_out.device + ) + _, _, us_alphas, us_peaks = self.calc_predictor_timestamp( + encoder_out, encoder_out_lens, token_num=text_lengths + ) + results = [] ibest_writer = None if kwargs.get("output_dir") is not None: @@ -188,20 +204,28 @@ class MonotonicAligner(torch.nn.Module): self.writer = DatadirWriter(kwargs.get("output_dir")) ibest_writer = self.writer["tp_res"] - for i, (us_alpha, us_peak, token_int) in enumerate(zip(us_alphas, us_peaks, text_token_int_list)): + for i, (us_alpha, us_peak, token_int) in enumerate( + zip(us_alphas, us_peaks, text_token_int_list) + ): token = tokenizer.ids2tokens(token_int) - timestamp_str, timestamp = ts_prediction_lfr6_standard(us_alpha[:encoder_out_lens[i] * 3], - us_peak[:encoder_out_lens[i] * 3], - copy.copy(token)) - text_postprocessed, time_stamp_postprocessed, _ = postprocess_utils.sentence_postprocess(token, timestamp) - result_i = {"key": key[i], "text": text_postprocessed, - "timestamp": time_stamp_postprocessed, - } + timestamp_str, timestamp = ts_prediction_lfr6_standard( + us_alpha[: encoder_out_lens[i] * 3], + us_peak[: encoder_out_lens[i] * 3], + copy.copy(token), + ) + text_postprocessed, time_stamp_postprocessed, _ = ( + postprocess_utils.sentence_postprocess(token, timestamp) + ) + result_i = { + "key": key[i], + "text": text_postprocessed, + "timestamp": time_stamp_postprocessed, + } results.append(result_i) if ibest_writer: # ibest_writer["token"][key[i]] = " ".join(token) ibest_writer["timestamp_list"][key[i]] = time_stamp_postprocessed ibest_writer["timestamp_str"][key[i]] = timestamp_str - - return results, meta_data \ No newline at end of file + + return results, meta_data diff --git a/funasr/models/mossformer/e2e_ss.py b/funasr/models/mossformer/e2e_ss.py index 1a46b3ffd..40d30ca89 100644 --- a/funasr/models/mossformer/e2e_ss.py +++ b/funasr/models/mossformer/e2e_ss.py @@ -4,7 +4,7 @@ import torch.nn as nn import torch.nn.functional as F import copy from funasr.models.base_model import FunASRModel -from funasr.models.encoder.mossformer_encoder import MossFormerEncoder, MossFormer_MaskNet +from funasr.models.encoder.mossformer_encoder import MossFormerEncoder, MossFormer_MaskNet from funasr.models.decoder.mossformer_decoder import MossFormerDecoder @@ -48,7 +48,9 @@ class MossFormer(FunASRModel): super(MossFormer, self).__init__() self.num_spks = num_spks # Encoding - self.enc = MossFormerEncoder(kernel_size=kernel_size, out_channels=in_channels, in_channels=1) + self.enc = MossFormerEncoder( + kernel_size=kernel_size, out_channels=in_channels, in_channels=1 + ) ##Compute Mask self.mask_net = MossFormer_MaskNet( @@ -62,12 +64,13 @@ class MossFormer(FunASRModel): max_length=max_length, ) self.dec = MossFormerDecoder( - in_channels=out_channels, - out_channels=1, - kernel_size=kernel_size, - stride = kernel_size//2, - bias=False + in_channels=out_channels, + out_channels=1, + kernel_size=kernel_size, + stride=kernel_size // 2, + bias=False, ) + def forward(self, input): x = self.enc(input) mask = self.mask_net(x) @@ -76,10 +79,7 @@ class MossFormer(FunASRModel): # Decoding est_source = torch.cat( - [ - self.dec(sep_x[i]).unsqueeze(-1) - for i in range(self.num_spks) - ], + [self.dec(sep_x[i]).unsqueeze(-1) for i in range(self.num_spks)], dim=-1, ) T_origin = input.size(1) @@ -91,5 +91,5 @@ class MossFormer(FunASRModel): out = [] for spk in range(self.num_spks): - out.append(est_source[:,:,spk]) + out.append(est_source[:, :, spk]) return out diff --git a/funasr/models/mossformer/mossformer.py b/funasr/models/mossformer/mossformer.py index f1e8e2861..30aab481a 100644 --- a/funasr/models/mossformer/mossformer.py +++ b/funasr/models/mossformer/mossformer.py @@ -7,17 +7,21 @@ from einops import rearrange def identity(t, *args, **kwargs): return t + def append_dims(x, num_dims): if num_dims <= 0: return x return x.view(*x.shape, *((1,) * num_dims)) + def exists(val): return val is not None + def default(val, d): return val if exists(val) else d + def padding_to_multiple_of(n, mult): remainder = n % mult if remainder == 0: @@ -26,7 +30,8 @@ def padding_to_multiple_of(n, mult): class Transpose(nn.Module): - """ Wrapper class of torch.transpose() for Sequential module. """ + """Wrapper class of torch.transpose() for Sequential module.""" + def __init__(self, shape: tuple): super(Transpose, self).__init__() self.shape = shape @@ -51,17 +56,20 @@ class DepthwiseConv1d(nn.Module): Returns: outputs - **outputs** (batch, out_channels, time): Tensor produces by depthwise 1-D convolution. """ + def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: int, - stride: int = 1, - padding: int = 0, - bias: bool = False, + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + padding: int = 0, + bias: bool = False, ) -> None: super(DepthwiseConv1d, self).__init__() - assert out_channels % in_channels == 0, "out_channels should be constant multiple of in_channels" + assert ( + out_channels % in_channels == 0 + ), "out_channels should be constant multiple of in_channels" self.conv = nn.Conv1d( in_channels=in_channels, out_channels=out_channels, @@ -90,12 +98,13 @@ class ConvModule(nn.Module): Outputs: outputs outputs (batch, time, dim): Tensor produces by conformer convolution module. """ + def __init__( - self, - in_channels: int, - kernel_size: int = 17, - expansion_factor: int = 2, - dropout_p: float = 0.1, + self, + in_channels: int, + kernel_size: int = 17, + expansion_factor: int = 2, + dropout_p: float = 0.1, ) -> None: super(ConvModule, self).__init__() assert (kernel_size - 1) % 2 == 0, "kernel_size should be a odd number for 'SAME' padding" @@ -103,7 +112,9 @@ class ConvModule(nn.Module): self.sequential = nn.Sequential( Transpose(shape=(1, 2)), - DepthwiseConv1d(in_channels, in_channels, kernel_size, stride=1, padding=(kernel_size - 1) // 2), + DepthwiseConv1d( + in_channels, in_channels, kernel_size, stride=1, padding=(kernel_size - 1) // 2 + ), ) def forward(self, inputs): @@ -111,33 +122,28 @@ class ConvModule(nn.Module): class OffsetScale(nn.Module): - def __init__(self, dim, heads = 1): + def __init__(self, dim, heads=1): super().__init__() self.gamma = nn.Parameter(torch.ones(heads, dim)) self.beta = nn.Parameter(torch.zeros(heads, dim)) - nn.init.normal_(self.gamma, std = 0.02) + nn.init.normal_(self.gamma, std=0.02) def forward(self, x): - out = einsum('... d, h d -> ... h d', x, self.gamma) + self.beta - return out.unbind(dim = -2) + out = einsum("... d, h d -> ... h d", x, self.gamma) + self.beta + return out.unbind(dim=-2) class FFConvM(nn.Module): - def __init__( - self, - dim_in, - dim_out, - norm_klass = nn.LayerNorm, - dropout = 0.1 - ): + def __init__(self, dim_in, dim_out, norm_klass=nn.LayerNorm, dropout=0.1): super().__init__() self.mdl = nn.Sequential( norm_klass(dim_in), nn.Linear(dim_in, dim_out), nn.SiLU(), ConvModule(dim_out), - nn.Dropout(dropout) + nn.Dropout(dropout), ) + def forward( self, x, @@ -151,17 +157,17 @@ class FLASH_ShareA_FFConvM(nn.Module): self, *, dim, - group_size = 256, - query_key_dim = 128, - expansion_factor = 1., - causal = False, - dropout = 0.1, - rotary_pos_emb = None, - norm_klass = nn.LayerNorm, - shift_tokens = True + group_size=256, + query_key_dim=128, + expansion_factor=1.0, + causal=False, + dropout=0.1, + rotary_pos_emb=None, + norm_klass=nn.LayerNorm, + shift_tokens=True ): super().__init__() - hidden_dim = int(dim * expansion_factor) + hidden_dim = int(dim * expansion_factor) self.group_size = group_size self.causal = causal self.shift_tokens = shift_tokens @@ -171,38 +177,32 @@ class FLASH_ShareA_FFConvM(nn.Module): # norm self.dropout = nn.Dropout(dropout) # projections - - self.to_hidden = FFConvM( - dim_in = dim, - dim_out = hidden_dim, - norm_klass = norm_klass, - dropout = dropout, - ) - self.to_qk = FFConvM( - dim_in = dim, - dim_out = query_key_dim, - norm_klass = norm_klass, - dropout = dropout, - ) - self.qk_offset_scale = OffsetScale(query_key_dim, heads = 4) + self.to_hidden = FFConvM( + dim_in=dim, + dim_out=hidden_dim, + norm_klass=norm_klass, + dropout=dropout, + ) + self.to_qk = FFConvM( + dim_in=dim, + dim_out=query_key_dim, + norm_klass=norm_klass, + dropout=dropout, + ) + + self.qk_offset_scale = OffsetScale(query_key_dim, heads=4) self.to_out = FFConvM( - dim_in = dim*2, - dim_out = dim, - norm_klass = norm_klass, - dropout = dropout, - ) - - self.gateActivate=nn.Sigmoid() + dim_in=dim * 2, + dim_out=dim, + norm_klass=norm_klass, + dropout=dropout, + ) - def forward( - self, - x, - *, - mask = None - ): + self.gateActivate = nn.Sigmoid() + def forward(self, x, *, mask=None): """ b - batch n - sequence length (within groups) @@ -213,95 +213,105 @@ class FLASH_ShareA_FFConvM(nn.Module): j - sequence dimension (target) """ - normed_x = x + normed_x = x # do token shift - a great, costless trick from an independent AI researcher in Shenzhen residual = x if self.shift_tokens: - x_shift, x_pass = normed_x.chunk(2, dim = -1) - x_shift = F.pad(x_shift, (0, 0, 1, -1), value = 0.) - normed_x = torch.cat((x_shift, x_pass), dim = -1) + x_shift, x_pass = normed_x.chunk(2, dim=-1) + x_shift = F.pad(x_shift, (0, 0, 1, -1), value=0.0) + normed_x = torch.cat((x_shift, x_pass), dim=-1) # initial projections - v, u = self.to_hidden(normed_x).chunk(2, dim = -1) + v, u = self.to_hidden(normed_x).chunk(2, dim=-1) qk = self.to_qk(normed_x) # offset and scale quad_q, lin_q, quad_k, lin_k = self.qk_offset_scale(qk) att_v, att_u = self.cal_attention(x, quad_q, lin_q, quad_k, lin_k, v, u) - out = (att_u*v ) * self.gateActivate(att_v*u) + out = (att_u * v) * self.gateActivate(att_v * u) x = x + self.to_out(out) return x - def cal_attention(self, x, quad_q, lin_q, quad_k, lin_k, v, u, mask = None): + def cal_attention(self, x, quad_q, lin_q, quad_k, lin_k, v, u, mask=None): b, n, device, g = x.shape[0], x.shape[-2], x.device, self.group_size if exists(mask): - lin_mask = rearrange(mask, '... -> ... 1') - lin_k = lin_k.masked_fill(~lin_mask, 0.) + lin_mask = rearrange(mask, "... -> ... 1") + lin_k = lin_k.masked_fill(~lin_mask, 0.0) # rotate queries and keys if exists(self.rotary_pos_emb): - quad_q, lin_q, quad_k, lin_k = map(self.rotary_pos_emb.rotate_queries_or_keys, (quad_q, lin_q, quad_k, lin_k)) + quad_q, lin_q, quad_k, lin_k = map( + self.rotary_pos_emb.rotate_queries_or_keys, (quad_q, lin_q, quad_k, lin_k) + ) # padding for groups padding = padding_to_multiple_of(n, g) if padding > 0: - quad_q, quad_k, lin_q, lin_k, v, u = map(lambda t: F.pad(t, (0, 0, 0, padding), value = 0.), (quad_q, quad_k, lin_q, lin_k, v, u)) + quad_q, quad_k, lin_q, lin_k, v, u = map( + lambda t: F.pad(t, (0, 0, 0, padding), value=0.0), + (quad_q, quad_k, lin_q, lin_k, v, u), + ) - mask = default(mask, torch.ones((b, n), device = device, dtype = torch.bool)) - mask = F.pad(mask, (0, padding), value = False) + mask = default(mask, torch.ones((b, n), device=device, dtype=torch.bool)) + mask = F.pad(mask, (0, padding), value=False) # group along sequence - quad_q, quad_k, lin_q, lin_k, v, u = map(lambda t: rearrange(t, 'b (g n) d -> b g n d', n = self.group_size), (quad_q, quad_k, lin_q, lin_k, v, u)) + quad_q, quad_k, lin_q, lin_k, v, u = map( + lambda t: rearrange(t, "b (g n) d -> b g n d", n=self.group_size), + (quad_q, quad_k, lin_q, lin_k, v, u), + ) if exists(mask): - mask = rearrange(mask, 'b (g j) -> b g 1 j', j = g) + mask = rearrange(mask, "b (g j) -> b g 1 j", j=g) # calculate quadratic attention output - sim = einsum('... i d, ... j d -> ... i j', quad_q, quad_k) / g + sim = einsum("... i d, ... j d -> ... i j", quad_q, quad_k) / g attn = F.relu(sim) ** 2 attn = self.dropout(attn) if exists(mask): - attn = attn.masked_fill(~mask, 0.) + attn = attn.masked_fill(~mask, 0.0) if self.causal: - causal_mask = torch.ones((g, g), dtype = torch.bool, device = device).triu(1) - attn = attn.masked_fill(causal_mask, 0.) + causal_mask = torch.ones((g, g), dtype=torch.bool, device=device).triu(1) + attn = attn.masked_fill(causal_mask, 0.0) - quad_out_v = einsum('... i j, ... j d -> ... i d', attn, v) - quad_out_u = einsum('... i j, ... j d -> ... i d', attn, u) + quad_out_v = einsum("... i j, ... j d -> ... i d", attn, v) + quad_out_u = einsum("... i j, ... j d -> ... i d", attn, u) # calculate linear attention output if self.causal: - lin_kv = einsum('b g n d, b g n e -> b g d e', lin_k, v) / g + lin_kv = einsum("b g n d, b g n e -> b g d e", lin_k, v) / g # exclusive cumulative sum along group dimension - lin_kv = lin_kv.cumsum(dim = 1) - lin_kv = F.pad(lin_kv, (0, 0, 0, 0, 1, -1), value = 0.) - lin_out_v = einsum('b g d e, b g n d -> b g n e', lin_kv, lin_q) + lin_kv = lin_kv.cumsum(dim=1) + lin_kv = F.pad(lin_kv, (0, 0, 0, 0, 1, -1), value=0.0) + lin_out_v = einsum("b g d e, b g n d -> b g n e", lin_kv, lin_q) - lin_ku = einsum('b g n d, b g n e -> b g d e', lin_k, u) / g + lin_ku = einsum("b g n d, b g n e -> b g d e", lin_k, u) / g # exclusive cumulative sum along group dimension - lin_ku = lin_ku.cumsum(dim = 1) - lin_ku = F.pad(lin_ku, (0, 0, 0, 0, 1, -1), value = 0.) - lin_out_u = einsum('b g d e, b g n d -> b g n e', lin_ku, lin_q) + lin_ku = lin_ku.cumsum(dim=1) + lin_ku = F.pad(lin_ku, (0, 0, 0, 0, 1, -1), value=0.0) + lin_out_u = einsum("b g d e, b g n d -> b g n e", lin_ku, lin_q) else: - lin_kv = einsum('b g n d, b g n e -> b d e', lin_k, v) / n - lin_out_v = einsum('b g n d, b d e -> b g n e', lin_q, lin_kv) + lin_kv = einsum("b g n d, b g n e -> b d e", lin_k, v) / n + lin_out_v = einsum("b g n d, b d e -> b g n e", lin_q, lin_kv) - lin_ku = einsum('b g n d, b g n e -> b d e', lin_k, u) / n - lin_out_u = einsum('b g n d, b d e -> b g n e', lin_q, lin_ku) + lin_ku = einsum("b g n d, b g n e -> b d e", lin_k, u) / n + lin_out_u = einsum("b g n d, b d e -> b g n e", lin_q, lin_ku) # fold back groups into full sequence, and excise out padding - return map(lambda t: rearrange(t, 'b g n d -> b (g n) d')[:, :n], (quad_out_v+lin_out_v, quad_out_u+lin_out_u)) - + return map( + lambda t: rearrange(t, "b g n d -> b (g n) d")[:, :n], + (quad_out_v + lin_out_v, quad_out_u + lin_out_u), + ) diff --git a/funasr/models/mossformer/mossformer_decoder.py b/funasr/models/mossformer/mossformer_decoder.py index e0189f717..56ad2ecb3 100644 --- a/funasr/models/mossformer/mossformer_decoder.py +++ b/funasr/models/mossformer/mossformer_decoder.py @@ -40,9 +40,7 @@ class MossFormerDecoder(nn.ConvTranspose1d): """ if x.dim() not in [2, 3]: - raise RuntimeError( - "{} accept 3/4D tensor as input".format(self.__name__) - ) + raise RuntimeError("{} accept 3/4D tensor as input".format(self.__name__)) x = super().forward(x if x.dim() == 3 else torch.unsqueeze(x, 1)) if torch.squeeze(x).dim() == 1: @@ -50,4 +48,3 @@ class MossFormerDecoder(nn.ConvTranspose1d): else: x = torch.squeeze(x) return x - diff --git a/funasr/models/mossformer/mossformer_encoder.py b/funasr/models/mossformer/mossformer_encoder.py index a28c960e8..888aef831 100644 --- a/funasr/models/mossformer/mossformer_encoder.py +++ b/funasr/models/mossformer/mossformer_encoder.py @@ -1,18 +1,20 @@ import torch import torch.nn as nn import torch.nn.functional as F + try: from rotary_embedding_torch import RotaryEmbedding except: - print("If you want use mossformer, lease install rotary_embedding_torch by: \n pip install -U rotary_embedding_torch") + print( + "If you want use mossformer, lease install rotary_embedding_torch by: \n pip install -U rotary_embedding_torch" + ) from funasr.models.transformer.layer_norm import GlobalLayerNorm, CumulativeLayerNorm, ScaleNorm from funasr.models.transformer.embedding import ScaledSinuEmbedding from funasr.models.transformer.mossformer import FLASH_ShareA_FFConvM def select_norm(norm, dim, shape): - """Just a wrapper to select the normalization type. - """ + """Just a wrapper to select the normalization type.""" if norm == "gln": return GlobalLayerNorm(dim, shape, elementwise_affine=True) @@ -30,37 +32,50 @@ class MossformerBlock(nn.Module): *, dim, depth, - group_size = 256, - query_key_dim = 128, - expansion_factor = 4., - causal = False, - attn_dropout = 0.1, - norm_type = 'scalenorm', - shift_tokens = True + group_size=256, + query_key_dim=128, + expansion_factor=4.0, + causal=False, + attn_dropout=0.1, + norm_type="scalenorm", + shift_tokens=True ): super().__init__() - assert norm_type in ('scalenorm', 'layernorm'), 'norm_type must be one of scalenorm or layernorm' + assert norm_type in ( + "scalenorm", + "layernorm", + ), "norm_type must be one of scalenorm or layernorm" - if norm_type == 'scalenorm': + if norm_type == "scalenorm": norm_klass = ScaleNorm - elif norm_type == 'layernorm': + elif norm_type == "layernorm": norm_klass = nn.LayerNorm self.group_size = group_size - rotary_pos_emb = RotaryEmbedding(dim = min(32, query_key_dim)) + rotary_pos_emb = RotaryEmbedding(dim=min(32, query_key_dim)) # max rotary embedding dimensions of 32, partial Rotary embeddings, from Wang et al - GPT-J - self.layers = nn.ModuleList([FLASH_ShareA_FFConvM(dim = dim, group_size = group_size, query_key_dim = query_key_dim, expansion_factor = expansion_factor, causal = causal, dropout = attn_dropout, rotary_pos_emb = rotary_pos_emb, norm_klass = norm_klass, shift_tokens = shift_tokens) for _ in range(depth)]) + self.layers = nn.ModuleList( + [ + FLASH_ShareA_FFConvM( + dim=dim, + group_size=group_size, + query_key_dim=query_key_dim, + expansion_factor=expansion_factor, + causal=causal, + dropout=attn_dropout, + rotary_pos_emb=rotary_pos_emb, + norm_klass=norm_klass, + shift_tokens=shift_tokens, + ) + for _ in range(depth) + ] + ) - def forward( - self, - x, - *, - mask = None - ): + def forward(self, x, *, mask=None): ii = 0 for flash in self.layers: - x = flash(x, mask = mask) + x = flash(x, mask=mask) ii = ii + 1 return x @@ -119,25 +134,19 @@ class MossFormer_MaskNet(nn.Module): self.pos_enc = ScaledSinuEmbedding(out_channels) self.mdl = Computation_Block( - num_blocks, - out_channels, - norm, - skip_around_intra=skip_around_intra, - ) - - self.conv1d_out = nn.Conv1d( - out_channels, out_channels * num_spks, kernel_size=1 + num_blocks, + out_channels, + norm, + skip_around_intra=skip_around_intra, ) + + self.conv1d_out = nn.Conv1d(out_channels, out_channels * num_spks, kernel_size=1) self.conv1_decoder = nn.Conv1d(out_channels, in_channels, 1, bias=False) self.prelu = nn.PReLU() self.activation = nn.ReLU() # gated output layer - self.output = nn.Sequential( - nn.Conv1d(out_channels, out_channels, 1), nn.Tanh() - ) - self.output_gate = nn.Sequential( - nn.Conv1d(out_channels, out_channels, 1), nn.Sigmoid() - ) + self.output = nn.Sequential(nn.Conv1d(out_channels, out_channels, 1), nn.Tanh()) + self.output_gate = nn.Sequential(nn.Conv1d(out_channels, out_channels, 1), nn.Sigmoid()) def forward(self, x): """Returns the output tensor. @@ -165,18 +174,17 @@ class MossFormer_MaskNet(nn.Module): # [B, N, L] x = self.conv1d_encoder(x) if self.use_global_pos_enc: - #x = self.pos_enc(x.transpose(1, -1)).transpose(1, -1) + x * ( + # x = self.pos_enc(x.transpose(1, -1)).transpose(1, -1) + x * ( # x.size(1) ** 0.5) base = x x = x.transpose(1, -1) emb = self.pos_enc(x) - emb = emb.transpose(0, -1) - #print('base: {}, emb: {}'.format(base.shape, emb.shape)) + emb = emb.transpose(0, -1) + # print('base: {}, emb: {}'.format(base.shape, emb.shape)) x = base + emb - # [B, N, S] - #for i in range(self.num_modules): + # for i in range(self.num_modules): # x = self.dual_mdl[i](x) x = self.mdl(x) x = self.prelu(x) @@ -264,6 +272,7 @@ class MossFormerEncoder(nn.Module): return x + class MossFormerM(nn.Module): """This class implements the transformer encoder. @@ -293,27 +302,28 @@ class MossFormerM(nn.Module): >>> output.shape torch.Size([8, 60, 512]) """ + def __init__( self, num_blocks, d_model=None, causal=False, - group_size = 256, - query_key_dim = 128, - expansion_factor = 4., - attn_dropout = 0.1 + group_size=256, + query_key_dim=128, + expansion_factor=4.0, + attn_dropout=0.1, ): super().__init__() self.mossformerM = MossformerBlock( - dim=d_model, - depth=num_blocks, - group_size=group_size, - query_key_dim=query_key_dim, - expansion_factor=expansion_factor, - causal=causal, - attn_dropout=attn_dropout - ) + dim=d_model, + depth=num_blocks, + group_size=group_size, + query_key_dim=query_key_dim, + expansion_factor=expansion_factor, + causal=causal, + attn_dropout=attn_dropout, + ) self.norm = nn.LayerNorm(d_model, eps=1e-6) def forward( @@ -371,7 +381,7 @@ class Computation_Block(nn.Module): super(Computation_Block, self).__init__() ##MossFormer2M: MossFormer with recurrence - #self.intra_mdl = MossFormer2M(num_blocks=num_blocks, d_model=out_channels) + # self.intra_mdl = MossFormer2M(num_blocks=num_blocks, d_model=out_channels) ##MossFormerM: the orignal MossFormer self.intra_mdl = MossFormerM(num_blocks=num_blocks, d_model=out_channels) self.skip_around_intra = skip_around_intra @@ -396,12 +406,12 @@ class Computation_Block(nn.Module): Output tensor of dimension [B, N, S]. where, B = Batchsize, N = number of filters - S = sequence time index + S = sequence time index """ B, N, S = x.shape # intra RNN # [B, S, N] - intra = x.permute(0, 2, 1).contiguous() #.view(B, S, N) + intra = x.permute(0, 2, 1).contiguous() # .view(B, S, N) intra = self.intra_mdl(intra) @@ -416,4 +426,3 @@ class Computation_Block(nn.Module): out = intra return out - diff --git a/funasr/models/normalize/global_mvn.py b/funasr/models/normalize/global_mvn.py index 8df7c3372..55a728706 100644 --- a/funasr/models/normalize/global_mvn.py +++ b/funasr/models/normalize/global_mvn.py @@ -8,6 +8,7 @@ import torch from funasr.models.transformer.utils.nets_utils import make_pad_mask from funasr.register import tables + @tables.register("normalize_classes", "GlobalMVN") class GlobalMVN(torch.nn.Module): """Apply global mean and variance normalization @@ -112,4 +113,4 @@ class GlobalMVN(torch.nn.Module): if norm_means: x += self.mean x.masked_fill_(make_pad_mask(ilens, x, 1), 0.0) - return x, ilens \ No newline at end of file + return x, ilens diff --git a/funasr/models/normalize/utterance_mvn.py b/funasr/models/normalize/utterance_mvn.py index 9558402ee..8216328c8 100644 --- a/funasr/models/normalize/utterance_mvn.py +++ b/funasr/models/normalize/utterance_mvn.py @@ -5,6 +5,7 @@ import torch from funasr.models.transformer.utils.nets_utils import make_pad_mask from funasr.register import tables + @tables.register("normalize_classes", "UtteranceMVN") class UtteranceMVN(torch.nn.Module): def __init__( diff --git a/funasr/models/paraformer/cif_predictor.py b/funasr/models/paraformer/cif_predictor.py index d538e2164..8b1a9bbdb 100644 --- a/funasr/models/paraformer/cif_predictor.py +++ b/funasr/models/paraformer/cif_predictor.py @@ -12,9 +12,20 @@ from funasr.train_utils.device_funcs import to_device from funasr.models.transformer.utils.nets_utils import make_pad_mask from torch.cuda.amp import autocast + @tables.register("predictor_classes", "CifPredictor") class CifPredictor(torch.nn.Module): - def __init__(self, idim, l_order, r_order, threshold=1.0, dropout=0.1, smooth_factor=1.0, noise_threshold=0, tail_threshold=0.45): + def __init__( + self, + idim, + l_order, + r_order, + threshold=1.0, + dropout=0.1, + smooth_factor=1.0, + noise_threshold=0, + tail_threshold=0.45, + ): super().__init__() self.pad = torch.nn.ConstantPad1d((l_order, r_order), 0) @@ -26,9 +37,16 @@ class CifPredictor(torch.nn.Module): self.noise_threshold = noise_threshold self.tail_threshold = tail_threshold - def forward(self, hidden, target_label=None, mask=None, ignore_id=-1, mask_chunk_predictor=None, - target_label_length=None): - + def forward( + self, + hidden, + target_label=None, + mask=None, + ignore_id=-1, + mask_chunk_predictor=None, + target_label_length=None, + ): + with autocast(False): h = hidden context = h.transpose(1, 2) @@ -58,14 +76,16 @@ class CifPredictor(torch.nn.Module): if target_length is not None: alphas *= (target_length / token_num)[:, None].repeat(1, alphas.size(1)) elif self.tail_threshold > 0.0: - hidden, alphas, token_num = self.tail_process_fn(hidden, alphas, token_num, mask=mask) - + hidden, alphas, token_num = self.tail_process_fn( + hidden, alphas, token_num, mask=mask + ) + acoustic_embeds, cif_peak = cif(hidden, alphas, self.threshold) - + if target_length is None and self.tail_threshold > 0.0: token_num_int = torch.max(token_num).type(torch.int32).item() acoustic_embeds = acoustic_embeds[:, :token_num_int, :] - + return acoustic_embeds, token_num, alphas, cif_peak def tail_process_fn(self, hidden, alphas, token_num=None, mask=None): @@ -91,10 +111,9 @@ class CifPredictor(torch.nn.Module): return hidden, alphas, token_num_floor - - def gen_frame_alignments(self, - alphas: torch.Tensor = None, - encoder_sequence_length: torch.Tensor = None): + def gen_frame_alignments( + self, alphas: torch.Tensor = None, encoder_sequence_length: torch.Tensor = None + ): batch_size, maximum_length = alphas.size() int_type = torch.int32 @@ -117,11 +136,15 @@ class CifPredictor(torch.nn.Module): index_div = torch.floor(torch.true_divide(alphas_cumsum, index)).type(int_type) index_div_bool_zeros = index_div.eq(0) index_div_bool_zeros_count = torch.sum(index_div_bool_zeros, dim=-1) + 1 - index_div_bool_zeros_count = torch.clamp(index_div_bool_zeros_count, 0, encoder_sequence_length.max()) + index_div_bool_zeros_count = torch.clamp( + index_div_bool_zeros_count, 0, encoder_sequence_length.max() + ) token_num_mask = (~make_pad_mask(token_num, maxlen=max_token_num)).to(token_num.device) index_div_bool_zeros_count *= token_num_mask - index_div_bool_zeros_count_tile = index_div_bool_zeros_count[:, :, None].repeat(1, 1, maximum_length) + index_div_bool_zeros_count_tile = index_div_bool_zeros_count[:, :, None].repeat( + 1, 1, maximum_length + ) ones = torch.ones_like(index_div_bool_zeros_count_tile) zeros = torch.zeros_like(index_div_bool_zeros_count_tile) ones = torch.cumsum(ones, dim=2) @@ -132,29 +155,36 @@ class CifPredictor(torch.nn.Module): index_div_bool_zeros_count_tile = 1 - index_div_bool_zeros_count_tile_bool.type(int_type) index_div_bool_zeros_count_tile_out = torch.sum(index_div_bool_zeros_count_tile, dim=1) index_div_bool_zeros_count_tile_out = index_div_bool_zeros_count_tile_out.type(int_type) - predictor_mask = (~make_pad_mask(encoder_sequence_length, maxlen=encoder_sequence_length.max())).type( - int_type).to(encoder_sequence_length.device) + predictor_mask = ( + (~make_pad_mask(encoder_sequence_length, maxlen=encoder_sequence_length.max())) + .type(int_type) + .to(encoder_sequence_length.device) + ) index_div_bool_zeros_count_tile_out = index_div_bool_zeros_count_tile_out * predictor_mask predictor_alignments = index_div_bool_zeros_count_tile_out - predictor_alignments_length = predictor_alignments.sum(-1).type(encoder_sequence_length.dtype) + predictor_alignments_length = predictor_alignments.sum(-1).type( + encoder_sequence_length.dtype + ) return predictor_alignments.detach(), predictor_alignments_length.detach() + @tables.register("predictor_classes", "CifPredictorV2") class CifPredictorV2(torch.nn.Module): - def __init__(self, - idim, - l_order, - r_order, - threshold=1.0, - dropout=0.1, - smooth_factor=1.0, - noise_threshold=0, - tail_threshold=0.0, - tf2torch_tensor_name_prefix_torch="predictor", - tf2torch_tensor_name_prefix_tf="seq2seq/cif", - tail_mask=True, - ): + def __init__( + self, + idim, + l_order, + r_order, + threshold=1.0, + dropout=0.1, + smooth_factor=1.0, + noise_threshold=0, + tail_threshold=0.0, + tf2torch_tensor_name_prefix_torch="predictor", + tf2torch_tensor_name_prefix_tf="seq2seq/cif", + tail_mask=True, + ): super().__init__() self.pad = torch.nn.ConstantPad1d((l_order, r_order), 0) @@ -169,16 +199,23 @@ class CifPredictorV2(torch.nn.Module): self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf self.tail_mask = tail_mask - def forward(self, hidden, target_label=None, mask=None, ignore_id=-1, mask_chunk_predictor=None, - target_label_length=None): - + def forward( + self, + hidden, + target_label=None, + mask=None, + ignore_id=-1, + mask_chunk_predictor=None, + target_label_length=None, + ): + with autocast(False): h = hidden context = h.transpose(1, 2) queries = self.pad(context) output = torch.relu(self.cif_conv1d(queries)) output = output.transpose(1, 2) - + output = self.cif_output(output) alphas = torch.sigmoid(output) alphas = torch.nn.functional.relu(alphas * self.smooth_factor - self.noise_threshold) @@ -200,10 +237,14 @@ class CifPredictorV2(torch.nn.Module): alphas *= (target_length / token_num)[:, None].repeat(1, alphas.size(1)) elif self.tail_threshold > 0.0: if self.tail_mask: - hidden, alphas, token_num = self.tail_process_fn(hidden, alphas, token_num, mask=mask) + hidden, alphas, token_num = self.tail_process_fn( + hidden, alphas, token_num, mask=mask + ) else: - hidden, alphas, token_num = self.tail_process_fn(hidden, alphas, token_num, mask=None) - + hidden, alphas, token_num = self.tail_process_fn( + hidden, alphas, token_num, mask=None + ) + acoustic_embeds, cif_peak = cif(hidden, alphas, self.threshold) if target_length is None and self.tail_threshold > 0.0: token_num_int = torch.max(token_num).type(torch.int32).item() @@ -232,9 +273,9 @@ class CifPredictorV2(torch.nn.Module): cache_hiddens = [] if cache is not None and "chunk_size" in cache: - alphas[:, :cache["chunk_size"][0]] = 0.0 + alphas[:, : cache["chunk_size"][0]] = 0.0 if not is_final: - alphas[:, sum(cache["chunk_size"][:2]):] = 0.0 + alphas[:, sum(cache["chunk_size"][:2]) :] = 0.0 if cache is not None and "cif_alphas" in cache and "cif_hidden" in cache: cache["cif_hidden"] = to_device(cache["cif_hidden"], device=hidden.device) cache["cif_alphas"] = to_device(cache["cif_alphas"], device=alphas.device) @@ -284,10 +325,12 @@ class CifPredictorV2(torch.nn.Module): max_token_len = max(token_length) if max_token_len == 0: - return hidden, torch.stack(token_length, 0), None, None + return hidden, torch.stack(token_length, 0), None, None list_ls = [] for b in range(batch_size): - pad_frames = torch.zeros((max_token_len - token_length[b], hidden_size), device=alphas.device) + pad_frames = torch.zeros( + (max_token_len - token_length[b], hidden_size), device=alphas.device + ) if token_length[b] == 0: list_ls.append(pad_frames) else: @@ -300,7 +343,6 @@ class CifPredictorV2(torch.nn.Module): cache["cif_hidden"] = torch.unsqueeze(cache["cif_hidden"], axis=0) return torch.stack(list_ls, 0), torch.stack(token_length, 0), None, None - def tail_process_fn(self, hidden, alphas, token_num=None, mask=None): b, t, d = hidden.size() tail_threshold = self.tail_threshold @@ -327,9 +369,9 @@ class CifPredictorV2(torch.nn.Module): return hidden, alphas, token_num_floor - def gen_frame_alignments(self, - alphas: torch.Tensor = None, - encoder_sequence_length: torch.Tensor = None): + def gen_frame_alignments( + self, alphas: torch.Tensor = None, encoder_sequence_length: torch.Tensor = None + ): batch_size, maximum_length = alphas.size() int_type = torch.int32 @@ -352,11 +394,15 @@ class CifPredictorV2(torch.nn.Module): index_div = torch.floor(torch.true_divide(alphas_cumsum, index)).type(int_type) index_div_bool_zeros = index_div.eq(0) index_div_bool_zeros_count = torch.sum(index_div_bool_zeros, dim=-1) + 1 - index_div_bool_zeros_count = torch.clamp(index_div_bool_zeros_count, 0, encoder_sequence_length.max()) + index_div_bool_zeros_count = torch.clamp( + index_div_bool_zeros_count, 0, encoder_sequence_length.max() + ) token_num_mask = (~make_pad_mask(token_num, maxlen=max_token_num)).to(token_num.device) index_div_bool_zeros_count *= token_num_mask - index_div_bool_zeros_count_tile = index_div_bool_zeros_count[:, :, None].repeat(1, 1, maximum_length) + index_div_bool_zeros_count_tile = index_div_bool_zeros_count[:, :, None].repeat( + 1, 1, maximum_length + ) ones = torch.ones_like(index_div_bool_zeros_count_tile) zeros = torch.zeros_like(index_div_bool_zeros_count_tile) ones = torch.cumsum(ones, dim=2) @@ -367,19 +413,25 @@ class CifPredictorV2(torch.nn.Module): index_div_bool_zeros_count_tile = 1 - index_div_bool_zeros_count_tile_bool.type(int_type) index_div_bool_zeros_count_tile_out = torch.sum(index_div_bool_zeros_count_tile, dim=1) index_div_bool_zeros_count_tile_out = index_div_bool_zeros_count_tile_out.type(int_type) - predictor_mask = (~make_pad_mask(encoder_sequence_length, maxlen=encoder_sequence_length.max())).type( - int_type).to(encoder_sequence_length.device) + predictor_mask = ( + (~make_pad_mask(encoder_sequence_length, maxlen=encoder_sequence_length.max())) + .type(int_type) + .to(encoder_sequence_length.device) + ) index_div_bool_zeros_count_tile_out = index_div_bool_zeros_count_tile_out * predictor_mask predictor_alignments = index_div_bool_zeros_count_tile_out - predictor_alignments_length = predictor_alignments.sum(-1).type(encoder_sequence_length.dtype) + predictor_alignments_length = predictor_alignments.sum(-1).type( + encoder_sequence_length.dtype + ) return predictor_alignments.detach(), predictor_alignments_length.detach() + @tables.register("predictor_classes", "CifPredictorV2Export") class CifPredictorV2Export(torch.nn.Module): def __init__(self, model, **kwargs): super().__init__() - + self.pad = model.pad self.cif_conv1d = model.cif_conv1d self.cif_output = model.cif_output @@ -387,27 +439,31 @@ class CifPredictorV2Export(torch.nn.Module): self.smooth_factor = model.smooth_factor self.noise_threshold = model.noise_threshold self.tail_threshold = model.tail_threshold - - def forward(self, hidden: torch.Tensor, - mask: torch.Tensor, - ): + + def forward( + self, + hidden: torch.Tensor, + mask: torch.Tensor, + ): alphas, token_num = self.forward_cnn(hidden, mask) mask = mask.transpose(-1, -2).float() mask = mask.squeeze(-1) hidden, alphas, token_num = self.tail_process_fn(hidden, alphas, mask=mask) acoustic_embeds, cif_peak = cif_export(hidden, alphas, self.threshold) - + return acoustic_embeds, token_num, alphas, cif_peak - - def forward_cnn(self, hidden: torch.Tensor, - mask: torch.Tensor, - ): + + def forward_cnn( + self, + hidden: torch.Tensor, + mask: torch.Tensor, + ): h = hidden context = h.transpose(1, 2) queries = self.pad(context) output = torch.relu(self.cif_conv1d(queries)) output = output.transpose(1, 2) - + output = self.cif_output(output) alphas = torch.sigmoid(output) alphas = torch.nn.functional.relu(alphas * self.smooth_factor - self.noise_threshold) @@ -415,67 +471,70 @@ class CifPredictorV2Export(torch.nn.Module): alphas = alphas * mask alphas = alphas.squeeze(-1) token_num = alphas.sum(-1) - + return alphas, token_num - + def tail_process_fn(self, hidden, alphas, token_num=None, mask=None): b, t, d = hidden.size() tail_threshold = self.tail_threshold - + zeros_t = torch.zeros((b, 1), dtype=torch.float32, device=alphas.device) ones_t = torch.ones_like(zeros_t) - + mask_1 = torch.cat([mask, zeros_t], dim=1) mask_2 = torch.cat([ones_t, mask], dim=1) mask = mask_2 - mask_1 tail_threshold = mask * tail_threshold alphas = torch.cat([alphas, zeros_t], dim=1) alphas = torch.add(alphas, tail_threshold) - + zeros = torch.zeros((b, 1, d), dtype=hidden.dtype).to(hidden.device) hidden = torch.cat([hidden, zeros], dim=1) token_num = alphas.sum(dim=-1) token_num_floor = torch.floor(token_num) - + return hidden, alphas, token_num_floor + @torch.jit.script def cif_export(hidden, alphas, threshold: float): batch_size, len_time, hidden_size = hidden.size() threshold = torch.tensor([threshold], dtype=alphas.dtype).to(alphas.device) - + # loop varss integrate = torch.zeros([batch_size], dtype=alphas.dtype, device=hidden.device) frame = torch.zeros([batch_size, hidden_size], dtype=hidden.dtype, device=hidden.device) # intermediate vars along time list_fires = [] list_frames = [] - + for t in range(len_time): alpha = alphas[:, t] - distribution_completion = torch.ones([batch_size], dtype=alphas.dtype, device=hidden.device) - integrate - + distribution_completion = ( + torch.ones([batch_size], dtype=alphas.dtype, device=hidden.device) - integrate + ) + integrate += alpha list_fires.append(integrate) - + fire_place = integrate >= threshold - integrate = torch.where(fire_place, - integrate - torch.ones([batch_size], dtype=alphas.dtype, device=hidden.device), - integrate) - cur = torch.where(fire_place, - distribution_completion, - alpha) + integrate = torch.where( + fire_place, + integrate - torch.ones([batch_size], dtype=alphas.dtype, device=hidden.device), + integrate, + ) + cur = torch.where(fire_place, distribution_completion, alpha) remainds = alpha - cur - + frame += cur[:, None] * hidden[:, t, :] list_frames.append(frame) - frame = torch.where(fire_place[:, None].repeat(1, hidden_size), - remainds[:, None] * hidden[:, t, :], - frame) - + frame = torch.where( + fire_place[:, None].repeat(1, hidden_size), remainds[:, None] * hidden[:, t, :], frame + ) + fires = torch.stack(list_fires, 1) frames = torch.stack(list_frames, 1) - + fire_idxs = fires >= threshold frame_fires = torch.zeros_like(hidden) max_label_len = frames[0, fire_idxs[0]].size(0) @@ -483,7 +542,7 @@ def cif_export(hidden, alphas, threshold: float): frame_fire = frames[b, fire_idxs[b]] frame_len = frame_fire.size(0) frame_fires[b, :frame_len, :] = frame_fire - + if frame_len >= max_label_len: max_label_len = frame_len frame_fires = frame_fires[:, :max_label_len, :] @@ -495,7 +554,7 @@ class mae_loss(torch.nn.Module): def __init__(self, normalize_length=False): super(mae_loss, self).__init__() self.normalize_length = normalize_length - self.criterion = torch.nn.L1Loss(reduction='sum') + self.criterion = torch.nn.L1Loss(reduction="sum") def forward(self, token_length, pre_token_length): loss_token_normalizer = token_length.size(0) @@ -524,19 +583,17 @@ def cif(hidden, alphas, threshold): list_fires.append(integrate) fire_place = integrate >= threshold - integrate = torch.where(fire_place, - integrate - torch.ones([batch_size], device=hidden.device), - integrate) - cur = torch.where(fire_place, - distribution_completion, - alpha) + integrate = torch.where( + fire_place, integrate - torch.ones([batch_size], device=hidden.device), integrate + ) + cur = torch.where(fire_place, distribution_completion, alpha) remainds = alpha - cur frame += cur[:, None] * hidden[:, t, :] list_frames.append(frame) - frame = torch.where(fire_place[:, None].repeat(1, hidden_size), - remainds[:, None] * hidden[:, t, :], - frame) + frame = torch.where( + fire_place[:, None].repeat(1, hidden_size), remainds[:, None] * hidden[:, t, :], frame + ) fires = torch.stack(list_fires, 1) frames = torch.stack(list_frames, 1) @@ -566,10 +623,11 @@ def cif_wo_hidden(alphas, threshold): list_fires.append(integrate) fire_place = integrate >= threshold - integrate = torch.where(fire_place, - integrate - torch.ones([batch_size], device=alphas.device)*threshold, - integrate) + integrate = torch.where( + fire_place, + integrate - torch.ones([batch_size], device=alphas.device) * threshold, + integrate, + ) fires = torch.stack(list_fires, 1) return fires - diff --git a/funasr/models/paraformer/decoder.py b/funasr/models/paraformer/decoder.py index b75d21dbd..7edd91a2a 100644 --- a/funasr/models/paraformer/decoder.py +++ b/funasr/models/paraformer/decoder.py @@ -17,7 +17,10 @@ from funasr.models.transformer.utils.nets_utils import make_pad_mask from funasr.models.transformer.decoder import BaseTransformerDecoder from funasr.models.transformer.positionwise_feed_forward import PositionwiseFeedForward from funasr.models.sanm.positionwise_feed_forward import PositionwiseFeedForwardDecoderSANM -from funasr.models.sanm.attention import MultiHeadedAttentionSANMDecoder, MultiHeadedAttentionCrossAtt +from funasr.models.sanm.attention import ( + MultiHeadedAttentionSANMDecoder, + MultiHeadedAttentionCrossAtt, +) class DecoderLayerSANM(torch.nn.Module): @@ -69,7 +72,7 @@ class DecoderLayerSANM(torch.nn.Module): if self.concat_after: self.concat_linear1 = torch.nn.Linear(size + size, size) self.concat_linear2 = torch.nn.Linear(size + size, size) - self.reserve_attn=False + self.reserve_attn = False self.attn_mat = [] def forward(self, tgt, tgt_mask, memory, memory_mask=None, cache=None): @@ -116,7 +119,7 @@ class DecoderLayerSANM(torch.nn.Module): # x = residual + self.dropout(self.src_attn(x, memory, memory_mask)) return x, tgt_mask, memory, memory_mask, cache - + def get_attn_mat(self, tgt, tgt_mask, memory, memory_mask=None, cache=None): residual = tgt tgt = self.norm1(tgt) @@ -173,10 +176,11 @@ class DecoderLayerSANM(torch.nn.Module): x = residual + self.dropout(self.src_attn(x, memory, memory_mask)) - return x, tgt_mask, memory, memory_mask, cache - def forward_chunk(self, tgt, memory, fsmn_cache=None, opt_cache=None, chunk_size=None, look_back=0): + def forward_chunk( + self, tgt, memory, fsmn_cache=None, opt_cache=None, chunk_size=None, look_back=0 + ): """Compute decoded features. Args: @@ -224,6 +228,7 @@ class ParaformerSANMDecoder(BaseTransformerDecoder): Paraformer: Fast and Accurate Parallel Transformer for Non-autoregressive End-to-End Speech Recognition https://arxiv.org/abs/2006.01713 """ + def __init__( self, vocab_size: int, @@ -303,7 +308,13 @@ class ParaformerSANMDecoder(BaseTransformerDecoder): attention_dim, self_attention_dropout_rate, kernel_size, sanm_shfit=sanm_shfit ), MultiHeadedAttentionCrossAtt( - attention_heads, attention_dim, src_attention_dropout_rate, lora_list, lora_rank, lora_alpha, lora_dropout + attention_heads, + attention_dim, + src_attention_dropout_rate, + lora_list, + lora_rank, + lora_alpha, + lora_dropout, ), PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate), dropout_rate, @@ -374,7 +385,7 @@ class ParaformerSANMDecoder(BaseTransformerDecoder): """ tgt = ys_in_pad tgt_mask = myutils.sequence_mask(ys_in_lens, device=tgt.device)[:, :, None] - + memory = hs_pad memory_mask = myutils.sequence_mask(hlens, device=memory.device)[:, None, :] if chunk_mask is not None: @@ -383,16 +394,10 @@ class ParaformerSANMDecoder(BaseTransformerDecoder): memory_mask = torch.cat((memory_mask, memory_mask[:, -2:-1, :]), dim=1) x = tgt - x, tgt_mask, memory, memory_mask, _ = self.decoders( - x, tgt_mask, memory, memory_mask - ) + x, tgt_mask, memory, memory_mask, _ = self.decoders(x, tgt_mask, memory, memory_mask) if self.decoders2 is not None: - x, tgt_mask, memory, memory_mask, _ = self.decoders2( - x, tgt_mask, memory, memory_mask - ) - x, tgt_mask, memory, memory_mask, _ = self.decoders3( - x, tgt_mask, memory, memory_mask - ) + x, tgt_mask, memory, memory_mask, _ = self.decoders2(x, tgt_mask, memory, memory_mask) + x, tgt_mask, memory, memory_mask, _ = self.decoders3(x, tgt_mask, memory, memory_mask) if self.normalize_before: hidden = self.after_norm(x) @@ -407,12 +412,12 @@ class ParaformerSANMDecoder(BaseTransformerDecoder): def score(self, ys, state, x): """Score.""" - ys_mask = myutils.sequence_mask(torch.tensor([len(ys)], dtype=torch.int32), device=x.device)[:, :, None] - logp, state = self.forward_one_step( - ys.unsqueeze(0), ys_mask, x.unsqueeze(0), cache=state - ) + ys_mask = myutils.sequence_mask( + torch.tensor([len(ys)], dtype=torch.int32), device=x.device + )[:, :, None] + logp, state = self.forward_one_step(ys.unsqueeze(0), ys_mask, x.unsqueeze(0), cache=state) return logp.squeeze(0), state - + def forward_asf2( self, hs_pad: torch.Tensor, @@ -430,7 +435,7 @@ class ParaformerSANMDecoder(BaseTransformerDecoder): tgt, tgt_mask, memory, memory_mask, _ = self.decoders[0](tgt, tgt_mask, memory, memory_mask) attn_mat = self.model.decoders[1].get_attn_mat(tgt, tgt_mask, memory, memory_mask) return attn_mat - + def forward_asf6( self, hs_pad: torch.Tensor, @@ -494,22 +499,24 @@ class ParaformerSANMDecoder(BaseTransformerDecoder): for i in range(self.att_layer_num): decoder = self.decoders[i] x, memory, fsmn_cache[i], opt_cache[i] = decoder.forward_chunk( - x, memory, fsmn_cache=fsmn_cache[i], opt_cache=opt_cache[i], - chunk_size=cache["chunk_size"], look_back=cache["decoder_chunk_look_back"] + x, + memory, + fsmn_cache=fsmn_cache[i], + opt_cache=opt_cache[i], + chunk_size=cache["chunk_size"], + look_back=cache["decoder_chunk_look_back"], ) if self.num_blocks - self.att_layer_num > 1: for i in range(self.num_blocks - self.att_layer_num): j = i + self.att_layer_num decoder = self.decoders2[i] - x, memory, fsmn_cache[j], _ = decoder.forward_chunk( + x, memory, fsmn_cache[j], _ = decoder.forward_chunk( x, memory, fsmn_cache=fsmn_cache[j] ) for decoder in self.decoders3: - x, memory, _, _ = decoder.forward_chunk( - x, memory - ) + x, memory, _, _ = decoder.forward_chunk(x, memory) if self.normalize_before: x = self.after_norm(x) if self.output_layer is not None: @@ -581,22 +588,19 @@ class ParaformerSANMDecoder(BaseTransformerDecoder): return y, new_cache + class DecoderLayerSANMExport(torch.nn.Module): - def __init__( - self, - model - ): + def __init__(self, model): super().__init__() self.self_attn = model.self_attn self.src_attn = model.src_attn self.feed_forward = model.feed_forward self.norm1 = model.norm1 - self.norm2 = model.norm2 if hasattr(model, 'norm2') else None - self.norm3 = model.norm3 if hasattr(model, 'norm3') else None + self.norm2 = model.norm2 if hasattr(model, "norm2") else None + self.norm3 = model.norm3 if hasattr(model, "norm3") else None self.size = model.size - def forward(self, tgt, tgt_mask, memory, memory_mask=None, cache=None): residual = tgt @@ -614,9 +618,8 @@ class DecoderLayerSANMExport(torch.nn.Module): x = self.norm3(x) x = residual + self.src_attn(x, memory, memory_mask) - return x, tgt_mask, memory, memory_mask, cache - + def get_attn_mat(self, tgt, tgt_mask, memory, memory_mask=None, cache=None): residual = tgt tgt = self.norm1(tgt) @@ -636,45 +639,39 @@ class DecoderLayerSANMExport(torch.nn.Module): @tables.register("decoder_classes", "ParaformerSANMDecoderExport") class ParaformerSANMDecoderExport(torch.nn.Module): - def __init__(self, model, - max_seq_len=512, - model_name='decoder', - onnx: bool = True, - **kwargs - ): + def __init__(self, model, max_seq_len=512, model_name="decoder", onnx: bool = True, **kwargs): super().__init__() # self.embed = model.embed #Embedding(model.embed, max_seq_len) from funasr.utils.torch_function import sequence_mask - + self.model = model self.make_pad_mask = sequence_mask(max_seq_len, flip=False) - + from funasr.models.sanm.attention import MultiHeadedAttentionSANMDecoderExport from funasr.models.sanm.attention import MultiHeadedAttentionCrossAttExport - - + for i, d in enumerate(self.model.decoders): if isinstance(d.self_attn, MultiHeadedAttentionSANMDecoder): d.self_attn = MultiHeadedAttentionSANMDecoderExport(d.self_attn) if isinstance(d.src_attn, MultiHeadedAttentionCrossAtt): d.src_attn = MultiHeadedAttentionCrossAttExport(d.src_attn) self.model.decoders[i] = DecoderLayerSANMExport(d) - + if self.model.decoders2 is not None: for i, d in enumerate(self.model.decoders2): if isinstance(d.self_attn, MultiHeadedAttentionSANMDecoder): d.self_attn = MultiHeadedAttentionSANMDecoderExport(d.self_attn) self.model.decoders2[i] = DecoderLayerSANMExport(d) - + for i, d in enumerate(self.model.decoders3): self.model.decoders3[i] = DecoderLayerSANMExport(d) - + self.output_layer = model.output_layer self.after_norm = model.after_norm self.model_name = model_name - + def prepare_mask(self, mask): mask_3d_btd = mask[:, :, None] if len(mask.shape) == 2: @@ -682,9 +679,9 @@ class ParaformerSANMDecoderExport(torch.nn.Module): elif len(mask.shape) == 3: mask_4d_bhlt = 1 - mask[:, None, :] mask_4d_bhlt = mask_4d_bhlt * -10000.0 - + return mask_3d_btd, mask_4d_bhlt - + def forward( self, hs_pad: torch.Tensor, @@ -694,31 +691,27 @@ class ParaformerSANMDecoderExport(torch.nn.Module): return_hidden: bool = False, return_both: bool = False, ): - + tgt = ys_in_pad tgt_mask = self.make_pad_mask(ys_in_lens) tgt_mask, _ = self.prepare_mask(tgt_mask) # tgt_mask = myutils.sequence_mask(ys_in_lens, device=tgt.device)[:, :, None] - + memory = hs_pad memory_mask = self.make_pad_mask(hlens) _, memory_mask = self.prepare_mask(memory_mask) # memory_mask = myutils.sequence_mask(hlens, device=memory.device)[:, None, :] - + x = tgt - x, tgt_mask, memory, memory_mask, _ = self.model.decoders( - x, tgt_mask, memory, memory_mask - ) + x, tgt_mask, memory, memory_mask, _ = self.model.decoders(x, tgt_mask, memory, memory_mask) if self.model.decoders2 is not None: x, tgt_mask, memory, memory_mask, _ = self.model.decoders2( x, tgt_mask, memory, memory_mask ) - x, tgt_mask, memory, memory_mask, _ = self.model.decoders3( - x, tgt_mask, memory, memory_mask - ) + x, tgt_mask, memory, memory_mask, _ = self.model.decoders3(x, tgt_mask, memory, memory_mask) hidden = self.after_norm(x) # x = self.output_layer(x) - + if self.output_layer is not None and return_hidden is False: x = self.output_layer(hidden) return x, ys_in_lens @@ -726,7 +719,7 @@ class ParaformerSANMDecoderExport(torch.nn.Module): x = self.output_layer(hidden) return x, hidden, ys_in_lens return hidden, ys_in_lens - + def forward_asf2( self, hs_pad: torch.Tensor, @@ -742,10 +735,12 @@ class ParaformerSANMDecoderExport(torch.nn.Module): memory_mask = myutils.sequence_mask(hlens, device=memory.device)[:, None, :] _, memory_mask = self.prepare_mask(memory_mask) - tgt, tgt_mask, memory, memory_mask, _ = self.model.decoders[0](tgt, tgt_mask, memory, memory_mask) + tgt, tgt_mask, memory, memory_mask, _ = self.model.decoders[0]( + tgt, tgt_mask, memory, memory_mask + ) attn_mat = self.model.decoders[1].get_attn_mat(tgt, tgt_mask, memory, memory_mask) return attn_mat - + def forward_asf6( self, hs_pad: torch.Tensor, @@ -761,15 +756,25 @@ class ParaformerSANMDecoderExport(torch.nn.Module): memory_mask = myutils.sequence_mask(hlens, device=memory.device)[:, None, :] _, memory_mask = self.prepare_mask(memory_mask) - tgt, tgt_mask, memory, memory_mask, _ = self.model.decoders[0](tgt, tgt_mask, memory, memory_mask) - tgt, tgt_mask, memory, memory_mask, _ = self.model.decoders[1](tgt, tgt_mask, memory, memory_mask) - tgt, tgt_mask, memory, memory_mask, _ = self.model.decoders[2](tgt, tgt_mask, memory, memory_mask) - tgt, tgt_mask, memory, memory_mask, _ = self.model.decoders[3](tgt, tgt_mask, memory, memory_mask) - tgt, tgt_mask, memory, memory_mask, _ = self.model.decoders[4](tgt, tgt_mask, memory, memory_mask) + tgt, tgt_mask, memory, memory_mask, _ = self.model.decoders[0]( + tgt, tgt_mask, memory, memory_mask + ) + tgt, tgt_mask, memory, memory_mask, _ = self.model.decoders[1]( + tgt, tgt_mask, memory, memory_mask + ) + tgt, tgt_mask, memory, memory_mask, _ = self.model.decoders[2]( + tgt, tgt_mask, memory, memory_mask + ) + tgt, tgt_mask, memory, memory_mask, _ = self.model.decoders[3]( + tgt, tgt_mask, memory, memory_mask + ) + tgt, tgt_mask, memory, memory_mask, _ = self.model.decoders[4]( + tgt, tgt_mask, memory, memory_mask + ) attn_mat = self.model.decoders[5].get_attn_mat(tgt, tgt_mask, memory, memory_mask) return attn_mat - - ''' + + """ def get_dummy_inputs(self, enc_size): tgt = torch.LongTensor([0]).unsqueeze(0) memory = torch.randn(1, 100, enc_size) @@ -818,14 +823,12 @@ class ParaformerSANMDecoderExport(torch.nn.Module): for d in range(cache_num) }) return ret - ''' - + """ + + @tables.register("decoder_classes", "ParaformerSANMDecoderOnlineExport") class ParaformerSANMDecoderOnlineExport(torch.nn.Module): - def __init__(self, model, - max_seq_len=512, - model_name='decoder', - onnx: bool = True, **kwargs): + def __init__(self, model, max_seq_len=512, model_name="decoder", onnx: bool = True, **kwargs): super().__init__() # self.embed = model.embed #Embedding(model.embed, max_seq_len) self.model = model @@ -833,7 +836,7 @@ class ParaformerSANMDecoderOnlineExport(torch.nn.Module): from funasr.utils.torch_function import sequence_mask self.model = model - + self.make_pad_mask = sequence_mask(max_seq_len, flip=False) from funasr.models.sanm.attention import MultiHeadedAttentionSANMDecoderExport @@ -854,11 +857,11 @@ class ParaformerSANMDecoderOnlineExport(torch.nn.Module): for i, d in enumerate(self.model.decoders3): self.model.decoders3[i] = DecoderLayerSANMExport(d) - + self.output_layer = model.output_layer self.after_norm = model.after_norm self.model_name = model_name - + def prepare_mask(self, mask): mask_3d_btd = mask[:, :, None] if len(mask.shape) == 2: @@ -866,9 +869,9 @@ class ParaformerSANMDecoderOnlineExport(torch.nn.Module): elif len(mask.shape) == 3: mask_4d_bhlt = 1 - mask[:, None, :] mask_4d_bhlt = mask_4d_bhlt * -10000.0 - + return mask_3d_btd, mask_4d_bhlt - + def forward( self, hs_pad: torch.Tensor, @@ -877,17 +880,17 @@ class ParaformerSANMDecoderOnlineExport(torch.nn.Module): ys_in_lens: torch.Tensor, *args, ): - + tgt = ys_in_pad tgt_mask = self.make_pad_mask(ys_in_lens) tgt_mask, _ = self.prepare_mask(tgt_mask) # tgt_mask = myutils.sequence_mask(ys_in_lens, device=tgt.device)[:, :, None] - + memory = hs_pad memory_mask = self.make_pad_mask(hlens) _, memory_mask = self.prepare_mask(memory_mask) # memory_mask = myutils.sequence_mask(hlens, device=memory.device)[:, None, :] - + x = tgt out_caches = list() for i, decoder in enumerate(self.model.decoders): @@ -903,76 +906,75 @@ class ParaformerSANMDecoderOnlineExport(torch.nn.Module): x, tgt_mask, memory, memory_mask, cache=in_cache ) out_caches.append(out_cache) - x, tgt_mask, memory, memory_mask, _ = self.model.decoders3( - x, tgt_mask, memory, memory_mask - ) + x, tgt_mask, memory, memory_mask, _ = self.model.decoders3(x, tgt_mask, memory, memory_mask) x = self.after_norm(x) x = self.output_layer(x) - + return x, out_caches - + def get_dummy_inputs(self, enc_size): enc = torch.randn(2, 100, enc_size).type(torch.float32) enc_len = torch.tensor([30, 100], dtype=torch.int32) acoustic_embeds = torch.randn(2, 10, enc_size).type(torch.float32) acoustic_embeds_len = torch.tensor([5, 10], dtype=torch.int32) cache_num = len(self.model.decoders) - if hasattr(self.model, 'decoders2') and self.model.decoders2 is not None: + if hasattr(self.model, "decoders2") and self.model.decoders2 is not None: cache_num += len(self.model.decoders2) cache = [ - torch.zeros((2, self.model.decoders[0].size, self.model.decoders[0].self_attn.kernel_size - 1), - dtype=torch.float32) + torch.zeros( + (2, self.model.decoders[0].size, self.model.decoders[0].self_attn.kernel_size - 1), + dtype=torch.float32, + ) for _ in range(cache_num) ] return (enc, enc_len, acoustic_embeds, acoustic_embeds_len, *cache) - + def get_input_names(self): cache_num = len(self.model.decoders) - if hasattr(self.model, 'decoders2') and self.model.decoders2 is not None: + if hasattr(self.model, "decoders2") and self.model.decoders2 is not None: cache_num += len(self.model.decoders2) - return ['enc', 'enc_len', 'acoustic_embeds', 'acoustic_embeds_len'] \ - + ['in_cache_%d' % i for i in range(cache_num)] - + return ["enc", "enc_len", "acoustic_embeds", "acoustic_embeds_len"] + [ + "in_cache_%d" % i for i in range(cache_num) + ] + def get_output_names(self): cache_num = len(self.model.decoders) - if hasattr(self.model, 'decoders2') and self.model.decoders2 is not None: + if hasattr(self.model, "decoders2") and self.model.decoders2 is not None: cache_num += len(self.model.decoders2) - return ['logits', 'sample_ids'] \ - + ['out_cache_%d' % i for i in range(cache_num)] - + return ["logits", "sample_ids"] + ["out_cache_%d" % i for i in range(cache_num)] + def get_dynamic_axes(self): ret = { - 'enc': { - 0: 'batch_size', - 1: 'enc_length' + "enc": {0: "batch_size", 1: "enc_length"}, + "acoustic_embeds": {0: "batch_size", 1: "token_length"}, + "enc_len": { + 0: "batch_size", }, - 'acoustic_embeds': { - 0: 'batch_size', - 1: 'token_length' + "acoustic_embeds_len": { + 0: "batch_size", }, - 'enc_len': { - 0: 'batch_size', - }, - 'acoustic_embeds_len': { - 0: 'batch_size', - }, - } cache_num = len(self.model.decoders) - if hasattr(self.model, 'decoders2') and self.model.decoders2 is not None: + if hasattr(self.model, "decoders2") and self.model.decoders2 is not None: cache_num += len(self.model.decoders2) - ret.update({ - 'in_cache_%d' % d: { - 0: 'batch_size', + ret.update( + { + "in_cache_%d" + % d: { + 0: "batch_size", + } + for d in range(cache_num) } - for d in range(cache_num) - }) - ret.update({ - 'out_cache_%d' % d: { - 0: 'batch_size', + ) + ret.update( + { + "out_cache_%d" + % d: { + 0: "batch_size", + } + for d in range(cache_num) } - for d in range(cache_num) - }) + ) return ret @@ -983,23 +985,24 @@ class ParaformerSANDecoder(BaseTransformerDecoder): Paraformer: Fast and Accurate Parallel Transformer for Non-autoregressive End-to-End Speech Recognition https://arxiv.org/abs/2006.01713 """ + def __init__( - self, - vocab_size: int, - encoder_output_size: int, - attention_heads: int = 4, - linear_units: int = 2048, - num_blocks: int = 6, - dropout_rate: float = 0.1, - positional_dropout_rate: float = 0.1, - self_attention_dropout_rate: float = 0.0, - src_attention_dropout_rate: float = 0.0, - input_layer: str = "embed", - use_output_layer: bool = True, - pos_enc_class=PositionalEncoding, - normalize_before: bool = True, - concat_after: bool = False, - embeds_id: int = -1, + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + embeds_id: int = -1, ): super().__init__( vocab_size=vocab_size, @@ -1017,12 +1020,8 @@ class ParaformerSANDecoder(BaseTransformerDecoder): num_blocks, lambda lnum: DecoderLayer( attention_dim, - MultiHeadedAttention( - attention_heads, attention_dim, self_attention_dropout_rate - ), - MultiHeadedAttention( - attention_heads, attention_dim, src_attention_dropout_rate - ), + MultiHeadedAttention(attention_heads, attention_dim, self_attention_dropout_rate), + MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate), PositionwiseFeedForward(attention_dim, linear_units, dropout_rate), dropout_rate, normalize_before, @@ -1033,11 +1032,11 @@ class ParaformerSANDecoder(BaseTransformerDecoder): self.attention_dim = attention_dim def forward( - self, - hs_pad: torch.Tensor, - hlens: torch.Tensor, - ys_in_pad: torch.Tensor, - ys_in_lens: torch.Tensor, + self, + hs_pad: torch.Tensor, + hlens: torch.Tensor, + ys_in_pad: torch.Tensor, + ys_in_lens: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: """Forward decoder. @@ -1060,23 +1059,17 @@ class ParaformerSANDecoder(BaseTransformerDecoder): tgt_mask = (~make_pad_mask(ys_in_lens)[:, None, :]).to(tgt.device) memory = hs_pad - memory_mask = (~make_pad_mask(hlens, maxlen=memory.size(1)))[:, None, :].to( - memory.device - ) + memory_mask = (~make_pad_mask(hlens, maxlen=memory.size(1)))[:, None, :].to(memory.device) # Padding for Longformer if memory_mask.shape[-1] != memory.shape[1]: padlen = memory.shape[1] - memory_mask.shape[-1] - memory_mask = torch.nn.functional.pad( - memory_mask, (0, padlen), "constant", False - ) + memory_mask = torch.nn.functional.pad(memory_mask, (0, padlen), "constant", False) # x = self.embed(tgt) x = tgt embeds_outputs = None for layer_id, decoder in enumerate(self.decoders): - x, tgt_mask, memory, memory_mask = decoder( - x, tgt_mask, memory, memory_mask - ) + x, tgt_mask, memory, memory_mask = decoder(x, tgt_mask, memory, memory_mask) if layer_id == self.embeds_id: embeds_outputs = x if self.normalize_before: @@ -1090,36 +1083,38 @@ class ParaformerSANDecoder(BaseTransformerDecoder): else: return x, olens + @tables.register("decoder_classes", "ParaformerDecoderSANExport") class ParaformerDecoderSANExport(torch.nn.Module): - def __init__(self, model, - max_seq_len=512, - model_name='decoder', - onnx: bool = True, ): + def __init__( + self, + model, + max_seq_len=512, + model_name="decoder", + onnx: bool = True, + ): super().__init__() # self.embed = model.embed #Embedding(model.embed, max_seq_len) self.model = model - from funasr.utils.torch_function import sequence_mask self.model = model self.make_pad_mask = sequence_mask(max_seq_len, flip=False) - from funasr.models.transformer.decoder import DecoderLayerExport from funasr.models.transformer.attention import MultiHeadedAttentionExport - + for i, d in enumerate(self.model.decoders): if isinstance(d.src_attn, MultiHeadedAttention): d.src_attn = MultiHeadedAttentionExport(d.src_attn) self.model.decoders[i] = DecoderLayerExport(d) - + self.output_layer = model.output_layer self.after_norm = model.after_norm self.model_name = model_name - + def prepare_mask(self, mask): mask_3d_btd = mask[:, :, None] if len(mask.shape) == 2: @@ -1127,9 +1122,9 @@ class ParaformerDecoderSANExport(torch.nn.Module): elif len(mask.shape) == 3: mask_4d_bhlt = 1 - mask[:, None, :] mask_4d_bhlt = mask_4d_bhlt * -10000.0 - + return mask_3d_btd, mask_4d_bhlt - + def forward( self, hs_pad: torch.Tensor, @@ -1137,72 +1132,62 @@ class ParaformerDecoderSANExport(torch.nn.Module): ys_in_pad: torch.Tensor, ys_in_lens: torch.Tensor, ): - + tgt = ys_in_pad tgt_mask = self.make_pad_mask(ys_in_lens) tgt_mask, _ = self.prepare_mask(tgt_mask) # tgt_mask = myutils.sequence_mask(ys_in_lens, device=tgt.device)[:, :, None] - + memory = hs_pad memory_mask = self.make_pad_mask(hlens) _, memory_mask = self.prepare_mask(memory_mask) # memory_mask = myutils.sequence_mask(hlens, device=memory.device)[:, None, :] - + x = tgt - x, tgt_mask, memory, memory_mask = self.model.decoders( - x, tgt_mask, memory, memory_mask - ) + x, tgt_mask, memory, memory_mask = self.model.decoders(x, tgt_mask, memory, memory_mask) x = self.after_norm(x) x = self.output_layer(x) - + return x, ys_in_lens - + def get_dummy_inputs(self, enc_size): tgt = torch.LongTensor([0]).unsqueeze(0) memory = torch.randn(1, 100, enc_size) pre_acoustic_embeds = torch.randn(1, 1, enc_size) cache_num = len(self.model.decoders) + len(self.model.decoders2) cache = [ - torch.zeros((1, self.model.decoders[0].size, self.model.decoders[0].self_attn.kernel_size)) + torch.zeros( + (1, self.model.decoders[0].size, self.model.decoders[0].self_attn.kernel_size) + ) for _ in range(cache_num) ] return (tgt, memory, pre_acoustic_embeds, cache) - + def is_optimizable(self): return True - + def get_input_names(self): cache_num = len(self.model.decoders) + len(self.model.decoders2) - return ['tgt', 'memory', 'pre_acoustic_embeds'] \ - + ['cache_%d' % i for i in range(cache_num)] - + return ["tgt", "memory", "pre_acoustic_embeds"] + ["cache_%d" % i for i in range(cache_num)] + def get_output_names(self): cache_num = len(self.model.decoders) + len(self.model.decoders2) - return ['y'] \ - + ['out_cache_%d' % i for i in range(cache_num)] - + return ["y"] + ["out_cache_%d" % i for i in range(cache_num)] + def get_dynamic_axes(self): ret = { - 'tgt': { - 0: 'tgt_batch', - 1: 'tgt_length' + "tgt": {0: "tgt_batch", 1: "tgt_length"}, + "memory": {0: "memory_batch", 1: "memory_length"}, + "pre_acoustic_embeds": { + 0: "acoustic_embeds_batch", + 1: "acoustic_embeds_length", }, - 'memory': { - 0: 'memory_batch', - 1: 'memory_length' - }, - 'pre_acoustic_embeds': { - 0: 'acoustic_embeds_batch', - 1: 'acoustic_embeds_length', - } } cache_num = len(self.model.decoders) + len(self.model.decoders2) - ret.update({ - 'cache_%d' % d: { - 0: 'cache_%d_batch' % d, - 2: 'cache_%d_length' % d + ret.update( + { + "cache_%d" % d: {0: "cache_%d_batch" % d, 2: "cache_%d_length" % d} + for d in range(cache_num) } - for d in range(cache_num) - }) + ) return ret - \ No newline at end of file diff --git a/funasr/models/paraformer/export_meta.py b/funasr/models/paraformer/export_meta.py index 4d491e945..5c1b6c08d 100644 --- a/funasr/models/paraformer/export_meta.py +++ b/funasr/models/paraformer/export_meta.py @@ -9,36 +9,36 @@ from funasr.register import tables def export_rebuild_model(model, **kwargs): - model.device = kwargs.get("device") - is_onnx = kwargs.get("type", "onnx") == "onnx" - encoder_class = tables.encoder_classes.get(kwargs["encoder"]+"Export") - model.encoder = encoder_class(model.encoder, onnx=is_onnx) - - predictor_class = tables.predictor_classes.get(kwargs["predictor"]+"Export") - model.predictor = predictor_class(model.predictor, onnx=is_onnx) + model.device = kwargs.get("device") + is_onnx = kwargs.get("type", "onnx") == "onnx" + encoder_class = tables.encoder_classes.get(kwargs["encoder"] + "Export") + model.encoder = encoder_class(model.encoder, onnx=is_onnx) + predictor_class = tables.predictor_classes.get(kwargs["predictor"] + "Export") + model.predictor = predictor_class(model.predictor, onnx=is_onnx) - decoder_class = tables.decoder_classes.get(kwargs["decoder"]+"Export") - model.decoder = decoder_class(model.decoder, onnx=is_onnx) - - from funasr.utils.torch_function import sequence_mask - model.make_pad_mask = sequence_mask(kwargs['max_seq_len'], flip=False) - - model.forward = types.MethodType(export_forward, model) - model.export_dummy_inputs = types.MethodType(export_dummy_inputs, model) - model.export_input_names = types.MethodType(export_input_names, model) - model.export_output_names = types.MethodType(export_output_names, model) - model.export_dynamic_axes = types.MethodType(export_dynamic_axes, model) - model.export_name = types.MethodType(export_name, model) - - return model + decoder_class = tables.decoder_classes.get(kwargs["decoder"] + "Export") + model.decoder = decoder_class(model.decoder, onnx=is_onnx) + + from funasr.utils.torch_function import sequence_mask + + model.make_pad_mask = sequence_mask(kwargs["max_seq_len"], flip=False) + + model.forward = types.MethodType(export_forward, model) + model.export_dummy_inputs = types.MethodType(export_dummy_inputs, model) + model.export_input_names = types.MethodType(export_input_names, model) + model.export_output_names = types.MethodType(export_output_names, model) + model.export_dynamic_axes = types.MethodType(export_dynamic_axes, model) + model.export_name = types.MethodType(export_name, model) + + return model def export_forward( - self, - speech: torch.Tensor, - speech_lengths: torch.Tensor, - ): + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, +): # a. To device batch = {"speech": speech, "speech_lengths": speech_lengths} # batch = to_device(batch, device=self.device) @@ -54,6 +54,7 @@ def export_forward( return decoder_out, pre_token_length + def export_dummy_inputs(self): speech = torch.randn(2, 30, 560) speech_lengths = torch.tensor([6, 30], dtype=torch.int32) @@ -61,25 +62,24 @@ def export_dummy_inputs(self): def export_input_names(self): - return ['speech', 'speech_lengths'] + return ["speech", "speech_lengths"] + def export_output_names(self): - return ['logits', 'token_num'] + return ["logits", "token_num"] + def export_dynamic_axes(self): return { - 'speech': { - 0: 'batch_size', - 1: 'feats_length' - }, - 'speech_lengths': { - 0: 'batch_size', - }, - 'logits': { - 0: 'batch_size', - 1: 'logits_length' + "speech": {0: "batch_size", 1: "feats_length"}, + "speech_lengths": { + 0: "batch_size", }, + "logits": {0: "batch_size", 1: "logits_length"}, } -def export_name(self, ): - return "model.onnx" \ No newline at end of file + +def export_name( + self, +): + return "model.onnx" diff --git a/funasr/models/paraformer/model.py b/funasr/models/paraformer/model.py index 6c7957c5f..0d9bb2be1 100644 --- a/funasr/models/paraformer/model.py +++ b/funasr/models/paraformer/model.py @@ -31,7 +31,7 @@ class Paraformer(torch.nn.Module): Paraformer: Fast and Accurate Parallel Transformer for Non-autoregressive End-to-End Speech Recognition https://arxiv.org/abs/2206.08317 """ - + def __init__( self, specaug: Optional[str] = None, @@ -91,17 +91,15 @@ class Paraformer(torch.nn.Module): **decoder_conf, ) if ctc_weight > 0.0: - + if ctc_conf is None: ctc_conf = {} - - ctc = CTC( - odim=vocab_size, encoder_output_size=encoder_output_size, **ctc_conf - ) + + ctc = CTC(odim=vocab_size, encoder_output_size=encoder_output_size, **ctc_conf) if predictor is not None: predictor_class = tables.predictor_classes.get(predictor) predictor = predictor_class(**predictor_conf) - + # note that eos is the same as sos (equivalent ID) self.blank_id = blank_id self.sos = sos if sos is not None else vocab_size - 1 @@ -131,7 +129,7 @@ class Paraformer(torch.nn.Module): self.decoder = None else: self.decoder = decoder - + self.criterion_att = LabelSmoothingLoss( size=vocab_size, padding_idx=ignore_id, @@ -156,16 +154,15 @@ class Paraformer(torch.nn.Module): self.sampling_ratio = sampling_ratio self.criterion_pre = mae_loss(normalize_length=length_normalized_loss) - self.share_embedding = share_embedding if self.share_embedding: self.decoder.embed = None - + self.use_1st_decoder_loss = use_1st_decoder_loss self.length_normalized_loss = length_normalized_loss self.beam_search = None self.error_calculator = None - + def forward( self, speech: torch.Tensor, @@ -185,40 +182,41 @@ class Paraformer(torch.nn.Module): text_lengths = text_lengths[:, 0] if len(speech_lengths.size()) > 1: speech_lengths = speech_lengths[:, 0] - + batch_size = speech.shape[0] - + # Encoder encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) - loss_ctc, cer_ctc = None, None loss_pre = None stats = dict() - + # decoder: CTC branch if self.ctc_weight != 0.0: loss_ctc, cer_ctc = self._calc_ctc_loss( encoder_out, encoder_out_lens, text, text_lengths ) - + # Collect CTC branch stats stats["loss_ctc"] = loss_ctc.detach() if loss_ctc is not None else None stats["cer_ctc"] = cer_ctc - # decoder: Attention decoder branch loss_att, acc_att, cer_att, wer_att, loss_pre, pre_loss_att = self._calc_att_loss( encoder_out, encoder_out_lens, text, text_lengths ) - + # 3. CTC-Att loss definition if self.ctc_weight == 0.0: loss = loss_att + loss_pre * self.predictor_weight else: - loss = self.ctc_weight * loss_ctc + (1 - self.ctc_weight) * loss_att + loss_pre * self.predictor_weight - - + loss = ( + self.ctc_weight * loss_ctc + + (1 - self.ctc_weight) * loss_att + + loss_pre * self.predictor_weight + ) + # Collect Attn branch stats stats["loss_att"] = loss_att.detach() if loss_att is not None else None stats["pre_loss_att"] = pre_loss_att.detach() if pre_loss_att is not None else None @@ -226,19 +224,21 @@ class Paraformer(torch.nn.Module): stats["cer"] = cer_att stats["wer"] = wer_att stats["loss_pre"] = loss_pre.detach().cpu() if loss_pre is not None else None - + stats["loss"] = torch.clone(loss.detach()) stats["batch_size"] = batch_size - + # force_gatherable: to-device and to-tensor if scalar for DataParallel if self.length_normalized_loss: batch_size = (text_lengths + self.predictor_bias).sum() loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) return loss, stats, weight - def encode( - self, speech: torch.Tensor, speech_lengths: torch.Tensor, **kwargs, + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + **kwargs, ) -> Tuple[torch.Tensor, torch.Tensor]: """Encoder. Note that this method is used by asr_inference.py Args: @@ -251,11 +251,10 @@ class Paraformer(torch.nn.Module): # Data augmentation if self.specaug is not None and self.training: speech, speech_lengths = self.specaug(speech, speech_lengths) - + # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN if self.normalize is not None: speech, speech_lengths = self.normalize(speech, speech_lengths) - # Forward encoder encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths) @@ -263,21 +262,22 @@ class Paraformer(torch.nn.Module): encoder_out = encoder_out[0] return encoder_out, encoder_out_lens - + def calc_predictor(self, encoder_out, encoder_out_lens): - - encoder_out_mask = (~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :]).to( - encoder_out.device) - pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = self.predictor(encoder_out, None, - encoder_out_mask, - ignore_id=self.ignore_id) - return pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index - - def cal_decoder_with_predictor(self, encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens): - - decoder_outs = self.decoder( - encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens + + encoder_out_mask = ( + ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :] + ).to(encoder_out.device) + pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = self.predictor( + encoder_out, None, encoder_out_mask, ignore_id=self.ignore_id ) + return pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index + + def cal_decoder_with_predictor( + self, encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens + ): + + decoder_outs = self.decoder(encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens) decoder_out = decoder_outs[0] decoder_out = torch.log_softmax(decoder_out, dim=-1) return decoder_out, ys_pad_lens @@ -289,30 +289,31 @@ class Paraformer(torch.nn.Module): ys_pad: torch.Tensor, ys_pad_lens: torch.Tensor, ): - encoder_out_mask = (~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :]).to( - encoder_out.device) + encoder_out_mask = ( + ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :] + ).to(encoder_out.device) if self.predictor_bias == 1: _, ys_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id) ys_pad_lens = ys_pad_lens + self.predictor_bias - pre_acoustic_embeds, pre_token_length, _, pre_peak_index = self.predictor(encoder_out, ys_pad, encoder_out_mask, - ignore_id=self.ignore_id) - + pre_acoustic_embeds, pre_token_length, _, pre_peak_index = self.predictor( + encoder_out, ys_pad, encoder_out_mask, ignore_id=self.ignore_id + ) + # 0. sampler decoder_out_1st = None pre_loss_att = None if self.sampling_ratio > 0.0: - sematic_embeds, decoder_out_1st = self.sampler(encoder_out, encoder_out_lens, ys_pad, ys_pad_lens, - pre_acoustic_embeds) + sematic_embeds, decoder_out_1st = self.sampler( + encoder_out, encoder_out_lens, ys_pad, ys_pad_lens, pre_acoustic_embeds + ) else: sematic_embeds = pre_acoustic_embeds - + # 1. Forward decoder - decoder_outs = self.decoder( - encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens - ) + decoder_outs = self.decoder(encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens) decoder_out, _ = decoder_outs[0], decoder_outs[1] - + if decoder_out_1st is None: decoder_out_1st = decoder_out # 2. Compute attention loss @@ -323,19 +324,21 @@ class Paraformer(torch.nn.Module): ignore_label=self.ignore_id, ) loss_pre = self.criterion_pre(ys_pad_lens.type_as(pre_token_length), pre_token_length) - + # Compute cer/wer using attention-decoder if self.training or self.error_calculator is None: cer_att, wer_att = None, None else: ys_hat = decoder_out_1st.argmax(dim=-1) cer_att, wer_att = self.error_calculator(ys_hat.cpu(), ys_pad.cpu()) - + return loss_att, acc_att, cer_att, wer_att, loss_pre, pre_loss_att - + def sampler(self, encoder_out, encoder_out_lens, ys_pad, ys_pad_lens, pre_acoustic_embeds): - - tgt_mask = (~make_pad_mask(ys_pad_lens, maxlen=ys_pad_lens.max())[:, :, None]).to(ys_pad.device) + + tgt_mask = (~make_pad_mask(ys_pad_lens, maxlen=ys_pad_lens.max())[:, :, None]).to( + ys_pad.device + ) ys_pad_masked = ys_pad * tgt_mask[:, :, 0] if self.share_embedding: ys_pad_embed = self.decoder.output_layer.weight[ys_pad_masked] @@ -353,19 +356,24 @@ class Paraformer(torch.nn.Module): input_mask = torch.ones_like(nonpad_positions) bsz, seq_len = ys_pad.size() for li in range(bsz): - target_num = (((seq_lens[li] - same_num[li].sum()).float()) * self.sampling_ratio).long() + target_num = ( + ((seq_lens[li] - same_num[li].sum()).float()) * self.sampling_ratio + ).long() if target_num > 0: - input_mask[li].scatter_(dim=0, - index=torch.randperm(seq_lens[li])[:target_num].to(input_mask.device), - value=0) + input_mask[li].scatter_( + dim=0, + index=torch.randperm(seq_lens[li])[:target_num].to(input_mask.device), + value=0, + ) input_mask = input_mask.eq(1) input_mask = input_mask.masked_fill(~nonpad_positions, False) input_mask_expand_dim = input_mask.unsqueeze(2).to(pre_acoustic_embeds.device) - - sematic_embeds = pre_acoustic_embeds.masked_fill(~input_mask_expand_dim, 0) + ys_pad_embed.masked_fill( - input_mask_expand_dim, 0) + + sematic_embeds = pre_acoustic_embeds.masked_fill( + ~input_mask_expand_dim, 0 + ) + ys_pad_embed.masked_fill(input_mask_expand_dim, 0) return sematic_embeds * tgt_mask, decoder_out * tgt_mask - + def _calc_ctc_loss( self, encoder_out: torch.Tensor, @@ -375,7 +383,7 @@ class Paraformer(torch.nn.Module): ): # Calc CTC loss loss_ctc = self.ctc(encoder_out, encoder_out_lens, ys_pad, ys_pad_lens) - + # Calc CER using CTC cer_ctc = None if not self.training and self.error_calculator is not None: @@ -383,33 +391,30 @@ class Paraformer(torch.nn.Module): cer_ctc = self.error_calculator(ys_hat.cpu(), ys_pad.cpu(), is_ctc=True) return loss_ctc, cer_ctc - - def init_beam_search(self, - **kwargs, - ): + def init_beam_search( + self, + **kwargs, + ): from funasr.models.paraformer.search import BeamSearchPara from funasr.models.transformer.scorers.ctc import CTCPrefixScorer from funasr.models.transformer.scorers.length_bonus import LengthBonus - + # 1. Build ASR model scorers = {} - + if self.ctc != None: ctc = CTCPrefixScorer(ctc=self.ctc, eos=self.eos) - scorers.update( - ctc=ctc - ) + scorers.update(ctc=ctc) token_list = kwargs.get("token_list") scorers.update( length_bonus=LengthBonus(len(token_list)), ) - # 3. Build ngram model # ngram is not supported now ngram = None scorers["ngram"] = ngram - + weights = dict( decoder=1.0 - kwargs.get("decoding_ctc_weight"), ctc=kwargs.get("decoding_ctc_weight", 0.0), @@ -432,25 +437,30 @@ class Paraformer(torch.nn.Module): # if isinstance(scorer, torch.nn.Module): # scorer.to(device=kwargs.get("device", "cpu"), dtype=getattr(torch, kwargs.get("dtype", "float32"))).eval() self.beam_search = beam_search - - def inference(self, - data_in, - data_lengths=None, - key: list=None, - tokenizer=None, - frontend=None, - **kwargs, - ): + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): # init beamsearch is_use_ctc = kwargs.get("decoding_ctc_weight", 0.0) > 0.00001 and self.ctc != None - is_use_lm = kwargs.get("lm_weight", 0.0) > 0.00001 and kwargs.get("lm_file", None) is not None + is_use_lm = ( + kwargs.get("lm_weight", 0.0) > 0.00001 and kwargs.get("lm_file", None) is not None + ) if self.beam_search is None and (is_use_lm or is_use_ctc): logging.info("enable beam_search") self.init_beam_search(**kwargs) self.nbest = kwargs.get("nbest", 1) - + meta_data = {} - if isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank": # fbank + if ( + isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank" + ): # fbank speech, speech_lengths = data_in, data_lengths if len(speech.shape) < 3: speech = speech[None, :, :] @@ -461,14 +471,24 @@ class Paraformer(torch.nn.Module): else: # extract fbank feats time1 = time.perf_counter() - audio_sample_list = load_audio_text_image_video(data_in, fs=frontend.fs, audio_fs=kwargs.get("fs", 16000), data_type=kwargs.get("data_type", "sound"), tokenizer=tokenizer) + audio_sample_list = load_audio_text_image_video( + data_in, + fs=frontend.fs, + audio_fs=kwargs.get("fs", 16000), + data_type=kwargs.get("data_type", "sound"), + tokenizer=tokenizer, + ) time2 = time.perf_counter() meta_data["load_data"] = f"{time2 - time1:0.3f}" - speech, speech_lengths = extract_fbank(audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend) + speech, speech_lengths = extract_fbank( + audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend + ) time3 = time.perf_counter() meta_data["extract_feat"] = f"{time3 - time2:0.3f}" - meta_data["batch_data_time"] = speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 - + meta_data["batch_data_time"] = ( + speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 + ) + speech = speech.to(device=kwargs["device"]) speech_lengths = speech_lengths.to(device=kwargs["device"]) # Encoder @@ -477,33 +497,40 @@ class Paraformer(torch.nn.Module): encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) if isinstance(encoder_out, tuple): encoder_out = encoder_out[0] - + # predictor predictor_outs = self.calc_predictor(encoder_out, encoder_out_lens) - pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = predictor_outs[0], predictor_outs[1], \ - predictor_outs[2], predictor_outs[3] + pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = ( + predictor_outs[0], + predictor_outs[1], + predictor_outs[2], + predictor_outs[3], + ) pre_token_length = pre_token_length.round().long() if torch.max(pre_token_length) < 1: return [] - decoder_outs = self.cal_decoder_with_predictor(encoder_out, encoder_out_lens, pre_acoustic_embeds, - pre_token_length) + decoder_outs = self.cal_decoder_with_predictor( + encoder_out, encoder_out_lens, pre_acoustic_embeds, pre_token_length + ) decoder_out, ys_pad_lens = decoder_outs[0], decoder_outs[1] - results = [] b, n, d = decoder_out.size() if isinstance(key[0], (list, tuple)): key = key[0] if len(key) < b: - key = key*b + key = key * b for i in range(b): - x = encoder_out[i, :encoder_out_lens[i], :] - am_scores = decoder_out[i, :pre_token_length[i], :] + x = encoder_out[i, : encoder_out_lens[i], :] + am_scores = decoder_out[i, : pre_token_length[i], :] if self.beam_search is not None: nbest_hyps = self.beam_search( - x=x, am_scores=am_scores, maxlenratio=kwargs.get("maxlenratio", 0.0), minlenratio=kwargs.get("minlenratio", 0.0) + x=x, + am_scores=am_scores, + maxlenratio=kwargs.get("maxlenratio", 0.0), + minlenratio=kwargs.get("minlenratio", 0.0), ) - + nbest_hyps = nbest_hyps[: self.nbest] else: @@ -511,9 +538,7 @@ class Paraformer(torch.nn.Module): score = am_scores.max(dim=-1)[0] score = torch.sum(score, dim=-1) # pad with mask tokens to ensure compatibility with sos/eos tokens - yseq = torch.tensor( - [self.sos] + yseq.tolist() + [self.eos], device=yseq.device - ) + yseq = torch.tensor([self.sos] + yseq.tolist() + [self.eos], device=yseq.device) nbest_hyps = [Hypothesis(yseq=yseq, score=score)] for nbest_idx, hyp in enumerate(nbest_hyps): ibest_writer = None @@ -527,17 +552,21 @@ class Paraformer(torch.nn.Module): token_int = hyp.yseq[1:last_pos] else: token_int = hyp.yseq[1:last_pos].tolist() - + # remove blank symbol id, which is assumed to be 0 - token_int = list(filter(lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int)) - + token_int = list( + filter( + lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int + ) + ) + if tokenizer is not None: # Change integer-ids to tokens token = tokenizer.ids2tokens(token_int) text_postprocessed = tokenizer.tokens2text(token) if not hasattr(tokenizer, "bpemodel"): text_postprocessed, _ = postprocess_utils.sentence_postprocess(token) - + result_i = {"key": key[i], "text": text_postprocessed} if ibest_writer is not None: @@ -547,13 +576,13 @@ class Paraformer(torch.nn.Module): else: result_i = {"key": key[i], "token_int": token_int} results.append(result_i) - + return results, meta_data def export(self, **kwargs): from .export_meta import export_rebuild_model - if 'max_seq_len' not in kwargs: - kwargs['max_seq_len'] = 512 + + if "max_seq_len" not in kwargs: + kwargs["max_seq_len"] = 512 models = export_rebuild_model(model=self, **kwargs) return models - diff --git a/funasr/models/paraformer/search.py b/funasr/models/paraformer/search.py index 31f4de214..16e13dd9d 100644 --- a/funasr/models/paraformer/search.py +++ b/funasr/models/paraformer/search.py @@ -9,7 +9,10 @@ from itertools import chain from typing import Any, Dict, List, NamedTuple, Tuple, Union from funasr.metrics.common import end_detect -from funasr.models.transformer.scorers.scorer_interface import PartialScorerInterface, ScorerInterface +from funasr.models.transformer.scorers.scorer_interface import ( + PartialScorerInterface, + ScorerInterface, +) class Hypothesis(NamedTuple): @@ -317,9 +320,7 @@ class BeamSearchPara(torch.nn.Module): Hypothesis( score=weighted_scores[j], yseq=self.append_token(hyp.yseq, j), - scores=self.merge_scores( - hyp.scores, scores, j, part_scores, part_j - ), + scores=self.merge_scores(hyp.scores, scores, j, part_scores, part_j), states=self.merge_states(states, part_states, part_j), ) ) @@ -331,7 +332,11 @@ class BeamSearchPara(torch.nn.Module): return best_hyps def forward( - self, x: torch.Tensor, am_scores: torch.Tensor, maxlenratio: float = 0.0, minlenratio: float = 0.0 + self, + x: torch.Tensor, + am_scores: torch.Tensor, + maxlenratio: float = 0.0, + minlenratio: float = 0.0, ) -> List[Hypothesis]: """Perform beam search. @@ -375,8 +380,7 @@ class BeamSearchPara(torch.nn.Module): # check the number of hypotheses reaching to eos if len(nbest_hyps) == 0: logging.warning( - "there is no N-best results, perform recognition " - "again with smaller minlenratio." + "there is no N-best results, perform recognition " "again with smaller minlenratio." ) return ( [] @@ -387,17 +391,13 @@ class BeamSearchPara(torch.nn.Module): # report the best result best = nbest_hyps[0] for k, v in best.scores.items(): - logging.info( - f"{v:6.2f} * {self.weights[k]:3} = {v * self.weights[k]:6.2f} for {k}" - ) + logging.info(f"{v:6.2f} * {self.weights[k]:3} = {v * self.weights[k]:6.2f} for {k}") logging.info(f"total log probability: {best.score:.2f}") logging.info(f"normalized log probability: {best.score / len(best.yseq):.2f}") logging.info(f"total number of ended hypotheses: {len(nbest_hyps)}") if self.token_list is not None: logging.info( - "best hypo: " - + "".join([self.token_list[x.item()] for x in best.yseq[1:-1]]) - + "\n" + "best hypo: " + "".join([self.token_list[x.item()] for x in best.yseq[1:-1]]) + "\n" ) return nbest_hyps @@ -432,8 +432,7 @@ class BeamSearchPara(torch.nn.Module): if i == maxlen - 1: logging.info("adding in the last position in the loop") running_hyps = [ - h._replace(yseq=self.append_token(h.yseq, self.eos)) - for h in running_hyps + h._replace(yseq=self.append_token(h.yseq, self.eos)) for h in running_hyps ] # add ended hypotheses to a final list, and removed them from current hypotheses diff --git a/funasr/models/paraformer_streaming/export_meta.py b/funasr/models/paraformer_streaming/export_meta.py index 0193dc83d..9740ecccf 100644 --- a/funasr/models/paraformer_streaming/export_meta.py +++ b/funasr/models/paraformer_streaming/export_meta.py @@ -9,29 +9,29 @@ from funasr.register import tables def export_rebuild_model(model, **kwargs): - model.device = kwargs.get("device") - is_onnx = kwargs.get("type", "onnx") == "onnx" - encoder_class = tables.encoder_classes.get(kwargs["encoder"]+"Export") - model.encoder = encoder_class(model.encoder, onnx=is_onnx) - - predictor_class = tables.predictor_classes.get(kwargs["predictor"]+"Export") - model.predictor = predictor_class(model.predictor, onnx=is_onnx) + model.device = kwargs.get("device") + is_onnx = kwargs.get("type", "onnx") == "onnx" + encoder_class = tables.encoder_classes.get(kwargs["encoder"] + "Export") + model.encoder = encoder_class(model.encoder, onnx=is_onnx) + predictor_class = tables.predictor_classes.get(kwargs["predictor"] + "Export") + model.predictor = predictor_class(model.predictor, onnx=is_onnx) - decoder_class = tables.decoder_classes.get(kwargs["decoder"]+"Export") - model.decoder = decoder_class(model.decoder, onnx=is_onnx) - - from funasr.utils.torch_function import sequence_mask - model.make_pad_mask = sequence_mask(kwargs['max_seq_len'], flip=False) - - model.forward = types.MethodType(export_forward, model) - model.export_dummy_inputs = types.MethodType(export_dummy_inputs, model) - model.export_input_names = types.MethodType(export_input_names, model) - model.export_output_names = types.MethodType(export_output_names, model) - model.export_dynamic_axes = types.MethodType(export_dynamic_axes, model) - model.export_name = types.MethodType(export_name, model) - - return model + decoder_class = tables.decoder_classes.get(kwargs["decoder"] + "Export") + model.decoder = decoder_class(model.decoder, onnx=is_onnx) + + from funasr.utils.torch_function import sequence_mask + + model.make_pad_mask = sequence_mask(kwargs["max_seq_len"], flip=False) + + model.forward = types.MethodType(export_forward, model) + model.export_dummy_inputs = types.MethodType(export_dummy_inputs, model) + model.export_input_names = types.MethodType(export_input_names, model) + model.export_output_names = types.MethodType(export_output_names, model) + model.export_dynamic_axes = types.MethodType(export_dynamic_axes, model) + model.export_name = types.MethodType(export_name, model) + + return model def export_rebuild_model(model, **kwargs): @@ -39,24 +39,25 @@ def export_rebuild_model(model, **kwargs): is_onnx = kwargs.get("type", "onnx") == "onnx" encoder_class = tables.encoder_classes.get(kwargs["encoder"] + "Export") model.encoder = encoder_class(model.encoder, onnx=is_onnx) - + predictor_class = tables.predictor_classes.get(kwargs["predictor"] + "Export") model.predictor = predictor_class(model.predictor, onnx=is_onnx) - + if kwargs["decoder"] == "ParaformerSANMDecoder": kwargs["decoder"] = "ParaformerSANMDecoderOnline" decoder_class = tables.decoder_classes.get(kwargs["decoder"] + "Export") model.decoder = decoder_class(model.decoder, onnx=is_onnx) - + from funasr.utils.torch_function import sequence_mask - + model.make_pad_mask = sequence_mask(max_seq_len=None, flip=False) - + import copy import types + encoder_model = copy.copy(model) decoder_model = copy.copy(model) - + # encoder encoder_model.forward = types.MethodType(export_encoder_forward, encoder_model) encoder_model.export_dummy_inputs = types.MethodType(export_encoder_dummy_inputs, encoder_model) @@ -64,7 +65,7 @@ def export_rebuild_model(model, **kwargs): encoder_model.export_output_names = types.MethodType(export_encoder_output_names, encoder_model) encoder_model.export_dynamic_axes = types.MethodType(export_encoder_dynamic_axes, encoder_model) encoder_model.export_name = types.MethodType(export_encoder_name, encoder_model) - + # decoder decoder_model.forward = types.MethodType(export_decoder_forward, decoder_model) decoder_model.export_dummy_inputs = types.MethodType(export_decoder_dummy_inputs, decoder_model) @@ -72,7 +73,7 @@ def export_rebuild_model(model, **kwargs): decoder_model.export_output_names = types.MethodType(export_decoder_output_names, decoder_model) decoder_model.export_dynamic_axes = types.MethodType(export_decoder_dynamic_axes, decoder_model) decoder_model.export_name = types.MethodType(export_decoder_name, decoder_model) - + return encoder_model, decoder_model @@ -84,11 +85,11 @@ def export_encoder_forward( # a. To device batch = {"speech": speech, "speech_lengths": speech_lengths, "online": True} # batch = to_device(batch, device=self.device) - + enc, enc_len = self.encoder(**batch) mask = self.make_pad_mask(enc_len)[:, None, :] alphas, _ = self.predictor.forward_cnn(enc, mask) - + return enc, enc_len, alphas @@ -99,33 +100,24 @@ def export_encoder_dummy_inputs(self): def export_encoder_input_names(self): - return ['speech', 'speech_lengths'] + return ["speech", "speech_lengths"] def export_encoder_output_names(self): - return ['enc', 'enc_len', 'alphas'] + return ["enc", "enc_len", "alphas"] def export_encoder_dynamic_axes(self): return { - 'speech': { - 0: 'batch_size', - 1: 'feats_length' + "speech": {0: "batch_size", 1: "feats_length"}, + "speech_lengths": { + 0: "batch_size", }, - 'speech_lengths': { - 0: 'batch_size', - }, - 'enc': { - 0: 'batch_size', - 1: 'feats_length' - }, - 'enc_len': { - 0: 'batch_size', - }, - 'alphas': { - 0: 'batch_size', - 1: 'feats_length' + "enc": {0: "batch_size", 1: "feats_length"}, + "enc_len": { + 0: "batch_size", }, + "alphas": {0: "batch_size", 1: "feats_length"}, } @@ -141,9 +133,11 @@ def export_decoder_forward( acoustic_embeds_len: torch.Tensor, *args, ): - decoder_out, out_caches = self.decoder(enc, enc_len, acoustic_embeds, acoustic_embeds_len, *args) + decoder_out, out_caches = self.decoder( + enc, enc_len, acoustic_embeds, acoustic_embeds_len, *args + ) sample_ids = decoder_out.argmax(dim=-1) - + return decoder_out, sample_ids, out_caches @@ -165,4 +159,4 @@ def export_decoder_dynamic_axes(self): def export_decoder_name(self): - return "decoder.onnx" \ No newline at end of file + return "decoder.onnx" diff --git a/funasr/models/paraformer_streaming/model.py b/funasr/models/paraformer_streaming/model.py index 499b48752..f28761494 100644 --- a/funasr/models/paraformer_streaming/model.py +++ b/funasr/models/paraformer_streaming/model.py @@ -41,28 +41,33 @@ class ParaformerStreaming(Paraformer): Paraformer: Fast and Accurate Parallel Transformer for Non-autoregressive End-to-End Speech Recognition https://arxiv.org/abs/2206.08317 """ - + def __init__( self, *args, **kwargs, ): - + super().__init__(*args, **kwargs) - + # import pdb; # pdb.set_trace() self.sampling_ratio = kwargs.get("sampling_ratio", 0.2) - self.scama_mask = None - if hasattr(self.encoder, "overlap_chunk_cls") and self.encoder.overlap_chunk_cls is not None: - from funasr.models.scama.chunk_utilis import build_scama_mask_for_cross_attention_decoder - self.build_scama_mask_for_cross_attention_decoder_fn = build_scama_mask_for_cross_attention_decoder + if ( + hasattr(self.encoder, "overlap_chunk_cls") + and self.encoder.overlap_chunk_cls is not None + ): + from funasr.models.scama.chunk_utilis import ( + build_scama_mask_for_cross_attention_decoder, + ) + + self.build_scama_mask_for_cross_attention_decoder_fn = ( + build_scama_mask_for_cross_attention_decoder + ) self.decoder_attention_chunk_type = kwargs.get("decoder_attention_chunk_type", "chunk") - - def forward( self, speech: torch.Tensor, @@ -85,49 +90,52 @@ class ParaformerStreaming(Paraformer): text_lengths = text_lengths[:, 0] if len(speech_lengths.size()) > 1: speech_lengths = speech_lengths[:, 0] - + batch_size = speech.shape[0] - + # Encoder if hasattr(self.encoder, "overlap_chunk_cls"): ind = self.encoder.overlap_chunk_cls.random_choice(self.training, decoding_ind) encoder_out, encoder_out_lens = self.encode(speech, speech_lengths, ind=ind) else: encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) - + loss_ctc, cer_ctc = None, None loss_pre = None stats = dict() - + # decoder: CTC branch if self.ctc_weight > 0.0: if hasattr(self.encoder, "overlap_chunk_cls"): - encoder_out_ctc, encoder_out_lens_ctc = self.encoder.overlap_chunk_cls.remove_chunk(encoder_out, - encoder_out_lens, - chunk_outs=None) + encoder_out_ctc, encoder_out_lens_ctc = self.encoder.overlap_chunk_cls.remove_chunk( + encoder_out, encoder_out_lens, chunk_outs=None + ) else: encoder_out_ctc, encoder_out_lens_ctc = encoder_out, encoder_out_lens - + loss_ctc, cer_ctc = self._calc_ctc_loss( encoder_out_ctc, encoder_out_lens_ctc, text, text_lengths ) # Collect CTC branch stats stats["loss_ctc"] = loss_ctc.detach() if loss_ctc is not None else None stats["cer_ctc"] = cer_ctc - + # decoder: Attention decoder branch loss_att, acc_att, cer_att, wer_att, loss_pre, pre_loss_att = self._calc_att_predictor_loss( encoder_out, encoder_out_lens, text, text_lengths ) - + # 3. CTC-Att loss definition if self.ctc_weight == 0.0: loss = loss_att + loss_pre * self.predictor_weight else: - loss = self.ctc_weight * loss_ctc + ( - 1 - self.ctc_weight) * loss_att + loss_pre * self.predictor_weight - + loss = ( + self.ctc_weight * loss_ctc + + (1 - self.ctc_weight) * loss_att + + loss_pre * self.predictor_weight + ) + # Collect Attn branch stats stats["loss_att"] = loss_att.detach() if loss_att is not None else None stats["pre_loss_att"] = pre_loss_att.detach() if pre_loss_att is not None else None @@ -135,17 +143,21 @@ class ParaformerStreaming(Paraformer): stats["cer"] = cer_att stats["wer"] = wer_att stats["loss_pre"] = loss_pre.detach().cpu() if loss_pre is not None else None - + stats["loss"] = torch.clone(loss.detach()) - + # force_gatherable: to-device and to-tensor if scalar for DataParallel if self.length_normalized_loss: batch_size = (text_lengths + self.predictor_bias).sum() loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) return loss, stats, weight - + def encode_chunk( - self, speech: torch.Tensor, speech_lengths: torch.Tensor, cache: dict = None, **kwargs, + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + cache: dict = None, + **kwargs, ) -> Tuple[torch.Tensor, torch.Tensor]: """Frontend + Encoder. Note that this method is used by asr_inference.py Args: @@ -154,22 +166,24 @@ class ParaformerStreaming(Paraformer): ind: int """ with autocast(False): - + # Data augmentation if self.specaug is not None and self.training: speech, speech_lengths = self.specaug(speech, speech_lengths) - + # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN if self.normalize is not None: speech, speech_lengths = self.normalize(speech, speech_lengths) - + # Forward encoder - encoder_out, encoder_out_lens, _ = self.encoder.forward_chunk(speech, speech_lengths, cache=cache["encoder"]) + encoder_out, encoder_out_lens, _ = self.encoder.forward_chunk( + speech, speech_lengths, cache=cache["encoder"] + ) if isinstance(encoder_out, tuple): encoder_out = encoder_out[0] - + return encoder_out, torch.tensor([encoder_out.size(1)]) - + def _calc_att_predictor_loss( self, encoder_out: torch.Tensor, @@ -177,41 +191,49 @@ class ParaformerStreaming(Paraformer): ys_pad: torch.Tensor, ys_pad_lens: torch.Tensor, ): - encoder_out_mask = (~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :]).to( - encoder_out.device) + encoder_out_mask = ( + ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :] + ).to(encoder_out.device) if self.predictor_bias == 1: _, ys_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id) ys_pad_lens = ys_pad_lens + self.predictor_bias mask_chunk_predictor = None if self.encoder.overlap_chunk_cls is not None: - mask_chunk_predictor = self.encoder.overlap_chunk_cls.get_mask_chunk_predictor(None, - device=encoder_out.device, - batch_size=encoder_out.size( - 0)) - mask_shfit_chunk = self.encoder.overlap_chunk_cls.get_mask_shfit_chunk(None, device=encoder_out.device, - batch_size=encoder_out.size(0)) + mask_chunk_predictor = self.encoder.overlap_chunk_cls.get_mask_chunk_predictor( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + mask_shfit_chunk = self.encoder.overlap_chunk_cls.get_mask_shfit_chunk( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) encoder_out = encoder_out * mask_shfit_chunk - pre_acoustic_embeds, pre_token_length, pre_alphas, _ = self.predictor(encoder_out, - ys_pad, - encoder_out_mask, - ignore_id=self.ignore_id, - mask_chunk_predictor=mask_chunk_predictor, - target_label_length=ys_pad_lens, - ) - predictor_alignments, predictor_alignments_len = self.predictor.gen_frame_alignments(pre_alphas, - encoder_out_lens) - + pre_acoustic_embeds, pre_token_length, pre_alphas, _ = self.predictor( + encoder_out, + ys_pad, + encoder_out_mask, + ignore_id=self.ignore_id, + mask_chunk_predictor=mask_chunk_predictor, + target_label_length=ys_pad_lens, + ) + predictor_alignments, predictor_alignments_len = self.predictor.gen_frame_alignments( + pre_alphas, encoder_out_lens + ) + scama_mask = None - if self.encoder.overlap_chunk_cls is not None and self.decoder_attention_chunk_type == 'chunk': + if ( + self.encoder.overlap_chunk_cls is not None + and self.decoder_attention_chunk_type == "chunk" + ): encoder_chunk_size = self.encoder.overlap_chunk_cls.chunk_size_pad_shift_cur attention_chunk_center_bias = 0 attention_chunk_size = encoder_chunk_size - decoder_att_look_back_factor = self.encoder.overlap_chunk_cls.decoder_att_look_back_factor_cur - mask_shift_att_chunk_decoder = self.encoder.overlap_chunk_cls. \ - get_mask_shift_att_chunk_decoder(None, - device=encoder_out.device, - batch_size=encoder_out.size(0) - ) + decoder_att_look_back_factor = ( + self.encoder.overlap_chunk_cls.decoder_att_look_back_factor_cur + ) + mask_shift_att_chunk_decoder = ( + self.encoder.overlap_chunk_cls.get_mask_shift_att_chunk_decoder( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + ) scama_mask = self.build_scama_mask_for_cross_attention_decoder_fn( predictor_alignments=predictor_alignments, encoder_sequence_length=encoder_out_lens, @@ -228,31 +250,41 @@ class ParaformerStreaming(Paraformer): is_training=self.training, ) elif self.encoder.overlap_chunk_cls is not None: - encoder_out, encoder_out_lens = self.encoder.overlap_chunk_cls.remove_chunk(encoder_out, - encoder_out_lens, - chunk_outs=None) + encoder_out, encoder_out_lens = self.encoder.overlap_chunk_cls.remove_chunk( + encoder_out, encoder_out_lens, chunk_outs=None + ) # 0. sampler decoder_out_1st = None pre_loss_att = None if self.sampling_ratio > 0.0: if self.use_1st_decoder_loss: - sematic_embeds, decoder_out_1st, pre_loss_att = \ - self.sampler_with_grad(encoder_out, encoder_out_lens, ys_pad, - ys_pad_lens, pre_acoustic_embeds, scama_mask) + sematic_embeds, decoder_out_1st, pre_loss_att = self.sampler_with_grad( + encoder_out, + encoder_out_lens, + ys_pad, + ys_pad_lens, + pre_acoustic_embeds, + scama_mask, + ) else: - sematic_embeds, decoder_out_1st = \ - self.sampler(encoder_out, encoder_out_lens, ys_pad, - ys_pad_lens, pre_acoustic_embeds, scama_mask) + sematic_embeds, decoder_out_1st = self.sampler( + encoder_out, + encoder_out_lens, + ys_pad, + ys_pad_lens, + pre_acoustic_embeds, + scama_mask, + ) else: sematic_embeds = pre_acoustic_embeds - + # 1. Forward decoder decoder_outs = self.decoder( encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens, scama_mask ) decoder_out, _ = decoder_outs[0], decoder_outs[1] - + if decoder_out_1st is None: decoder_out_1st = decoder_out # 2. Compute attention loss @@ -263,19 +295,29 @@ class ParaformerStreaming(Paraformer): ignore_label=self.ignore_id, ) loss_pre = self.criterion_pre(ys_pad_lens.type_as(pre_token_length), pre_token_length) - + # Compute cer/wer using attention-decoder if self.training or self.error_calculator is None: cer_att, wer_att = None, None else: ys_hat = decoder_out_1st.argmax(dim=-1) cer_att, wer_att = self.error_calculator(ys_hat.cpu(), ys_pad.cpu()) - + return loss_att, acc_att, cer_att, wer_att, loss_pre, pre_loss_att - - def sampler(self, encoder_out, encoder_out_lens, ys_pad, ys_pad_lens, pre_acoustic_embeds, chunk_mask=None): - - tgt_mask = (~make_pad_mask(ys_pad_lens, maxlen=ys_pad_lens.max())[:, :, None]).to(ys_pad.device) + + def sampler( + self, + encoder_out, + encoder_out_lens, + ys_pad, + ys_pad_lens, + pre_acoustic_embeds, + chunk_mask=None, + ): + + tgt_mask = (~make_pad_mask(ys_pad_lens, maxlen=ys_pad_lens.max())[:, :, None]).to( + ys_pad.device + ) ys_pad_masked = ys_pad * tgt_mask[:, :, 0] if self.share_embedding: ys_pad_embed = self.decoder.output_layer.weight[ys_pad_masked] @@ -293,52 +335,65 @@ class ParaformerStreaming(Paraformer): input_mask = torch.ones_like(nonpad_positions) bsz, seq_len = ys_pad.size() for li in range(bsz): - target_num = (((seq_lens[li] - same_num[li].sum()).float()) * self.sampling_ratio).long() + target_num = ( + ((seq_lens[li] - same_num[li].sum()).float()) * self.sampling_ratio + ).long() if target_num > 0: - input_mask[li].scatter_(dim=0, index=torch.randperm(seq_lens[li])[:target_num].cuda(), value=0) + input_mask[li].scatter_( + dim=0, index=torch.randperm(seq_lens[li])[:target_num].cuda(), value=0 + ) input_mask = input_mask.eq(1) input_mask = input_mask.masked_fill(~nonpad_positions, False) input_mask_expand_dim = input_mask.unsqueeze(2).to(pre_acoustic_embeds.device) - - sematic_embeds = pre_acoustic_embeds.masked_fill(~input_mask_expand_dim, 0) + ys_pad_embed.masked_fill( - input_mask_expand_dim, 0) + + sematic_embeds = pre_acoustic_embeds.masked_fill( + ~input_mask_expand_dim, 0 + ) + ys_pad_embed.masked_fill(input_mask_expand_dim, 0) return sematic_embeds * tgt_mask, decoder_out * tgt_mask - def calc_predictor(self, encoder_out, encoder_out_lens): - - encoder_out_mask = (~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :]).to( - encoder_out.device) + + encoder_out_mask = ( + ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :] + ).to(encoder_out.device) mask_chunk_predictor = None if self.encoder.overlap_chunk_cls is not None: - mask_chunk_predictor = self.encoder.overlap_chunk_cls.get_mask_chunk_predictor(None, - device=encoder_out.device, - batch_size=encoder_out.size( - 0)) - mask_shfit_chunk = self.encoder.overlap_chunk_cls.get_mask_shfit_chunk(None, device=encoder_out.device, - batch_size=encoder_out.size(0)) + mask_chunk_predictor = self.encoder.overlap_chunk_cls.get_mask_chunk_predictor( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + mask_shfit_chunk = self.encoder.overlap_chunk_cls.get_mask_shfit_chunk( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) encoder_out = encoder_out * mask_shfit_chunk - pre_acoustic_embeds, pre_token_length, pre_alphas, pre_peak_index = self.predictor(encoder_out, - None, - encoder_out_mask, - ignore_id=self.ignore_id, - mask_chunk_predictor=mask_chunk_predictor, - target_label_length=None, - ) - predictor_alignments, predictor_alignments_len = self.predictor.gen_frame_alignments(pre_alphas, - encoder_out_lens + 1 if self.predictor.tail_threshold > 0.0 else encoder_out_lens) - + pre_acoustic_embeds, pre_token_length, pre_alphas, pre_peak_index = self.predictor( + encoder_out, + None, + encoder_out_mask, + ignore_id=self.ignore_id, + mask_chunk_predictor=mask_chunk_predictor, + target_label_length=None, + ) + predictor_alignments, predictor_alignments_len = self.predictor.gen_frame_alignments( + pre_alphas, + encoder_out_lens + 1 if self.predictor.tail_threshold > 0.0 else encoder_out_lens, + ) + scama_mask = None - if self.encoder.overlap_chunk_cls is not None and self.decoder_attention_chunk_type == 'chunk': + if ( + self.encoder.overlap_chunk_cls is not None + and self.decoder_attention_chunk_type == "chunk" + ): encoder_chunk_size = self.encoder.overlap_chunk_cls.chunk_size_pad_shift_cur attention_chunk_center_bias = 0 attention_chunk_size = encoder_chunk_size - decoder_att_look_back_factor = self.encoder.overlap_chunk_cls.decoder_att_look_back_factor_cur - mask_shift_att_chunk_decoder = self.encoder.overlap_chunk_cls. \ - get_mask_shift_att_chunk_decoder(None, - device=encoder_out.device, - batch_size=encoder_out.size(0) - ) + decoder_att_look_back_factor = ( + self.encoder.overlap_chunk_cls.decoder_att_look_back_factor_cur + ) + mask_shift_att_chunk_decoder = ( + self.encoder.overlap_chunk_cls.get_mask_shift_att_chunk_decoder( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + ) scama_mask = self.build_scama_mask_for_cross_attention_decoder_fn( predictor_alignments=predictor_alignments, encoder_sequence_length=encoder_out_lens, @@ -355,30 +410,32 @@ class ParaformerStreaming(Paraformer): is_training=self.training, ) self.scama_mask = scama_mask - + return pre_acoustic_embeds, pre_token_length, pre_alphas, pre_peak_index - + def calc_predictor_chunk(self, encoder_out, encoder_out_lens, cache=None, **kwargs): is_final = kwargs.get("is_final", False) return self.predictor.forward_chunk(encoder_out, cache["encoder"], is_final=is_final) - - def cal_decoder_with_predictor(self, encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens): + + def cal_decoder_with_predictor( + self, encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens + ): decoder_outs = self.decoder( encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens, self.scama_mask ) decoder_out = decoder_outs[0] decoder_out = torch.log_softmax(decoder_out, dim=-1) return decoder_out, ys_pad_lens - - def cal_decoder_with_predictor_chunk(self, encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens, cache=None): - decoder_outs = self.decoder.forward_chunk( - encoder_out, sematic_embeds, cache["decoder"] - ) + + def cal_decoder_with_predictor_chunk( + self, encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens, cache=None + ): + decoder_outs = self.decoder.forward_chunk(encoder_out, sematic_embeds, cache["decoder"]) decoder_out = decoder_outs decoder_out = torch.log_softmax(decoder_out, dim=-1) return decoder_out, ys_pad_lens - + def init_cache(self, cache: dict = {}, **kwargs): chunk_size = kwargs.get("chunk_size", [0, 10, 5]) encoder_chunk_look_back = kwargs.get("encoder_chunk_look_back", 0) @@ -387,51 +444,67 @@ class ParaformerStreaming(Paraformer): enc_output_size = kwargs["encoder_conf"]["output_size"] feats_dims = kwargs["frontend_conf"]["n_mels"] * kwargs["frontend_conf"]["lfr_m"] - cache_encoder = {"start_idx": 0, "cif_hidden": torch.zeros((batch_size, 1, enc_output_size)), - "cif_alphas": torch.zeros((batch_size, 1)), "chunk_size": chunk_size, - "encoder_chunk_look_back": encoder_chunk_look_back, "last_chunk": False, "opt": None, - "feats": torch.zeros((batch_size, chunk_size[0] + chunk_size[2], feats_dims)), - "tail_chunk": False} + cache_encoder = { + "start_idx": 0, + "cif_hidden": torch.zeros((batch_size, 1, enc_output_size)), + "cif_alphas": torch.zeros((batch_size, 1)), + "chunk_size": chunk_size, + "encoder_chunk_look_back": encoder_chunk_look_back, + "last_chunk": False, + "opt": None, + "feats": torch.zeros((batch_size, chunk_size[0] + chunk_size[2], feats_dims)), + "tail_chunk": False, + } cache["encoder"] = cache_encoder - - cache_decoder = {"decode_fsmn": None, "decoder_chunk_look_back": decoder_chunk_look_back, "opt": None, - "chunk_size": chunk_size} + + cache_decoder = { + "decode_fsmn": None, + "decoder_chunk_look_back": decoder_chunk_look_back, + "opt": None, + "chunk_size": chunk_size, + } cache["decoder"] = cache_decoder cache["frontend"] = {} cache["prev_samples"] = torch.empty(0) - + return cache - - def generate_chunk(self, - speech, - speech_lengths=None, - key: list = None, - tokenizer=None, - frontend=None, - **kwargs, - ): + + def generate_chunk( + self, + speech, + speech_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): cache = kwargs.get("cache", {}) speech = speech.to(device=kwargs["device"]) speech_lengths = speech_lengths.to(device=kwargs["device"]) - + # Encoder - encoder_out, encoder_out_lens = self.encode_chunk(speech, speech_lengths, cache=cache, is_final=kwargs.get("is_final", False)) + encoder_out, encoder_out_lens = self.encode_chunk( + speech, speech_lengths, cache=cache, is_final=kwargs.get("is_final", False) + ) if isinstance(encoder_out, tuple): encoder_out = encoder_out[0] - + # predictor - predictor_outs = self.calc_predictor_chunk(encoder_out, encoder_out_lens, cache=cache, is_final=kwargs.get("is_final", False)) - pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = predictor_outs[0], predictor_outs[1], \ - predictor_outs[2], predictor_outs[3] + predictor_outs = self.calc_predictor_chunk( + encoder_out, encoder_out_lens, cache=cache, is_final=kwargs.get("is_final", False) + ) + pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = ( + predictor_outs[0], + predictor_outs[1], + predictor_outs[2], + predictor_outs[3], + ) pre_token_length = pre_token_length.round().long() if torch.max(pre_token_length) < 1: return [] - decoder_outs = self.cal_decoder_with_predictor_chunk(encoder_out, - encoder_out_lens, - pre_acoustic_embeds, - pre_token_length, - cache=cache - ) + decoder_outs = self.cal_decoder_with_predictor_chunk( + encoder_out, encoder_out_lens, pre_acoustic_embeds, pre_token_length, cache=cache + ) decoder_out, ys_pad_lens = decoder_outs[0], decoder_outs[1] results = [] @@ -439,123 +512,144 @@ class ParaformerStreaming(Paraformer): if isinstance(key[0], (list, tuple)): key = key[0] for i in range(b): - x = encoder_out[i, :encoder_out_lens[i], :] - am_scores = decoder_out[i, :pre_token_length[i], :] + x = encoder_out[i, : encoder_out_lens[i], :] + am_scores = decoder_out[i, : pre_token_length[i], :] if self.beam_search is not None: nbest_hyps = self.beam_search( - x=x, am_scores=am_scores, maxlenratio=kwargs.get("maxlenratio", 0.0), - minlenratio=kwargs.get("minlenratio", 0.0) + x=x, + am_scores=am_scores, + maxlenratio=kwargs.get("maxlenratio", 0.0), + minlenratio=kwargs.get("minlenratio", 0.0), ) - + nbest_hyps = nbest_hyps[: self.nbest] else: - + yseq = am_scores.argmax(dim=-1) score = am_scores.max(dim=-1)[0] score = torch.sum(score, dim=-1) # pad with mask tokens to ensure compatibility with sos/eos tokens - yseq = torch.tensor( - [self.sos] + yseq.tolist() + [self.eos], device=yseq.device - ) + yseq = torch.tensor([self.sos] + yseq.tolist() + [self.eos], device=yseq.device) nbest_hyps = [Hypothesis(yseq=yseq, score=score)] for nbest_idx, hyp in enumerate(nbest_hyps): - + # remove sos/eos and get results last_pos = -1 if isinstance(hyp.yseq, list): token_int = hyp.yseq[1:last_pos] else: token_int = hyp.yseq[1:last_pos].tolist() - + # remove blank symbol id, which is assumed to be 0 - token_int = list(filter(lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int)) - + token_int = list( + filter( + lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int + ) + ) # Change integer-ids to tokens token = tokenizer.ids2tokens(token_int) # text = tokenizer.tokens2text(token) - + result_i = token - results.extend(result_i) - + return results - - def inference(self, - data_in, - data_lengths=None, - key: list = None, - tokenizer=None, - frontend=None, - cache: dict={}, - **kwargs, - ): + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + cache: dict = {}, + **kwargs, + ): # init beamsearch is_use_ctc = kwargs.get("decoding_ctc_weight", 0.0) > 0.00001 and self.ctc != None - is_use_lm = kwargs.get("lm_weight", 0.0) > 0.00001 and kwargs.get("lm_file", None) is not None + is_use_lm = ( + kwargs.get("lm_weight", 0.0) > 0.00001 and kwargs.get("lm_file", None) is not None + ) if self.beam_search is None and (is_use_lm or is_use_ctc): logging.info("enable beam_search") self.init_beam_search(**kwargs) self.nbest = kwargs.get("nbest", 1) - + if len(cache) == 0: self.init_cache(cache, **kwargs) - - + meta_data = {} chunk_size = kwargs.get("chunk_size", [0, 10, 5]) chunk_stride_samples = int(chunk_size[1] * 960) # 600ms - + time1 = time.perf_counter() cfg = {"is_final": kwargs.get("is_final", False)} - audio_sample_list = load_audio_text_image_video(data_in, - fs=frontend.fs, - audio_fs=kwargs.get("fs", 16000), - data_type=kwargs.get("data_type", "sound"), - tokenizer=tokenizer, - cache=cfg, - ) - _is_final = cfg["is_final"] # if data_in is a file or url, set is_final=True - + audio_sample_list = load_audio_text_image_video( + data_in, + fs=frontend.fs, + audio_fs=kwargs.get("fs", 16000), + data_type=kwargs.get("data_type", "sound"), + tokenizer=tokenizer, + cache=cfg, + ) + _is_final = cfg["is_final"] # if data_in is a file or url, set is_final=True + time2 = time.perf_counter() meta_data["load_data"] = f"{time2 - time1:0.3f}" assert len(audio_sample_list) == 1, "batch_size must be set 1" - + audio_sample = torch.cat((cache["prev_samples"], audio_sample_list[0])) - + n = int(len(audio_sample) // chunk_stride_samples + int(_is_final)) - m = int(len(audio_sample) % chunk_stride_samples * (1-int(_is_final))) + m = int(len(audio_sample) % chunk_stride_samples * (1 - int(_is_final))) tokens = [] for i in range(n): - kwargs["is_final"] = _is_final and i == n -1 - audio_sample_i = audio_sample[i*chunk_stride_samples:(i+1)*chunk_stride_samples] + kwargs["is_final"] = _is_final and i == n - 1 + audio_sample_i = audio_sample[i * chunk_stride_samples : (i + 1) * chunk_stride_samples] if kwargs["is_final"] and len(audio_sample_i) < 960: cache["encoder"]["tail_chunk"] = True speech = cache["encoder"]["feats"] - speech_lengths = torch.tensor([speech.shape[1]], dtype=torch.int64).to(speech.device) + speech_lengths = torch.tensor([speech.shape[1]], dtype=torch.int64).to( + speech.device + ) else: # extract fbank feats - speech, speech_lengths = extract_fbank([audio_sample_i], data_type=kwargs.get("data_type", "sound"), - frontend=frontend, cache=cache["frontend"], is_final=kwargs["is_final"]) + speech, speech_lengths = extract_fbank( + [audio_sample_i], + data_type=kwargs.get("data_type", "sound"), + frontend=frontend, + cache=cache["frontend"], + is_final=kwargs["is_final"], + ) time3 = time.perf_counter() meta_data["extract_feat"] = f"{time3 - time2:0.3f}" - meta_data["batch_data_time"] = speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 - - tokens_i = self.generate_chunk(speech, speech_lengths, key=key, tokenizer=tokenizer, cache=cache, frontend=frontend, **kwargs) + meta_data["batch_data_time"] = ( + speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 + ) + + tokens_i = self.generate_chunk( + speech, + speech_lengths, + key=key, + tokenizer=tokenizer, + cache=cache, + frontend=frontend, + **kwargs, + ) tokens.extend(tokens_i) - + text_postprocessed, _ = postprocess_utils.sentence_postprocess(tokens) - + result_i = {"key": key[0], "text": text_postprocessed} result = [result_i] - - + cache["prev_samples"] = audio_sample[:-m] if _is_final: self.init_cache(cache, **kwargs) - + if kwargs.get("output_dir"): if not hasattr(self, "writer"): self.writer = DatadirWriter(kwargs.get("output_dir")) @@ -564,8 +658,9 @@ class ParaformerStreaming(Paraformer): ibest_writer["text"][key[0]] = text_postprocessed return result, meta_data - + def export(self, **kwargs): from .export_meta import export_rebuild_model + models = export_rebuild_model(model=self, **kwargs) - return models \ No newline at end of file + return models diff --git a/funasr/models/qwen_audio/audio.py b/funasr/models/qwen_audio/audio.py index 1a3789773..a1aad79ea 100644 --- a/funasr/models/qwen_audio/audio.py +++ b/funasr/models/qwen_audio/audio.py @@ -15,11 +15,11 @@ from typing import Optional, Union from funasr.register import tables - def exact_div(x, y): assert x % y == 0 return x // y + # hard-coded audio hyperparameters SAMPLE_RATE = 16000 N_FFT = 400 @@ -34,30 +34,37 @@ FRAMES_PER_SECOND = exact_div(SAMPLE_RATE, HOP_LENGTH) # 10ms per audio frame TOKENS_PER_SECOND = exact_div(SAMPLE_RATE, N_SAMPLES_PER_TOKEN) # 20ms per audio token - def get_T_after_cnn(L_in, dilation=1): - for (padding, kernel_size, stride) in eval("[(1,3,1)] + [(1,3,2)] "): + for padding, kernel_size, stride in eval("[(1,3,1)] + [(1,3,2)] "): L_out = L_in + 2 * padding - dilation * (kernel_size - 1) - 1 L_out = 1 + L_out // stride L_in = L_out return L_out + def load_bytesio_audio(content, sr: int = SAMPLE_RATE): cmd = [ "ffmpeg", "-nostdin", - "-threads", "0", - "-i", "pipe:", - "-f", "s16le", - "-ac", "1", - "-acodec", "pcm_s16le", - "-ar", str(sr), - "pipe:" + "-threads", + "0", + "-i", + "pipe:", + "-f", + "s16le", + "-ac", + "1", + "-acodec", + "pcm_s16le", + "-ar", + str(sr), + "pipe:", ] p = Popen(cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE, bufsize=-1) out, _ = p.communicate(input=content) return np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0 + def load_audio(file: str, sr: int = SAMPLE_RATE): """ Open an audio file and read as mono waveform, resampling as necessary @@ -104,9 +111,7 @@ def pad_or_trim(array, length: int = N_SAMPLES, *, axis: int = -1): """ if torch.is_tensor(array): if array.shape[axis] > length: - array = array.index_select( - dim=axis, index=torch.arange(length, device=array.device) - ) + array = array.index_select(dim=axis, index=torch.arange(length, device=array.device)) if array.shape[axis] < length: pad_widths = [(0, 0)] * array.ndim @@ -123,15 +128,14 @@ def pad_or_trim(array, length: int = N_SAMPLES, *, axis: int = -1): return array + def trim(array, length: int = N_SAMPLES, *, axis: int = -1): """ Pad or trim the audio array to N_SAMPLES, as expected by the encoder. """ if torch.is_tensor(array): if array.shape[axis] > length: - array = array.index_select( - dim=axis, index=torch.arange(length, device=array.device) - ) + array = array.index_select(dim=axis, index=torch.arange(length, device=array.device)) else: if array.shape[axis] > length: array = array.take(indices=range(length), axis=axis) @@ -151,7 +155,7 @@ def mel_filters(device, n_mels: int = N_MELS) -> torch.Tensor: """ assert n_mels == 80, f"Unsupported n_mels: {n_mels}" with np.load( - os.path.join(os.path.dirname(__file__), "mel_filters.npz") # todo + os.path.join(os.path.dirname(__file__), "mel_filters.npz") # todo # os.path.join("assets", "mel_filters.npz") ) as f: return torch.from_numpy(f[f"mel_{n_mels}"]).to(device) @@ -227,8 +231,6 @@ class LayerNorm(nn.LayerNorm): return super().forward(x).type(x.dtype) - - class Linear(nn.Linear): def forward(self, x: Tensor) -> Tensor: return F.linear( @@ -239,9 +241,7 @@ class Linear(nn.Linear): class Conv1d(nn.Conv1d): - def _conv_forward( - self, x: Tensor, weight: Tensor, bias: Optional[Tensor] - ) -> Tensor: + def _conv_forward(self, x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> Tensor: return super()._conv_forward( x, weight.to(x.dtype), None if bias is None else bias.to(x.dtype) ) @@ -287,9 +287,7 @@ class MultiHeadAttention(nn.Module): wv, qk = self.qkv_attention(q, k, v, mask) return self.out(wv), qk - def qkv_attention( - self, q: Tensor, k: Tensor, v: Tensor, mask: Optional[Tensor] = None - ): + def qkv_attention(self, q: Tensor, k: Tensor, v: Tensor, mask: Optional[Tensor] = None): n_batch, n_ctx, n_state = q.shape scale = (n_state // self.n_head) ** -0.25 q = q.view(*q.shape[:2], self.n_head, -1).permute(0, 2, 1, 3) * scale @@ -311,15 +309,11 @@ class ResidualAttentionBlock(nn.Module): self.attn = MultiHeadAttention(n_state, n_head) self.attn_ln = LayerNorm(n_state) - self.cross_attn = ( - MultiHeadAttention(n_state, n_head) if cross_attention else None - ) + self.cross_attn = MultiHeadAttention(n_state, n_head) if cross_attention else None self.cross_attn_ln = LayerNorm(n_state) if cross_attention else None n_mlp = n_state * 4 - self.mlp = nn.Sequential( - Linear(n_state, n_mlp), nn.GELU(), Linear(n_mlp, n_state) - ) + self.mlp = nn.Sequential(Linear(n_state, n_mlp), nn.GELU(), Linear(n_mlp, n_state)) self.mlp_ln = LayerNorm(n_state) def forward( @@ -335,19 +329,20 @@ class ResidualAttentionBlock(nn.Module): x = x + self.mlp(self.mlp_ln(x)) return x + @tables.register("encoder_classes", "QwenAudioEncoder") class QwenAudioEncoder(nn.Module): def __init__( - self, - n_mels: int, - n_ctx: int, - n_state: int, - n_head: int, - n_layer: int, - output_dim: int = 512, - avg_pool: bool = True, - add_audio_bos_eos_token: bool = True, - **kwargs + self, + n_mels: int, + n_ctx: int, + n_state: int, + n_head: int, + n_layer: int, + output_dim: int = 512, + avg_pool: bool = True, + add_audio_bos_eos_token: bool = True, + **kwargs, ): super().__init__() self.conv1 = Conv1d(n_mels, n_state, kernel_size=3, padding=1) @@ -371,15 +366,14 @@ class QwenAudioEncoder(nn.Module): self.output_dim = output_dim self.n_head = n_head - def forward(self, x: Tensor, padding_mask: Tensor=None, audio_lengths: Tensor=None): + def forward(self, x: Tensor, padding_mask: Tensor = None, audio_lengths: Tensor = None): """ x : torch.Tensor, shape = (batch_size, n_mels, n_ctx) the mel spectrogram of the audio """ - x = x.to(dtype=self.conv1.weight.dtype, - device=self.conv1.weight.device) + x = x.to(dtype=self.conv1.weight.dtype, device=self.conv1.weight.device) if audio_lengths is not None: - input_mel_len = audio_lengths[:,0] * 2 + input_mel_len = audio_lengths[:, 0] * 2 max_mel_len_in_batch = input_mel_len.max() x = x[:, :, :max_mel_len_in_batch] x = F.gelu(self.conv1(x)) @@ -388,35 +382,36 @@ class QwenAudioEncoder(nn.Module): bsz = x.size(0) src_len = x.size(1) - self.input_positional_embedding = self.positional_embedding[:src_len] - assert x.shape[1:] == self.input_positional_embedding.shape, f"incorrect audio shape: {x.shape[1:], self.input_positional_embedding.shape}" + assert ( + x.shape[1:] == self.input_positional_embedding.shape + ), f"incorrect audio shape: {x.shape[1:], self.input_positional_embedding.shape}" x = (x + self.input_positional_embedding).to(x.dtype) if padding_mask is not None: - padding_mask = padding_mask.to(dtype=self.conv1.weight.dtype, - device=self.conv1.weight.device) + padding_mask = padding_mask.to( + dtype=self.conv1.weight.dtype, device=self.conv1.weight.device + ) batch_src_len = padding_mask.size(1) x = x[:, :batch_src_len, :] - padding_mask = padding_mask.view( - bsz, -1, batch_src_len - ) + padding_mask = padding_mask.view(bsz, -1, batch_src_len) padding_mask_ = padding_mask.all(1) x[padding_mask_] = 0 - key_padding_mask = padding_mask_.view(bsz, 1, 1, batch_src_len). \ - expand(-1, self.n_head, -1, -1).reshape(bsz, self.n_head, 1, batch_src_len) + key_padding_mask = ( + padding_mask_.view(bsz, 1, 1, batch_src_len) + .expand(-1, self.n_head, -1, -1) + .reshape(bsz, self.n_head, 1, batch_src_len) + ) new_padding_mask = torch.zeros_like(key_padding_mask, dtype=x.dtype) padding_mask = new_padding_mask.masked_fill(key_padding_mask, float("-inf")) for block in self.blocks: x = block(x, mask=padding_mask) - if self.avg_pooler: x = x.permute(0, 2, 1) x = self.avg_pooler(x) x = x.permute(0, 2, 1) - x = self.ln_post(x) x = self.proj(x) @@ -430,15 +425,16 @@ class QwenAudioEncoder(nn.Module): def encode(self, input_audios: Tensor, input_audio_lengths: Tensor, audio_span_tokens: List): real_input_audio_lens = input_audio_lengths[:, 0].tolist() max_len_in_batch = max(real_input_audio_lens) - padding_mask = torch.ones([input_audios.size(0), max_len_in_batch]).to(dtype=self.conv1.weight.dtype, - device=self.conv1.weight.device) + padding_mask = torch.ones([input_audios.size(0), max_len_in_batch]).to( + dtype=self.conv1.weight.dtype, device=self.conv1.weight.device + ) for index in range(len(input_audios)): - padding_mask[index, :input_audio_lengths[index][0].item()] = 0 - x, bos, eos = self(input_audios, padding_mask,input_audio_lengths) + padding_mask[index, : input_audio_lengths[index][0].item()] = 0 + x, bos, eos = self(input_audios, padding_mask, input_audio_lengths) output_audios = [] for i in range(len(audio_span_tokens)): audio_span = audio_span_tokens[i] - audio = x[i][:audio_span-2] + audio = x[i][: audio_span - 2] if bos is not None: audio = torch.concat([bos, audio, eos]) assert len(audio) == audio_span diff --git a/funasr/models/qwen_audio/model.py b/funasr/models/qwen_audio/model.py index f981b67fa..b0af4563c 100644 --- a/funasr/models/qwen_audio/model.py +++ b/funasr/models/qwen_audio/model.py @@ -13,6 +13,7 @@ from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank from funasr.register import tables + @tables.register("model_classes", "Qwen/Qwen-Audio") @tables.register("model_classes", "Qwen-Audio") @tables.register("model_classes", "Qwen/QwenAudio") @@ -24,52 +25,60 @@ class QwenAudioWarp(nn.Module): https://arxiv.org/abs/2311.07919 Modified from https://github.com/QwenLM/Qwen-Audio """ + def __init__(self, *args, **kwargs): super().__init__() from transformers import AutoModelForCausalLM, AutoTokenizer from transformers.generation import GenerationConfig model_or_path = kwargs.get("model_path", "QwenAudio") - model = AutoModelForCausalLM.from_pretrained(model_or_path, device_map="cpu", - trust_remote_code=True) + model = AutoModelForCausalLM.from_pretrained( + model_or_path, device_map="cpu", trust_remote_code=True + ) tokenizer = AutoTokenizer.from_pretrained(model_or_path, trust_remote_code=True) - self.model = model self.tokenizer = tokenizer - - def forward(self, ): + + def forward( + self, + ): pass - def inference(self, - data_in, - data_lengths=None, - key: list = None, - tokenizer=None, - frontend=None, - **kwargs, - ): + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): if kwargs.get("batch_size", 1) > 1: raise NotImplementedError("batch decoding is not implemented") - meta_data = {} # meta_data["batch_data_time"] = -1 - prompt = kwargs.get("prompt", "<|startoftranscription|><|en|><|transcribe|><|en|><|notimestamps|><|wo_itn|>") + prompt = kwargs.get( + "prompt", "<|startoftranscription|><|en|><|transcribe|><|en|><|notimestamps|><|wo_itn|>" + ) query = f"{prompt}" audio_info = self.tokenizer.process_audio(query) - inputs = self.tokenizer(query, return_tensors='pt', audio_info=audio_info) + inputs = self.tokenizer(query, return_tensors="pt", audio_info=audio_info) inputs = inputs.to(self.model.device) pred = self.model.generate(**inputs, audio_info=audio_info) - response = self.tokenizer.decode(pred.cpu()[0], skip_special_tokens=False, audio_info=audio_info) + response = self.tokenizer.decode( + pred.cpu()[0], skip_special_tokens=False, audio_info=audio_info + ) results = [] result_i = {"key": key[0], "text": response} - + results.append(result_i) - + return results, meta_data + @tables.register("model_classes", "Qwen/Qwen-Audio-Chat") @tables.register("model_classes", "Qwen/QwenAudioChat") @tables.register("model_classes", "Qwen-Audio-Chat") @@ -85,34 +94,35 @@ class QwenAudioChatWarp(nn.Module): super().__init__() from transformers import AutoModelForCausalLM, AutoTokenizer from transformers.generation import GenerationConfig + model_or_path = kwargs.get("model_path", "QwenAudio") bf16 = kwargs.get("bf16", False) fp16 = kwargs.get("fp16", False) - model = AutoModelForCausalLM.from_pretrained(model_or_path, - device_map="cpu", - bf16=bf16, - fp16=fp16, - trust_remote_code=True) + model = AutoModelForCausalLM.from_pretrained( + model_or_path, device_map="cpu", bf16=bf16, fp16=fp16, trust_remote_code=True + ) tokenizer = AutoTokenizer.from_pretrained(model_or_path, trust_remote_code=True) - + self.model = model self.tokenizer = tokenizer - - def forward(self, ): + + def forward( + self, + ): pass - - def inference(self, - data_in, - data_lengths=None, - key: list = None, - tokenizer=None, - frontend=None, - **kwargs, - ): + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): if kwargs.get("batch_size", 1) > 1: raise NotImplementedError("batch decoding is not implemented") - - + meta_data = {} prompt = kwargs.get("prompt", "what does the person say?") @@ -120,10 +130,12 @@ class QwenAudioChatWarp(nn.Module): history = cache.get("history", None) if data_in[0] is not None: # 1st dialogue turn - query = self.tokenizer.from_list_format([ - {'audio': data_in[0]}, # Either a local path or an url - {'text': prompt}, - ]) + query = self.tokenizer.from_list_format( + [ + {"audio": data_in[0]}, # Either a local path or an url + {"text": prompt}, + ] + ) else: query = prompt response, history = self.model.chat(self.tokenizer, query=query, history=history) @@ -133,7 +145,7 @@ class QwenAudioChatWarp(nn.Module): results = [] result_i = {"key": key[0], "text": response} - + results.append(result_i) - + return results, meta_data diff --git a/funasr/models/rwkv_bat/rwkv.py b/funasr/models/rwkv_bat/rwkv.py index bd218a282..c777cf2b8 100644 --- a/funasr/models/rwkv_bat/rwkv.py +++ b/funasr/models/rwkv_bat/rwkv.py @@ -76,6 +76,7 @@ class RWKV(torch.nn.Module): x = x + self.dropout_ffn(ffn) return x, state + class RWKVDecoderLayer(torch.nn.Module): """RWKV module. diff --git a/funasr/models/rwkv_bat/rwkv_attention.py b/funasr/models/rwkv_bat/rwkv_attention.py index c085874e4..59bf0ffbd 100644 --- a/funasr/models/rwkv_bat/rwkv_attention.py +++ b/funasr/models/rwkv_bat/rwkv_attention.py @@ -13,6 +13,7 @@ from typing import List, Optional, Tuple, Union wkv_kernel_encoder = None wkv_kernel_decoder = None + class WKVLinearAttentionEncoder(torch.autograd.Function): """WKVLinearAttention function definition.""" @@ -44,8 +45,7 @@ class WKVLinearAttentionEncoder(torch.autograd.Function): ) assert batch * dim % min(dim, 32) == 0, ( - f"batch size ({batch}) by dimension ({dim}) should be a multiple of " - f"{min(dim, 32)}" + f"batch size ({batch}) by dimension ({dim}) should be a multiple of " f"{min(dim, 32)}" ) ctx.input_dtype = key.dtype @@ -124,6 +124,7 @@ class WKVLinearAttentionEncoder(torch.autograd.Function): grad_value, ) + class WKVLinearAttentionDecoder(torch.autograd.Function): """WKVLinearAttention function definition.""" @@ -155,8 +156,7 @@ class WKVLinearAttentionDecoder(torch.autograd.Function): ) assert batch * dim % min(dim, 32) == 0, ( - f"batch size ({batch}) by dimension ({dim}) should be a multiple of " - f"{min(dim, 32)}" + f"batch size ({batch}) by dimension ({dim}) should be a multiple of " f"{min(dim, 32)}" ) ctx.input_dtype = key.dtype @@ -235,6 +235,7 @@ class WKVLinearAttentionDecoder(torch.autograd.Function): grad_value, ) + def load_encoder_wkv_kernel(context_size: int) -> None: """Load WKV CUDA kernel. @@ -280,6 +281,7 @@ def load_encoder_wkv_kernel(context_size: int) -> None: ) wkv_kernel_encoder.context_size = context_size + def load_decoder_wkv_kernel(context_size: int) -> None: """Load WKV CUDA kernel. @@ -325,6 +327,7 @@ def load_decoder_wkv_kernel(context_size: int) -> None: ) wkv_kernel_decoder.context_size = context_size + class SelfAttention(torch.nn.Module): """SelfAttention module definition. @@ -406,17 +409,13 @@ class SelfAttention(torch.nn.Module): with torch.no_grad(): self.time_decay.data = decay_speed - self.time_first.data = torch.ones_like( - self.time_first * math.log(0.3) + zigzag - ) + self.time_first.data = torch.ones_like(self.time_first * math.log(0.3) + zigzag) self.time_mix_key.data = torch.pow(time_weight, ratio_1_to_almost0) self.time_mix_value.data = ( torch.pow(time_weight, ratio_1_to_almost0) + 0.3 * ratio_0_to_1 ) - self.time_mix_receptance.data = torch.pow( - time_weight, 0.5 * ratio_1_to_almost0 - ) + self.time_mix_receptance.data = torch.pow(time_weight, 0.5 * ratio_1_to_almost0) @torch.no_grad() def wkv_linear_attention( @@ -485,13 +484,7 @@ class DecoderSelfAttention(SelfAttention): num_blocks: int, ) -> None: """Construct a SelfAttention object.""" - super().__init__( - size, - attention_size, - block_id, - dropout_rate, - num_blocks - ) + super().__init__(size, attention_size, block_id, dropout_rate, num_blocks) # load_decoder_wkv_kernel(context_size) def forward( @@ -509,15 +502,11 @@ class DecoderSelfAttention(SelfAttention): x: SelfAttention output sequences. (B, U, size) """ - shifted_x = ( - self.time_shift(x) if state is None else state[1][..., self.block_id] - ) + shifted_x = self.time_shift(x) if state is None else state[1][..., self.block_id] key = x * self.time_mix_key + shifted_x * (1 - self.time_mix_key) value = x * self.time_mix_value + shifted_x * (1 - self.time_mix_value) - receptance = x * self.time_mix_receptance + shifted_x * ( - 1 - self.time_mix_receptance - ) + receptance = x * self.time_mix_receptance + shifted_x * (1 - self.time_mix_receptance) key = self.proj_key(key) value = self.proj_value(value) @@ -545,6 +534,7 @@ class DecoderSelfAttention(SelfAttention): return x, state + class EncoderSelfAttention(SelfAttention): """SelfAttention module definition. @@ -567,13 +557,7 @@ class EncoderSelfAttention(SelfAttention): num_blocks: int, ) -> None: """Construct a SelfAttention object.""" - super().__init__( - size, - attention_size, - block_id, - dropout_rate, - num_blocks - ) + super().__init__(size, attention_size, block_id, dropout_rate, num_blocks) # load_encoder_wkv_kernel(context_size) def forward( @@ -591,15 +575,11 @@ class EncoderSelfAttention(SelfAttention): x: SelfAttention output sequences. (B, U, size) """ - shifted_x = ( - self.time_shift(x) if state is None else state[1][..., self.block_id] - ) + shifted_x = self.time_shift(x) if state is None else state[1][..., self.block_id] key = x * self.time_mix_key + shifted_x * (1 - self.time_mix_key) value = x * self.time_mix_value + shifted_x * (1 - self.time_mix_value) - receptance = x * self.time_mix_receptance + shifted_x * ( - 1 - self.time_mix_receptance - ) + receptance = x * self.time_mix_receptance + shifted_x * (1 - self.time_mix_receptance) key = self.proj_key(key) value = self.proj_value(value) @@ -626,4 +606,3 @@ class EncoderSelfAttention(SelfAttention): x = self.proj_output(receptance * wkv) return x, state - diff --git a/funasr/models/rwkv_bat/rwkv_encoder.py b/funasr/models/rwkv_bat/rwkv_encoder.py index c0e5f4255..a27088d69 100644 --- a/funasr/models/rwkv_bat/rwkv_encoder.py +++ b/funasr/models/rwkv_bat/rwkv_encoder.py @@ -44,7 +44,7 @@ class RWKVEncoder(torch.nn.Module): att_dropout_rate: float = 0.0, ffn_dropout_rate: float = 0.0, dropout_rate: float = 0.0, - subsampling_factor: int =4, + subsampling_factor: int = 4, time_reduction_factor: int = 1, kernel: int = 3, **kwargs, @@ -54,7 +54,7 @@ class RWKVEncoder(torch.nn.Module): self.embed = RWKVConvInput( input_size, - [output_size//4, output_size//2, output_size], + [output_size // 4, output_size // 2, output_size], subsampling_factor, conv_kernel_size=kernel, output_size=output_size, @@ -64,7 +64,7 @@ class RWKVEncoder(torch.nn.Module): linear_size = output_size * 4 if linear_size is None else linear_size attention_size = output_size if attention_size is None else attention_size - + self.rwkv_blocks = torch.nn.ModuleList( [ RWKV( @@ -122,12 +122,12 @@ class RWKVEncoder(torch.nn.Module): x, _ = block(x) else: x = self.rwkv_infer(x) - + x = self.final_norm(x) if self.time_reduction_factor > 1: - x = x[:,::self.time_reduction_factor,:] - olens = torch.floor_divide(olens-1, self.time_reduction_factor) + 1 + x = x[:, :: self.time_reduction_factor, :] + olens = torch.floor_divide(olens - 1, self.time_reduction_factor) + 1 return x, olens, None @@ -135,9 +135,7 @@ class RWKVEncoder(torch.nn.Module): batch_size = xs_pad.shape[0] - hidden_sizes = [ - self._output_size for i in range(5) - ] + hidden_sizes = [self._output_size for i in range(5)] state = [ torch.zeros( @@ -152,7 +150,7 @@ class RWKVEncoder(torch.nn.Module): xs_out = [] for t in range(xs_pad.shape[1]): - x_t = xs_pad[:,t,:] + x_t = xs_pad[:, t, :] for idx, block in enumerate(self.rwkv_blocks): x_t, state = block(x_t, state=state) xs_out.append(x_t) diff --git a/funasr/models/rwkv_bat/rwkv_feed_forward.py b/funasr/models/rwkv_bat/rwkv_feed_forward.py index 32949ab02..a8056bf70 100644 --- a/funasr/models/rwkv_bat/rwkv_feed_forward.py +++ b/funasr/models/rwkv_bat/rwkv_feed_forward.py @@ -72,14 +72,10 @@ class FeedForward(torch.nn.Module): state: Decoder hidden state. [5 x (B, 1, size, N)] """ - shifted_x = ( - self.time_shift(x) if state is None else state[0][..., self.block_id] - ) + shifted_x = self.time_shift(x) if state is None else state[0][..., self.block_id] key = x * self.time_mix_key + shifted_x * (1 - self.time_mix_key) - receptance = x * self.time_mix_receptance + shifted_x * ( - 1 - self.time_mix_receptance - ) + receptance = x * self.time_mix_receptance + shifted_x * (1 - self.time_mix_receptance) key = torch.square(torch.relu(self.proj_key(key))) value = self.proj_value(self.dropout(key)) diff --git a/funasr/models/rwkv_bat/rwkv_subsampling.py b/funasr/models/rwkv_bat/rwkv_subsampling.py index a688acaff..5108ae68d 100644 --- a/funasr/models/rwkv_bat/rwkv_subsampling.py +++ b/funasr/models/rwkv_bat/rwkv_subsampling.py @@ -62,18 +62,50 @@ class RWKVConvInput(torch.nn.Module): conv_size1, conv_size2, conv_size3 = conv_size self.conv = torch.nn.Sequential( - torch.nn.Conv2d(1, conv_size1, conv_kernel_size, stride=1, padding=(conv_kernel_size-1)//2), - torch.nn.ReLU(), - torch.nn.Conv2d(conv_size1, conv_size1, conv_kernel_size, stride=[1, 2], padding=(conv_kernel_size-1)//2), - torch.nn.ReLU(), - torch.nn.Conv2d(conv_size1, conv_size2, conv_kernel_size, stride=1, padding=(conv_kernel_size-1)//2), - torch.nn.ReLU(), - torch.nn.Conv2d(conv_size2, conv_size2, conv_kernel_size, stride=[1, 2], padding=(conv_kernel_size-1)//2), - torch.nn.ReLU(), - torch.nn.Conv2d(conv_size2, conv_size3, conv_kernel_size, stride=1, padding=(conv_kernel_size-1)//2), - torch.nn.ReLU(), - torch.nn.Conv2d(conv_size3, conv_size3, conv_kernel_size, stride=[1, 2], padding=(conv_kernel_size-1)//2), - torch.nn.ReLU(), + torch.nn.Conv2d( + 1, conv_size1, conv_kernel_size, stride=1, padding=(conv_kernel_size - 1) // 2 + ), + torch.nn.ReLU(), + torch.nn.Conv2d( + conv_size1, + conv_size1, + conv_kernel_size, + stride=[1, 2], + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + torch.nn.Conv2d( + conv_size1, + conv_size2, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + torch.nn.Conv2d( + conv_size2, + conv_size2, + conv_kernel_size, + stride=[1, 2], + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + torch.nn.Conv2d( + conv_size2, + conv_size3, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + torch.nn.Conv2d( + conv_size3, + conv_size3, + conv_kernel_size, + stride=[1, 2], + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), ) output_proj = conv_size3 * ((input_size // 2) // 2) @@ -90,18 +122,50 @@ class RWKVConvInput(torch.nn.Module): kernel_1 = int(subsampling_factor / 2) self.conv = torch.nn.Sequential( - torch.nn.Conv2d(1, conv_size1, conv_kernel_size, stride=1, padding=(conv_kernel_size-1)//2), - torch.nn.ReLU(), - torch.nn.Conv2d(conv_size1, conv_size1, conv_kernel_size, stride=[kernel_1, 2], padding=(conv_kernel_size-1)//2), - torch.nn.ReLU(), - torch.nn.Conv2d(conv_size1, conv_size2, conv_kernel_size, stride=1, padding=(conv_kernel_size-1)//2), - torch.nn.ReLU(), - torch.nn.Conv2d(conv_size2, conv_size2, conv_kernel_size, stride=[2, 2], padding=(conv_kernel_size-1)//2), - torch.nn.ReLU(), - torch.nn.Conv2d(conv_size2, conv_size3, conv_kernel_size, stride=1, padding=(conv_kernel_size-1)//2), - torch.nn.ReLU(), - torch.nn.Conv2d(conv_size3, conv_size3, conv_kernel_size, stride=1, padding=(conv_kernel_size-1)//2), - torch.nn.ReLU(), + torch.nn.Conv2d( + 1, conv_size1, conv_kernel_size, stride=1, padding=(conv_kernel_size - 1) // 2 + ), + torch.nn.ReLU(), + torch.nn.Conv2d( + conv_size1, + conv_size1, + conv_kernel_size, + stride=[kernel_1, 2], + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + torch.nn.Conv2d( + conv_size1, + conv_size2, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + torch.nn.Conv2d( + conv_size2, + conv_size2, + conv_kernel_size, + stride=[2, 2], + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + torch.nn.Conv2d( + conv_size2, + conv_size3, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), + torch.nn.Conv2d( + conv_size3, + conv_size3, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), + torch.nn.ReLU(), ) output_proj = conv_size3 * ((input_size // 2) // 2) @@ -137,30 +201,32 @@ class RWKVConvInput(torch.nn.Module): olens = max(mask.eq(0).sum(1)) b, t, f = x.size() - x = x.unsqueeze(1) # (b. 1. t. f) + x = x.unsqueeze(1) # (b. 1. t. f) if chunk_size is not None: max_input_length = int( - chunk_size * self.subsampling_factor * (math.ceil(float(t) / (chunk_size * self.subsampling_factor) )) + chunk_size + * self.subsampling_factor + * (math.ceil(float(t) / (chunk_size * self.subsampling_factor))) ) x = map(lambda inputs: pad_to_len(inputs, max_input_length, 1), x) x = list(x) x = torch.stack(x, dim=0) - N_chunks = max_input_length // ( chunk_size * self.subsampling_factor) + N_chunks = max_input_length // (chunk_size * self.subsampling_factor) x = x.view(b * N_chunks, 1, chunk_size * self.subsampling_factor, f) x = self.conv(x) _, c, _, f = x.size() if chunk_size is not None: - x = x.transpose(1, 2).contiguous().view(b, -1, c * f)[:,:olens,:] + x = x.transpose(1, 2).contiguous().view(b, -1, c * f)[:, :olens, :] else: x = x.transpose(1, 2).contiguous().view(b, -1, c * f) if self.output is not None: x = self.output(x) - return x, mask[:,:olens][:,:x.size(1)] + return x, mask[:, :olens][:, : x.size(1)] def create_new_vgg_mask(self, mask: torch.Tensor) -> torch.Tensor: """Create a new mask for VGG output sequences. @@ -170,9 +236,9 @@ class RWKVConvInput(torch.nn.Module): mask: Mask of output sequences. (B, sub(T)) """ if self.subsampling_factor > 1: - return mask[:, ::2][:, ::self.stride_1] + return mask[:, ::2][:, :: self.stride_1] else: - return mask + return mask def get_size_before_subsampling(self, size: int) -> int: """Return the original size before subsampling for a given size. diff --git a/funasr/models/sa_asr/attention.py b/funasr/models/sa_asr/attention.py index 2cce9ecb0..74e687a52 100644 --- a/funasr/models/sa_asr/attention.py +++ b/funasr/models/sa_asr/attention.py @@ -16,9 +16,8 @@ from funasr.models.transformer.utils.nets_utils import make_pad_mask import funasr.models.lora.layers as lora - class CosineDistanceAttention(nn.Module): - """ Compute Cosine Distance between spk decoder output and speaker profile + """Compute Cosine Distance between spk decoder output and speaker profile Args: profile_path: speaker profile file path (.npy file) """ @@ -35,16 +34,16 @@ class CosineDistanceAttention(nn.Module): """ x = spk_decoder_out.unsqueeze(2) # (B, L, 1, D) if profile_lens is not None: - + mask = (make_pad_mask(profile_lens)[:, None, :]).to(profile.device) - min_value = float( - numpy.finfo(torch.tensor(0, dtype=x.dtype).numpy().dtype).min + min_value = float(numpy.finfo(torch.tensor(0, dtype=x.dtype).numpy().dtype).min) + weights_not_softmax = F.cosine_similarity(x, profile.unsqueeze(1), dim=-1).masked_fill( + mask, min_value ) - weights_not_softmax=F.cosine_similarity(x, profile.unsqueeze(1), dim=-1).masked_fill(mask, min_value) weights = self.softmax(weights_not_softmax).masked_fill(mask, 0.0) # (B, L, N) else: x = x[:, -1:, :, :] - weights_not_softmax=F.cosine_similarity(x, profile.unsqueeze(1).to(x.device), dim=-1) + weights_not_softmax = F.cosine_similarity(x, profile.unsqueeze(1).to(x.device), dim=-1) weights = self.softmax(weights_not_softmax) # (B, 1, N) spk_embedding = torch.matmul(weights, profile.to(weights.device)) # (B, L, D) diff --git a/funasr/models/sa_asr/beam_search_sa_asr.py b/funasr/models/sa_asr/beam_search_sa_asr.py index 6454bc0ec..d953b4c6a 100755 --- a/funasr/models/sa_asr/beam_search_sa_asr.py +++ b/funasr/models/sa_asr/beam_search_sa_asr.py @@ -21,7 +21,7 @@ class Hypothesis(NamedTuple): """Hypothesis data type.""" yseq: torch.Tensor - spk_weigths : List + spk_weigths: List score: Union[float, torch.Tensor] = 0 scores: Dict[str, Union[float, torch.Tensor]] = dict() states: Dict[str, Any] = dict() @@ -153,7 +153,11 @@ class BeamSearch(torch.nn.Module): return torch.cat((xs, x)) def score_full( - self, hyp: Hypothesis, asr_enc: torch.Tensor, spk_enc: torch.Tensor, profile: torch.Tensor, + self, + hyp: Hypothesis, + asr_enc: torch.Tensor, + spk_enc: torch.Tensor, + profile: torch.Tensor, ) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: """Score new hypothesis by `self.full_scorers`. @@ -173,13 +177,20 @@ class BeamSearch(torch.nn.Module): states = dict() for k, d in self.full_scorers.items(): if isinstance(d, AbsDecoder): - scores[k], spk_weigths, states[k] = d.score(hyp.yseq, hyp.states[k], asr_enc, spk_enc, profile) + scores[k], spk_weigths, states[k] = d.score( + hyp.yseq, hyp.states[k], asr_enc, spk_enc, profile + ) else: scores[k], states[k] = d.score(hyp.yseq, hyp.states[k], asr_enc) return scores, spk_weigths, states def score_partial( - self, hyp: Hypothesis, ids: torch.Tensor, asr_enc: torch.Tensor, spk_enc: torch.Tensor, profile: torch.Tensor, + self, + hyp: Hypothesis, + ids: torch.Tensor, + asr_enc: torch.Tensor, + spk_enc: torch.Tensor, + profile: torch.Tensor, ) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: """Score new hypothesis by `self.part_scorers`. @@ -200,7 +211,9 @@ class BeamSearch(torch.nn.Module): states = dict() for k, d in self.part_scorers.items(): if isinstance(d, AbsDecoder): - scores[k], states[k] = d.score_partial(hyp.yseq, ids, hyp.states[k], asr_enc, spk_enc, profile) + scores[k], states[k] = d.score_partial( + hyp.yseq, ids, hyp.states[k], asr_enc, spk_enc, profile + ) else: scores[k], states[k] = d.score_partial(hyp.yseq, ids, hyp.states[k], asr_enc) return scores, states @@ -288,7 +301,11 @@ class BeamSearch(torch.nn.Module): return new_states def search( - self, running_hyps: List[Hypothesis], asr_enc: torch.Tensor, spk_enc: torch.Tensor, profile: torch.Tensor + self, + running_hyps: List[Hypothesis], + asr_enc: torch.Tensor, + spk_enc: torch.Tensor, + profile: torch.Tensor, ) -> List[Hypothesis]: """Search new tokens for running hypotheses and encoded speech x. @@ -330,11 +347,9 @@ class BeamSearch(torch.nn.Module): Hypothesis( score=weighted_scores[j], yseq=self.append_token(hyp.yseq, j), - scores=self.merge_scores( - hyp.scores, scores, j, part_scores, part_j - ), + scores=self.merge_scores(hyp.scores, scores, j, part_scores, part_j), states=self.merge_states(states, part_states, part_j), - spk_weigths=hyp.spk_weigths+[spk_weigths], + spk_weigths=hyp.spk_weigths + [spk_weigths], ) ) @@ -345,7 +360,12 @@ class BeamSearch(torch.nn.Module): return best_hyps def forward( - self, asr_enc: torch.Tensor, spk_enc: torch.Tensor, profile: torch.Tensor, maxlenratio: float = 0.0, minlenratio: float = 0.0 + self, + asr_enc: torch.Tensor, + spk_enc: torch.Tensor, + profile: torch.Tensor, + maxlenratio: float = 0.0, + minlenratio: float = 0.0, ) -> List[Hypothesis]: """Perform beam search. @@ -377,7 +397,7 @@ class BeamSearch(torch.nn.Module): for i in range(maxlen): logging.debug("position " + str(i)) best = self.search(running_hyps, asr_enc, spk_enc, profile) - #import pdb;pdb.set_trace() + # import pdb;pdb.set_trace() # post process of one iteration running_hyps = self.post_process(i, maxlen, maxlenratio, best, ended_hyps) # end detection @@ -394,29 +414,26 @@ class BeamSearch(torch.nn.Module): # check the number of hypotheses reaching to eos if len(nbest_hyps) == 0: logging.warning( - "there is no N-best results, perform recognition " - "again with smaller minlenratio." + "there is no N-best results, perform recognition " "again with smaller minlenratio." ) return ( [] if minlenratio < 0.1 - else self.forward(asr_enc, spk_enc, profile, maxlenratio, max(0.0, minlenratio - 0.1)) + else self.forward( + asr_enc, spk_enc, profile, maxlenratio, max(0.0, minlenratio - 0.1) + ) ) # report the best result best = nbest_hyps[0] for k, v in best.scores.items(): - logging.info( - f"{v:6.2f} * {self.weights[k]:3} = {v * self.weights[k]:6.2f} for {k}" - ) + logging.info(f"{v:6.2f} * {self.weights[k]:3} = {v * self.weights[k]:6.2f} for {k}") logging.info(f"total log probability: {best.score:.2f}") logging.info(f"normalized log probability: {best.score / len(best.yseq):.2f}") logging.info(f"total number of ended hypotheses: {len(nbest_hyps)}") if self.token_list is not None: logging.info( - "best hypo: " - + "".join([self.token_list[x] for x in best.yseq[1:-1]]) - + "\n" + "best hypo: " + "".join([self.token_list[x] for x in best.yseq[1:-1]]) + "\n" ) return nbest_hyps @@ -444,15 +461,13 @@ class BeamSearch(torch.nn.Module): logging.debug(f"the number of running hypotheses: {len(running_hyps)}") if self.token_list is not None: logging.debug( - "best hypo: " - + "".join([self.token_list[x] for x in running_hyps[0].yseq[1:]]) + "best hypo: " + "".join([self.token_list[x] for x in running_hyps[0].yseq[1:]]) ) # add eos in the final loop to avoid that there are no ended hyps if i == maxlen - 1: logging.info("adding in the last position in the loop") running_hyps = [ - h._replace(yseq=self.append_token(h.yseq, self.eos)) - for h in running_hyps + h._replace(yseq=self.append_token(h.yseq, self.eos)) for h in running_hyps ] # add ended hypotheses to a final list, and removed them from current hypotheses diff --git a/funasr/models/sa_asr/e2e_sa_asr.py b/funasr/models/sa_asr/e2e_sa_asr.py index 3bfecfc88..f4827b257 100644 --- a/funasr/models/sa_asr/e2e_sa_asr.py +++ b/funasr/models/sa_asr/e2e_sa_asr.py @@ -14,9 +14,7 @@ import torch import torch.nn.functional as F from funasr.layers.abs_normalize import AbsNormalize -from funasr.losses.label_smoothing_loss import ( - LabelSmoothingLoss, NllLoss # noqa: H301 -) +from funasr.losses.label_smoothing_loss import LabelSmoothingLoss, NllLoss # noqa: H301 from funasr.models.ctc import CTC from funasr.models.decoder.abs_decoder import AbsDecoder from funasr.models.encoder.abs_encoder import AbsEncoder @@ -43,28 +41,28 @@ class SAASRModel(FunASRModel): """CTC-attention hybrid Encoder-Decoder model""" def __init__( - self, - vocab_size: int, - max_spk_num: int, - token_list: Union[Tuple[str, ...], List[str]], - frontend: Optional[AbsFrontend], - specaug: Optional[AbsSpecAug], - normalize: Optional[AbsNormalize], - asr_encoder: AbsEncoder, - spk_encoder: torch.nn.Module, - decoder: AbsDecoder, - ctc: CTC, - spk_weight: float = 0.5, - ctc_weight: float = 0.5, - interctc_weight: float = 0.0, - ignore_id: int = -1, - lsm_weight: float = 0.0, - length_normalized_loss: bool = False, - report_cer: bool = True, - report_wer: bool = True, - sym_space: str = "", - sym_blank: str = "", - extract_feats_in_collect_stats: bool = True, + self, + vocab_size: int, + max_spk_num: int, + token_list: Union[Tuple[str, ...], List[str]], + frontend: Optional[AbsFrontend], + specaug: Optional[AbsSpecAug], + normalize: Optional[AbsNormalize], + asr_encoder: AbsEncoder, + spk_encoder: torch.nn.Module, + decoder: AbsDecoder, + ctc: CTC, + spk_weight: float = 0.5, + ctc_weight: float = 0.5, + interctc_weight: float = 0.0, + ignore_id: int = -1, + lsm_weight: float = 0.0, + length_normalized_loss: bool = False, + report_cer: bool = True, + report_wer: bool = True, + sym_space: str = "", + sym_blank: str = "", + extract_feats_in_collect_stats: bool = True, ): assert 0.0 <= ctc_weight <= 1.0, ctc_weight assert 0.0 <= interctc_weight < 1.0, interctc_weight @@ -75,7 +73,7 @@ class SAASRModel(FunASRModel): self.sos = 1 self.eos = 2 self.vocab_size = vocab_size - self.max_spk_num=max_spk_num + self.max_spk_num = max_spk_num self.ignore_id = ignore_id self.spk_weight = spk_weight self.ctc_weight = ctc_weight @@ -97,7 +95,6 @@ class SAASRModel(FunASRModel): self.error_calculator = None - # we set self.decoder = None in the CTC mode since # self.decoder parameters were never used and PyTorch complained # and threw an Exception in the multi-GPU experiment. @@ -133,15 +130,15 @@ class SAASRModel(FunASRModel): self.extract_feats_in_collect_stats = extract_feats_in_collect_stats def forward( - self, - speech: torch.Tensor, - speech_lengths: torch.Tensor, - text: torch.Tensor, - text_lengths: torch.Tensor, - profile: torch.Tensor, - profile_lengths: torch.Tensor, - text_id: torch.Tensor, - text_id_lengths: torch.Tensor + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + text: torch.Tensor, + text_lengths: torch.Tensor, + profile: torch.Tensor, + profile_lengths: torch.Tensor, + text_id: torch.Tensor, + text_id_lengths: torch.Tensor, ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]: """Frontend + Encoder + Decoder + Calc loss @@ -156,10 +153,7 @@ class SAASRModel(FunASRModel): assert text_lengths.dim() == 1, text_lengths.shape # Check that batch_size is unified assert ( - speech.shape[0] - == speech_lengths.shape[0] - == text.shape[0] - == text_lengths.shape[0] + speech.shape[0] == speech_lengths.shape[0] == text.shape[0] == text_lengths.shape[0] ), (speech.shape, speech_lengths.shape, text.shape, text_lengths.shape) batch_size = speech.shape[0] @@ -183,7 +177,6 @@ class SAASRModel(FunASRModel): asr_encoder_out, encoder_out_lens, text, text_lengths ) - # Intermediate CTC (optional) loss_interctc = 0.0 if self.interctc_weight != 0.0 and intermediate_outs is not None: @@ -204,15 +197,20 @@ class SAASRModel(FunASRModel): loss_interctc = loss_interctc / len(intermediate_outs) # calculate whole encoder loss - loss_ctc = ( - 1 - self.interctc_weight - ) * loss_ctc + self.interctc_weight * loss_interctc - + loss_ctc = (1 - self.interctc_weight) * loss_ctc + self.interctc_weight * loss_interctc # 2b. Attention decoder branch if self.ctc_weight != 1.0: loss_att, loss_spk, acc_att, acc_spk, cer_att, wer_att = self._calc_att_loss( - asr_encoder_out, spk_encoder_out, encoder_out_lens, text, text_lengths, profile, profile_lengths, text_id, text_id_lengths + asr_encoder_out, + spk_encoder_out, + encoder_out_lens, + text, + text_lengths, + profile, + profile_lengths, + text_id, + text_id_lengths, ) # 3. CTC-Att loss definition @@ -228,7 +226,6 @@ class SAASRModel(FunASRModel): else: loss = self.spk_weight * loss_spk + (1 - self.spk_weight) * loss_asr - stats = dict( loss=loss.detach(), loss_asr=loss_asr.detach(), @@ -247,11 +244,11 @@ class SAASRModel(FunASRModel): return loss, stats, weight def collect_feats( - self, - speech: torch.Tensor, - speech_lengths: torch.Tensor, - text: torch.Tensor, - text_lengths: torch.Tensor, + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + text: torch.Tensor, + text_lengths: torch.Tensor, ) -> Dict[str, torch.Tensor]: if self.extract_feats_in_collect_stats: feats, feats_lengths = self._extract_feats(speech, speech_lengths) @@ -266,7 +263,7 @@ class SAASRModel(FunASRModel): return {"feats": feats, "feats_lengths": feats_lengths} def encode( - self, speech: torch.Tensor, speech_lengths: torch.Tensor + self, speech: torch.Tensor, speech_lengths: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: """Frontend + Encoder. Note that this method is used by asr_inference.py @@ -291,9 +288,7 @@ class SAASRModel(FunASRModel): # feats: (Batch, Length, Dim) # -> encoder_out: (Batch, Length2, Dim2) if self.asr_encoder.interctc_use_conditioning: - encoder_out, encoder_out_lens, _ = self.asr_encoder( - feats, feats_lengths, ctc=self.ctc - ) + encoder_out, encoder_out_lens, _ = self.asr_encoder(feats, feats_lengths, ctc=self.ctc) else: encoder_out, encoder_out_lens, _ = self.asr_encoder(feats, feats_lengths) intermediate_outs = None @@ -303,10 +298,12 @@ class SAASRModel(FunASRModel): encoder_out_spk_ori = self.spk_encoder(feats_raw, feats_lengths)[0] # import ipdb;ipdb.set_trace() - if encoder_out_spk_ori.size(1)!=encoder_out.size(1): - encoder_out_spk=F.interpolate(encoder_out_spk_ori.transpose(-2,-1), size=(encoder_out.size(1)), mode='nearest').transpose(-2,-1) + if encoder_out_spk_ori.size(1) != encoder_out.size(1): + encoder_out_spk = F.interpolate( + encoder_out_spk_ori.transpose(-2, -1), size=(encoder_out.size(1)), mode="nearest" + ).transpose(-2, -1) else: - encoder_out_spk=encoder_out_spk_ori + encoder_out_spk = encoder_out_spk_ori assert encoder_out.size(0) == speech.size(0), ( encoder_out.size(), @@ -327,7 +324,7 @@ class SAASRModel(FunASRModel): return encoder_out, encoder_out_lens, encoder_out_spk def _extract_feats( - self, speech: torch.Tensor, speech_lengths: torch.Tensor + self, speech: torch.Tensor, speech_lengths: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: assert speech_lengths.dim() == 1, speech_lengths.shape @@ -346,11 +343,11 @@ class SAASRModel(FunASRModel): return feats, feats_lengths def nll( - self, - encoder_out: torch.Tensor, - encoder_out_lens: torch.Tensor, - ys_pad: torch.Tensor, - ys_pad_lens: torch.Tensor, + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_pad_lens: torch.Tensor, ) -> torch.Tensor: """Compute negative log likelihood(nll) from transformer-decoder @@ -384,12 +381,12 @@ class SAASRModel(FunASRModel): return nll def batchify_nll( - self, - encoder_out: torch.Tensor, - encoder_out_lens: torch.Tensor, - ys_pad: torch.Tensor, - ys_pad_lens: torch.Tensor, - batch_size: int = 100, + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_pad_lens: torch.Tensor, + batch_size: int = 100, ): """Compute negative log likelihood(nll) from transformer-decoder @@ -431,28 +428,34 @@ class SAASRModel(FunASRModel): return nll def _calc_att_loss( - self, - asr_encoder_out: torch.Tensor, - spk_encoder_out: torch.Tensor, - encoder_out_lens: torch.Tensor, - ys_pad: torch.Tensor, - ys_pad_lens: torch.Tensor, - profile: torch.Tensor, - profile_lens: torch.Tensor, - text_id: torch.Tensor, - text_id_lengths: torch.Tensor + self, + asr_encoder_out: torch.Tensor, + spk_encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_pad_lens: torch.Tensor, + profile: torch.Tensor, + profile_lens: torch.Tensor, + text_id: torch.Tensor, + text_id_lengths: torch.Tensor, ): ys_in_pad, ys_out_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id) ys_in_lens = ys_pad_lens + 1 # 1. Forward decoder decoder_out, weights_no_pad, _ = self.decoder( - asr_encoder_out, spk_encoder_out, encoder_out_lens, ys_in_pad, ys_in_lens, profile, profile_lens + asr_encoder_out, + spk_encoder_out, + encoder_out_lens, + ys_in_pad, + ys_in_lens, + profile, + profile_lens, ) - spk_num_no_pad=weights_no_pad.size(-1) - pad=(0,self.max_spk_num-spk_num_no_pad) - weights=F.pad(weights_no_pad, pad, mode='constant', value=0) + spk_num_no_pad = weights_no_pad.size(-1) + pad = (0, self.max_spk_num - spk_num_no_pad) + weights = F.pad(weights_no_pad, pad, mode="constant", value=0) # pre_id=weights.argmax(-1) # pre_text=decoder_out.argmax(-1) @@ -467,7 +470,7 @@ class SAASRModel(FunASRModel): loss_att = self.criterion_att(decoder_out, ys_out_pad) loss_spk = self.criterion_spk(torch.log(weights), text_id) - acc_spk= th_accuracy( + acc_spk = th_accuracy( weights.view(-1, self.max_spk_num), text_id, ignore_label=self.ignore_id, @@ -488,11 +491,11 @@ class SAASRModel(FunASRModel): return loss_att, loss_spk, acc_att, acc_spk, cer_att, wer_att def _calc_ctc_loss( - self, - encoder_out: torch.Tensor, - encoder_out_lens: torch.Tensor, - ys_pad: torch.Tensor, - ys_pad_lens: torch.Tensor, + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_pad_lens: torch.Tensor, ): # Calc CTC loss loss_ctc = self.ctc(encoder_out, encoder_out_lens, ys_pad, ys_pad_lens) diff --git a/funasr/models/sa_asr/transformer_decoder.py b/funasr/models/sa_asr/transformer_decoder.py index b34a3aacf..be33fa527 100644 --- a/funasr/models/sa_asr/transformer_decoder.py +++ b/funasr/models/sa_asr/transformer_decoder.py @@ -29,6 +29,7 @@ from funasr.models.transformer.scorers.scorer_interface import BatchScorerInterf from funasr.register import tables + class DecoderLayer(nn.Module): """Single decoder layer module. @@ -52,14 +53,14 @@ class DecoderLayer(nn.Module): """ def __init__( - self, - size, - self_attn, - src_attn, - feed_forward, - dropout_rate, - normalize_before=True, - concat_after=False, + self, + size, + self_attn, + src_attn, + feed_forward, + dropout_rate, + normalize_before=True, + concat_after=False, ): """Construct an DecoderLayer object.""" super(DecoderLayer, self).__init__() @@ -116,9 +117,7 @@ class DecoderLayer(nn.Module): tgt_q_mask = tgt_mask[:, -1:, :] if self.concat_after: - tgt_concat = torch.cat( - (tgt_q, self.self_attn(tgt_q, tgt, tgt, tgt_q_mask)), dim=-1 - ) + tgt_concat = torch.cat((tgt_q, self.self_attn(tgt_q, tgt, tgt, tgt_q_mask)), dim=-1) x = residual + self.concat_linear1(tgt_concat) else: x = residual + self.dropout(self.self_attn(tgt_q, tgt, tgt, tgt_q_mask)) @@ -129,9 +128,7 @@ class DecoderLayer(nn.Module): if self.normalize_before: x = self.norm2(x) if self.concat_after: - x_concat = torch.cat( - (x, self.src_attn(x, memory, memory, memory_mask)), dim=-1 - ) + x_concat = torch.cat((x, self.src_attn(x, memory, memory, memory_mask)), dim=-1) x = residual + self.concat_linear2(x_concat) else: x = residual + self.dropout(self.src_attn(x, memory, memory, memory_mask)) @@ -174,15 +171,15 @@ class BaseTransformerDecoder(nn.Module, BatchScorerInterface): """ def __init__( - self, - vocab_size: int, - encoder_output_size: int, - dropout_rate: float = 0.1, - positional_dropout_rate: float = 0.1, - input_layer: str = "embed", - use_output_layer: bool = True, - pos_enc_class=PositionalEncoding, - normalize_before: bool = True, + self, + vocab_size: int, + encoder_output_size: int, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, ): super().__init__() attention_dim = encoder_output_size @@ -215,11 +212,11 @@ class BaseTransformerDecoder(nn.Module, BatchScorerInterface): self.decoders = None def forward( - self, - hs_pad: torch.Tensor, - hlens: torch.Tensor, - ys_in_pad: torch.Tensor, - ys_in_lens: torch.Tensor, + self, + hs_pad: torch.Tensor, + hlens: torch.Tensor, + ys_in_pad: torch.Tensor, + ys_in_lens: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: """Forward decoder. @@ -247,20 +244,14 @@ class BaseTransformerDecoder(nn.Module, BatchScorerInterface): tgt_mask = tgt_mask & m memory = hs_pad - memory_mask = (~make_pad_mask(hlens, maxlen=memory.size(1)))[:, None, :].to( - memory.device - ) + memory_mask = (~make_pad_mask(hlens, maxlen=memory.size(1)))[:, None, :].to(memory.device) # Padding for Longformer if memory_mask.shape[-1] != memory.shape[1]: padlen = memory.shape[1] - memory_mask.shape[-1] - memory_mask = torch.nn.functional.pad( - memory_mask, (0, padlen), "constant", False - ) + memory_mask = torch.nn.functional.pad(memory_mask, (0, padlen), "constant", False) x = self.embed(tgt) - x, tgt_mask, memory, memory_mask = self.decoders( - x, tgt_mask, memory, memory_mask - ) + x, tgt_mask, memory, memory_mask = self.decoders(x, tgt_mask, memory, memory_mask) if self.normalize_before: x = self.after_norm(x) if self.output_layer is not None: @@ -270,11 +261,11 @@ class BaseTransformerDecoder(nn.Module, BatchScorerInterface): return x, olens def forward_one_step( - self, - tgt: torch.Tensor, - tgt_mask: torch.Tensor, - memory: torch.Tensor, - cache: List[torch.Tensor] = None, + self, + tgt: torch.Tensor, + tgt_mask: torch.Tensor, + memory: torch.Tensor, + cache: List[torch.Tensor] = None, ) -> Tuple[torch.Tensor, List[torch.Tensor]]: """Forward one step. @@ -294,9 +285,7 @@ class BaseTransformerDecoder(nn.Module, BatchScorerInterface): cache = [None] * len(self.decoders) new_cache = [] for c, decoder in zip(cache, self.decoders): - x, tgt_mask, memory, memory_mask = decoder( - x, tgt_mask, memory, None, cache=c - ) + x, tgt_mask, memory, memory_mask = decoder(x, tgt_mask, memory, None, cache=c) new_cache.append(x) if self.normalize_before: @@ -311,13 +300,11 @@ class BaseTransformerDecoder(nn.Module, BatchScorerInterface): def score(self, ys, state, x): """Score.""" ys_mask = subsequent_mask(len(ys), device=x.device).unsqueeze(0) - logp, state = self.forward_one_step( - ys.unsqueeze(0), ys_mask, x.unsqueeze(0), cache=state - ) + logp, state = self.forward_one_step(ys.unsqueeze(0), ys_mask, x.unsqueeze(0), cache=state) return logp.squeeze(0), state def batch_score( - self, ys: torch.Tensor, states: List[Any], xs: torch.Tensor + self, ys: torch.Tensor, states: List[Any], xs: torch.Tensor ) -> Tuple[torch.Tensor, List[Any]]: """Score new token batch. @@ -341,8 +328,7 @@ class BaseTransformerDecoder(nn.Module, BatchScorerInterface): else: # transpose state of [batch, layer] into [layer, batch] batch_state = [ - torch.stack([states[b][i] for b in range(n_batch)]) - for i in range(n_layers) + torch.stack([states[b][i] for b in range(n_batch)]) for i in range(n_layers) ] # batch decoding @@ -353,24 +339,25 @@ class BaseTransformerDecoder(nn.Module, BatchScorerInterface): state_list = [[states[i][b] for i in range(n_layers)] for b in range(n_batch)] return logp, state_list + @tables.register("decoder_classes", "TransformerDecoder") class TransformerDecoder(BaseTransformerDecoder): def __init__( - self, - vocab_size: int, - encoder_output_size: int, - attention_heads: int = 4, - linear_units: int = 2048, - num_blocks: int = 6, - dropout_rate: float = 0.1, - positional_dropout_rate: float = 0.1, - self_attention_dropout_rate: float = 0.0, - src_attention_dropout_rate: float = 0.0, - input_layer: str = "embed", - use_output_layer: bool = True, - pos_enc_class=PositionalEncoding, - normalize_before: bool = True, - concat_after: bool = False, + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, ): super().__init__( vocab_size=vocab_size, @@ -388,12 +375,8 @@ class TransformerDecoder(BaseTransformerDecoder): num_blocks, lambda lnum: DecoderLayer( attention_dim, - MultiHeadedAttention( - attention_heads, attention_dim, self_attention_dropout_rate - ), - MultiHeadedAttention( - attention_heads, attention_dim, src_attention_dropout_rate - ), + MultiHeadedAttention(attention_heads, attention_dim, self_attention_dropout_rate), + MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate), PositionwiseFeedForward(attention_dim, linear_units, dropout_rate), dropout_rate, normalize_before, @@ -409,23 +392,24 @@ class ParaformerDecoderSAN(BaseTransformerDecoder): Paraformer: Fast and Accurate Parallel Transformer for Non-autoregressive End-to-End Speech Recognition https://arxiv.org/abs/2006.01713 """ + def __init__( - self, - vocab_size: int, - encoder_output_size: int, - attention_heads: int = 4, - linear_units: int = 2048, - num_blocks: int = 6, - dropout_rate: float = 0.1, - positional_dropout_rate: float = 0.1, - self_attention_dropout_rate: float = 0.0, - src_attention_dropout_rate: float = 0.0, - input_layer: str = "embed", - use_output_layer: bool = True, - pos_enc_class=PositionalEncoding, - normalize_before: bool = True, - concat_after: bool = False, - embeds_id: int = -1, + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + embeds_id: int = -1, ): super().__init__( vocab_size=vocab_size, @@ -443,12 +427,8 @@ class ParaformerDecoderSAN(BaseTransformerDecoder): num_blocks, lambda lnum: DecoderLayer( attention_dim, - MultiHeadedAttention( - attention_heads, attention_dim, self_attention_dropout_rate - ), - MultiHeadedAttention( - attention_heads, attention_dim, src_attention_dropout_rate - ), + MultiHeadedAttention(attention_heads, attention_dim, self_attention_dropout_rate), + MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate), PositionwiseFeedForward(attention_dim, linear_units, dropout_rate), dropout_rate, normalize_before, @@ -459,11 +439,11 @@ class ParaformerDecoderSAN(BaseTransformerDecoder): self.attention_dim = attention_dim def forward( - self, - hs_pad: torch.Tensor, - hlens: torch.Tensor, - ys_in_pad: torch.Tensor, - ys_in_lens: torch.Tensor, + self, + hs_pad: torch.Tensor, + hlens: torch.Tensor, + ys_in_pad: torch.Tensor, + ys_in_lens: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: """Forward decoder. @@ -486,23 +466,17 @@ class ParaformerDecoderSAN(BaseTransformerDecoder): tgt_mask = (~make_pad_mask(ys_in_lens)[:, None, :]).to(tgt.device) memory = hs_pad - memory_mask = (~make_pad_mask(hlens, maxlen=memory.size(1)))[:, None, :].to( - memory.device - ) + memory_mask = (~make_pad_mask(hlens, maxlen=memory.size(1)))[:, None, :].to(memory.device) # Padding for Longformer if memory_mask.shape[-1] != memory.shape[1]: padlen = memory.shape[1] - memory_mask.shape[-1] - memory_mask = torch.nn.functional.pad( - memory_mask, (0, padlen), "constant", False - ) + memory_mask = torch.nn.functional.pad(memory_mask, (0, padlen), "constant", False) # x = self.embed(tgt) x = tgt embeds_outputs = None for layer_id, decoder in enumerate(self.decoders): - x, tgt_mask, memory, memory_mask = decoder( - x, tgt_mask, memory, memory_mask - ) + x, tgt_mask, memory, memory_mask = decoder(x, tgt_mask, memory, memory_mask) if layer_id == self.embeds_id: embeds_outputs = x if self.normalize_before: @@ -516,27 +490,28 @@ class ParaformerDecoderSAN(BaseTransformerDecoder): else: return x, olens + @tables.register("decoder_classes", "LightweightConvolutionTransformerDecoder") class LightweightConvolutionTransformerDecoder(BaseTransformerDecoder): def __init__( - self, - vocab_size: int, - encoder_output_size: int, - attention_heads: int = 4, - linear_units: int = 2048, - num_blocks: int = 6, - dropout_rate: float = 0.1, - positional_dropout_rate: float = 0.1, - self_attention_dropout_rate: float = 0.0, - src_attention_dropout_rate: float = 0.0, - input_layer: str = "embed", - use_output_layer: bool = True, - pos_enc_class=PositionalEncoding, - normalize_before: bool = True, - concat_after: bool = False, - conv_wshare: int = 4, - conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11), - conv_usebias: int = False, + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + conv_wshare: int = 4, + conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11), + conv_usebias: int = False, ): if len(conv_kernel_length) != num_blocks: raise ValueError( @@ -567,9 +542,7 @@ class LightweightConvolutionTransformerDecoder(BaseTransformerDecoder): use_kernel_mask=True, use_bias=conv_usebias, ), - MultiHeadedAttention( - attention_heads, attention_dim, src_attention_dropout_rate - ), + MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate), PositionwiseFeedForward(attention_dim, linear_units, dropout_rate), dropout_rate, normalize_before, @@ -577,27 +550,28 @@ class LightweightConvolutionTransformerDecoder(BaseTransformerDecoder): ), ) + @tables.register("decoder_classes", "LightweightConvolution2DTransformerDecoder") class LightweightConvolution2DTransformerDecoder(BaseTransformerDecoder): def __init__( - self, - vocab_size: int, - encoder_output_size: int, - attention_heads: int = 4, - linear_units: int = 2048, - num_blocks: int = 6, - dropout_rate: float = 0.1, - positional_dropout_rate: float = 0.1, - self_attention_dropout_rate: float = 0.0, - src_attention_dropout_rate: float = 0.0, - input_layer: str = "embed", - use_output_layer: bool = True, - pos_enc_class=PositionalEncoding, - normalize_before: bool = True, - concat_after: bool = False, - conv_wshare: int = 4, - conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11), - conv_usebias: int = False, + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + conv_wshare: int = 4, + conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11), + conv_usebias: int = False, ): if len(conv_kernel_length) != num_blocks: raise ValueError( @@ -628,9 +602,7 @@ class LightweightConvolution2DTransformerDecoder(BaseTransformerDecoder): use_kernel_mask=True, use_bias=conv_usebias, ), - MultiHeadedAttention( - attention_heads, attention_dim, src_attention_dropout_rate - ), + MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate), PositionwiseFeedForward(attention_dim, linear_units, dropout_rate), dropout_rate, normalize_before, @@ -642,24 +614,24 @@ class LightweightConvolution2DTransformerDecoder(BaseTransformerDecoder): @tables.register("decoder_classes", "DynamicConvolutionTransformerDecoder") class DynamicConvolutionTransformerDecoder(BaseTransformerDecoder): def __init__( - self, - vocab_size: int, - encoder_output_size: int, - attention_heads: int = 4, - linear_units: int = 2048, - num_blocks: int = 6, - dropout_rate: float = 0.1, - positional_dropout_rate: float = 0.1, - self_attention_dropout_rate: float = 0.0, - src_attention_dropout_rate: float = 0.0, - input_layer: str = "embed", - use_output_layer: bool = True, - pos_enc_class=PositionalEncoding, - normalize_before: bool = True, - concat_after: bool = False, - conv_wshare: int = 4, - conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11), - conv_usebias: int = False, + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + conv_wshare: int = 4, + conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11), + conv_usebias: int = False, ): if len(conv_kernel_length) != num_blocks: raise ValueError( @@ -690,9 +662,7 @@ class DynamicConvolutionTransformerDecoder(BaseTransformerDecoder): use_kernel_mask=True, use_bias=conv_usebias, ), - MultiHeadedAttention( - attention_heads, attention_dim, src_attention_dropout_rate - ), + MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate), PositionwiseFeedForward(attention_dim, linear_units, dropout_rate), dropout_rate, normalize_before, @@ -700,27 +670,28 @@ class DynamicConvolutionTransformerDecoder(BaseTransformerDecoder): ), ) + @tables.register("decoder_classes", "DynamicConvolution2DTransformerDecoder") class DynamicConvolution2DTransformerDecoder(BaseTransformerDecoder): def __init__( - self, - vocab_size: int, - encoder_output_size: int, - attention_heads: int = 4, - linear_units: int = 2048, - num_blocks: int = 6, - dropout_rate: float = 0.1, - positional_dropout_rate: float = 0.1, - self_attention_dropout_rate: float = 0.0, - src_attention_dropout_rate: float = 0.0, - input_layer: str = "embed", - use_output_layer: bool = True, - pos_enc_class=PositionalEncoding, - normalize_before: bool = True, - concat_after: bool = False, - conv_wshare: int = 4, - conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11), - conv_usebias: int = False, + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + conv_wshare: int = 4, + conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11), + conv_usebias: int = False, ): if len(conv_kernel_length) != num_blocks: raise ValueError( @@ -751,9 +722,7 @@ class DynamicConvolution2DTransformerDecoder(BaseTransformerDecoder): use_kernel_mask=True, use_bias=conv_usebias, ), - MultiHeadedAttention( - attention_heads, attention_dim, src_attention_dropout_rate - ), + MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate), PositionwiseFeedForward(attention_dim, linear_units, dropout_rate), dropout_rate, normalize_before, diff --git a/funasr/models/sanm/attention.py b/funasr/models/sanm/attention.py index 1768bbd82..da8850fe5 100644 --- a/funasr/models/sanm/attention.py +++ b/funasr/models/sanm/attention.py @@ -25,14 +25,15 @@ def preprocess_for_attn(x, mask, cache, pad_fn, kernel_size): x = pad_fn(x) else: x = torch.cat((cache, x), dim=2) - cache = x[:, :, -(kernel_size-1):] + cache = x[:, :, -(kernel_size - 1) :] return x, cache torch_version = tuple([int(i) for i in torch.__version__.split(".")[:2]]) if torch_version >= (1, 8): import torch.fx - torch.fx.wrap('preprocess_for_attn') + + torch.fx.wrap("preprocess_for_attn") class MultiHeadedAttention(nn.Module): @@ -99,9 +100,7 @@ class MultiHeadedAttention(nn.Module): n_batch = value.size(0) if mask is not None: mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) - min_value = float( - numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min - ) + min_value = float(numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min) scores = scores.masked_fill(mask, min_value) self.attn = torch.softmax(scores, dim=-1).masked_fill( mask, 0.0 @@ -136,10 +135,6 @@ class MultiHeadedAttention(nn.Module): return self.forward_attention(v, scores, mask) - - - - class MultiHeadedAttentionSANM(nn.Module): """Multi-Head Attention layer. @@ -150,7 +145,19 @@ class MultiHeadedAttentionSANM(nn.Module): """ - def __init__(self, n_head, in_feat, n_feat, dropout_rate, kernel_size, sanm_shfit=0, lora_list=None, lora_rank=8, lora_alpha=16, lora_dropout=0.1): + def __init__( + self, + n_head, + in_feat, + n_feat, + dropout_rate, + kernel_size, + sanm_shfit=0, + lora_list=None, + lora_rank=8, + lora_alpha=16, + lora_dropout=0.1, + ): """Construct an MultiHeadedAttention object.""" super().__init__() assert n_feat % n_head == 0 @@ -162,21 +169,32 @@ class MultiHeadedAttentionSANM(nn.Module): # self.linear_v = nn.Linear(n_feat, n_feat) if lora_list is not None: if "o" in lora_list: - self.linear_out = lora.Linear(n_feat, n_feat, r=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout) + self.linear_out = lora.Linear( + n_feat, n_feat, r=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout + ) else: self.linear_out = nn.Linear(n_feat, n_feat) lora_qkv_list = ["q" in lora_list, "k" in lora_list, "v" in lora_list] if lora_qkv_list == [False, False, False]: self.linear_q_k_v = nn.Linear(in_feat, n_feat * 3) else: - self.linear_q_k_v = lora.MergedLinear(in_feat, n_feat * 3, r=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout, enable_lora=lora_qkv_list) + self.linear_q_k_v = lora.MergedLinear( + in_feat, + n_feat * 3, + r=lora_rank, + lora_alpha=lora_alpha, + lora_dropout=lora_dropout, + enable_lora=lora_qkv_list, + ) else: self.linear_out = nn.Linear(n_feat, n_feat) self.linear_q_k_v = nn.Linear(in_feat, n_feat * 3) self.attn = None self.dropout = nn.Dropout(p=dropout_rate) - self.fsmn_block = nn.Conv1d(n_feat, n_feat, kernel_size, stride=1, padding=0, groups=n_feat, bias=False) + self.fsmn_block = nn.Conv1d( + n_feat, n_feat, kernel_size, stride=1, padding=0, groups=n_feat, bias=False + ) # padding left_padding = (kernel_size - 1) // 2 if sanm_shfit > 0: @@ -219,9 +237,15 @@ class MultiHeadedAttentionSANM(nn.Module): b, t, d = x.size() q_k_v = self.linear_q_k_v(x) q, k, v = torch.split(q_k_v, int(self.h * self.d_k), dim=-1) - q_h = torch.reshape(q, (b, t, self.h, self.d_k)).transpose(1, 2) # (batch, head, time1, d_k) - k_h = torch.reshape(k, (b, t, self.h, self.d_k)).transpose(1, 2) # (batch, head, time2, d_k) - v_h = torch.reshape(v, (b, t, self.h, self.d_k)).transpose(1, 2) # (batch, head, time2, d_k) + q_h = torch.reshape(q, (b, t, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time1, d_k) + k_h = torch.reshape(k, (b, t, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time2, d_k) + v_h = torch.reshape(v, (b, t, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time2, d_k) return q_h, k_h, v_h, v @@ -245,9 +269,7 @@ class MultiHeadedAttentionSANM(nn.Module): mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) - min_value = float( - numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min - ) + min_value = float(numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min) scores = scores.masked_fill(mask, min_value) self.attn = torch.softmax(scores, dim=-1).masked_fill( mask, 0.0 @@ -301,19 +323,21 @@ class MultiHeadedAttentionSANM(nn.Module): q_h, k_h, v_h, v = self.forward_qkv(x) if chunk_size is not None and look_back > 0 or look_back == -1: if cache is not None: - k_h_stride = k_h[:, :, :-(chunk_size[2]), :] - v_h_stride = v_h[:, :, :-(chunk_size[2]), :] + k_h_stride = k_h[:, :, : -(chunk_size[2]), :] + v_h_stride = v_h[:, :, : -(chunk_size[2]), :] k_h = torch.cat((cache["k"], k_h), dim=2) v_h = torch.cat((cache["v"], v_h), dim=2) cache["k"] = torch.cat((cache["k"], k_h_stride), dim=2) cache["v"] = torch.cat((cache["v"], v_h_stride), dim=2) if look_back != -1: - cache["k"] = cache["k"][:, :, -(look_back * chunk_size[1]):, :] - cache["v"] = cache["v"][:, :, -(look_back * chunk_size[1]):, :] + cache["k"] = cache["k"][:, :, -(look_back * chunk_size[1]) :, :] + cache["v"] = cache["v"][:, :, -(look_back * chunk_size[1]) :, :] else: - cache_tmp = {"k": k_h[:, :, :-(chunk_size[2]), :], - "v": v_h[:, :, :-(chunk_size[2]), :]} + cache_tmp = { + "k": k_h[:, :, : -(chunk_size[2]), :], + "v": v_h[:, :, : -(chunk_size[2]), :], + } cache = cache_tmp fsmn_memory = self.forward_fsmn(v, None) q_h = q_h * self.d_k ** (-0.5) @@ -321,6 +345,7 @@ class MultiHeadedAttentionSANM(nn.Module): att_outs = self.forward_attention(v_h, scores, None) return att_outs + fsmn_memory, cache + class MultiHeadedAttentionSANMExport(nn.Module): def __init__(self, model): super().__init__() @@ -338,7 +363,7 @@ class MultiHeadedAttentionSANMExport(nn.Module): mask_3d_btd, mask_4d_bhlt = mask q_h, k_h, v_h, v = self.forward_qkv(x) fsmn_memory = self.forward_fsmn(v, mask_3d_btd) - q_h = q_h * self.d_k**(-0.5) + q_h = q_h * self.d_k ** (-0.5) scores = torch.matmul(q_h, k_h.transpose(-2, -1)) att_outs = self.forward_attention(v_h, scores, mask_4d_bhlt) return att_outs + fsmn_memory @@ -397,7 +422,7 @@ class MultiHeadedAttentionSANMExport(nn.Module): mask_3d_btd, mask_4d_bhlt = mask q_h, k_h, v_h, v = self.forward_qkv(x) fsmn_memory = self.forward_fsmn(v, mask_3d_btd) - q_h = q_h * self.d_k**(-0.5) + q_h = q_h * self.d_k ** (-0.5) scores = torch.matmul(q_h, k_h.transpose(-2, -1)) att_outs = self.forward_attention(v_h, scores, mask_4d_bhlt) return att_outs + fsmn_memory @@ -439,7 +464,6 @@ class MultiHeadedAttentionSANMExport(nn.Module): return self.linear_out(context_layer) # (batch, time1, d_model) - class MultiHeadedAttentionSANMDecoder(nn.Module): """Multi-Head Attention layer. @@ -456,8 +480,9 @@ class MultiHeadedAttentionSANMDecoder(nn.Module): self.dropout = nn.Dropout(p=dropout_rate) - self.fsmn_block = nn.Conv1d(n_feat, n_feat, - kernel_size, stride=1, padding=0, groups=n_feat, bias=False) + self.fsmn_block = nn.Conv1d( + n_feat, n_feat, kernel_size, stride=1, padding=0, groups=n_feat, bias=False + ) # padding # padding left_padding = (kernel_size - 1) // 2 @@ -468,17 +493,17 @@ class MultiHeadedAttentionSANMDecoder(nn.Module): self.kernel_size = kernel_size def forward(self, inputs, mask, cache=None, mask_shfit_chunk=None): - ''' + """ :param x: (#batch, time1, size). :param mask: Mask tensor (#batch, 1, time) :return: - ''' + """ # print("in fsmn, inputs", inputs.size()) b, t, d = inputs.size() # logging.info( # "mask: {}".format(mask.size())) if mask is not None: - mask = torch.reshape(mask, (b ,-1, 1)) + mask = torch.reshape(mask, (b, -1, 1)) # logging.info("in fsmn, mask: {}, {}".format(mask.size(), mask[0:100:50, :, :])) if mask_shfit_chunk is not None: # logging.info("in fsmn, mask_fsmn: {}, {}".format(mask_shfit_chunk.size(), mask_shfit_chunk[0:100:50, :, :])) @@ -502,7 +527,7 @@ class MultiHeadedAttentionSANMDecoder(nn.Module): # if t < self.kernel_size: # x = self.pad_fn(x) x = torch.cat((cache[:, :, 1:], x), dim=2) - x = x[:, :, -(self.kernel_size+t-1):] + x = x[:, :, -(self.kernel_size + t - 1) :] # print("in fsmn, cache is not None, x_cat", x.size()) cache = x x = self.fsmn_block(x) @@ -517,6 +542,7 @@ class MultiHeadedAttentionSANMDecoder(nn.Module): x = x * mask return x, cache + class MultiHeadedAttentionSANMDecoderExport(nn.Module): def __init__(self, model): super().__init__() @@ -545,7 +571,17 @@ class MultiHeadedAttentionCrossAtt(nn.Module): """ - def __init__(self, n_head, n_feat, dropout_rate, lora_list=None, lora_rank=8, lora_alpha=16, lora_dropout=0.1, encoder_output_size=None): + def __init__( + self, + n_head, + n_feat, + dropout_rate, + lora_list=None, + lora_rank=8, + lora_alpha=16, + lora_dropout=0.1, + encoder_output_size=None, + ): """Construct an MultiHeadedAttention object.""" super().__init__() assert n_feat % n_head == 0 @@ -554,22 +590,36 @@ class MultiHeadedAttentionCrossAtt(nn.Module): self.h = n_head if lora_list is not None: if "q" in lora_list: - self.linear_q = lora.Linear(n_feat, n_feat, r=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout) + self.linear_q = lora.Linear( + n_feat, n_feat, r=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout + ) else: self.linear_q = nn.Linear(n_feat, n_feat) lora_kv_list = ["k" in lora_list, "v" in lora_list] if lora_kv_list == [False, False]: - self.linear_k_v = nn.Linear(n_feat if encoder_output_size is None else encoder_output_size, n_feat*2) + self.linear_k_v = nn.Linear( + n_feat if encoder_output_size is None else encoder_output_size, n_feat * 2 + ) else: - self.linear_k_v = lora.MergedLinear(n_feat if encoder_output_size is None else encoder_output_size, n_feat * 2, - r=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout, enable_lora=lora_kv_list) + self.linear_k_v = lora.MergedLinear( + n_feat if encoder_output_size is None else encoder_output_size, + n_feat * 2, + r=lora_rank, + lora_alpha=lora_alpha, + lora_dropout=lora_dropout, + enable_lora=lora_kv_list, + ) if "o" in lora_list: - self.linear_out = lora.Linear(n_feat, n_feat, r=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout) + self.linear_out = lora.Linear( + n_feat, n_feat, r=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout + ) else: self.linear_out = nn.Linear(n_feat, n_feat) else: self.linear_q = nn.Linear(n_feat, n_feat) - self.linear_k_v = nn.Linear(n_feat if encoder_output_size is None else encoder_output_size, n_feat*2) + self.linear_k_v = nn.Linear( + n_feat if encoder_output_size is None else encoder_output_size, n_feat * 2 + ) self.linear_out = nn.Linear(n_feat, n_feat) self.attn = None self.dropout = nn.Dropout(p=dropout_rate) @@ -592,13 +642,18 @@ class MultiHeadedAttentionCrossAtt(nn.Module): # print("in forward_qkv, x", x.size()) b = x.size(0) q = self.linear_q(x) - q_h = torch.reshape(q, (b, -1, self.h, self.d_k)).transpose(1, 2) # (batch, head, time1, d_k) + q_h = torch.reshape(q, (b, -1, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time1, d_k) k_v = self.linear_k_v(memory) - k, v = torch.split(k_v, int(self.h*self.d_k), dim=-1) - k_h = torch.reshape(k, (b, -1, self.h, self.d_k)).transpose(1, 2) # (batch, head, time2, d_k) - v_h = torch.reshape(v, (b, -1, self.h, self.d_k)).transpose(1, 2) # (batch, head, time2, d_k) - + k, v = torch.split(k_v, int(self.h * self.d_k), dim=-1) + k_h = torch.reshape(k, (b, -1, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time2, d_k) + v_h = torch.reshape(v, (b, -1, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time2, d_k) return q_h, k_h, v_h @@ -618,9 +673,7 @@ class MultiHeadedAttentionCrossAtt(nn.Module): n_batch = value.size(0) if mask is not None: mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) - min_value = float( - numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min - ) + min_value = float(numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min) # logging.info( # "scores: {}, mask_size: {}".format(scores.size(), mask.size())) scores = scores.masked_fill(mask, min_value) @@ -676,16 +729,19 @@ class MultiHeadedAttentionCrossAtt(nn.Module): if cache is not None: k_h = torch.cat((cache["k"], k_h), dim=2) v_h = torch.cat((cache["v"], v_h), dim=2) - cache["k"] = k_h[:, :, -(look_back * chunk_size[1]):, :] - cache["v"] = v_h[:, :, -(look_back * chunk_size[1]):, :] + cache["k"] = k_h[:, :, -(look_back * chunk_size[1]) :, :] + cache["v"] = v_h[:, :, -(look_back * chunk_size[1]) :, :] else: - cache_tmp = {"k": k_h[:, :, -(look_back * chunk_size[1]):, :], - "v": v_h[:, :, -(look_back * chunk_size[1]):, :]} + cache_tmp = { + "k": k_h[:, :, -(look_back * chunk_size[1]) :, :], + "v": v_h[:, :, -(look_back * chunk_size[1]) :, :], + } cache = cache_tmp q_h = q_h * self.d_k ** (-0.5) scores = torch.matmul(q_h, k_h.transpose(-2, -1)) return self.forward_attention(v_h, scores, None), cache + class MultiHeadedAttentionCrossAttExport(nn.Module): def __init__(self, model): super().__init__() @@ -726,7 +782,8 @@ class MultiHeadedAttentionCrossAttExport(nn.Module): context_layer = context_layer.permute(0, 2, 1, 3).contiguous() new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) context_layer = context_layer.view(new_context_layer_shape) - if ret_attn: return self.linear_out(context_layer), self.attn + if ret_attn: + return self.linear_out(context_layer), self.attn return self.linear_out(context_layer) # (batch, time1, d_model) @@ -769,9 +826,15 @@ class MultiHeadSelfAttention(nn.Module): b, t, d = x.size() q_k_v = self.linear_q_k_v(x) q, k, v = torch.split(q_k_v, int(self.h * self.d_k), dim=-1) - q_h = torch.reshape(q, (b, t, self.h, self.d_k)).transpose(1, 2) # (batch, head, time1, d_k) - k_h = torch.reshape(k, (b, t, self.h, self.d_k)).transpose(1, 2) # (batch, head, time2, d_k) - v_h = torch.reshape(v, (b, t, self.h, self.d_k)).transpose(1, 2) # (batch, head, time2, d_k) + q_h = torch.reshape(q, (b, t, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time1, d_k) + k_h = torch.reshape(k, (b, t, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time2, d_k) + v_h = torch.reshape(v, (b, t, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time2, d_k) return q_h, k_h, v_h, v @@ -795,9 +858,7 @@ class MultiHeadSelfAttention(nn.Module): mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) - min_value = float( - numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min - ) + min_value = float(numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min) scores = scores.masked_fill(mask, min_value) self.attn = torch.softmax(scores, dim=-1).masked_fill( mask, 0.0 diff --git a/funasr/models/sanm/decoder.py b/funasr/models/sanm/decoder.py index 35752829f..1a4fb26e2 100644 --- a/funasr/models/sanm/decoder.py +++ b/funasr/models/sanm/decoder.py @@ -13,7 +13,10 @@ import numpy as np from funasr.models.scama import utils as myutils from funasr.models.transformer.decoder import BaseTransformerDecoder -from funasr.models.sanm.attention import MultiHeadedAttentionSANMDecoder, MultiHeadedAttentionCrossAtt +from funasr.models.sanm.attention import ( + MultiHeadedAttentionSANMDecoder, + MultiHeadedAttentionCrossAtt, +) from funasr.models.transformer.embedding import PositionalEncoding from funasr.models.transformer.layer_norm import LayerNorm from funasr.models.sanm.positionwise_feed_forward import PositionwiseFeedForwardDecoderSANM @@ -21,6 +24,7 @@ from funasr.models.transformer.utils.repeat import repeat from funasr.register import tables + class DecoderLayerSANM(nn.Module): """Single decoder layer module. @@ -151,10 +155,11 @@ class DecoderLayerSANM(nn.Module): x = residual + self.dropout(self.src_attn(x, memory, memory_mask)) - return x, tgt_mask, memory, memory_mask, cache - def forward_chunk(self, tgt, memory, fsmn_cache=None, opt_cache=None, chunk_size=None, look_back=0): + def forward_chunk( + self, tgt, memory, fsmn_cache=None, opt_cache=None, chunk_size=None, look_back=0 + ): """Compute decoded features. Args: @@ -202,7 +207,7 @@ class FsmnDecoder(BaseTransformerDecoder): San-m: Memory equipped self-attention for end-to-end speech recognition https://arxiv.org/abs/2006.01713 """ - + def __init__( self, vocab_size: int, @@ -240,7 +245,7 @@ class FsmnDecoder(BaseTransformerDecoder): ) if attention_dim is None: attention_dim = encoder_output_size - + if input_layer == "embed": self.embed = torch.nn.Sequential( torch.nn.Embedding(vocab_size, attention_dim), @@ -255,7 +260,7 @@ class FsmnDecoder(BaseTransformerDecoder): ) else: raise ValueError(f"only 'embed' or 'linear' is supported: {input_layer}") - + self.normalize_before = normalize_before if self.normalize_before: self.after_norm = LayerNorm(attention_dim) @@ -263,7 +268,7 @@ class FsmnDecoder(BaseTransformerDecoder): self.output_layer = torch.nn.Linear(attention_dim, vocab_size) else: self.output_layer = None - + self.att_layer_num = att_layer_num self.num_blocks = num_blocks if sanm_shfit is None: @@ -276,7 +281,10 @@ class FsmnDecoder(BaseTransformerDecoder): attention_dim, self_attention_dropout_rate, kernel_size, sanm_shfit=sanm_shfit ), MultiHeadedAttentionCrossAtt( - attention_heads, attention_dim, src_attention_dropout_rate, encoder_output_size=encoder_output_size + attention_heads, + attention_dim, + src_attention_dropout_rate, + encoder_output_size=encoder_output_size, ), PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate), dropout_rate, @@ -292,7 +300,10 @@ class FsmnDecoder(BaseTransformerDecoder): lambda lnum: DecoderLayerSANM( attention_dim, MultiHeadedAttentionSANMDecoder( - attention_dim, self_attention_dropout_rate, kernel_size, sanm_shfit=sanm_shfit + attention_dim, + self_attention_dropout_rate, + kernel_size, + sanm_shfit=sanm_shfit, ), None, PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate), @@ -301,7 +312,7 @@ class FsmnDecoder(BaseTransformerDecoder): concat_after, ), ) - + self.decoders3 = repeat( 1, lambda lnum: DecoderLayerSANM( @@ -321,8 +332,12 @@ class FsmnDecoder(BaseTransformerDecoder): attention_dim + encoder_output_size, None, None, - PositionwiseFeedForwardDecoderSANM(attention_dim + encoder_output_size, linear_units, dropout_rate, - adim=attention_dim), + PositionwiseFeedForwardDecoderSANM( + attention_dim + encoder_output_size, + linear_units, + dropout_rate, + adim=attention_dim, + ), dropout_rate, normalize_before, concat_after, @@ -334,7 +349,7 @@ class FsmnDecoder(BaseTransformerDecoder): self.tf2torch_tensor_name_prefix_torch = tf2torch_tensor_name_prefix_torch self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf self.embed_tensor_name_prefix_tf = embed_tensor_name_prefix_tf - + def forward( self, hs_pad: torch.Tensor, @@ -363,47 +378,54 @@ class FsmnDecoder(BaseTransformerDecoder): """ tgt = ys_in_pad tgt_mask = myutils.sequence_mask(ys_in_lens, device=tgt.device)[:, :, None] - + memory = hs_pad memory_mask = myutils.sequence_mask(hlens, device=memory.device)[:, None, :] if chunk_mask is not None: memory_mask = memory_mask * chunk_mask if tgt_mask.size(1) != memory_mask.size(1): memory_mask = torch.cat((memory_mask, memory_mask[:, -2:-1, :]), dim=1) - + x = self.embed(tgt) - + if pre_acoustic_embeds is not None and self.concat_embeds: x = torch.cat((x, pre_acoustic_embeds), dim=-1) x, _, _, _, _ = self.embed_concat_ffn(x, None, None, None, None) - - x, tgt_mask, memory, memory_mask, _ = self.decoders( - x, tgt_mask, memory, memory_mask - ) + + x, tgt_mask, memory, memory_mask, _ = self.decoders(x, tgt_mask, memory, memory_mask) if self.decoders2 is not None: - x, tgt_mask, memory, memory_mask, _ = self.decoders2( - x, tgt_mask, memory, memory_mask - ) - x, tgt_mask, memory, memory_mask, _ = self.decoders3( - x, tgt_mask, memory, memory_mask - ) + x, tgt_mask, memory, memory_mask, _ = self.decoders2(x, tgt_mask, memory, memory_mask) + x, tgt_mask, memory, memory_mask, _ = self.decoders3(x, tgt_mask, memory, memory_mask) if self.normalize_before: x = self.after_norm(x) if self.output_layer is not None: x = self.output_layer(x) - + olens = tgt_mask.sum(1) return x, olens - - def score(self, ys, state, x, x_mask=None, pre_acoustic_embeds: torch.Tensor = None, ): + + def score( + self, + ys, + state, + x, + x_mask=None, + pre_acoustic_embeds: torch.Tensor = None, + ): """Score.""" - ys_mask = myutils.sequence_mask(torch.tensor([len(ys)], dtype=torch.int32), device=x.device)[:, :, None] + ys_mask = myutils.sequence_mask( + torch.tensor([len(ys)], dtype=torch.int32), device=x.device + )[:, :, None] logp, state = self.forward_one_step( - ys.unsqueeze(0), ys_mask, x.unsqueeze(0), memory_mask=x_mask, pre_acoustic_embeds=pre_acoustic_embeds, - cache=state + ys.unsqueeze(0), + ys_mask, + x.unsqueeze(0), + memory_mask=x_mask, + pre_acoustic_embeds=pre_acoustic_embeds, + cache=state, ) return logp.squeeze(0), state - + def forward_one_step( self, tgt: torch.Tensor, @@ -426,15 +448,15 @@ class FsmnDecoder(BaseTransformerDecoder): y, cache: NN output value and cache per `self.decoders`. y.shape` is (batch, maxlen_out, token) """ - + x = tgt[:, -1:] tgt_mask = None x = self.embed(x) - + if pre_acoustic_embeds is not None and self.concat_embeds: x = torch.cat((x, pre_acoustic_embeds), dim=-1) x, _, _, _, _ = self.embed_concat_ffn(x, None, None, None, None) - + if cache is None: cache_layer_num = len(self.decoders) if self.decoders2 is not None: @@ -449,7 +471,7 @@ class FsmnDecoder(BaseTransformerDecoder): x, tgt_mask, memory, memory_mask, cache=c ) new_cache.append(c_ret) - + if self.num_blocks - self.att_layer_num >= 1: for i in range(self.num_blocks - self.att_layer_num): j = i + self.att_layer_num @@ -459,12 +481,12 @@ class FsmnDecoder(BaseTransformerDecoder): x, tgt_mask, memory, memory_mask, cache=c ) new_cache.append(c_ret) - + for decoder in self.decoders3: x, tgt_mask, memory, memory_mask, _ = decoder.forward_one_step( x, tgt_mask, memory, None, cache=None ) - + if self.normalize_before: y = self.after_norm(x[:, -1]) else: @@ -472,7 +494,5 @@ class FsmnDecoder(BaseTransformerDecoder): if self.output_layer is not None: y = self.output_layer(y) y = torch.log_softmax(y, dim=-1) - + return y, new_cache - - \ No newline at end of file diff --git a/funasr/models/sanm/encoder.py b/funasr/models/sanm/encoder.py index f5748187a..dc30a942e 100644 --- a/funasr/models/sanm/encoder.py +++ b/funasr/models/sanm/encoder.py @@ -17,7 +17,10 @@ import numpy as np from funasr.train_utils.device_funcs import to_device from funasr.models.transformer.utils.nets_utils import make_pad_mask from funasr.models.sanm.attention import MultiHeadedAttention, MultiHeadedAttentionSANM -from funasr.models.transformer.embedding import SinusoidalPositionEncoder, StreamSinusoidalPositionEncoder +from funasr.models.transformer.embedding import ( + SinusoidalPositionEncoder, + StreamSinusoidalPositionEncoder, +) from funasr.models.transformer.layer_norm import LayerNorm from funasr.models.transformer.utils.multi_layer_conv import Conv1dLinear from funasr.models.transformer.utils.multi_layer_conv import MultiLayeredConv1d @@ -37,6 +40,7 @@ from funasr.models.ctc.ctc import CTC from funasr.register import tables + class EncoderLayerSANM(nn.Module): def __init__( self, @@ -96,7 +100,18 @@ class EncoderLayerSANM(nn.Module): x = self.norm1(x) if self.concat_after: - x_concat = torch.cat((x, self.self_attn(x, mask, mask_shfit_chunk=mask_shfit_chunk, mask_att_chunk_encoder=mask_att_chunk_encoder)), dim=-1) + x_concat = torch.cat( + ( + x, + self.self_attn( + x, + mask, + mask_shfit_chunk=mask_shfit_chunk, + mask_att_chunk_encoder=mask_att_chunk_encoder, + ), + ), + dim=-1, + ) if self.in_size == self.size: x = residual + stoch_layer_coeff * self.concat_linear(x_concat) else: @@ -104,11 +119,21 @@ class EncoderLayerSANM(nn.Module): else: if self.in_size == self.size: x = residual + stoch_layer_coeff * self.dropout( - self.self_attn(x, mask, mask_shfit_chunk=mask_shfit_chunk, mask_att_chunk_encoder=mask_att_chunk_encoder) + self.self_attn( + x, + mask, + mask_shfit_chunk=mask_shfit_chunk, + mask_att_chunk_encoder=mask_att_chunk_encoder, + ) ) else: x = stoch_layer_coeff * self.dropout( - self.self_attn(x, mask, mask_shfit_chunk=mask_shfit_chunk, mask_att_chunk_encoder=mask_att_chunk_encoder) + self.self_attn( + x, + mask, + mask_shfit_chunk=mask_shfit_chunk, + mask_att_chunk_encoder=mask_att_chunk_encoder, + ) ) if not self.normalize_before: x = self.norm1(x) @@ -158,6 +183,7 @@ class EncoderLayerSANM(nn.Module): return x, cache + @tables.register("encoder_classes", "SANMEncoder") class SANMEncoder(nn.Module): """ @@ -185,8 +211,8 @@ class SANMEncoder(nn.Module): padding_idx: int = -1, interctc_layer_idx: List[int] = [], interctc_use_conditioning: bool = False, - kernel_size : int = 11, - sanm_shfit : int = 0, + kernel_size: int = 11, + sanm_shfit: int = 0, lora_list: List[str] = None, lora_rank: int = 8, lora_alpha: int = 16, @@ -306,7 +332,7 @@ class SANMEncoder(nn.Module): ) self.encoders = repeat( - num_blocks-1, + num_blocks - 1, lambda lnum: EncoderLayerSANM( output_size, output_size, @@ -349,7 +375,7 @@ class SANMEncoder(nn.Module): position embedded tensor and mask """ masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device) - xs_pad = xs_pad * self.output_size()**0.5 + xs_pad = xs_pad * self.output_size() ** 0.5 if self.embed is None: xs_pad = xs_pad elif ( @@ -408,15 +434,16 @@ class SANMEncoder(nn.Module): return feats cache["feats"] = to_device(cache["feats"], device=feats.device) overlap_feats = torch.cat((cache["feats"], feats), dim=1) - cache["feats"] = overlap_feats[:, -(cache["chunk_size"][0] + cache["chunk_size"][2]):, :] + cache["feats"] = overlap_feats[:, -(cache["chunk_size"][0] + cache["chunk_size"][2]) :, :] return overlap_feats - def forward_chunk(self, - xs_pad: torch.Tensor, - ilens: torch.Tensor, - cache: dict = None, - ctc: CTC = None, - ): + def forward_chunk( + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, + cache: dict = None, + ctc: CTC = None, + ): xs_pad *= self.output_size() ** 0.5 if self.embed is None: xs_pad = xs_pad @@ -456,6 +483,7 @@ class SANMEncoder(nn.Module): return (xs_pad, intermediate_outs), None, None return xs_pad, ilens, None + class EncoderLayerSANMExport(nn.Module): def __init__( self, @@ -484,6 +512,7 @@ class EncoderLayerSANMExport(nn.Module): return x, mask + @tables.register("encoder_classes", "SANMEncoderChunkOptExport") @tables.register("encoder_classes", "SANMEncoderExport") class SANMEncoderExport(nn.Module): @@ -492,7 +521,7 @@ class SANMEncoderExport(nn.Module): model, max_seq_len=512, feats_dim=560, - model_name='encoder', + model_name="encoder", onnx: bool = True, ): super().__init__() @@ -503,28 +532,27 @@ class SANMEncoderExport(nn.Module): self.feats_dim = feats_dim self._output_size = model._output_size - from funasr.utils.torch_function import sequence_mask - self.make_pad_mask = sequence_mask(max_seq_len, flip=False) - + from funasr.models.sanm.attention import MultiHeadedAttentionSANMExport - if hasattr(model, 'encoders0'): + + if hasattr(model, "encoders0"): for i, d in enumerate(self.model.encoders0): if isinstance(d.self_attn, MultiHeadedAttentionSANM): d.self_attn = MultiHeadedAttentionSANMExport(d.self_attn) self.model.encoders0[i] = EncoderLayerSANMExport(d) - + for i, d in enumerate(self.model.encoders): if isinstance(d.self_attn, MultiHeadedAttentionSANM): d.self_attn = MultiHeadedAttentionSANMExport(d.self_attn) self.model.encoders[i] = EncoderLayerSANMExport(d) - + self.model_name = model_name self.num_heads = model.encoders[0].self_attn.h self.hidden_size = model.encoders[0].self_attn.linear_out.out_features - + def prepare_mask(self, mask): mask_3d_btd = mask[:, :, None] if len(mask.shape) == 2: @@ -532,57 +560,45 @@ class SANMEncoderExport(nn.Module): elif len(mask.shape) == 3: mask_4d_bhlt = 1 - mask[:, None, :] mask_4d_bhlt = mask_4d_bhlt * -10000.0 - + return mask_3d_btd, mask_4d_bhlt - - def forward(self, - speech: torch.Tensor, - speech_lengths: torch.Tensor, - online: bool = False - ): + + def forward(self, speech: torch.Tensor, speech_lengths: torch.Tensor, online: bool = False): if not online: - speech = speech * self._output_size ** 0.5 + speech = speech * self._output_size**0.5 mask = self.make_pad_mask(speech_lengths) mask = self.prepare_mask(mask) if self.embed is None: xs_pad = speech else: xs_pad = self.embed(speech) - + encoder_outs = self.model.encoders0(xs_pad, mask) xs_pad, masks = encoder_outs[0], encoder_outs[1] - + encoder_outs = self.model.encoders(xs_pad, mask) xs_pad, masks = encoder_outs[0], encoder_outs[1] - + xs_pad = self.model.after_norm(xs_pad) - + return xs_pad, speech_lengths - + def get_output_size(self): return self.model.encoders[0].size - + def get_dummy_inputs(self): feats = torch.randn(1, 100, self.feats_dim) - return (feats) - + return feats + def get_input_names(self): - return ['feats'] - + return ["feats"] + def get_output_names(self): - return ['encoder_out', 'encoder_out_lens', 'predictor_weight'] - + return ["encoder_out", "encoder_out_lens", "predictor_weight"] + def get_dynamic_axes(self): return { - 'feats': { - 1: 'feats_length' - }, - 'encoder_out': { - 1: 'enc_out_length' - }, - 'predictor_weight': { - 1: 'pre_out_length' - } - + "feats": {1: "feats_length"}, + "encoder_out": {1: "enc_out_length"}, + "predictor_weight": {1: "pre_out_length"}, } - diff --git a/funasr/models/sanm/model.py b/funasr/models/sanm/model.py index 0cef54061..b3f84b2b9 100644 --- a/funasr/models/sanm/model.py +++ b/funasr/models/sanm/model.py @@ -10,6 +10,7 @@ import torch from funasr.models.transformer.model import Transformer from funasr.register import tables + @tables.register("model_classes", "SANM") class SANM(Transformer): """ diff --git a/funasr/models/sanm/multihead_att.py b/funasr/models/sanm/multihead_att.py index 5ef36eda2..c7d97961d 100644 --- a/funasr/models/sanm/multihead_att.py +++ b/funasr/models/sanm/multihead_att.py @@ -22,7 +22,7 @@ class MultiHeadedAttentionSANMExport(nn.Module): mask_3d_btd, mask_4d_bhlt = mask q_h, k_h, v_h, v = self.forward_qkv(x) fsmn_memory = self.forward_fsmn(v, mask_3d_btd) - q_h = q_h * self.d_k**(-0.5) + q_h = q_h * self.d_k ** (-0.5) scores = torch.matmul(q_h, k_h.transpose(-2, -1)) att_outs = self.forward_attention(v_h, scores, mask_4d_bhlt) return att_outs + fsmn_memory @@ -71,14 +71,15 @@ def preprocess_for_attn(x, mask, cache, pad_fn, kernel_size): x = pad_fn(x) else: x = torch.cat((cache, x), dim=2) - cache = x[:, :, -(kernel_size-1):] + cache = x[:, :, -(kernel_size - 1) :] return x, cache torch_version = tuple([int(i) for i in torch.__version__.split(".")[:2]]) if torch_version >= (1, 8): import torch.fx - torch.fx.wrap('preprocess_for_attn') + + torch.fx.wrap("preprocess_for_attn") class MultiHeadedAttentionSANMDecoderExport(nn.Module): @@ -153,7 +154,7 @@ class OnnxMultiHeadedAttention(nn.Module): self.linear_out = model.linear_out self.attn = None self.all_head_size = self.h * self.d_k - + def forward(self, query, key, value, mask): q, k, v = self.forward_qkv(query, key, value) scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) @@ -172,13 +173,13 @@ class OnnxMultiHeadedAttention(nn.Module): k = self.transpose_for_scores(k) v = self.transpose_for_scores(v) return q, k, v - + def forward_attention(self, value, scores, mask): scores = scores + mask self.attn = torch.softmax(scores, dim=-1) context_layer = torch.matmul(self.attn, value) # (batch, head, time1, d_k) - + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) context_layer = context_layer.view(new_context_layer_shape) @@ -191,12 +192,12 @@ class OnnxRelPosMultiHeadedAttention(OnnxMultiHeadedAttention): self.linear_pos = model.linear_pos self.pos_bias_u = model.pos_bias_u self.pos_bias_v = model.pos_bias_v - + def forward(self, query, key, value, pos_emb, mask): q, k, v = self.forward_qkv(query, key, value) q = q.transpose(1, 2) # (batch, time1, head, d_k) - p = self.transpose_for_scores(self.linear_pos(pos_emb)) # (batch, head, time1, d_k) + p = self.transpose_for_scores(self.linear_pos(pos_emb)) # (batch, head, time1, d_k) # (batch, head, time1, d_k) q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2) @@ -214,9 +215,7 @@ class OnnxRelPosMultiHeadedAttention(OnnxMultiHeadedAttention): matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) matrix_bd = self.rel_shift(matrix_bd) - scores = (matrix_ac + matrix_bd) / math.sqrt( - self.d_k - ) # (batch, head, time1, time2) + scores = (matrix_ac + matrix_bd) / math.sqrt(self.d_k) # (batch, head, time1, time2) return self.forward_attention(v, scores, mask) @@ -235,9 +234,8 @@ class OnnxRelPosMultiHeadedAttention(OnnxMultiHeadedAttention): self.attn = torch.softmax(scores, dim=-1) context_layer = torch.matmul(self.attn, value) # (batch, head, time1, d_k) - + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) context_layer = context_layer.view(new_context_layer_shape) return self.linear_out(context_layer) # (batch, time1, d_model) - diff --git a/funasr/models/sanm/positionwise_feed_forward.py b/funasr/models/sanm/positionwise_feed_forward.py index 7125854e5..e28ba8825 100644 --- a/funasr/models/sanm/positionwise_feed_forward.py +++ b/funasr/models/sanm/positionwise_feed_forward.py @@ -9,7 +9,6 @@ import torch from funasr.models.transformer.layer_norm import LayerNorm - class PositionwiseFeedForwardDecoderSANM(torch.nn.Module): """Positionwise feed forward layer. diff --git a/funasr/models/scama/beam_search.py b/funasr/models/scama/beam_search.py index b8aa876b5..259f483b2 100644 --- a/funasr/models/scama/beam_search.py +++ b/funasr/models/scama/beam_search.py @@ -33,7 +33,6 @@ class Hypothesis(NamedTuple): )._asdict() - class BeamSearchScama(torch.nn.Module): """Beam search implementation.""" @@ -151,7 +150,8 @@ class BeamSearchScama(torch.nn.Module): return torch.cat((xs, x)) def score_full( - self, hyp: Hypothesis, + self, + hyp: Hypothesis, x: torch.Tensor, x_mask: torch.Tensor = None, pre_acoustic_embeds: torch.Tensor = None, @@ -173,7 +173,9 @@ class BeamSearchScama(torch.nn.Module): scores = dict() states = dict() for k, d in self.full_scorers.items(): - scores[k], states[k] = d.score(hyp.yseq, hyp.states[k], x, x_mask=x_mask, pre_acoustic_embeds=pre_acoustic_embeds) + scores[k], states[k] = d.score( + hyp.yseq, hyp.states[k], x, x_mask=x_mask, pre_acoustic_embeds=pre_acoustic_embeds + ) return scores, states def score_partial( @@ -283,7 +285,8 @@ class BeamSearchScama(torch.nn.Module): return new_states def search( - self, running_hyps: List[Hypothesis], + self, + running_hyps: List[Hypothesis], x: torch.Tensor, x_mask: torch.Tensor = None, pre_acoustic_embeds: torch.Tensor = None, @@ -303,7 +306,9 @@ class BeamSearchScama(torch.nn.Module): for hyp in running_hyps: # scoring weighted_scores = torch.zeros(self.n_vocab, dtype=x.dtype, device=x.device) - scores, states = self.score_full(hyp, x, x_mask=x_mask, pre_acoustic_embeds=pre_acoustic_embeds) + scores, states = self.score_full( + hyp, x, x_mask=x_mask, pre_acoustic_embeds=pre_acoustic_embeds + ) for k in self.full_scorers: weighted_scores += self.weights[k] * scores[k] # partial scoring @@ -327,9 +332,7 @@ class BeamSearchScama(torch.nn.Module): Hypothesis( score=weighted_scores[j], yseq=self.append_token(hyp.yseq, j), - scores=self.merge_scores( - hyp.scores, scores, j, part_scores, part_j - ), + scores=self.merge_scores(hyp.scores, scores, j, part_scores, part_j), states=self.merge_states(states, part_states, part_j), ) ) @@ -341,7 +344,8 @@ class BeamSearchScama(torch.nn.Module): return best_hyps def forward( - self, x: torch.Tensor, + self, + x: torch.Tensor, scama_mask: torch.Tensor = None, pre_acoustic_embeds: torch.Tensor = None, maxlenratio: float = 0.0, @@ -386,20 +390,26 @@ class BeamSearchScama(torch.nn.Module): mask_enc = None if scama_mask is not None: token_num_predictor = scama_mask.size(1) - token_id_slice = min(i, token_num_predictor-1) - mask_enc = scama_mask[:, token_id_slice:token_id_slice+1, :] + token_id_slice = min(i, token_num_predictor - 1) + mask_enc = scama_mask[:, token_id_slice : token_id_slice + 1, :] # if mask_enc.size(1) == 0: # mask_enc = scama_mask[:, -2:-1, :] # # mask_enc = torch.zeros_like(mask_enc) pre_acoustic_embeds_cur = None if pre_acoustic_embeds is not None: b, t, d = pre_acoustic_embeds.size() - pad = torch.zeros((b, 1, d), dtype=pre_acoustic_embeds.dtype).to(device=pre_acoustic_embeds.device) + pad = torch.zeros((b, 1, d), dtype=pre_acoustic_embeds.dtype).to( + device=pre_acoustic_embeds.device + ) pre_acoustic_embeds = torch.cat((pre_acoustic_embeds, pad), dim=1) token_id_slice = min(i, t) - pre_acoustic_embeds_cur = pre_acoustic_embeds[:, token_id_slice:token_id_slice+1, :] + pre_acoustic_embeds_cur = pre_acoustic_embeds[ + :, token_id_slice : token_id_slice + 1, : + ] - best = self.search(running_hyps, x, x_mask=mask_enc, pre_acoustic_embeds=pre_acoustic_embeds_cur) + best = self.search( + running_hyps, x, x_mask=mask_enc, pre_acoustic_embeds=pre_acoustic_embeds_cur + ) # post process of one iteration running_hyps = self.post_process(i, maxlen, maxlenratio, best, ended_hyps) # end detection @@ -416,8 +426,7 @@ class BeamSearchScama(torch.nn.Module): # check the number of hypotheses reaching to eos if len(nbest_hyps) == 0: logging.warning( - "there is no N-best results, perform recognition " - "again with smaller minlenratio." + "there is no N-best results, perform recognition " "again with smaller minlenratio." ) return ( [] @@ -431,17 +440,13 @@ class BeamSearchScama(torch.nn.Module): logging.debug("nbest: y: {}, yseq: {}, score: {}".format(x.yseq, yseq, x.score)) best = nbest_hyps[0] for k, v in best.scores.items(): - logging.info( - f"{v:6.2f} * {self.weights[k]:3} = {v * self.weights[k]:6.2f} for {k}" - ) + logging.info(f"{v:6.2f} * {self.weights[k]:3} = {v * self.weights[k]:6.2f} for {k}") logging.info(f"total log probability: {best.score:.2f}") logging.info(f"normalized log probability: {best.score / len(best.yseq):.2f}") logging.info(f"total number of ended hypotheses: {len(nbest_hyps)}") if self.token_list is not None: logging.info( - "best hypo: " - + "".join([self.token_list[x] for x in best.yseq[1:-1]]) - + "\n" + "best hypo: " + "".join([self.token_list[x] for x in best.yseq[1:-1]]) + "\n" ) return nbest_hyps @@ -469,15 +474,13 @@ class BeamSearchScama(torch.nn.Module): logging.debug(f"the number of running hypotheses: {len(running_hyps)}") if self.token_list is not None: logging.debug( - "best hypo: " - + "".join([self.token_list[x] for x in running_hyps[0].yseq[1:]]) + "best hypo: " + "".join([self.token_list[x] for x in running_hyps[0].yseq[1:]]) ) # add eos in the final loop to avoid that there are no ended hyps if i == maxlen - 1: logging.info("adding in the last position in the loop") running_hyps = [ - h._replace(yseq=self.append_token(h.yseq, self.eos)) - for h in running_hyps + h._replace(yseq=self.append_token(h.yseq, self.eos)) for h in running_hyps ] # add ended hypotheses to a final list, and removed them from current hypotheses @@ -495,6 +498,7 @@ class BeamSearchScama(torch.nn.Module): remained_hyps.append(hyp) return remained_hyps + class BeamSearchScamaStreaming(torch.nn.Module): """Beam search implementation.""" @@ -612,11 +616,12 @@ class BeamSearchScamaStreaming(torch.nn.Module): return torch.cat((xs, x)) def score_full( - self, hyp: Hypothesis, + self, + hyp: Hypothesis, x: torch.Tensor, x_mask: torch.Tensor = None, pre_acoustic_embeds: torch.Tensor = None, - cache: dict={}, + cache: dict = {}, ) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: """Score new hypothesis by `self.full_scorers`. @@ -635,7 +640,14 @@ class BeamSearchScamaStreaming(torch.nn.Module): scores = dict() states = dict() for k, d in self.full_scorers.items(): - scores[k], states[k] = d.score(hyp.yseq, hyp.states[k], x, x_mask=x_mask, pre_acoustic_embeds=pre_acoustic_embeds, cache=cache) + scores[k], states[k] = d.score( + hyp.yseq, + hyp.states[k], + x, + x_mask=x_mask, + pre_acoustic_embeds=pre_acoustic_embeds, + cache=cache, + ) return scores, states def score_partial( @@ -745,11 +757,12 @@ class BeamSearchScamaStreaming(torch.nn.Module): return new_states def search( - self, running_hyps: List[Hypothesis], + self, + running_hyps: List[Hypothesis], x: torch.Tensor, x_mask: torch.Tensor = None, pre_acoustic_embeds: torch.Tensor = None, - cache: dict={}, + cache: dict = {}, ) -> List[Hypothesis]: """Search new tokens for running hypotheses and encoded speech x. @@ -766,7 +779,9 @@ class BeamSearchScamaStreaming(torch.nn.Module): for hyp in running_hyps: # scoring weighted_scores = torch.zeros(self.n_vocab, dtype=x.dtype, device=x.device) - scores, states = self.score_full(hyp, x, x_mask=x_mask, pre_acoustic_embeds=pre_acoustic_embeds, cache=cache) + scores, states = self.score_full( + hyp, x, x_mask=x_mask, pre_acoustic_embeds=pre_acoustic_embeds, cache=cache + ) for k in self.full_scorers: weighted_scores += self.weights[k] * scores[k] # partial scoring @@ -790,9 +805,7 @@ class BeamSearchScamaStreaming(torch.nn.Module): Hypothesis( score=weighted_scores[j], yseq=self.append_token(hyp.yseq, j), - scores=self.merge_scores( - hyp.scores, scores, j, part_scores, part_j - ), + scores=self.merge_scores(hyp.scores, scores, j, part_scores, part_j), states=self.merge_states(states, part_states, part_j), ) ) @@ -804,14 +817,15 @@ class BeamSearchScamaStreaming(torch.nn.Module): return best_hyps def forward( - self, x: torch.Tensor, + self, + x: torch.Tensor, scama_mask: torch.Tensor = None, pre_acoustic_embeds: torch.Tensor = None, maxlenratio: float = 0.0, minlenratio: float = 0.0, maxlen: int = None, minlen: int = 0, - cache:dict={}, + cache: dict = {}, ) -> List[Hypothesis]: """Perform beam search. @@ -859,12 +873,22 @@ class BeamSearchScamaStreaming(torch.nn.Module): pre_acoustic_embeds_cur = None if pre_acoustic_embeds is not None: b, t, d = pre_acoustic_embeds.size() - pad = torch.zeros((b, 1, d), dtype=pre_acoustic_embeds.dtype).to(device=pre_acoustic_embeds.device) + pad = torch.zeros((b, 1, d), dtype=pre_acoustic_embeds.dtype).to( + device=pre_acoustic_embeds.device + ) pre_acoustic_embeds = torch.cat((pre_acoustic_embeds, pad), dim=1) token_id_slice = min(i, t) - pre_acoustic_embeds_cur = pre_acoustic_embeds[:, token_id_slice:token_id_slice+1, :] + pre_acoustic_embeds_cur = pre_acoustic_embeds[ + :, token_id_slice : token_id_slice + 1, : + ] - best = self.search(running_hyps, x, x_mask=mask_enc, pre_acoustic_embeds=pre_acoustic_embeds_cur, cache=cache["decoder"]) + best = self.search( + running_hyps, + x, + x_mask=mask_enc, + pre_acoustic_embeds=pre_acoustic_embeds_cur, + cache=cache["decoder"], + ) # post process of one iteration running_hyps = self.post_process(i, maxlen, maxlenratio, best, ended_hyps) # end detection @@ -881,8 +905,7 @@ class BeamSearchScamaStreaming(torch.nn.Module): # check the number of hypotheses reaching to eos if len(nbest_hyps) == 0: logging.warning( - "there is no N-best results, perform recognition " - "again with smaller minlenratio." + "there is no N-best results, perform recognition " "again with smaller minlenratio." ) return ( [] @@ -896,17 +919,13 @@ class BeamSearchScamaStreaming(torch.nn.Module): logging.debug("nbest: y: {}, yseq: {}, score: {}".format(x.yseq, yseq, x.score)) best = nbest_hyps[0] for k, v in best.scores.items(): - logging.info( - f"{v:6.2f} * {self.weights[k]:3} = {v * self.weights[k]:6.2f} for {k}" - ) + logging.info(f"{v:6.2f} * {self.weights[k]:3} = {v * self.weights[k]:6.2f} for {k}") logging.info(f"total log probability: {best.score:.2f}") logging.info(f"normalized log probability: {best.score / len(best.yseq):.2f}") logging.info(f"total number of ended hypotheses: {len(nbest_hyps)}") if self.token_list is not None: logging.info( - "best hypo: " - + "".join([self.token_list[x] for x in best.yseq[1:-1]]) - + "\n" + "best hypo: " + "".join([self.token_list[x] for x in best.yseq[1:-1]]) + "\n" ) return nbest_hyps @@ -934,15 +953,13 @@ class BeamSearchScamaStreaming(torch.nn.Module): logging.debug(f"the number of running hypotheses: {len(running_hyps)}") if self.token_list is not None: logging.debug( - "best hypo: " - + "".join([self.token_list[x] for x in running_hyps[0].yseq[1:]]) + "best hypo: " + "".join([self.token_list[x] for x in running_hyps[0].yseq[1:]]) ) # add eos in the final loop to avoid that there are no ended hyps if i == maxlen - 1: logging.info("adding in the last position in the loop") running_hyps = [ - h._replace(yseq=self.append_token(h.yseq, self.eos)) - for h in running_hyps + h._replace(yseq=self.append_token(h.yseq, self.eos)) for h in running_hyps ] # add ended hypotheses to a final list, and removed them from current hypotheses diff --git a/funasr/models/scama/chunk_utilis.py b/funasr/models/scama/chunk_utilis.py index 245d28211..2fe3fa4e3 100644 --- a/funasr/models/scama/chunk_utilis.py +++ b/funasr/models/scama/chunk_utilis.py @@ -7,14 +7,16 @@ from funasr.models.scama.utils import sequence_mask from funasr.models.transformer.utils.nets_utils import make_pad_mask -class overlap_chunk(): +class overlap_chunk: """ Author: Speech Lab of DAMO Academy, Alibaba Group San-m: Memory equipped self-attention for end-to-end speech recognition https://arxiv.org/abs/2006.01713 """ - def __init__(self, + + def __init__( + self, chunk_size: tuple = (16,), stride: tuple = (10,), pad_left: tuple = (0,), @@ -24,10 +26,25 @@ class overlap_chunk(): ): pad_left = self.check_chunk_size_args(chunk_size, pad_left) - encoder_att_look_back_factor = self.check_chunk_size_args(chunk_size, encoder_att_look_back_factor) - decoder_att_look_back_factor = self.check_chunk_size_args(chunk_size, decoder_att_look_back_factor) - self.chunk_size, self.stride, self.pad_left, self.encoder_att_look_back_factor, self.decoder_att_look_back_factor \ - = chunk_size, stride, pad_left, encoder_att_look_back_factor, decoder_att_look_back_factor + encoder_att_look_back_factor = self.check_chunk_size_args( + chunk_size, encoder_att_look_back_factor + ) + decoder_att_look_back_factor = self.check_chunk_size_args( + chunk_size, decoder_att_look_back_factor + ) + ( + self.chunk_size, + self.stride, + self.pad_left, + self.encoder_att_look_back_factor, + self.decoder_att_look_back_factor, + ) = ( + chunk_size, + stride, + pad_left, + encoder_att_look_back_factor, + decoder_att_look_back_factor, + ) self.shfit_fsmn = shfit_fsmn self.x_add_mask = None self.x_rm_mask = None @@ -37,23 +54,50 @@ class overlap_chunk(): self.mask_att_chunk_encoder = None self.mask_shift_att_chunk_decoder = None self.chunk_outs = None - self.chunk_size_cur, self.stride_cur, self.pad_left_cur, self.encoder_att_look_back_factor_cur, self.chunk_size_pad_shift_cur \ - = None, None, None, None, None + ( + self.chunk_size_cur, + self.stride_cur, + self.pad_left_cur, + self.encoder_att_look_back_factor_cur, + self.chunk_size_pad_shift_cur, + ) = (None, None, None, None, None) def check_chunk_size_args(self, chunk_size, x): if len(x) < len(chunk_size): x = [x[0] for i in chunk_size] return x - def get_chunk_size(self, - ind: int = 0 - ): + def get_chunk_size(self, ind: int = 0): # with torch.no_grad: - chunk_size, stride, pad_left, encoder_att_look_back_factor, decoder_att_look_back_factor = \ - self.chunk_size[ind], self.stride[ind], self.pad_left[ind], self.encoder_att_look_back_factor[ind], self.decoder_att_look_back_factor[ind] - self.chunk_size_cur, self.stride_cur, self.pad_left_cur, self.encoder_att_look_back_factor_cur, self.chunk_size_pad_shift_cur, self.decoder_att_look_back_factor_cur \ - = chunk_size, stride, pad_left, encoder_att_look_back_factor, chunk_size + self.shfit_fsmn, decoder_att_look_back_factor - return self.chunk_size_cur, self.stride_cur, self.pad_left_cur, self.encoder_att_look_back_factor_cur, self.chunk_size_pad_shift_cur + chunk_size, stride, pad_left, encoder_att_look_back_factor, decoder_att_look_back_factor = ( + self.chunk_size[ind], + self.stride[ind], + self.pad_left[ind], + self.encoder_att_look_back_factor[ind], + self.decoder_att_look_back_factor[ind], + ) + ( + self.chunk_size_cur, + self.stride_cur, + self.pad_left_cur, + self.encoder_att_look_back_factor_cur, + self.chunk_size_pad_shift_cur, + self.decoder_att_look_back_factor_cur, + ) = ( + chunk_size, + stride, + pad_left, + encoder_att_look_back_factor, + chunk_size + self.shfit_fsmn, + decoder_att_look_back_factor, + ) + return ( + self.chunk_size_cur, + self.stride_cur, + self.pad_left_cur, + self.encoder_att_look_back_factor_cur, + self.chunk_size_pad_shift_cur, + ) def random_choice(self, training=True, decoding_ind=None): chunk_num = len(self.chunk_size) @@ -65,25 +109,31 @@ class overlap_chunk(): return ind - - - def gen_chunk_mask(self, x_len, ind=0, num_units=1, num_units_predictor=1): with torch.no_grad(): x_len = x_len.cpu().numpy() x_len_max = x_len.max() - chunk_size, stride, pad_left, encoder_att_look_back_factor, chunk_size_pad_shift = self.get_chunk_size(ind) + chunk_size, stride, pad_left, encoder_att_look_back_factor, chunk_size_pad_shift = ( + self.get_chunk_size(ind) + ) shfit_fsmn = self.shfit_fsmn pad_right = chunk_size - stride - pad_left - chunk_num_batch = np.ceil(x_len/stride).astype(np.int32) - x_len_chunk = (chunk_num_batch-1) * chunk_size_pad_shift + shfit_fsmn + pad_left + 0 + x_len - (chunk_num_batch-1) * stride + chunk_num_batch = np.ceil(x_len / stride).astype(np.int32) + x_len_chunk = ( + (chunk_num_batch - 1) * chunk_size_pad_shift + + shfit_fsmn + + pad_left + + 0 + + x_len + - (chunk_num_batch - 1) * stride + ) x_len_chunk = x_len_chunk.astype(x_len.dtype) x_len_chunk_max = x_len_chunk.max() - chunk_num = int(math.ceil(x_len_max/stride)) + chunk_num = int(math.ceil(x_len_max / stride)) dtype = np.int32 max_len_for_x_mask_tmp = max(chunk_size, x_len_max + pad_left) x_add_mask = np.zeros([0, max_len_for_x_mask_tmp], dtype=dtype) @@ -91,7 +141,7 @@ class overlap_chunk(): mask_shfit_chunk = np.zeros([0, num_units], dtype=dtype) mask_chunk_predictor = np.zeros([0, num_units_predictor], dtype=dtype) mask_shift_att_chunk_decoder = np.zeros([0, 1], dtype=dtype) - mask_att_chunk_encoder = np.zeros([0, chunk_num*chunk_size_pad_shift], dtype=dtype) + mask_att_chunk_encoder = np.zeros([0, chunk_num * chunk_size_pad_shift], dtype=dtype) for chunk_ids in range(chunk_num): # x_mask add fsmn_padding = np.zeros((shfit_fsmn, max_len_for_x_mask_tmp), dtype=dtype) @@ -104,15 +154,19 @@ class overlap_chunk(): x_add_mask = np.concatenate([x_add_mask, x_add_mask_fsmn], axis=0) # x_mask rm - fsmn_padding = np.zeros((max_len_for_x_mask_tmp, shfit_fsmn),dtype=dtype) - padding_mask_left = np.zeros((max_len_for_x_mask_tmp, pad_left),dtype=dtype) + fsmn_padding = np.zeros((max_len_for_x_mask_tmp, shfit_fsmn), dtype=dtype) + padding_mask_left = np.zeros((max_len_for_x_mask_tmp, pad_left), dtype=dtype) padding_mask_right = np.zeros((max_len_for_x_mask_tmp, pad_right), dtype=dtype) x_mask_cur = np.diag(np.ones(stride, dtype=dtype)) - x_mask_cur_pad_top = np.zeros((chunk_ids*stride, stride), dtype=dtype) + x_mask_cur_pad_top = np.zeros((chunk_ids * stride, stride), dtype=dtype) x_mask_cur_pad_bottom = np.zeros((max_len_for_x_mask_tmp, stride), dtype=dtype) - x_rm_mask_cur = np.concatenate([x_mask_cur_pad_top, x_mask_cur, x_mask_cur_pad_bottom], axis=0) + x_rm_mask_cur = np.concatenate( + [x_mask_cur_pad_top, x_mask_cur, x_mask_cur_pad_bottom], axis=0 + ) x_rm_mask_cur = x_rm_mask_cur[:max_len_for_x_mask_tmp, :stride] - x_rm_mask_cur_fsmn = np.concatenate([fsmn_padding, padding_mask_left, x_rm_mask_cur, padding_mask_right], axis=1) + x_rm_mask_cur_fsmn = np.concatenate( + [fsmn_padding, padding_mask_left, x_rm_mask_cur, padding_mask_right], axis=1 + ) x_rm_mask = np.concatenate([x_rm_mask, x_rm_mask_cur_fsmn], axis=1) # fsmn_padding_mask @@ -124,22 +178,26 @@ class overlap_chunk(): # predictor mask zeros_1 = np.zeros([shfit_fsmn + pad_left, num_units_predictor], dtype=dtype) ones_2 = np.ones([stride, num_units_predictor], dtype=dtype) - zeros_3 = np.zeros([chunk_size - stride - pad_left, num_units_predictor], dtype=dtype) + zeros_3 = np.zeros( + [chunk_size - stride - pad_left, num_units_predictor], dtype=dtype + ) ones_zeros = np.concatenate([ones_2, zeros_3], axis=0) mask_chunk_predictor_cur = np.concatenate([zeros_1, ones_zeros], axis=0) - mask_chunk_predictor = np.concatenate([mask_chunk_predictor, mask_chunk_predictor_cur], axis=0) + mask_chunk_predictor = np.concatenate( + [mask_chunk_predictor, mask_chunk_predictor_cur], axis=0 + ) # encoder att mask - zeros_1_top = np.zeros([shfit_fsmn, chunk_num*chunk_size_pad_shift], dtype=dtype) + zeros_1_top = np.zeros([shfit_fsmn, chunk_num * chunk_size_pad_shift], dtype=dtype) zeros_2_num = max(chunk_ids - encoder_att_look_back_factor, 0) - zeros_2 = np.zeros([chunk_size, zeros_2_num*chunk_size_pad_shift], dtype=dtype) + zeros_2 = np.zeros([chunk_size, zeros_2_num * chunk_size_pad_shift], dtype=dtype) encoder_att_look_back_num = max(chunk_ids - zeros_2_num, 0) zeros_2_left = np.zeros([chunk_size, shfit_fsmn], dtype=dtype) ones_2_mid = np.ones([stride, stride], dtype=dtype) - zeros_2_bottom = np.zeros([chunk_size-stride, stride], dtype=dtype) - zeros_2_right = np.zeros([chunk_size, chunk_size-stride], dtype=dtype) + zeros_2_bottom = np.zeros([chunk_size - stride, stride], dtype=dtype) + zeros_2_right = np.zeros([chunk_size, chunk_size - stride], dtype=dtype) ones_2 = np.concatenate([ones_2_mid, zeros_2_bottom], axis=0) ones_2 = np.concatenate([zeros_2_left, ones_2, zeros_2_right], axis=1) ones_2 = np.tile(ones_2, [1, encoder_att_look_back_num]) @@ -149,21 +207,25 @@ class overlap_chunk(): ones_3 = np.concatenate([zeros_3_left, ones_3_right], axis=1) zeros_remain_num = max(chunk_num - 1 - chunk_ids, 0) - zeros_remain = np.zeros([chunk_size, zeros_remain_num*chunk_size_pad_shift], dtype=dtype) + zeros_remain = np.zeros( + [chunk_size, zeros_remain_num * chunk_size_pad_shift], dtype=dtype + ) ones2_bottom = np.concatenate([zeros_2, ones_2, ones_3, zeros_remain], axis=1) mask_att_chunk_encoder_cur = np.concatenate([zeros_1_top, ones2_bottom], axis=0) - mask_att_chunk_encoder = np.concatenate([mask_att_chunk_encoder, mask_att_chunk_encoder_cur], axis=0) - + mask_att_chunk_encoder = np.concatenate( + [mask_att_chunk_encoder, mask_att_chunk_encoder_cur], axis=0 + ) # decoder fsmn_shift_att_mask zeros_1 = np.zeros([shfit_fsmn, 1]) ones_1 = np.ones([chunk_size, 1]) mask_shift_att_chunk_decoder_cur = np.concatenate([zeros_1, ones_1], axis=0) mask_shift_att_chunk_decoder = np.concatenate( - [mask_shift_att_chunk_decoder, mask_shift_att_chunk_decoder_cur], axis=0) + [mask_shift_att_chunk_decoder, mask_shift_att_chunk_decoder_cur], axis=0 + ) - self.x_add_mask = x_add_mask[:x_len_chunk_max, :x_len_max+pad_left] + self.x_add_mask = x_add_mask[:x_len_chunk_max, : x_len_max + pad_left] self.x_len_chunk = x_len_chunk self.x_rm_mask = x_rm_mask[:x_len_max, :x_len_chunk_max] self.x_len = x_len @@ -171,18 +233,19 @@ class overlap_chunk(): self.mask_chunk_predictor = mask_chunk_predictor[:x_len_chunk_max, :] self.mask_att_chunk_encoder = mask_att_chunk_encoder[:x_len_chunk_max, :x_len_chunk_max] self.mask_shift_att_chunk_decoder = mask_shift_att_chunk_decoder[:x_len_chunk_max, :] - self.chunk_outs = (self.x_add_mask, + self.chunk_outs = ( + self.x_add_mask, self.x_len_chunk, self.x_rm_mask, self.x_len, self.mask_shfit_chunk, self.mask_chunk_predictor, self.mask_att_chunk_encoder, - self.mask_shift_att_chunk_decoder) + self.mask_shift_att_chunk_decoder, + ) return self.chunk_outs - def split_chunk(self, x, x_len, chunk_outs): """ :param x: (b, t, d) @@ -190,10 +253,9 @@ class overlap_chunk(): :param ind: int :return: """ - x = x[:, :x_len.max(), :] + x = x[:, : x_len.max(), :] b, t, d = x.size() - x_len_mask = (~make_pad_mask(x_len, maxlen=t)).to( - x.device) + x_len_mask = (~make_pad_mask(x_len, maxlen=t)).to(x.device) x *= x_len_mask[:, :, None] x_add_mask = self.get_x_add_mask(chunk_outs, x.device, dtype=x.dtype) @@ -209,10 +271,9 @@ class overlap_chunk(): return x_chunk, x_len_chunk def remove_chunk(self, x_chunk, x_len_chunk, chunk_outs): - x_chunk = x_chunk[:, :x_len_chunk.max(), :] + x_chunk = x_chunk[:, : x_len_chunk.max(), :] b, t, d = x_chunk.size() - x_len_chunk_mask = (~make_pad_mask(x_len_chunk, maxlen=t)).to( - x_chunk.device) + x_len_chunk_mask = (~make_pad_mask(x_len_chunk, maxlen=t)).to(x_chunk.device) x_chunk *= x_len_chunk_mask[:, :, None] x_rm_mask = self.get_x_rm_mask(chunk_outs, x_chunk.device, dtype=x_chunk.dtype) @@ -224,77 +285,104 @@ class overlap_chunk(): return x, x_len - def get_x_add_mask(self, chunk_outs=None, device='cpu', idx=0, dtype=torch.float32): + def get_x_add_mask(self, chunk_outs=None, device="cpu", idx=0, dtype=torch.float32): with torch.no_grad(): - x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx] + x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx] x = torch.from_numpy(x).type(dtype).to(device) return x - def get_x_len_chunk(self, chunk_outs=None, device='cpu', idx=1, dtype=torch.float32): + def get_x_len_chunk(self, chunk_outs=None, device="cpu", idx=1, dtype=torch.float32): with torch.no_grad(): - x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx] + x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx] x = torch.from_numpy(x).type(dtype).to(device) return x - - def get_x_rm_mask(self, chunk_outs=None, device='cpu', idx=2, dtype=torch.float32): + def get_x_rm_mask(self, chunk_outs=None, device="cpu", idx=2, dtype=torch.float32): with torch.no_grad(): - x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx] + x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx] x = torch.from_numpy(x).type(dtype).to(device) return x - def get_x_len(self, chunk_outs=None, device='cpu', idx=3, dtype=torch.float32): + def get_x_len(self, chunk_outs=None, device="cpu", idx=3, dtype=torch.float32): with torch.no_grad(): - x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx] + x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx] x = torch.from_numpy(x).type(dtype).to(device) return x - - def get_mask_shfit_chunk(self, chunk_outs=None, device='cpu', batch_size=1, num_units=1, idx=4, dtype=torch.float32): + def get_mask_shfit_chunk( + self, chunk_outs=None, device="cpu", batch_size=1, num_units=1, idx=4, dtype=torch.float32 + ): with torch.no_grad(): - x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx] - x = np.tile(x[None, :, :, ], [batch_size, 1, num_units]) + x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx] + x = np.tile( + x[ + None, + :, + :, + ], + [batch_size, 1, num_units], + ) x = torch.from_numpy(x).type(dtype).to(device) return x - def get_mask_chunk_predictor(self, chunk_outs=None, device='cpu', batch_size=1, num_units=1, idx=5, dtype=torch.float32): + def get_mask_chunk_predictor( + self, chunk_outs=None, device="cpu", batch_size=1, num_units=1, idx=5, dtype=torch.float32 + ): with torch.no_grad(): - x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx] - x = np.tile(x[None, :, :, ], [batch_size, 1, num_units]) + x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx] + x = np.tile( + x[ + None, + :, + :, + ], + [batch_size, 1, num_units], + ) x = torch.from_numpy(x).type(dtype).to(device) return x - def get_mask_att_chunk_encoder(self, chunk_outs=None, device='cpu', batch_size=1, idx=6, dtype=torch.float32): + def get_mask_att_chunk_encoder( + self, chunk_outs=None, device="cpu", batch_size=1, idx=6, dtype=torch.float32 + ): with torch.no_grad(): - x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx] - x = np.tile(x[None, :, :, ], [batch_size, 1, 1]) + x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx] + x = np.tile( + x[ + None, + :, + :, + ], + [batch_size, 1, 1], + ) x = torch.from_numpy(x).type(dtype).to(device) return x - def get_mask_shift_att_chunk_decoder(self, chunk_outs=None, device='cpu', batch_size=1, idx=7, dtype=torch.float32): + def get_mask_shift_att_chunk_decoder( + self, chunk_outs=None, device="cpu", batch_size=1, idx=7, dtype=torch.float32 + ): with torch.no_grad(): - x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx] + x = chunk_outs[idx] if chunk_outs is not None else self.chunk_outs[idx] x = np.tile(x[None, None, :, 0], [batch_size, 1, 1]) x = torch.from_numpy(x).type(dtype).to(device) return x - def build_scama_mask_for_cross_attention_decoder( - predictor_alignments: torch.Tensor, - encoder_sequence_length: torch.Tensor, - chunk_size: int = 5, - encoder_chunk_size: int = 5, - attention_chunk_center_bias: int = 0, - attention_chunk_size: int = 1, - attention_chunk_type: str = 'chunk', - step=None, - predictor_mask_chunk_hopping: torch.Tensor = None, - decoder_att_look_back_factor: int = 1, - mask_shift_att_chunk_decoder: torch.Tensor = None, - target_length: torch.Tensor = None, - is_training=True, - dtype: torch.dtype = torch.float32): + predictor_alignments: torch.Tensor, + encoder_sequence_length: torch.Tensor, + chunk_size: int = 5, + encoder_chunk_size: int = 5, + attention_chunk_center_bias: int = 0, + attention_chunk_size: int = 1, + attention_chunk_type: str = "chunk", + step=None, + predictor_mask_chunk_hopping: torch.Tensor = None, + decoder_att_look_back_factor: int = 1, + mask_shift_att_chunk_decoder: torch.Tensor = None, + target_length: torch.Tensor = None, + is_training=True, + dtype: torch.dtype = torch.float32, +): with torch.no_grad(): device = predictor_alignments.device batch_size, chunk_num = predictor_alignments.size() @@ -304,85 +392,102 @@ def build_scama_mask_for_cross_attention_decoder( target_length = predictor_alignments.sum(dim=-1).type(encoder_sequence_length.dtype) maximum_target_length = target_length.max() predictor_alignments_cumsum = torch.cumsum(predictor_alignments, dim=1) - predictor_alignments_cumsum = predictor_alignments_cumsum[:, None, :].repeat(1, maximum_target_length, 1) - - + predictor_alignments_cumsum = predictor_alignments_cumsum[:, None, :].repeat( + 1, maximum_target_length, 1 + ) + index = torch.ones([batch_size, maximum_target_length], dtype=int_type).to(device) index = torch.cumsum(index, dim=1) index = index[:, :, None].repeat(1, 1, chunk_num) - + index_div = torch.floor(torch.divide(predictor_alignments_cumsum, index)).type(int_type) index_div_bool_zeros = index_div == 0 index_div_bool_zeros_count = torch.sum(index_div_bool_zeros.type(int_type), dim=-1) + 1 - + index_div_bool_zeros_count = torch.clip(index_div_bool_zeros_count, min=1, max=chunk_num) - + index_div_bool_zeros_count *= chunk_size index_div_bool_zeros_count += attention_chunk_center_bias - index_div_bool_zeros_count = torch.clip(index_div_bool_zeros_count-1, min=0, max=maximum_encoder_length) + index_div_bool_zeros_count = torch.clip( + index_div_bool_zeros_count - 1, min=0, max=maximum_encoder_length + ) index_div_bool_zeros_count_ori = index_div_bool_zeros_count - - index_div_bool_zeros_count = (torch.floor(index_div_bool_zeros_count / encoder_chunk_size)+1)*encoder_chunk_size + + index_div_bool_zeros_count = ( + torch.floor(index_div_bool_zeros_count / encoder_chunk_size) + 1 + ) * encoder_chunk_size max_len_chunk = math.ceil(maximum_encoder_length / encoder_chunk_size) * encoder_chunk_size - + mask_flip, mask_flip2 = None, None if attention_chunk_size is not None: index_div_bool_zeros_count_beg = index_div_bool_zeros_count - attention_chunk_size - index_div_bool_zeros_count_beg = torch.clip(index_div_bool_zeros_count_beg, 0, max_len_chunk) - index_div_bool_zeros_count_beg_mask = sequence_mask(index_div_bool_zeros_count_beg, maxlen=max_len_chunk, dtype=int_type, device=device) + index_div_bool_zeros_count_beg = torch.clip( + index_div_bool_zeros_count_beg, 0, max_len_chunk + ) + index_div_bool_zeros_count_beg_mask = sequence_mask( + index_div_bool_zeros_count_beg, maxlen=max_len_chunk, dtype=int_type, device=device + ) mask_flip = 1 - index_div_bool_zeros_count_beg_mask - attention_chunk_size2 = attention_chunk_size * (decoder_att_look_back_factor+1) + attention_chunk_size2 = attention_chunk_size * (decoder_att_look_back_factor + 1) index_div_bool_zeros_count_beg = index_div_bool_zeros_count - attention_chunk_size2 - - index_div_bool_zeros_count_beg = torch.clip(index_div_bool_zeros_count_beg, 0, max_len_chunk) - index_div_bool_zeros_count_beg_mask = sequence_mask(index_div_bool_zeros_count_beg, maxlen=max_len_chunk, dtype=int_type, device=device) + + index_div_bool_zeros_count_beg = torch.clip( + index_div_bool_zeros_count_beg, 0, max_len_chunk + ) + index_div_bool_zeros_count_beg_mask = sequence_mask( + index_div_bool_zeros_count_beg, maxlen=max_len_chunk, dtype=int_type, device=device + ) mask_flip2 = 1 - index_div_bool_zeros_count_beg_mask - - mask = sequence_mask(index_div_bool_zeros_count, maxlen=max_len_chunk, dtype=dtype, device=device) - + + mask = sequence_mask( + index_div_bool_zeros_count, maxlen=max_len_chunk, dtype=dtype, device=device + ) + if predictor_mask_chunk_hopping is not None: - b, k, t = mask.size() - predictor_mask_chunk_hopping = predictor_mask_chunk_hopping[:, None, :, 0].repeat(1, k, 1) - - mask_mask_flip = mask - if mask_flip is not None: - mask_mask_flip = mask_flip * mask - - def _fn(): - mask_sliced = mask[:b, :k, encoder_chunk_size:t] - zero_pad_right = torch.zeros([b, k, encoder_chunk_size], dtype=mask_sliced.dtype).to(device) - mask_sliced = torch.cat([mask_sliced, zero_pad_right], dim=2) - _, _, tt = predictor_mask_chunk_hopping.size() - pad_right_p = max_len_chunk - tt - predictor_mask_chunk_hopping_pad = torch.nn.functional.pad(predictor_mask_chunk_hopping, [0, pad_right_p], "constant", 0) - masked = mask_sliced * predictor_mask_chunk_hopping_pad - - mask_true = mask_mask_flip + masked - return mask_true - - mask = _fn() if t > chunk_size else mask_mask_flip - - - + b, k, t = mask.size() + predictor_mask_chunk_hopping = predictor_mask_chunk_hopping[:, None, :, 0].repeat( + 1, k, 1 + ) + + mask_mask_flip = mask + if mask_flip is not None: + mask_mask_flip = mask_flip * mask + + def _fn(): + mask_sliced = mask[:b, :k, encoder_chunk_size:t] + zero_pad_right = torch.zeros( + [b, k, encoder_chunk_size], dtype=mask_sliced.dtype + ).to(device) + mask_sliced = torch.cat([mask_sliced, zero_pad_right], dim=2) + _, _, tt = predictor_mask_chunk_hopping.size() + pad_right_p = max_len_chunk - tt + predictor_mask_chunk_hopping_pad = torch.nn.functional.pad( + predictor_mask_chunk_hopping, [0, pad_right_p], "constant", 0 + ) + masked = mask_sliced * predictor_mask_chunk_hopping_pad + + mask_true = mask_mask_flip + masked + return mask_true + + mask = _fn() if t > chunk_size else mask_mask_flip + if mask_flip2 is not None: mask *= mask_flip2 - - mask_target = sequence_mask(target_length, maxlen=maximum_target_length, dtype=mask.dtype, device=device) + + mask_target = sequence_mask( + target_length, maxlen=maximum_target_length, dtype=mask.dtype, device=device + ) mask = mask[:, :maximum_target_length, :] * mask_target[:, :, None] - - - - mask_len = sequence_mask(encoder_sequence_length, maxlen=maximum_encoder_length, dtype=mask.dtype, device=device) + + mask_len = sequence_mask( + encoder_sequence_length, maxlen=maximum_encoder_length, dtype=mask.dtype, device=device + ) mask = mask[:, :, :maximum_encoder_length] * mask_len[:, None, :] - - - - - if attention_chunk_type == 'full': + + if attention_chunk_type == "full": mask = torch.ones_like(mask).to(device) if mask_shift_att_chunk_decoder is not None: mask = mask * mask_shift_att_chunk_decoder mask = mask[:, :maximum_target_length, :maximum_encoder_length].type(dtype).to(device) return mask - diff --git a/funasr/models/scama/decoder.py b/funasr/models/scama/decoder.py index 8257f59f6..31b235778 100644 --- a/funasr/models/scama/decoder.py +++ b/funasr/models/scama/decoder.py @@ -13,7 +13,10 @@ import numpy as np from funasr.models.scama import utils as myutils from funasr.models.transformer.decoder import BaseTransformerDecoder -from funasr.models.sanm.attention import MultiHeadedAttentionSANMDecoder, MultiHeadedAttentionCrossAtt +from funasr.models.sanm.attention import ( + MultiHeadedAttentionSANMDecoder, + MultiHeadedAttentionCrossAtt, +) from funasr.models.transformer.embedding import PositionalEncoding from funasr.models.transformer.layer_norm import LayerNorm from funasr.models.sanm.positionwise_feed_forward import PositionwiseFeedForwardDecoderSANM @@ -21,6 +24,7 @@ from funasr.models.transformer.utils.repeat import repeat from funasr.register import tables + class DecoderLayerSANM(nn.Module): """Single decoder layer module. @@ -151,10 +155,11 @@ class DecoderLayerSANM(nn.Module): x = residual + self.dropout(self.src_attn(x, memory, memory_mask)) - return x, tgt_mask, memory, memory_mask, cache - def forward_chunk(self, tgt, memory, fsmn_cache=None, opt_cache=None, chunk_size=None, look_back=0): + def forward_chunk( + self, tgt, memory, fsmn_cache=None, opt_cache=None, chunk_size=None, look_back=0 + ): """Compute decoded features. Args: @@ -194,6 +199,7 @@ class DecoderLayerSANM(nn.Module): return x, memory, fsmn_cache, opt_cache + @tables.register("decoder_classes", "FsmnDecoderSCAMAOpt") class FsmnDecoderSCAMAOpt(BaseTransformerDecoder): """ @@ -201,31 +207,31 @@ class FsmnDecoderSCAMAOpt(BaseTransformerDecoder): SCAMA: Streaming chunk-aware multihead attention for online end-to-end speech recognition https://arxiv.org/abs/2006.01712 """ - + def __init__( - self, - vocab_size: int, - encoder_output_size: int, - attention_heads: int = 4, - linear_units: int = 2048, - num_blocks: int = 6, - dropout_rate: float = 0.1, - positional_dropout_rate: float = 0.1, - self_attention_dropout_rate: float = 0.0, - src_attention_dropout_rate: float = 0.0, - input_layer: str = "embed", - use_output_layer: bool = True, - pos_enc_class=PositionalEncoding, - normalize_before: bool = True, - concat_after: bool = False, - att_layer_num: int = 6, - kernel_size: int = 21, - sanm_shfit: int = None, - concat_embeds: bool = False, - attention_dim: int = None, - tf2torch_tensor_name_prefix_torch: str = "decoder", - tf2torch_tensor_name_prefix_tf: str = "seq2seq/decoder", - embed_tensor_name_prefix_tf: str = None, + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + att_layer_num: int = 6, + kernel_size: int = 21, + sanm_shfit: int = None, + concat_embeds: bool = False, + attention_dim: int = None, + tf2torch_tensor_name_prefix_torch: str = "decoder", + tf2torch_tensor_name_prefix_tf: str = "seq2seq/decoder", + embed_tensor_name_prefix_tf: str = None, ): super().__init__( vocab_size=vocab_size, @@ -275,7 +281,10 @@ class FsmnDecoderSCAMAOpt(BaseTransformerDecoder): attention_dim, self_attention_dropout_rate, kernel_size, sanm_shfit=sanm_shfit ), MultiHeadedAttentionCrossAtt( - attention_heads, attention_dim, src_attention_dropout_rate, encoder_output_size=encoder_output_size + attention_heads, + attention_dim, + src_attention_dropout_rate, + encoder_output_size=encoder_output_size, ), PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate), dropout_rate, @@ -291,7 +300,10 @@ class FsmnDecoderSCAMAOpt(BaseTransformerDecoder): lambda lnum: DecoderLayerSANM( attention_dim, MultiHeadedAttentionSANMDecoder( - attention_dim, self_attention_dropout_rate, kernel_size, sanm_shfit=sanm_shfit + attention_dim, + self_attention_dropout_rate, + kernel_size, + sanm_shfit=sanm_shfit, ), None, PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate), @@ -320,8 +332,12 @@ class FsmnDecoderSCAMAOpt(BaseTransformerDecoder): attention_dim + encoder_output_size, None, None, - PositionwiseFeedForwardDecoderSANM(attention_dim + encoder_output_size, linear_units, dropout_rate, - adim=attention_dim), + PositionwiseFeedForwardDecoderSANM( + attention_dim + encoder_output_size, + linear_units, + dropout_rate, + adim=attention_dim, + ), dropout_rate, normalize_before, concat_after, @@ -335,13 +351,13 @@ class FsmnDecoderSCAMAOpt(BaseTransformerDecoder): self.embed_tensor_name_prefix_tf = embed_tensor_name_prefix_tf def forward( - self, - hs_pad: torch.Tensor, - hlens: torch.Tensor, - ys_in_pad: torch.Tensor, - ys_in_lens: torch.Tensor, - chunk_mask: torch.Tensor = None, - pre_acoustic_embeds: torch.Tensor = None, + self, + hs_pad: torch.Tensor, + hlens: torch.Tensor, + ys_in_pad: torch.Tensor, + ys_in_lens: torch.Tensor, + chunk_mask: torch.Tensor = None, + pre_acoustic_embeds: torch.Tensor = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """Forward decoder. @@ -376,16 +392,10 @@ class FsmnDecoderSCAMAOpt(BaseTransformerDecoder): x = torch.cat((x, pre_acoustic_embeds), dim=-1) x, _, _, _, _ = self.embed_concat_ffn(x, None, None, None, None) - x, tgt_mask, memory, memory_mask, _ = self.decoders( - x, tgt_mask, memory, memory_mask - ) + x, tgt_mask, memory, memory_mask, _ = self.decoders(x, tgt_mask, memory, memory_mask) if self.decoders2 is not None: - x, tgt_mask, memory, memory_mask, _ = self.decoders2( - x, tgt_mask, memory, memory_mask - ) - x, tgt_mask, memory, memory_mask, _ = self.decoders3( - x, tgt_mask, memory, memory_mask - ) + x, tgt_mask, memory, memory_mask, _ = self.decoders2(x, tgt_mask, memory, memory_mask) + x, tgt_mask, memory, memory_mask, _ = self.decoders3(x, tgt_mask, memory, memory_mask) if self.normalize_before: x = self.after_norm(x) if self.output_layer is not None: @@ -394,23 +404,36 @@ class FsmnDecoderSCAMAOpt(BaseTransformerDecoder): olens = tgt_mask.sum(1) return x, olens - def score(self, ys, state, x, x_mask=None, pre_acoustic_embeds: torch.Tensor = None, ): + def score( + self, + ys, + state, + x, + x_mask=None, + pre_acoustic_embeds: torch.Tensor = None, + ): """Score.""" - ys_mask = myutils.sequence_mask(torch.tensor([len(ys)], dtype=torch.int32), device=x.device)[:, :, None] + ys_mask = myutils.sequence_mask( + torch.tensor([len(ys)], dtype=torch.int32), device=x.device + )[:, :, None] logp, state = self.forward_one_step( - ys.unsqueeze(0), ys_mask, x.unsqueeze(0), memory_mask=x_mask, pre_acoustic_embeds=pre_acoustic_embeds, - cache=state + ys.unsqueeze(0), + ys_mask, + x.unsqueeze(0), + memory_mask=x_mask, + pre_acoustic_embeds=pre_acoustic_embeds, + cache=state, ) return logp.squeeze(0), state def forward_one_step( - self, - tgt: torch.Tensor, - tgt_mask: torch.Tensor, - memory: torch.Tensor, - memory_mask: torch.Tensor = None, - pre_acoustic_embeds: torch.Tensor = None, - cache: List[torch.Tensor] = None, + self, + tgt: torch.Tensor, + tgt_mask: torch.Tensor, + memory: torch.Tensor, + memory_mask: torch.Tensor = None, + pre_acoustic_embeds: torch.Tensor = None, + cache: List[torch.Tensor] = None, ) -> Tuple[torch.Tensor, List[torch.Tensor]]: """Forward one step. @@ -473,4 +496,3 @@ class FsmnDecoderSCAMAOpt(BaseTransformerDecoder): y = torch.log_softmax(y, dim=-1) return y, new_cache - diff --git a/funasr/models/scama/encoder.py b/funasr/models/scama/encoder.py index 2c676b23b..e1fe9242e 100644 --- a/funasr/models/scama/encoder.py +++ b/funasr/models/scama/encoder.py @@ -17,7 +17,10 @@ import numpy as np from funasr.train_utils.device_funcs import to_device from funasr.models.transformer.utils.nets_utils import make_pad_mask from funasr.models.sanm.attention import MultiHeadedAttention, MultiHeadedAttentionSANM -from funasr.models.transformer.embedding import SinusoidalPositionEncoder, StreamSinusoidalPositionEncoder +from funasr.models.transformer.embedding import ( + SinusoidalPositionEncoder, + StreamSinusoidalPositionEncoder, +) from funasr.models.transformer.layer_norm import LayerNorm from funasr.models.transformer.utils.multi_layer_conv import Conv1dLinear from funasr.models.transformer.utils.multi_layer_conv import MultiLayeredConv1d @@ -37,6 +40,7 @@ from funasr.models.ctc.ctc import CTC from funasr.register import tables + class EncoderLayerSANM(nn.Module): def __init__( self, @@ -96,7 +100,18 @@ class EncoderLayerSANM(nn.Module): x = self.norm1(x) if self.concat_after: - x_concat = torch.cat((x, self.self_attn(x, mask, mask_shfit_chunk=mask_shfit_chunk, mask_att_chunk_encoder=mask_att_chunk_encoder)), dim=-1) + x_concat = torch.cat( + ( + x, + self.self_attn( + x, + mask, + mask_shfit_chunk=mask_shfit_chunk, + mask_att_chunk_encoder=mask_att_chunk_encoder, + ), + ), + dim=-1, + ) if self.in_size == self.size: x = residual + stoch_layer_coeff * self.concat_linear(x_concat) else: @@ -104,11 +119,21 @@ class EncoderLayerSANM(nn.Module): else: if self.in_size == self.size: x = residual + stoch_layer_coeff * self.dropout( - self.self_attn(x, mask, mask_shfit_chunk=mask_shfit_chunk, mask_att_chunk_encoder=mask_att_chunk_encoder) + self.self_attn( + x, + mask, + mask_shfit_chunk=mask_shfit_chunk, + mask_att_chunk_encoder=mask_att_chunk_encoder, + ) ) else: x = stoch_layer_coeff * self.dropout( - self.self_attn(x, mask, mask_shfit_chunk=mask_shfit_chunk, mask_att_chunk_encoder=mask_att_chunk_encoder) + self.self_attn( + x, + mask, + mask_shfit_chunk=mask_shfit_chunk, + mask_att_chunk_encoder=mask_att_chunk_encoder, + ) ) if not self.normalize_before: x = self.norm1(x) @@ -168,34 +193,34 @@ class SANMEncoderChunkOpt(nn.Module): """ def __init__( - self, - input_size: int, - output_size: int = 256, - attention_heads: int = 4, - linear_units: int = 2048, - num_blocks: int = 6, - dropout_rate: float = 0.1, - positional_dropout_rate: float = 0.1, - attention_dropout_rate: float = 0.0, - input_layer: Optional[str] = "conv2d", - pos_enc_class=SinusoidalPositionEncoder, - normalize_before: bool = True, - concat_after: bool = False, - positionwise_layer_type: str = "linear", - positionwise_conv_kernel_size: int = 1, - padding_idx: int = -1, - interctc_layer_idx: List[int] = [], - interctc_use_conditioning: bool = False, - kernel_size: int = 11, - sanm_shfit: int = 0, - selfattention_layer_type: str = "sanm", - chunk_size: Union[int, Sequence[int]] = (16,), - stride: Union[int, Sequence[int]] = (10,), - pad_left: Union[int, Sequence[int]] = (0,), - encoder_att_look_back_factor: Union[int, Sequence[int]] = (1,), - decoder_att_look_back_factor: Union[int, Sequence[int]] = (1,), - tf2torch_tensor_name_prefix_torch: str = "encoder", - tf2torch_tensor_name_prefix_tf: str = "seq2seq/encoder", + self, + input_size: int, + output_size: int = 256, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + attention_dropout_rate: float = 0.0, + input_layer: Optional[str] = "conv2d", + pos_enc_class=SinusoidalPositionEncoder, + normalize_before: bool = True, + concat_after: bool = False, + positionwise_layer_type: str = "linear", + positionwise_conv_kernel_size: int = 1, + padding_idx: int = -1, + interctc_layer_idx: List[int] = [], + interctc_use_conditioning: bool = False, + kernel_size: int = 11, + sanm_shfit: int = 0, + selfattention_layer_type: str = "sanm", + chunk_size: Union[int, Sequence[int]] = (16,), + stride: Union[int, Sequence[int]] = (10,), + pad_left: Union[int, Sequence[int]] = (0,), + encoder_att_look_back_factor: Union[int, Sequence[int]] = (1,), + decoder_att_look_back_factor: Union[int, Sequence[int]] = (1,), + tf2torch_tensor_name_prefix_torch: str = "encoder", + tf2torch_tensor_name_prefix_tf: str = "seq2seq/encoder", ): super().__init__() self._output_size = output_size @@ -334,12 +359,12 @@ class SANMEncoderChunkOpt(nn.Module): return self._output_size def forward( - self, - xs_pad: torch.Tensor, - ilens: torch.Tensor, - prev_states: torch.Tensor = None, - ctc: CTC = None, - ind: int = 0, + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, + prev_states: torch.Tensor = None, + ctc: CTC = None, + ind: int = 0, ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: """Embed positions in tensor. @@ -355,10 +380,10 @@ class SANMEncoderChunkOpt(nn.Module): if self.embed is None: xs_pad = xs_pad elif ( - isinstance(self.embed, Conv2dSubsampling) - or isinstance(self.embed, Conv2dSubsampling2) - or isinstance(self.embed, Conv2dSubsampling6) - or isinstance(self.embed, Conv2dSubsampling8) + isinstance(self.embed, Conv2dSubsampling) + or isinstance(self.embed, Conv2dSubsampling2) + or isinstance(self.embed, Conv2dSubsampling6) + or isinstance(self.embed, Conv2dSubsampling8) ): short_status, limit_size = check_short_utt(self.embed, xs_pad.size(1)) if short_status: @@ -378,21 +403,26 @@ class SANMEncoderChunkOpt(nn.Module): chunk_outs = self.overlap_chunk_cls.gen_chunk_mask(ilens, ind) xs_pad, ilens = self.overlap_chunk_cls.split_chunk(xs_pad, ilens, chunk_outs=chunk_outs) masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device) - mask_shfit_chunk = self.overlap_chunk_cls.get_mask_shfit_chunk(chunk_outs, xs_pad.device, xs_pad.size(0), - dtype=xs_pad.dtype) - mask_att_chunk_encoder = self.overlap_chunk_cls.get_mask_att_chunk_encoder(chunk_outs, xs_pad.device, - xs_pad.size(0), - dtype=xs_pad.dtype) + mask_shfit_chunk = self.overlap_chunk_cls.get_mask_shfit_chunk( + chunk_outs, xs_pad.device, xs_pad.size(0), dtype=xs_pad.dtype + ) + mask_att_chunk_encoder = self.overlap_chunk_cls.get_mask_att_chunk_encoder( + chunk_outs, xs_pad.device, xs_pad.size(0), dtype=xs_pad.dtype + ) encoder_outs = self.encoders0(xs_pad, masks, None, mask_shfit_chunk, mask_att_chunk_encoder) xs_pad, masks = encoder_outs[0], encoder_outs[1] intermediate_outs = [] if len(self.interctc_layer_idx) == 0: - encoder_outs = self.encoders(xs_pad, masks, None, mask_shfit_chunk, mask_att_chunk_encoder) + encoder_outs = self.encoders( + xs_pad, masks, None, mask_shfit_chunk, mask_att_chunk_encoder + ) xs_pad, masks = encoder_outs[0], encoder_outs[1] else: for layer_idx, encoder_layer in enumerate(self.encoders): - encoder_outs = encoder_layer(xs_pad, masks, None, mask_shfit_chunk, mask_att_chunk_encoder) + encoder_outs = encoder_layer( + xs_pad, masks, None, mask_shfit_chunk, mask_att_chunk_encoder + ) xs_pad, masks = encoder_outs[0], encoder_outs[1] if layer_idx + 1 in self.interctc_layer_idx: encoder_out = xs_pad @@ -420,15 +450,16 @@ class SANMEncoderChunkOpt(nn.Module): return feats cache["feats"] = to_device(cache["feats"], device=feats.device) overlap_feats = torch.cat((cache["feats"], feats), dim=1) - cache["feats"] = overlap_feats[:, -(cache["chunk_size"][0] + cache["chunk_size"][2]):, :] + cache["feats"] = overlap_feats[:, -(cache["chunk_size"][0] + cache["chunk_size"][2]) :, :] return overlap_feats - def forward_chunk(self, - xs_pad: torch.Tensor, - ilens: torch.Tensor, - cache: dict = None, - **kwargs, - ): + def forward_chunk( + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, + cache: dict = None, + **kwargs, + ): is_final = kwargs.get("is_final", False) xs_pad *= self.output_size() ** 0.5 if self.embed is None: @@ -446,12 +477,19 @@ class SANMEncoderChunkOpt(nn.Module): new_cache = cache["opt"] for layer_idx, encoder_layer in enumerate(self.encoders0): - encoder_outs = encoder_layer.forward_chunk(xs_pad, new_cache[layer_idx], cache["chunk_size"], cache["encoder_chunk_look_back"]) + encoder_outs = encoder_layer.forward_chunk( + xs_pad, new_cache[layer_idx], cache["chunk_size"], cache["encoder_chunk_look_back"] + ) xs_pad, new_cache[0] = encoder_outs[0], encoder_outs[1] for layer_idx, encoder_layer in enumerate(self.encoders): - encoder_outs = encoder_layer.forward_chunk(xs_pad, new_cache[layer_idx+len(self.encoders0)], cache["chunk_size"], cache["encoder_chunk_look_back"]) - xs_pad, new_cache[layer_idx+len(self.encoders0)] = encoder_outs[0], encoder_outs[1] + encoder_outs = encoder_layer.forward_chunk( + xs_pad, + new_cache[layer_idx + len(self.encoders0)], + cache["chunk_size"], + cache["encoder_chunk_look_back"], + ) + xs_pad, new_cache[layer_idx + len(self.encoders0)] = encoder_outs[0], encoder_outs[1] if self.normalize_before: xs_pad = self.after_norm(xs_pad) @@ -459,4 +497,3 @@ class SANMEncoderChunkOpt(nn.Module): cache["opt"] = new_cache return xs_pad, ilens, None - diff --git a/funasr/models/scama/model.py b/funasr/models/scama/model.py index 32e16bded..c15f435a4 100644 --- a/funasr/models/scama/model.py +++ b/funasr/models/scama/model.py @@ -35,6 +35,7 @@ else: def autocast(enabled=True): yield + @tables.register("model_classes", "SCAMA") class SCAMA(nn.Module): """ @@ -77,11 +78,11 @@ class SCAMA(nn.Module): if specaug is not None: specaug_class = tables.specaug_classes.get(specaug) specaug = specaug_class(**specaug_conf) - + if normalize is not None: normalize_class = tables.normalize_classes.get(normalize) normalize = normalize_class(**normalize_conf) - + encoder_class = tables.encoder_classes.get(encoder) encoder = encoder_class(input_size=input_size, **encoder_conf) encoder_output_size = encoder.output_size() @@ -93,13 +94,11 @@ class SCAMA(nn.Module): **decoder_conf, ) if ctc_weight > 0.0: - + if ctc_conf is None: ctc_conf = {} - - ctc = CTC( - odim=vocab_size, encoder_output_size=encoder_output_size, **ctc_conf - ) + + ctc = CTC(odim=vocab_size, encoder_output_size=encoder_output_size, **ctc_conf) predictor_class = tables.predictor_classes.get(predictor) predictor = predictor_class(**predictor_conf) @@ -111,12 +110,11 @@ class SCAMA(nn.Module): self.vocab_size = vocab_size self.ignore_id = ignore_id self.ctc_weight = ctc_weight - + self.specaug = specaug self.normalize = normalize - - self.encoder = encoder + self.encoder = encoder if ctc_weight == 1.0: self.decoder = None @@ -134,7 +132,7 @@ class SCAMA(nn.Module): self.ctc = None else: self.ctc = ctc - + self.predictor = predictor self.predictor_weight = predictor_weight self.predictor_bias = predictor_bias @@ -148,10 +146,15 @@ class SCAMA(nn.Module): self.length_normalized_loss = length_normalized_loss self.beam_search = None self.error_calculator = None - + if self.encoder.overlap_chunk_cls is not None: - from funasr.models.scama.chunk_utilis import build_scama_mask_for_cross_attention_decoder - self.build_scama_mask_for_cross_attention_decoder_fn = build_scama_mask_for_cross_attention_decoder + from funasr.models.scama.chunk_utilis import ( + build_scama_mask_for_cross_attention_decoder, + ) + + self.build_scama_mask_for_cross_attention_decoder_fn = ( + build_scama_mask_for_cross_attention_decoder + ) self.decoder_attention_chunk_type = kwargs.get("decoder_attention_chunk_type", "chunk") def forward( @@ -175,55 +178,56 @@ class SCAMA(nn.Module): text_lengths = text_lengths[:, 0] if len(speech_lengths.size()) > 1: speech_lengths = speech_lengths[:, 0] - + batch_size = speech.shape[0] - + # Encoder ind = self.encoder.overlap_chunk_cls.random_choice(self.training, decoding_ind) encoder_out, encoder_out_lens = self.encode(speech, speech_lengths, ind=ind) - loss_ctc, cer_ctc = None, None loss_pre = None stats = dict() - + # decoder: CTC branch - + if self.ctc_weight > 0.0: - encoder_out_ctc, encoder_out_lens_ctc = self.encoder.overlap_chunk_cls.remove_chunk(encoder_out, - encoder_out_lens, - chunk_outs=None) + encoder_out_ctc, encoder_out_lens_ctc = self.encoder.overlap_chunk_cls.remove_chunk( + encoder_out, encoder_out_lens, chunk_outs=None + ) - loss_ctc, cer_ctc = self._calc_ctc_loss( encoder_out_ctc, encoder_out_lens_ctc, text, text_lengths ) # Collect CTC branch stats stats["loss_ctc"] = loss_ctc.detach() if loss_ctc is not None else None stats["cer_ctc"] = cer_ctc - + # decoder: Attention decoder branch loss_att, acc_att, cer_att, wer_att, loss_pre = self._calc_att_predictor_loss( encoder_out, encoder_out_lens, text, text_lengths ) - + # 3. CTC-Att loss definition if self.ctc_weight == 0.0: loss = loss_att + loss_pre * self.predictor_weight else: - loss = self.ctc_weight * loss_ctc + ( - 1 - self.ctc_weight) * loss_att + loss_pre * self.predictor_weight - + loss = ( + self.ctc_weight * loss_ctc + + (1 - self.ctc_weight) * loss_att + + loss_pre * self.predictor_weight + ) + # Collect Attn branch stats stats["loss_att"] = loss_att.detach() if loss_att is not None else None stats["acc"] = acc_att stats["cer"] = cer_att stats["wer"] = wer_att stats["loss_pre"] = loss_pre.detach().cpu() if loss_pre is not None else None - + stats["loss"] = torch.clone(loss.detach()) - + # force_gatherable: to-device and to-tensor if scalar for DataParallel if self.length_normalized_loss: batch_size = (text_lengths + self.predictor_bias).sum() @@ -231,7 +235,10 @@ class SCAMA(nn.Module): return loss, stats, weight def encode( - self, speech: torch.Tensor, speech_lengths: torch.Tensor, **kwargs, + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + **kwargs, ) -> Tuple[torch.Tensor, torch.Tensor]: """Encoder. Note that this method is used by asr_inference.py Args: @@ -240,24 +247,28 @@ class SCAMA(nn.Module): ind: int """ with autocast(False): - + # Data augmentation if self.specaug is not None and self.training: speech, speech_lengths = self.specaug(speech, speech_lengths) - + # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN if self.normalize is not None: speech, speech_lengths = self.normalize(speech, speech_lengths) - + # Forward encoder encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths) if isinstance(encoder_out, tuple): encoder_out = encoder_out[0] - + return encoder_out, encoder_out_lens def encode_chunk( - self, speech: torch.Tensor, speech_lengths: torch.Tensor, cache: dict = None, **kwargs, + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + cache: dict = None, + **kwargs, ) -> Tuple[torch.Tensor, torch.Tensor]: """Frontend + Encoder. Note that this method is used by asr_inference.py Args: @@ -266,20 +277,22 @@ class SCAMA(nn.Module): ind: int """ with autocast(False): - + # Data augmentation if self.specaug is not None and self.training: speech, speech_lengths = self.specaug(speech, speech_lengths) - + # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN if self.normalize is not None: speech, speech_lengths = self.normalize(speech, speech_lengths) - + # Forward encoder - encoder_out, encoder_out_lens, _ = self.encoder.forward_chunk(speech, speech_lengths, cache=cache["encoder"]) + encoder_out, encoder_out_lens, _ = self.encoder.forward_chunk( + speech, speech_lengths, cache=cache["encoder"] + ) if isinstance(encoder_out, tuple): encoder_out = encoder_out[0] - + return encoder_out, torch.tensor([encoder_out.size(1)]) def calc_predictor_chunk(self, encoder_out, encoder_out_lens, cache=None, **kwargs): @@ -297,36 +310,44 @@ class SCAMA(nn.Module): ys_in_pad, ys_out_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id) ys_in_lens = ys_pad_lens + 1 - encoder_out_mask = sequence_mask(encoder_out_lens, maxlen=encoder_out.size(1), dtype=encoder_out.dtype, - device=encoder_out.device)[:, None, :] + encoder_out_mask = sequence_mask( + encoder_out_lens, + maxlen=encoder_out.size(1), + dtype=encoder_out.dtype, + device=encoder_out.device, + )[:, None, :] mask_chunk_predictor = None if self.encoder.overlap_chunk_cls is not None: - mask_chunk_predictor = self.encoder.overlap_chunk_cls.get_mask_chunk_predictor(None, - device=encoder_out.device, - batch_size=encoder_out.size( - 0)) - mask_shfit_chunk = self.encoder.overlap_chunk_cls.get_mask_shfit_chunk(None, device=encoder_out.device, - batch_size=encoder_out.size(0)) + mask_chunk_predictor = self.encoder.overlap_chunk_cls.get_mask_chunk_predictor( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + mask_shfit_chunk = self.encoder.overlap_chunk_cls.get_mask_shfit_chunk( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) encoder_out = encoder_out * mask_shfit_chunk - pre_acoustic_embeds, pre_token_length, pre_alphas, _ = self.predictor(encoder_out, - ys_out_pad, - encoder_out_mask, - ignore_id=self.ignore_id, - mask_chunk_predictor=mask_chunk_predictor, - target_label_length=ys_in_lens, - ) - predictor_alignments, predictor_alignments_len = self.predictor.gen_frame_alignments(pre_alphas, - encoder_out_lens) - + pre_acoustic_embeds, pre_token_length, pre_alphas, _ = self.predictor( + encoder_out, + ys_out_pad, + encoder_out_mask, + ignore_id=self.ignore_id, + mask_chunk_predictor=mask_chunk_predictor, + target_label_length=ys_in_lens, + ) + predictor_alignments, predictor_alignments_len = self.predictor.gen_frame_alignments( + pre_alphas, encoder_out_lens + ) encoder_chunk_size = self.encoder.overlap_chunk_cls.chunk_size_pad_shift_cur attention_chunk_center_bias = 0 attention_chunk_size = encoder_chunk_size - decoder_att_look_back_factor = self.encoder.overlap_chunk_cls.decoder_att_look_back_factor_cur - mask_shift_att_chunk_decoder = self.encoder.overlap_chunk_cls.get_mask_shift_att_chunk_decoder(None, - device=encoder_out.device, - batch_size=encoder_out.size( - 0)) + decoder_att_look_back_factor = ( + self.encoder.overlap_chunk_cls.decoder_att_look_back_factor_cur + ) + mask_shift_att_chunk_decoder = ( + self.encoder.overlap_chunk_cls.get_mask_shift_att_chunk_decoder( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + ) scama_mask = self.build_scama_mask_for_cross_attention_decoder_fn( predictor_alignments=predictor_alignments, encoder_sequence_length=encoder_out_lens, @@ -343,7 +364,6 @@ class SCAMA(nn.Module): is_training=self.training, ) - # try: # 1. Forward decoder decoder_out, _ = self.decoder( @@ -353,7 +373,6 @@ class SCAMA(nn.Module): ys_in_lens, chunk_mask=scama_mask, pre_acoustic_embeds=pre_acoustic_embeds, - ) # 2. Compute attention loss @@ -385,36 +404,44 @@ class SCAMA(nn.Module): # ys_in_lens = ys_pad_lens + 1 ys_out_pad, ys_in_lens = None, None - encoder_out_mask = sequence_mask(encoder_out_lens, maxlen=encoder_out.size(1), dtype=encoder_out.dtype, - device=encoder_out.device)[:, None, :] + encoder_out_mask = sequence_mask( + encoder_out_lens, + maxlen=encoder_out.size(1), + dtype=encoder_out.dtype, + device=encoder_out.device, + )[:, None, :] mask_chunk_predictor = None - mask_chunk_predictor = self.encoder.overlap_chunk_cls.get_mask_chunk_predictor(None, - device=encoder_out.device, - batch_size=encoder_out.size( - 0)) - mask_shfit_chunk = self.encoder.overlap_chunk_cls.get_mask_shfit_chunk(None, device=encoder_out.device, - batch_size=encoder_out.size(0)) + mask_chunk_predictor = self.encoder.overlap_chunk_cls.get_mask_chunk_predictor( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + mask_shfit_chunk = self.encoder.overlap_chunk_cls.get_mask_shfit_chunk( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) encoder_out = encoder_out * mask_shfit_chunk - pre_acoustic_embeds, pre_token_length, pre_alphas, _ = self.predictor(encoder_out, - ys_out_pad, - encoder_out_mask, - ignore_id=self.ignore_id, - mask_chunk_predictor=mask_chunk_predictor, - target_label_length=ys_in_lens, - ) - predictor_alignments, predictor_alignments_len = self.predictor.gen_frame_alignments(pre_alphas, - encoder_out_lens) - + pre_acoustic_embeds, pre_token_length, pre_alphas, _ = self.predictor( + encoder_out, + ys_out_pad, + encoder_out_mask, + ignore_id=self.ignore_id, + mask_chunk_predictor=mask_chunk_predictor, + target_label_length=ys_in_lens, + ) + predictor_alignments, predictor_alignments_len = self.predictor.gen_frame_alignments( + pre_alphas, encoder_out_lens + ) encoder_chunk_size = self.encoder.overlap_chunk_cls.chunk_size_pad_shift_cur attention_chunk_center_bias = 0 attention_chunk_size = encoder_chunk_size - decoder_att_look_back_factor = self.encoder.overlap_chunk_cls.decoder_att_look_back_factor_cur - mask_shift_att_chunk_decoder = self.encoder.overlap_chunk_cls.get_mask_shift_att_chunk_decoder(None, - device=encoder_out.device, - batch_size=encoder_out.size( - 0)) + decoder_att_look_back_factor = ( + self.encoder.overlap_chunk_cls.decoder_att_look_back_factor_cur + ) + mask_shift_att_chunk_decoder = ( + self.encoder.overlap_chunk_cls.get_mask_shift_att_chunk_decoder( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + ) scama_mask = self.build_scama_mask_for_cross_attention_decoder_fn( predictor_alignments=predictor_alignments, encoder_sequence_length=encoder_out_lens, @@ -431,37 +458,41 @@ class SCAMA(nn.Module): is_training=self.training, ) - return pre_acoustic_embeds, pre_token_length, predictor_alignments, predictor_alignments_len, scama_mask + return ( + pre_acoustic_embeds, + pre_token_length, + predictor_alignments, + predictor_alignments_len, + scama_mask, + ) - def init_beam_search(self, - **kwargs, - ): + def init_beam_search( + self, + **kwargs, + ): from funasr.models.scama.beam_search import BeamSearchScamaStreaming - from funasr.models.transformer.scorers.ctc import CTCPrefixScorer from funasr.models.transformer.scorers.length_bonus import LengthBonus - + # 1. Build ASR model scorers = {} - + if self.ctc != None: ctc = CTCPrefixScorer(ctc=self.ctc, eos=self.eos) - scorers.update( - ctc=ctc - ) + scorers.update(ctc=ctc) token_list = kwargs.get("token_list") scorers.update( decoder=self.decoder, length_bonus=LengthBonus(len(token_list)), ) - + # 3. Build ngram model # ngram is not supported now ngram = None scorers["ngram"] = ngram - + weights = dict( decoder=1.0 - kwargs.get("decoding_ctc_weight", 0.0), ctc=kwargs.get("decoding_ctc_weight", 0.0), @@ -469,7 +500,7 @@ class SCAMA(nn.Module): ngram=kwargs.get("ngram_weight", 0.0), length_bonus=kwargs.get("penalty", 0.0), ) - + beam_search = BeamSearchScamaStreaming( beam_size=kwargs.get("beam_size", 2), weights=weights, @@ -486,38 +517,43 @@ class SCAMA(nn.Module): # scorer.to(device=kwargs.get("device", "cpu"), dtype=getattr(torch, kwargs.get("dtype", "float32"))).eval() self.beam_search = beam_search - def generate_chunk(self, - speech, - speech_lengths=None, - key: list = None, - tokenizer=None, - frontend=None, - **kwargs, - ): + def generate_chunk( + self, + speech, + speech_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): cache = kwargs.get("cache", {}) speech = speech.to(device=kwargs["device"]) speech_lengths = speech_lengths.to(device=kwargs["device"]) - + # Encoder - encoder_out, encoder_out_lens = self.encode_chunk(speech, speech_lengths, cache=cache, - is_final=kwargs.get("is_final", False)) + encoder_out, encoder_out_lens = self.encode_chunk( + speech, speech_lengths, cache=cache, is_final=kwargs.get("is_final", False) + ) if isinstance(encoder_out, tuple): encoder_out = encoder_out[0] if "running_hyps" not in cache: running_hyps = self.beam_search.init_hyp(encoder_out) cache["running_hyps"] = running_hyps - - - # predictor - predictor_outs = self.calc_predictor_chunk(encoder_out, - encoder_out_lens, - cache=cache, - is_final=kwargs.get("is_final", False), - ) - pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = predictor_outs[0], predictor_outs[1], \ - predictor_outs[2], predictor_outs[3] - pre_token_length = pre_token_length.round().long() + # predictor + predictor_outs = self.calc_predictor_chunk( + encoder_out, + encoder_out_lens, + cache=cache, + is_final=kwargs.get("is_final", False), + ) + pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = ( + predictor_outs[0], + predictor_outs[1], + predictor_outs[2], + predictor_outs[3], + ) + pre_token_length = pre_token_length.round().long() if torch.max(pre_token_length) < 1: return [] @@ -527,7 +563,12 @@ class SCAMA(nn.Module): minlen = max(0, minlen - kwargs.get("token_num_relax", 5)) # c. Passed the encoder result and the beam search nbest_hyps = self.beam_search( - x=encoder_out[0], scama_mask=None, pre_acoustic_embeds=pre_acoustic_embeds, maxlen=int(maxlen), minlen=int(minlen), cache=cache, + x=encoder_out[0], + scama_mask=None, + pre_acoustic_embeds=pre_acoustic_embeds, + maxlen=int(maxlen), + minlen=int(minlen), + cache=cache, ) cache["running_hyps"] = nbest_hyps @@ -544,125 +585,155 @@ class SCAMA(nn.Module): else: token_int = hyp.yseq[1:last_pos].tolist() - # remove blank symbol id, which is assumed to be 0 - token_int = list(filter(lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int)) - + token_int = list( + filter( + lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int + ) + ) + # Change integer-ids to tokens token = tokenizer.ids2tokens(token_int) # text = tokenizer.tokens2text(token) - + result_i = token - + results.extend(result_i) - + return results def init_cache(self, cache: dict = {}, **kwargs): device = kwargs.get("device", "cuda") - + chunk_size = kwargs.get("chunk_size", [0, 10, 5]) encoder_chunk_look_back = kwargs.get("encoder_chunk_look_back", 0) decoder_chunk_look_back = kwargs.get("decoder_chunk_look_back", 0) batch_size = 1 - + enc_output_size = kwargs["encoder_conf"]["output_size"] feats_dims = kwargs["frontend_conf"]["n_mels"] * kwargs["frontend_conf"]["lfr_m"] - cache_encoder = {"start_idx": 0, "cif_hidden": torch.zeros((batch_size, 1, enc_output_size)).to(device=device), - "cif_alphas": torch.zeros((batch_size, 1)).to(device=device), "chunk_size": chunk_size, - "encoder_chunk_look_back": encoder_chunk_look_back, "last_chunk": False, "opt": None, - "feats": torch.zeros((batch_size, chunk_size[0] + chunk_size[2], feats_dims)).to(device=device), - "tail_chunk": False} + cache_encoder = { + "start_idx": 0, + "cif_hidden": torch.zeros((batch_size, 1, enc_output_size)).to(device=device), + "cif_alphas": torch.zeros((batch_size, 1)).to(device=device), + "chunk_size": chunk_size, + "encoder_chunk_look_back": encoder_chunk_look_back, + "last_chunk": False, + "opt": None, + "feats": torch.zeros((batch_size, chunk_size[0] + chunk_size[2], feats_dims)).to( + device=device + ), + "tail_chunk": False, + } cache["encoder"] = cache_encoder - - cache_decoder = {"decode_fsmn": None, "decoder_chunk_look_back": decoder_chunk_look_back, "opt": None, - "chunk_size": chunk_size} + + cache_decoder = { + "decode_fsmn": None, + "decoder_chunk_look_back": decoder_chunk_look_back, + "opt": None, + "chunk_size": chunk_size, + } cache["decoder"] = cache_decoder cache["frontend"] = {} - cache["prev_samples"] = torch.empty(0).to(device=device) return cache - def inference(self, - data_in, - data_lengths=None, - key: list = None, - tokenizer=None, - frontend=None, - cache: dict = {}, - **kwargs, - ): - + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + cache: dict = {}, + **kwargs, + ): + # init beamsearch is_use_ctc = kwargs.get("decoding_ctc_weight", 0.0) > 0.00001 and self.ctc != None - is_use_lm = kwargs.get("lm_weight", 0.0) > 0.00001 and kwargs.get("lm_file", None) is not None + is_use_lm = ( + kwargs.get("lm_weight", 0.0) > 0.00001 and kwargs.get("lm_file", None) is not None + ) if self.beam_search is None: - logging.info("enable beam_search") self.init_beam_search(**kwargs) self.nbest = kwargs.get("nbest", 1) - + if len(cache) == 0: self.init_cache(cache, **kwargs) - + meta_data = {} chunk_size = kwargs.get("chunk_size", [0, 10, 5]) chunk_stride_samples = int(chunk_size[1] * 960) # 600ms - + time1 = time.perf_counter() cfg = {"is_final": kwargs.get("is_final", False)} - audio_sample_list = load_audio_text_image_video(data_in, - fs=frontend.fs, - audio_fs=kwargs.get("fs", 16000), - data_type=kwargs.get("data_type", "sound"), - tokenizer=tokenizer, - cache=cfg, - ) + audio_sample_list = load_audio_text_image_video( + data_in, + fs=frontend.fs, + audio_fs=kwargs.get("fs", 16000), + data_type=kwargs.get("data_type", "sound"), + tokenizer=tokenizer, + cache=cfg, + ) _is_final = cfg["is_final"] # if data_in is a file or url, set is_final=True - + time2 = time.perf_counter() meta_data["load_data"] = f"{time2 - time1:0.3f}" assert len(audio_sample_list) == 1, "batch_size must be set 1" - + audio_sample = torch.cat((cache["prev_samples"], audio_sample_list[0])) - + n = int(len(audio_sample) // chunk_stride_samples + int(_is_final)) m = int(len(audio_sample) % chunk_stride_samples * (1 - int(_is_final))) tokens = [] for i in range(n): kwargs["is_final"] = _is_final and i == n - 1 - audio_sample_i = audio_sample[i * chunk_stride_samples:(i + 1) * chunk_stride_samples] - + audio_sample_i = audio_sample[i * chunk_stride_samples : (i + 1) * chunk_stride_samples] + # extract fbank feats - speech, speech_lengths = extract_fbank([audio_sample_i], data_type=kwargs.get("data_type", "sound"), - frontend=frontend, cache=cache["frontend"], - is_final=kwargs["is_final"]) + speech, speech_lengths = extract_fbank( + [audio_sample_i], + data_type=kwargs.get("data_type", "sound"), + frontend=frontend, + cache=cache["frontend"], + is_final=kwargs["is_final"], + ) time3 = time.perf_counter() meta_data["extract_feat"] = f"{time3 - time2:0.3f}" - meta_data["batch_data_time"] = speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 - - tokens_i = self.generate_chunk(speech, speech_lengths, key=key, tokenizer=tokenizer, cache=cache, - frontend=frontend, **kwargs) + meta_data["batch_data_time"] = ( + speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 + ) + + tokens_i = self.generate_chunk( + speech, + speech_lengths, + key=key, + tokenizer=tokenizer, + cache=cache, + frontend=frontend, + **kwargs, + ) tokens.extend(tokens_i) - + text_postprocessed, _ = postprocess_utils.sentence_postprocess(tokens) - + result_i = {"key": key[0], "text": text_postprocessed} result = [result_i] - + cache["prev_samples"] = audio_sample[:-m] if _is_final: self.init_cache(cache, **kwargs) - + if kwargs.get("output_dir"): writer = DatadirWriter(kwargs.get("output_dir")) ibest_writer = writer[f"{1}best_recog"] ibest_writer["token"][key[0]] = " ".join(tokens) ibest_writer["text"][key[0]] = text_postprocessed - + return result, meta_data diff --git a/funasr/models/scama/utils.py b/funasr/models/scama/utils.py index 88325964a..c3f7bc37c 100644 --- a/funasr/models/scama/utils.py +++ b/funasr/models/scama/utils.py @@ -15,6 +15,7 @@ def sequence_mask(lengths, maxlen=None, dtype=torch.float32, device=None): return mask.type(dtype).to(device) if device is not None else mask.type(dtype) + def apply_cmvn(inputs, mvn): device = inputs.device dtype = inputs.dtype @@ -27,15 +28,13 @@ def apply_cmvn(inputs, mvn): return inputs.type(torch.float32) - - -def drop_and_add(inputs: torch.Tensor, - outputs: torch.Tensor, - training: bool, - dropout_rate: float = 0.1, - stoch_layer_coeff: float = 1.0): - - +def drop_and_add( + inputs: torch.Tensor, + outputs: torch.Tensor, + training: bool, + dropout_rate: float = 0.1, + stoch_layer_coeff: float = 1.0, +): outputs = F.dropout(outputs, p=dropout_rate, training=training, inplace=True) outputs *= stoch_layer_coeff @@ -51,8 +50,8 @@ def drop_and_add(inputs: torch.Tensor, def proc_tf_vocab(vocab_path): with open(vocab_path, encoding="utf-8") as f: token_list = [line.rstrip() for line in f] - if '' not in token_list: - token_list.append('') + if "" not in token_list: + token_list.append("") return token_list @@ -60,12 +59,12 @@ def gen_config_for_tfmodel(config_path, vocab_path, output_dir): token_list = proc_tf_vocab(vocab_path) with open(config_path, encoding="utf-8") as f: config = yaml.safe_load(f) - - config['token_list'] = token_list - + + config["token_list"] = token_list + if not os.path.exists(output_dir): os.makedirs(output_dir) - + with open(os.path.join(output_dir, "config.yaml"), "w", encoding="utf-8") as f: yaml_no_alias_safe_dump(config, f, indent=4, sort_keys=False) @@ -78,15 +77,13 @@ class NoAliasSafeDumper(yaml.SafeDumper): def yaml_no_alias_safe_dump(data, stream=None, **kwargs): """Safe-dump in yaml with no anchor/alias""" - return yaml.dump( - data, stream, allow_unicode=True, Dumper=NoAliasSafeDumper, **kwargs - ) + return yaml.dump(data, stream, allow_unicode=True, Dumper=NoAliasSafeDumper, **kwargs) -if __name__ == '__main__': +if __name__ == "__main__": import sys - + config_path = sys.argv[1] vocab_path = sys.argv[2] output_dir = sys.argv[3] - gen_config_for_tfmodel(config_path, vocab_path, output_dir) \ No newline at end of file + gen_config_for_tfmodel(config_path, vocab_path, output_dir) diff --git a/funasr/models/seaco_paraformer/export_meta.py b/funasr/models/seaco_paraformer/export_meta.py index 260b625a9..6d8096f59 100644 --- a/funasr/models/seaco_paraformer/export_meta.py +++ b/funasr/models/seaco_paraformer/export_meta.py @@ -9,173 +9,189 @@ from funasr.register import tables class ContextualEmbedderExport(torch.nn.Module): - def __init__(self, - model, - max_seq_len=512, - feats_dim=560, - **kwargs,): + def __init__( + self, + model, + max_seq_len=512, + feats_dim=560, + **kwargs, + ): super().__init__() - self.embedding = model.decoder.embed # model.bias_embed + self.embedding = model.decoder.embed # model.bias_embed model.bias_encoder.batch_first = False self.bias_encoder = model.bias_encoder - + def forward(self, hotword): - hotword = self.embedding(hotword).transpose(0, 1) # batch second + hotword = self.embedding(hotword).transpose(0, 1) # batch second hw_embed, (_, _) = self.bias_encoder(hotword) return hw_embed - + def export_dummy_inputs(self): - hotword = torch.tensor([ - [10, 11, 12, 13, 14, 10, 11, 12, 13, 14], - [100, 101, 0, 0, 0, 0, 0, 0, 0, 0], - [1, 0, 0, 0, 0, 0, 0, 0, 0, 0], - [10, 11, 12, 13, 14, 10, 11, 12, 13, 14], - [100, 101, 0, 0, 0, 0, 0, 0, 0, 0], - [1, 0, 0, 0, 0, 0, 0, 0, 0, 0], - ], - dtype=torch.int32) + hotword = torch.tensor( + [ + [10, 11, 12, 13, 14, 10, 11, 12, 13, 14], + [100, 101, 0, 0, 0, 0, 0, 0, 0, 0], + [1, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [10, 11, 12, 13, 14, 10, 11, 12, 13, 14], + [100, 101, 0, 0, 0, 0, 0, 0, 0, 0], + [1, 0, 0, 0, 0, 0, 0, 0, 0, 0], + ], + dtype=torch.int32, + ) # hotword_length = torch.tensor([10, 2, 1], dtype=torch.int32) - return (hotword) + return hotword def export_input_names(self): - return ['hotword'] + return ["hotword"] def export_output_names(self): - return ['hw_embed'] + return ["hw_embed"] def export_dynamic_axes(self): return { - 'hotword': { - 0: 'num_hotwords', + "hotword": { + 0: "num_hotwords", }, - 'hw_embed': { - 0: 'num_hotwords', + "hw_embed": { + 0: "num_hotwords", }, } - + def export_name(self): - return 'model_eb.onnx' - + return "model_eb.onnx" + def export_rebuild_model(model, **kwargs): - model.device = kwargs.get("device") - is_onnx = kwargs.get("type", "onnx") == "onnx" - encoder_class = tables.encoder_classes.get(kwargs["encoder"]+"Export") - model.encoder = encoder_class(model.encoder, onnx=is_onnx) - - predictor_class = tables.predictor_classes.get(kwargs["predictor"]+"Export") - model.predictor = predictor_class(model.predictor, onnx=is_onnx) + model.device = kwargs.get("device") + is_onnx = kwargs.get("type", "onnx") == "onnx" + encoder_class = tables.encoder_classes.get(kwargs["encoder"] + "Export") + model.encoder = encoder_class(model.encoder, onnx=is_onnx) - # before decoder convert into export class - embedder_class = ContextualEmbedderExport - embedder_model = embedder_class(model, onnx=is_onnx) + predictor_class = tables.predictor_classes.get(kwargs["predictor"] + "Export") + model.predictor = predictor_class(model.predictor, onnx=is_onnx) - decoder_class = tables.decoder_classes.get(kwargs["decoder"]+"Export") - model.decoder = decoder_class(model.decoder, onnx=is_onnx) - - seaco_decoder_class = tables.decoder_classes.get(kwargs["seaco_decoder"]+"Export") - model.seaco_decoder = seaco_decoder_class(model.seaco_decoder, onnx=is_onnx) - - from funasr.utils.torch_function import sequence_mask - model.make_pad_mask = sequence_mask(kwargs["max_seq_len"], flip=False) - - from funasr.utils.torch_function import sequence_mask - model.make_pad_mask = sequence_mask(kwargs["max_seq_len"], flip=False) - model.feats_dim = 560 - model.NOBIAS = 8377 + # before decoder convert into export class + embedder_class = ContextualEmbedderExport + embedder_model = embedder_class(model, onnx=is_onnx) - import copy - import types - backbone_model = copy.copy(model) + decoder_class = tables.decoder_classes.get(kwargs["decoder"] + "Export") + model.decoder = decoder_class(model.decoder, onnx=is_onnx) - # backbone - backbone_model.forward = types.MethodType(export_backbone_forward, backbone_model) - backbone_model.export_dummy_inputs = types.MethodType(export_backbone_dummy_inputs, backbone_model) - backbone_model.export_input_names = types.MethodType(export_backbone_input_names, backbone_model) - backbone_model.export_output_names = types.MethodType(export_backbone_output_names, backbone_model) - backbone_model.export_dynamic_axes = types.MethodType(export_backbone_dynamic_axes, backbone_model) - backbone_model.export_name = types.MethodType(export_backbone_name, backbone_model) - - return backbone_model, embedder_model + seaco_decoder_class = tables.decoder_classes.get(kwargs["seaco_decoder"] + "Export") + model.seaco_decoder = seaco_decoder_class(model.seaco_decoder, onnx=is_onnx) + + from funasr.utils.torch_function import sequence_mask + + model.make_pad_mask = sequence_mask(kwargs["max_seq_len"], flip=False) + + from funasr.utils.torch_function import sequence_mask + + model.make_pad_mask = sequence_mask(kwargs["max_seq_len"], flip=False) + model.feats_dim = 560 + model.NOBIAS = 8377 + + import copy + import types + + backbone_model = copy.copy(model) + + # backbone + backbone_model.forward = types.MethodType(export_backbone_forward, backbone_model) + backbone_model.export_dummy_inputs = types.MethodType( + export_backbone_dummy_inputs, backbone_model + ) + backbone_model.export_input_names = types.MethodType( + export_backbone_input_names, backbone_model + ) + backbone_model.export_output_names = types.MethodType( + export_backbone_output_names, backbone_model + ) + backbone_model.export_dynamic_axes = types.MethodType( + export_backbone_dynamic_axes, backbone_model + ) + backbone_model.export_name = types.MethodType(export_backbone_name, backbone_model) + + return backbone_model, embedder_model def export_backbone_forward( - self, - speech: torch.Tensor, - speech_lengths: torch.Tensor, - bias_embed: torch.Tensor, - # lmbd: float, - ): - # a. To device - batch = {"speech": speech, "speech_lengths": speech_lengths} + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + bias_embed: torch.Tensor, + # lmbd: float, +): + # a. To device + batch = {"speech": speech, "speech_lengths": speech_lengths} - enc, enc_len = self.encoder(**batch) - mask = self.make_pad_mask(enc_len)[:, None, :] - pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = self.predictor(enc, mask) - pre_token_length = pre_token_length.floor().type(torch.int32) + enc, enc_len = self.encoder(**batch) + mask = self.make_pad_mask(enc_len)[:, None, :] + pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = self.predictor(enc, mask) + pre_token_length = pre_token_length.floor().type(torch.int32) - decoder_out, decoder_hidden, _ = self.decoder(enc, enc_len, pre_acoustic_embeds, pre_token_length, return_hidden=True, return_both=True) - decoder_out = torch.log_softmax(decoder_out, dim=-1) - # seaco forward - B, N, D = bias_embed.shape - _contextual_length = torch.ones(B) * N + decoder_out, decoder_hidden, _ = self.decoder( + enc, enc_len, pre_acoustic_embeds, pre_token_length, return_hidden=True, return_both=True + ) + decoder_out = torch.log_softmax(decoder_out, dim=-1) + # seaco forward + B, N, D = bias_embed.shape + _contextual_length = torch.ones(B) * N - # ASF - hotword_scores = self.seaco_decoder.forward_asf6(bias_embed, _contextual_length, decoder_hidden, pre_token_length) - hotword_scores = hotword_scores[0].sum(0).sum(0) - # _ = self.decoder2(bias_embed, _contextual_length, decoder_hidden, pre_token_length) - # hotword_scores = self.decoder2.model.decoders[-1].attn_mat[0][0].sum(0).sum(0) - dec_filter = torch.sort(hotword_scores, descending=True)[1][:51] - contextual_info = bias_embed[:,dec_filter] - num_hot_word = contextual_info.shape[1] - _contextual_length = torch.Tensor([num_hot_word]).int().repeat(B).to(enc.device) + # ASF + hotword_scores = self.seaco_decoder.forward_asf6( + bias_embed, _contextual_length, decoder_hidden, pre_token_length + ) + hotword_scores = hotword_scores[0].sum(0).sum(0) + # _ = self.decoder2(bias_embed, _contextual_length, decoder_hidden, pre_token_length) + # hotword_scores = self.decoder2.model.decoders[-1].attn_mat[0][0].sum(0).sum(0) + dec_filter = torch.sort(hotword_scores, descending=True)[1][:51] + contextual_info = bias_embed[:, dec_filter] + num_hot_word = contextual_info.shape[1] + _contextual_length = torch.Tensor([num_hot_word]).int().repeat(B).to(enc.device) + + # again + cif_attended, _ = self.seaco_decoder( + contextual_info, _contextual_length, pre_acoustic_embeds, pre_token_length + ) + dec_attended, _ = self.seaco_decoder( + contextual_info, _contextual_length, decoder_hidden, pre_token_length + ) + merged = cif_attended + dec_attended + dha_output = self.hotword_output_layer(merged) + dha_pred = torch.log_softmax(dha_output, dim=-1) + # merging logits + dha_ids = dha_pred.max(-1)[-1] + dha_mask = (dha_ids == self.NOBIAS).int().unsqueeze(-1) + decoder_out = decoder_out * dha_mask + dha_pred * (1 - dha_mask) + return decoder_out, pre_token_length, alphas - # again - cif_attended, _ = self.seaco_decoder(contextual_info, _contextual_length, pre_acoustic_embeds, pre_token_length) - dec_attended, _ = self.seaco_decoder(contextual_info, _contextual_length, decoder_hidden, pre_token_length) - merged = cif_attended + dec_attended - dha_output = self.hotword_output_layer(merged) - dha_pred = torch.log_softmax(dha_output, dim=-1) - # merging logits - dha_ids = dha_pred.max(-1)[-1] - dha_mask = (dha_ids == self.NOBIAS).int().unsqueeze(-1) - decoder_out = decoder_out * dha_mask + dha_pred * (1-dha_mask) - return decoder_out, pre_token_length, alphas def export_backbone_dummy_inputs(self): - speech = torch.randn(2, 30, self.feats_dim) - speech_lengths = torch.tensor([15, 30], dtype=torch.int32) - bias_embed = torch.randn(2, 1, 512) - return (speech, speech_lengths, bias_embed) + speech = torch.randn(2, 30, self.feats_dim) + speech_lengths = torch.tensor([15, 30], dtype=torch.int32) + bias_embed = torch.randn(2, 1, 512) + return (speech, speech_lengths, bias_embed) + def export_backbone_input_names(self): - return ['speech', 'speech_lengths', 'bias_embed'] + return ["speech", "speech_lengths", "bias_embed"] + def export_backbone_output_names(self): - return ['logits', 'token_num', 'alphas'] + return ["logits", "token_num", "alphas"] + def export_backbone_dynamic_axes(self): - return { - 'speech': { - 0: 'batch_size', - 1: 'feats_length' - }, - 'speech_lengths': { - 0: 'batch_size', - }, - 'bias_embed': { - 0: 'batch_size', - 1: 'num_hotwords' - }, - 'logits': { - 0: 'batch_size', - 1: 'logits_length' - }, - 'pre_acoustic_embeds': { - 1: 'feats_length1' - } - } - + return { + "speech": {0: "batch_size", 1: "feats_length"}, + "speech_lengths": { + 0: "batch_size", + }, + "bias_embed": {0: "batch_size", 1: "num_hotwords"}, + "logits": {0: "batch_size", 1: "logits_length"}, + "pre_acoustic_embeds": {1: "feats_length1"}, + } + + def export_backbone_name(self): - return 'model.onnx' - \ No newline at end of file + return "model.onnx" diff --git a/funasr/models/seaco_paraformer/model.py b/funasr/models/seaco_paraformer/model.py index b28de94f7..3b6595cb8 100644 --- a/funasr/models/seaco_paraformer/model.py +++ b/funasr/models/seaco_paraformer/model.py @@ -47,35 +47,37 @@ class SeacoParaformer(BiCifParaformer, Paraformer): SeACo-Paraformer: A Non-Autoregressive ASR System with Flexible and Effective Hotword Customization Ability https://arxiv.org/abs/2308.03266 """ - + def __init__( self, *args, **kwargs, ): super().__init__(*args, **kwargs) - + self.inner_dim = kwargs.get("inner_dim", 256) self.bias_encoder_type = kwargs.get("bias_encoder_type", "lstm") bias_encoder_dropout_rate = kwargs.get("bias_encoder_dropout_rate", 0.0) bias_encoder_bid = kwargs.get("bias_encoder_bid", False) seaco_lsm_weight = kwargs.get("seaco_lsm_weight", 0.0) seaco_length_normalized_loss = kwargs.get("seaco_length_normalized_loss", True) - + # bias encoder - if self.bias_encoder_type == 'lstm': - self.bias_encoder = torch.nn.LSTM(self.inner_dim, - self.inner_dim, - 2, - batch_first=True, - dropout=bias_encoder_dropout_rate, - bidirectional=bias_encoder_bid) + if self.bias_encoder_type == "lstm": + self.bias_encoder = torch.nn.LSTM( + self.inner_dim, + self.inner_dim, + 2, + batch_first=True, + dropout=bias_encoder_dropout_rate, + bidirectional=bias_encoder_bid, + ) if bias_encoder_bid: - self.lstm_proj = torch.nn.Linear(self.inner_dim*2, self.inner_dim) + self.lstm_proj = torch.nn.Linear(self.inner_dim * 2, self.inner_dim) else: self.lstm_proj = None # self.bias_embed = torch.nn.Embedding(self.vocab_size, self.inner_dim) - elif self.bias_encoder_type == 'mean': + elif self.bias_encoder_type == "mean": self.bias_embed = torch.nn.Embedding(self.vocab_size, self.inner_dim) else: logging.error("Unsupport bias encoder type: {}".format(self.bias_encoder_type)) @@ -101,7 +103,7 @@ class SeacoParaformer(BiCifParaformer, Paraformer): self.seaco_weight = kwargs.get("seaco_weight", 0.01) self.NO_BIAS = kwargs.get("NO_BIAS", 8377) self.predictor_name = kwargs.get("predictor") - + def forward( self, speech: torch.Tensor, @@ -111,7 +113,7 @@ class SeacoParaformer(BiCifParaformer, Paraformer): **kwargs, ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]: """Frontend + Encoder + Decoder + Calc loss - + Args: speech: (Batch, Length, ...) speech_lengths: (Batch, ) @@ -124,38 +126,36 @@ class SeacoParaformer(BiCifParaformer, Paraformer): speech_lengths = speech_lengths[:, 0] # Check that batch_size is unified assert ( - speech.shape[0] - == speech_lengths.shape[0] - == text.shape[0] - == text_lengths.shape[0] + speech.shape[0] == speech_lengths.shape[0] == text.shape[0] == text_lengths.shape[0] ), (speech.shape, speech_lengths.shape, text.shape, text_lengths.shape) - + hotword_pad = kwargs.get("hotword_pad") hotword_lengths = kwargs.get("hotword_lengths") seaco_label_pad = kwargs.get("seaco_label_pad") if len(hotword_lengths.size()) > 1: hotword_lengths = hotword_lengths[:, 0] - + batch_size = speech.shape[0] # for data-parallel text = text[:, : text_lengths.max()] - speech = speech[:, :speech_lengths.max()] - + speech = speech[:, : speech_lengths.max()] + # 1. Encoder encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) if self.predictor_bias == 1: _, ys_pad = add_sos_eos(text, self.sos, self.eos, self.ignore_id) ys_lengths = text_lengths + self.predictor_bias - stats = dict() - loss_seaco = self._calc_seaco_loss(encoder_out, - encoder_out_lens, - ys_pad, - ys_lengths, - hotword_pad, - hotword_lengths, - seaco_label_pad, - ) + stats = dict() + loss_seaco = self._calc_seaco_loss( + encoder_out, + encoder_out_lens, + ys_pad, + ys_lengths, + hotword_pad, + hotword_lengths, + seaco_label_pad, + ) if self.train_decoder: loss_att, acc_att, _, _, _ = self._calc_att_loss( encoder_out, encoder_out_lens, text, text_lengths @@ -165,7 +165,7 @@ class SeacoParaformer(BiCifParaformer, Paraformer): stats["acc_att"] = acc_att else: loss = loss_seaco - + stats["loss_seaco"] = torch.clone(loss_seaco.detach()) stats["loss"] = torch.clone(loss.detach()) @@ -177,97 +177,148 @@ class SeacoParaformer(BiCifParaformer, Paraformer): def _merge(self, cif_attended, dec_attended): return cif_attended + dec_attended - + def calc_predictor(self, encoder_out, encoder_out_lens): - encoder_out_mask = (~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :]).to( - encoder_out.device) - predictor_outs = self.predictor(encoder_out, None, encoder_out_mask, ignore_id=self.ignore_id) + encoder_out_mask = ( + ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :] + ).to(encoder_out.device) + predictor_outs = self.predictor( + encoder_out, None, encoder_out_mask, ignore_id=self.ignore_id + ) return predictor_outs[:4] - + def _calc_seaco_loss( - self, - encoder_out: torch.Tensor, - encoder_out_lens: torch.Tensor, - ys_pad: torch.Tensor, - ys_lengths: torch.Tensor, - hotword_pad: torch.Tensor, - hotword_lengths: torch.Tensor, - seaco_label_pad: torch.Tensor, - ): + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_lengths: torch.Tensor, + hotword_pad: torch.Tensor, + hotword_lengths: torch.Tensor, + seaco_label_pad: torch.Tensor, + ): # predictor forward - encoder_out_mask = (~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :]).to( - encoder_out.device) - pre_acoustic_embeds = self.predictor(encoder_out, ys_pad, encoder_out_mask, ignore_id=self.ignore_id)[0] + encoder_out_mask = ( + ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :] + ).to(encoder_out.device) + pre_acoustic_embeds = self.predictor( + encoder_out, ys_pad, encoder_out_mask, ignore_id=self.ignore_id + )[0] # decoder forward - decoder_out, _ = self.decoder(encoder_out, encoder_out_lens, pre_acoustic_embeds, ys_lengths, return_hidden=True) - selected = self._hotword_representation(hotword_pad, - hotword_lengths) - contextual_info = selected.squeeze(0).repeat(encoder_out.shape[0], 1, 1).to(encoder_out.device) + decoder_out, _ = self.decoder( + encoder_out, encoder_out_lens, pre_acoustic_embeds, ys_lengths, return_hidden=True + ) + selected = self._hotword_representation(hotword_pad, hotword_lengths) + contextual_info = ( + selected.squeeze(0).repeat(encoder_out.shape[0], 1, 1).to(encoder_out.device) + ) num_hot_word = contextual_info.shape[1] - _contextual_length = torch.Tensor([num_hot_word]).int().repeat(encoder_out.shape[0]).to(encoder_out.device) + _contextual_length = ( + torch.Tensor([num_hot_word]).int().repeat(encoder_out.shape[0]).to(encoder_out.device) + ) # dha core - cif_attended, _ = self.seaco_decoder(contextual_info, _contextual_length, pre_acoustic_embeds, ys_lengths) - dec_attended, _ = self.seaco_decoder(contextual_info, _contextual_length, decoder_out, ys_lengths) + cif_attended, _ = self.seaco_decoder( + contextual_info, _contextual_length, pre_acoustic_embeds, ys_lengths + ) + dec_attended, _ = self.seaco_decoder( + contextual_info, _contextual_length, decoder_out, ys_lengths + ) merged = self._merge(cif_attended, dec_attended) - dha_output = self.hotword_output_layer(merged[:, :-1]) # remove the last token in loss calculation + dha_output = self.hotword_output_layer( + merged[:, :-1] + ) # remove the last token in loss calculation loss_att = self.criterion_seaco(dha_output, seaco_label_pad) return loss_att - def _seaco_decode_with_ASF(self, - encoder_out, - encoder_out_lens, - sematic_embeds, - ys_pad_lens, - hw_list, - nfilter=50, - seaco_weight=1.0): + def _seaco_decode_with_ASF( + self, + encoder_out, + encoder_out_lens, + sematic_embeds, + ys_pad_lens, + hw_list, + nfilter=50, + seaco_weight=1.0, + ): # decoder forward - decoder_out, decoder_hidden, _ = self.decoder(encoder_out, encoder_out_lens, sematic_embeds, ys_pad_lens, return_hidden=True, return_both=True) + decoder_out, decoder_hidden, _ = self.decoder( + encoder_out, + encoder_out_lens, + sematic_embeds, + ys_pad_lens, + return_hidden=True, + return_both=True, + ) decoder_pred = torch.log_softmax(decoder_out, dim=-1) if hw_list is not None: hw_lengths = [len(i) for i in hw_list] hw_list_ = [torch.Tensor(i).long() for i in hw_list] hw_list_pad = pad_list(hw_list_, 0).to(encoder_out.device) - selected = self._hotword_representation(hw_list_pad, torch.Tensor(hw_lengths).int().to(encoder_out.device)) + selected = self._hotword_representation( + hw_list_pad, torch.Tensor(hw_lengths).int().to(encoder_out.device) + ) - contextual_info = selected.squeeze(0).repeat(encoder_out.shape[0], 1, 1).to(encoder_out.device) + contextual_info = ( + selected.squeeze(0).repeat(encoder_out.shape[0], 1, 1).to(encoder_out.device) + ) num_hot_word = contextual_info.shape[1] - _contextual_length = torch.Tensor([num_hot_word]).int().repeat(encoder_out.shape[0]).to(encoder_out.device) + _contextual_length = ( + torch.Tensor([num_hot_word]) + .int() + .repeat(encoder_out.shape[0]) + .to(encoder_out.device) + ) # ASF Core if nfilter > 0 and nfilter < num_hot_word: - hotword_scores = self.seaco_decoder.forward_asf6(contextual_info, _contextual_length, decoder_hidden, ys_pad_lens) + hotword_scores = self.seaco_decoder.forward_asf6( + contextual_info, _contextual_length, decoder_hidden, ys_pad_lens + ) hotword_scores = hotword_scores[0].sum(0).sum(0) # hotword_scores /= torch.sqrt(torch.tensor(hw_lengths)[:-1].float()).to(hotword_scores.device) - dec_filter = torch.topk(hotword_scores, min(nfilter, num_hot_word-1))[1].tolist() + dec_filter = torch.topk(hotword_scores, min(nfilter, num_hot_word - 1))[1].tolist() add_filter = dec_filter - add_filter.append(len(hw_list_pad)-1) + add_filter.append(len(hw_list_pad) - 1) # filter hotword embedding selected = selected[add_filter] # again - contextual_info = selected.squeeze(0).repeat(encoder_out.shape[0], 1, 1).to(encoder_out.device) + contextual_info = ( + selected.squeeze(0).repeat(encoder_out.shape[0], 1, 1).to(encoder_out.device) + ) num_hot_word = contextual_info.shape[1] - _contextual_length = torch.Tensor([num_hot_word]).int().repeat(encoder_out.shape[0]).to(encoder_out.device) - + _contextual_length = ( + torch.Tensor([num_hot_word]) + .int() + .repeat(encoder_out.shape[0]) + .to(encoder_out.device) + ) + # SeACo Core - cif_attended, _ = self.seaco_decoder(contextual_info, _contextual_length, sematic_embeds, ys_pad_lens) - dec_attended, _ = self.seaco_decoder(contextual_info, _contextual_length, decoder_hidden, ys_pad_lens) + cif_attended, _ = self.seaco_decoder( + contextual_info, _contextual_length, sematic_embeds, ys_pad_lens + ) + dec_attended, _ = self.seaco_decoder( + contextual_info, _contextual_length, decoder_hidden, ys_pad_lens + ) merged = self._merge(cif_attended, dec_attended) - dha_output = self.hotword_output_layer(merged) # remove the last token in loss calculation + dha_output = self.hotword_output_layer( + merged + ) # remove the last token in loss calculation dha_pred = torch.log_softmax(dha_output, dim=-1) + def _merge_res(dec_output, dha_output): lmbd = torch.Tensor([seaco_weight] * dha_output.shape[0]) - dha_ids = dha_output.max(-1)[-1]# [0] + dha_ids = dha_output.max(-1)[-1] # [0] dha_mask = (dha_ids == self.NO_BIAS).int().unsqueeze(-1) a = (1 - lmbd) / lmbd b = 1 / lmbd a, b = a.to(dec_output.device), b.to(dec_output.device) dha_mask = (dha_mask + a.reshape(-1, 1, 1)) / b.reshape(-1, 1, 1) # logits = dec_output * dha_mask + dha_output[:,:,:-1] * (1-dha_mask) - logits = dec_output * dha_mask + dha_output[:,:,:] * (1-dha_mask) + logits = dec_output * dha_mask + dha_output[:, :, :] * (1 - dha_mask) return logits merged_pred = _merge_res(decoder_pred, dha_pred) @@ -275,13 +326,11 @@ class SeacoParaformer(BiCifParaformer, Paraformer): else: return decoder_pred - def _hotword_representation(self, - hotword_pad, - hotword_lengths): - if self.bias_encoder_type != 'lstm': + def _hotword_representation(self, hotword_pad, hotword_lengths): + if self.bias_encoder_type != "lstm": logging.error("Unsupported bias encoder type") - - ''' + + """ hw_embed = self.decoder.embed(hotword_pad) hw_embed, (_, _) = self.bias_encoder(hw_embed) if self.lstm_proj is not None: @@ -289,11 +338,16 @@ class SeacoParaformer(BiCifParaformer, Paraformer): _ind = np.arange(0, hw_embed.shape[0]).tolist() selected = hw_embed[_ind, [i-1 for i in hotword_lengths.detach().cpu().tolist()]] return selected - ''' + """ # hw_embed = self.sac_embedding(hotword_pad) hw_embed = self.decoder.embed(hotword_pad) - hw_embed = torch.nn.utils.rnn.pack_padded_sequence(hw_embed, hotword_lengths.cpu().type(torch.int64), batch_first=True, enforce_sorted=False) + hw_embed = torch.nn.utils.rnn.pack_padded_sequence( + hw_embed, + hotword_lengths.cpu().type(torch.int64), + batch_first=True, + enforce_sorted=False, + ) packed_rnn_output, _ = self.bias_encoder(hw_embed) rnn_output = torch.nn.utils.rnn.pad_packed_sequence(packed_rnn_output, batch_first=True)[0] if self.lstm_proj is not None: @@ -301,50 +355,59 @@ class SeacoParaformer(BiCifParaformer, Paraformer): else: hw_hidden = rnn_output _ind = np.arange(0, hw_hidden.shape[0]).tolist() - selected = hw_hidden[_ind, [i-1 for i in hotword_lengths.detach().cpu().tolist()]] - return selected - - def inference(self, - data_in, - data_lengths=None, - key: list = None, - tokenizer=None, - frontend=None, - **kwargs, - ): - + selected = hw_hidden[_ind, [i - 1 for i in hotword_lengths.detach().cpu().tolist()]] + return selected + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): + # init beamsearch is_use_ctc = kwargs.get("decoding_ctc_weight", 0.0) > 0.00001 and self.ctc != None - is_use_lm = kwargs.get("lm_weight", 0.0) > 0.00001 and kwargs.get("lm_file", None) is not None + is_use_lm = ( + kwargs.get("lm_weight", 0.0) > 0.00001 and kwargs.get("lm_file", None) is not None + ) if self.beam_search is None and (is_use_lm or is_use_ctc): logging.info("enable beam_search") self.init_beam_search(**kwargs) self.nbest = kwargs.get("nbest", 1) meta_data = {} - + # extract fbank feats time1 = time.perf_counter() - audio_sample_list = load_audio_text_image_video(data_in, fs=frontend.fs, audio_fs=kwargs.get("fs", 16000)) + audio_sample_list = load_audio_text_image_video( + data_in, fs=frontend.fs, audio_fs=kwargs.get("fs", 16000) + ) time2 = time.perf_counter() meta_data["load_data"] = f"{time2 - time1:0.3f}" - speech, speech_lengths = extract_fbank(audio_sample_list, data_type=kwargs.get("data_type", "sound"), - frontend=frontend) + speech, speech_lengths = extract_fbank( + audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend + ) time3 = time.perf_counter() meta_data["extract_feat"] = f"{time3 - time2:0.3f}" - meta_data[ - "batch_data_time"] = speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 - + meta_data["batch_data_time"] = ( + speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 + ) + speech = speech.to(device=kwargs["device"]) speech_lengths = speech_lengths.to(device=kwargs["device"]) - + # hotword - self.hotword_list = self.generate_hotwords_list(kwargs.get("hotword", None), tokenizer=tokenizer, frontend=frontend) - + self.hotword_list = self.generate_hotwords_list( + kwargs.get("hotword", None), tokenizer=tokenizer, frontend=frontend + ) + # Encoder encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) if isinstance(encoder_out, tuple): encoder_out = encoder_out[0] - + # predictor predictor_outs = self.calc_predictor(encoder_out, encoder_out_lens) pre_acoustic_embeds, pre_token_length = predictor_outs[0], predictor_outs[1] @@ -352,42 +415,43 @@ class SeacoParaformer(BiCifParaformer, Paraformer): if torch.max(pre_token_length) < 1: return ([],) - decoder_out = self._seaco_decode_with_ASF(encoder_out, - encoder_out_lens, - pre_acoustic_embeds, - pre_token_length, - hw_list=self.hotword_list - ) + decoder_out = self._seaco_decode_with_ASF( + encoder_out, + encoder_out_lens, + pre_acoustic_embeds, + pre_token_length, + hw_list=self.hotword_list, + ) # decoder_out, _ = decoder_outs[0], decoder_outs[1] if self.predictor_name == "CifPredictorV3": - _, _, us_alphas, us_peaks = self.calc_predictor_timestamp(encoder_out, - encoder_out_lens, - pre_token_length) + _, _, us_alphas, us_peaks = self.calc_predictor_timestamp( + encoder_out, encoder_out_lens, pre_token_length + ) else: us_alphas = None - + results = [] b, n, d = decoder_out.size() for i in range(b): - x = encoder_out[i, :encoder_out_lens[i], :] - am_scores = decoder_out[i, :pre_token_length[i], :] + x = encoder_out[i, : encoder_out_lens[i], :] + am_scores = decoder_out[i, : pre_token_length[i], :] if self.beam_search is not None: nbest_hyps = self.beam_search( - x=x, am_scores=am_scores, maxlenratio=kwargs.get("maxlenratio", 0.0), - minlenratio=kwargs.get("minlenratio", 0.0) + x=x, + am_scores=am_scores, + maxlenratio=kwargs.get("maxlenratio", 0.0), + minlenratio=kwargs.get("minlenratio", 0.0), ) - + nbest_hyps = nbest_hyps[: self.nbest] else: - + yseq = am_scores.argmax(dim=-1) score = am_scores.max(dim=-1)[0] score = torch.sum(score, dim=-1) # pad with mask tokens to ensure compatibility with sos/eos tokens - yseq = torch.tensor( - [self.sos] + yseq.tolist() + [self.eos], device=yseq.device - ) + yseq = torch.tensor([self.sos] + yseq.tolist() + [self.eos], device=yseq.device) nbest_hyps = [Hypothesis(yseq=yseq, score=score)] for nbest_idx, hyp in enumerate(nbest_hyps): ibest_writer = None @@ -395,31 +459,40 @@ class SeacoParaformer(BiCifParaformer, Paraformer): if not hasattr(self, "writer"): self.writer = DatadirWriter(kwargs.get("output_dir")) ibest_writer = self.writer[f"{nbest_idx + 1}best_recog"] - + # remove sos/eos and get results last_pos = -1 if isinstance(hyp.yseq, list): token_int = hyp.yseq[1:last_pos] else: token_int = hyp.yseq[1:last_pos].tolist() - + # remove blank symbol id, which is assumed to be 0 token_int = list( - filter(lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int)) - + filter( + lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int + ) + ) + if tokenizer is not None: # Change integer-ids to tokens token = tokenizer.ids2tokens(token_int) text = tokenizer.tokens2text(token) if us_alphas is not None: - _, timestamp = ts_prediction_lfr6_standard(us_alphas[i][:encoder_out_lens[i] * 3], - us_peaks[i][:encoder_out_lens[i] * 3], - copy.copy(token), - vad_offset=kwargs.get("begin_time", 0)) - text_postprocessed, time_stamp_postprocessed, _ = \ + _, timestamp = ts_prediction_lfr6_standard( + us_alphas[i][: encoder_out_lens[i] * 3], + us_peaks[i][: encoder_out_lens[i] * 3], + copy.copy(token), + vad_offset=kwargs.get("begin_time", 0), + ) + text_postprocessed, time_stamp_postprocessed, _ = ( postprocess_utils.sentence_postprocess(token, timestamp) - result_i = {"key": key[i], "text": text_postprocessed, - "timestamp": time_stamp_postprocessed} + ) + result_i = { + "key": key[i], + "text": text_postprocessed, + "timestamp": time_stamp_postprocessed, + } if ibest_writer is not None: ibest_writer["token"][key[i]] = " ".join(token) ibest_writer["timestamp"][key[i]] = time_stamp_postprocessed @@ -433,7 +506,7 @@ class SeacoParaformer(BiCifParaformer, Paraformer): else: result_i = {"key": key[i], "token_int": token_int} results.append(result_i) - + return results, meta_data def generate_hotwords_list(self, hotword_list_or_file, tokenizer=None, frontend=None): @@ -448,9 +521,9 @@ class SeacoParaformer(BiCifParaformer, Paraformer): value = s[1:] seg_dict[key] = " ".join(value) return seg_dict - + def seg_tokenize(txt, seg_dict): - pattern = re.compile(r'^[\u4E00-\u9FA50-9]+$') + pattern = re.compile(r"^[\u4E00-\u9FA50-9]+$") out_txt = "" for word in txt: word = word.lower() @@ -466,11 +539,11 @@ class SeacoParaformer(BiCifParaformer, Paraformer): else: out_txt += "" + " " return out_txt.strip().split() - + seg_dict = None if frontend.cmvn_file is not None: model_dir = os.path.dirname(frontend.cmvn_file) - seg_dict_file = os.path.join(model_dir, 'seg_dict') + seg_dict_file = os.path.join(model_dir, "seg_dict") if os.path.exists(seg_dict_file): seg_dict = load_seg_dict(seg_dict_file) else: @@ -479,11 +552,11 @@ class SeacoParaformer(BiCifParaformer, Paraformer): if hotword_list_or_file is None: hotword_list = None # for local txt inputs - elif os.path.exists(hotword_list_or_file) and hotword_list_or_file.endswith('.txt'): + elif os.path.exists(hotword_list_or_file) and hotword_list_or_file.endswith(".txt"): logging.info("Attempting to parse hotwords from local txt...") hotword_list = [] hotword_str_list = [] - with codecs.open(hotword_list_or_file, 'r') as fin: + with codecs.open(hotword_list_or_file, "r") as fin: for line in fin.readlines(): hw = line.strip() hw_list = hw.split() @@ -492,11 +565,14 @@ class SeacoParaformer(BiCifParaformer, Paraformer): hotword_str_list.append(hw) hotword_list.append(tokenizer.tokens2ids(hw_list)) hotword_list.append([self.sos]) - hotword_str_list.append('') - logging.info("Initialized hotword list from file: {}, hotword list: {}." - .format(hotword_list_or_file, hotword_str_list)) + hotword_str_list.append("") + logging.info( + "Initialized hotword list from file: {}, hotword list: {}.".format( + hotword_list_or_file, hotword_str_list + ) + ) # for url, download and generate txt - elif hotword_list_or_file.startswith('http'): + elif hotword_list_or_file.startswith("http"): logging.info("Attempting to parse hotwords from url...") work_dir = tempfile.TemporaryDirectory().name if not os.path.exists(work_dir): @@ -507,7 +583,7 @@ class SeacoParaformer(BiCifParaformer, Paraformer): hotword_list_or_file = text_file_path hotword_list = [] hotword_str_list = [] - with codecs.open(hotword_list_or_file, 'r') as fin: + with codecs.open(hotword_list_or_file, "r") as fin: for line in fin.readlines(): hw = line.strip() hw_list = hw.split() @@ -516,11 +592,14 @@ class SeacoParaformer(BiCifParaformer, Paraformer): hotword_str_list.append(hw) hotword_list.append(tokenizer.tokens2ids(hw_list)) hotword_list.append([self.sos]) - hotword_str_list.append('') - logging.info("Initialized hotword list from file: {}, hotword list: {}." - .format(hotword_list_or_file, hotword_str_list)) + hotword_str_list.append("") + logging.info( + "Initialized hotword list from file: {}, hotword list: {}.".format( + hotword_list_or_file, hotword_str_list + ) + ) # for text str input - elif not hotword_list_or_file.endswith('.txt'): + elif not hotword_list_or_file.endswith(".txt"): logging.info("Attempting to parse hotwords as str...") hotword_list = [] hotword_str_list = [] @@ -531,7 +610,7 @@ class SeacoParaformer(BiCifParaformer, Paraformer): hw_list = seg_tokenize(hw_list, seg_dict) hotword_list.append(tokenizer.tokens2ids(hw_list)) hotword_list.append([self.sos]) - hotword_str_list.append('') + hotword_str_list.append("") logging.info("Hotword list: {}.".format(hotword_str_list)) else: hotword_list = None @@ -541,9 +620,9 @@ class SeacoParaformer(BiCifParaformer, Paraformer): self, **kwargs, ): - if 'max_seq_len' not in kwargs: - kwargs['max_seq_len'] = 512 + if "max_seq_len" not in kwargs: + kwargs["max_seq_len"] = 512 from .export_meta import export_rebuild_model + models = export_rebuild_model(model=self, **kwargs) return models - diff --git a/funasr/models/sense_voice/decoder.py b/funasr/models/sense_voice/decoder.py index 9fdb3bdf6..5d34ff3f9 100644 --- a/funasr/models/sense_voice/decoder.py +++ b/funasr/models/sense_voice/decoder.py @@ -31,125 +31,129 @@ class Linear(nn.Linear): ) - def sense_voice_decode_forward( - self, - x: torch.Tensor, - xa: torch.Tensor, - kv_cache: Optional[dict] = None, - **kwargs, + self, + x: torch.Tensor, + xa: torch.Tensor, + kv_cache: Optional[dict] = None, + **kwargs, ): - """Forward decoder. + """Forward decoder. - Args: - hs_pad: encoded memory, float32 (batch, maxlen_in, feat) - hlens: (batch) - ys_in_pad: - input token ids, int64 (batch, maxlen_out) - if input_layer == "embed" - input tensor (batch, maxlen_out, #mels) in the other cases - ys_in_lens: (batch) - Returns: - (tuple): tuple containing: + Args: + hs_pad: encoded memory, float32 (batch, maxlen_in, feat) + hlens: (batch) + ys_in_pad: + input token ids, int64 (batch, maxlen_out) + if input_layer == "embed" + input tensor (batch, maxlen_out, #mels) in the other cases + ys_in_lens: (batch) + Returns: + (tuple): tuple containing: - x: decoded token score before softmax (batch, maxlen_out, token) - if use_output_layer is True, - olens: (batch, ) - """ - # import pdb;pdb.set_trace() - use_padmask = self.use_padmask - hlens = kwargs.get("hlens", None) - - ys_in_lens = kwargs.get("ys_in_lens", None) - - offset = next(iter(kv_cache.values())).shape[1] if kv_cache else 0 - tgt, memory = x, xa - tgt[tgt == -1] = 0 - tgt = ( - self.token_embedding(tgt) - + self.positional_embedding[offset: offset + tgt.size(1)] - ) - # tgt = self.dropout(tgt) - - x = tgt.to(memory.dtype) - - if use_padmask and hlens is not None: - memory_mask = (~make_pad_mask(hlens)[:, None, :]).to(memory.device) - else: - memory_mask = None - - for layer, block in enumerate(self.blocks): - x = block(x, memory, mask=self.mask, memory_mask=memory_mask, is_pad_mask=False, is_pad_memory_mask=True) - - x = self.ln(x) - x = ( - x @ torch.transpose(self.token_embedding.weight.to(x.dtype), 0, 1) - ).float() - - return x + x: decoded token score before softmax (batch, maxlen_out, token) + if use_output_layer is True, + olens: (batch, ) + """ + # import pdb;pdb.set_trace() + use_padmask = self.use_padmask + hlens = kwargs.get("hlens", None) + + ys_in_lens = kwargs.get("ys_in_lens", None) + + offset = next(iter(kv_cache.values())).shape[1] if kv_cache else 0 + tgt, memory = x, xa + tgt[tgt == -1] = 0 + tgt = self.token_embedding(tgt) + self.positional_embedding[offset : offset + tgt.size(1)] + # tgt = self.dropout(tgt) + + x = tgt.to(memory.dtype) + + if use_padmask and hlens is not None: + memory_mask = (~make_pad_mask(hlens)[:, None, :]).to(memory.device) + else: + memory_mask = None + + for layer, block in enumerate(self.blocks): + x = block( + x, + memory, + mask=self.mask, + memory_mask=memory_mask, + is_pad_mask=False, + is_pad_memory_mask=True, + ) + + x = self.ln(x) + x = (x @ torch.transpose(self.token_embedding.weight.to(x.dtype), 0, 1)).float() + + return x class MultiHeadAttention(nn.Module): - def __init__(self, n_state: int, n_head: int): - super().__init__() - self.n_head = n_head - self.query = Linear(n_state, n_state) - self.key = Linear(n_state, n_state, bias=False) - self.value = Linear(n_state, n_state) - self.out = Linear(n_state, n_state) - - def forward( - self, - x: Tensor, - xa: Optional[Tensor] = None, - mask: Optional[Tensor] = None, - kv_cache: Optional[dict] = None, - **kwargs, - ): - is_pad_mask = kwargs.get("is_pad_mask", False) - - q = self.query(x) - - if kv_cache is None or xa is None or self.key not in kv_cache: - # hooks, if installed (i.e. kv_cache is not None), will prepend the cached kv tensors; - # otherwise, perform key/value projections for self- or cross-attention as usual. - k = self.key(x if xa is None else xa) - v = self.value(x if xa is None else xa) - else: - # for cross-attention, calculate keys and values once and reuse in subsequent calls. - k = kv_cache[self.key] - v = kv_cache[self.value] - - wv, qk = self.qkv_attention(q, k, v, mask, is_pad_mask=is_pad_mask) - return self.out(wv), qk - - def qkv_attention( - self, q: Tensor, k: Tensor, v: Tensor, mask: Optional[Tensor] = None, **kwargs, - ): - is_pad_mask = kwargs.get("is_pad_mask", False) - n_batch, n_ctx, n_state = q.shape - scale = (n_state // self.n_head) ** -0.25 - q = q.view(*q.shape[:2], self.n_head, -1).permute(0, 2, 1, 3) * scale - k = k.view(*k.shape[:2], self.n_head, -1).permute(0, 2, 3, 1) * scale - v = v.view(*v.shape[:2], self.n_head, -1).permute(0, 2, 1, 3) - - qk = q @ k - if mask is not None: - if not is_pad_mask: - qk = qk + mask[:n_ctx, :n_ctx] - else: - mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) - min_value = float( - np.finfo(torch.tensor(0, dtype=qk.dtype).numpy().dtype).min - ) - qk = qk.masked_fill(mask, min_value) - - qk = qk.float() - - w = F.softmax(qk, dim=-1).to(q.dtype) - if mask is not None and is_pad_mask: - w = w.masked_fill(mask, 0.0) - return (w @ v).permute(0, 2, 1, 3).flatten(start_dim=2), qk.detach() + def __init__(self, n_state: int, n_head: int): + super().__init__() + self.n_head = n_head + self.query = Linear(n_state, n_state) + self.key = Linear(n_state, n_state, bias=False) + self.value = Linear(n_state, n_state) + self.out = Linear(n_state, n_state) + + def forward( + self, + x: Tensor, + xa: Optional[Tensor] = None, + mask: Optional[Tensor] = None, + kv_cache: Optional[dict] = None, + **kwargs, + ): + is_pad_mask = kwargs.get("is_pad_mask", False) + + q = self.query(x) + + if kv_cache is None or xa is None or self.key not in kv_cache: + # hooks, if installed (i.e. kv_cache is not None), will prepend the cached kv tensors; + # otherwise, perform key/value projections for self- or cross-attention as usual. + k = self.key(x if xa is None else xa) + v = self.value(x if xa is None else xa) + else: + # for cross-attention, calculate keys and values once and reuse in subsequent calls. + k = kv_cache[self.key] + v = kv_cache[self.value] + + wv, qk = self.qkv_attention(q, k, v, mask, is_pad_mask=is_pad_mask) + return self.out(wv), qk + + def qkv_attention( + self, + q: Tensor, + k: Tensor, + v: Tensor, + mask: Optional[Tensor] = None, + **kwargs, + ): + is_pad_mask = kwargs.get("is_pad_mask", False) + n_batch, n_ctx, n_state = q.shape + scale = (n_state // self.n_head) ** -0.25 + q = q.view(*q.shape[:2], self.n_head, -1).permute(0, 2, 1, 3) * scale + k = k.view(*k.shape[:2], self.n_head, -1).permute(0, 2, 3, 1) * scale + v = v.view(*v.shape[:2], self.n_head, -1).permute(0, 2, 1, 3) + + qk = q @ k + if mask is not None: + if not is_pad_mask: + qk = qk + mask[:n_ctx, :n_ctx] + else: + mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) + min_value = float(np.finfo(torch.tensor(0, dtype=qk.dtype).numpy().dtype).min) + qk = qk.masked_fill(mask, min_value) + + qk = qk.float() + + w = F.softmax(qk, dim=-1).to(q.dtype) + if mask is not None and is_pad_mask: + w = w.masked_fill(mask, 0.0) + return (w @ v).permute(0, 2, 1, 3).flatten(start_dim=2), qk.detach() from funasr.models.sense_voice.rwkv_v6 import RWKVLayer @@ -157,149 +161,152 @@ from omegaconf import OmegaConf class ResidualAttentionBlockRWKV(nn.Module): - def __init__(self, n_state: int, n_head: int, cross_attention: bool = False, layer_id=0, **kwargs): - super().__init__() - - rwkv_cfg = kwargs.get("rwkv_cfg", {}) - args = OmegaConf.create(rwkv_cfg) - self.attn = RWKVLayer(args=args, layer_id=layer_id) - if args.get("datatype", "bf16") == "bf16": - self.attn.to(torch.bfloat16) - - self.ln0 = None - if layer_id == 0 and not args.get("ln0", True): - self.ln0 = LayerNorm(args.n_embd) - if args.get("init_rwkv", True): - print("init_rwkv") - layer_id = 0 - scale = ((1 + layer_id) / args.get("n_layer")) ** 0.7 - nn.init.constant_(self.ln0.weight, scale) - self.layer_id = layer_id - self.args = args - - self.ln1 = None - if not args.get("ln1", True): - self.ln1 = LayerNorm(args.n_embd) - # init - if args.get("init_rwkv", True): - print("init_rwkv") - scale = ((1 + layer_id) / args.get("n_layer")) ** 0.7 - nn.init.constant_(self.ln1.weight, scale) - - self.cross_attn = ( - MultiHeadAttention(n_state, n_head) if cross_attention else None - ) - self.cross_attn_ln = LayerNorm(n_state) if cross_attention else None - - n_mlp = n_state * 4 - self.mlp = nn.Sequential( - Linear(n_state, n_mlp), nn.GELU(), Linear(n_mlp, n_state) - ) - self.mlp_ln = LayerNorm(n_state) - - def forward( - self, - x: Tensor, - xa: Optional[Tensor] = None, - mask: Optional[Tensor] = None, - kv_cache: Optional[dict] = None, - **kwargs, - ): - is_pad_mask = kwargs.get("is_pad_mask", False) - is_pad_memory_mask = kwargs.get("is_pad_memory_mask", False) - - if self.layer_id == 0 and self.ln0 is not None: - x = self.ln0(x) - - if self.ln1 is None: - x = x + self.attn(x, mask=mask, kv_cache=kv_cache, is_pad_mask=is_pad_mask)[0] - else: - x = x + self.attn(self.ln1(x), mask=mask, kv_cache=kv_cache, is_pad_mask=is_pad_mask)[0] - - if self.cross_attn: - x = x + self.cross_attn(self.cross_attn_ln(x), xa, kv_cache=kv_cache, is_pad_mask=is_pad_memory_mask)[0] - x = x + self.mlp(self.mlp_ln(x)) - - return x + def __init__( + self, n_state: int, n_head: int, cross_attention: bool = False, layer_id=0, **kwargs + ): + super().__init__() + + rwkv_cfg = kwargs.get("rwkv_cfg", {}) + args = OmegaConf.create(rwkv_cfg) + self.attn = RWKVLayer(args=args, layer_id=layer_id) + if args.get("datatype", "bf16") == "bf16": + self.attn.to(torch.bfloat16) + + self.ln0 = None + if layer_id == 0 and not args.get("ln0", True): + self.ln0 = LayerNorm(args.n_embd) + if args.get("init_rwkv", True): + print("init_rwkv") + layer_id = 0 + scale = ((1 + layer_id) / args.get("n_layer")) ** 0.7 + nn.init.constant_(self.ln0.weight, scale) + self.layer_id = layer_id + self.args = args + + self.ln1 = None + if not args.get("ln1", True): + self.ln1 = LayerNorm(args.n_embd) + # init + if args.get("init_rwkv", True): + print("init_rwkv") + scale = ((1 + layer_id) / args.get("n_layer")) ** 0.7 + nn.init.constant_(self.ln1.weight, scale) + + self.cross_attn = MultiHeadAttention(n_state, n_head) if cross_attention else None + self.cross_attn_ln = LayerNorm(n_state) if cross_attention else None + + n_mlp = n_state * 4 + self.mlp = nn.Sequential(Linear(n_state, n_mlp), nn.GELU(), Linear(n_mlp, n_state)) + self.mlp_ln = LayerNorm(n_state) + + def forward( + self, + x: Tensor, + xa: Optional[Tensor] = None, + mask: Optional[Tensor] = None, + kv_cache: Optional[dict] = None, + **kwargs, + ): + is_pad_mask = kwargs.get("is_pad_mask", False) + is_pad_memory_mask = kwargs.get("is_pad_memory_mask", False) + + if self.layer_id == 0 and self.ln0 is not None: + x = self.ln0(x) + + if self.ln1 is None: + x = x + self.attn(x, mask=mask, kv_cache=kv_cache, is_pad_mask=is_pad_mask)[0] + else: + x = x + self.attn(self.ln1(x), mask=mask, kv_cache=kv_cache, is_pad_mask=is_pad_mask)[0] + + if self.cross_attn: + x = ( + x + + self.cross_attn( + self.cross_attn_ln(x), xa, kv_cache=kv_cache, is_pad_mask=is_pad_memory_mask + )[0] + ) + x = x + self.mlp(self.mlp_ln(x)) + + return x + @tables.register("decoder_classes", "SenseVoiceDecoder") class SenseVoiceDecoder(nn.Module): - def __init__( - self, n_vocab: int, n_ctx: int, n_state: int, n_head: int, n_layer: int, **kwargs - ): - super().__init__() - - self.token_embedding = nn.Embedding(n_vocab, n_state) - self.positional_embedding = nn.Parameter(torch.empty(n_ctx, n_state)) + def __init__(self, n_vocab: int, n_ctx: int, n_state: int, n_head: int, n_layer: int, **kwargs): + super().__init__() + self.token_embedding = nn.Embedding(n_vocab, n_state) + self.positional_embedding = nn.Parameter(torch.empty(n_ctx, n_state)) - self.blocks = nn.ModuleList( - [ - ResidualAttentionBlockRWKV(n_state, n_head, cross_attention=True, layer_id=i, **kwargs) - for i in range(n_layer) - ] - ) - self.ln = LayerNorm(n_state) - - mask = torch.empty(n_ctx, n_ctx).fill_(-np.inf).triu_(1) - self.register_buffer("mask", mask, persistent=False) - - self.use_padmask = kwargs.get("use_padmask", True) + self.blocks = nn.ModuleList( + [ + ResidualAttentionBlockRWKV( + n_state, n_head, cross_attention=True, layer_id=i, **kwargs + ) + for i in range(n_layer) + ] + ) + self.ln = LayerNorm(n_state) + mask = torch.empty(n_ctx, n_ctx).fill_(-np.inf).triu_(1) + self.register_buffer("mask", mask, persistent=False) - - def forward( - self, - x: torch.Tensor, - xa: torch.Tensor, - kv_cache: Optional[dict] = None, - **kwargs, - ): - """Forward decoder. - - Args: - hs_pad: encoded memory, float32 (batch, maxlen_in, feat) - hlens: (batch) - ys_in_pad: - input token ids, int64 (batch, maxlen_out) - if input_layer == "embed" - input tensor (batch, maxlen_out, #mels) in the other cases - ys_in_lens: (batch) - Returns: - (tuple): tuple containing: - - x: decoded token score before softmax (batch, maxlen_out, token) - if use_output_layer is True, - olens: (batch, ) - """ - # import pdb;pdb.set_trace() - use_padmask = self.use_padmask - hlens = kwargs.get("hlens", None) - - ys_in_lens = kwargs.get("ys_in_lens", None) - - offset = next(iter(kv_cache.values())).shape[1] if kv_cache else 0 - tgt, memory = x, xa - tgt[tgt == -1] = 0 - tgt = ( - self.token_embedding(tgt) - + self.positional_embedding[offset: offset + tgt.size(1)] - ) - # tgt = self.dropout(tgt) - - x = tgt.to(memory.dtype) - - if use_padmask and hlens is not None: - memory_mask = (~make_pad_mask(hlens)[:, None, :]).to(memory.device) - else: - memory_mask = None - - for layer, block in enumerate(self.blocks): - x = block(x, memory, mask=self.mask, memory_mask=memory_mask, is_pad_mask=False, is_pad_memory_mask=True) - - x = self.ln(x) - x = ( - x @ torch.transpose(self.token_embedding.weight.to(x.dtype), 0, 1) - ).float() - - return x + self.use_padmask = kwargs.get("use_padmask", True) + + def forward( + self, + x: torch.Tensor, + xa: torch.Tensor, + kv_cache: Optional[dict] = None, + **kwargs, + ): + """Forward decoder. + + Args: + hs_pad: encoded memory, float32 (batch, maxlen_in, feat) + hlens: (batch) + ys_in_pad: + input token ids, int64 (batch, maxlen_out) + if input_layer == "embed" + input tensor (batch, maxlen_out, #mels) in the other cases + ys_in_lens: (batch) + Returns: + (tuple): tuple containing: + + x: decoded token score before softmax (batch, maxlen_out, token) + if use_output_layer is True, + olens: (batch, ) + """ + # import pdb;pdb.set_trace() + use_padmask = self.use_padmask + hlens = kwargs.get("hlens", None) + + ys_in_lens = kwargs.get("ys_in_lens", None) + + offset = next(iter(kv_cache.values())).shape[1] if kv_cache else 0 + tgt, memory = x, xa + tgt[tgt == -1] = 0 + tgt = self.token_embedding(tgt) + self.positional_embedding[offset : offset + tgt.size(1)] + # tgt = self.dropout(tgt) + + x = tgt.to(memory.dtype) + + if use_padmask and hlens is not None: + memory_mask = (~make_pad_mask(hlens)[:, None, :]).to(memory.device) + else: + memory_mask = None + + for layer, block in enumerate(self.blocks): + x = block( + x, + memory, + mask=self.mask, + memory_mask=memory_mask, + is_pad_mask=False, + is_pad_memory_mask=True, + ) + + x = self.ln(x) + x = (x @ torch.transpose(self.token_embedding.weight.to(x.dtype), 0, 1)).float() + + return x diff --git a/funasr/models/sense_voice/encoder.py b/funasr/models/sense_voice/encoder.py index 3870c52cf..d464f1c70 100644 --- a/funasr/models/sense_voice/encoder.py +++ b/funasr/models/sense_voice/encoder.py @@ -8,60 +8,50 @@ from funasr.models.transformer.utils.nets_utils import make_pad_mask def sense_voice_encode_forward( - self, - x: torch.Tensor, - ilens: torch.Tensor = None, - **kwargs, + self, + x: torch.Tensor, + ilens: torch.Tensor = None, + **kwargs, ): - use_padmask = self.use_padmask - x = F.gelu(self.conv1(x)) - x = F.gelu(self.conv2(x)) - x = x.permute(0, 2, 1) - - n_frames = x.size(1) - max_pos = self.positional_embedding.size(0) - max_pos = n_frames if n_frames < max_pos else max_pos - x = (x[:, :max_pos, :] + self.positional_embedding[None, :max_pos, :]).to(x.dtype) - - - if ilens is not None: - if self.downsample_rate == 4: - olens = ( - 1 - + ( - ilens - - self.conv1.kernel_size[0] - + 2 * self.conv1.padding[0] - ) - // self.conv1.stride[0] - ) - else: - olens = ilens - olens = ( - 1 - + ( - olens - - self.conv2.kernel_size[0] - + 2 * self.conv2.padding[0] - ) - // self.conv2.stride[0] - ) - olens = torch.clamp(olens, max=max_pos) - else: - olens = None - - if use_padmask and olens is not None: - padding_mask = (~make_pad_mask(olens)[:, None, :]).to(x.device) - else: - padding_mask = None - - for layer, block in enumerate(self.blocks): - x = block(x, mask=padding_mask, is_pad_mask=True) - + use_padmask = self.use_padmask + x = F.gelu(self.conv1(x)) + x = F.gelu(self.conv2(x)) + x = x.permute(0, 2, 1) - x = self.ln_post(x) - - if ilens is None: - return x - else: - return x, olens + n_frames = x.size(1) + max_pos = self.positional_embedding.size(0) + max_pos = n_frames if n_frames < max_pos else max_pos + x = (x[:, :max_pos, :] + self.positional_embedding[None, :max_pos, :]).to(x.dtype) + + if ilens is not None: + if self.downsample_rate == 4: + olens = ( + 1 + + (ilens - self.conv1.kernel_size[0] + 2 * self.conv1.padding[0]) + // self.conv1.stride[0] + ) + else: + olens = ilens + olens = ( + 1 + + (olens - self.conv2.kernel_size[0] + 2 * self.conv2.padding[0]) + // self.conv2.stride[0] + ) + olens = torch.clamp(olens, max=max_pos) + else: + olens = None + + if use_padmask and olens is not None: + padding_mask = (~make_pad_mask(olens)[:, None, :]).to(x.device) + else: + padding_mask = None + + for layer, block in enumerate(self.blocks): + x = block(x, mask=padding_mask, is_pad_mask=True) + + x = self.ln_post(x) + + if ilens is None: + return x + else: + return x, olens diff --git a/funasr/models/sense_voice/model.py b/funasr/models/sense_voice/model.py index 41ad88060..4cbb490b4 100644 --- a/funasr/models/sense_voice/model.py +++ b/funasr/models/sense_voice/model.py @@ -18,32 +18,32 @@ from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank from funasr.register import tables - - @tables.register("model_classes", "SenseVoice") class SenseVoice(nn.Module): def __init__(self, *args, **kwargs): super().__init__() - + dims = kwargs.get("dims", {}) dims = whisper.model.ModelDimensions(**dims) model = whisper.model.Whisper(dims=dims) - + # encoder model.encoder.downsample_rate = kwargs.get("downsample_rate", 4) model.encoder.use_padmask = kwargs.get("use_padmask", True) from .encoder import sense_voice_encode_forward + model.encoder.forward = types.MethodType(sense_voice_encode_forward, model.encoder) - + # decoder model.decoder.use_padmask = kwargs.get("use_padmask", True) from .decoder import sense_voice_decode_forward + model.decoder.forward = types.MethodType(sense_voice_decode_forward, model.decoder) - + self.model = model - + self.encoder_output_size = self.model.dims.n_audio_state - + self.activation_checkpoint = kwargs.get("activation_checkpoint", False) self.ignore_id = kwargs.get("ignore_id", -1) self.vocab_size = kwargs.get("vocab_size", -1) @@ -54,14 +54,13 @@ class SenseVoice(nn.Module): smoothing=kwargs.get("lsm_weight", 0.0), normalize_length=self.length_normalized_loss, ) - + specaug = kwargs.get("specaug", None) if specaug is not None: specaug_class = tables.specaug_classes.get(specaug) specaug = specaug_class(**kwargs.get("specaug_conf", {})) self.specaug = specaug - def forward( self, speech: torch.Tensor, @@ -71,19 +70,22 @@ class SenseVoice(nn.Module): **kwargs, ): target_mask = kwargs.get("target_mask", None) - + # import pdb; # pdb.set_trace() if len(text_lengths.size()) > 1: text_lengths = text_lengths[:, 0] if len(speech_lengths.size()) > 1: speech_lengths = speech_lengths[:, 0] - + batch_size = speech.shape[0] if self.activation_checkpoint: from torch.utils.checkpoint import checkpoint - encoder_out, encoder_out_lens = checkpoint(self.encode, speech, speech_lengths, use_reentrant=False) + + encoder_out, encoder_out_lens = checkpoint( + self.encode, speech, speech_lengths, use_reentrant=False + ) else: encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) @@ -95,7 +97,7 @@ class SenseVoice(nn.Module): stats["acc"] = acc_att stats["loss"] = torch.clone(loss.detach()) stats["batch_size"] = batch_size - + # force_gatherable: to-device and to-tensor if scalar for DataParallel if self.length_normalized_loss: batch_size = int((text_lengths + 1).sum()) @@ -103,8 +105,11 @@ class SenseVoice(nn.Module): return loss, stats, weight def encode( - self, speech: torch.Tensor, speech_lengths: torch.Tensor, **kwargs, - ) : + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + **kwargs, + ): """Encoder. Note that this method is used by asr_inference.py Args: speech: (Batch, Length, ...) @@ -117,224 +122,11 @@ class SenseVoice(nn.Module): if self.specaug is not None and self.training: speech, speech_lengths = self.specaug(speech, speech_lengths) - # Forward encoder encoder_out, encoder_out_lens = self.model.encoder(speech.permute(0, 2, 1), speech_lengths) - + return encoder_out, encoder_out_lens - - def _calc_att_loss( - self, - encoder_out: torch.Tensor, - encoder_out_lens: torch.Tensor, - ys_pad: torch.Tensor, - ys_pad_lens: torch.Tensor, - **kwargs, - ): - target_mask = kwargs.get("target_mask", None) - stats = {} - - # 1. Forward decoder - decoder_out = self.model.decoder( - x=ys_pad, xa=encoder_out, hlens=encoder_out_lens, ys_in_lens=ys_pad_lens - ) - - # 2. Compute attention loss - mask = torch.ones_like(ys_pad) * (-1) - ys_pad_mask = (ys_pad * target_mask + mask * (1-target_mask)).to(torch.int64) - ys_pad_mask[ys_pad_mask == 0] = -1 - loss_att = self.criterion_att(decoder_out[:, :-1, :], ys_pad_mask[:, 1:]) - - with torch.no_grad(): - preds = torch.argmax(decoder_out, -1) - acc_att = compute_accuracy(preds[:, :-1], ys_pad_mask[:, 1:], ignore_label=self.ignore_id) - - return loss_att, acc_att, None, None - - - def inference(self, - data_in, - data_lengths=None, - key: list = None, - tokenizer=None, - frontend=None, - **kwargs, - ): - if kwargs.get("batch_size", 1) > 1: - raise NotImplementedError("batch decoding is not implemented") - - if frontend is None and not hasattr(self, "frontend"): - frontend_class = tables.frontend_classes.get("WhisperFrontend") - frontend = frontend_class(n_mels=self.model.dims.n_mels, do_pad_trim=kwargs.get("do_pad_trim", True)) - self.frontend = frontend - else: - frontend = frontend if frontend is not None else self.frontend - - meta_data = {} - if isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank": # fbank - speech, speech_lengths = data_in, data_lengths - if len(speech.shape) < 3: - speech = speech[None, :, :] - if speech_lengths is None: - speech_lengths = speech.shape[1] - else: - # extract fbank feats - time1 = time.perf_counter() - audio_sample_list = load_audio_text_image_video(data_in, fs=frontend.fs if hasattr(frontend, "fs") else 16000, audio_fs=kwargs.get("fs", 16000), - data_type=kwargs.get("data_type", "sound"), - tokenizer=tokenizer) - time2 = time.perf_counter() - meta_data["load_data"] = f"{time2 - time1:0.3f}" - speech, speech_lengths = extract_fbank(audio_sample_list, data_type=kwargs.get("data_type", "sound"), - frontend=frontend) - time3 = time.perf_counter() - meta_data["extract_feat"] = f"{time3 - time2:0.3f}" - frame_shift = frontend.frame_shift if hasattr(frontend, "frame_shift") else 10 - lfr_n = frontend.lfr_n if hasattr(frontend, "lfr_n") else 1 - meta_data["batch_data_time"] = speech_lengths.sum().item() * frame_shift * lfr_n / 1000 - - speech = speech.to(device=kwargs["device"])[0, :, :] - speech_lengths = speech_lengths.to(device=kwargs["device"]) - - DecodingOptions = kwargs.get("DecodingOptions", {}) - task = DecodingOptions.get("task", "ASR") - if isinstance(task, str): - task = [task] - task = "".join([f"<|{x}|>" for x in task]) - initial_prompt = kwargs.get("initial_prompt", f"<|startoftranscript|>{task}") - DecodingOptions["initial_prompt"] = initial_prompt - - language = DecodingOptions.get("language", None) - language = None if language == "auto" else language - DecodingOptions["language"] = language - - DecodingOptions["vocab_path"] = kwargs["tokenizer_conf"].get("vocab_path", None) - - - if "without_timestamps" not in DecodingOptions: - DecodingOptions["without_timestamps"] = True - - - options = whisper.DecodingOptions(**DecodingOptions) - - result = whisper.decode(self.model, speech, options) - text = f"{result.text}" - results = [] - result_i = {"key": key[0], "text": text} - - results.append(result_i) - - return results, meta_data - - -@tables.register("model_classes", "SenseVoiceRWKV") -class SenseVoiceRWKV(nn.Module): - def __init__(self, *args, **kwargs): - super().__init__() - - dims = kwargs.get("dims", {}) - dims = whisper.model.ModelDimensions(**dims) - model = whisper.model.Whisper(dims=dims) - - # encoder - model.encoder.downsample_rate = kwargs.get("downsample_rate", 4) - model.encoder.use_padmask = kwargs.get("use_padmask", True) - from .encoder import sense_voice_encode_forward - model.encoder.forward = types.MethodType(sense_voice_encode_forward, model.encoder) - - # decoder - del model.decoder - decoder = kwargs.get("decoder", "SenseVoiceDecoder") - decoder_class = tables.decoder_classes.get(decoder) - decoder = decoder_class(n_vocab=dims.n_vocab, - n_ctx=dims.n_text_ctx, - n_state=dims.n_text_state, - n_head=dims.n_text_head, - n_layer=dims.n_text_layer, - **kwargs.get("decoder_conf")) - model.decoder = decoder - - self.model = model - - self.encoder_output_size = self.model.dims.n_audio_state - - self.activation_checkpoint = kwargs.get("activation_checkpoint", False) - self.ignore_id = kwargs.get("ignore_id", -1) - self.vocab_size = kwargs.get("vocab_size", -1) - self.length_normalized_loss = kwargs.get("length_normalized_loss", True) - self.criterion_att = LabelSmoothingLoss( - size=self.vocab_size, - padding_idx=self.ignore_id, - smoothing=kwargs.get("lsm_weight", 0.0), - normalize_length=self.length_normalized_loss, - ) - - specaug = kwargs.get("specaug", None) - if specaug is not None: - specaug_class = tables.specaug_classes.get(specaug) - specaug = specaug_class(**kwargs.get("specaug_conf", {})) - self.specaug = specaug - - def forward( - self, - speech: torch.Tensor, - speech_lengths: torch.Tensor, - text: torch.Tensor, - text_lengths: torch.Tensor, - **kwargs, - ): - target_mask = kwargs.get("target_mask", None) - - # import pdb; - # pdb.set_trace() - if len(text_lengths.size()) > 1: - text_lengths = text_lengths[:, 0] - if len(speech_lengths.size()) > 1: - speech_lengths = speech_lengths[:, 0] - - batch_size = speech.shape[0] - - if self.activation_checkpoint: - from torch.utils.checkpoint import checkpoint - encoder_out, encoder_out_lens = checkpoint(self.encode, speech, speech_lengths, use_reentrant=False) - else: - encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) - - loss_att, acc_att, cer_att, wer_att = self._calc_att_loss( - encoder_out, encoder_out_lens, text, text_lengths, target_mask=target_mask - ) - loss = loss_att - stats = {} - stats["acc"] = acc_att - stats["loss"] = torch.clone(loss.detach()) - stats["batch_size"] = batch_size - - # force_gatherable: to-device and to-tensor if scalar for DataParallel - if self.length_normalized_loss: - batch_size = int((text_lengths + 1).sum()) - loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) - return loss, stats, weight - - def encode( - self, speech: torch.Tensor, speech_lengths: torch.Tensor, **kwargs, - ): - """Encoder. Note that this method is used by asr_inference.py - Args: - speech: (Batch, Length, ...) - speech_lengths: (Batch, ) - ind: int - """ - with autocast(False): - # Data augmentation - if self.specaug is not None and self.training: - speech, speech_lengths = self.specaug(speech, speech_lengths) - - # Forward encoder - encoder_out, encoder_out_lens = self.model.encoder(speech.permute(0, 2, 1), speech_lengths) - - return encoder_out, encoder_out_lens - def _calc_att_loss( self, encoder_out: torch.Tensor, @@ -345,44 +137,51 @@ class SenseVoiceRWKV(nn.Module): ): target_mask = kwargs.get("target_mask", None) stats = {} - + # 1. Forward decoder decoder_out = self.model.decoder( x=ys_pad, xa=encoder_out, hlens=encoder_out_lens, ys_in_lens=ys_pad_lens ) - + # 2. Compute attention loss mask = torch.ones_like(ys_pad) * (-1) ys_pad_mask = (ys_pad * target_mask + mask * (1 - target_mask)).to(torch.int64) ys_pad_mask[ys_pad_mask == 0] = -1 loss_att = self.criterion_att(decoder_out[:, :-1, :], ys_pad_mask[:, 1:]) - + with torch.no_grad(): preds = torch.argmax(decoder_out, -1) - acc_att = compute_accuracy(preds[:, :-1], ys_pad_mask[:, 1:], ignore_label=self.ignore_id) - + acc_att = compute_accuracy( + preds[:, :-1], ys_pad_mask[:, 1:], ignore_label=self.ignore_id + ) + return loss_att, acc_att, None, None - - def inference(self, - data_in, - data_lengths=None, - key: list = None, - tokenizer=None, - frontend=None, - **kwargs, - ): + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): if kwargs.get("batch_size", 1) > 1: raise NotImplementedError("batch decoding is not implemented") - + if frontend is None and not hasattr(self, "frontend"): frontend_class = tables.frontend_classes.get("WhisperFrontend") - frontend = frontend_class(n_mels=self.model.dims.n_mels, do_pad_trim=kwargs.get("do_pad_trim", True)) + frontend = frontend_class( + n_mels=self.model.dims.n_mels, do_pad_trim=kwargs.get("do_pad_trim", True) + ) self.frontend = frontend else: frontend = frontend if frontend is not None else self.frontend - + meta_data = {} - if isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank": # fbank + if ( + isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank" + ): # fbank speech, speech_lengths = data_in, data_lengths if len(speech.shape) < 3: speech = speech[None, :, :] @@ -391,24 +190,27 @@ class SenseVoiceRWKV(nn.Module): else: # extract fbank feats time1 = time.perf_counter() - audio_sample_list = load_audio_text_image_video(data_in, - fs=frontend.fs if hasattr(frontend, "fs") else 16000, - audio_fs=kwargs.get("fs", 16000), - data_type=kwargs.get("data_type", "sound"), - tokenizer=tokenizer) + audio_sample_list = load_audio_text_image_video( + data_in, + fs=frontend.fs if hasattr(frontend, "fs") else 16000, + audio_fs=kwargs.get("fs", 16000), + data_type=kwargs.get("data_type", "sound"), + tokenizer=tokenizer, + ) time2 = time.perf_counter() meta_data["load_data"] = f"{time2 - time1:0.3f}" - speech, speech_lengths = extract_fbank(audio_sample_list, data_type=kwargs.get("data_type", "sound"), - frontend=frontend) + speech, speech_lengths = extract_fbank( + audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend + ) time3 = time.perf_counter() meta_data["extract_feat"] = f"{time3 - time2:0.3f}" frame_shift = frontend.frame_shift if hasattr(frontend, "frame_shift") else 10 lfr_n = frontend.lfr_n if hasattr(frontend, "lfr_n") else 1 meta_data["batch_data_time"] = speech_lengths.sum().item() * frame_shift * lfr_n / 1000 - + speech = speech.to(device=kwargs["device"])[0, :, :] speech_lengths = speech_lengths.to(device=kwargs["device"]) - + DecodingOptions = kwargs.get("DecodingOptions", {}) task = DecodingOptions.get("task", "ASR") if isinstance(task, str): @@ -416,23 +218,252 @@ class SenseVoiceRWKV(nn.Module): task = "".join([f"<|{x}|>" for x in task]) initial_prompt = kwargs.get("initial_prompt", f"<|startoftranscript|>{task}") DecodingOptions["initial_prompt"] = initial_prompt - + language = DecodingOptions.get("language", None) language = None if language == "auto" else language DecodingOptions["language"] = language - + DecodingOptions["vocab_path"] = kwargs["tokenizer_conf"].get("vocab_path", None) - + if "without_timestamps" not in DecodingOptions: DecodingOptions["without_timestamps"] = True - + options = whisper.DecodingOptions(**DecodingOptions) - + result = whisper.decode(self.model, speech, options) text = f"{result.text}" results = [] result_i = {"key": key[0], "text": text} - + results.append(result_i) - + + return results, meta_data + + +@tables.register("model_classes", "SenseVoiceRWKV") +class SenseVoiceRWKV(nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + + dims = kwargs.get("dims", {}) + dims = whisper.model.ModelDimensions(**dims) + model = whisper.model.Whisper(dims=dims) + + # encoder + model.encoder.downsample_rate = kwargs.get("downsample_rate", 4) + model.encoder.use_padmask = kwargs.get("use_padmask", True) + from .encoder import sense_voice_encode_forward + + model.encoder.forward = types.MethodType(sense_voice_encode_forward, model.encoder) + + # decoder + del model.decoder + decoder = kwargs.get("decoder", "SenseVoiceDecoder") + decoder_class = tables.decoder_classes.get(decoder) + decoder = decoder_class( + n_vocab=dims.n_vocab, + n_ctx=dims.n_text_ctx, + n_state=dims.n_text_state, + n_head=dims.n_text_head, + n_layer=dims.n_text_layer, + **kwargs.get("decoder_conf"), + ) + model.decoder = decoder + + self.model = model + + self.encoder_output_size = self.model.dims.n_audio_state + + self.activation_checkpoint = kwargs.get("activation_checkpoint", False) + self.ignore_id = kwargs.get("ignore_id", -1) + self.vocab_size = kwargs.get("vocab_size", -1) + self.length_normalized_loss = kwargs.get("length_normalized_loss", True) + self.criterion_att = LabelSmoothingLoss( + size=self.vocab_size, + padding_idx=self.ignore_id, + smoothing=kwargs.get("lsm_weight", 0.0), + normalize_length=self.length_normalized_loss, + ) + + specaug = kwargs.get("specaug", None) + if specaug is not None: + specaug_class = tables.specaug_classes.get(specaug) + specaug = specaug_class(**kwargs.get("specaug_conf", {})) + self.specaug = specaug + + def forward( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + text: torch.Tensor, + text_lengths: torch.Tensor, + **kwargs, + ): + target_mask = kwargs.get("target_mask", None) + + # import pdb; + # pdb.set_trace() + if len(text_lengths.size()) > 1: + text_lengths = text_lengths[:, 0] + if len(speech_lengths.size()) > 1: + speech_lengths = speech_lengths[:, 0] + + batch_size = speech.shape[0] + + if self.activation_checkpoint: + from torch.utils.checkpoint import checkpoint + + encoder_out, encoder_out_lens = checkpoint( + self.encode, speech, speech_lengths, use_reentrant=False + ) + else: + encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) + + loss_att, acc_att, cer_att, wer_att = self._calc_att_loss( + encoder_out, encoder_out_lens, text, text_lengths, target_mask=target_mask + ) + loss = loss_att + stats = {} + stats["acc"] = acc_att + stats["loss"] = torch.clone(loss.detach()) + stats["batch_size"] = batch_size + + # force_gatherable: to-device and to-tensor if scalar for DataParallel + if self.length_normalized_loss: + batch_size = int((text_lengths + 1).sum()) + loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) + return loss, stats, weight + + def encode( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + **kwargs, + ): + """Encoder. Note that this method is used by asr_inference.py + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + ind: int + """ + with autocast(False): + # Data augmentation + if self.specaug is not None and self.training: + speech, speech_lengths = self.specaug(speech, speech_lengths) + + # Forward encoder + encoder_out, encoder_out_lens = self.model.encoder(speech.permute(0, 2, 1), speech_lengths) + + return encoder_out, encoder_out_lens + + def _calc_att_loss( + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_pad_lens: torch.Tensor, + **kwargs, + ): + target_mask = kwargs.get("target_mask", None) + stats = {} + + # 1. Forward decoder + decoder_out = self.model.decoder( + x=ys_pad, xa=encoder_out, hlens=encoder_out_lens, ys_in_lens=ys_pad_lens + ) + + # 2. Compute attention loss + mask = torch.ones_like(ys_pad) * (-1) + ys_pad_mask = (ys_pad * target_mask + mask * (1 - target_mask)).to(torch.int64) + ys_pad_mask[ys_pad_mask == 0] = -1 + loss_att = self.criterion_att(decoder_out[:, :-1, :], ys_pad_mask[:, 1:]) + + with torch.no_grad(): + preds = torch.argmax(decoder_out, -1) + acc_att = compute_accuracy( + preds[:, :-1], ys_pad_mask[:, 1:], ignore_label=self.ignore_id + ) + + return loss_att, acc_att, None, None + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): + if kwargs.get("batch_size", 1) > 1: + raise NotImplementedError("batch decoding is not implemented") + + if frontend is None and not hasattr(self, "frontend"): + frontend_class = tables.frontend_classes.get("WhisperFrontend") + frontend = frontend_class( + n_mels=self.model.dims.n_mels, do_pad_trim=kwargs.get("do_pad_trim", True) + ) + self.frontend = frontend + else: + frontend = frontend if frontend is not None else self.frontend + + meta_data = {} + if ( + isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank" + ): # fbank + speech, speech_lengths = data_in, data_lengths + if len(speech.shape) < 3: + speech = speech[None, :, :] + if speech_lengths is None: + speech_lengths = speech.shape[1] + else: + # extract fbank feats + time1 = time.perf_counter() + audio_sample_list = load_audio_text_image_video( + data_in, + fs=frontend.fs if hasattr(frontend, "fs") else 16000, + audio_fs=kwargs.get("fs", 16000), + data_type=kwargs.get("data_type", "sound"), + tokenizer=tokenizer, + ) + time2 = time.perf_counter() + meta_data["load_data"] = f"{time2 - time1:0.3f}" + speech, speech_lengths = extract_fbank( + audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend + ) + time3 = time.perf_counter() + meta_data["extract_feat"] = f"{time3 - time2:0.3f}" + frame_shift = frontend.frame_shift if hasattr(frontend, "frame_shift") else 10 + lfr_n = frontend.lfr_n if hasattr(frontend, "lfr_n") else 1 + meta_data["batch_data_time"] = speech_lengths.sum().item() * frame_shift * lfr_n / 1000 + + speech = speech.to(device=kwargs["device"])[0, :, :] + speech_lengths = speech_lengths.to(device=kwargs["device"]) + + DecodingOptions = kwargs.get("DecodingOptions", {}) + task = DecodingOptions.get("task", "ASR") + if isinstance(task, str): + task = [task] + task = "".join([f"<|{x}|>" for x in task]) + initial_prompt = kwargs.get("initial_prompt", f"<|startoftranscript|>{task}") + DecodingOptions["initial_prompt"] = initial_prompt + + language = DecodingOptions.get("language", None) + language = None if language == "auto" else language + DecodingOptions["language"] = language + + DecodingOptions["vocab_path"] = kwargs["tokenizer_conf"].get("vocab_path", None) + + if "without_timestamps" not in DecodingOptions: + DecodingOptions["without_timestamps"] = True + + options = whisper.DecodingOptions(**DecodingOptions) + + result = whisper.decode(self.model, speech, options) + text = f"{result.text}" + results = [] + result_i = {"key": key[0], "text": text} + + results.append(result_i) + return results, meta_data diff --git a/funasr/models/sense_voice/rwkv_v6.py b/funasr/models/sense_voice/rwkv_v6.py index 6eb53fcde..36269a108 100644 --- a/funasr/models/sense_voice/rwkv_v6.py +++ b/funasr/models/sense_voice/rwkv_v6.py @@ -4,22 +4,22 @@ import os, math, gc, importlib import torch + # torch._C._jit_set_profiling_executor(True) # torch._C._jit_set_profiling_mode(True) import torch.nn as nn from torch.nn import functional as F - def __nop(ob): - return ob + return ob MyModule = nn.Module MyFunction = __nop if "RWKV_JIT_ON" in os.environ and os.environ["RWKV_JIT_ON"] == "1": - MyModule = torch.jit.ScriptModule - MyFunction = torch.jit.script_method + MyModule = torch.jit.ScriptModule + MyFunction = torch.jit.script_method ######################################################################################################## # CUDA Kernel @@ -27,399 +27,445 @@ if "RWKV_JIT_ON" in os.environ and os.environ["RWKV_JIT_ON"] == "1": wkv6_cuda = None -def load_rwkv_kernel(HEAD_SIZE: int=64, RWKV_CTXLEN: int=512,): - from torch.utils.cpp_extension import load - global wkv6_cuda - - - if wkv6_cuda is not None: - return - - absolute_file_path = os.path.abspath(__file__) - cur_dir = os.path.dirname(absolute_file_path) - wkv6_cuda = load(name="wkv6", sources=[f"{cur_dir}/cuda/wkv6_op.cpp", f"{cur_dir}/cuda/wkv6_cuda.cu"], - verbose=True, extra_cuda_cflags=["-res-usage", "--use_fast_math", "-O3", "-Xptxas -O3", - "--extra-device-vectorization", f"-D_N_={HEAD_SIZE}", - f"-D_T_={RWKV_CTXLEN}"]) + +def load_rwkv_kernel( + HEAD_SIZE: int = 64, + RWKV_CTXLEN: int = 512, +): + from torch.utils.cpp_extension import load + + global wkv6_cuda + + if wkv6_cuda is not None: + return + + absolute_file_path = os.path.abspath(__file__) + cur_dir = os.path.dirname(absolute_file_path) + wkv6_cuda = load( + name="wkv6", + sources=[f"{cur_dir}/cuda/wkv6_op.cpp", f"{cur_dir}/cuda/wkv6_cuda.cu"], + verbose=True, + extra_cuda_cflags=[ + "-res-usage", + "--use_fast_math", + "-O3", + "-Xptxas -O3", + "--extra-device-vectorization", + f"-D_N_={HEAD_SIZE}", + f"-D_T_={RWKV_CTXLEN}", + ], + ) + # dtype = torch.float dtype = torch.bfloat16 + + class WKV_6(torch.autograd.Function): - @staticmethod - def forward(ctx, B, T, C, H, r, k, v, w, u): - with torch.no_grad(): - # assert r.dtype == torch.bfloat16 - # assert k.dtype == torch.bfloat16 - # assert v.dtype == torch.bfloat16 - # assert w.dtype == torch.bfloat16 - # assert u.dtype == torch.bfloat16 - # assert HEAD_SIZE == C // H - ctx.B = B - ctx.T = T - ctx.C = C - ctx.H = H - assert r.is_contiguous() - assert k.is_contiguous() - assert v.is_contiguous() - assert w.is_contiguous() - assert u.is_contiguous() - ew = (-torch.exp(w.float())).contiguous() - ctx.save_for_backward(r, k, v, ew, u) - y = torch.empty((B, T, C), device=r.device, dtype=dtype, - memory_format=torch.contiguous_format) # .uniform_(-100, 100) - wkv6_cuda.forward(B, T, C, H, r, k, v, ew, u, y) - return y - - @staticmethod - def backward(ctx, gy): - with torch.no_grad(): - # assert gy.dtype == torch.bfloat16 - B = ctx.B - T = ctx.T - C = ctx.C - H = ctx.H - assert gy.is_contiguous() - r, k, v, ew, u = ctx.saved_tensors - gr = torch.empty((B, T, C), device=gy.device, requires_grad=False, dtype=dtype, - memory_format=torch.contiguous_format) # .uniform_(-100, 100) - gk = torch.empty((B, T, C), device=gy.device, requires_grad=False, dtype=dtype, - memory_format=torch.contiguous_format) # .uniform_(-100, 100) - gv = torch.empty((B, T, C), device=gy.device, requires_grad=False, dtype=dtype, - memory_format=torch.contiguous_format) # .uniform_(-100, 100) - gw = torch.empty((B, T, C), device=gy.device, requires_grad=False, dtype=dtype, - memory_format=torch.contiguous_format) # .uniform_(-100, 100) - gu = torch.empty((B, C), device=gy.device, requires_grad=False, dtype=dtype, - memory_format=torch.contiguous_format) # .uniform_(-100, 100) - wkv6_cuda.backward(B, T, C, H, r, k, v, ew, u, gy, gr, gk, gv, gw, gu) - gu = torch.sum(gu, 0).view(H, C // H) - return (None, None, None, None, gr, gk, gv, gw, gu) + @staticmethod + def forward(ctx, B, T, C, H, r, k, v, w, u): + with torch.no_grad(): + # assert r.dtype == torch.bfloat16 + # assert k.dtype == torch.bfloat16 + # assert v.dtype == torch.bfloat16 + # assert w.dtype == torch.bfloat16 + # assert u.dtype == torch.bfloat16 + # assert HEAD_SIZE == C // H + ctx.B = B + ctx.T = T + ctx.C = C + ctx.H = H + assert r.is_contiguous() + assert k.is_contiguous() + assert v.is_contiguous() + assert w.is_contiguous() + assert u.is_contiguous() + ew = (-torch.exp(w.float())).contiguous() + ctx.save_for_backward(r, k, v, ew, u) + y = torch.empty( + (B, T, C), device=r.device, dtype=dtype, memory_format=torch.contiguous_format + ) # .uniform_(-100, 100) + wkv6_cuda.forward(B, T, C, H, r, k, v, ew, u, y) + return y + + @staticmethod + def backward(ctx, gy): + with torch.no_grad(): + # assert gy.dtype == torch.bfloat16 + B = ctx.B + T = ctx.T + C = ctx.C + H = ctx.H + assert gy.is_contiguous() + r, k, v, ew, u = ctx.saved_tensors + gr = torch.empty( + (B, T, C), + device=gy.device, + requires_grad=False, + dtype=dtype, + memory_format=torch.contiguous_format, + ) # .uniform_(-100, 100) + gk = torch.empty( + (B, T, C), + device=gy.device, + requires_grad=False, + dtype=dtype, + memory_format=torch.contiguous_format, + ) # .uniform_(-100, 100) + gv = torch.empty( + (B, T, C), + device=gy.device, + requires_grad=False, + dtype=dtype, + memory_format=torch.contiguous_format, + ) # .uniform_(-100, 100) + gw = torch.empty( + (B, T, C), + device=gy.device, + requires_grad=False, + dtype=dtype, + memory_format=torch.contiguous_format, + ) # .uniform_(-100, 100) + gu = torch.empty( + (B, C), + device=gy.device, + requires_grad=False, + dtype=dtype, + memory_format=torch.contiguous_format, + ) # .uniform_(-100, 100) + wkv6_cuda.backward(B, T, C, H, r, k, v, ew, u, gy, gr, gk, gv, gw, gu) + gu = torch.sum(gu, 0).view(H, C // H) + return (None, None, None, None, gr, gk, gv, gw, gu) def RUN_CUDA_RWKV6(B, T, C, H, r, k, v, w, u): - return WKV_6.apply(B, T, C, H, r, k, v, w, u) + return WKV_6.apply(B, T, C, H, r, k, v, w, u) class RWKV_Tmix_x060(MyModule): - def __init__(self, args, layer_id): - super().__init__() - self.args = args - - load_rwkv_kernel(args.head_size_a, args.ctx_len) - - self.layer_id = layer_id - - self.head_size = args.head_size_a - self.n_head = args.dim_att // self.head_size - assert args.dim_att % self.n_head == 0 - - with torch.no_grad(): - ratio_0_to_1 = layer_id / (args.n_layer - 1) # 0 to 1 - ratio_1_to_almost0 = 1.0 - (layer_id / args.n_layer) # 1 to ~0 - ddd = torch.ones(1, 1, args.n_embd) - for i in range(args.n_embd): - ddd[0, 0, i] = i / args.n_embd - - # fancy time_mix - self.time_maa_x = nn.Parameter(1.0 - torch.pow(ddd, ratio_1_to_almost0)) - self.time_maa_w = nn.Parameter(1.0 - torch.pow(ddd, ratio_1_to_almost0)) - self.time_maa_k = nn.Parameter(1.0 - torch.pow(ddd, ratio_1_to_almost0)) - self.time_maa_v = nn.Parameter(1.0 - (torch.pow(ddd, ratio_1_to_almost0) + 0.3 * ratio_0_to_1)) - self.time_maa_r = nn.Parameter(1.0 - torch.pow(ddd, 0.5 * ratio_1_to_almost0)) - self.time_maa_g = nn.Parameter(1.0 - torch.pow(ddd, 0.5 * ratio_1_to_almost0)) - - D_MIX_LORA = 32 # generate TIME_MIX for w,k,v,r,g - self.time_maa_w1 = nn.Parameter(torch.zeros(args.n_embd, D_MIX_LORA * 5)) - self.time_maa_w2 = nn.Parameter(torch.zeros(5, D_MIX_LORA, args.n_embd).uniform_(-0.01, 0.01)) - - # fancy time_decay - decay_speed = torch.ones(args.dim_att) - for n in range(args.dim_att): - decay_speed[n] = -6 + 5 * (n / (args.dim_att - 1)) ** (0.7 + 1.3 * ratio_0_to_1) - self.time_decay = nn.Parameter(decay_speed.reshape(1, 1, args.dim_att)) - - D_DECAY_LORA = 64 - self.time_decay_w1 = nn.Parameter(torch.zeros(args.n_embd, D_DECAY_LORA)) - self.time_decay_w2 = nn.Parameter(torch.zeros(D_DECAY_LORA, args.dim_att).uniform_(-0.01, 0.01)) - - tmp = torch.zeros(args.dim_att) - for n in range(args.dim_att): - zigzag = ((n + 1) % 3 - 1) * 0.1 - tmp[n] = ratio_0_to_1 * (1 - (n / (args.dim_att - 1))) + zigzag - - self.time_faaaa = nn.Parameter(tmp.reshape(self.n_head, self.head_size)) - - self.time_shift = nn.ZeroPad2d((0, 0, 1, -1)) - self.receptance = nn.Linear(args.n_embd, args.dim_att, bias=False) - self.key = nn.Linear(args.n_embd, args.dim_att, bias=False) - - self.value = nn.Linear(args.n_embd, args.dim_att, bias=False) - self.output = nn.Linear(args.dim_att, args.n_embd, bias=False) - self.gate = nn.Linear(args.n_embd, args.dim_att, bias=False) - self.ln_x = nn.GroupNorm(self.n_head, args.dim_att, eps=(1e-5) * (args.head_size_divisor ** 2)) - - @MyFunction - def jit_func(self, x): - B, T, C = x.size() - - xx = self.time_shift(x) - x - - xxx = x + xx * self.time_maa_x - xxx = torch.tanh(xxx @ self.time_maa_w1).view(B * T, 5, -1).transpose(0, 1) - xxx = torch.bmm(xxx, self.time_maa_w2).view(5, B, T, -1) - mw, mk, mv, mr, mg = xxx.unbind(dim=0) - - xw = x + xx * (self.time_maa_w + mw) - xk = x + xx * (self.time_maa_k + mk) - xv = x + xx * (self.time_maa_v + mv) - xr = x + xx * (self.time_maa_r + mr) - xg = x + xx * (self.time_maa_g + mg) - - r = self.receptance(xr) - k = self.key(xk) - v = self.value(xv) - g = F.silu(self.gate(xg)) - - ww = torch.tanh(xw @ self.time_decay_w1) @ self.time_decay_w2 - w = self.time_decay + ww - - return r, k, v, g, w - - @MyFunction - def jit_func_2(self, x, g): - B, T, C = x.size() - x = x.view(B * T, C) - - x = self.ln_x(x).view(B, T, C) - x = self.output(x * g) - return x - - def forward(self, x): - B, T, C = x.size() - H = self.n_head - - r, k, v, g, w = self.jit_func(x) - x = RUN_CUDA_RWKV6(B, T, C, H, r, k, v, w, u=self.time_faaaa) - - return self.jit_func_2(x, g) + def __init__(self, args, layer_id): + super().__init__() + self.args = args + + load_rwkv_kernel(args.head_size_a, args.ctx_len) + + self.layer_id = layer_id + + self.head_size = args.head_size_a + self.n_head = args.dim_att // self.head_size + assert args.dim_att % self.n_head == 0 + + with torch.no_grad(): + ratio_0_to_1 = layer_id / (args.n_layer - 1) # 0 to 1 + ratio_1_to_almost0 = 1.0 - (layer_id / args.n_layer) # 1 to ~0 + ddd = torch.ones(1, 1, args.n_embd) + for i in range(args.n_embd): + ddd[0, 0, i] = i / args.n_embd + + # fancy time_mix + self.time_maa_x = nn.Parameter(1.0 - torch.pow(ddd, ratio_1_to_almost0)) + self.time_maa_w = nn.Parameter(1.0 - torch.pow(ddd, ratio_1_to_almost0)) + self.time_maa_k = nn.Parameter(1.0 - torch.pow(ddd, ratio_1_to_almost0)) + self.time_maa_v = nn.Parameter( + 1.0 - (torch.pow(ddd, ratio_1_to_almost0) + 0.3 * ratio_0_to_1) + ) + self.time_maa_r = nn.Parameter(1.0 - torch.pow(ddd, 0.5 * ratio_1_to_almost0)) + self.time_maa_g = nn.Parameter(1.0 - torch.pow(ddd, 0.5 * ratio_1_to_almost0)) + + D_MIX_LORA = 32 # generate TIME_MIX for w,k,v,r,g + self.time_maa_w1 = nn.Parameter(torch.zeros(args.n_embd, D_MIX_LORA * 5)) + self.time_maa_w2 = nn.Parameter( + torch.zeros(5, D_MIX_LORA, args.n_embd).uniform_(-0.01, 0.01) + ) + + # fancy time_decay + decay_speed = torch.ones(args.dim_att) + for n in range(args.dim_att): + decay_speed[n] = -6 + 5 * (n / (args.dim_att - 1)) ** (0.7 + 1.3 * ratio_0_to_1) + self.time_decay = nn.Parameter(decay_speed.reshape(1, 1, args.dim_att)) + + D_DECAY_LORA = 64 + self.time_decay_w1 = nn.Parameter(torch.zeros(args.n_embd, D_DECAY_LORA)) + self.time_decay_w2 = nn.Parameter( + torch.zeros(D_DECAY_LORA, args.dim_att).uniform_(-0.01, 0.01) + ) + + tmp = torch.zeros(args.dim_att) + for n in range(args.dim_att): + zigzag = ((n + 1) % 3 - 1) * 0.1 + tmp[n] = ratio_0_to_1 * (1 - (n / (args.dim_att - 1))) + zigzag + + self.time_faaaa = nn.Parameter(tmp.reshape(self.n_head, self.head_size)) + + self.time_shift = nn.ZeroPad2d((0, 0, 1, -1)) + self.receptance = nn.Linear(args.n_embd, args.dim_att, bias=False) + self.key = nn.Linear(args.n_embd, args.dim_att, bias=False) + + self.value = nn.Linear(args.n_embd, args.dim_att, bias=False) + self.output = nn.Linear(args.dim_att, args.n_embd, bias=False) + self.gate = nn.Linear(args.n_embd, args.dim_att, bias=False) + self.ln_x = nn.GroupNorm( + self.n_head, args.dim_att, eps=(1e-5) * (args.head_size_divisor**2) + ) + + @MyFunction + def jit_func(self, x): + B, T, C = x.size() + + xx = self.time_shift(x) - x + + xxx = x + xx * self.time_maa_x + xxx = torch.tanh(xxx @ self.time_maa_w1).view(B * T, 5, -1).transpose(0, 1) + xxx = torch.bmm(xxx, self.time_maa_w2).view(5, B, T, -1) + mw, mk, mv, mr, mg = xxx.unbind(dim=0) + + xw = x + xx * (self.time_maa_w + mw) + xk = x + xx * (self.time_maa_k + mk) + xv = x + xx * (self.time_maa_v + mv) + xr = x + xx * (self.time_maa_r + mr) + xg = x + xx * (self.time_maa_g + mg) + + r = self.receptance(xr) + k = self.key(xk) + v = self.value(xv) + g = F.silu(self.gate(xg)) + + ww = torch.tanh(xw @ self.time_decay_w1) @ self.time_decay_w2 + w = self.time_decay + ww + + return r, k, v, g, w + + @MyFunction + def jit_func_2(self, x, g): + B, T, C = x.size() + x = x.view(B * T, C) + + x = self.ln_x(x).view(B, T, C) + x = self.output(x * g) + return x + + def forward(self, x): + B, T, C = x.size() + H = self.n_head + + r, k, v, g, w = self.jit_func(x) + x = RUN_CUDA_RWKV6(B, T, C, H, r, k, v, w, u=self.time_faaaa) + + return self.jit_func_2(x, g) + class RWKV_CMix_x060(MyModule): - def __init__(self, args, layer_id): - super().__init__() - self.args = args - self.layer_id = layer_id - self.time_shift = nn.ZeroPad2d((0, 0, 1, -1)) - - with torch.no_grad(): # fancy init of time_mix - ratio_1_to_almost0 = 1.0 - (layer_id / args.n_layer) # 1 to ~0 - ddd = torch.ones(1, 1, args.n_embd) - for i in range(args.n_embd): - ddd[0, 0, i] = i / args.n_embd - self.time_maa_k = nn.Parameter(1.0 - torch.pow(ddd, ratio_1_to_almost0)) - self.time_maa_r = nn.Parameter(1.0 - torch.pow(ddd, ratio_1_to_almost0)) - - self.key = nn.Linear(args.n_embd, args.dim_ffn, bias=False) - self.receptance = nn.Linear(args.n_embd, args.n_embd, bias=False) - self.value = nn.Linear(args.dim_ffn, args.n_embd, bias=False) - - @MyFunction - def forward(self, x): - xx = self.time_shift(x) - x - xk = x + xx * self.time_maa_k - xr = x + xx * self.time_maa_r - - k = self.key(xk) - k = torch.relu(k) ** 2 - kv = self.value(k) - return torch.sigmoid(self.receptance(xr)) * kv + def __init__(self, args, layer_id): + super().__init__() + self.args = args + self.layer_id = layer_id + self.time_shift = nn.ZeroPad2d((0, 0, 1, -1)) + + with torch.no_grad(): # fancy init of time_mix + ratio_1_to_almost0 = 1.0 - (layer_id / args.n_layer) # 1 to ~0 + ddd = torch.ones(1, 1, args.n_embd) + for i in range(args.n_embd): + ddd[0, 0, i] = i / args.n_embd + self.time_maa_k = nn.Parameter(1.0 - torch.pow(ddd, ratio_1_to_almost0)) + self.time_maa_r = nn.Parameter(1.0 - torch.pow(ddd, ratio_1_to_almost0)) + + self.key = nn.Linear(args.n_embd, args.dim_ffn, bias=False) + self.receptance = nn.Linear(args.n_embd, args.n_embd, bias=False) + self.value = nn.Linear(args.dim_ffn, args.n_embd, bias=False) + + @MyFunction + def forward(self, x): + xx = self.time_shift(x) - x + xk = x + xx * self.time_maa_k + xr = x + xx * self.time_maa_r + + k = self.key(xk) + k = torch.relu(k) ** 2 + kv = self.value(k) + return torch.sigmoid(self.receptance(xr)) * kv class Block(nn.Module): - def __init__(self, args, layer_id): - super().__init__() - self.args = args - self.layer_id = layer_id - - self.ln1 = nn.LayerNorm(args.n_embd) - self.ln2 = nn.LayerNorm(args.n_embd) - - if self.layer_id == 0: - self.ln0 = nn.LayerNorm(args.n_embd) + def __init__(self, args, layer_id): + super().__init__() + self.args = args + self.layer_id = layer_id + self.ln1 = nn.LayerNorm(args.n_embd) + self.ln2 = nn.LayerNorm(args.n_embd) - self.att = RWKV_Tmix_x060(args, layer_id) - - self.ffn = RWKV_CMix_x060(args, layer_id) + if self.layer_id == 0: + self.ln0 = nn.LayerNorm(args.n_embd) + self.att = RWKV_Tmix_x060(args, layer_id) - if args.dropout > 0: - self.drop0 = nn.Dropout(p=args.dropout) - self.drop1 = nn.Dropout(p=args.dropout) - - def forward(self, x, x_emb=None): - args = self.args - B, T, C = x.size() - if self.layer_id == 0: - x = self.ln0(x) + self.ffn = RWKV_CMix_x060(args, layer_id) - - if self.args.dropout == 0: - if self.layer_id == 0 and args.pre_ffn > 0: - x = x + self.ffnPre(self.ln1(x)) - else: - x = x + self.att(self.ln1(x)) - x = x + self.ffn(self.ln2(x)) - else: - if self.layer_id == 0 and args.pre_ffn > 0: - x = self.drop0(x + self.ffnPre(self.ln1(x))) - else: - x = self.drop0(x + self.att(self.ln1(x))) - x = self.drop1(x + self.ffn(self.ln2(x))) - - return x + if args.dropout > 0: + self.drop0 = nn.Dropout(p=args.dropout) + self.drop1 = nn.Dropout(p=args.dropout) + + def forward(self, x, x_emb=None): + args = self.args + B, T, C = x.size() + if self.layer_id == 0: + x = self.ln0(x) + + if self.args.dropout == 0: + if self.layer_id == 0 and args.pre_ffn > 0: + x = x + self.ffnPre(self.ln1(x)) + else: + x = x + self.att(self.ln1(x)) + x = x + self.ffn(self.ln2(x)) + else: + if self.layer_id == 0 and args.pre_ffn > 0: + x = self.drop0(x + self.ffnPre(self.ln1(x))) + else: + x = self.drop0(x + self.att(self.ln1(x))) + x = self.drop1(x + self.ffn(self.ln2(x))) + + return x class RWKVLayer(nn.Module): - def __init__(self, args, layer_id): - super().__init__() - self.args = args - self.layer_id = layer_id - if args.dim_ffn is None: - args.dim_ffn = int((args.n_embd * 3.5) // 32 * 32) - self.ln0 = None - if self.layer_id == 0 and args.get("ln0", True): - self.ln0 = nn.LayerNorm(args.n_embd) - - self.ln1 = None - if args.get("ln1", True): - self.ln1 = nn.LayerNorm(args.n_embd) - self.ln2 = nn.LayerNorm(args.n_embd) - + def __init__(self, args, layer_id): + super().__init__() + self.args = args + self.layer_id = layer_id + if args.dim_ffn is None: + args.dim_ffn = int((args.n_embd * 3.5) // 32 * 32) + self.ln0 = None + if self.layer_id == 0 and args.get("ln0", True): + self.ln0 = nn.LayerNorm(args.n_embd) - self.att = RWKV_Tmix_x060(args, layer_id) - - self.ffn = RWKV_CMix_x060(args, layer_id) - - if args.dropout > 0: - self.drop0 = nn.Dropout(p=args.dropout) - self.drop1 = nn.Dropout(p=args.dropout) - - # init - if args.get("init_rwkv", True): - print("init_rwkv") - nn.init.orthogonal_(self.att.receptance.weight, gain=1) - nn.init.orthogonal_(self.att.key.weight, gain=0.1) - nn.init.orthogonal_(self.att.value.weight, gain=1) - nn.init.orthogonal_(self.att.gate.weight, gain=0.1) - nn.init.zeros_(self.att.output.weight) - - nn.init.orthogonal_(self.ffn.key.weight, gain=1) - nn.init.zeros_(self.ffn.value.weight) - nn.init.zeros_(self.ffn.receptance.weight) - scale = ((1 + layer_id) / args.get("n_layer")) ** 0.7 - nn.init.constant_(self.ln2.weight, scale) - if self.ln0 is not None: - nn.init.constant_(self.ln0.weight, scale) - if self.ln1 is not None: - nn.init.constant_(self.ln1.weight, scale) - - def forward(self, x, x_emb=None, mask=None, **kwargs): - - args = self.args - if args.get("datatype", "bf16") == "bf16": - x = x.bfloat16() - B, T, C = x.size() - if self.layer_id == 0 and self.ln0 is not None: - x = self.ln0(x) - - if self.args.dropout == 0: - if self.ln1 is None: - x = x + self.att(x) - else: - x = x + self.att(self.ln1(x)) - x = x + self.ffn(self.ln2(x)) - else: - if self.ln1 is None: - x = self.drop0(x + self.att(x)) - else: - x = self.drop0(x + self.att(self.ln1(x))) - x = self.drop1(x + self.ffn(self.ln2(x))) - - if args.get("datatype", "bf16") == "bf16": - x = x.to(torch.float32) - return x + self.ln1 = None + if args.get("ln1", True): + self.ln1 = nn.LayerNorm(args.n_embd) + self.ln2 = nn.LayerNorm(args.n_embd) + + self.att = RWKV_Tmix_x060(args, layer_id) + + self.ffn = RWKV_CMix_x060(args, layer_id) + + if args.dropout > 0: + self.drop0 = nn.Dropout(p=args.dropout) + self.drop1 = nn.Dropout(p=args.dropout) + + # init + if args.get("init_rwkv", True): + print("init_rwkv") + nn.init.orthogonal_(self.att.receptance.weight, gain=1) + nn.init.orthogonal_(self.att.key.weight, gain=0.1) + nn.init.orthogonal_(self.att.value.weight, gain=1) + nn.init.orthogonal_(self.att.gate.weight, gain=0.1) + nn.init.zeros_(self.att.output.weight) + + nn.init.orthogonal_(self.ffn.key.weight, gain=1) + nn.init.zeros_(self.ffn.value.weight) + nn.init.zeros_(self.ffn.receptance.weight) + scale = ((1 + layer_id) / args.get("n_layer")) ** 0.7 + nn.init.constant_(self.ln2.weight, scale) + if self.ln0 is not None: + nn.init.constant_(self.ln0.weight, scale) + if self.ln1 is not None: + nn.init.constant_(self.ln1.weight, scale) + + def forward(self, x, x_emb=None, mask=None, **kwargs): + + args = self.args + if args.get("datatype", "bf16") == "bf16": + x = x.bfloat16() + B, T, C = x.size() + if self.layer_id == 0 and self.ln0 is not None: + x = self.ln0(x) + + if self.args.dropout == 0: + if self.ln1 is None: + x = x + self.att(x) + else: + x = x + self.att(self.ln1(x)) + x = x + self.ffn(self.ln2(x)) + else: + if self.ln1 is None: + x = self.drop0(x + self.att(x)) + else: + x = self.drop0(x + self.att(self.ln1(x))) + x = self.drop1(x + self.ffn(self.ln2(x))) + + if args.get("datatype", "bf16") == "bf16": + x = x.to(torch.float32) + return x class RWKV(nn.Module): - def __init__(self, args): - super().__init__() - self.args = args - if not hasattr(args, 'dim_att'): - args.dim_att = args.n_embd - if not hasattr(args, 'dim_ffn'): - if '-f4' in os.environ["RWKV_MY_TESTING"]: - args.dim_ffn = int((args.n_embd * 4) // 32 * 32) - else: - args.dim_ffn = int((args.n_embd * 3.5) // 32 * 32) # default = 3.5x emb size - if not hasattr(args, 'tiny_att_layer'): - args.tiny_att_layer = -1 - if not hasattr(args, 'tiny_att_dim'): - args.tiny_att_dim = -1 - assert args.n_embd % 32 == 0 - assert args.dim_att % 32 == 0 - assert args.dim_ffn % 32 == 0 - - self.emb = nn.Embedding(args.vocab_size, args.n_embd) - - self.blocks = nn.ModuleList([Block(args, i) for i in range(args.n_layer)]) - - self.ln_out = nn.LayerNorm(args.n_embd) - self.head = nn.Linear(args.n_embd, args.vocab_size, bias=False) + def __init__(self, args): + super().__init__() + self.args = args + if not hasattr(args, "dim_att"): + args.dim_att = args.n_embd + if not hasattr(args, "dim_ffn"): + if "-f4" in os.environ["RWKV_MY_TESTING"]: + args.dim_ffn = int((args.n_embd * 4) // 32 * 32) + else: + args.dim_ffn = int((args.n_embd * 3.5) // 32 * 32) # default = 3.5x emb size + if not hasattr(args, "tiny_att_layer"): + args.tiny_att_layer = -1 + if not hasattr(args, "tiny_att_dim"): + args.tiny_att_dim = -1 + assert args.n_embd % 32 == 0 + assert args.dim_att % 32 == 0 + assert args.dim_ffn % 32 == 0 + self.emb = nn.Embedding(args.vocab_size, args.n_embd) - if args.dropout > 0: - self.drop0 = nn.Dropout(p=args.dropout) + self.blocks = nn.ModuleList([Block(args, i) for i in range(args.n_layer)]) + self.ln_out = nn.LayerNorm(args.n_embd) + self.head = nn.Linear(args.n_embd, args.vocab_size, bias=False) - def forward(self, idx): - args = self.args - B, T = idx.size() - assert T <= args.ctx_len, "Cannot forward, model ctx_len is exhausted." - - x = self.emb(idx) - x_emb = x - - if args.dropout > 0: - x = self.drop0(x) - if args.tiny_att_dim > 0: - for block in self.blocks: - if args.grad_cp == 1: - x = deepspeed.checkpointing.checkpoint(block, x, x_emb) - else: - x = block(x, x_emb) - else: - for block in self.blocks: - if args.grad_cp == 1: - x = deepspeed.checkpointing.checkpoint(block, x) - else: - x = block(x) - - x = self.ln_out(x) - - if args.head_qk > 0: - q = self.head_q(x)[:, :T, :] - k = self.head_k(x)[:, :T, :] - c = (q @ k.transpose(-2, -1)) * (1.0 / args.head_qk) - c = c.masked_fill(self.copy_mask[:T, :T] == 0, 0) - - if "32" in os.environ["RWKV_FLOAT_MODE"]: - c = c @ F.one_hot(idx, num_classes=args.vocab_size) - elif os.environ["RWKV_FLOAT_MODE"] == "fp16": - c = c @ F.one_hot(idx, num_classes=args.vocab_size).half() - elif os.environ["RWKV_FLOAT_MODE"] == "bf16": - c = c @ F.one_hot(idx, num_classes=args.vocab_size).bfloat16() - - x = self.head(x) + c - else: - x = self.head(x) - - return x + if args.dropout > 0: + self.drop0 = nn.Dropout(p=args.dropout) + + def forward(self, idx): + args = self.args + B, T = idx.size() + assert T <= args.ctx_len, "Cannot forward, model ctx_len is exhausted." + + x = self.emb(idx) + x_emb = x + + if args.dropout > 0: + x = self.drop0(x) + if args.tiny_att_dim > 0: + for block in self.blocks: + if args.grad_cp == 1: + x = deepspeed.checkpointing.checkpoint(block, x, x_emb) + else: + x = block(x, x_emb) + else: + for block in self.blocks: + if args.grad_cp == 1: + x = deepspeed.checkpointing.checkpoint(block, x) + else: + x = block(x) + + x = self.ln_out(x) + + if args.head_qk > 0: + q = self.head_q(x)[:, :T, :] + k = self.head_k(x)[:, :T, :] + c = (q @ k.transpose(-2, -1)) * (1.0 / args.head_qk) + c = c.masked_fill(self.copy_mask[:T, :T] == 0, 0) + + if "32" in os.environ["RWKV_FLOAT_MODE"]: + c = c @ F.one_hot(idx, num_classes=args.vocab_size) + elif os.environ["RWKV_FLOAT_MODE"] == "fp16": + c = c @ F.one_hot(idx, num_classes=args.vocab_size).half() + elif os.environ["RWKV_FLOAT_MODE"] == "bf16": + c = c @ F.one_hot(idx, num_classes=args.vocab_size).bfloat16() + + x = self.head(x) + c + else: + x = self.head(x) + + return x diff --git a/funasr/models/sense_voice/whisper_lib/__init__.py b/funasr/models/sense_voice/whisper_lib/__init__.py index d7fbba36f..855b8cf26 100644 --- a/funasr/models/sense_voice/whisper_lib/__init__.py +++ b/funasr/models/sense_voice/whisper_lib/__init__.py @@ -136,13 +136,9 @@ def load_model( checkpoint_file = open(name, "rb").read() if in_memory else name alignment_heads = None else: - raise RuntimeError( - f"Model {name} not found; available models = {available_models()}" - ) + raise RuntimeError(f"Model {name} not found; available models = {available_models()}") - with ( - io.BytesIO(checkpoint_file) if in_memory else open(checkpoint_file, "rb") - ) as fp: + with io.BytesIO(checkpoint_file) if in_memory else open(checkpoint_file, "rb") as fp: checkpoint = torch.load(fp, map_location=device) del checkpoint_file diff --git a/funasr/models/sense_voice/whisper_lib/audio.py b/funasr/models/sense_voice/whisper_lib/audio.py index 52da32cab..2f688e90b 100644 --- a/funasr/models/sense_voice/whisper_lib/audio.py +++ b/funasr/models/sense_voice/whisper_lib/audio.py @@ -68,9 +68,7 @@ def pad_or_trim(array, length: int = N_SAMPLES, *, axis: int = -1): """ if torch.is_tensor(array): if array.shape[axis] > length: - array = array.index_select( - dim=axis, index=torch.arange(length, device=array.device) - ) + array = array.index_select(dim=axis, index=torch.arange(length, device=array.device)) if array.shape[axis] < length: pad_widths = [(0, 0)] * array.ndim @@ -89,7 +87,7 @@ def pad_or_trim(array, length: int = N_SAMPLES, *, axis: int = -1): @lru_cache(maxsize=None) -def mel_filters(device, n_mels: int, filters_path: str=None) -> torch.Tensor: +def mel_filters(device, n_mels: int, filters_path: str = None) -> torch.Tensor: """ load the mel filterbank matrix for projecting STFT into a Mel spectrogram. Allows decoupling librosa dependency; saved using: diff --git a/funasr/models/sense_voice/whisper_lib/decoding.py b/funasr/models/sense_voice/whisper_lib/decoding.py index 203efe804..382a180ae 100644 --- a/funasr/models/sense_voice/whisper_lib/decoding.py +++ b/funasr/models/sense_voice/whisper_lib/decoding.py @@ -19,7 +19,11 @@ if TYPE_CHECKING: @torch.no_grad() def detect_language( - model: "Whisper", mel: Tensor, tokenizer: Tokenizer = None, initial_prompt = None, x = None, + model: "Whisper", + mel: Tensor, + tokenizer: Tokenizer = None, + initial_prompt=None, + x=None, ) -> Tuple[Tensor, List[dict]]: """ Detect the spoken language in the audio, and return them as list of strings, along with the ids @@ -34,16 +38,9 @@ def detect_language( list of dictionaries containing the probability distribution over all languages. """ if tokenizer is None: - tokenizer = get_tokenizer( - model.is_multilingual, num_languages=model.num_languages - ) - if ( - tokenizer.language is None - or tokenizer.language_token not in tokenizer.sot_sequence - ): - raise ValueError( - "This model doesn't have language tokens so it can't perform lang id" - ) + tokenizer = get_tokenizer(model.is_multilingual, num_languages=model.num_languages) + if tokenizer.language is None or tokenizer.language_token not in tokenizer.sot_sequence: + raise ValueError("This model doesn't have language tokens so it can't perform lang id") single = mel.ndim == 2 if single: @@ -59,17 +56,19 @@ def detect_language( # FIX(funasr): sense vocie # x = torch.tensor([[tokenizer.sot]] * n_audio).to(mel.device) # [n_audio, 1] if x is None: - x = torch.tensor([tokenizer.encode(initial_prompt, allowed_special="all")] * n_audio).to(mel.device) # [n_audio, 1] + x = torch.tensor([tokenizer.encode(initial_prompt, allowed_special="all")] * n_audio).to( + mel.device + ) # [n_audio, 1] else: x = x.to(mel.device) - logits = model.logits(x[:,:-1], mel)[:, -1] + logits = model.logits(x[:, :-1], mel)[:, -1] # collect detected languages; suppress all non-language tokens mask = torch.ones(logits.shape[-1], dtype=torch.bool) mask[list(tokenizer.all_language_tokens)] = False mask[tokenizer.no_speech] = False - + logits[:, mask] = -np.inf language_tokens = logits.argmax(dim=-1) language_token_probs = logits.softmax(dim=-1).cpu() @@ -77,7 +76,10 @@ def detect_language( language_probs = [ { c: language_token_probs[i, j].item() - for j, c in zip(list(tokenizer.all_language_tokens) + [tokenizer.no_speech], list(tokenizer.all_language_codes) + ["nospeech"]) + for j, c in zip( + list(tokenizer.all_language_tokens) + [tokenizer.no_speech], + list(tokenizer.all_language_codes) + ["nospeech"], + ) } for i in range(n_audio) ] @@ -120,13 +122,15 @@ class DecodingOptions: gain_event: bool = False # this will suppress blank outputs gain_tokens_bg: Optional[Union[str, List[int]]] = "<|Speech|><|BGM|><|Applause|><|Laughter|>" - gain_tokens_ed: Optional[Union[str, List[int]]] = "<|/Speech|><|/BGM|><|/Applause|><|/Laughter|>" - gain_tokens_score: List[float] = field(default_factory=lambda: [1, 1, 25.0, 5.0]) #[25, 5] + gain_tokens_ed: Optional[Union[str, List[int]]] = ( + "<|/Speech|><|/BGM|><|/Applause|><|/Laughter|>" + ) + gain_tokens_score: List[float] = field(default_factory=lambda: [1, 1, 25.0, 5.0]) # [25, 5] use_emo_threshold: bool = False # this will suppress blank outputs emo_unk_token: Optional[Union[str, List[int]]] = "<|SPECIAL_TOKEN_1|>" emo_target_tokens: Optional[Union[str, List[int]]] = "<|HAPPY|><|SAD|><|ANGRY|>" - emo_target_threshold: List[float] = field(default_factory=lambda: [0.1, 0.1, 0.1]) #[25, 5] + emo_target_threshold: List[float] = field(default_factory=lambda: [0.1, 0.1, 0.1]) # [25, 5] # timestamp sampling options without_timestamps: bool = False # use <|notimestamps|> to sample text tokens only @@ -203,9 +207,7 @@ class PyTorchInference(Inference): class SequenceRanker: - def rank( - self, tokens: List[List[Tensor]], sum_logprobs: List[List[float]] - ) -> List[int]: + def rank(self, tokens: List[List[Tensor]], sum_logprobs: List[List[float]]) -> List[int]: """ Given a list of groups of samples and their cumulative log probabilities, return the indices of the samples in each group to select as the final result @@ -243,9 +245,7 @@ class TokenDecoder: def reset(self): """Initialize any stateful variables for decoding a new sequence""" - def update( - self, tokens: Tensor, logits: Tensor, sum_logprobs: Tensor - ) -> Tuple[Tensor, bool]: + def update(self, tokens: Tensor, logits: Tensor, sum_logprobs: Tensor) -> Tuple[Tensor, bool]: """Specify how to select the next token, based on the current trace and logits Parameters @@ -300,9 +300,7 @@ class GreedyDecoder(TokenDecoder): self.temperature = temperature self.eot = eot - def update( - self, tokens: Tensor, logits: Tensor, sum_logprobs: Tensor - ) -> Tuple[Tensor, bool]: + def update(self, tokens: Tensor, logits: Tensor, sum_logprobs: Tensor) -> Tuple[Tensor, bool]: if self.temperature == 0: next_tokens = logits.argmax(dim=-1) else: @@ -339,16 +337,12 @@ class BeamSearchDecoder(TokenDecoder): self.max_candidates: int = round(beam_size * self.patience) self.finished_sequences = None - assert ( - self.max_candidates > 0 - ), f"Invalid beam size ({beam_size}) or patience ({patience})" + assert self.max_candidates > 0, f"Invalid beam size ({beam_size}) or patience ({patience})" def reset(self): self.finished_sequences = None - def update( - self, tokens: Tensor, logits: Tensor, sum_logprobs: Tensor - ) -> Tuple[Tensor, bool]: + def update(self, tokens: Tensor, logits: Tensor, sum_logprobs: Tensor) -> Tuple[Tensor, bool]: if tokens.shape[0] % self.beam_size != 0: raise ValueError(f"{tokens.shape}[0] % {self.beam_size} != 0") @@ -392,9 +386,7 @@ class BeamSearchDecoder(TokenDecoder): # add newly finished sequences to self.finished_sequences assert len(self.finished_sequences) == len(finished_sequences) - for previously_finished, newly_finished in zip( - self.finished_sequences, finished_sequences - ): + for previously_finished, newly_finished in zip(self.finished_sequences, finished_sequences): for seq in sorted(newly_finished, key=newly_finished.get, reverse=True): if len(previously_finished) >= self.max_candidates: break # the candidate list is full @@ -402,8 +394,7 @@ class BeamSearchDecoder(TokenDecoder): # mark as completed if all audio has enough number of samples completed = all( - len(sequences) >= self.max_candidates - for sequences in self.finished_sequences + len(sequences) >= self.max_candidates for sequences in self.finished_sequences ) return tokens, completed @@ -411,9 +402,7 @@ class BeamSearchDecoder(TokenDecoder): # collect all finished sequences, including patience, and add unfinished ones if not enough sum_logprobs = sum_logprobs.cpu() for i, sequences in enumerate(self.finished_sequences): - if ( - len(sequences) < self.beam_size - ): # when not enough sequences are finished + if len(sequences) < self.beam_size: # when not enough sequences are finished for j in list(np.argsort(sum_logprobs[i]))[::-1]: sequence = preceding_tokens[i, j].tolist() + [self.eot] sequences[tuple(sequence)] = sum_logprobs[i][j].item() @@ -421,8 +410,7 @@ class BeamSearchDecoder(TokenDecoder): break tokens: List[List[Tensor]] = [ - [torch.tensor(seq) for seq in sequences.keys()] - for sequences in self.finished_sequences + [torch.tensor(seq) for seq in sequences.keys()] for sequences in self.finished_sequences ] sum_logprobs: List[List[float]] = [ list(sequences.values()) for sequences in self.finished_sequences @@ -463,8 +451,11 @@ class SuppressTokens(LogitFilter): def apply(self, logits: Tensor, tokens: Tensor): logits[:, self.suppress_tokens] = -np.inf + class GainEventToken(LogitFilter): - def __init__(self, bg_tokens: Sequence[int], ed_tokens:Sequence[int], gain_values: Sequence[float]): + def __init__( + self, bg_tokens: Sequence[int], ed_tokens: Sequence[int], gain_values: Sequence[float] + ): self.bg_tokens = list(bg_tokens) self.ed_tokens = list(ed_tokens) self.gain_value = [np.log(max(ga, 1e-9)) for ga in gain_values] @@ -477,13 +468,16 @@ class GainEventToken(LogitFilter): sum_bg = sum([1 if x == bg else 0 for x in tokens[i]]) sum_ed = sum([1 if x == ed else 0 for x in tokens[i]]) logits[i, bg] += ga - if sum_bg > sum_ed or tokens[i,-1] in [bg, ed]: + if sum_bg > sum_ed or tokens[i, -1] in [bg, ed]: logits[i, bg] = -np.inf if sum_bg <= sum_ed: logits[i, ed] = -np.inf + class ThresholdEmoToken(LogitFilter): - def __init__(self, unk_tokens: Sequence[int], emo_tokens:Sequence[int], th_values: Sequence[float]): + def __init__( + self, unk_tokens: Sequence[int], emo_tokens: Sequence[int], th_values: Sequence[float] + ): self.unk_token = list(unk_tokens)[0] self.emo_tokens = list(emo_tokens) self.th_values = list(th_values) @@ -493,7 +487,7 @@ class ThresholdEmoToken(LogitFilter): for i in range(len(tokens)): for emo, th in zip(self.emo_tokens, self.th_values): if logits[i].argmax() == emo and logits[i].softmax(dim=-1)[emo] < th: - logits[i, self.unk_token] = max(logits[i, emo], logits[i, self.unk_token]) + logits[i, self.unk_token] = max(logits[i, emo], logits[i, self.unk_token]) logits[i, emo] = -np.inf # for bg, ed, ga in zip(self.bg_tokens, self.ed_tokens, self.gain_value): @@ -526,12 +520,8 @@ class ApplyTimestampRules(LogitFilter): for k in range(tokens.shape[0]): sampled_tokens = tokens[k, self.sample_begin :] seq = [t for t in sampled_tokens.tolist()] - last_was_timestamp = ( - len(seq) >= 1 and seq[-1] >= self.tokenizer.timestamp_begin - ) - penultimate_was_timestamp = ( - len(seq) < 2 or seq[-2] >= self.tokenizer.timestamp_begin - ) + last_was_timestamp = len(seq) >= 1 and seq[-1] >= self.tokenizer.timestamp_begin + penultimate_was_timestamp = len(seq) < 2 or seq[-2] >= self.tokenizer.timestamp_begin if last_was_timestamp: if penultimate_was_timestamp: # has to be non-timestamp @@ -539,9 +529,7 @@ class ApplyTimestampRules(LogitFilter): else: # cannot be normal text tokens logits[k, : self.tokenizer.eot] = -np.inf - timestamps = sampled_tokens[ - sampled_tokens.ge(self.tokenizer.timestamp_begin) - ] + timestamps = sampled_tokens[sampled_tokens.ge(self.tokenizer.timestamp_begin)] if timestamps.numel() > 0: # timestamps shouldn't decrease; forbid timestamp tokens smaller than the last # also force each segment to have a nonzero length, to prevent infinite looping @@ -557,17 +545,13 @@ class ApplyTimestampRules(LogitFilter): # apply the `max_initial_timestamp` option if self.max_initial_timestamp_index is not None: - last_allowed = ( - self.tokenizer.timestamp_begin + self.max_initial_timestamp_index - ) + last_allowed = self.tokenizer.timestamp_begin + self.max_initial_timestamp_index logits[:, last_allowed + 1 :] = -np.inf # if sum of probability over timestamps is above any other token, sample timestamp logprobs = F.log_softmax(logits.float(), dim=-1) for k in range(tokens.shape[0]): - timestamp_logprob = logprobs[k, self.tokenizer.timestamp_begin :].logsumexp( - dim=-1 - ) + timestamp_logprob = logprobs[k, self.tokenizer.timestamp_begin :].logsumexp(dim=-1) max_text_token_logprob = logprobs[k, : self.tokenizer.timestamp_begin].max() if timestamp_logprob > max_text_token_logprob: logits[k, : self.tokenizer.timestamp_begin] = -np.inf @@ -588,7 +572,7 @@ class DecodingTask: num_languages=model.num_languages, language=language, task=options.task, - vocab_path=options.vocab_path + vocab_path=options.vocab_path, ) self.tokenizer: Tokenizer = tokenizer self.options: DecodingOptions = self._verify_options(options) @@ -626,30 +610,28 @@ class DecodingTask: if self.options.suppress_tokens: self.logit_filters.append(SuppressTokens(self._get_suppress_tokens())) if self.options.gain_event: - self.logit_filters.append(GainEventToken( - self.tokenizer.encode(self.options.gain_tokens_bg, allowed_special="all"), - self.tokenizer.encode(self.options.gain_tokens_ed, allowed_special="all"), - self.options.gain_tokens_score + self.logit_filters.append( + GainEventToken( + self.tokenizer.encode(self.options.gain_tokens_bg, allowed_special="all"), + self.tokenizer.encode(self.options.gain_tokens_ed, allowed_special="all"), + self.options.gain_tokens_score, ) ) if self.options.use_emo_threshold: - self.logit_filters.append(ThresholdEmoToken( - self.tokenizer.encode(self.options.emo_unk_token, allowed_special="all"), - self.tokenizer.encode(self.options.emo_target_tokens, allowed_special="all"), - self.options.emo_target_threshold + self.logit_filters.append( + ThresholdEmoToken( + self.tokenizer.encode(self.options.emo_unk_token, allowed_special="all"), + self.tokenizer.encode(self.options.emo_target_tokens, allowed_special="all"), + self.options.emo_target_threshold, ) ) if not options.without_timestamps: precision = CHUNK_LENGTH / model.dims.n_audio_ctx # usually 0.02 seconds max_initial_timestamp_index = None if options.max_initial_timestamp: - max_initial_timestamp_index = round( - self.options.max_initial_timestamp / precision - ) + max_initial_timestamp_index = round(self.options.max_initial_timestamp / precision) self.logit_filters.append( - ApplyTimestampRules( - tokenizer, self.sample_begin, max_initial_timestamp_index - ) + ApplyTimestampRules(tokenizer, self.sample_begin, max_initial_timestamp_index) ) def _verify_options(self, options: DecodingOptions) -> DecodingOptions: @@ -660,9 +642,7 @@ class DecodingTask: raise ValueError("best_of with greedy sampling (T=0) is not compatible") if options.patience is not None and options.beam_size is None: raise ValueError("patience requires beam_size to be given") - if options.length_penalty is not None and not ( - 0 <= options.length_penalty <= 1 - ): + if options.length_penalty is not None and not (0 <= options.length_penalty <= 1): raise ValueError("length_penalty (alpha) should be a value between 0 and 1") return options @@ -672,9 +652,7 @@ class DecodingTask: if prefix := self.options.prefix: prefix_tokens = ( - self.tokenizer.encode(" " + prefix.strip()) - if isinstance(prefix, str) - else prefix + self.tokenizer.encode(" " + prefix.strip()) if isinstance(prefix, str) else prefix ) if self.sample_len is not None: max_prefix_len = self.n_ctx // 2 - self.sample_len @@ -683,16 +661,10 @@ class DecodingTask: if prompt := self.options.prompt: prompt_tokens = ( - self.tokenizer.encode(" " + prompt.strip()) - if isinstance(prompt, str) - else prompt + self.tokenizer.encode(" " + prompt.strip()) if isinstance(prompt, str) else prompt ) - tokens = ( - [self.tokenizer.sot_prev] - + prompt_tokens[-(self.n_ctx // 2 - 1) :] - + tokens - ) - #FIX(funasr): sense vocie + tokens = [self.tokenizer.sot_prev] + prompt_tokens[-(self.n_ctx // 2 - 1) :] + tokens + # FIX(funasr): sense vocie if initial_prompt := self.options.initial_prompt: if self.options.language is not None: initial_prompt = f"{initial_prompt}<|{self.options.language}|>" @@ -701,7 +673,6 @@ class DecodingTask: tokens = self.tokenizer.encode(initial_prompt, allowed_special="all") tokens += [0] - return tuple(tokens) def _get_suppress_tokens(self) -> Tuple[int]: @@ -746,12 +717,8 @@ class DecodingTask: else: audio_features = self.model.encoder(mel) - if audio_features.dtype != ( - torch.float16 if self.options.fp16 else torch.float32 - ): - return TypeError( - f"audio_features has an incorrect dtype: {audio_features.dtype}" - ) + if audio_features.dtype != (torch.float16 if self.options.fp16 else torch.float32): + return TypeError(f"audio_features has an incorrect dtype: {audio_features.dtype}") return audio_features @@ -766,15 +733,18 @@ class DecodingTask: languages = [max(probs, key=probs.get) for probs in lang_probs] # FIX(funasr): sense vocie # if self.options.language is None: - # tokens[:, self.sot_index + 1] = lang_tokens # write language tokens + # tokens[:, self.sot_index + 1] = lang_tokens # write language tokens if self.options.language is None: # tokens[:, self.sot_index + 1] = lang_tokens # write language tokens languages = "".join([f"<|{language}|>" for language in languages]) n_audio = audio_features.shape[0] - lang_tokens = torch.tensor([self.tokenizer.encode(languages, allowed_special="all")] * n_audio).to( - audio_features.device) # [n_audio, 1] - + lang_tokens = torch.tensor( + [self.tokenizer.encode(languages, allowed_special="all")] * n_audio + ).to( + audio_features.device + ) # [n_audio, 1] + tokens[:, -1:] = lang_tokens[:, :] languages = [languages] @@ -789,9 +759,7 @@ class DecodingTask: for i in range(self.sample_len): logits = self.inference.logits(tokens, audio_features) - if ( - i == 0 and self.tokenizer.no_speech is not None - ): # save no_speech_probs + if i == 0 and self.tokenizer.no_speech is not None: # save no_speech_probs probs_at_sot = logits[:, self.sot_index].float().softmax(dim=-1) no_speech_probs = probs_at_sot[:, self.tokenizer.no_speech].tolist() @@ -825,12 +793,8 @@ class DecodingTask: languages, language_probs = self._detect_language(audio_features, tokens) if self.options.task == "lang_id": return [ - DecodingResult( - audio_features=features, language=language, language_probs=probs - ) - for features, language, probs in zip( - audio_features, languages, language_probs - ) + DecodingResult(audio_features=features, language=language, language_probs=probs) + for features, language, probs in zip(audio_features, languages, language_probs) ] # repeat text tensors by the group size, for beam search or best-of-n sampling @@ -850,8 +814,7 @@ class DecodingTask: # get the final candidates for each group, and slice between the first sampled token and EOT tokens, sum_logprobs = self.decoder.finalize(tokens, sum_logprobs) tokens: List[List[Tensor]] = [ - [t[self.sample_begin : (t == tokenizer.eot).nonzero()[0, 0]] for t in s] - for s in tokens + [t[self.sample_begin : (t == tokenizer.eot).nonzero()[0, 0]] for t in s] for s in tokens ] # select the top-ranked sample in each group @@ -860,9 +823,7 @@ class DecodingTask: texts: List[str] = [tokenizer.decode(t).strip() for t in tokens] sum_logprobs: List[float] = [lp[i] for i, lp in zip(selected, sum_logprobs)] - avg_logprobs: List[float] = [ - lp / (len(t) + 1) for t, lp in zip(tokens, sum_logprobs) - ] + avg_logprobs: List[float] = [lp / (len(t) + 1) for t, lp in zip(tokens, sum_logprobs)] fields = ( texts, @@ -886,9 +847,7 @@ class DecodingTask: temperature=self.options.temperature, compression_ratio=compression_ratio(text), ) - for text, language, tokens, features, avg_logprob, no_speech_prob in zip( - *fields - ) + for text, language, tokens, features, avg_logprob, no_speech_prob in zip(*fields) ] diff --git a/funasr/models/sense_voice/whisper_lib/model.py b/funasr/models/sense_voice/whisper_lib/model.py index 5f7caebc2..40939df69 100644 --- a/funasr/models/sense_voice/whisper_lib/model.py +++ b/funasr/models/sense_voice/whisper_lib/model.py @@ -42,9 +42,7 @@ class Linear(nn.Linear): class Conv1d(nn.Conv1d): - def _conv_forward( - self, x: Tensor, weight: Tensor, bias: Optional[Tensor] - ) -> Tensor: + def _conv_forward(self, x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> Tensor: return super()._conv_forward( x, weight.to(x.dtype), None if bias is None else bias.to(x.dtype) ) @@ -94,7 +92,12 @@ class MultiHeadAttention(nn.Module): return self.out(wv), qk def qkv_attention( - self, q: Tensor, k: Tensor, v: Tensor, mask: Optional[Tensor] = None, **kwargs, + self, + q: Tensor, + k: Tensor, + v: Tensor, + mask: Optional[Tensor] = None, + **kwargs, ): is_pad_mask = kwargs.get("is_pad_mask", False) n_batch, n_ctx, n_state = q.shape @@ -109,11 +112,9 @@ class MultiHeadAttention(nn.Module): qk = qk + mask[:n_ctx, :n_ctx] else: mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) - min_value = float( - np.finfo(torch.tensor(0, dtype=qk.dtype).numpy().dtype).min - ) + min_value = float(np.finfo(torch.tensor(0, dtype=qk.dtype).numpy().dtype).min) qk = qk.masked_fill(mask, min_value) - + qk = qk.float() w = F.softmax(qk, dim=-1).to(q.dtype) @@ -129,15 +130,11 @@ class ResidualAttentionBlock(nn.Module): self.attn = MultiHeadAttention(n_state, n_head) self.attn_ln = LayerNorm(n_state) - self.cross_attn = ( - MultiHeadAttention(n_state, n_head) if cross_attention else None - ) + self.cross_attn = MultiHeadAttention(n_state, n_head) if cross_attention else None self.cross_attn_ln = LayerNorm(n_state) if cross_attention else None n_mlp = n_state * 4 - self.mlp = nn.Sequential( - Linear(n_state, n_mlp), nn.GELU(), Linear(n_mlp, n_state) - ) + self.mlp = nn.Sequential(Linear(n_state, n_mlp), nn.GELU(), Linear(n_mlp, n_state)) self.mlp_ln = LayerNorm(n_state) def forward( @@ -152,15 +149,18 @@ class ResidualAttentionBlock(nn.Module): is_pad_memory_mask = kwargs.get("is_pad_memory_mask", False) x = x + self.attn(self.attn_ln(x), mask=mask, kv_cache=kv_cache, is_pad_mask=is_pad_mask)[0] if self.cross_attn: - x = x + self.cross_attn(self.cross_attn_ln(x), xa, kv_cache=kv_cache, is_pad_mask=is_pad_memory_mask)[0] + x = ( + x + + self.cross_attn( + self.cross_attn_ln(x), xa, kv_cache=kv_cache, is_pad_mask=is_pad_memory_mask + )[0] + ) x = x + self.mlp(self.mlp_ln(x)) return x class AudioEncoder(nn.Module): - def __init__( - self, n_mels: int, n_ctx: int, n_state: int, n_head: int, n_layer: int - ): + def __init__(self, n_mels: int, n_ctx: int, n_state: int, n_head: int, n_layer: int): super().__init__() self.conv1 = Conv1d(n_mels, n_state, kernel_size=3, stride=2, padding=1) self.conv2 = Conv1d(n_state, n_state, kernel_size=3, stride=2, padding=1) @@ -184,7 +184,6 @@ class AudioEncoder(nn.Module): # x = (x + self.positional_embedding).to(x.dtype) x = (x + self.positional_embedding[: x.size(1), :]).to(x.dtype) - for block in self.blocks: x = block(x) @@ -193,19 +192,14 @@ class AudioEncoder(nn.Module): class TextDecoder(nn.Module): - def __init__( - self, n_vocab: int, n_ctx: int, n_state: int, n_head: int, n_layer: int - ): + def __init__(self, n_vocab: int, n_ctx: int, n_state: int, n_head: int, n_layer: int): super().__init__() self.token_embedding = nn.Embedding(n_vocab, n_state) self.positional_embedding = nn.Parameter(torch.empty(n_ctx, n_state)) self.blocks: Iterable[ResidualAttentionBlock] = nn.ModuleList( - [ - ResidualAttentionBlock(n_state, n_head, cross_attention=True) - for _ in range(n_layer) - ] + [ResidualAttentionBlock(n_state, n_head, cross_attention=True) for _ in range(n_layer)] ) self.ln = LayerNorm(n_state) @@ -220,19 +214,14 @@ class TextDecoder(nn.Module): the encoded audio features to be attended on """ offset = next(iter(kv_cache.values())).shape[1] if kv_cache else 0 - x = ( - self.token_embedding(x) - + self.positional_embedding[offset : offset + x.shape[-1]] - ) + x = self.token_embedding(x) + self.positional_embedding[offset : offset + x.shape[-1]] x = x.to(xa.dtype) for block in self.blocks: x = block(x, xa, mask=self.mask, kv_cache=kv_cache) x = self.ln(x) - logits = ( - x @ torch.transpose(self.token_embedding.weight.to(x.dtype), 0, 1) - ).float() + logits = (x @ torch.transpose(self.token_embedding.weight.to(x.dtype), 0, 1)).float() return logits @@ -257,21 +246,15 @@ class Whisper(nn.Module): ) # use the last half among the decoder layers for time alignment by default; # to use a specific set of heads, see `set_alignment_heads()` below. - all_heads = torch.zeros( - self.dims.n_text_layer, self.dims.n_text_head, dtype=torch.bool - ) + all_heads = torch.zeros(self.dims.n_text_layer, self.dims.n_text_head, dtype=torch.bool) all_heads[self.dims.n_text_layer // 2 :] = True # self.register_buffer("alignment_heads", all_heads.to_sparse(), persistent=False) # alignment_heads_dense = model.get_buffer("alignment_heads").to_dense() # model.register_buffer("alignment_heads", alignment_heads_dense, persistent=False) def set_alignment_heads(self, dump: bytes): - array = np.frombuffer( - gzip.decompress(base64.b85decode(dump)), dtype=bool - ).copy() - mask = torch.from_numpy(array).reshape( - self.dims.n_text_layer, self.dims.n_text_head - ) + array = np.frombuffer(gzip.decompress(base64.b85decode(dump)), dtype=bool).copy() + mask = torch.from_numpy(array).reshape(self.dims.n_text_layer, self.dims.n_text_head) self.register_buffer("alignment_heads", mask.to_sparse(), persistent=False) def embed_audio(self, mel: torch.Tensor): @@ -280,9 +263,7 @@ class Whisper(nn.Module): def logits(self, tokens: torch.Tensor, audio_features: torch.Tensor): return self.decoder(tokens, audio_features) - def forward( - self, mel: torch.Tensor, tokens: torch.Tensor - ) -> Dict[str, torch.Tensor]: + def forward(self, mel: torch.Tensor, tokens: torch.Tensor) -> Dict[str, torch.Tensor]: return self.decoder(tokens, self.encoder(mel)) @property @@ -332,4 +313,4 @@ class Whisper(nn.Module): detect_language = detect_language_function transcribe = transcribe_function - decode = decode_function \ No newline at end of file + decode = decode_function diff --git a/funasr/models/sense_voice/whisper_lib/normalizers/basic.py b/funasr/models/sense_voice/whisper_lib/normalizers/basic.py index a82403203..76addc520 100644 --- a/funasr/models/sense_voice/whisper_lib/normalizers/basic.py +++ b/funasr/models/sense_voice/whisper_lib/normalizers/basic.py @@ -30,15 +30,19 @@ def remove_symbols_and_diacritics(s: str, keep=""): and drop any diacritics (category 'Mn' and some manual mappings) """ return "".join( - c - if c in keep - else ADDITIONAL_DIACRITICS[c] - if c in ADDITIONAL_DIACRITICS - else "" - if unicodedata.category(c) == "Mn" - else " " - if unicodedata.category(c)[0] in "MSP" - else c + ( + c + if c in keep + else ( + ADDITIONAL_DIACRITICS[c] + if c in ADDITIONAL_DIACRITICS + else ( + "" + if unicodedata.category(c) == "Mn" + else " " if unicodedata.category(c)[0] in "MSP" else c + ) + ) + ) for c in unicodedata.normalize("NFKD", s) ) @@ -48,16 +52,13 @@ def remove_symbols(s: str): Replace any other markers, symbols, punctuations with a space, keeping diacritics """ return "".join( - " " if unicodedata.category(c)[0] in "MSP" else c - for c in unicodedata.normalize("NFKC", s) + " " if unicodedata.category(c)[0] in "MSP" else c for c in unicodedata.normalize("NFKC", s) ) class BasicTextNormalizer: def __init__(self, remove_diacritics: bool = False, split_letters: bool = False): - self.clean = ( - remove_symbols_and_diacritics if remove_diacritics else remove_symbols - ) + self.clean = remove_symbols_and_diacritics if remove_diacritics else remove_symbols self.split_letters = split_letters def __call__(self, s: str): @@ -69,8 +70,6 @@ class BasicTextNormalizer: if self.split_letters: s = " ".join(regex.findall(r"\X", s, regex.U)) - s = re.sub( - r"\s+", " ", s - ) # replace any successive whitespace characters with a space + s = re.sub(r"\s+", " ", s) # replace any successive whitespace characters with a space return s diff --git a/funasr/models/sense_voice/whisper_lib/normalizers/english.py b/funasr/models/sense_voice/whisper_lib/normalizers/english.py index 4932042bc..cb4653098 100644 --- a/funasr/models/sense_voice/whisper_lib/normalizers/english.py +++ b/funasr/models/sense_voice/whisper_lib/normalizers/english.py @@ -84,8 +84,7 @@ class EnglishNumberNormalizer: name.replace("y", "ies"): (value, "s") for name, value in self.tens.items() } self.tens_ordinal = { - name.replace("y", "ieth"): (value, "th") - for name, value in self.tens.items() + name.replace("y", "ieth"): (value, "th") for name, value in self.tens.items() } self.tens_suffixed = {**self.tens_plural, **self.tens_ordinal} @@ -132,8 +131,7 @@ class EnglishNumberNormalizer: "cents": "¢", } self.prefixes = set( - list(self.preceding_prefixers.values()) - + list(self.following_prefixers.values()) + list(self.preceding_prefixers.values()) + list(self.following_prefixers.values()) ) self.suffixers = { "per": {"cent": "%"}, @@ -223,9 +221,7 @@ class EnglishNumberNormalizer: if value is None: value = ones elif isinstance(value, str) or prev in self.ones: - if ( - prev in self.tens and ones < 10 - ): # replace the last zero with the digit + if prev in self.tens and ones < 10: # replace the last zero with the digit assert value[-1] == "0" value = value[:-1] + str(ones) else: diff --git a/funasr/models/sense_voice/whisper_lib/timing.py b/funasr/models/sense_voice/whisper_lib/timing.py index b695ead0a..ba9cb1399 100644 --- a/funasr/models/sense_voice/whisper_lib/timing.py +++ b/funasr/models/sense_voice/whisper_lib/timing.py @@ -27,9 +27,7 @@ def median_filter(x: torch.Tensor, filter_width: int): # `F.pad` does not support 1D or 2D inputs for reflect padding but supports 3D and 4D x = x[None, None, :] - assert ( - filter_width > 0 and filter_width % 2 == 1 - ), "`filter_width` should be an odd number" + assert filter_width > 0 and filter_width % 2 == 1, "`filter_width` should be an odd number" result = None x = F.pad(x, (filter_width // 2, filter_width // 2, 0, 0), mode="reflect") @@ -111,9 +109,7 @@ def dtw_cuda(x, BLOCK_SIZE=1024): M, N = x.shape assert M < BLOCK_SIZE, f"M should be smaller than {BLOCK_SIZE=}" - x_skew = ( - F.pad(x, (0, M + 1), value=np.inf).flatten()[: M * (N + M)].reshape(M, N + M) - ) + x_skew = F.pad(x, (0, M + 1), value=np.inf).flatten()[: M * (N + M)].reshape(M, N + M) x_skew = x_skew.T.contiguous() cost = torch.ones(N + M + 2, M + 2) * np.inf cost[0, 0] = 0 @@ -132,9 +128,7 @@ def dtw_cuda(x, BLOCK_SIZE=1024): BLOCK_SIZE=BLOCK_SIZE, ) - trace = trace.T.flatten()[: (M + 1) * (M + N + 3)].reshape(M + 1, M + N + 3)[ - :, : N + 1 - ] + trace = trace.T.flatten()[: (M + 1) * (M + N + 3)].reshape(M + 1, M + N + 3)[:, : N + 1] return backtrace(trace.cpu().numpy()) @@ -228,8 +222,7 @@ def find_alignment( start_times = jump_times[word_boundaries[:-1]] end_times = jump_times[word_boundaries[1:]] word_probabilities = [ - np.mean(text_token_probs[i:j]) - for i, j in zip(word_boundaries[:-1], word_boundaries[1:]) + np.mean(text_token_probs[i:j]) for i, j in zip(word_boundaries[:-1], word_boundaries[1:]) ] return [ @@ -290,8 +283,7 @@ def add_word_timestamps( return text_tokens_per_segment = [ - [token for token in segment["tokens"] if token < tokenizer.eot] - for segment in segments + [token for token in segment["tokens"] if token < tokenizer.eot] for segment in segments ] text_tokens = list(itertools.chain.from_iterable(text_tokens_per_segment)) @@ -346,38 +338,22 @@ def add_word_timestamps( # twice the median word duration. if words[0]["end"] - last_speech_timestamp > median_duration * 4 and ( words[0]["end"] - words[0]["start"] > max_duration - or ( - len(words) > 1 - and words[1]["end"] - words[0]["start"] > max_duration * 2 - ) + or (len(words) > 1 and words[1]["end"] - words[0]["start"] > max_duration * 2) ): - if ( - len(words) > 1 - and words[1]["end"] - words[1]["start"] > max_duration - ): + if len(words) > 1 and words[1]["end"] - words[1]["start"] > max_duration: boundary = max(words[1]["end"] / 2, words[1]["end"] - max_duration) words[0]["end"] = words[1]["start"] = boundary words[0]["start"] = max(0, words[0]["end"] - max_duration) # prefer the segment-level start timestamp if the first word is too long. - if ( - segment["start"] < words[0]["end"] - and segment["start"] - 0.5 > words[0]["start"] - ): - words[0]["start"] = max( - 0, min(words[0]["end"] - median_duration, segment["start"]) - ) + if segment["start"] < words[0]["end"] and segment["start"] - 0.5 > words[0]["start"]: + words[0]["start"] = max(0, min(words[0]["end"] - median_duration, segment["start"])) else: segment["start"] = words[0]["start"] # prefer the segment-level end timestamp if the last word is too long. - if ( - segment["end"] > words[-1]["start"] - and segment["end"] + 0.5 < words[-1]["end"] - ): - words[-1]["end"] = max( - words[-1]["start"] + median_duration, segment["end"] - ) + if segment["end"] > words[-1]["start"] and segment["end"] + 0.5 < words[-1]["end"]: + words[-1]["end"] = max(words[-1]["start"] + median_duration, segment["end"]) else: segment["end"] = words[-1]["end"] diff --git a/funasr/models/sense_voice/whisper_lib/tokenizer.py b/funasr/models/sense_voice/whisper_lib/tokenizer.py index 4334a13d2..5b276c2a0 100644 --- a/funasr/models/sense_voice/whisper_lib/tokenizer.py +++ b/funasr/models/sense_voice/whisper_lib/tokenizer.py @@ -179,7 +179,7 @@ class Tokenizer: langs = tuple(LANGUAGES.keys())[: self.num_languages] sot_sequence = [sot] if self.language is not None: - if self.language == 'nospeech': + if self.language == "nospeech": sot_sequence.append(self.no_speech) else: sot_sequence.append(sot + 1 + langs.index(self.language)) @@ -204,7 +204,7 @@ class Tokenizer: This method decodes given tokens with timestamps tokens annotated, e.g. "<|1.08|>". """ return self.encoding.decode(token_ids, **kwargs) - + def get_vocab_size(self) -> int: return self.encoding.n_vocab @@ -291,9 +291,7 @@ class Tokenizer: keeping basic punctuations like commas, periods, question marks, exclamation points, etc. """ symbols = list('"#()*+/:;<=>@[\\]^_`{|}~「」『』') - symbols += ( - "<< >> <<< >>> -- --- -( -[ (' (\" (( )) ((( ))) [[ ]] {{ }} ♪♪ ♪♪♪".split() - ) + symbols += "<< >> <<< >>> -- --- -( -[ (' (\" (( )) ((( ))) [[ ]] {{ }} ♪♪ ♪♪♪".split() # symbols that may be a single token or multiple tokens depending on the tokenizer. # In case they're multiple tokens, suppress the first token, which is safe because: @@ -368,7 +366,7 @@ class Tokenizer: @lru_cache(maxsize=None) -def get_encoding(name: str = "gpt2", num_languages: int = 99, vocab_path:str=None): +def get_encoding(name: str = "gpt2", num_languages: int = 99, vocab_path: str = None): if vocab_path is None: vocab_path = os.path.join(os.path.dirname(__file__), "assets", f"{name}.tiktoken") @@ -379,7 +377,7 @@ def get_encoding(name: str = "gpt2", num_languages: int = 99, vocab_path:str=Non n_vocab = len(ranks) special_tokens = {} - if False: #name == "gpt2" or name == "multilingual": + if False: # name == "gpt2" or name == "multilingual": specials = [ "<|endoftext|>", "<|startoftranscript|>", @@ -437,7 +435,7 @@ def get_tokenizer( if language not in LANGUAGES: if language in TO_LANGUAGE_CODE: language = TO_LANGUAGE_CODE[language] - elif language == 'nospeech': + elif language == "nospeech": pass else: raise ValueError(f"Unsupported language: {language}") @@ -453,10 +451,6 @@ def get_tokenizer( if encoding_path is not None: encoding_name = encoding_path - encoding = get_encoding(name=encoding_name, num_languages=num_languages, vocab_path=vocab_path) - - return Tokenizer( - encoding=encoding, num_languages=num_languages, language=language, task=task - ) + return Tokenizer(encoding=encoding, num_languages=num_languages, language=language, task=task) diff --git a/funasr/models/sense_voice/whisper_lib/transcribe.py b/funasr/models/sense_voice/whisper_lib/transcribe.py index 1c075a201..5c5f49d80 100644 --- a/funasr/models/sense_voice/whisper_lib/transcribe.py +++ b/funasr/models/sense_voice/whisper_lib/transcribe.py @@ -146,9 +146,7 @@ def transcribe( _, probs = model.detect_language(mel_segment) decode_options["language"] = max(probs, key=probs.get) if verbose is not None: - print( - f"Detected language: {LANGUAGES[decode_options['language']].title()}" - ) + print(f"Detected language: {LANGUAGES[decode_options['language']].title()}") language: str = decode_options["language"] task: str = decode_options.get("task", "transcribe") @@ -176,9 +174,7 @@ def transcribe( warnings.warn("Word-level timestamps on translations may not be reliable.") def decode_with_fallback(segment: torch.Tensor) -> DecodingResult: - temperatures = ( - [temperature] if isinstance(temperature, (int, float)) else temperature - ) + temperatures = [temperature] if isinstance(temperature, (int, float)) else temperature decode_result = None for t in temperatures: @@ -200,10 +196,7 @@ def transcribe( and decode_result.compression_ratio > compression_ratio_threshold ): needs_fallback = True # too repetitive - if ( - logprob_threshold is not None - and decode_result.avg_logprob < logprob_threshold - ): + if logprob_threshold is not None and decode_result.avg_logprob < logprob_threshold: needs_fallback = True # average log probability is too low if ( no_speech_threshold is not None @@ -217,9 +210,7 @@ def transcribe( clip_idx = 0 seek = seek_clips[clip_idx][0] - input_stride = exact_div( - N_FRAMES, model.dims.n_audio_ctx - ) # mel frames per output token: 2 + input_stride = exact_div(N_FRAMES, model.dims.n_audio_ctx) # mel frames per output token: 2 time_precision = ( input_stride * HOP_LENGTH / SAMPLE_RATE ) # time per output token: 0.02 (seconds) @@ -233,9 +224,7 @@ def transcribe( else: initial_prompt_tokens = [] - def new_segment( - *, start: float, end: float, tokens: torch.Tensor, result: DecodingResult - ): + def new_segment(*, start: float, end: float, tokens: torch.Tensor, result: DecodingResult): tokens = tokens.tolist() text_tokens = [token for token in tokens if token < tokenizer.eot] return { @@ -251,9 +240,7 @@ def transcribe( } # show the progress bar when verbose is False (if True, transcribed text will be printed) - with tqdm.tqdm( - total=content_frames, unit="frames", disable=verbose is not False - ) as pbar: + with tqdm.tqdm(total=content_frames, unit="frames", disable=verbose is not False) as pbar: last_speech_timestamp = 0.0 # NOTE: This loop is obscurely flattened to make the diff readable. # A later commit should turn this into a simpler nested loop. @@ -282,10 +269,7 @@ def transcribe( if no_speech_threshold is not None: # no voice activity check should_skip = result.no_speech_prob > no_speech_threshold - if ( - logprob_threshold is not None - and result.avg_logprob > logprob_threshold - ): + if logprob_threshold is not None and result.avg_logprob > logprob_threshold: # don't skip if the logprob is high enough, despite the no_speech_prob should_skip = False @@ -334,12 +318,8 @@ def transcribe( last_slice = 0 for current_slice in slices: sliced_tokens = tokens[last_slice:current_slice] - start_timestamp_pos = ( - sliced_tokens[0].item() - tokenizer.timestamp_begin - ) - end_timestamp_pos = ( - sliced_tokens[-1].item() - tokenizer.timestamp_begin - ) + start_timestamp_pos = sliced_tokens[0].item() - tokenizer.timestamp_begin + end_timestamp_pos = sliced_tokens[-1].item() - tokenizer.timestamp_begin current_segments.append( new_segment( start=time_offset + start_timestamp_pos * time_precision, @@ -355,21 +335,14 @@ def transcribe( seek += segment_size else: # otherwise, ignore the unfinished segment and seek to the last timestamp - last_timestamp_pos = ( - tokens[last_slice - 1].item() - tokenizer.timestamp_begin - ) + last_timestamp_pos = tokens[last_slice - 1].item() - tokenizer.timestamp_begin seek += last_timestamp_pos * input_stride else: duration = segment_duration timestamps = tokens[timestamp_tokens.nonzero().flatten()] - if ( - len(timestamps) > 0 - and timestamps[-1].item() != tokenizer.timestamp_begin - ): + if len(timestamps) > 0 and timestamps[-1].item() != tokenizer.timestamp_begin: # no consecutive timestamps but it has a timestamp; use the last one. - last_timestamp_pos = ( - timestamps[-1].item() - tokenizer.timestamp_begin - ) + last_timestamp_pos = timestamps[-1].item() - tokenizer.timestamp_begin duration = last_timestamp_pos * time_precision current_segments.append( @@ -427,9 +400,7 @@ def transcribe( if not segment["words"]: continue if is_segment_anomaly(segment): - next_segment = next_words_segment( - current_segments[si + 1 :] - ) + next_segment = next_words_segment(current_segments[si + 1 :]) if next_segment is not None: hal_next_start = next_segment["words"][0]["start"] else: @@ -446,8 +417,7 @@ def transcribe( ) if silence_before and silence_after: seek = round( - max(time_offset + 1, segment["start"]) - * FRAMES_PER_SECOND + max(time_offset + 1, segment["start"]) * FRAMES_PER_SECOND ) if content_duration - segment["end"] < threshold: seek = content_frames @@ -475,9 +445,7 @@ def transcribe( all_segments.extend( [ {"id": i, **segment} - for i, segment in enumerate( - current_segments, start=len(all_segments) - ) + for i, segment in enumerate(current_segments, start=len(all_segments)) ] ) all_tokens.extend( diff --git a/funasr/models/sense_voice/whisper_lib/triton_ops.py b/funasr/models/sense_voice/whisper_lib/triton_ops.py index edd456414..9919595ce 100644 --- a/funasr/models/sense_voice/whisper_lib/triton_ops.py +++ b/funasr/models/sense_voice/whisper_lib/triton_ops.py @@ -11,9 +11,7 @@ except ImportError: @triton.jit -def dtw_kernel( - cost, trace, x, x_stride, cost_stride, trace_stride, N, M, BLOCK_SIZE: tl.constexpr -): +def dtw_kernel(cost, trace, x, x_stride, cost_stride, trace_stride, N, M, BLOCK_SIZE: tl.constexpr): offsets = tl.arange(0, BLOCK_SIZE) mask = offsets < M @@ -43,9 +41,7 @@ def dtw_kernel( @lru_cache(maxsize=None) def median_kernel(filter_width: int): @triton.jit - def kernel( - y, x, x_stride, y_stride, BLOCK_SIZE: tl.constexpr - ): # x.shape[-1] == filter_width + def kernel(y, x, x_stride, y_stride, BLOCK_SIZE: tl.constexpr): # x.shape[-1] == filter_width row_idx = tl.program_id(0) offsets = tl.arange(0, BLOCK_SIZE) mask = offsets < y_stride @@ -63,10 +59,7 @@ def median_kernel(filter_width: int): kernel.src = kernel.src.replace( " LOAD_ALL_ROWS_HERE", "\n".join( - [ - f" row{i} = tl.load(x_ptr + offsets + {i}, mask=mask)" - for i in range(filter_width) - ] + [f" row{i} = tl.load(x_ptr + offsets + {i}, mask=mask)" for i in range(filter_width)] ), ) kernel.src = kernel.src.replace( diff --git a/funasr/models/sense_voice/whisper_lib/utils.py b/funasr/models/sense_voice/whisper_lib/utils.py index 9b9b13862..5fc61256f 100644 --- a/funasr/models/sense_voice/whisper_lib/utils.py +++ b/funasr/models/sense_voice/whisper_lib/utils.py @@ -47,9 +47,7 @@ def compression_ratio(text) -> float: return len(text_bytes) / len(zlib.compress(text_bytes)) -def format_timestamp( - seconds: float, always_include_hours: bool = False, decimal_marker: str = "." -): +def format_timestamp(seconds: float, always_include_hours: bool = False, decimal_marker: str = "."): assert seconds >= 0, "non-negative timestamp expected" milliseconds = round(seconds * 1000.0) @@ -63,9 +61,7 @@ def format_timestamp( milliseconds -= seconds * 1_000 hours_marker = f"{hours:02d}:" if always_include_hours or hours > 0 else "" - return ( - f"{hours_marker}{minutes:02d}:{seconds:02d}{decimal_marker}{milliseconds:03d}" - ) + return f"{hours_marker}{minutes:02d}:{seconds:02d}{decimal_marker}{milliseconds:03d}" def get_start(segments: List[dict]) -> Optional[float]: @@ -88,30 +84,22 @@ class ResultWriter: def __init__(self, output_dir: str): self.output_dir = output_dir - def __call__( - self, result: dict, audio_path: str, options: Optional[dict] = None, **kwargs - ): + def __call__(self, result: dict, audio_path: str, options: Optional[dict] = None, **kwargs): audio_basename = os.path.basename(audio_path) audio_basename = os.path.splitext(audio_basename)[0] - output_path = os.path.join( - self.output_dir, audio_basename + "." + self.extension - ) + output_path = os.path.join(self.output_dir, audio_basename + "." + self.extension) with open(output_path, "w", encoding="utf-8") as f: self.write_result(result, file=f, options=options, **kwargs) - def write_result( - self, result: dict, file: TextIO, options: Optional[dict] = None, **kwargs - ): + def write_result(self, result: dict, file: TextIO, options: Optional[dict] = None, **kwargs): raise NotImplementedError class WriteTXT(ResultWriter): extension: str = "txt" - def write_result( - self, result: dict, file: TextIO, options: Optional[dict] = None, **kwargs - ): + def write_result(self, result: dict, file: TextIO, options: Optional[dict] = None, **kwargs): for segment in result["segments"]: print(segment["text"].strip(), file=file, flush=True) @@ -156,17 +144,10 @@ class SubtitlesWriter(ResultWriter): segment["words"][chunk_index : chunk_index + words_count] ): timing = original_timing.copy() - long_pause = ( - not preserve_segments and timing["start"] - last > 3.0 - ) + long_pause = not preserve_segments and timing["start"] - last > 3.0 has_room = line_len + len(timing["word"]) <= max_line_width seg_break = i == 0 and len(subtitle) > 0 and preserve_segments - if ( - line_len > 0 - and has_room - and not long_pause - and not seg_break - ): + if line_len > 0 and has_room and not long_pause and not seg_break: # line continuation line_len += len(timing["word"]) else: @@ -209,9 +190,7 @@ class SubtitlesWriter(ResultWriter): yield start, end, "".join( [ - re.sub(r"^(\s*)(.*)$", r"\1\2", word) - if j == i - else word + re.sub(r"^(\s*)(.*)$", r"\1\2", word) if j == i else word for j, word in enumerate(all_words) ] ) @@ -238,9 +217,7 @@ class WriteVTT(SubtitlesWriter): always_include_hours: bool = False decimal_marker: str = "." - def write_result( - self, result: dict, file: TextIO, options: Optional[dict] = None, **kwargs - ): + def write_result(self, result: dict, file: TextIO, options: Optional[dict] = None, **kwargs): print("WEBVTT\n", file=file) for start, end, text in self.iterate_result(result, options, **kwargs): print(f"{start} --> {end}\n{text}\n", file=file, flush=True) @@ -251,9 +228,7 @@ class WriteSRT(SubtitlesWriter): always_include_hours: bool = True decimal_marker: str = "," - def write_result( - self, result: dict, file: TextIO, options: Optional[dict] = None, **kwargs - ): + def write_result(self, result: dict, file: TextIO, options: Optional[dict] = None, **kwargs): for i, (start, end, text) in enumerate( self.iterate_result(result, options, **kwargs), start=1 ): @@ -272,9 +247,7 @@ class WriteTSV(ResultWriter): extension: str = "tsv" - def write_result( - self, result: dict, file: TextIO, options: Optional[dict] = None, **kwargs - ): + def write_result(self, result: dict, file: TextIO, options: Optional[dict] = None, **kwargs): print("start", "end", "text", sep="\t", file=file) for segment in result["segments"]: print(round(1000 * segment["start"]), file=file, end="\t") @@ -285,15 +258,11 @@ class WriteTSV(ResultWriter): class WriteJSON(ResultWriter): extension: str = "json" - def write_result( - self, result: dict, file: TextIO, options: Optional[dict] = None, **kwargs - ): + def write_result(self, result: dict, file: TextIO, options: Optional[dict] = None, **kwargs): json.dump(result, file) -def get_writer( - output_format: str, output_dir: str -) -> Callable[[dict, TextIO, dict], None]: +def get_writer(output_format: str, output_dir: str) -> Callable[[dict, TextIO, dict], None]: writers = { "txt": WriteTXT, "vtt": WriteVTT, @@ -305,9 +274,7 @@ def get_writer( if output_format == "all": all_writers = [writer(output_dir) for writer in writers.values()] - def write_all( - result: dict, file: TextIO, options: Optional[dict] = None, **kwargs - ): + def write_all(result: dict, file: TextIO, options: Optional[dict] = None, **kwargs): for writer in all_writers: writer(result, file, options, **kwargs) diff --git a/funasr/models/sond/attention.py b/funasr/models/sond/attention.py index 290ab034c..1af553477 100644 --- a/funasr/models/sond/attention.py +++ b/funasr/models/sond/attention.py @@ -17,6 +17,7 @@ import torch.nn.functional as F from funasr.models.transformer.utils.nets_utils import make_pad_mask import funasr.models.lora.layers as lora + class MultiHeadedAttention(nn.Module): """Multi-Head Attention layer. @@ -81,9 +82,7 @@ class MultiHeadedAttention(nn.Module): n_batch = value.size(0) if mask is not None: mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) - min_value = float( - numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min - ) + min_value = float(numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min) scores = scores.masked_fill(mask, min_value) self.attn = torch.softmax(scores, dim=-1).masked_fill( mask, 0.0 @@ -118,7 +117,6 @@ class MultiHeadedAttention(nn.Module): return self.forward_attention(v, scores, mask) - class RelPositionMultiHeadedAttention(MultiHeadedAttention): """Multi-Head Attention layer with relative position encoding (new implementation). @@ -164,7 +162,7 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention): x_padded = x_padded.view(*x.size()[:2], x.size(3) + 1, x.size(2)) x = x_padded[:, :, 1:].view_as(x)[ :, :, :, : x.size(-1) // 2 + 1 - ] # only keep the positions from 0 to time2 + ] # only keep the positions from 0 to time2 if self.zero_triu: ones = torch.ones((x.size(2), x.size(3)), device=x.device) @@ -211,17 +209,11 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention): matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) matrix_bd = self.rel_shift(matrix_bd) - scores = (matrix_ac + matrix_bd) / math.sqrt( - self.d_k - ) # (batch, head, time1, time2) + scores = (matrix_ac + matrix_bd) / math.sqrt(self.d_k) # (batch, head, time1, time2) return self.forward_attention(v, scores, mask) - - - - class MultiHeadSelfAttention(nn.Module): """Multi-Head Attention layer. @@ -261,9 +253,15 @@ class MultiHeadSelfAttention(nn.Module): b, t, d = x.size() q_k_v = self.linear_q_k_v(x) q, k, v = torch.split(q_k_v, int(self.h * self.d_k), dim=-1) - q_h = torch.reshape(q, (b, t, self.h, self.d_k)).transpose(1, 2) # (batch, head, time1, d_k) - k_h = torch.reshape(k, (b, t, self.h, self.d_k)).transpose(1, 2) # (batch, head, time2, d_k) - v_h = torch.reshape(v, (b, t, self.h, self.d_k)).transpose(1, 2) # (batch, head, time2, d_k) + q_h = torch.reshape(q, (b, t, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time1, d_k) + k_h = torch.reshape(k, (b, t, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time2, d_k) + v_h = torch.reshape(v, (b, t, self.h, self.d_k)).transpose( + 1, 2 + ) # (batch, head, time2, d_k) return q_h, k_h, v_h, v @@ -287,9 +285,7 @@ class MultiHeadSelfAttention(nn.Module): mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) - min_value = float( - numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min - ) + min_value = float(numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min) scores = scores.masked_fill(mask, min_value) self.attn = torch.softmax(scores, dim=-1).masked_fill( mask, 0.0 @@ -324,5 +320,3 @@ class MultiHeadSelfAttention(nn.Module): scores = torch.matmul(q_h, k_h.transpose(-2, -1)) att_outs = self.forward_attention(v_h, scores, mask, mask_att_chunk_encoder) return att_outs - - diff --git a/funasr/models/sond/e2e_diar_sond.py b/funasr/models/sond/e2e_diar_sond.py index 21d1d5965..0f5a7c27c 100644 --- a/funasr/models/sond/e2e_diar_sond.py +++ b/funasr/models/sond/e2e_diar_sond.py @@ -96,8 +96,12 @@ class DiarSondModel(FunASRModel): ) self.criterion_bce = SequenceBinaryCrossEntropy(normalize_length=length_normalized_loss) self.pse_embedding = self.generate_pse_embedding() - self.power_weight = torch.from_numpy(2 ** np.arange(max_spk_num)[np.newaxis, np.newaxis, :]).float() - self.int_token_arr = torch.from_numpy(np.array(self.token_list).astype(int)[np.newaxis, np.newaxis, :]).int() + self.power_weight = torch.from_numpy( + 2 ** np.arange(max_spk_num)[np.newaxis, np.newaxis, :] + ).float() + self.int_token_arr = torch.from_numpy( + np.array(self.token_list).astype(int)[np.newaxis, np.newaxis, :] + ).int() self.speaker_discrimination_loss_weight = speaker_discrimination_loss_weight self.inter_score_loss_weight = inter_score_loss_weight self.forward_steps = 0 @@ -107,8 +111,11 @@ class DiarSondModel(FunASRModel): def get_regularize_parameters(self): to_regularize_parameters, normal_parameters = [], [] for name, param in self.named_parameters(): - if ("encoder" in name and "weight" in name and "bn" not in name and - ("conv2" in name or "conv1" in name or "conv_sc" in name or "dense" in name) + if ( + "encoder" in name + and "weight" in name + and "bn" not in name + and ("conv2" in name or "conv1" in name or "conv_sc" in name or "dense" in name) ): to_regularize_parameters.append((name, param)) else: @@ -128,8 +135,11 @@ class DiarSondModel(FunASRModel): raw_profile: B, N, D raw_binary_labels: B, T, N """ - assert raw_profile.shape[1] == raw_binary_labels.shape[2], \ - "Num profile: {}, Num label: {}".format(raw_profile.shape[1], raw_binary_labels.shape[-1]) + assert ( + raw_profile.shape[1] == raw_binary_labels.shape[2] + ), "Num profile: {}, Num label: {}".format( + raw_profile.shape[1], raw_binary_labels.shape[-1] + ) profile = torch.clone(raw_profile) binary_labels = torch.clone(raw_binary_labels) bsz, num_spk = profile.shape[0], profile.shape[1] @@ -187,21 +197,24 @@ class DiarSondModel(FunASRModel): # 0b. augment profiles if self.profileaug is not None and self.training: speech, profile, binary_labels = self.profileaug( - speech, speech_lengths, - profile, profile_lengths, - binary_labels, binary_labels_lengths + speech, + speech_lengths, + profile, + profile_lengths, + binary_labels, + binary_labels_lengths, ) # 1. Calculate power-set encoding (PSE) labels - pad_bin_labels = F.pad(binary_labels, (0, self.max_spk_num - binary_labels.shape[2]), "constant", 0.0) + pad_bin_labels = F.pad( + binary_labels, (0, self.max_spk_num - binary_labels.shape[2]), "constant", 0.0 + ) raw_pse_labels = torch.sum(pad_bin_labels * self.power_weight, dim=2, keepdim=True) pse_labels = torch.argmax((raw_pse_labels.int() == self.int_token_arr).float(), dim=2) # 2. Network forward pred, inter_outputs = self.prediction_forward( - speech, speech_lengths, - profile, profile_lengths, - return_inter_outputs=True + speech, speech_lengths, profile, profile_lengths, return_inter_outputs=True ) (speech, speech_lengths), (profile, profile_lengths), (ci_score, cd_score) = inter_outputs @@ -219,13 +232,20 @@ class DiarSondModel(FunASRModel): loss_diar = self.classification_loss(pred, pse_labels, binary_labels_lengths) loss_spk_dis = self.speaker_discrimination_loss(profile, profile_lengths) - loss_inter_ci, loss_inter_cd = self.internal_score_loss(cd_score, ci_score, pse_labels, binary_labels_lengths) + loss_inter_ci, loss_inter_cd = self.internal_score_loss( + cd_score, ci_score, pse_labels, binary_labels_lengths + ) regularizer_loss = None if self.model_regularizer_weight > 0 and self.to_regularize_parameters is not None: regularizer_loss = self.calculate_regularizer_loss() - label_mask = make_pad_mask(binary_labels_lengths, maxlen=pse_labels.shape[1]).to(pse_labels.device) - loss = (loss_diar + self.speaker_discrimination_loss_weight * loss_spk_dis - + self.inter_score_loss_weight * (loss_inter_ci + loss_inter_cd)) + label_mask = make_pad_mask(binary_labels_lengths, maxlen=pse_labels.shape[1]).to( + pse_labels.device + ) + loss = ( + loss_diar + + self.speaker_discrimination_loss_weight * loss_spk_dis + + self.inter_score_loss_weight * (loss_inter_ci + loss_inter_cd) + ) # if regularizer_loss is not None: # loss = loss + regularizer_loss * self.model_regularizer_weight @@ -242,7 +262,7 @@ class DiarSondModel(FunASRModel): ) = self.calc_diarization_error( pred=F.embedding(pred.argmax(dim=2) * (~label_mask), self.pse_embedding), label=F.embedding(pse_labels * (~label_mask), self.pse_embedding), - length=binary_labels_lengths + length=binary_labels_lengths, ) if speech_scored > 0 and num_frames > 0: @@ -285,36 +305,34 @@ class DiarSondModel(FunASRModel): return regularizer_loss def classification_loss( - self, - predictions: torch.Tensor, - labels: torch.Tensor, - prediction_lengths: torch.Tensor + self, predictions: torch.Tensor, labels: torch.Tensor, prediction_lengths: torch.Tensor ) -> torch.Tensor: mask = make_pad_mask(prediction_lengths, maxlen=labels.shape[1]) - pad_labels = labels.masked_fill( - mask.to(predictions.device), - value=self.ignore_id - ) + pad_labels = labels.masked_fill(mask.to(predictions.device), value=self.ignore_id) loss = self.criterion_diar(predictions.contiguous(), pad_labels) return loss def speaker_discrimination_loss( - self, - profile: torch.Tensor, - profile_lengths: torch.Tensor + self, profile: torch.Tensor, profile_lengths: torch.Tensor ) -> torch.Tensor: - profile_mask = (torch.linalg.norm(profile, ord=2, dim=2, keepdim=True) > 0).float() # (B, N, 1) + profile_mask = ( + torch.linalg.norm(profile, ord=2, dim=2, keepdim=True) > 0 + ).float() # (B, N, 1) mask = torch.matmul(profile_mask, profile_mask.transpose(1, 2)) # (B, N, N) mask = mask * (1.0 - torch.eye(self.max_spk_num).unsqueeze(0).to(mask)) eps = 1e-12 - coding_norm = torch.linalg.norm( - profile * profile_mask + (1 - profile_mask) * eps, - dim=2, keepdim=True - ) * profile_mask + coding_norm = ( + torch.linalg.norm( + profile * profile_mask + (1 - profile_mask) * eps, dim=2, keepdim=True + ) + * profile_mask + ) # profile: Batch, N, dim - cos_theta = F.cosine_similarity(profile.unsqueeze(2), profile.unsqueeze(1), dim=-1, eps=eps) * mask + cos_theta = ( + F.cosine_similarity(profile.unsqueeze(2), profile.unsqueeze(1), dim=-1, eps=eps) * mask + ) cos_theta = torch.clip(cos_theta, -1 + eps, 1 - eps) loss = (F.relu(mask * coding_norm * (cos_theta - 0.0))).sum() / mask.sum() @@ -322,20 +340,17 @@ class DiarSondModel(FunASRModel): def calculate_multi_labels(self, pse_labels, pse_labels_lengths): mask = make_pad_mask(pse_labels_lengths, maxlen=pse_labels.shape[1]) - padding_labels = pse_labels.masked_fill( - mask.to(pse_labels.device), - value=0 - ).to(pse_labels) + padding_labels = pse_labels.masked_fill(mask.to(pse_labels.device), value=0).to(pse_labels) multi_labels = F.embedding(padding_labels, self.pse_embedding) return multi_labels def internal_score_loss( - self, - cd_score: torch.Tensor, - ci_score: torch.Tensor, - pse_labels: torch.Tensor, - pse_labels_lengths: torch.Tensor + self, + cd_score: torch.Tensor, + ci_score: torch.Tensor, + pse_labels: torch.Tensor, + pse_labels_lengths: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: multi_labels = self.calculate_multi_labels(pse_labels, pse_labels_lengths) ci_loss = self.criterion_bce(ci_score, multi_labels, pse_labels_lengths) @@ -355,13 +370,15 @@ class DiarSondModel(FunASRModel): return {"feats": feats, "feats_lengths": feats_lengths} def encode_speaker( - self, - profile: torch.Tensor, - profile_lengths: torch.Tensor, + self, + profile: torch.Tensor, + profile_lengths: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: with autocast(False): if profile.shape[1] < self.max_spk_num: - profile = F.pad(profile, [0, 0, 0, self.max_spk_num-profile.shape[1], 0, 0], "constant", 0.0) + profile = F.pad( + profile, [0, 0, 0, self.max_spk_num - profile.shape[1], 0, 0], "constant", 0.0 + ) profile_mask = (torch.linalg.norm(profile, ord=2, dim=2, keepdim=True) > 0).float() profile = F.normalize(profile, dim=2) if self.speaker_encoder is not None: @@ -371,9 +388,9 @@ class DiarSondModel(FunASRModel): return profile, profile_lengths def encode_speech( - self, - speech: torch.Tensor, - speech_lengths: torch.Tensor, + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: if self.encoder is not None and self.inputs_type == "raw": speech, speech_lengths = self.encode(speech, speech_lengths) @@ -384,24 +401,21 @@ class DiarSondModel(FunASRModel): return speech, speech_lengths @staticmethod - def concate_speech_ivc( - speech: torch.Tensor, - ivc: torch.Tensor - ) -> torch.Tensor: + def concate_speech_ivc(speech: torch.Tensor, ivc: torch.Tensor) -> torch.Tensor: nn, tt = ivc.shape[1], speech.shape[1] - speech = speech.unsqueeze(dim=1) # B x 1 x T x D + speech = speech.unsqueeze(dim=1) # B x 1 x T x D speech = speech.expand(-1, nn, -1, -1) # B x N x T x D - ivc = ivc.unsqueeze(dim=2) # B x N x 1 x D - ivc = ivc.expand(-1, -1, tt, -1) # B x N x T x D + ivc = ivc.unsqueeze(dim=2) # B x N x 1 x D + ivc = ivc.expand(-1, -1, tt, -1) # B x N x T x D sd_in = torch.cat([speech, ivc], dim=3) # B x N x T x 2D return sd_in def calc_similarity( - self, - speech_encoder_outputs: torch.Tensor, - speaker_encoder_outputs: torch.Tensor, - seq_len: torch.Tensor = None, - spk_len: torch.Tensor = None, + self, + speech_encoder_outputs: torch.Tensor, + speaker_encoder_outputs: torch.Tensor, + seq_len: torch.Tensor = None, + spk_len: torch.Tensor = None, ) -> Tuple[torch.Tensor, torch.Tensor]: bb, tt = speech_encoder_outputs.shape[0], speech_encoder_outputs.shape[1] d_sph, d_spk = speech_encoder_outputs.shape[2], speaker_encoder_outputs.shape[2] @@ -430,12 +444,12 @@ class DiarSondModel(FunASRModel): return logits def prediction_forward( - self, - speech: torch.Tensor, - speech_lengths: torch.Tensor, - profile: torch.Tensor, - profile_lengths: torch.Tensor, - return_inter_outputs: bool = False, + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + profile: torch.Tensor, + profile_lengths: torch.Tensor, + return_inter_outputs: bool = False, ) -> [torch.Tensor, Optional[list]]: # speech encoding speech, speech_lengths = self.encode_speech(speech, speech_lengths) @@ -448,7 +462,11 @@ class DiarSondModel(FunASRModel): logits = self.post_net_forward(similarity, speech_lengths) if return_inter_outputs: - return logits, [(speech, speech_lengths), (profile, profile_lengths), (ci_simi, cd_simi)] + return logits, [ + (speech, speech_lengths), + (profile, profile_lengths), + (ci_simi, cd_simi), + ] return logits def encode( diff --git a/funasr/models/sond/encoder/ci_scorers.py b/funasr/models/sond/encoder/ci_scorers.py index c5f45b901..a60560b7c 100644 --- a/funasr/models/sond/encoder/ci_scorers.py +++ b/funasr/models/sond/encoder/ci_scorers.py @@ -7,9 +7,9 @@ class DotScorer(torch.nn.Module): super().__init__() def forward( - self, - xs_pad: torch.Tensor, - spk_emb: torch.Tensor, + self, + xs_pad: torch.Tensor, + spk_emb: torch.Tensor, ): # xs_pad: B, T, D # spk_emb: B, N, D @@ -22,9 +22,9 @@ class CosScorer(torch.nn.Module): super().__init__() def forward( - self, - xs_pad: torch.Tensor, - spk_emb: torch.Tensor, + self, + xs_pad: torch.Tensor, + spk_emb: torch.Tensor, ): # xs_pad: B, T, D # spk_emb: B, N, D diff --git a/funasr/models/sond/encoder/conv_encoder.py b/funasr/models/sond/encoder/conv_encoder.py index 218116025..c4f7098c7 100644 --- a/funasr/models/sond/encoder/conv_encoder.py +++ b/funasr/models/sond/encoder/conv_encoder.py @@ -17,24 +17,24 @@ from funasr.models.transformer.utils.repeat import repeat class EncoderLayer(nn.Module): def __init__( - self, - input_units, - num_units, - kernel_size=3, - activation="tanh", - stride=1, - include_batch_norm=False, - residual=False + self, + input_units, + num_units, + kernel_size=3, + activation="tanh", + stride=1, + include_batch_norm=False, + residual=False, ): super().__init__() left_padding = math.ceil((kernel_size - stride) / 2) right_padding = kernel_size - stride - left_padding self.conv_padding = nn.ConstantPad1d((left_padding, right_padding), 0.0) self.conv1d = nn.Conv1d( - input_units, - num_units, - kernel_size, - stride, + input_units, + num_units, + kernel_size, + stride, ) self.activation = self.get_activation(activation) if include_batch_norm: @@ -71,23 +71,23 @@ class ConvEncoder(AbsEncoder): """ def __init__( - self, - num_layers, - input_units, - num_units, - kernel_size=3, - dropout_rate=0.3, - position_encoder=None, - activation='tanh', - auxiliary_states=True, - out_units=None, - out_norm=False, - out_residual=False, - include_batchnorm=False, - regularization_weight=0.0, - stride=1, - tf2torch_tensor_name_prefix_torch: str = "speaker_encoder", - tf2torch_tensor_name_prefix_tf: str = "EAND/speaker_encoder", + self, + num_layers, + input_units, + num_units, + kernel_size=3, + dropout_rate=0.3, + position_encoder=None, + activation="tanh", + auxiliary_states=True, + out_units=None, + out_norm=False, + out_residual=False, + include_batchnorm=False, + regularization_weight=0.0, + stride=1, + tf2torch_tensor_name_prefix_torch: str = "speaker_encoder", + tf2torch_tensor_name_prefix_tf: str = "EAND/speaker_encoder", ): super().__init__() self._output_size = num_units @@ -125,8 +125,8 @@ class ConvEncoder(AbsEncoder): activation, self.stride[lnum], include_batchnorm, - residual=True if lnum > 0 else False - ) + residual=True if lnum > 0 else False, + ), ) if self.out_units is not None: @@ -137,7 +137,7 @@ class ConvEncoder(AbsEncoder): num_units, out_units, kernel_size, - ) + ) if self.out_norm: self.after_norm = LayerNorm(out_units) @@ -172,4 +172,3 @@ class ConvEncoder(AbsEncoder): outputs = outputs + inputs return outputs, ilens, None - diff --git a/funasr/models/sond/encoder/ecapa_tdnn_encoder.py b/funasr/models/sond/encoder/ecapa_tdnn_encoder.py index 878a3c032..1af8b7033 100644 --- a/funasr/models/sond/encoder/ecapa_tdnn_encoder.py +++ b/funasr/models/sond/encoder/ecapa_tdnn_encoder.py @@ -39,9 +39,7 @@ class _BatchNorm1d(nn.Module): if x.ndim == 3: x = x.reshape(shape_or[0] * shape_or[1], shape_or[2]) else: - x = x.reshape( - shape_or[0] * shape_or[1], shape_or[3], shape_or[2] - ) + x = x.reshape(shape_or[0] * shape_or[1], shape_or[3], shape_or[2]) elif not self.skip_transpose: x = x.transpose(-1, 1) @@ -105,9 +103,7 @@ class _Conv1d(nn.Module): x = x.unsqueeze(1) if self.padding == "same": - x = self._manage_padding( - x, self.kernel_size, self.dilation, self.stride - ) + x = self._manage_padding(x, self.kernel_size, self.dilation, self.stride) elif self.padding == "causal": num_pad = (self.kernel_size - 1) * self.dilation @@ -117,10 +113,7 @@ class _Conv1d(nn.Module): pass else: - raise ValueError( - "Padding must be 'same', 'valid' or 'causal'. Got " - + self.padding - ) + raise ValueError("Padding must be 'same', 'valid' or 'causal'. Got " + self.padding) wx = self.conv(x) @@ -133,7 +126,11 @@ class _Conv1d(nn.Module): return wx def _manage_padding( - self, x, kernel_size: int, dilation: int, stride: int, + self, + x, + kernel_size: int, + dilation: int, + stride: int, ): # Detecting input shape L_in = x.shape[-1] @@ -147,8 +144,7 @@ class _Conv1d(nn.Module): return x def _check_input_shape(self, shape): - """Checks the input shape and returns the number of input channels. - """ + """Checks the input shape and returns the number of input channels.""" if len(shape) == 2: self.unsqueeze = True @@ -158,15 +154,12 @@ class _Conv1d(nn.Module): elif len(shape) == 3: in_channels = shape[2] else: - raise ValueError( - "conv1d expects 2d, 3d inputs. Got " + str(len(shape)) - ) + raise ValueError("conv1d expects 2d, 3d inputs. Got " + str(len(shape))) # Kernel size must be odd if self.kernel_size % 2 == 0: raise ValueError( - "The field kernel size must be an odd number. Got %s." - % (self.kernel_size) + "The field kernel size must be an odd number. Got %s." % (self.kernel_size) ) return in_channels @@ -200,9 +193,9 @@ def length_to_mask(length, max_len=None, dtype=None, device=None): if max_len is None: max_len = length.max().long().item() # using arange to generate mask - mask = torch.arange( - max_len, device=length.device, dtype=length.dtype - ).expand(len(length), max_len) < length.unsqueeze(1) + mask = torch.arange(max_len, device=length.device, dtype=length.dtype).expand( + len(length), max_len + ) < length.unsqueeze(1) if dtype is None: dtype = length.dtype @@ -264,9 +257,7 @@ class Res2NetBlock(torch.nn.Module): torch.Size([8, 120, 64]) """ - def __init__( - self, in_channels, out_channels, scale=8, kernel_size=3, dilation=1 - ): + def __init__(self, in_channels, out_channels, scale=8, kernel_size=3, dilation=1): super(Res2NetBlock, self).__init__() assert in_channels % scale == 0 assert out_channels % scale == 0 @@ -326,13 +317,9 @@ class SEBlock(nn.Module): def __init__(self, in_channels, se_channels, out_channels): super(SEBlock, self).__init__() - self.conv1 = Conv1d( - in_channels=in_channels, out_channels=se_channels, kernel_size=1 - ) + self.conv1 = Conv1d(in_channels=in_channels, out_channels=se_channels, kernel_size=1) self.relu = torch.nn.ReLU(inplace=True) - self.conv2 = Conv1d( - in_channels=se_channels, out_channels=out_channels, kernel_size=1 - ) + self.conv2 = Conv1d(in_channels=se_channels, out_channels=out_channels, kernel_size=1) self.sigmoid = torch.nn.Sigmoid() def forward(self, x, lengths=None): @@ -382,9 +369,7 @@ class AttentiveStatisticsPooling(nn.Module): else: self.tdnn = TDNNBlock(channels, attention_channels, 1, 1) self.tanh = nn.Tanh() - self.conv = Conv1d( - in_channels=attention_channels, out_channels=channels, kernel_size=1 - ) + self.conv = Conv1d(in_channels=attention_channels, out_channels=channels, kernel_size=1) def forward(self, x, lengths=None): """Calculates mean and std for a batch (input tensor). @@ -398,9 +383,7 @@ class AttentiveStatisticsPooling(nn.Module): def _compute_statistics(x, m, dim=2, eps=self.eps): mean = (m * x).sum(dim) - std = torch.sqrt( - (m * (x - mean.unsqueeze(dim)).pow(2)).sum(dim).clamp(eps) - ) + std = torch.sqrt((m * (x - mean.unsqueeze(dim)).pow(2)).sum(dim).clamp(eps)) return mean, std if lengths is None: @@ -638,9 +621,12 @@ class ECAPA_TDNN(torch.nn.Module): for i in range(num_chunk): # B x C st, ed = i * self.window_shift, i * self.window_shift + self.window_size - x = self.asp(x[:, :, st: ed], - lengths=torch.clamp(lengths - i, 0, self.window_size) - if lengths is not None else None) + x = self.asp( + x[:, :, st:ed], + lengths=( + torch.clamp(lengths - i, 0, self.window_size) if lengths is not None else None + ), + ) x = self.asp_bn(x) x = self.fc(x) stat_list.append(x) diff --git a/funasr/models/sond/encoder/fsmn_encoder.py b/funasr/models/sond/encoder/fsmn_encoder.py index fb87ee8d7..9ec9912a4 100644 --- a/funasr/models/sond/encoder/fsmn_encoder.py +++ b/funasr/models/sond/encoder/fsmn_encoder.py @@ -18,16 +18,17 @@ from funasr.models.transformer.utils.multi_layer_conv import FsmnFeedForward class FsmnBlock(torch.nn.Module): def __init__( - self, - n_feat, - dropout_rate, - kernel_size, - fsmn_shift=0, + self, + n_feat, + dropout_rate, + kernel_size, + fsmn_shift=0, ): super().__init__() self.dropout = nn.Dropout(p=dropout_rate) - self.fsmn_block = nn.Conv1d(n_feat, n_feat, kernel_size, stride=1, - padding=0, groups=n_feat, bias=False) + self.fsmn_block = nn.Conv1d( + n_feat, n_feat, kernel_size, stride=1, padding=0, groups=n_feat, bias=False + ) # padding left_padding = (kernel_size - 1) // 2 if fsmn_shift > 0: @@ -53,14 +54,7 @@ class FsmnBlock(torch.nn.Module): class EncoderLayer(torch.nn.Module): - def __init__( - self, - in_size, - size, - feed_forward, - fsmn_block, - dropout_rate=0.0 - ): + def __init__(self, in_size, size, feed_forward, fsmn_block, dropout_rate=0.0): super().__init__() self.in_size = in_size self.size = size @@ -69,9 +63,7 @@ class EncoderLayer(torch.nn.Module): self.dropout = nn.Dropout(dropout_rate) def forward( - self, - xs_pad: torch.Tensor, - mask: torch.Tensor + self, xs_pad: torch.Tensor, mask: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: # xs_pad in Batch, Time, Dim @@ -86,24 +78,24 @@ class EncoderLayer(torch.nn.Module): class FsmnEncoder(AbsEncoder): - """Encoder using Fsmn - """ + """Encoder using Fsmn""" - def __init__(self, - in_units, - filter_size, - fsmn_num_layers, - dnn_num_layers, - num_memory_units=512, - ffn_inner_dim=2048, - dropout_rate=0.0, - shift=0, - position_encoder=None, - sample_rate=1, - out_units=None, - tf2torch_tensor_name_prefix_torch="post_net", - tf2torch_tensor_name_prefix_tf="EAND/post_net" - ): + def __init__( + self, + in_units, + filter_size, + fsmn_num_layers, + dnn_num_layers, + num_memory_units=512, + ffn_inner_dim=2048, + dropout_rate=0.0, + shift=0, + position_encoder=None, + sample_rate=1, + out_units=None, + tf2torch_tensor_name_prefix_torch="post_net", + tf2torch_tensor_name_prefix_tf="EAND/post_net", + ): """Initializes the parameters of the encoder. Args: @@ -148,14 +140,9 @@ class FsmnEncoder(AbsEncoder): ffn_inner_dim, num_memory_units, 1, - dropout_rate - ), - FsmnBlock( - num_memory_units, dropout_rate, - filter_size, - self.shift[lnum] - ) + ), + FsmnBlock(num_memory_units, dropout_rate, filter_size, self.shift[lnum]), ), ) @@ -167,7 +154,7 @@ class FsmnEncoder(AbsEncoder): num_memory_units, 1, dropout_rate, - ) + ), ) if out_units is not None: self.conv1d = nn.Conv1d(num_memory_units, out_units, 1, 1) @@ -176,10 +163,7 @@ class FsmnEncoder(AbsEncoder): return self.num_memory_units def forward( - self, - xs_pad: torch.Tensor, - ilens: torch.Tensor, - prev_states: torch.Tensor = None + self, xs_pad: torch.Tensor, ilens: torch.Tensor, prev_states: torch.Tensor = None ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: inputs = xs_pad if self.position_encoder is not None: @@ -194,4 +178,3 @@ class FsmnEncoder(AbsEncoder): inputs = self.conv1d(inputs.transpose(1, 2)).transpose(1, 2) return inputs, ilens, None - diff --git a/funasr/models/sond/encoder/resnet34_encoder.py b/funasr/models/sond/encoder/resnet34_encoder.py index 8bfe49142..ba577a3a4 100644 --- a/funasr/models/sond/encoder/resnet34_encoder.py +++ b/funasr/models/sond/encoder/resnet34_encoder.py @@ -64,8 +64,9 @@ class BasicBlock(torch.nn.Module): self.num_layer = num_layer for i in range(num_layer): - layer = BasicLayer(in_filters if i == 0 else filters, filters, - stride if i == 0 else 1, bn_momentum) + layer = BasicLayer( + in_filters if i == 0 else filters, filters, stride if i == 0 else 1, bn_momentum + ) self.add_module("layer_{}".format(i), layer) def forward(self, xs_pad, ilens): @@ -78,14 +79,14 @@ class BasicBlock(torch.nn.Module): class ResNet34(AbsEncoder): def __init__( - self, - input_size, - use_head_conv=True, - batchnorm_momentum=0.5, - use_head_maxpool=False, - num_nodes_pooling_layer=256, - layers_in_block=(3, 4, 6, 3), - filters_in_block=(32, 64, 128, 256), + self, + input_size, + use_head_conv=True, + batchnorm_momentum=0.5, + use_head_maxpool=False, + num_nodes_pooling_layer=256, + layers_in_block=(3, 4, 6, 3), + filters_in_block=(32, 64, 128, 256), ): super(ResNet34, self).__init__() @@ -98,8 +99,12 @@ class ResNet34(AbsEncoder): pre_filters = filters_in_block[0] if use_head_conv: - self.pre_conv = torch.nn.Conv2d(1, pre_filters, 3, 1, 1, bias=False, padding_mode="zeros") - self.pre_conv_bn = torch.nn.BatchNorm2d(pre_filters, eps=1e-3, momentum=batchnorm_momentum) + self.pre_conv = torch.nn.Conv2d( + 1, pre_filters, 3, 1, 1, bias=False, padding_mode="zeros" + ) + self.pre_conv_bn = torch.nn.BatchNorm2d( + pre_filters, eps=1e-3, momentum=batchnorm_momentum + ) if use_head_maxpool: self.head_maxpool = torch.nn.MaxPool2d(3, 1, padding=1) @@ -108,17 +113,21 @@ class ResNet34(AbsEncoder): if i == 0: in_filters = pre_filters if self.use_head_conv else 1 else: - in_filters = filters_in_block[i-1] + in_filters = filters_in_block[i - 1] - block = BasicBlock(in_filters, - filters=filters_in_block[i], - num_layer=layers_in_block[i], - stride=1 if i == 0 else 2, - bn_momentum=batchnorm_momentum) + block = BasicBlock( + in_filters, + filters=filters_in_block[i], + num_layer=layers_in_block[i], + stride=1 if i == 0 else 2, + bn_momentum=batchnorm_momentum, + ) self.add_module("block_{}".format(i), block) self.resnet0_dense = torch.nn.Conv2d(filters_in_block[-1], num_nodes_pooling_layer, 1) - self.resnet0_bn = torch.nn.BatchNorm2d(num_nodes_pooling_layer, eps=1e-3, momentum=batchnorm_momentum) + self.resnet0_bn = torch.nn.BatchNorm2d( + num_nodes_pooling_layer, eps=1e-3, momentum=batchnorm_momentum + ) self.time_ds_ratio = 8 @@ -126,15 +135,15 @@ class ResNet34(AbsEncoder): return self.num_nodes_pooling_layer def forward( - self, - xs_pad: torch.Tensor, - ilens: torch.Tensor, - prev_states: torch.Tensor = None + self, xs_pad: torch.Tensor, ilens: torch.Tensor, prev_states: torch.Tensor = None ) -> Tuple[torch.Tensor, torch.Tensor]: features = xs_pad - assert features.size(-1) == self.input_size, \ - "Dimension of features {} doesn't match the input_size {}.".format(features.size(-1), self.input_size) + assert ( + features.size(-1) == self.input_size + ), "Dimension of features {} doesn't match the input_size {}.".format( + features.size(-1), self.input_size + ) features = torch.unsqueeze(features, dim=1) if self.use_head_conv: features = self.pre_conv(features) @@ -155,21 +164,22 @@ class ResNet34(AbsEncoder): return features, resnet_out_lens + # Note: For training, this implement is not equivalent to tf because of the kernel_regularizer in tf.layers. # TODO: implement kernel_regularizer in torch with munal loss addition or weigth_decay in the optimizer class ResNet34_SP_L2Reg(AbsEncoder): def __init__( - self, - input_size, - use_head_conv=True, - batchnorm_momentum=0.5, - use_head_maxpool=False, - num_nodes_pooling_layer=256, - layers_in_block=(3, 4, 6, 3), - filters_in_block=(32, 64, 128, 256), - tf2torch_tensor_name_prefix_torch="encoder", - tf2torch_tensor_name_prefix_tf="EAND/speech_encoder", - tf_train_steps=720000, + self, + input_size, + use_head_conv=True, + batchnorm_momentum=0.5, + use_head_maxpool=False, + num_nodes_pooling_layer=256, + layers_in_block=(3, 4, 6, 3), + filters_in_block=(32, 64, 128, 256), + tf2torch_tensor_name_prefix_torch="encoder", + tf2torch_tensor_name_prefix_tf="EAND/speech_encoder", + tf_train_steps=720000, ): super(ResNet34_SP_L2Reg, self).__init__() @@ -185,8 +195,12 @@ class ResNet34_SP_L2Reg(AbsEncoder): pre_filters = filters_in_block[0] if use_head_conv: - self.pre_conv = torch.nn.Conv2d(1, pre_filters, 3, 1, 1, bias=False, padding_mode="zeros") - self.pre_conv_bn = torch.nn.BatchNorm2d(pre_filters, eps=1e-3, momentum=batchnorm_momentum) + self.pre_conv = torch.nn.Conv2d( + 1, pre_filters, 3, 1, 1, bias=False, padding_mode="zeros" + ) + self.pre_conv_bn = torch.nn.BatchNorm2d( + pre_filters, eps=1e-3, momentum=batchnorm_momentum + ) if use_head_maxpool: self.head_maxpool = torch.nn.MaxPool2d(3, 1, padding=1) @@ -195,17 +209,23 @@ class ResNet34_SP_L2Reg(AbsEncoder): if i == 0: in_filters = pre_filters if self.use_head_conv else 1 else: - in_filters = filters_in_block[i-1] + in_filters = filters_in_block[i - 1] - block = BasicBlock(in_filters, - filters=filters_in_block[i], - num_layer=layers_in_block[i], - stride=1 if i == 0 else 2, - bn_momentum=batchnorm_momentum) + block = BasicBlock( + in_filters, + filters=filters_in_block[i], + num_layer=layers_in_block[i], + stride=1 if i == 0 else 2, + bn_momentum=batchnorm_momentum, + ) self.add_module("block_{}".format(i), block) - self.resnet0_dense = torch.nn.Conv1d(filters_in_block[-1] * input_size // 8, num_nodes_pooling_layer, 1) - self.resnet0_bn = torch.nn.BatchNorm1d(num_nodes_pooling_layer, eps=1e-3, momentum=batchnorm_momentum) + self.resnet0_dense = torch.nn.Conv1d( + filters_in_block[-1] * input_size // 8, num_nodes_pooling_layer, 1 + ) + self.resnet0_bn = torch.nn.BatchNorm1d( + num_nodes_pooling_layer, eps=1e-3, momentum=batchnorm_momentum + ) self.time_ds_ratio = 8 @@ -213,15 +233,15 @@ class ResNet34_SP_L2Reg(AbsEncoder): return self.num_nodes_pooling_layer def forward( - self, - xs_pad: torch.Tensor, - ilens: torch.Tensor, - prev_states: torch.Tensor = None + self, xs_pad: torch.Tensor, ilens: torch.Tensor, prev_states: torch.Tensor = None ) -> Tuple[torch.Tensor, torch.Tensor]: features = xs_pad - assert features.size(-1) == self.input_size, \ - "Dimension of features {} doesn't match the input_size {}.".format(features.size(-1), self.input_size) + assert ( + features.size(-1) == self.input_size + ), "Dimension of features {} doesn't match the input_size {}.".format( + features.size(-1), self.input_size + ) features = torch.unsqueeze(features, dim=1) if self.use_head_conv: features = self.pre_conv(features) @@ -238,7 +258,7 @@ class ResNet34_SP_L2Reg(AbsEncoder): # B, C, T, F bb, cc, tt, ff = resnet_outs.shape - resnet_outs = torch.reshape(resnet_outs.permute(0, 3, 1, 2), [bb, ff*cc, tt]) + resnet_outs = torch.reshape(resnet_outs.permute(0, 3, 1, 2), [bb, ff * cc, tt]) features = self.resnet0_dense(resnet_outs) features = F.relu(features) features = self.resnet0_bn(features) @@ -248,22 +268,22 @@ class ResNet34_SP_L2Reg(AbsEncoder): class ResNet34Diar(ResNet34): def __init__( - self, - input_size, - embedding_node="resnet1_dense", - use_head_conv=True, - batchnorm_momentum=0.5, - use_head_maxpool=False, - num_nodes_pooling_layer=256, - layers_in_block=(3, 4, 6, 3), - filters_in_block=(32, 64, 128, 256), - num_nodes_resnet1=256, - num_nodes_last_layer=256, - pooling_type="window_shift", - pool_size=20, - stride=1, - tf2torch_tensor_name_prefix_torch="encoder", - tf2torch_tensor_name_prefix_tf="seq2seq/speech_encoder" + self, + input_size, + embedding_node="resnet1_dense", + use_head_conv=True, + batchnorm_momentum=0.5, + use_head_maxpool=False, + num_nodes_pooling_layer=256, + layers_in_block=(3, 4, 6, 3), + filters_in_block=(32, 64, 128, 256), + num_nodes_resnet1=256, + num_nodes_last_layer=256, + pooling_type="window_shift", + pool_size=20, + stride=1, + tf2torch_tensor_name_prefix_torch="encoder", + tf2torch_tensor_name_prefix_tf="seq2seq/speech_encoder", ): """ Author: Speech Lab, Alibaba Group, China @@ -291,10 +311,14 @@ class ResNet34Diar(ResNet34): self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf self.resnet1_dense = torch.nn.Linear(num_nodes_pooling_layer * 2, num_nodes_resnet1) - self.resnet1_bn = torch.nn.BatchNorm1d(num_nodes_resnet1, eps=1e-3, momentum=batchnorm_momentum) + self.resnet1_bn = torch.nn.BatchNorm1d( + num_nodes_resnet1, eps=1e-3, momentum=batchnorm_momentum + ) self.resnet2_dense = torch.nn.Linear(num_nodes_resnet1, num_nodes_last_layer) - self.resnet2_bn = torch.nn.BatchNorm1d(num_nodes_last_layer, eps=1e-3, momentum=batchnorm_momentum) + self.resnet2_bn = torch.nn.BatchNorm1d( + num_nodes_last_layer, eps=1e-3, momentum=batchnorm_momentum + ) def output_size(self) -> int: if self.embedding_node.startswith("resnet1"): @@ -305,19 +329,21 @@ class ResNet34Diar(ResNet34): return self.num_nodes_pooling_layer def forward( - self, - xs_pad: torch.Tensor, - ilens: torch.Tensor, - prev_states: torch.Tensor = None, + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, + prev_states: torch.Tensor = None, ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: endpoints = OrderedDict() res_out, ilens = super().forward(xs_pad, ilens) endpoints["resnet0_bn"] = res_out if self.pooling_type == "frame_gsp": - features = statistic_pooling(res_out, ilens, (3, )) + features = statistic_pooling(res_out, ilens, (3,)) else: - features, ilens = windowed_statistic_pooling(res_out, ilens, (2, 3), self.pool_size, self.stride) + features, ilens = windowed_statistic_pooling( + res_out, ilens, (2, 3), self.pool_size, self.stride + ) features = features.transpose(1, 2) endpoints["pooling"] = features @@ -340,22 +366,22 @@ class ResNet34Diar(ResNet34): class ResNet34SpL2RegDiar(ResNet34_SP_L2Reg): def __init__( - self, - input_size, - embedding_node="resnet1_dense", - use_head_conv=True, - batchnorm_momentum=0.5, - use_head_maxpool=False, - num_nodes_pooling_layer=256, - layers_in_block=(3, 4, 6, 3), - filters_in_block=(32, 64, 128, 256), - num_nodes_resnet1=256, - num_nodes_last_layer=256, - pooling_type="window_shift", - pool_size=20, - stride=1, - tf2torch_tensor_name_prefix_torch="encoder", - tf2torch_tensor_name_prefix_tf="seq2seq/speech_encoder" + self, + input_size, + embedding_node="resnet1_dense", + use_head_conv=True, + batchnorm_momentum=0.5, + use_head_maxpool=False, + num_nodes_pooling_layer=256, + layers_in_block=(3, 4, 6, 3), + filters_in_block=(32, 64, 128, 256), + num_nodes_resnet1=256, + num_nodes_last_layer=256, + pooling_type="window_shift", + pool_size=20, + stride=1, + tf2torch_tensor_name_prefix_torch="encoder", + tf2torch_tensor_name_prefix_tf="seq2seq/speech_encoder", ): """ Author: Speech Lab, Alibaba Group, China @@ -383,10 +409,14 @@ class ResNet34SpL2RegDiar(ResNet34_SP_L2Reg): self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf self.resnet1_dense = torch.nn.Linear(num_nodes_pooling_layer * 2, num_nodes_resnet1) - self.resnet1_bn = torch.nn.BatchNorm1d(num_nodes_resnet1, eps=1e-3, momentum=batchnorm_momentum) + self.resnet1_bn = torch.nn.BatchNorm1d( + num_nodes_resnet1, eps=1e-3, momentum=batchnorm_momentum + ) self.resnet2_dense = torch.nn.Linear(num_nodes_resnet1, num_nodes_last_layer) - self.resnet2_bn = torch.nn.BatchNorm1d(num_nodes_last_layer, eps=1e-3, momentum=batchnorm_momentum) + self.resnet2_bn = torch.nn.BatchNorm1d( + num_nodes_last_layer, eps=1e-3, momentum=batchnorm_momentum + ) def output_size(self) -> int: if self.embedding_node.startswith("resnet1"): @@ -397,19 +427,21 @@ class ResNet34SpL2RegDiar(ResNet34_SP_L2Reg): return self.num_nodes_pooling_layer def forward( - self, - xs_pad: torch.Tensor, - ilens: torch.Tensor, - prev_states: torch.Tensor = None, + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, + prev_states: torch.Tensor = None, ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: endpoints = OrderedDict() res_out, ilens = super().forward(xs_pad, ilens) endpoints["resnet0_bn"] = res_out if self.pooling_type == "frame_gsp": - features = statistic_pooling(res_out, ilens, (2, )) + features = statistic_pooling(res_out, ilens, (2,)) else: - features, ilens = windowed_statistic_pooling(res_out, ilens, (2, ), self.pool_size, self.stride) + features, ilens = windowed_statistic_pooling( + res_out, ilens, (2,), self.pool_size, self.stride + ) features = features.transpose(1, 2) endpoints["pooling"] = features @@ -428,4 +460,3 @@ class ResNet34SpL2RegDiar(ResNet34_SP_L2Reg): endpoints["resnet2_bn"] = features return endpoints[self.embedding_node], ilens, None - diff --git a/funasr/models/sond/encoder/self_attention_encoder.py b/funasr/models/sond/encoder/self_attention_encoder.py index f3c473644..2e979b1be 100644 --- a/funasr/models/sond/encoder/self_attention_encoder.py +++ b/funasr/models/sond/encoder/self_attention_encoder.py @@ -87,7 +87,9 @@ class EncoderLayer(nn.Module): x = self.norm1(x) if self.concat_after: - x_concat = torch.cat((x, self.self_attn(x, mask, mask_att_chunk_encoder=mask_att_chunk_encoder)), dim=-1) + x_concat = torch.cat( + (x, self.self_attn(x, mask, mask_att_chunk_encoder=mask_att_chunk_encoder)), dim=-1 + ) if self.in_size == self.size: x = residual + stoch_layer_coeff * self.concat_linear(x_concat) else: @@ -207,32 +209,36 @@ class SelfAttentionEncoder(AbsEncoder): self.encoders = repeat( num_blocks, - lambda lnum: EncoderLayer( - output_size, - output_size, - MultiHeadSelfAttention( - attention_heads, + lambda lnum: ( + EncoderLayer( output_size, output_size, - attention_dropout_rate, - ), - positionwise_layer(*positionwise_layer_args), - dropout_rate, - normalize_before, - concat_after, - ) if lnum > 0 else EncoderLayer( - input_size, - output_size, - MultiHeadSelfAttention( - attention_heads, - input_size if input_layer == "pe" or input_layer == "null" else output_size, + MultiHeadSelfAttention( + attention_heads, + output_size, + output_size, + attention_dropout_rate, + ), + positionwise_layer(*positionwise_layer_args), + dropout_rate, + normalize_before, + concat_after, + ) + if lnum > 0 + else EncoderLayer( + input_size, output_size, - attention_dropout_rate, - ), - positionwise_layer(*positionwise_layer_args), - dropout_rate, - normalize_before, - concat_after, + MultiHeadSelfAttention( + attention_heads, + input_size if input_layer == "pe" or input_layer == "null" else output_size, + output_size, + attention_dropout_rate, + ), + positionwise_layer(*positionwise_layer_args), + dropout_rate, + normalize_before, + concat_after, + ) ), ) if self.normalize_before: @@ -270,7 +276,7 @@ class SelfAttentionEncoder(AbsEncoder): position embedded tensor and mask """ masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device) - xs_pad = xs_pad * self.output_size()**0.5 + xs_pad = xs_pad * self.output_size() ** 0.5 if self.embed is None: xs_pad = xs_pad elif ( @@ -325,4 +331,3 @@ class SelfAttentionEncoder(AbsEncoder): if len(intermediate_outs) > 0: return (xs_pad, intermediate_outs), olens, None return xs_pad, olens, None - diff --git a/funasr/models/sond/label_aggregation.py b/funasr/models/sond/label_aggregation.py index 7ec06aae1..dddc4c4bf 100644 --- a/funasr/models/sond/label_aggregation.py +++ b/funasr/models/sond/label_aggregation.py @@ -90,9 +90,7 @@ class LabelAggregateMaxPooling(torch.nn.Module): self.hop_length = hop_length def extra_repr(self): - return ( - f"hop_length={self.hop_length}, " - ) + return f"hop_length={self.hop_length}, " def forward( self, input: torch.Tensor, ilens: torch.Tensor = None @@ -107,7 +105,9 @@ class LabelAggregateMaxPooling(torch.nn.Module): """ - output = F.max_pool1d(input.transpose(1, 2), self.hop_length, self.hop_length).transpose(1, 2) + output = F.max_pool1d(input.transpose(1, 2), self.hop_length, self.hop_length).transpose( + 1, 2 + ) olens = ilens // self.hop_length - return output.to(input.dtype), olens \ No newline at end of file + return output.to(input.dtype), olens diff --git a/funasr/models/sond/pooling/pooling_layers.py b/funasr/models/sond/pooling/pooling_layers.py index 0aa10fefb..0a426a94f 100644 --- a/funasr/models/sond/pooling/pooling_layers.py +++ b/funasr/models/sond/pooling/pooling_layers.py @@ -59,8 +59,8 @@ class TSTP(nn.Module): class ASTP(nn.Module): - """ Attentive statistics pooling: Channel- and context-dependent - statistics pooling, first used in ECAPA_TDNN. + """Attentive statistics pooling: Channel- and context-dependent + statistics pooling, first used in ECAPA_TDNN. """ def __init__(self, in_dim, bottleneck_dim=128, global_context_att=False): @@ -71,14 +71,15 @@ class ASTP(nn.Module): # need to transpose inputs. if global_context_att: self.linear1 = nn.Conv1d( - in_dim * 3, bottleneck_dim, - kernel_size=1) # equals W and b in the paper + in_dim * 3, bottleneck_dim, kernel_size=1 + ) # equals W and b in the paper else: self.linear1 = nn.Conv1d( - in_dim, bottleneck_dim, - kernel_size=1) # equals W and b in the paper - self.linear2 = nn.Conv1d(bottleneck_dim, in_dim, - kernel_size=1) # equals V and k in the paper + in_dim, bottleneck_dim, kernel_size=1 + ) # equals W and b in the paper + self.linear2 = nn.Conv1d( + bottleneck_dim, in_dim, kernel_size=1 + ) # equals V and k in the paper def forward(self, x): """ @@ -92,17 +93,15 @@ class ASTP(nn.Module): if self.global_context_att: context_mean = torch.mean(x, dim=-1, keepdim=True).expand_as(x) - context_std = torch.sqrt( - torch.var(x, dim=-1, keepdim=True) + 1e-10).expand_as(x) + context_std = torch.sqrt(torch.var(x, dim=-1, keepdim=True) + 1e-10).expand_as(x) x_in = torch.cat((x, context_mean, context_std), dim=1) else: x_in = x # DON'T use ReLU here! ReLU may be hard to converge. - alpha = torch.tanh( - self.linear1(x_in)) # alpha = F.relu(self.linear1(x_in)) + alpha = torch.tanh(self.linear1(x_in)) # alpha = F.relu(self.linear1(x_in)) alpha = torch.softmax(self.linear2(alpha), dim=2) mean = torch.sum(alpha * x, dim=2) - var = torch.sum(alpha * (x ** 2), dim=2) - mean ** 2 + var = torch.sum(alpha * (x**2), dim=2) - mean**2 std = torch.sqrt(var.clamp(min=1e-10)) return torch.cat([mean, std], dim=1) diff --git a/funasr/models/sond/pooling/statistic_pooling.py b/funasr/models/sond/pooling/statistic_pooling.py index a3d7e34b0..91db2393e 100644 --- a/funasr/models/sond/pooling/statistic_pooling.py +++ b/funasr/models/sond/pooling/statistic_pooling.py @@ -7,11 +7,12 @@ import math VAR2STD_EPSILON = 1e-12 + class StatisticPooling(torch.nn.Module): def __init__(self, pooling_dim: Union[int, Tuple] = 2, eps=1e-12): super(StatisticPooling, self).__init__() if isinstance(pooling_dim, int): - pooling_dim = (pooling_dim, ) + pooling_dim = (pooling_dim,) self.pooling_dim = pooling_dim self.eps = eps @@ -22,11 +23,13 @@ class StatisticPooling(torch.nn.Module): masks = torch.ones_like(xs_pad).to(xs_pad) else: masks = make_non_pad_mask(ilens, xs_pad, length_dim=2).to(xs_pad) - mean = (torch.sum(xs_pad, dim=self.pooling_dim, keepdim=True) / - torch.sum(masks, dim=self.pooling_dim, keepdim=True)) + mean = torch.sum(xs_pad, dim=self.pooling_dim, keepdim=True) / torch.sum( + masks, dim=self.pooling_dim, keepdim=True + ) squared_difference = torch.pow(xs_pad - mean, 2.0) - variance = (torch.sum(squared_difference, dim=self.pooling_dim, keepdim=True) / - torch.sum(masks, dim=self.pooling_dim, keepdim=True)) + variance = torch.sum(squared_difference, dim=self.pooling_dim, keepdim=True) / torch.sum( + masks, dim=self.pooling_dim, keepdim=True + ) for i in reversed(self.pooling_dim): mean, variance = torch.squeeze(mean, dim=i), torch.squeeze(variance, dim=i) @@ -40,9 +43,7 @@ class StatisticPooling(torch.nn.Module): def statistic_pooling( - xs_pad: torch.Tensor, - ilens: torch.Tensor = None, - pooling_dim: Tuple = (2, 3) + xs_pad: torch.Tensor, ilens: torch.Tensor = None, pooling_dim: Tuple = (2, 3) ) -> torch.Tensor: # xs_pad in (Batch, Channel, Time, Frequency) @@ -50,11 +51,13 @@ def statistic_pooling( seq_mask = torch.ones_like(xs_pad).to(xs_pad) else: seq_mask = make_non_pad_mask(ilens, xs_pad, length_dim=2).to(xs_pad) - mean = (torch.sum(xs_pad, dim=pooling_dim, keepdim=True) / - torch.sum(seq_mask, dim=pooling_dim, keepdim=True)) + mean = torch.sum(xs_pad, dim=pooling_dim, keepdim=True) / torch.sum( + seq_mask, dim=pooling_dim, keepdim=True + ) squared_difference = torch.pow(xs_pad - mean, 2.0) - variance = (torch.sum(squared_difference, dim=pooling_dim, keepdim=True) / - torch.sum(seq_mask, dim=pooling_dim, keepdim=True)) + variance = torch.sum(squared_difference, dim=pooling_dim, keepdim=True) / torch.sum( + seq_mask, dim=pooling_dim, keepdim=True + ) for i in reversed(pooling_dim): mean, variance = torch.squeeze(mean, dim=i), torch.squeeze(variance, dim=i) @@ -68,11 +71,11 @@ def statistic_pooling( def windowed_statistic_pooling( - xs_pad: torch.Tensor, - ilens: torch.Tensor = None, - pooling_dim: Tuple = (2, 3), - pooling_size: int = 20, - pooling_stride: int = 1 + xs_pad: torch.Tensor, + ilens: torch.Tensor = None, + pooling_dim: Tuple = (2, 3), + pooling_size: int = 20, + pooling_stride: int = 1, ) -> Tuple[torch.Tensor, int]: # xs_pad in (Batch, Channel, Time, Frequency) @@ -87,8 +90,8 @@ def windowed_statistic_pooling( for i in range(num_chunk): # B x C - st, ed = i*pooling_stride, i*pooling_stride+pooling_size - stat = statistic_pooling(features[:, :, st: ed], pooling_dim=pooling_dim) + st, ed = i * pooling_stride, i * pooling_stride + pooling_size + stat = statistic_pooling(features[:, :, st:ed], pooling_dim=pooling_dim) stat_list.append(stat.unsqueeze(2)) # B x C x T diff --git a/funasr/models/sond/sv_decoder.py b/funasr/models/sond/sv_decoder.py index ee6b0e5e7..7d8484db9 100644 --- a/funasr/models/sond/sv_decoder.py +++ b/funasr/models/sond/sv_decoder.py @@ -5,19 +5,23 @@ from funasr.models.decoder.abs_decoder import AbsDecoder class DenseDecoder(AbsDecoder): def __init__( - self, - vocab_size, - encoder_output_size, - num_nodes_resnet1: int = 256, - num_nodes_last_layer: int = 256, - batchnorm_momentum: float = 0.5, + self, + vocab_size, + encoder_output_size, + num_nodes_resnet1: int = 256, + num_nodes_last_layer: int = 256, + batchnorm_momentum: float = 0.5, ): super(DenseDecoder, self).__init__() self.resnet1_dense = torch.nn.Linear(encoder_output_size, num_nodes_resnet1) - self.resnet1_bn = torch.nn.BatchNorm1d(num_nodes_resnet1, eps=1e-3, momentum=batchnorm_momentum) + self.resnet1_bn = torch.nn.BatchNorm1d( + num_nodes_resnet1, eps=1e-3, momentum=batchnorm_momentum + ) self.resnet2_dense = torch.nn.Linear(num_nodes_resnet1, num_nodes_last_layer) - self.resnet2_bn = torch.nn.BatchNorm1d(num_nodes_last_layer, eps=1e-3, momentum=batchnorm_momentum) + self.resnet2_bn = torch.nn.BatchNorm1d( + num_nodes_last_layer, eps=1e-3, momentum=batchnorm_momentum + ) self.output_dense = torch.nn.Linear(num_nodes_last_layer, vocab_size, bias=False) diff --git a/funasr/models/specaug/mask_along_axis.py b/funasr/models/specaug/mask_along_axis.py index 416c4ea7b..d4f237b3d 100644 --- a/funasr/models/specaug/mask_along_axis.py +++ b/funasr/models/specaug/mask_along_axis.py @@ -66,6 +66,7 @@ def mask_along_axis( spec = spec.view(*org_size) return spec, spec_lengths + def mask_along_axis_lfr( spec: torch.Tensor, spec_lengths: torch.Tensor, @@ -150,8 +151,7 @@ class MaskAlongAxis(torch.nn.Module): mask_width_range = (0, mask_width_range) if len(mask_width_range) != 2: raise TypeError( - f"mask_width_range must be a tuple of int and int values: " - f"{mask_width_range}", + f"mask_width_range must be a tuple of int and int values: " f"{mask_width_range}", ) assert mask_width_range[1] > mask_width_range[0] @@ -271,6 +271,7 @@ class MaskAlongAxisVariableMaxWidth(torch.nn.Module): ) return spec, spec_lengths + class MaskAlongAxisLFR(torch.nn.Module): def __init__( self, @@ -284,8 +285,7 @@ class MaskAlongAxisLFR(torch.nn.Module): mask_width_range = (0, mask_width_range) if len(mask_width_range) != 2: raise TypeError( - f"mask_width_range must be a tuple of int and int values: " - f"{mask_width_range}", + f"mask_width_range must be a tuple of int and int values: " f"{mask_width_range}", ) assert mask_width_range[1] > mask_width_range[0] @@ -333,4 +333,4 @@ class MaskAlongAxisLFR(torch.nn.Module): num_mask=self.num_mask, replace_with_zero=self.replace_with_zero, lfr_rate=self.lfr_rate, - ) \ No newline at end of file + ) diff --git a/funasr/models/specaug/profileaug.py b/funasr/models/specaug/profileaug.py index 448c579d6..669d323f8 100644 --- a/funasr/models/specaug/profileaug.py +++ b/funasr/models/specaug/profileaug.py @@ -4,6 +4,7 @@ import torch from torch.nn import functional as F import torch.nn as nn + class ProfileAug(nn.Module): """ Implement the augmentation for profiles including: @@ -11,15 +12,16 @@ class ProfileAug(nn.Module): - Merge aug: merge two profiles into one, labels are also merged into one, the other set to zero - Disturb aug: disturb some profile with others to simulate the inaccurate clustering centroids. """ + def __init__( - self, - apply_split_aug: bool = True, - split_aug_prob: float = 0.05, - apply_merge_aug: bool = True, - merge_aug_prob: float = 0.2, - apply_disturb_aug: bool = True, - disturb_aug_prob: float = 0.4, - disturb_alpha: float = 0.2, + self, + apply_split_aug: bool = True, + split_aug_prob: float = 0.05, + apply_merge_aug: bool = True, + merge_aug_prob: float = 0.2, + apply_disturb_aug: bool = True, + disturb_aug_prob: float = 0.4, + disturb_alpha: float = 0.2, ) -> None: super().__init__() self.apply_split_aug = apply_split_aug @@ -46,8 +48,9 @@ class ProfileAug(nn.Module): to_cover_idx = pad_spk_idx[torch.randint(len(pad_spk_idx), ())] disturb_vec = torch.randn((dim,)).to(profile) disturb_vec = F.normalize(disturb_vec, dim=-1) - profile[idx, to_cover_idx] = F.normalize(profile[idx, split_spk_idx] + - self.disturb_alpha * disturb_vec) + profile[idx, to_cover_idx] = F.normalize( + profile[idx, split_spk_idx] + self.disturb_alpha * disturb_vec + ) mask[idx, split_spk_idx] = 0 mask[idx, to_cover_idx] = 0 return profile, binary_labels, mask @@ -62,15 +65,19 @@ class ProfileAug(nn.Module): valid_spk_idx = torch.nonzero(profile_norm[idx] * mask[idx]) if len(valid_spk_idx) == 0: continue - to_merge = torch.randint(len(valid_spk_idx), (2, )) + to_merge = torch.randint(len(valid_spk_idx), (2,)) spk_idx_1, spk_idx_2 = valid_spk_idx[to_merge[0]], valid_spk_idx[to_merge[1]] # merge profile profile[idx, spk_idx_1] = profile[idx, spk_idx_1] + profile[idx, spk_idx_2] profile[idx, spk_idx_1] = F.normalize(profile[idx, spk_idx_1], dim=-1) profile[idx, spk_idx_2] = 0 # merge binary labels - binary_labels[idx, :, spk_idx_1] = binary_labels[idx, :, spk_idx_1] + binary_labels[idx, :, spk_idx_2] - binary_labels[idx, :, spk_idx_1] = (binary_labels[idx, :, spk_idx_1] > 0).to(binary_labels) + binary_labels[idx, :, spk_idx_1] = ( + binary_labels[idx, :, spk_idx_1] + binary_labels[idx, :, spk_idx_2] + ) + binary_labels[idx, :, spk_idx_1] = (binary_labels[idx, :, spk_idx_1] > 0).to( + binary_labels + ) binary_labels[idx, :, spk_idx_2] = 0 mask[idx, spk_idx_1] = 0 @@ -92,30 +99,44 @@ class ProfileAug(nn.Module): to_disturb_idx = pos_spk_idx[torch.randint(len(pos_spk_idx), ())] disturb_idx = valid_spk_idx[torch.randint(len(valid_spk_idx), ())] alpha = self.disturb_alpha * torch.rand(()).item() - profile[idx, to_disturb_idx] = ((1 - alpha) * profile[idx, to_disturb_idx] - + alpha * profile[idx, disturb_idx]) + profile[idx, to_disturb_idx] = (1 - alpha) * profile[ + idx, to_disturb_idx + ] + alpha * profile[idx, disturb_idx] profile[idx, to_disturb_idx] = F.normalize(profile[idx, to_disturb_idx], dim=-1) mask[idx, to_disturb_idx] = 0 return profile, binary_labels, mask def forward( - self, - speech: torch.Tensor, speech_lengths: torch.Tensor = None, - profile: torch.Tensor = None, profile_lengths: torch.Tensor = None, - binary_labels: torch.Tensor = None, labels_length: torch.Tensor = None + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor = None, + profile: torch.Tensor = None, + profile_lengths: torch.Tensor = None, + binary_labels: torch.Tensor = None, + labels_length: torch.Tensor = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: # copy inputs to avoid inplace-operation - speech, profile, binary_labels = torch.clone(speech), torch.clone(profile), torch.clone(binary_labels) + speech, profile, binary_labels = ( + torch.clone(speech), + torch.clone(profile), + torch.clone(binary_labels), + ) profile = F.normalize(profile, dim=-1) profile_mask = torch.ones(profile.shape[:2]).to(profile) if self.apply_disturb_aug: - profile, binary_labels, profile_mask = self.disturb_aug(profile, binary_labels, profile_mask) + profile, binary_labels, profile_mask = self.disturb_aug( + profile, binary_labels, profile_mask + ) if self.apply_split_aug: - profile, binary_labels, profile_mask = self.split_aug(profile, binary_labels, profile_mask) + profile, binary_labels, profile_mask = self.split_aug( + profile, binary_labels, profile_mask + ) if self.apply_merge_aug: - profile, binary_labels, profile_mask = self.merge_aug(profile, binary_labels, profile_mask) + profile, binary_labels, profile_mask = self.merge_aug( + profile, binary_labels, profile_mask + ) return speech, profile, binary_labels diff --git a/funasr/models/specaug/specaug.py b/funasr/models/specaug/specaug.py index 49f83e2e4..75a4b0b28 100644 --- a/funasr/models/specaug/specaug.py +++ b/funasr/models/specaug/specaug.py @@ -1,4 +1,5 @@ """SpecAugment module.""" + from typing import Optional from typing import Sequence from typing import Union @@ -11,6 +12,7 @@ from funasr.register import tables import torch.nn as nn + @tables.register("specaug_classes", "SpecAug") class SpecAug(nn.Module): """Implementation of SpecAug. @@ -40,9 +42,7 @@ class SpecAug(nn.Module): num_time_mask: int = 2, ): if not apply_time_warp and not apply_time_mask and not apply_freq_mask: - raise ValueError( - "Either one of time_warp, time_mask, or freq_mask should be applied" - ) + raise ValueError("Either one of time_warp, time_mask, or freq_mask should be applied") if ( apply_time_mask and (time_mask_width_range is not None) @@ -101,6 +101,7 @@ class SpecAug(nn.Module): x, x_lengths = self.time_mask(x, x_lengths) return x, x_lengths + @tables.register("specaug_classes", "SpecAugLFR") class SpecAugLFR(nn.Module): """Implementation of SpecAug. @@ -122,9 +123,7 @@ class SpecAugLFR(nn.Module): num_time_mask: int = 2, ): if not apply_time_warp and not apply_time_mask and not apply_freq_mask: - raise ValueError( - "Either one of time_warp, time_mask, or freq_mask should be applied" - ) + raise ValueError("Either one of time_warp, time_mask, or freq_mask should be applied") if ( apply_time_mask and (time_mask_width_range is not None) @@ -149,7 +148,7 @@ class SpecAugLFR(nn.Module): dim="freq", mask_width_range=freq_mask_width_range, num_mask=num_freq_mask, - lfr_rate=lfr_rate+1, + lfr_rate=lfr_rate + 1, ) else: @@ -184,4 +183,4 @@ class SpecAugLFR(nn.Module): x, x_lengths = self.freq_mask(x, x_lengths) if self.time_mask is not None: x, x_lengths = self.time_mask(x, x_lengths) - return x, x_lengths \ No newline at end of file + return x, x_lengths diff --git a/funasr/models/specaug/time_warp.py b/funasr/models/specaug/time_warp.py index d42820df7..ef3532f6a 100644 --- a/funasr/models/specaug/time_warp.py +++ b/funasr/models/specaug/time_warp.py @@ -1,4 +1,5 @@ """Time warp module.""" + import torch from funasr.models.transformer.utils.nets_utils import pad_list diff --git a/funasr/models/transducer/beam_search_transducer.py b/funasr/models/transducer/beam_search_transducer.py index f599615c0..aa9e34e1e 100644 --- a/funasr/models/transducer/beam_search_transducer.py +++ b/funasr/models/transducer/beam_search_transducer.py @@ -92,21 +92,18 @@ class BeamSearchTransducer: self.vocab_size = decoder.vocab_size - assert beam_size <= self.vocab_size, ( - "beam_size (%d) should be smaller than or equal to vocabulary size (%d)." - % ( - beam_size, - self.vocab_size, - ) + assert ( + beam_size <= self.vocab_size + ), "beam_size (%d) should be smaller than or equal to vocabulary size (%d)." % ( + beam_size, + self.vocab_size, ) self.beam_size = beam_size if search_type == "default": self.search_algorithm = self.default_beam_search elif search_type == "tsd": - assert max_sym_exp > 1, "max_sym_exp (%d) should be greater than one." % ( - max_sym_exp - ) + assert max_sym_exp > 1, "max_sym_exp (%d) should be greater than one." % (max_sym_exp) self.max_sym_exp = max_sym_exp self.search_algorithm = self.time_sync_decoding @@ -130,9 +127,7 @@ class BeamSearchTransducer: self.search_algorithm = self.modified_adaptive_expansion_search else: - raise NotImplementedError( - "Specified search type (%s) is not supported." % search_type - ) + raise NotImplementedError("Specified search type (%s) is not supported." % search_type) self.use_lm = lm is not None @@ -244,17 +239,12 @@ class BeamSearchTransducer: k_expansions = [] for i, hyp in enumerate(hyps): - hyp_i = [ - (int(k), hyp.score + float(v)) - for k, v in zip(topk_idx[i], topk_logp[i]) - ] + hyp_i = [(int(k), hyp.score + float(v)) for k, v in zip(topk_idx[i], topk_logp[i])] k_best_exp = max(hyp_i, key=lambda x: x[1])[1] k_expansions.append( sorted( - filter( - lambda x: (k_best_exp - self.expansion_gamma) <= x[1], hyp_i - ), + filter(lambda x: (k_best_exp - self.expansion_gamma) <= x[1], hyp_i), key=lambda x: x[1], reverse=True, ) @@ -342,9 +332,7 @@ class BeamSearchTransducer: if self.use_lm: lm_scores, lm_state = self.lm.score( - torch.LongTensor( - [self.sos] + max_hyp.yseq[1:], device=self.decoder.device - ), + torch.LongTensor([self.sos] + max_hyp.yseq[1:], device=self.decoder.device), max_hyp.lm_state, None, ) @@ -376,7 +364,7 @@ class BeamSearchTransducer: break return kept_hyps - + def align_length_sync_decoding( self, enc_out: torch.Tensor, diff --git a/funasr/models/transducer/joint_network.py b/funasr/models/transducer/joint_network.py index 7d424dbc9..955d6a0cf 100644 --- a/funasr/models/transducer/joint_network.py +++ b/funasr/models/transducer/joint_network.py @@ -39,9 +39,7 @@ class JointNetwork(torch.nn.Module): self.lin_out = torch.nn.Linear(joint_space_size, output_size) - self.joint_activation = get_activation( - joint_activation_type - ) + self.joint_activation = get_activation(joint_activation_type) def forward( self, diff --git a/funasr/models/transducer/model.py b/funasr/models/transducer/model.py index 8b876d26d..37925fcf6 100644 --- a/funasr/models/transducer/model.py +++ b/funasr/models/transducer/model.py @@ -98,24 +98,24 @@ class Transducer(torch.nn.Module): decoder_output_size, **joint_network_conf, ) - + self.criterion_transducer = None self.error_calculator = None - + self.use_auxiliary_ctc = auxiliary_ctc_weight > 0 self.use_auxiliary_lm_loss = auxiliary_lm_loss_weight > 0 - + if self.use_auxiliary_ctc: self.ctc_lin = torch.nn.Linear(encoder.output_size(), vocab_size) self.ctc_dropout_rate = auxiliary_ctc_dropout_rate - + if self.use_auxiliary_lm_loss: self.lm_lin = torch.nn.Linear(decoder.output_size, vocab_size) self.lm_loss_smoothing = auxiliary_lm_loss_smoothing - + self.transducer_weight = transducer_weight self.fastemit_lambda = fastemit_lambda - + self.auxiliary_ctc_weight = auxiliary_ctc_weight self.auxiliary_lm_loss_weight = auxiliary_lm_loss_weight self.blank_id = blank_id @@ -141,7 +141,7 @@ class Transducer(torch.nn.Module): self.beam_search = None self.ctc = None self.ctc_weight = 0.0 - + def forward( self, speech: torch.Tensor, @@ -161,29 +161,31 @@ class Transducer(torch.nn.Module): text_lengths = text_lengths[:, 0] if len(speech_lengths.size()) > 1: speech_lengths = speech_lengths[:, 0] - + batch_size = speech.shape[0] # 1. Encoder encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) - if hasattr(self.encoder, 'overlap_chunk_cls') and self.encoder.overlap_chunk_cls is not None: - encoder_out, encoder_out_lens = self.encoder.overlap_chunk_cls.remove_chunk(encoder_out, encoder_out_lens, - chunk_outs=None) + if ( + hasattr(self.encoder, "overlap_chunk_cls") + and self.encoder.overlap_chunk_cls is not None + ): + encoder_out, encoder_out_lens = self.encoder.overlap_chunk_cls.remove_chunk( + encoder_out, encoder_out_lens, chunk_outs=None + ) # 2. Transducer-related I/O preparation decoder_in, target, t_len, u_len = get_transducer_task_io( text, encoder_out_lens, ignore_id=self.ignore_id, ) - + # 3. Decoder self.decoder.set_device(encoder_out.device) decoder_out = self.decoder(decoder_in, u_len) - + # 4. Joint Network - joint_out = self.joint_network( - encoder_out.unsqueeze(2), decoder_out.unsqueeze(1) - ) - + joint_out = self.joint_network(encoder_out.unsqueeze(2), decoder_out.unsqueeze(1)) + # 5. Losses loss_trans, cer_trans, wer_trans = self._calc_transducer_loss( encoder_out, @@ -192,9 +194,9 @@ class Transducer(torch.nn.Module): t_len, u_len, ) - + loss_ctc, loss_lm = 0.0, 0.0 - + if self.use_auxiliary_ctc: loss_ctc = self._calc_ctc_loss( encoder_out, @@ -202,16 +204,16 @@ class Transducer(torch.nn.Module): t_len, u_len, ) - + if self.use_auxiliary_lm_loss: loss_lm = self._calc_lm_loss(decoder_out, target) - + loss = ( self.transducer_weight * loss_trans + self.auxiliary_ctc_weight * loss_ctc + self.auxiliary_lm_loss_weight * loss_lm ) - + stats = dict( loss=loss.detach(), loss_transducer=loss_trans.detach(), @@ -220,14 +222,17 @@ class Transducer(torch.nn.Module): cer_transducer=cer_trans, wer_transducer=wer_trans, ) - + # force_gatherable: to-device and to-tensor if scalar for DataParallel loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) - + return loss, stats, weight def encode( - self, speech: torch.Tensor, speech_lengths: torch.Tensor, **kwargs, + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + **kwargs, ) -> Tuple[torch.Tensor, torch.Tensor]: """Frontend + Encoder. Note that this method is used by asr_inference.py Args: @@ -240,11 +245,11 @@ class Transducer(torch.nn.Module): # Data augmentation if self.specaug is not None and self.training: speech, speech_lengths = self.specaug(speech, speech_lengths) - + # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN if self.normalize is not None: speech, speech_lengths = self.normalize(speech, speech_lengths) - + # Forward encoder # feats: (Batch, Length, Dim) # -> encoder_out: (Batch, Length2, Dim2) @@ -253,12 +258,12 @@ class Transducer(torch.nn.Module): if isinstance(encoder_out, tuple): intermediate_outs = encoder_out[1] encoder_out = encoder_out[0] - + if intermediate_outs is not None: return (encoder_out, intermediate_outs), encoder_out_lens - + return encoder_out, encoder_out_lens - + def _calc_transducer_loss( self, encoder_out: torch.Tensor, @@ -285,17 +290,17 @@ class Transducer(torch.nn.Module): if self.criterion_transducer is None: try: from warp_rnnt import rnnt_loss as RNNTLoss + self.criterion_transducer = RNNTLoss - + except ImportError: logging.error( - "warp-rnnt was not installed." - "Please consult the installation documentation." + "warp-rnnt was not installed." "Please consult the installation documentation." ) exit(1) - + log_probs = torch.log_softmax(joint_out, dim=-1) - + loss_transducer = self.criterion_transducer( log_probs, target, @@ -306,11 +311,11 @@ class Transducer(torch.nn.Module): fastemit_lambda=self.fastemit_lambda, gather=True, ) - + if not self.training and (self.report_cer or self.report_wer): if self.error_calculator is None: from funasr.metrics import ErrorCalculatorTransducer as ErrorCalculator - + self.error_calculator = ErrorCalculator( self.decoder, self.joint_network, @@ -320,13 +325,13 @@ class Transducer(torch.nn.Module): report_cer=self.report_cer, report_wer=self.report_wer, ) - + cer_transducer, wer_transducer = self.error_calculator(encoder_out, target, t_len) - + return loss_transducer, cer_transducer, wer_transducer - + return loss_transducer, None, None - + def _calc_ctc_loss( self, encoder_out: torch.Tensor, @@ -346,14 +351,12 @@ class Transducer(torch.nn.Module): loss_ctc: CTC loss value. """ - ctc_in = self.ctc_lin( - torch.nn.functional.dropout(encoder_out, p=self.ctc_dropout_rate) - ) + ctc_in = self.ctc_lin(torch.nn.functional.dropout(encoder_out, p=self.ctc_dropout_rate)) ctc_in = torch.log_softmax(ctc_in.transpose(0, 1), dim=-1) - + target_mask = target != 0 ctc_target = target[target_mask].cpu() - + with torch.backends.cudnn.flags(deterministic=True): loss_ctc = torch.nn.functional.ctc_loss( ctc_in, @@ -364,9 +367,9 @@ class Transducer(torch.nn.Module): reduction="sum", ) loss_ctc /= target.size(0) - + return loss_ctc - + def _calc_lm_loss( self, decoder_out: torch.Tensor, @@ -384,40 +387,37 @@ class Transducer(torch.nn.Module): """ lm_loss_in = self.lm_lin(decoder_out[:, :-1, :]).view(-1, self.vocab_size) lm_target = target.view(-1).type(torch.int64) - + with torch.no_grad(): true_dist = lm_loss_in.clone() true_dist.fill_(self.lm_loss_smoothing / (self.vocab_size - 1)) - + # Ignore blank ID (0) ignore = lm_target == 0 lm_target = lm_target.masked_fill(ignore, 0) - + true_dist.scatter_(1, lm_target.unsqueeze(1), (1 - self.lm_loss_smoothing)) - + loss_lm = torch.nn.functional.kl_div( torch.log_softmax(lm_loss_in, dim=1), true_dist, reduction="none", ) - loss_lm = loss_lm.masked_fill(ignore.unsqueeze(1), 0).sum() / decoder_out.size( - 0 - ) - + loss_lm = loss_lm.masked_fill(ignore.unsqueeze(1), 0).sum() / decoder_out.size(0) + return loss_lm - - def init_beam_search(self, - **kwargs, - ): - + + def init_beam_search( + self, + **kwargs, + ): + # 1. Build ASR model scorers = {} - + if self.ctc != None: ctc = CTCPrefixScorer(ctc=self.ctc, eos=self.eos) - scorers.update( - ctc=ctc - ) + scorers.update(ctc=ctc) token_list = kwargs.get("token_list") scorers.update( length_bonus=LengthBonus(len(token_list)), @@ -427,7 +427,7 @@ class Transducer(torch.nn.Module): # ngram is not supported now ngram = None scorers["ngram"] = ngram - + beam_search = BeamSearchTransducer( self.decoder, self.joint_network, @@ -439,37 +439,46 @@ class Transducer(torch.nn.Module): # if isinstance(scorer, torch.nn.Module): # scorer.to(device=kwargs.get("device", "cpu"), dtype=getattr(torch, kwargs.get("dtype", "float32"))).eval() self.beam_search = beam_search - - def inference(self, - data_in: list, - data_lengths: list=None, - key: list=None, - tokenizer=None, - **kwargs, - ): - + + def inference( + self, + data_in: list, + data_lengths: list = None, + key: list = None, + tokenizer=None, + **kwargs, + ): + if kwargs.get("batch_size", 1) > 1: raise NotImplementedError("batch decoding is not implemented") - + # init beamsearch is_use_ctc = kwargs.get("decoding_ctc_weight", 0.0) > 0.00001 and self.ctc != None - is_use_lm = kwargs.get("lm_weight", 0.0) > 0.00001 and kwargs.get("lm_file", None) is not None + is_use_lm = ( + kwargs.get("lm_weight", 0.0) > 0.00001 and kwargs.get("lm_file", None) is not None + ) # if self.beam_search is None and (is_use_lm or is_use_ctc): logging.info("enable beam_search") self.init_beam_search(**kwargs) self.nbest = kwargs.get("nbest", 1) - + meta_data = {} # extract fbank feats time1 = time.perf_counter() - audio_sample_list = load_audio_text_image_video(data_in, fs=self.frontend.fs, audio_fs=kwargs.get("fs", 16000)) + audio_sample_list = load_audio_text_image_video( + data_in, fs=self.frontend.fs, audio_fs=kwargs.get("fs", 16000) + ) time2 = time.perf_counter() meta_data["load_data"] = f"{time2 - time1:0.3f}" - speech, speech_lengths = extract_fbank(audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=self.frontend) + speech, speech_lengths = extract_fbank( + audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=self.frontend + ) time3 = time.perf_counter() meta_data["extract_feat"] = f"{time3 - time2:0.3f}" - meta_data["batch_data_time"] = speech_lengths.sum().item() * self.frontend.frame_shift * self.frontend.lfr_n / 1000 - + meta_data["batch_data_time"] = ( + speech_lengths.sum().item() * self.frontend.frame_shift * self.frontend.lfr_n / 1000 + ) + speech = speech.to(device=kwargs["device"]) speech_lengths = speech_lengths.to(device=kwargs["device"]) @@ -477,7 +486,7 @@ class Transducer(torch.nn.Module): encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) if isinstance(encoder_out, tuple): encoder_out = encoder_out[0] - + # c. Passed the encoder result and the beam search nbest_hyps = self.beam_search(encoder_out[0], is_final=True) nbest_hyps = nbest_hyps[: self.nbest] @@ -495,25 +504,33 @@ class Transducer(torch.nn.Module): # remove sos/eos and get results last_pos = -1 if isinstance(hyp.yseq, list): - token_int = hyp.yseq#[1:last_pos] + token_int = hyp.yseq # [1:last_pos] else: - token_int = hyp.yseq#[1:last_pos].tolist() - + token_int = hyp.yseq # [1:last_pos].tolist() + # remove blank symbol id, which is assumed to be 0 - token_int = list(filter(lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int)) - + token_int = list( + filter( + lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int + ) + ) + # Change integer-ids to tokens token = tokenizer.ids2tokens(token_int) text = tokenizer.tokens2text(token) - + text_postprocessed, _ = postprocess_utils.sentence_postprocess(token) - result_i = {"key": key[i], "token": token, "text": text, "text_postprocessed": text_postprocessed} + result_i = { + "key": key[i], + "token": token, + "text": text, + "text_postprocessed": text_postprocessed, + } results.append(result_i) - + if ibest_writer is not None: ibest_writer["token"][key[i]] = " ".join(token) ibest_writer["text"][key[i]] = text ibest_writer["text_postprocessed"][key[i]] = text_postprocessed - - return results, meta_data + return results, meta_data diff --git a/funasr/models/transducer/rnn_decoder.py b/funasr/models/transducer/rnn_decoder.py index b999d9c2d..4a0ffa3c3 100644 --- a/funasr/models/transducer/rnn_decoder.py +++ b/funasr/models/transducer/rnn_decoder.py @@ -78,11 +78,10 @@ def build_attention_list( ) att_list.append(att) else: - raise ValueError( - "Number of encoders needs to be more than one. {}".format(num_encs) - ) + raise ValueError("Number of encoders needs to be more than one. {}".format(num_encs)) return att_list + @tables.register("decoder_classes", "rnn_decoder") class RNNDecoder(nn.Module): def __init__( @@ -125,16 +124,20 @@ class RNNDecoder(nn.Module): self.decoder = torch.nn.ModuleList() self.dropout_dec = torch.nn.ModuleList() self.decoder += [ - torch.nn.LSTMCell(hidden_size + eprojs, hidden_size) - if self.dtype == "lstm" - else torch.nn.GRUCell(hidden_size + eprojs, hidden_size) + ( + torch.nn.LSTMCell(hidden_size + eprojs, hidden_size) + if self.dtype == "lstm" + else torch.nn.GRUCell(hidden_size + eprojs, hidden_size) + ) ] self.dropout_dec += [torch.nn.Dropout(p=dropout)] for _ in range(1, self.dlayers): self.decoder += [ - torch.nn.LSTMCell(hidden_size, hidden_size) - if self.dtype == "lstm" - else torch.nn.GRUCell(hidden_size, hidden_size) + ( + torch.nn.LSTMCell(hidden_size, hidden_size) + if self.dtype == "lstm" + else torch.nn.GRUCell(hidden_size, hidden_size) + ) ] self.dropout_dec += [torch.nn.Dropout(p=dropout)] # NOTE: dropout is applied only for the vertical connections @@ -145,9 +148,7 @@ class RNNDecoder(nn.Module): else: self.output = torch.nn.Linear(hidden_size, vocab_size) - self.att_list = build_attention_list( - eprojs=eprojs, dunits=hidden_size, **att_conf - ) + self.att_list = build_attention_list(eprojs=eprojs, dunits=hidden_size, **att_conf) def zero_state(self, hs_pad): return hs_pad.new_zeros(hs_pad.size(0), self.dunits) @@ -163,9 +164,7 @@ class RNNDecoder(nn.Module): else: z_list[0] = self.decoder[0](ey, z_prev[0]) for i in range(1, self.dlayers): - z_list[i] = self.decoder[i]( - self.dropout_dec[i - 1](z_list[i - 1]), z_prev[i] - ) + z_list[i] = self.decoder[i](self.dropout_dec[i - 1](z_list[i - 1]), z_prev[i]) return z_list, c_list def forward(self, hs_pad, hlens, ys_in_pad, ys_in_lens, strm_idx=0): @@ -315,13 +314,9 @@ class RNNDecoder(nn.Module): state["a_prev"][self.num_encs], ) ey = torch.cat((ey, att_c), dim=1) # utt(1) x (zdim + hdim) - z_list, c_list = self.rnn_forward( - ey, z_list, c_list, state["z_prev"], state["c_prev"] - ) + z_list, c_list = self.rnn_forward(ey, z_list, c_list, state["z_prev"], state["c_prev"]) if self.context_residual: - logits = self.output( - torch.cat((self.dropout_dec[-1](z_list[-1]), att_c), dim=-1) - ) + logits = self.output(torch.cat((self.dropout_dec[-1](z_list[-1]), att_c), dim=-1)) else: logits = self.output(self.dropout_dec[-1](z_list[-1])) logp = F.log_softmax(logits, dim=1).squeeze(0) diff --git a/funasr/models/transducer/rnnt_decoder.py b/funasr/models/transducer/rnnt_decoder.py index 26ca1f2e2..bc1787ac8 100644 --- a/funasr/models/transducer/rnnt_decoder.py +++ b/funasr/models/transducer/rnnt_decoder.py @@ -42,7 +42,6 @@ class RNNTDecoder(torch.nn.Module): """Construct a RNNDecoder object.""" super().__init__() - if rnn_type not in ("lstm", "gru"): raise ValueError(f"Not supported: rnn_type={rnn_type}") @@ -51,9 +50,7 @@ class RNNTDecoder(torch.nn.Module): rnn_class = torch.nn.LSTM if rnn_type == "lstm" else torch.nn.GRU - self.rnn = torch.nn.ModuleList( - [rnn_class(embed_size, hidden_size, 1, batch_first=True)] - ) + self.rnn = torch.nn.ModuleList([rnn_class(embed_size, hidden_size, 1, batch_first=True)]) for _ in range(1, num_layers): self.rnn += [rnn_class(hidden_size, hidden_size, 1, batch_first=True)] @@ -77,9 +74,9 @@ class RNNTDecoder(torch.nn.Module): time_mask_width_range=3, num_time_mask=4, apply_freq_mask=False, - apply_time_warp=False + apply_time_warp=False, ) - + def forward( self, labels: torch.Tensor, @@ -128,13 +125,11 @@ class RNNTDecoder(torch.nn.Module): for layer in range(self.dlayers): if self.dtype == "lstm": - x, (h_next[layer : layer + 1], c_next[layer : layer + 1]) = self.rnn[ - layer - ](x, hx=(h_prev[layer : layer + 1], c_prev[layer : layer + 1])) - else: - x, h_next[layer : layer + 1] = self.rnn[layer]( - x, hx=h_prev[layer : layer + 1] + x, (h_next[layer : layer + 1], c_next[layer : layer + 1]) = self.rnn[layer]( + x, hx=(h_prev[layer : layer + 1], c_prev[layer : layer + 1]) ) + else: + x, h_next[layer : layer + 1] = self.rnn[layer](x, hx=h_prev[layer : layer + 1]) x = self.dropout_rnn[layer](x) @@ -203,9 +198,7 @@ class RNNTDecoder(torch.nn.Module): """ self.device = device - def init_state( - self, batch_size: int - ) -> Tuple[torch.Tensor, Optional[torch.tensor]]: + def init_state(self, batch_size: int) -> Tuple[torch.Tensor, Optional[torch.tensor]]: """Initialize decoder states. Args: @@ -267,7 +260,5 @@ class RNNTDecoder(torch.nn.Module): """ return ( torch.cat([s[0] for s in new_states], dim=1), - torch.cat([s[1] for s in new_states], dim=1) - if self.dtype == "lstm" - else None, + torch.cat([s[1] for s in new_states], dim=1) if self.dtype == "lstm" else None, ) diff --git a/funasr/models/transformer/attention.py b/funasr/models/transformer/attention.py index 695023d47..e4add5cf8 100644 --- a/funasr/models/transformer/attention.py +++ b/funasr/models/transformer/attention.py @@ -17,6 +17,7 @@ import torch.nn.functional as F from funasr.models.transformer.utils.nets_utils import make_pad_mask import funasr.models.lora.layers as lora + class MultiHeadedAttention(nn.Module): """Multi-Head Attention layer. @@ -81,9 +82,7 @@ class MultiHeadedAttention(nn.Module): n_batch = value.size(0) if mask is not None: mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) - min_value = float( - numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min - ) + min_value = float(numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min) scores = scores.masked_fill(mask, min_value) self.attn = torch.softmax(scores, dim=-1).masked_fill( mask, 0.0 @@ -129,17 +128,17 @@ class MultiHeadedAttentionExport(nn.Module): self.linear_out = model.linear_out self.attn = None self.all_head_size = self.h * self.d_k - + def forward(self, query, key, value, mask): q, k, v = self.forward_qkv(query, key, value) scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) return self.forward_attention(v, scores, mask) - + def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor: new_x_shape = x.size()[:-1] + (self.h, self.d_k) x = x.view(new_x_shape) return x.permute(0, 2, 1, 3) - + def forward_qkv(self, query, key, value): q = self.linear_q(query) k = self.linear_k(key) @@ -148,13 +147,13 @@ class MultiHeadedAttentionExport(nn.Module): k = self.transpose_for_scores(k) v = self.transpose_for_scores(v) return q, k, v - + def forward_attention(self, value, scores, mask): scores = scores + mask - + self.attn = torch.softmax(scores, dim=-1) context_layer = torch.matmul(self.attn, value) # (batch, head, time1, d_k) - + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) context_layer = context_layer.view(new_context_layer_shape) @@ -167,51 +166,49 @@ class RelPosMultiHeadedAttentionExport(MultiHeadedAttentionExport): self.linear_pos = model.linear_pos self.pos_bias_u = model.pos_bias_u self.pos_bias_v = model.pos_bias_v - + def forward(self, query, key, value, pos_emb, mask): q, k, v = self.forward_qkv(query, key, value) q = q.transpose(1, 2) # (batch, time1, head, d_k) - + p = self.transpose_for_scores(self.linear_pos(pos_emb)) # (batch, head, time1, d_k) - + # (batch, head, time1, d_k) q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2) # (batch, head, time1, d_k) q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2) - + # compute attention score # first compute matrix a and matrix c # as described in https://arxiv.org/abs/1901.02860 Section 3.3 # (batch, head, time1, time2) matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1)) - + # compute matrix b and matrix d # (batch, head, time1, time1) matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) matrix_bd = self.rel_shift(matrix_bd) - - scores = (matrix_ac + matrix_bd) / math.sqrt( - self.d_k - ) # (batch, head, time1, time2) - + + scores = (matrix_ac + matrix_bd) / math.sqrt(self.d_k) # (batch, head, time1, time2) + return self.forward_attention(v, scores, mask) - + def rel_shift(self, x): zero_pad = torch.zeros((*x.size()[:3], 1), device=x.device, dtype=x.dtype) x_padded = torch.cat([zero_pad, x], dim=-1) - + x_padded = x_padded.view(*x.size()[:2], x.size(3) + 1, x.size(2)) x = x_padded[:, :, 1:].view_as(x)[ :, :, :, : x.size(-1) // 2 + 1 - ] # only keep the positions from 0 to time2 + ] # only keep the positions from 0 to time2 return x - + def forward_attention(self, value, scores, mask): scores = scores + mask - + self.attn = torch.softmax(scores, dim=-1) context_layer = torch.matmul(self.attn, value) # (batch, head, time1, d_k) - + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) context_layer = context_layer.view(new_context_layer_shape) @@ -306,9 +303,7 @@ class LegacyRelPositionMultiHeadedAttention(MultiHeadedAttention): matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) matrix_bd = self.rel_shift(matrix_bd) - scores = (matrix_ac + matrix_bd) / math.sqrt( - self.d_k - ) # (batch, head, time1, time2) + scores = (matrix_ac + matrix_bd) / math.sqrt(self.d_k) # (batch, head, time1, time2) return self.forward_attention(v, scores, mask) @@ -358,7 +353,7 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention): x_padded = x_padded.view(*x.size()[:2], x.size(3) + 1, x.size(2)) x = x_padded[:, :, 1:].view_as(x)[ :, :, :, : x.size(-1) // 2 + 1 - ] # only keep the positions from 0 to time2 + ] # only keep the positions from 0 to time2 if self.zero_triu: ones = torch.ones((x.size(2), x.size(3)), device=x.device) @@ -405,9 +400,7 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention): matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) matrix_bd = self.rel_shift(matrix_bd) - scores = (matrix_ac + matrix_bd) / math.sqrt( - self.d_k - ) # (batch, head, time1, time2) + scores = (matrix_ac + matrix_bd) / math.sqrt(self.d_k) # (batch, head, time1, time2) return self.forward_attention(v, scores, mask) @@ -552,21 +545,9 @@ class RelPositionMultiHeadedAttentionChunk(torch.nn.Module): """ n_batch = query.size(0) - q = ( - self.linear_q(query) - .view(n_batch, -1, self.num_heads, self.d_k) - .transpose(1, 2) - ) - k = ( - self.linear_k(key) - .view(n_batch, -1, self.num_heads, self.d_k) - .transpose(1, 2) - ) - v = ( - self.linear_v(value) - .view(n_batch, -1, self.num_heads, self.d_k) - .transpose(1, 2) - ) + q = self.linear_q(query).view(n_batch, -1, self.num_heads, self.d_k).transpose(1, 2) + k = self.linear_k(key).view(n_batch, -1, self.num_heads, self.d_k).transpose(1, 2) + v = self.linear_v(value).view(n_batch, -1, self.num_heads, self.d_k).transpose(1, 2) return q, k, v @@ -597,9 +578,7 @@ class RelPositionMultiHeadedAttentionChunk(torch.nn.Module): attn_output = torch.matmul(attn_output, value) attn_output = self.linear_out( - attn_output.transpose(1, 2) - .contiguous() - .view(batch_size, -1, self.num_heads * self.d_k) + attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.d_k) ) return attn_output @@ -629,4 +608,3 @@ class RelPositionMultiHeadedAttentionChunk(torch.nn.Module): q, k, v = self.forward_qkv(query, key, value) scores = self.compute_att_score(q, k, pos_enc, left_context=left_context) return self.forward_attention(v, scores, mask, chunk_mask=chunk_mask) - diff --git a/funasr/models/transformer/decoder.py b/funasr/models/transformer/decoder.py index 1e88a2579..7d849f349 100644 --- a/funasr/models/transformer/decoder.py +++ b/funasr/models/transformer/decoder.py @@ -28,6 +28,7 @@ from funasr.models.transformer.scorers.scorer_interface import BatchScorerInterf from funasr.register import tables + class DecoderLayer(nn.Module): """Single decoder layer module. @@ -51,14 +52,14 @@ class DecoderLayer(nn.Module): """ def __init__( - self, - size, - self_attn, - src_attn, - feed_forward, - dropout_rate, - normalize_before=True, - concat_after=False, + self, + size, + self_attn, + src_attn, + feed_forward, + dropout_rate, + normalize_before=True, + concat_after=False, ): """Construct an DecoderLayer object.""" super(DecoderLayer, self).__init__() @@ -115,9 +116,7 @@ class DecoderLayer(nn.Module): tgt_q_mask = tgt_mask[:, -1:, :] if self.concat_after: - tgt_concat = torch.cat( - (tgt_q, self.self_attn(tgt_q, tgt, tgt, tgt_q_mask)), dim=-1 - ) + tgt_concat = torch.cat((tgt_q, self.self_attn(tgt_q, tgt, tgt, tgt_q_mask)), dim=-1) x = residual + self.concat_linear1(tgt_concat) else: x = residual + self.dropout(self.self_attn(tgt_q, tgt, tgt, tgt_q_mask)) @@ -128,9 +127,7 @@ class DecoderLayer(nn.Module): if self.normalize_before: x = self.norm2(x) if self.concat_after: - x_concat = torch.cat( - (x, self.src_attn(x, memory, memory, memory_mask)), dim=-1 - ) + x_concat = torch.cat((x, self.src_attn(x, memory, memory, memory_mask)), dim=-1) x = residual + self.concat_linear2(x_concat) else: x = residual + self.dropout(self.src_attn(x, memory, memory, memory_mask)) @@ -159,23 +156,23 @@ class DecoderLayerExport(nn.Module): self.norm1 = model.norm1 self.norm2 = model.norm2 self.norm3 = model.norm3 - + def forward(self, tgt, tgt_mask, memory, memory_mask, cache=None): residual = tgt tgt = self.norm1(tgt) tgt_q = tgt tgt_q_mask = tgt_mask x = residual + self.self_attn(tgt_q, tgt, tgt, tgt_q_mask) - + residual = x x = self.norm2(x) - + x = residual + self.src_attn(x, memory, memory, memory_mask) - + residual = x x = self.norm3(x) x = residual + self.feed_forward(x) - + return x, tgt_mask, memory, memory_mask @@ -202,15 +199,15 @@ class BaseTransformerDecoder(nn.Module, BatchScorerInterface): """ def __init__( - self, - vocab_size: int, - encoder_output_size: int, - dropout_rate: float = 0.1, - positional_dropout_rate: float = 0.1, - input_layer: str = "embed", - use_output_layer: bool = True, - pos_enc_class=PositionalEncoding, - normalize_before: bool = True, + self, + vocab_size: int, + encoder_output_size: int, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, ): super().__init__() attention_dim = encoder_output_size @@ -243,11 +240,11 @@ class BaseTransformerDecoder(nn.Module, BatchScorerInterface): self.decoders = None def forward( - self, - hs_pad: torch.Tensor, - hlens: torch.Tensor, - ys_in_pad: torch.Tensor, - ys_in_lens: torch.Tensor, + self, + hs_pad: torch.Tensor, + hlens: torch.Tensor, + ys_in_pad: torch.Tensor, + ys_in_lens: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: """Forward decoder. @@ -275,20 +272,14 @@ class BaseTransformerDecoder(nn.Module, BatchScorerInterface): tgt_mask = tgt_mask & m memory = hs_pad - memory_mask = (~make_pad_mask(hlens, maxlen=memory.size(1)))[:, None, :].to( - memory.device - ) + memory_mask = (~make_pad_mask(hlens, maxlen=memory.size(1)))[:, None, :].to(memory.device) # Padding for Longformer if memory_mask.shape[-1] != memory.shape[1]: padlen = memory.shape[1] - memory_mask.shape[-1] - memory_mask = torch.nn.functional.pad( - memory_mask, (0, padlen), "constant", False - ) + memory_mask = torch.nn.functional.pad(memory_mask, (0, padlen), "constant", False) x = self.embed(tgt) - x, tgt_mask, memory, memory_mask = self.decoders( - x, tgt_mask, memory, memory_mask - ) + x, tgt_mask, memory, memory_mask = self.decoders(x, tgt_mask, memory, memory_mask) if self.normalize_before: x = self.after_norm(x) if self.output_layer is not None: @@ -298,11 +289,11 @@ class BaseTransformerDecoder(nn.Module, BatchScorerInterface): return x, olens def forward_one_step( - self, - tgt: torch.Tensor, - tgt_mask: torch.Tensor, - memory: torch.Tensor, - cache: List[torch.Tensor] = None, + self, + tgt: torch.Tensor, + tgt_mask: torch.Tensor, + memory: torch.Tensor, + cache: List[torch.Tensor] = None, ) -> Tuple[torch.Tensor, List[torch.Tensor]]: """Forward one step. @@ -322,9 +313,7 @@ class BaseTransformerDecoder(nn.Module, BatchScorerInterface): cache = [None] * len(self.decoders) new_cache = [] for c, decoder in zip(cache, self.decoders): - x, tgt_mask, memory, memory_mask = decoder( - x, tgt_mask, memory, None, cache=c - ) + x, tgt_mask, memory, memory_mask = decoder(x, tgt_mask, memory, None, cache=c) new_cache.append(x) if self.normalize_before: @@ -339,13 +328,11 @@ class BaseTransformerDecoder(nn.Module, BatchScorerInterface): def score(self, ys, state, x): """Score.""" ys_mask = subsequent_mask(len(ys), device=x.device).unsqueeze(0) - logp, state = self.forward_one_step( - ys.unsqueeze(0), ys_mask, x.unsqueeze(0), cache=state - ) + logp, state = self.forward_one_step(ys.unsqueeze(0), ys_mask, x.unsqueeze(0), cache=state) return logp.squeeze(0), state def batch_score( - self, ys: torch.Tensor, states: List[Any], xs: torch.Tensor + self, ys: torch.Tensor, states: List[Any], xs: torch.Tensor ) -> Tuple[torch.Tensor, List[Any]]: """Score new token batch. @@ -369,8 +356,7 @@ class BaseTransformerDecoder(nn.Module, BatchScorerInterface): else: # transpose state of [batch, layer] into [layer, batch] batch_state = [ - torch.stack([states[b][i] for b in range(n_batch)]) - for i in range(n_layers) + torch.stack([states[b][i] for b in range(n_batch)]) for i in range(n_layers) ] # batch decoding @@ -381,24 +367,25 @@ class BaseTransformerDecoder(nn.Module, BatchScorerInterface): state_list = [[states[i][b] for i in range(n_layers)] for b in range(n_batch)] return logp, state_list + @tables.register("decoder_classes", "TransformerDecoder") class TransformerDecoder(BaseTransformerDecoder): def __init__( - self, - vocab_size: int, - encoder_output_size: int, - attention_heads: int = 4, - linear_units: int = 2048, - num_blocks: int = 6, - dropout_rate: float = 0.1, - positional_dropout_rate: float = 0.1, - self_attention_dropout_rate: float = 0.0, - src_attention_dropout_rate: float = 0.0, - input_layer: str = "embed", - use_output_layer: bool = True, - pos_enc_class=PositionalEncoding, - normalize_before: bool = True, - concat_after: bool = False, + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, ): super().__init__( vocab_size=vocab_size, @@ -416,12 +403,8 @@ class TransformerDecoder(BaseTransformerDecoder): num_blocks, lambda lnum: DecoderLayer( attention_dim, - MultiHeadedAttention( - attention_heads, attention_dim, self_attention_dropout_rate - ), - MultiHeadedAttention( - attention_heads, attention_dim, src_attention_dropout_rate - ), + MultiHeadedAttention(attention_heads, attention_dim, self_attention_dropout_rate), + MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate), PositionwiseFeedForward(attention_dim, linear_units, dropout_rate), dropout_rate, normalize_before, @@ -433,24 +416,24 @@ class TransformerDecoder(BaseTransformerDecoder): @tables.register("decoder_classes", "LightweightConvolutionTransformerDecoder") class LightweightConvolutionTransformerDecoder(BaseTransformerDecoder): def __init__( - self, - vocab_size: int, - encoder_output_size: int, - attention_heads: int = 4, - linear_units: int = 2048, - num_blocks: int = 6, - dropout_rate: float = 0.1, - positional_dropout_rate: float = 0.1, - self_attention_dropout_rate: float = 0.0, - src_attention_dropout_rate: float = 0.0, - input_layer: str = "embed", - use_output_layer: bool = True, - pos_enc_class=PositionalEncoding, - normalize_before: bool = True, - concat_after: bool = False, - conv_wshare: int = 4, - conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11), - conv_usebias: int = False, + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + conv_wshare: int = 4, + conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11), + conv_usebias: int = False, ): if len(conv_kernel_length) != num_blocks: raise ValueError( @@ -481,9 +464,7 @@ class LightweightConvolutionTransformerDecoder(BaseTransformerDecoder): use_kernel_mask=True, use_bias=conv_usebias, ), - MultiHeadedAttention( - attention_heads, attention_dim, src_attention_dropout_rate - ), + MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate), PositionwiseFeedForward(attention_dim, linear_units, dropout_rate), dropout_rate, normalize_before, @@ -491,27 +472,28 @@ class LightweightConvolutionTransformerDecoder(BaseTransformerDecoder): ), ) + @tables.register("decoder_classes", "LightweightConvolution2DTransformerDecoder") class LightweightConvolution2DTransformerDecoder(BaseTransformerDecoder): def __init__( - self, - vocab_size: int, - encoder_output_size: int, - attention_heads: int = 4, - linear_units: int = 2048, - num_blocks: int = 6, - dropout_rate: float = 0.1, - positional_dropout_rate: float = 0.1, - self_attention_dropout_rate: float = 0.0, - src_attention_dropout_rate: float = 0.0, - input_layer: str = "embed", - use_output_layer: bool = True, - pos_enc_class=PositionalEncoding, - normalize_before: bool = True, - concat_after: bool = False, - conv_wshare: int = 4, - conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11), - conv_usebias: int = False, + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + conv_wshare: int = 4, + conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11), + conv_usebias: int = False, ): if len(conv_kernel_length) != num_blocks: raise ValueError( @@ -542,9 +524,7 @@ class LightweightConvolution2DTransformerDecoder(BaseTransformerDecoder): use_kernel_mask=True, use_bias=conv_usebias, ), - MultiHeadedAttention( - attention_heads, attention_dim, src_attention_dropout_rate - ), + MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate), PositionwiseFeedForward(attention_dim, linear_units, dropout_rate), dropout_rate, normalize_before, @@ -556,24 +536,24 @@ class LightweightConvolution2DTransformerDecoder(BaseTransformerDecoder): @tables.register("decoder_classes", "DynamicConvolutionTransformerDecoder") class DynamicConvolutionTransformerDecoder(BaseTransformerDecoder): def __init__( - self, - vocab_size: int, - encoder_output_size: int, - attention_heads: int = 4, - linear_units: int = 2048, - num_blocks: int = 6, - dropout_rate: float = 0.1, - positional_dropout_rate: float = 0.1, - self_attention_dropout_rate: float = 0.0, - src_attention_dropout_rate: float = 0.0, - input_layer: str = "embed", - use_output_layer: bool = True, - pos_enc_class=PositionalEncoding, - normalize_before: bool = True, - concat_after: bool = False, - conv_wshare: int = 4, - conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11), - conv_usebias: int = False, + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + conv_wshare: int = 4, + conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11), + conv_usebias: int = False, ): if len(conv_kernel_length) != num_blocks: raise ValueError( @@ -604,9 +584,7 @@ class DynamicConvolutionTransformerDecoder(BaseTransformerDecoder): use_kernel_mask=True, use_bias=conv_usebias, ), - MultiHeadedAttention( - attention_heads, attention_dim, src_attention_dropout_rate - ), + MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate), PositionwiseFeedForward(attention_dim, linear_units, dropout_rate), dropout_rate, normalize_before, @@ -614,27 +592,28 @@ class DynamicConvolutionTransformerDecoder(BaseTransformerDecoder): ), ) + @tables.register("decoder_classes", "DynamicConvolution2DTransformerDecoder") class DynamicConvolution2DTransformerDecoder(BaseTransformerDecoder): def __init__( - self, - vocab_size: int, - encoder_output_size: int, - attention_heads: int = 4, - linear_units: int = 2048, - num_blocks: int = 6, - dropout_rate: float = 0.1, - positional_dropout_rate: float = 0.1, - self_attention_dropout_rate: float = 0.0, - src_attention_dropout_rate: float = 0.0, - input_layer: str = "embed", - use_output_layer: bool = True, - pos_enc_class=PositionalEncoding, - normalize_before: bool = True, - concat_after: bool = False, - conv_wshare: int = 4, - conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11), - conv_usebias: int = False, + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + conv_wshare: int = 4, + conv_kernel_length: Sequence[int] = (11, 11, 11, 11, 11, 11), + conv_usebias: int = False, ): if len(conv_kernel_length) != num_blocks: raise ValueError( @@ -665,9 +644,7 @@ class DynamicConvolution2DTransformerDecoder(BaseTransformerDecoder): use_kernel_mask=True, use_bias=conv_usebias, ), - MultiHeadedAttention( - attention_heads, attention_dim, src_attention_dropout_rate - ), + MultiHeadedAttention(attention_heads, attention_dim, src_attention_dropout_rate), PositionwiseFeedForward(attention_dim, linear_units, dropout_rate), dropout_rate, normalize_before, diff --git a/funasr/models/transformer/embedding.py b/funasr/models/transformer/embedding.py index 1995bbe79..7c86dab2e 100644 --- a/funasr/models/transformer/embedding.py +++ b/funasr/models/transformer/embedding.py @@ -11,6 +11,7 @@ import torch import torch.nn.functional as F from torch import einsum + def _pre_hook( state_dict, prefix, @@ -64,9 +65,7 @@ class PositionalEncoding(torch.nn.Module): return pe = torch.zeros(x.size(1), self.d_model) if self.reverse: - position = torch.arange( - x.size(1) - 1, -1, -1.0, dtype=torch.float32 - ).unsqueeze(1) + position = torch.arange(x.size(1) - 1, -1, -1.0, dtype=torch.float32).unsqueeze(1) else: position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1) div_term = torch.exp( @@ -170,9 +169,7 @@ class LearnableFourierPosEnc(torch.nn.Module): if self.gamma is None: self.gamma = self.d_model // 2 - assert ( - d_model % 2 == 0 - ), "d_model should be divisible by two in order to use this layer." + assert d_model % 2 == 0, "d_model should be divisible by two in order to use this layer." self.w_r = torch.nn.Parameter(torch.empty(1, d_model // 2)) self._reset() # init the weights @@ -185,9 +182,7 @@ class LearnableFourierPosEnc(torch.nn.Module): ) def _reset(self): - self.w_r.data = torch.normal( - 0, (1 / math.sqrt(self.gamma)), (1, self.d_model // 2) - ) + self.w_r.data = torch.normal(0, (1 / math.sqrt(self.gamma)), (1, self.d_model // 2)) def extend_pe(self, x): """Reset the positional encodings.""" @@ -384,45 +379,57 @@ class StreamPositionalEncoding(torch.nn.Module): x = x * self.xscale + self.pe[:, start_idx : start_idx + x.size(1)] return self.dropout(x) -class SinusoidalPositionEncoder(torch.nn.Module): - ''' - ''' +class SinusoidalPositionEncoder(torch.nn.Module): + """ """ + def __int__(self, d_model=80, dropout_rate=0.1): pass - def encode(self, positions: torch.Tensor = None, depth: int = None, dtype: torch.dtype = torch.float32): + def encode( + self, positions: torch.Tensor = None, depth: int = None, dtype: torch.dtype = torch.float32 + ): batch_size = positions.size(0) positions = positions.type(dtype) device = positions.device - log_timescale_increment = torch.log(torch.tensor([10000], dtype=dtype, device=device)) / (depth / 2 - 1) - inv_timescales = torch.exp(torch.arange(depth / 2, device=device).type(dtype) * (-log_timescale_increment)) + log_timescale_increment = torch.log(torch.tensor([10000], dtype=dtype, device=device)) / ( + depth / 2 - 1 + ) + inv_timescales = torch.exp( + torch.arange(depth / 2, device=device).type(dtype) * (-log_timescale_increment) + ) inv_timescales = torch.reshape(inv_timescales, [batch_size, -1]) - scaled_time = torch.reshape(positions, [1, -1, 1]) * torch.reshape(inv_timescales, [1, 1, -1]) + scaled_time = torch.reshape(positions, [1, -1, 1]) * torch.reshape( + inv_timescales, [1, 1, -1] + ) encoding = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=2) return encoding.type(dtype) def forward(self, x): batch_size, timesteps, input_dim = x.size() - positions = torch.arange(1, timesteps+1, device=x.device)[None, :] + positions = torch.arange(1, timesteps + 1, device=x.device)[None, :] position_encoding = self.encode(positions, input_dim, x.dtype).to(x.device) return x + position_encoding -class StreamSinusoidalPositionEncoder(torch.nn.Module): - ''' - ''' +class StreamSinusoidalPositionEncoder(torch.nn.Module): + """ """ + def __int__(self, d_model=80, dropout_rate=0.1): pass - def encode(self, positions: torch.Tensor = None, depth: int = None, dtype: torch.dtype = torch.float32): + def encode( + self, positions: torch.Tensor = None, depth: int = None, dtype: torch.dtype = torch.float32 + ): batch_size = positions.size(0) positions = positions.type(dtype) log_timescale_increment = torch.log(torch.tensor([10000], dtype=dtype)) / (depth / 2 - 1) inv_timescales = torch.exp(torch.arange(depth / 2).type(dtype) * (-log_timescale_increment)) inv_timescales = torch.reshape(inv_timescales, [batch_size, -1]) - scaled_time = torch.reshape(positions, [1, -1, 1]) * torch.reshape(inv_timescales, [1, 1, -1]) + scaled_time = torch.reshape(positions, [1, -1, 1]) * torch.reshape( + inv_timescales, [1, 1, -1] + ) encoding = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=2) return encoding.type(dtype) @@ -432,9 +439,10 @@ class StreamSinusoidalPositionEncoder(torch.nn.Module): if cache is not None: start_idx = cache["start_idx"] cache["start_idx"] += timesteps - positions = torch.arange(1, timesteps+start_idx+1)[None, :] + positions = torch.arange(1, timesteps + start_idx + 1)[None, :] position_encoding = self.encode(positions, input_dim, x.dtype).to(x.device) - return x + position_encoding[:, start_idx: start_idx + timesteps] + return x + position_encoding[:, start_idx : start_idx + timesteps] + class StreamingRelPositionalEncoding(torch.nn.Module): """Relative positional encoding. @@ -444,9 +452,7 @@ class StreamingRelPositionalEncoding(torch.nn.Module): dropout_rate: Dropout rate. """ - def __init__( - self, size: int, dropout_rate: float = 0.0, max_len: int = 5000 - ) -> None: + def __init__(self, size: int, dropout_rate: float = 0.0, max_len: int = 5000) -> None: """Construct a RelativePositionalEncoding object.""" super().__init__() @@ -477,8 +483,7 @@ class StreamingRelPositionalEncoding(torch.nn.Module): position = torch.arange(0, time1, dtype=torch.float32).unsqueeze(1) div_term = torch.exp( - torch.arange(0, self.size, 2, dtype=torch.float32) - * -(math.log(10000.0) / self.size) + torch.arange(0, self.size, 2, dtype=torch.float32) * -(math.log(10000.0) / self.size) ) pe_positive[:, 0::2] = torch.sin(position * div_term) @@ -489,9 +494,7 @@ class StreamingRelPositionalEncoding(torch.nn.Module): pe_negative[:, 1::2] = torch.cos(-1 * position * div_term) pe_negative = pe_negative[1:].unsqueeze(0) - self.pe = torch.cat([pe_positive, pe_negative], dim=1).to( - dtype=x.dtype, device=x.device - ) + self.pe = torch.cat([pe_positive, pe_negative], dim=1).to(dtype=x.dtype, device=x.device) def forward(self, x: torch.Tensor, left_context: int = 0) -> torch.Tensor: """Compute positional encoding. @@ -505,9 +508,7 @@ class StreamingRelPositionalEncoding(torch.nn.Module): time1 = x.size(1) + left_context - pos_enc = self.pe[ - :, self.pe.size(1) // 2 - time1 + 1 : self.pe.size(1) // 2 + x.size(1) - ] + pos_enc = self.pe[:, self.pe.size(1) // 2 - time1 + 1 : self.pe.size(1) // 2 + x.size(1)] pos_enc = self.dropout(pos_enc) return pos_enc @@ -516,14 +517,17 @@ class StreamingRelPositionalEncoding(torch.nn.Module): class ScaledSinuEmbedding(torch.nn.Module): def __init__(self, dim): super().__init__() - self.scale = torch.nn.Parameter(torch.ones(1,)) - inv_freq = 1. / (10000 ** (torch.arange(0, dim, 2).float() / dim)) - self.register_buffer('inv_freq', inv_freq) + self.scale = torch.nn.Parameter( + torch.ones( + 1, + ) + ) + inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim)) + self.register_buffer("inv_freq", inv_freq) def forward(self, x): n, device = x.shape[1], x.device - t = torch.arange(n, device = device).type_as(self.inv_freq) - sinu = einsum('i , j -> i j', t, self.inv_freq) - emb = torch.cat((sinu.sin(), sinu.cos()), dim = -1) + t = torch.arange(n, device=device).type_as(self.inv_freq) + sinu = einsum("i , j -> i j", t, self.inv_freq) + emb = torch.cat((sinu.sin(), sinu.cos()), dim=-1) return emb * self.scale - diff --git a/funasr/models/transformer/encoder.py b/funasr/models/transformer/encoder.py index 1f1486704..a6a85ae75 100644 --- a/funasr/models/transformer/encoder.py +++ b/funasr/models/transformer/encoder.py @@ -30,6 +30,7 @@ from funasr.models.transformer.utils.subsampling import check_short_utt from funasr.register import tables + class EncoderLayer(nn.Module): """Encoder layer module. @@ -53,14 +54,14 @@ class EncoderLayer(nn.Module): """ def __init__( - self, - size, - self_attn, - feed_forward, - dropout_rate, - normalize_before=True, - concat_after=False, - stochastic_depth_rate=0.0, + self, + size, + self_attn, + feed_forward, + dropout_rate, + normalize_before=True, + concat_after=False, + stochastic_depth_rate=0.0, ): """Construct an EncoderLayer object.""" super(EncoderLayer, self).__init__() @@ -118,9 +119,7 @@ class EncoderLayer(nn.Module): x_concat = torch.cat((x, self.self_attn(x_q, x, x, mask)), dim=-1) x = residual + stoch_layer_coeff * self.concat_linear(x_concat) else: - x = residual + stoch_layer_coeff * self.dropout( - self.self_attn(x_q, x, x, mask) - ) + x = residual + stoch_layer_coeff * self.dropout(self.self_attn(x_q, x, x, mask)) if not self.normalize_before: x = self.norm1(x) @@ -136,6 +135,7 @@ class EncoderLayer(nn.Module): return x, mask + @tables.register("encoder_classes", "TransformerEncoder") class TransformerEncoder(nn.Module): """Transformer encoder module. @@ -163,24 +163,24 @@ class TransformerEncoder(nn.Module): """ def __init__( - self, - input_size: int, - output_size: int = 256, - attention_heads: int = 4, - linear_units: int = 2048, - num_blocks: int = 6, - dropout_rate: float = 0.1, - positional_dropout_rate: float = 0.1, - attention_dropout_rate: float = 0.0, - input_layer: Optional[str] = "conv2d", - pos_enc_class=PositionalEncoding, - normalize_before: bool = True, - concat_after: bool = False, - positionwise_layer_type: str = "linear", - positionwise_conv_kernel_size: int = 1, - padding_idx: int = -1, - interctc_layer_idx: List[int] = [], - interctc_use_conditioning: bool = False, + self, + input_size: int, + output_size: int = 256, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + attention_dropout_rate: float = 0.0, + input_layer: Optional[str] = "conv2d", + pos_enc_class=PositionalEncoding, + normalize_before: bool = True, + concat_after: bool = False, + positionwise_layer_type: str = "linear", + positionwise_conv_kernel_size: int = 1, + padding_idx: int = -1, + interctc_layer_idx: List[int] = [], + interctc_use_conditioning: bool = False, ): super().__init__() self._output_size = output_size @@ -243,9 +243,7 @@ class TransformerEncoder(nn.Module): num_blocks, lambda lnum: EncoderLayer( output_size, - MultiHeadedAttention( - attention_heads, output_size, attention_dropout_rate - ), + MultiHeadedAttention(attention_heads, output_size, attention_dropout_rate), positionwise_layer(*positionwise_layer_args), dropout_rate, normalize_before, @@ -265,11 +263,11 @@ class TransformerEncoder(nn.Module): return self._output_size def forward( - self, - xs_pad: torch.Tensor, - ilens: torch.Tensor, - prev_states: torch.Tensor = None, - ctc: CTC = None, + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, + prev_states: torch.Tensor = None, + ctc: CTC = None, ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: """Embed positions in tensor. @@ -285,10 +283,10 @@ class TransformerEncoder(nn.Module): if self.embed is None: xs_pad = xs_pad elif ( - isinstance(self.embed, Conv2dSubsampling) - or isinstance(self.embed, Conv2dSubsampling2) - or isinstance(self.embed, Conv2dSubsampling6) - or isinstance(self.embed, Conv2dSubsampling8) + isinstance(self.embed, Conv2dSubsampling) + or isinstance(self.embed, Conv2dSubsampling2) + or isinstance(self.embed, Conv2dSubsampling6) + or isinstance(self.embed, Conv2dSubsampling8) ): short_status, limit_size = check_short_utt(self.embed, xs_pad.size(1)) if short_status: @@ -329,4 +327,3 @@ class TransformerEncoder(nn.Module): if len(intermediate_outs) > 0: return (xs_pad, intermediate_outs), olens, None return xs_pad, olens, None - diff --git a/funasr/models/transformer/layer_norm.py b/funasr/models/transformer/layer_norm.py index 86832306c..ce4f8877f 100644 --- a/funasr/models/transformer/layer_norm.py +++ b/funasr/models/transformer/layer_norm.py @@ -36,11 +36,7 @@ class LayerNorm(torch.nn.LayerNorm): """ if self.dim == -1: return super(LayerNorm, self).forward(x) - return ( - super(LayerNorm, self) - .forward(x.transpose(self.dim, -1)) - .transpose(self.dim, -1) - ) + return super(LayerNorm, self).forward(x.transpose(self.dim, -1)).transpose(self.dim, -1) class GlobalLayerNorm(nn.Module): @@ -97,10 +93,7 @@ class GlobalLayerNorm(nn.Module): mean = torch.mean(x, (1, 2), keepdim=True) var = torch.mean((x - mean) ** 2, (1, 2), keepdim=True) if self.elementwise_affine: - x = ( - self.weight * (x - mean) / torch.sqrt(var + self.eps) - + self.bias - ) + x = self.weight * (x - mean) / torch.sqrt(var + self.eps) + self.bias else: x = (x - mean) / torch.sqrt(var + self.eps) @@ -108,10 +101,7 @@ class GlobalLayerNorm(nn.Module): mean = torch.mean(x, (1, 2, 3), keepdim=True) var = torch.mean((x - mean) ** 2, (1, 2, 3), keepdim=True) if self.elementwise_affine: - x = ( - self.weight * (x - mean) / torch.sqrt(var + self.eps) - + self.bias - ) + x = self.weight * (x - mean) / torch.sqrt(var + self.eps) + self.bias else: x = (x - mean) / torch.sqrt(var + self.eps) return x @@ -165,13 +155,12 @@ class CumulativeLayerNorm(nn.LayerNorm): class ScaleNorm(nn.Module): - def __init__(self, dim, eps = 1e-5): + def __init__(self, dim, eps=1e-5): super().__init__() - self.scale = dim ** -0.5 + self.scale = dim**-0.5 self.eps = eps self.g = nn.Parameter(torch.ones(1)) def forward(self, x): - norm = torch.norm(x, dim = -1, keepdim = True) * self.scale - return x / norm.clamp(min = self.eps) * self.g - + norm = torch.norm(x, dim=-1, keepdim=True) * self.scale + return x / norm.clamp(min=self.eps) * self.g diff --git a/funasr/models/transformer/model.py b/funasr/models/transformer/model.py index c68ccd7ba..0d5ed2376 100644 --- a/funasr/models/transformer/model.py +++ b/funasr/models/transformer/model.py @@ -10,6 +10,7 @@ from funasr.losses.label_smoothing_loss import LabelSmoothingLoss from funasr.models.ctc.ctc import CTC from funasr.models.transformer.utils.add_sos_eos import add_sos_eos from funasr.metrics.compute_acc import th_accuracy + # from funasr.models.e2e_asr_common import ErrorCalculator from funasr.train_utils.device_funcs import force_gatherable from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank @@ -17,11 +18,11 @@ from funasr.utils import postprocess_utils from funasr.utils.datadir_writer import DatadirWriter from funasr.register import tables + @tables.register("model_classes", "Transformer") class Transformer(nn.Module): """CTC-attention hybrid Encoder-Decoder model""" - def __init__( self, specaug: str = None, @@ -74,14 +75,12 @@ class Transformer(nn.Module): **decoder_conf, ) if ctc_weight > 0.0: - + if ctc_conf is None: ctc_conf = {} - - ctc = CTC( - odim=vocab_size, encoder_output_size=encoder_output_size, **ctc_conf - ) - + + ctc = CTC(odim=vocab_size, encoder_output_size=encoder_output_size, **ctc_conf) + self.blank_id = blank_id self.sos = sos if sos is not None else vocab_size - 1 self.eos = eos if eos is not None else vocab_size - 1 @@ -105,7 +104,7 @@ class Transformer(nn.Module): self.decoder = None else: self.decoder = decoder - + self.criterion_att = LabelSmoothingLoss( size=vocab_size, padding_idx=ignore_id, @@ -123,14 +122,14 @@ class Transformer(nn.Module): self.ctc = None else: self.ctc = ctc - + self.share_embedding = share_embedding if self.share_embedding: self.decoder.embed = None - + self.length_normalized_loss = length_normalized_loss self.beam_search = None - + def forward( self, speech: torch.Tensor, @@ -152,30 +151,30 @@ class Transformer(nn.Module): text_lengths = text_lengths[:, 0] if len(speech_lengths.size()) > 1: speech_lengths = speech_lengths[:, 0] - + batch_size = speech.shape[0] - + # 1. Encoder encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) intermediate_outs = None if isinstance(encoder_out, tuple): intermediate_outs = encoder_out[1] encoder_out = encoder_out[0] - + loss_att, acc_att, cer_att, wer_att = None, None, None, None loss_ctc, cer_ctc = None, None stats = dict() - + # decoder: CTC branch if self.ctc_weight != 0.0: loss_ctc, cer_ctc = self._calc_ctc_loss( encoder_out, encoder_out_lens, text, text_lengths ) - + # Collect CTC branch stats stats["loss_ctc"] = loss_ctc.detach() if loss_ctc is not None else None stats["cer_ctc"] = cer_ctc - + # Intermediate CTC (optional) loss_interctc = 0.0 if self.interctc_weight != 0.0 and intermediate_outs is not None: @@ -186,25 +185,23 @@ class Transformer(nn.Module): intermediate_out, encoder_out_lens, text, text_lengths ) loss_interctc = loss_interctc + loss_ic - + # Collect Intermedaite CTC stats stats["loss_interctc_layer{}".format(layer_idx)] = ( loss_ic.detach() if loss_ic is not None else None ) stats["cer_interctc_layer{}".format(layer_idx)] = cer_ic - + loss_interctc = loss_interctc / len(intermediate_outs) - + # calculate whole encoder loss - loss_ctc = ( - 1 - self.interctc_weight - ) * loss_ctc + self.interctc_weight * loss_interctc - + loss_ctc = (1 - self.interctc_weight) * loss_ctc + self.interctc_weight * loss_interctc + # decoder: Attention decoder branch loss_att, acc_att, cer_att, wer_att = self._calc_att_loss( encoder_out, encoder_out_lens, text, text_lengths ) - + # 3. CTC-Att loss definition if self.ctc_weight == 0.0: loss = loss_att @@ -212,25 +209,27 @@ class Transformer(nn.Module): loss = loss_ctc else: loss = self.ctc_weight * loss_ctc + (1 - self.ctc_weight) * loss_att - + # Collect Attn branch stats stats["loss_att"] = loss_att.detach() if loss_att is not None else None stats["acc"] = acc_att stats["cer"] = cer_att stats["wer"] = wer_att - + # Collect total loss stats stats["loss"] = torch.clone(loss.detach()) - + # force_gatherable: to-device and to-tensor if scalar for DataParallel if self.length_normalized_loss: batch_size = int((text_lengths + 1).sum()) loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) return loss, stats, weight - def encode( - self, speech: torch.Tensor, speech_lengths: torch.Tensor, **kwargs, + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + **kwargs, ) -> Tuple[torch.Tensor, torch.Tensor]: """Frontend + Encoder. Note that this method is used by asr_inference.py Args: @@ -243,30 +242,28 @@ class Transformer(nn.Module): # Data augmentation if self.specaug is not None and self.training: speech, speech_lengths = self.specaug(speech, speech_lengths) - + # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN if self.normalize is not None: speech, speech_lengths = self.normalize(speech, speech_lengths) - + # Forward encoder # feats: (Batch, Length, Dim) # -> encoder_out: (Batch, Length2, Dim2) if self.encoder.interctc_use_conditioning: - encoder_out, encoder_out_lens, _ = self.encoder( - speech, speech_lengths, ctc=self.ctc - ) + encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths, ctc=self.ctc) else: encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths) intermediate_outs = None if isinstance(encoder_out, tuple): intermediate_outs = encoder_out[1] encoder_out = encoder_out[0] - + if intermediate_outs is not None: return (encoder_out, intermediate_outs), encoder_out_lens - + return encoder_out, encoder_out_lens - + def _calc_att_loss( self, encoder_out: torch.Tensor, @@ -276,12 +273,10 @@ class Transformer(nn.Module): ): ys_in_pad, ys_out_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id) ys_in_lens = ys_pad_lens + 1 - + # 1. Forward decoder - decoder_out, _ = self.decoder( - encoder_out, encoder_out_lens, ys_in_pad, ys_in_lens - ) - + decoder_out, _ = self.decoder(encoder_out, encoder_out_lens, ys_in_pad, ys_in_lens) + # 2. Compute attention loss loss_att = self.criterion_att(decoder_out, ys_out_pad) acc_att = th_accuracy( @@ -289,16 +284,16 @@ class Transformer(nn.Module): ys_out_pad, ignore_label=self.ignore_id, ) - + # Compute cer/wer using attention-decoder if self.training or self.error_calculator is None: cer_att, wer_att = None, None else: ys_hat = decoder_out.argmax(dim=-1) cer_att, wer_att = self.error_calculator(ys_hat.cpu(), ys_pad.cpu()) - + return loss_att, acc_att, cer_att, wer_att - + def _calc_ctc_loss( self, encoder_out: torch.Tensor, @@ -308,41 +303,39 @@ class Transformer(nn.Module): ): # Calc CTC loss loss_ctc = self.ctc(encoder_out, encoder_out_lens, ys_pad, ys_pad_lens) - + # Calc CER using CTC cer_ctc = None if not self.training and self.error_calculator is not None: ys_hat = self.ctc.argmax(encoder_out).data cer_ctc = self.error_calculator(ys_hat.cpu(), ys_pad.cpu(), is_ctc=True) return loss_ctc, cer_ctc - - def init_beam_search(self, - **kwargs, - ): + + def init_beam_search( + self, + **kwargs, + ): from funasr.models.transformer.search import BeamSearch from funasr.models.transformer.scorers.ctc import CTCPrefixScorer from funasr.models.transformer.scorers.length_bonus import LengthBonus - + # 1. Build ASR model scorers = {} - + if self.ctc != None: ctc = CTCPrefixScorer(ctc=self.ctc, eos=self.eos) - scorers.update( - ctc=ctc - ) + scorers.update(ctc=ctc) token_list = kwargs.get("token_list") scorers.update( decoder=self.decoder, length_bonus=LengthBonus(len(token_list)), ) - # 3. Build ngram model # ngram is not supported now ngram = None scorers["ngram"] = ngram - + weights = dict( decoder=1.0 - kwargs.get("decoding_ctc_weight", 0.5), ctc=kwargs.get("decoding_ctc_weight", 0.5), @@ -362,19 +355,20 @@ class Transformer(nn.Module): ) self.beam_search = beam_search - - def inference(self, - data_in, - data_lengths=None, - key: list=None, - tokenizer=None, - frontend=None, - **kwargs, - ): - + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): + if kwargs.get("batch_size", 1) > 1: raise NotImplementedError("batch decoding is not implemented") - + # init beamsearch if self.beam_search is None: logging.info("enable beam_search") @@ -382,7 +376,9 @@ class Transformer(nn.Module): self.nbest = kwargs.get("nbest", 1) meta_data = {} - if isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank": # fbank + if ( + isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank" + ): # fbank speech, speech_lengths = data_in, data_lengths if len(speech.shape) < 3: speech = speech[None, :, :] @@ -391,16 +387,23 @@ class Transformer(nn.Module): else: # extract fbank feats time1 = time.perf_counter() - audio_sample_list = load_audio_text_image_video(data_in, fs=frontend.fs, audio_fs=kwargs.get("fs", 16000), - data_type=kwargs.get("data_type", "sound"), - tokenizer=tokenizer) + audio_sample_list = load_audio_text_image_video( + data_in, + fs=frontend.fs, + audio_fs=kwargs.get("fs", 16000), + data_type=kwargs.get("data_type", "sound"), + tokenizer=tokenizer, + ) time2 = time.perf_counter() meta_data["load_data"] = f"{time2 - time1:0.3f}" - speech, speech_lengths = extract_fbank(audio_sample_list, data_type=kwargs.get("data_type", "sound"), - frontend=frontend) + speech, speech_lengths = extract_fbank( + audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend + ) time3 = time.perf_counter() meta_data["extract_feat"] = f"{time3 - time2:0.3f}" - meta_data["batch_data_time"] = speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 + meta_data["batch_data_time"] = ( + speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 + ) speech = speech.to(device=kwargs["device"]) speech_lengths = speech_lengths.to(device=kwargs["device"]) @@ -408,14 +411,15 @@ class Transformer(nn.Module): encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) if isinstance(encoder_out, tuple): encoder_out = encoder_out[0] - + # c. Passed the encoder result and the beam search nbest_hyps = self.beam_search( - x=encoder_out[0], maxlenratio=kwargs.get("maxlenratio", 0.0), minlenratio=kwargs.get("minlenratio", 0.0) + x=encoder_out[0], + maxlenratio=kwargs.get("maxlenratio", 0.0), + minlenratio=kwargs.get("minlenratio", 0.0), ) - - nbest_hyps = nbest_hyps[: self.nbest] + nbest_hyps = nbest_hyps[: self.nbest] results = [] b, n, d = encoder_out.size() @@ -427,28 +431,31 @@ class Transformer(nn.Module): if not hasattr(self, "writer"): self.writer = DatadirWriter(kwargs.get("output_dir")) ibest_writer = self.writer[f"{nbest_idx + 1}best_recog"] - + # remove sos/eos and get results last_pos = -1 if isinstance(hyp.yseq, list): token_int = hyp.yseq[1:last_pos] else: token_int = hyp.yseq[1:last_pos].tolist() - + # remove blank symbol id, which is assumed to be 0 - token_int = list(filter(lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int)) - + token_int = list( + filter( + lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int + ) + ) + # Change integer-ids to tokens token = tokenizer.ids2tokens(token_int) text = tokenizer.tokens2text(token) - + text_postprocessed, _ = postprocess_utils.sentence_postprocess(token) result_i = {"key": key[i], "token": token, "text": text_postprocessed} results.append(result_i) - + if ibest_writer is not None: ibest_writer["token"][key[i]] = " ".join(token) ibest_writer["text"][key[i]] = text_postprocessed - - return results, meta_data + return results, meta_data diff --git a/funasr/models/transformer/positionwise_feed_forward.py b/funasr/models/transformer/positionwise_feed_forward.py index 081ff5b7f..7cfa5f9d5 100644 --- a/funasr/models/transformer/positionwise_feed_forward.py +++ b/funasr/models/transformer/positionwise_feed_forward.py @@ -35,15 +35,14 @@ class PositionwiseFeedForward(torch.nn.Module): class PositionwiseFeedForwardDecoderSANMExport(torch.nn.Module): - def __init__(self, model): - super().__init__() - self.w_1 = model.w_1 - self.w_2 = model.w_2 - self.activation = model.activation - self.norm = model.norm - - def forward(self, x): - x = self.activation(self.w_1(x)) - x = self.w_2(self.norm(x)) - return x + def __init__(self, model): + super().__init__() + self.w_1 = model.w_1 + self.w_2 = model.w_2 + self.activation = model.activation + self.norm = model.norm + def forward(self, x): + x = self.activation(self.w_1(x)) + x = self.w_2(self.norm(x)) + return x diff --git a/funasr/models/transformer/scorers/ctc.py b/funasr/models/transformer/scorers/ctc.py index 6639a1312..73a14bdfa 100644 --- a/funasr/models/transformer/scorers/ctc.py +++ b/funasr/models/transformer/scorers/ctc.py @@ -80,9 +80,7 @@ class CTCPrefixScorer(BatchPartialScorerInterface): """ prev_score, state = state presub_score, new_st = self.impl(y.cpu(), ids.cpu(), state) - tscore = torch.as_tensor( - presub_score - prev_score, device=x.device, dtype=x.dtype - ) + tscore = torch.as_tensor(presub_score - prev_score, device=x.device, dtype=x.dtype) return tscore, (presub_score, new_st) def batch_init_state(self, x: torch.Tensor): diff --git a/funasr/models/transformer/scorers/ctc_prefix_score.py b/funasr/models/transformer/scorers/ctc_prefix_score.py index 0c67ecd09..ce0703e73 100644 --- a/funasr/models/transformer/scorers/ctc_prefix_score.py +++ b/funasr/models/transformer/scorers/ctc_prefix_score.py @@ -38,11 +38,7 @@ class CTCPrefixScoreTH(object): self.input_length = x.size(1) self.odim = x.size(2) self.dtype = x.dtype - self.device = ( - torch.device("cuda:%d" % x.get_device()) - if x.is_cuda - else torch.device("cpu") - ) + self.device = torch.device("cuda:%d" % x.get_device()) if x.is_cuda else torch.device("cpu") # Pad the rest of posteriors in the batch # TODO(takaaki-hori): need a better way without for-loops for i, l in enumerate(xlens): @@ -58,9 +54,7 @@ class CTCPrefixScoreTH(object): # Setup CTC windowing self.margin = margin if margin > 0: - self.frame_ids = torch.arange( - self.input_length, dtype=self.dtype, device=self.device - ) + self.frame_ids = torch.arange(self.input_length, dtype=self.dtype, device=self.device) # Base indices for index conversion self.idx_bh = None self.idx_b = torch.arange(self.batch, device=self.device) @@ -98,18 +92,12 @@ class CTCPrefixScoreTH(object): # select input dimensions for scoring if self.scoring_num > 0: - scoring_idmap = torch.full( - (n_bh, self.odim), -1, dtype=torch.long, device=self.device - ) + scoring_idmap = torch.full((n_bh, self.odim), -1, dtype=torch.long, device=self.device) snum = self.scoring_num if self.idx_bh is None or n_bh > len(self.idx_bh): self.idx_bh = torch.arange(n_bh, device=self.device).view(-1, 1) - scoring_idmap[self.idx_bh[:n_bh], scoring_ids] = torch.arange( - snum, device=self.device - ) - scoring_idx = ( - scoring_ids + self.idx_bo.repeat(1, n_hyps).view(-1, 1) - ).view(-1) + scoring_idmap[self.idx_bh[:n_bh], scoring_ids] = torch.arange(snum, device=self.device) + scoring_idx = (scoring_ids + self.idx_bo.repeat(1, n_hyps).view(-1, 1)).view(-1) x_ = torch.index_select( self.x.view(2, -1, self.batch * self.odim), 2, scoring_idx ).view(2, -1, n_bh, snum) @@ -156,9 +144,7 @@ class CTCPrefixScoreTH(object): # compute forward probabilities log(r_t^n(h)) and log(r_t^b(h)) for t in range(start, end): rp = r[t - 1] - rr = torch.stack([rp[0], log_phi[t - 1], rp[0], rp[1]]).view( - 2, 2, n_bh, snum - ) + rr = torch.stack([rp[0], log_phi[t - 1], rp[0], rp[1]]).view(2, 2, n_bh, snum) r[t] = torch.logsumexp(rr, 1) + x_[:, t] # compute log prefix probabilities log(psi) @@ -205,9 +191,7 @@ class CTCPrefixScoreTH(object): # convert ids to BHS space (S: scoring_num) if scoring_idmap is not None: snum = self.scoring_num - hyp_idx = (best_ids // self.odim + (self.idx_b * n_hyps).view(-1, 1)).view( - -1 - ) + hyp_idx = (best_ids // self.odim + (self.idx_b * n_hyps).view(-1, 1)).view(-1) label_ids = torch.fmod(best_ids, self.odim).view(-1) score_idx = scoring_idmap[hyp_idx, label_ids] score_idx[score_idx == -1] = 0 @@ -215,9 +199,7 @@ class CTCPrefixScoreTH(object): else: snum = self.odim # select forward probabilities - r_new = torch.index_select(r.view(-1, 2, n_bh * snum), 2, vidx).view( - -1, 2, n_bh - ) + r_new = torch.index_select(r.view(-1, 2, n_bh * snum), 2, vidx).view(-1, 2, n_bh) return r_new, s_new, f_min, f_max def extend_prob(self, x): @@ -322,9 +304,7 @@ class CTCPrefixScore(object): r[output_length - 1] = self.logzero # prepare forward probabilities for the last label - r_sum = self.xp.logaddexp( - r_prev[:, 0], r_prev[:, 1] - ) # log(r_t^n(g) + r_t^b(g)) + r_sum = self.xp.logaddexp(r_prev[:, 0], r_prev[:, 1]) # log(r_t^n(g) + r_t^b(g)) last = y[-1] if output_length > 0 and last in cs: log_phi = self.xp.ndarray((self.input_length, len(cs)), dtype=np.float32) @@ -339,9 +319,7 @@ class CTCPrefixScore(object): log_psi = r[start - 1, 0] for t in six.moves.range(start, self.input_length): r[t, 0] = self.xp.logaddexp(r[t - 1, 0], log_phi[t - 1]) + xs[t] - r[t, 1] = ( - self.xp.logaddexp(r[t - 1, 0], r[t - 1, 1]) + self.x[t, self.blank] - ) + r[t, 1] = self.xp.logaddexp(r[t - 1, 0], r[t - 1, 1]) + self.x[t, self.blank] log_psi = self.xp.logaddexp(log_psi, log_phi[t - 1] + xs[t]) # get P(...eos|X) that ends with the prefix itself diff --git a/funasr/models/transformer/scorers/length_bonus.py b/funasr/models/transformer/scorers/length_bonus.py index 8156f1442..3a7e9e073 100644 --- a/funasr/models/transformer/scorers/length_bonus.py +++ b/funasr/models/transformer/scorers/length_bonus.py @@ -1,4 +1,5 @@ """Length bonus module.""" + from typing import Any from typing import List from typing import Tuple @@ -54,8 +55,6 @@ class LengthBonus(BatchScorerInterface): """ return ( - torch.tensor([1.0], device=xs.device, dtype=xs.dtype).expand( - ys.shape[0], self.n - ), + torch.tensor([1.0], device=xs.device, dtype=xs.dtype).expand(ys.shape[0], self.n), None, ) diff --git a/funasr/models/transformer/scorers/scorer_interface.py b/funasr/models/transformer/scorers/scorer_interface.py index 946ec6be3..3acdf2f4d 100644 --- a/funasr/models/transformer/scorers/scorer_interface.py +++ b/funasr/models/transformer/scorers/scorer_interface.py @@ -51,9 +51,7 @@ class ScorerInterface: """ return None if state is None else state[i] - def score( - self, y: torch.Tensor, state: Any, x: torch.Tensor - ) -> Tuple[torch.Tensor, Any]: + def score(self, y: torch.Tensor, state: Any, x: torch.Tensor) -> Tuple[torch.Tensor, Any]: """Score new token (required). Args: diff --git a/funasr/models/transformer/search.py b/funasr/models/transformer/search.py index ab7ac7d78..c1d90b64e 100644 --- a/funasr/models/transformer/search.py +++ b/funasr/models/transformer/search.py @@ -13,6 +13,7 @@ from funasr.metrics.common import end_detect from funasr.models.transformer.scorers.scorer_interface import PartialScorerInterface from funasr.models.transformer.scorers.scorer_interface import ScorerInterface + class Hypothesis(NamedTuple): """Hypothesis data type.""" @@ -30,7 +31,6 @@ class Hypothesis(NamedTuple): )._asdict() - class BeamSearch(torch.nn.Module): """Beam search implementation.""" @@ -276,9 +276,7 @@ class BeamSearch(torch.nn.Module): new_states[k] = d.select_state(part_states[k], part_idx) return new_states - def search( - self, running_hyps: List[Hypothesis], x: torch.Tensor - ) -> List[Hypothesis]: + def search(self, running_hyps: List[Hypothesis], x: torch.Tensor) -> List[Hypothesis]: """Search new tokens for running hypotheses and encoded speech x. Args: @@ -318,9 +316,7 @@ class BeamSearch(torch.nn.Module): Hypothesis( score=weighted_scores[j], yseq=self.append_token(hyp.yseq, j), - scores=self.merge_scores( - hyp.scores, scores, j, part_scores, part_j - ), + scores=self.merge_scores(hyp.scores, scores, j, part_scores, part_j), states=self.merge_states(states, part_states, part_j), ) ) @@ -383,8 +379,7 @@ class BeamSearch(torch.nn.Module): # check the number of hypotheses reaching to eos if len(nbest_hyps) == 0: logging.warning( - "there is no N-best results, perform recognition " - "again with smaller minlenratio." + "there is no N-best results, perform recognition " "again with smaller minlenratio." ) return ( [] @@ -395,17 +390,13 @@ class BeamSearch(torch.nn.Module): # report the best result best = nbest_hyps[0] for k, v in best.scores.items(): - logging.info( - f"{v:6.2f} * {self.weights[k]:3} = {v * self.weights[k]:6.2f} for {k}" - ) + logging.info(f"{v:6.2f} * {self.weights[k]:3} = {v * self.weights[k]:6.2f} for {k}") logging.info(f"total log probability: {best.score:.2f}") logging.info(f"normalized log probability: {best.score / len(best.yseq):.2f}") logging.info(f"total number of ended hypotheses: {len(nbest_hyps)}") if self.token_list is not None: logging.info( - "best hypo: " - + "".join([self.token_list[x] for x in best.yseq[1:-1]]) - + "\n" + "best hypo: " + "".join([self.token_list[x] for x in best.yseq[1:-1]]) + "\n" ) return nbest_hyps @@ -433,15 +424,13 @@ class BeamSearch(torch.nn.Module): logging.debug(f"the number of running hypotheses: {len(running_hyps)}") if self.token_list is not None: logging.debug( - "best hypo: " - + "".join([self.token_list[x] for x in running_hyps[0].yseq[1:]]) + "best hypo: " + "".join([self.token_list[x] for x in running_hyps[0].yseq[1:]]) ) # add eos in the final loop to avoid that there are no ended hyps if i == maxlen - 1: logging.info("adding in the last position in the loop") running_hyps = [ - h._replace(yseq=self.append_token(h.yseq, self.eos)) - for h in running_hyps + h._replace(yseq=self.append_token(h.yseq, self.eos)) for h in running_hyps ] # add ended hypotheses to a final list, and removed them from current hypotheses diff --git a/funasr/models/transformer/utils/dynamic_conv.py b/funasr/models/transformer/utils/dynamic_conv.py index 8a2a0c1ea..f6850384b 100644 --- a/funasr/models/transformer/utils/dynamic_conv.py +++ b/funasr/models/transformer/utils/dynamic_conv.py @@ -95,9 +95,9 @@ class DynamicConvolution(nn.Module): weight_new = torch.zeros(B * H * T * (T + k - 1), dtype=weight.dtype) weight_new = weight_new.view(B, H, T, T + k - 1).fill_(float("-inf")) weight_new = weight_new.to(x.device) # B x H x T x T+k-1 - weight_new.as_strided( - (B, H, T, k), ((T + k - 1) * T * H, (T + k - 1) * T, T + k, 1) - ).copy_(weight) + weight_new.as_strided((B, H, T, k), ((T + k - 1) * T * H, (T + k - 1) * T, T + k, 1)).copy_( + weight + ) weight_new = weight_new.narrow(-1, int((k - 1) / 2), T) # B x H x T x T(k) if self.use_kernel_mask: kernel_mask = torch.tril(torch.ones(T, T, device=x.device)).unsqueeze(0) diff --git a/funasr/models/transformer/utils/dynamic_conv2d.py b/funasr/models/transformer/utils/dynamic_conv2d.py index f8a4dd6e9..29d4dd8b9 100644 --- a/funasr/models/transformer/utils/dynamic_conv2d.py +++ b/funasr/models/transformer/utils/dynamic_conv2d.py @@ -95,9 +95,7 @@ class DynamicConvolution2D(nn.Module): # convolution of frequency axis weight_f = self.linear_weight_f(x).view(B * T, 1, k) # B x T x k self.attn_f = weight_f.view(B, T, k).unsqueeze(1) - xf = F.conv1d( - x.view(1, B * T, C), weight_f, padding=self.padding_size, groups=B * T - ) + xf = F.conv1d(x.view(1, B * T, C), weight_f, padding=self.padding_size, groups=B * T) xf = xf.view(B, T, C) # get kernel of convolution @@ -107,9 +105,9 @@ class DynamicConvolution2D(nn.Module): weight_new = torch.zeros(B * H * T * (T + k - 1), dtype=weight.dtype) weight_new = weight_new.view(B, H, T, T + k - 1).fill_(float("-inf")) weight_new = weight_new.to(x.device) # B x H x T x T+k-1 - weight_new.as_strided( - (B, H, T, k), ((T + k - 1) * T * H, (T + k - 1) * T, T + k, 1) - ).copy_(weight) + weight_new.as_strided((B, H, T, k), ((T + k - 1) * T * H, (T + k - 1) * T, T + k, 1)).copy_( + weight + ) weight_new = weight_new.narrow(-1, int((k - 1) / 2), T) # B x H x T x T(k) if self.use_kernel_mask: kernel_mask = torch.tril(torch.ones(T, T, device=x.device)).unsqueeze(0) diff --git a/funasr/models/transformer/utils/lightconv.py b/funasr/models/transformer/utils/lightconv.py index b24940259..b6d131b5f 100644 --- a/funasr/models/transformer/utils/lightconv.py +++ b/funasr/models/transformer/utils/lightconv.py @@ -50,9 +50,7 @@ class LightweightConvolution(nn.Module): self.act = nn.GLU() # lightconv related - self.weight = nn.Parameter( - torch.Tensor(self.wshare, 1, kernel_size).uniform_(0, 1) - ) + self.weight = nn.Parameter(torch.Tensor(self.wshare, 1, kernel_size).uniform_(0, 1)) self.use_bias = use_bias if self.use_bias: self.bias = nn.Parameter(torch.Tensor(n_feat)) @@ -96,9 +94,7 @@ class LightweightConvolution(nn.Module): self.kernel_mask = self.kernel_mask.to(x.device) weight = weight.masked_fill(self.kernel_mask == 0.0, float("-inf")) weight = F.softmax(weight, dim=-1) - x = F.conv1d(x, weight, padding=self.padding_size, groups=self.wshare).view( - B, C, T - ) + x = F.conv1d(x, weight, padding=self.padding_size, groups=self.wshare).view(B, C, T) if self.use_bias: x = x + self.bias.view(1, -1, 1) x = x.transpose(1, 2) # B x T x C diff --git a/funasr/models/transformer/utils/lightconv2d.py b/funasr/models/transformer/utils/lightconv2d.py index 294d23244..60e5af53e 100644 --- a/funasr/models/transformer/utils/lightconv2d.py +++ b/funasr/models/transformer/utils/lightconv2d.py @@ -50,9 +50,7 @@ class LightweightConvolution2D(nn.Module): self.act = nn.GLU() # lightconv related - self.weight = nn.Parameter( - torch.Tensor(self.wshare, 1, kernel_size).uniform_(0, 1) - ) + self.weight = nn.Parameter(torch.Tensor(self.wshare, 1, kernel_size).uniform_(0, 1)) self.weight_f = nn.Parameter(torch.Tensor(1, 1, kernel_size).uniform_(0, 1)) self.use_bias = use_bias if self.use_bias: @@ -93,9 +91,9 @@ class LightweightConvolution2D(nn.Module): # convolution along frequency axis weight_f = F.softmax(self.weight_f, dim=-1) weight_f = F.dropout(weight_f, self.dropout_rate, training=self.training) - weight_new = torch.zeros( - B * T, 1, self.kernel_size, device=x.device, dtype=x.dtype - ).copy_(weight_f) + weight_new = torch.zeros(B * T, 1, self.kernel_size, device=x.device, dtype=x.dtype).copy_( + weight_f + ) xf = F.conv1d( x.view(1, B * T, C), weight_new, padding=self.padding_size, groups=B * T ).view(B, T, C) @@ -107,9 +105,7 @@ class LightweightConvolution2D(nn.Module): self.kernel_mask = self.kernel_mask.to(x.device) weight = weight.masked_fill(self.kernel_mask == 0.0, float("-inf")) weight = F.softmax(weight, dim=-1) - x = F.conv1d(x, weight, padding=self.padding_size, groups=self.wshare).view( - B, C, T - ) + x = F.conv1d(x, weight, padding=self.padding_size, groups=self.wshare).view(B, C, T) if self.use_bias: x = x + self.bias.view(1, -1, 1) x = x.transpose(1, 2) # B x T x C diff --git a/funasr/models/transformer/utils/mask.py b/funasr/models/transformer/utils/mask.py index a8c168bc1..e859e0e06 100644 --- a/funasr/models/transformer/utils/mask.py +++ b/funasr/models/transformer/utils/mask.py @@ -34,6 +34,7 @@ def target_mask(ys_in_pad, ignore_id): m = subsequent_mask(ys_mask.size(-1), device=ys_mask.device).unsqueeze(0) return ys_mask.unsqueeze(-2) & m + def vad_mask(size, vad_pos, device="cpu", dtype=torch.bool): """Create mask for decoder self-attention. @@ -46,7 +47,6 @@ def vad_mask(size, vad_pos, device="cpu", dtype=torch.bool): ret = torch.ones(size, size, device=device, dtype=dtype) if vad_pos <= 0 or vad_pos >= size: return ret - sub_corner = torch.zeros( - vad_pos - 1, size - vad_pos, device=device, dtype=dtype) - ret[0:vad_pos - 1, vad_pos:] = sub_corner + sub_corner = torch.zeros(vad_pos - 1, size - vad_pos, device=device, dtype=dtype) + ret[0 : vad_pos - 1, vad_pos:] = sub_corner return ret diff --git a/funasr/models/transformer/utils/multi_layer_conv.py b/funasr/models/transformer/utils/multi_layer_conv.py index 9d269ab56..4893da167 100644 --- a/funasr/models/transformer/utils/multi_layer_conv.py +++ b/funasr/models/transformer/utils/multi_layer_conv.py @@ -96,7 +96,7 @@ class FsmnFeedForward(torch.nn.Module): kernel_size, stride=1, padding=(kernel_size - 1) // 2, - bias=False + bias=False, ) self.norm = torch.nn.LayerNorm(hidden_chans) self.dropout = torch.nn.Dropout(dropout_rate) diff --git a/funasr/models/transformer/utils/nets_utils.py b/funasr/models/transformer/utils/nets_utils.py index ce151a0a0..29d23ee59 100644 --- a/funasr/models/transformer/utils/nets_utils.py +++ b/funasr/models/transformer/utils/nets_utils.py @@ -25,9 +25,7 @@ def to_device(m, x): elif isinstance(m, torch.Tensor): device = m.device else: - raise TypeError( - "Expected torch.nn.Module or torch.tensor, " f"bot got: {type(m)}" - ) + raise TypeError("Expected torch.nn.Module or torch.tensor, " f"bot got: {type(m)}") return x.to(device) @@ -215,9 +213,7 @@ def make_pad_mask(lengths, xs=None, length_dim=-1, maxlen=None): if length_dim < 0: length_dim = xs.dim() + length_dim # ind = (:, None, ..., None, :, , None, ..., None) - ind = tuple( - slice(None) if i in (0, length_dim) else None for i in range(xs.dim()) - ) + ind = tuple(slice(None) if i in (0, length_dim) else None for i in range(xs.dim())) mask = mask[ind].expand_as(xs).to(xs.device) return mask @@ -342,8 +338,6 @@ def mask_by_length(xs, lengths, fill=0): return ret - - def to_torch_tensor(x): """Change to torch.Tensor or ComplexTensor from numpy.ndarray. @@ -434,9 +428,9 @@ def get_subsample(train_args, mode, arch): return subsample elif ( - (mode == "asr" and arch in ("rnn", "rnn-t")) - or (mode == "mt" and arch == "rnn") - or (mode == "st" and arch == "rnn") + (mode == "asr" and arch in ("rnn", "rnn-t")) + or (mode == "mt" and arch == "rnn") + or (mode == "st" and arch == "rnn") ): subsample = np.ones(train_args.elayers + 1, dtype=np.int32) if train_args.etype.endswith("p") and not train_args.etype.startswith("vgg"): @@ -452,14 +446,10 @@ def get_subsample(train_args, mode, arch): return subsample elif mode == "asr" and arch == "rnn_mix": - subsample = np.ones( - train_args.elayers_sd + train_args.elayers + 1, dtype=np.int32 - ) + subsample = np.ones(train_args.elayers_sd + train_args.elayers + 1, dtype=np.int32) if train_args.etype.endswith("p") and not train_args.etype.startswith("vgg"): ss = train_args.subsample.split("_") - for j in range( - min(train_args.elayers_sd + train_args.elayers + 1, len(ss)) - ): + for j in range(min(train_args.elayers_sd + train_args.elayers + 1, len(ss))): subsample[j] = int(ss[j]) else: logging.warning( @@ -473,9 +463,7 @@ def get_subsample(train_args, mode, arch): subsample_list = [] for idx in range(train_args.num_encs): subsample = np.ones(train_args.elayers[idx] + 1, dtype=np.int32) - if train_args.etype[idx].endswith("p") and not train_args.etype[ - idx - ].startswith("vgg"): + if train_args.etype[idx].endswith("p") and not train_args.etype[idx].startswith("vgg"): ss = train_args.subsample[idx].split("_") for j in range(min(train_args.elayers[idx] + 1, len(ss))): subsample[j] = int(ss[j]) @@ -493,9 +481,7 @@ def get_subsample(train_args, mode, arch): raise ValueError("Invalid options: mode={}, arch={}".format(mode, arch)) -def rename_state_dict( - old_prefix: str, new_prefix: str, state_dict: Dict[str, torch.Tensor] -): +def rename_state_dict(old_prefix: str, new_prefix: str, state_dict: Dict[str, torch.Tensor]): """Replace keys of old prefix with new prefix in state dict.""" # need this list not to break the dict iterator old_keys = [k for k in state_dict if k.startswith(old_prefix)] @@ -506,6 +492,7 @@ def rename_state_dict( new_k = k.replace(old_prefix, new_prefix) state_dict[new_k] = v + class Swish(torch.nn.Module): """Swish activation definition. @@ -540,6 +527,7 @@ class Swish(torch.nn.Module): """Forward computation.""" return self.swish(x) + def get_activation(act): """Return activation function.""" @@ -553,6 +541,7 @@ def get_activation(act): return activation_funcs[act]() + class TooShortUttError(Exception): """Raised when the utt is too short for subsampling. @@ -613,9 +602,7 @@ def sub_factor_to_params(sub_factor: int, input_size: int) -> Tuple[int, int, in elif sub_factor == 6: return 5, 3, (((input_size - 1) // 2 - 2) // 3) else: - raise ValueError( - "subsampling_factor parameter should be set to either 2, 4 or 6." - ) + raise ValueError("subsampling_factor parameter should be set to either 2, 4 or 6.") def make_chunk_mask( @@ -651,6 +638,7 @@ def make_chunk_mask( return ~mask + def make_source_mask(lengths: torch.Tensor) -> torch.Tensor: """Create source mask for given lengths. @@ -735,6 +723,7 @@ def get_transducer_task_io( return decoder_in, target, t_len, u_len + def pad_to_len(t: torch.Tensor, pad_len: int, dim: int): """Pad the tensor `t` at `dim` to the length `pad_len` with right padding zeros.""" if t.size(dim) == pad_len: @@ -742,6 +731,4 @@ def pad_to_len(t: torch.Tensor, pad_len: int, dim: int): else: pad_size = list(t.shape) pad_size[dim] = pad_len - t.size(dim) - return torch.cat( - [t, torch.zeros(*pad_size, dtype=t.dtype, device=t.device)], dim=dim - ) + return torch.cat([t, torch.zeros(*pad_size, dtype=t.dtype, device=t.device)], dim=dim) diff --git a/funasr/models/transformer/utils/subsampling.py b/funasr/models/transformer/utils/subsampling.py index 088675e5d..b18bd3132 100644 --- a/funasr/models/transformer/utils/subsampling.py +++ b/funasr/models/transformer/utils/subsampling.py @@ -15,6 +15,7 @@ from funasr.models.transformer.utils.nets_utils import sub_factor_to_params, pad from typing import Optional, Tuple, Union import math + class TooShortUttError(Exception): """Raised when the utt is too short for subsampling. @@ -103,6 +104,7 @@ class Conv2dSubsampling(torch.nn.Module): raise NotImplementedError("Support only `-1` (for `reset_parameters`).") return self.out[key] + class Conv2dSubsamplingPad(torch.nn.Module): """Convolutional 2D subsampling (to 1/4 length). @@ -326,6 +328,7 @@ class Conv2dSubsampling8(torch.nn.Module): return x, None return x, x_mask[:, :, :-2:2][:, :, :-2:2][:, :, :-2:2] + class Conv1dSubsampling(torch.nn.Module): """Convolutional 1D subsampling (to 1/2 length). @@ -337,10 +340,16 @@ class Conv1dSubsampling(torch.nn.Module): """ - def __init__(self, idim, odim, kernel_size, stride, pad, - tf2torch_tensor_name_prefix_torch: str = "stride_conv", - tf2torch_tensor_name_prefix_tf: str = "seq2seq/proj_encoder/downsampling", - ): + def __init__( + self, + idim, + odim, + kernel_size, + stride, + pad, + tf2torch_tensor_name_prefix_torch: str = "stride_conv", + tf2torch_tensor_name_prefix_tf: str = "seq2seq/proj_encoder/downsampling", + ): super(Conv1dSubsampling, self).__init__() self.conv = torch.nn.Conv1d(idim, odim, kernel_size, stride) self.pad_fn = torch.nn.ConstantPad1d(pad, 0.0) @@ -353,13 +362,11 @@ class Conv1dSubsampling(torch.nn.Module): return self.odim def forward(self, x, x_len): - """Subsample x. - - """ + """Subsample x.""" x = x.transpose(1, 2) # (b, d ,t) x = self.pad_fn(x) - #x = F.relu(self.conv(x)) - x = F.leaky_relu(self.conv(x), negative_slope=0.) + # x = F.relu(self.conv(x)) + x = F.leaky_relu(self.conv(x), negative_slope=0.0) x = x.transpose(1, 2) # (b, t ,d) if x_len is None: @@ -395,14 +402,38 @@ class StreamingConvInput(torch.nn.Module): conv_size1, conv_size2 = conv_size self.conv = torch.nn.Sequential( - torch.nn.Conv2d(1, conv_size1, conv_kernel_size, stride=1, padding=(conv_kernel_size-1)//2), + torch.nn.Conv2d( + 1, + conv_size1, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), torch.nn.ReLU(), - torch.nn.Conv2d(conv_size1, conv_size1, conv_kernel_size, stride=1, padding=(conv_kernel_size-1)//2), + torch.nn.Conv2d( + conv_size1, + conv_size1, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), torch.nn.ReLU(), torch.nn.MaxPool2d((1, 2)), - torch.nn.Conv2d(conv_size1, conv_size2, conv_kernel_size, stride=1, padding=(conv_kernel_size-1)//2), + torch.nn.Conv2d( + conv_size1, + conv_size2, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), torch.nn.ReLU(), - torch.nn.Conv2d(conv_size2, conv_size2, conv_kernel_size, stride=1, padding=(conv_kernel_size-1)//2), + torch.nn.Conv2d( + conv_size2, + conv_size2, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), torch.nn.ReLU(), torch.nn.MaxPool2d((1, 2)), ) @@ -421,14 +452,38 @@ class StreamingConvInput(torch.nn.Module): kernel_1 = int(subsampling_factor / 2) self.conv = torch.nn.Sequential( - torch.nn.Conv2d(1, conv_size1, conv_kernel_size, stride=1, padding=(conv_kernel_size-1)//2), + torch.nn.Conv2d( + 1, + conv_size1, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), torch.nn.ReLU(), - torch.nn.Conv2d(conv_size1, conv_size1, conv_kernel_size, stride=1, padding=(conv_kernel_size-1)//2), + torch.nn.Conv2d( + conv_size1, + conv_size1, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), torch.nn.ReLU(), torch.nn.MaxPool2d((kernel_1, 2)), - torch.nn.Conv2d(conv_size1, conv_size2, conv_kernel_size, stride=1, padding=(conv_kernel_size-1)//2), + torch.nn.Conv2d( + conv_size1, + conv_size2, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), torch.nn.ReLU(), - torch.nn.Conv2d(conv_size2, conv_size2, conv_kernel_size, stride=1, padding=(conv_kernel_size-1)//2), + torch.nn.Conv2d( + conv_size2, + conv_size2, + conv_kernel_size, + stride=1, + padding=(conv_kernel_size - 1) // 2, + ), torch.nn.ReLU(), torch.nn.MaxPool2d((2, 2)), ) @@ -444,9 +499,9 @@ class StreamingConvInput(torch.nn.Module): else: if subsampling_factor == 1: self.conv = torch.nn.Sequential( - torch.nn.Conv2d(1, conv_size, 3, [1,2], [1,0]), + torch.nn.Conv2d(1, conv_size, 3, [1, 2], [1, 0]), torch.nn.ReLU(), - torch.nn.Conv2d(conv_size, conv_size, conv_kernel_size, [1,2], [1,0]), + torch.nn.Conv2d(conv_size, conv_size, conv_kernel_size, [1, 2], [1, 0]), torch.nn.ReLU(), ) @@ -465,9 +520,11 @@ class StreamingConvInput(torch.nn.Module): ) self.conv = torch.nn.Sequential( - torch.nn.Conv2d(1, conv_size, 3, 2, [1,0]), + torch.nn.Conv2d(1, conv_size, 3, 2, [1, 0]), torch.nn.ReLU(), - torch.nn.Conv2d(conv_size, conv_size, kernel_2, stride_2, [(kernel_2-1)//2, 0]), + torch.nn.Conv2d( + conv_size, conv_size, kernel_2, stride_2, [(kernel_2 - 1) // 2, 0] + ), torch.nn.ReLU(), ) @@ -505,30 +562,32 @@ class StreamingConvInput(torch.nn.Module): olens = max(mask.eq(0).sum(1)) b, t, f = x.size() - x = x.unsqueeze(1) # (b. 1. t. f) + x = x.unsqueeze(1) # (b. 1. t. f) if chunk_size is not None: max_input_length = int( - chunk_size * self.subsampling_factor * (math.ceil(float(t) / (chunk_size * self.subsampling_factor) )) + chunk_size + * self.subsampling_factor + * (math.ceil(float(t) / (chunk_size * self.subsampling_factor))) ) x = map(lambda inputs: pad_to_len(inputs, max_input_length, 1), x) x = list(x) x = torch.stack(x, dim=0) - N_chunks = max_input_length // ( chunk_size * self.subsampling_factor) + N_chunks = max_input_length // (chunk_size * self.subsampling_factor) x = x.view(b * N_chunks, 1, chunk_size * self.subsampling_factor, f) x = self.conv(x) _, c, _, f = x.size() if chunk_size is not None: - x = x.transpose(1, 2).contiguous().view(b, -1, c * f)[:,:olens,:] + x = x.transpose(1, 2).contiguous().view(b, -1, c * f)[:, :olens, :] else: x = x.transpose(1, 2).contiguous().view(b, -1, c * f) if self.output is not None: x = self.output(x) - return x, mask[:,:olens][:,:x.size(1)] + return x, mask[:, :olens][:, : x.size(1)] def create_new_vgg_mask(self, mask: torch.Tensor) -> torch.Tensor: """Create a new mask for VGG output sequences. @@ -538,8 +597,8 @@ class StreamingConvInput(torch.nn.Module): mask: Mask of output sequences. (B, sub(T)) """ if self.subsampling_factor > 1: - vgg1_t_len = mask.size(1) - (mask.size(1) % (self.subsampling_factor // 2 )) - mask = mask[:, :vgg1_t_len][:, ::self.subsampling_factor // 2] + vgg1_t_len = mask.size(1) - (mask.size(1) % (self.subsampling_factor // 2)) + mask = mask[:, :vgg1_t_len][:, :: self.subsampling_factor // 2] vgg2_t_len = mask.size(1) - (mask.size(1) % 2) mask = mask[:, :vgg2_t_len][:, ::2] @@ -556,7 +615,7 @@ class StreamingConvInput(torch.nn.Module): mask: Mask of output sequences. (B, sub(T)) """ if self.subsampling_factor > 1: - return mask[:, ::2][:, ::self.stride_2] + return mask[:, ::2][:, :: self.stride_2] else: return mask diff --git a/funasr/models/transformer/utils/vgg2l.py b/funasr/models/transformer/utils/vgg2l.py index fb45d0b5a..7466baccf 100644 --- a/funasr/models/transformer/utils/vgg2l.py +++ b/funasr/models/transformer/utils/vgg2l.py @@ -39,9 +39,7 @@ class VGG2L(torch.nn.Module): else: self.output = torch.nn.Linear(128 * ((idim // 2) // 2), odim) - def forward( - self, feats: torch.Tensor, feats_mask: torch.Tensor - ) -> Union[ + def forward(self, feats: torch.Tensor, feats_mask: torch.Tensor) -> Union[ Tuple[torch.Tensor, torch.Tensor], Tuple[Tuple[torch.Tensor, torch.Tensor], torch.Tensor], ]: @@ -62,9 +60,7 @@ class VGG2L(torch.nn.Module): b, c, t, f = vgg_output.size() - vgg_output = self.output( - vgg_output.transpose(1, 2).contiguous().view(b, t, c * f) - ) + vgg_output = self.output(vgg_output.transpose(1, 2).contiguous().view(b, t, c * f)) if feats_mask is not None: vgg_mask = self.create_new_mask(feats_mask) diff --git a/funasr/models/uniasr/beam_search.py b/funasr/models/uniasr/beam_search.py index 839a1f4e0..f422129d0 100644 --- a/funasr/models/uniasr/beam_search.py +++ b/funasr/models/uniasr/beam_search.py @@ -33,7 +33,6 @@ class Hypothesis(NamedTuple): )._asdict() - class BeamSearchScama(torch.nn.Module): """Beam search implementation.""" @@ -151,7 +150,8 @@ class BeamSearchScama(torch.nn.Module): return torch.cat((xs, x)) def score_full( - self, hyp: Hypothesis, + self, + hyp: Hypothesis, x: torch.Tensor, x_mask: torch.Tensor = None, pre_acoustic_embeds: torch.Tensor = None, @@ -173,7 +173,9 @@ class BeamSearchScama(torch.nn.Module): scores = dict() states = dict() for k, d in self.full_scorers.items(): - scores[k], states[k] = d.score(hyp.yseq, hyp.states[k], x, x_mask=x_mask, pre_acoustic_embeds=pre_acoustic_embeds) + scores[k], states[k] = d.score( + hyp.yseq, hyp.states[k], x, x_mask=x_mask, pre_acoustic_embeds=pre_acoustic_embeds + ) return scores, states def score_partial( @@ -283,7 +285,8 @@ class BeamSearchScama(torch.nn.Module): return new_states def search( - self, running_hyps: List[Hypothesis], + self, + running_hyps: List[Hypothesis], x: torch.Tensor, x_mask: torch.Tensor = None, pre_acoustic_embeds: torch.Tensor = None, @@ -303,7 +306,9 @@ class BeamSearchScama(torch.nn.Module): for hyp in running_hyps: # scoring weighted_scores = torch.zeros(self.n_vocab, dtype=x.dtype, device=x.device) - scores, states = self.score_full(hyp, x, x_mask=x_mask, pre_acoustic_embeds=pre_acoustic_embeds) + scores, states = self.score_full( + hyp, x, x_mask=x_mask, pre_acoustic_embeds=pre_acoustic_embeds + ) for k in self.full_scorers: weighted_scores += self.weights[k] * scores[k] # partial scoring @@ -327,9 +332,7 @@ class BeamSearchScama(torch.nn.Module): Hypothesis( score=weighted_scores[j], yseq=self.append_token(hyp.yseq, j), - scores=self.merge_scores( - hyp.scores, scores, j, part_scores, part_j - ), + scores=self.merge_scores(hyp.scores, scores, j, part_scores, part_j), states=self.merge_states(states, part_states, part_j), ) ) @@ -341,7 +344,8 @@ class BeamSearchScama(torch.nn.Module): return best_hyps def forward( - self, x: torch.Tensor, + self, + x: torch.Tensor, scama_mask: torch.Tensor = None, pre_acoustic_embeds: torch.Tensor = None, maxlenratio: float = 0.0, @@ -386,20 +390,26 @@ class BeamSearchScama(torch.nn.Module): mask_enc = None if scama_mask is not None: token_num_predictor = scama_mask.size(1) - token_id_slice = min(i, token_num_predictor-1) - mask_enc = scama_mask[:, token_id_slice:token_id_slice+1, :] + token_id_slice = min(i, token_num_predictor - 1) + mask_enc = scama_mask[:, token_id_slice : token_id_slice + 1, :] # if mask_enc.size(1) == 0: # mask_enc = scama_mask[:, -2:-1, :] # # mask_enc = torch.zeros_like(mask_enc) pre_acoustic_embeds_cur = None if pre_acoustic_embeds is not None: b, t, d = pre_acoustic_embeds.size() - pad = torch.zeros((b, 1, d), dtype=pre_acoustic_embeds.dtype).to(device=pre_acoustic_embeds.device) + pad = torch.zeros((b, 1, d), dtype=pre_acoustic_embeds.dtype).to( + device=pre_acoustic_embeds.device + ) pre_acoustic_embeds = torch.cat((pre_acoustic_embeds, pad), dim=1) token_id_slice = min(i, t) - pre_acoustic_embeds_cur = pre_acoustic_embeds[:, token_id_slice:token_id_slice+1, :] + pre_acoustic_embeds_cur = pre_acoustic_embeds[ + :, token_id_slice : token_id_slice + 1, : + ] - best = self.search(running_hyps, x, x_mask=mask_enc, pre_acoustic_embeds=pre_acoustic_embeds_cur) + best = self.search( + running_hyps, x, x_mask=mask_enc, pre_acoustic_embeds=pre_acoustic_embeds_cur + ) # post process of one iteration running_hyps = self.post_process(i, maxlen, maxlenratio, best, ended_hyps) # end detection @@ -416,8 +426,7 @@ class BeamSearchScama(torch.nn.Module): # check the number of hypotheses reaching to eos if len(nbest_hyps) == 0: logging.warning( - "there is no N-best results, perform recognition " - "again with smaller minlenratio." + "there is no N-best results, perform recognition " "again with smaller minlenratio." ) return ( [] @@ -431,17 +440,13 @@ class BeamSearchScama(torch.nn.Module): logging.debug("nbest: y: {}, yseq: {}, score: {}".format(x.yseq, yseq, x.score)) best = nbest_hyps[0] for k, v in best.scores.items(): - logging.info( - f"{v:6.2f} * {self.weights[k]:3} = {v * self.weights[k]:6.2f} for {k}" - ) + logging.info(f"{v:6.2f} * {self.weights[k]:3} = {v * self.weights[k]:6.2f} for {k}") logging.info(f"total log probability: {best.score:.2f}") logging.info(f"normalized log probability: {best.score / len(best.yseq):.2f}") logging.info(f"total number of ended hypotheses: {len(nbest_hyps)}") if self.token_list is not None: logging.info( - "best hypo: " - + "".join([self.token_list[x] for x in best.yseq[1:-1]]) - + "\n" + "best hypo: " + "".join([self.token_list[x] for x in best.yseq[1:-1]]) + "\n" ) return nbest_hyps @@ -469,15 +474,13 @@ class BeamSearchScama(torch.nn.Module): logging.debug(f"the number of running hypotheses: {len(running_hyps)}") if self.token_list is not None: logging.debug( - "best hypo: " - + "".join([self.token_list[x] for x in running_hyps[0].yseq[1:]]) + "best hypo: " + "".join([self.token_list[x] for x in running_hyps[0].yseq[1:]]) ) # add eos in the final loop to avoid that there are no ended hyps if i == maxlen - 1: logging.info("adding in the last position in the loop") running_hyps = [ - h._replace(yseq=self.append_token(h.yseq, self.eos)) - for h in running_hyps + h._replace(yseq=self.append_token(h.yseq, self.eos)) for h in running_hyps ] # add ended hypotheses to a final list, and removed them from current hypotheses diff --git a/funasr/models/uniasr/model.py b/funasr/models/uniasr/model.py index 6e564dca7..bde637777 100644 --- a/funasr/models/uniasr/model.py +++ b/funasr/models/uniasr/model.py @@ -22,6 +22,7 @@ from funasr.models.transformer.utils.nets_utils import make_pad_mask, pad_list from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank from funasr.models.scama.utils import sequence_mask + @tables.register("model_classes", "UniASR") class UniASR(torch.nn.Module): """ @@ -56,8 +57,8 @@ class UniASR(torch.nn.Module): ctc2: str = None, ctc2_conf: dict = None, ctc2_weight: float = 0.5, - decoder_attention_chunk_type: str = 'chunk', - decoder_attention_chunk_type2: str = 'chunk', + decoder_attention_chunk_type: str = "chunk", + decoder_attention_chunk_type2: str = "chunk", stride_conv=None, stride_conv_conf: dict = None, loss_weight_model1: float = 0.5, @@ -71,7 +72,6 @@ class UniASR(torch.nn.Module): length_normalized_loss: bool = False, share_embedding: bool = False, **kwargs, - ): super().__init__() @@ -81,7 +81,7 @@ class UniASR(torch.nn.Module): if normalize is not None: normalize_class = tables.normalize_classes.get(normalize) normalize = normalize_class(**normalize_conf) - + encoder_class = tables.encoder_classes.get(encoder) encoder = encoder_class(input_size=input_size, **encoder_conf) encoder_output_size = encoder.output_size() @@ -94,12 +94,14 @@ class UniASR(torch.nn.Module): ) predictor_class = tables.predictor_classes.get(predictor) predictor = predictor_class(**predictor_conf) - - from funasr.models.transformer.utils.subsampling import Conv1dSubsampling - stride_conv = Conv1dSubsampling(**stride_conv_conf, idim=input_size + encoder_output_size, - odim=input_size + encoder_output_size) + + stride_conv = Conv1dSubsampling( + **stride_conv_conf, + idim=input_size + encoder_output_size, + odim=input_size + encoder_output_size, + ) stride_conv_output_size = stride_conv.output_size() encoder_class = tables.encoder_classes.get(encoder2) @@ -115,8 +117,6 @@ class UniASR(torch.nn.Module): predictor_class = tables.predictor_classes.get(predictor2) predictor2 = predictor_class(**predictor2_conf) - - self.blank_id = blank_id self.sos = sos self.eos = eos @@ -127,7 +127,7 @@ class UniASR(torch.nn.Module): self.specaug = specaug self.normalize = normalize - + self.encoder = encoder self.error_calculator = None @@ -142,16 +142,20 @@ class UniASR(torch.nn.Module): smoothing=lsm_weight, normalize_length=length_normalized_loss, ) - + self.predictor = predictor self.predictor_weight = predictor_weight self.criterion_pre = mae_loss(normalize_length=length_normalized_loss) self.encoder1_encoder2_joint_training = kwargs.get("encoder1_encoder2_joint_training", True) - if self.encoder.overlap_chunk_cls is not None: - from funasr.models.scama.chunk_utilis import build_scama_mask_for_cross_attention_decoder - self.build_scama_mask_for_cross_attention_decoder_fn = build_scama_mask_for_cross_attention_decoder + from funasr.models.scama.chunk_utilis import ( + build_scama_mask_for_cross_attention_decoder, + ) + + self.build_scama_mask_for_cross_attention_decoder_fn = ( + build_scama_mask_for_cross_attention_decoder + ) self.decoder_attention_chunk_type = decoder_attention_chunk_type self.encoder2 = encoder2 @@ -164,8 +168,13 @@ class UniASR(torch.nn.Module): self.stride_conv = stride_conv self.loss_weight_model1 = loss_weight_model1 if self.encoder2.overlap_chunk_cls is not None: - from funasr.models.scama.chunk_utilis import build_scama_mask_for_cross_attention_decoder - self.build_scama_mask_for_cross_attention_decoder_fn2 = build_scama_mask_for_cross_attention_decoder + from funasr.models.scama.chunk_utilis import ( + build_scama_mask_for_cross_attention_decoder, + ) + + self.build_scama_mask_for_cross_attention_decoder_fn2 = ( + build_scama_mask_for_cross_attention_decoder + ) self.decoder_attention_chunk_type2 = decoder_attention_chunk_type2 self.length_normalized_loss = length_normalized_loss @@ -196,16 +205,16 @@ class UniASR(torch.nn.Module): batch_size = speech.shape[0] - ind = self.encoder.overlap_chunk_cls.random_choice(self.training, decoding_ind) # 1. Encoder if self.enable_maas_finetune: with torch.no_grad(): - speech_raw, encoder_out, encoder_out_lens = self.encode(speech, speech_lengths, ind=ind) + speech_raw, encoder_out, encoder_out_lens = self.encode( + speech, speech_lengths, ind=ind + ) else: speech_raw, encoder_out, encoder_out_lens = self.encode(speech, speech_lengths, ind=ind) - loss_att, acc_att, cer_att, wer_att = None, None, None, None loss_ctc, cer_ctc = None, None stats = dict() @@ -231,12 +240,11 @@ class UniASR(torch.nn.Module): stats["wer"] = wer_att stats["loss_pre"] = loss_pre.detach().cpu() if loss_pre is not None else None else: - + loss_att, acc_att, cer_att, wer_att, loss_pre = self._calc_att_predictor_loss( encoder_out, encoder_out_lens, text, text_lengths ) - loss = loss_att + loss_pre * self.predictor_weight # Collect Attn branch stats @@ -254,21 +262,23 @@ class UniASR(torch.nn.Module): # encoder2 if self.freeze_encoder2: with torch.no_grad(): - encoder_out, encoder_out_lens = self.encode2(encoder_out, encoder_out_lens, speech_raw, speech_lengths, ind=ind) + encoder_out, encoder_out_lens = self.encode2( + encoder_out, encoder_out_lens, speech_raw, speech_lengths, ind=ind + ) else: - encoder_out, encoder_out_lens = self.encode2(encoder_out, encoder_out_lens, speech_raw, speech_lengths, ind=ind) + encoder_out, encoder_out_lens = self.encode2( + encoder_out, encoder_out_lens, speech_raw, speech_lengths, ind=ind + ) intermediate_outs = None if isinstance(encoder_out, tuple): intermediate_outs = encoder_out[1] encoder_out = encoder_out[0] - loss_att, acc_att, cer_att, wer_att, loss_pre = self._calc_att_predictor_loss2( encoder_out, encoder_out_lens, text, text_lengths ) - loss = loss_att + loss_pre * self.predictor2_weight # Collect Attn branch stats @@ -277,7 +287,7 @@ class UniASR(torch.nn.Module): stats["cer2"] = cer_att stats["wer2"] = wer_att stats["loss_pre2"] = loss_pre.detach().cpu() if loss_pre is not None else None - + loss2 = loss loss = loss1 * self.loss_weight_model1 + loss2 * (1 - self.loss_weight_model1) @@ -288,7 +298,6 @@ class UniASR(torch.nn.Module): if self.length_normalized_loss: batch_size = int((text_lengths + 1).sum()) - loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) return loss, stats, weight @@ -312,7 +321,10 @@ class UniASR(torch.nn.Module): return {"feats": feats, "feats_lengths": feats_lengths} def encode( - self, speech: torch.Tensor, speech_lengths: torch.Tensor, **kwargs, + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + **kwargs, ): """Frontend + Encoder. Note that this method is used by asr_inference.py Args: @@ -324,13 +336,12 @@ class UniASR(torch.nn.Module): # Data augmentation if self.specaug is not None and self.training: speech, speech_lengths = self.specaug(speech, speech_lengths) - + # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN if self.normalize is not None: speech, speech_lengths = self.normalize(speech, speech_lengths) - - speech_raw = speech.clone().to(speech.device) + speech_raw = speech.clone().to(speech.device) # 4. Forward encoder encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths, ind=ind) @@ -375,10 +386,8 @@ class UniASR(torch.nn.Module): if isinstance(encoder_out, tuple): encoder_out = encoder_out[0] - return encoder_out, encoder_out_lens - def nll( self, encoder_out: torch.Tensor, @@ -472,9 +481,7 @@ class UniASR(torch.nn.Module): ys_in_lens = ys_pad_lens + 1 # 1. Forward decoder - decoder_out, _ = self.decoder( - encoder_out, encoder_out_lens, ys_in_pad, ys_in_lens - ) + decoder_out, _ = self.decoder(encoder_out, encoder_out_lens, ys_in_pad, ys_in_lens) # 2. Compute attention loss loss_att = self.criterion_att(decoder_out, ys_out_pad) @@ -503,37 +510,49 @@ class UniASR(torch.nn.Module): ys_in_pad, ys_out_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id) ys_in_lens = ys_pad_lens + 1 - encoder_out_mask = sequence_mask(encoder_out_lens, maxlen=encoder_out.size(1), dtype=encoder_out.dtype, - device=encoder_out.device)[:, None, :] + encoder_out_mask = sequence_mask( + encoder_out_lens, + maxlen=encoder_out.size(1), + dtype=encoder_out.dtype, + device=encoder_out.device, + )[:, None, :] mask_chunk_predictor = None if self.encoder.overlap_chunk_cls is not None: - mask_chunk_predictor = self.encoder.overlap_chunk_cls.get_mask_chunk_predictor(None, - device=encoder_out.device, - batch_size=encoder_out.size( - 0)) - mask_shfit_chunk = self.encoder.overlap_chunk_cls.get_mask_shfit_chunk(None, device=encoder_out.device, - batch_size=encoder_out.size(0)) + mask_chunk_predictor = self.encoder.overlap_chunk_cls.get_mask_chunk_predictor( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + mask_shfit_chunk = self.encoder.overlap_chunk_cls.get_mask_shfit_chunk( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) encoder_out = encoder_out * mask_shfit_chunk - pre_acoustic_embeds, pre_token_length, pre_alphas, _ = self.predictor(encoder_out, - ys_out_pad, - encoder_out_mask, - ignore_id=self.ignore_id, - mask_chunk_predictor=mask_chunk_predictor, - target_label_length=ys_in_lens, - ) - predictor_alignments, predictor_alignments_len = self.predictor.gen_frame_alignments(pre_alphas, - encoder_out_lens) + pre_acoustic_embeds, pre_token_length, pre_alphas, _ = self.predictor( + encoder_out, + ys_out_pad, + encoder_out_mask, + ignore_id=self.ignore_id, + mask_chunk_predictor=mask_chunk_predictor, + target_label_length=ys_in_lens, + ) + predictor_alignments, predictor_alignments_len = self.predictor.gen_frame_alignments( + pre_alphas, encoder_out_lens + ) scama_mask = None - if self.encoder.overlap_chunk_cls is not None and self.decoder_attention_chunk_type == 'chunk': + if ( + self.encoder.overlap_chunk_cls is not None + and self.decoder_attention_chunk_type == "chunk" + ): encoder_chunk_size = self.encoder.overlap_chunk_cls.chunk_size_pad_shift_cur attention_chunk_center_bias = 0 attention_chunk_size = encoder_chunk_size - decoder_att_look_back_factor = self.encoder.overlap_chunk_cls.decoder_att_look_back_factor_cur - mask_shift_att_chunk_decoder = self.encoder.overlap_chunk_cls.get_mask_shift_att_chunk_decoder(None, - device=encoder_out.device, - batch_size=encoder_out.size( - 0)) + decoder_att_look_back_factor = ( + self.encoder.overlap_chunk_cls.decoder_att_look_back_factor_cur + ) + mask_shift_att_chunk_decoder = ( + self.encoder.overlap_chunk_cls.get_mask_shift_att_chunk_decoder( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + ) scama_mask = self.build_scama_mask_for_cross_attention_decoder_fn( predictor_alignments=predictor_alignments, encoder_sequence_length=encoder_out_lens, @@ -550,8 +569,9 @@ class UniASR(torch.nn.Module): is_training=self.training, ) elif self.encoder.overlap_chunk_cls is not None: - encoder_out, encoder_out_lens = self.encoder.overlap_chunk_cls.remove_chunk(encoder_out, encoder_out_lens, - chunk_outs=None) + encoder_out, encoder_out_lens = self.encoder.overlap_chunk_cls.remove_chunk( + encoder_out, encoder_out_lens, chunk_outs=None + ) # try: # 1. Forward decoder decoder_out, _ = self.decoder( @@ -561,7 +581,6 @@ class UniASR(torch.nn.Module): ys_in_lens, chunk_mask=scama_mask, pre_acoustic_embeds=pre_acoustic_embeds, - ) # 2. Compute attention loss @@ -592,37 +611,49 @@ class UniASR(torch.nn.Module): ys_in_pad, ys_out_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id) ys_in_lens = ys_pad_lens + 1 - encoder_out_mask = sequence_mask(encoder_out_lens, maxlen=encoder_out.size(1), dtype=encoder_out.dtype, - device=encoder_out.device)[:, None, :] + encoder_out_mask = sequence_mask( + encoder_out_lens, + maxlen=encoder_out.size(1), + dtype=encoder_out.dtype, + device=encoder_out.device, + )[:, None, :] mask_chunk_predictor = None if self.encoder2.overlap_chunk_cls is not None: - mask_chunk_predictor = self.encoder2.overlap_chunk_cls.get_mask_chunk_predictor(None, - device=encoder_out.device, - batch_size=encoder_out.size( - 0)) - mask_shfit_chunk = self.encoder2.overlap_chunk_cls.get_mask_shfit_chunk(None, device=encoder_out.device, - batch_size=encoder_out.size(0)) + mask_chunk_predictor = self.encoder2.overlap_chunk_cls.get_mask_chunk_predictor( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + mask_shfit_chunk = self.encoder2.overlap_chunk_cls.get_mask_shfit_chunk( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) encoder_out = encoder_out * mask_shfit_chunk - pre_acoustic_embeds, pre_token_length, pre_alphas, _ = self.predictor2(encoder_out, - ys_out_pad, - encoder_out_mask, - ignore_id=self.ignore_id, - mask_chunk_predictor=mask_chunk_predictor, - target_label_length=ys_in_lens, - ) - predictor_alignments, predictor_alignments_len = self.predictor2.gen_frame_alignments(pre_alphas, - encoder_out_lens) + pre_acoustic_embeds, pre_token_length, pre_alphas, _ = self.predictor2( + encoder_out, + ys_out_pad, + encoder_out_mask, + ignore_id=self.ignore_id, + mask_chunk_predictor=mask_chunk_predictor, + target_label_length=ys_in_lens, + ) + predictor_alignments, predictor_alignments_len = self.predictor2.gen_frame_alignments( + pre_alphas, encoder_out_lens + ) scama_mask = None - if self.encoder2.overlap_chunk_cls is not None and self.decoder_attention_chunk_type2 == 'chunk': + if ( + self.encoder2.overlap_chunk_cls is not None + and self.decoder_attention_chunk_type2 == "chunk" + ): encoder_chunk_size = self.encoder2.overlap_chunk_cls.chunk_size_pad_shift_cur attention_chunk_center_bias = 0 attention_chunk_size = encoder_chunk_size - decoder_att_look_back_factor = self.encoder2.overlap_chunk_cls.decoder_att_look_back_factor_cur - mask_shift_att_chunk_decoder = self.encoder2.overlap_chunk_cls.get_mask_shift_att_chunk_decoder(None, - device=encoder_out.device, - batch_size=encoder_out.size( - 0)) + decoder_att_look_back_factor = ( + self.encoder2.overlap_chunk_cls.decoder_att_look_back_factor_cur + ) + mask_shift_att_chunk_decoder = ( + self.encoder2.overlap_chunk_cls.get_mask_shift_att_chunk_decoder( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + ) scama_mask = self.build_scama_mask_for_cross_attention_decoder_fn2( predictor_alignments=predictor_alignments, encoder_sequence_length=encoder_out_lens, @@ -639,8 +670,9 @@ class UniASR(torch.nn.Module): is_training=self.training, ) elif self.encoder2.overlap_chunk_cls is not None: - encoder_out, encoder_out_lens = self.encoder2.overlap_chunk_cls.remove_chunk(encoder_out, encoder_out_lens, - chunk_outs=None) + encoder_out, encoder_out_lens = self.encoder2.overlap_chunk_cls.remove_chunk( + encoder_out, encoder_out_lens, chunk_outs=None + ) # try: # 1. Forward decoder decoder_out, _ = self.decoder2( @@ -681,37 +713,49 @@ class UniASR(torch.nn.Module): # ys_in_lens = ys_pad_lens + 1 ys_out_pad, ys_in_lens = None, None - encoder_out_mask = sequence_mask(encoder_out_lens, maxlen=encoder_out.size(1), dtype=encoder_out.dtype, - device=encoder_out.device)[:, None, :] + encoder_out_mask = sequence_mask( + encoder_out_lens, + maxlen=encoder_out.size(1), + dtype=encoder_out.dtype, + device=encoder_out.device, + )[:, None, :] mask_chunk_predictor = None if self.encoder.overlap_chunk_cls is not None: - mask_chunk_predictor = self.encoder.overlap_chunk_cls.get_mask_chunk_predictor(None, - device=encoder_out.device, - batch_size=encoder_out.size( - 0)) - mask_shfit_chunk = self.encoder.overlap_chunk_cls.get_mask_shfit_chunk(None, device=encoder_out.device, - batch_size=encoder_out.size(0)) + mask_chunk_predictor = self.encoder.overlap_chunk_cls.get_mask_chunk_predictor( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + mask_shfit_chunk = self.encoder.overlap_chunk_cls.get_mask_shfit_chunk( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) encoder_out = encoder_out * mask_shfit_chunk - pre_acoustic_embeds, pre_token_length, pre_alphas, _ = self.predictor(encoder_out, - ys_out_pad, - encoder_out_mask, - ignore_id=self.ignore_id, - mask_chunk_predictor=mask_chunk_predictor, - target_label_length=ys_in_lens, - ) - predictor_alignments, predictor_alignments_len = self.predictor.gen_frame_alignments(pre_alphas, - encoder_out_lens) + pre_acoustic_embeds, pre_token_length, pre_alphas, _ = self.predictor( + encoder_out, + ys_out_pad, + encoder_out_mask, + ignore_id=self.ignore_id, + mask_chunk_predictor=mask_chunk_predictor, + target_label_length=ys_in_lens, + ) + predictor_alignments, predictor_alignments_len = self.predictor.gen_frame_alignments( + pre_alphas, encoder_out_lens + ) scama_mask = None - if self.encoder.overlap_chunk_cls is not None and self.decoder_attention_chunk_type == 'chunk': + if ( + self.encoder.overlap_chunk_cls is not None + and self.decoder_attention_chunk_type == "chunk" + ): encoder_chunk_size = self.encoder.overlap_chunk_cls.chunk_size_pad_shift_cur attention_chunk_center_bias = 0 attention_chunk_size = encoder_chunk_size - decoder_att_look_back_factor = self.encoder.overlap_chunk_cls.decoder_att_look_back_factor_cur - mask_shift_att_chunk_decoder = self.encoder.overlap_chunk_cls.get_mask_shift_att_chunk_decoder(None, - device=encoder_out.device, - batch_size=encoder_out.size( - 0)) + decoder_att_look_back_factor = ( + self.encoder.overlap_chunk_cls.decoder_att_look_back_factor_cur + ) + mask_shift_att_chunk_decoder = ( + self.encoder.overlap_chunk_cls.get_mask_shift_att_chunk_decoder( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + ) scama_mask = self.build_scama_mask_for_cross_attention_decoder_fn( predictor_alignments=predictor_alignments, encoder_sequence_length=encoder_out_lens, @@ -728,10 +772,17 @@ class UniASR(torch.nn.Module): is_training=self.training, ) elif self.encoder.overlap_chunk_cls is not None: - encoder_out, encoder_out_lens = self.encoder.overlap_chunk_cls.remove_chunk(encoder_out, encoder_out_lens, - chunk_outs=None) + encoder_out, encoder_out_lens = self.encoder.overlap_chunk_cls.remove_chunk( + encoder_out, encoder_out_lens, chunk_outs=None + ) - return pre_acoustic_embeds, pre_token_length, predictor_alignments, predictor_alignments_len, scama_mask + return ( + pre_acoustic_embeds, + pre_token_length, + predictor_alignments, + predictor_alignments_len, + scama_mask, + ) def calc_predictor_mask2( self, @@ -744,37 +795,49 @@ class UniASR(torch.nn.Module): # ys_in_lens = ys_pad_lens + 1 ys_out_pad, ys_in_lens = None, None - encoder_out_mask = sequence_mask(encoder_out_lens, maxlen=encoder_out.size(1), dtype=encoder_out.dtype, - device=encoder_out.device)[:, None, :] + encoder_out_mask = sequence_mask( + encoder_out_lens, + maxlen=encoder_out.size(1), + dtype=encoder_out.dtype, + device=encoder_out.device, + )[:, None, :] mask_chunk_predictor = None if self.encoder2.overlap_chunk_cls is not None: - mask_chunk_predictor = self.encoder2.overlap_chunk_cls.get_mask_chunk_predictor(None, - device=encoder_out.device, - batch_size=encoder_out.size( - 0)) - mask_shfit_chunk = self.encoder2.overlap_chunk_cls.get_mask_shfit_chunk(None, device=encoder_out.device, - batch_size=encoder_out.size(0)) + mask_chunk_predictor = self.encoder2.overlap_chunk_cls.get_mask_chunk_predictor( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + mask_shfit_chunk = self.encoder2.overlap_chunk_cls.get_mask_shfit_chunk( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) encoder_out = encoder_out * mask_shfit_chunk - pre_acoustic_embeds, pre_token_length, pre_alphas, _ = self.predictor2(encoder_out, - ys_out_pad, - encoder_out_mask, - ignore_id=self.ignore_id, - mask_chunk_predictor=mask_chunk_predictor, - target_label_length=ys_in_lens, - ) - predictor_alignments, predictor_alignments_len = self.predictor2.gen_frame_alignments(pre_alphas, - encoder_out_lens) + pre_acoustic_embeds, pre_token_length, pre_alphas, _ = self.predictor2( + encoder_out, + ys_out_pad, + encoder_out_mask, + ignore_id=self.ignore_id, + mask_chunk_predictor=mask_chunk_predictor, + target_label_length=ys_in_lens, + ) + predictor_alignments, predictor_alignments_len = self.predictor2.gen_frame_alignments( + pre_alphas, encoder_out_lens + ) scama_mask = None - if self.encoder2.overlap_chunk_cls is not None and self.decoder_attention_chunk_type2 == 'chunk': + if ( + self.encoder2.overlap_chunk_cls is not None + and self.decoder_attention_chunk_type2 == "chunk" + ): encoder_chunk_size = self.encoder2.overlap_chunk_cls.chunk_size_pad_shift_cur attention_chunk_center_bias = 0 attention_chunk_size = encoder_chunk_size - decoder_att_look_back_factor = self.encoder2.overlap_chunk_cls.decoder_att_look_back_factor_cur - mask_shift_att_chunk_decoder = self.encoder2.overlap_chunk_cls.get_mask_shift_att_chunk_decoder(None, - device=encoder_out.device, - batch_size=encoder_out.size( - 0)) + decoder_att_look_back_factor = ( + self.encoder2.overlap_chunk_cls.decoder_att_look_back_factor_cur + ) + mask_shift_att_chunk_decoder = ( + self.encoder2.overlap_chunk_cls.get_mask_shift_att_chunk_decoder( + None, device=encoder_out.device, batch_size=encoder_out.size(0) + ) + ) scama_mask = self.build_scama_mask_for_cross_attention_decoder_fn2( predictor_alignments=predictor_alignments, encoder_sequence_length=encoder_out_lens, @@ -791,14 +854,22 @@ class UniASR(torch.nn.Module): is_training=self.training, ) elif self.encoder2.overlap_chunk_cls is not None: - encoder_out, encoder_out_lens = self.encoder2.overlap_chunk_cls.remove_chunk(encoder_out, encoder_out_lens, - chunk_outs=None) + encoder_out, encoder_out_lens = self.encoder2.overlap_chunk_cls.remove_chunk( + encoder_out, encoder_out_lens, chunk_outs=None + ) - return pre_acoustic_embeds, pre_token_length, predictor_alignments, predictor_alignments_len, scama_mask + return ( + pre_acoustic_embeds, + pre_token_length, + predictor_alignments, + predictor_alignments_len, + scama_mask, + ) - def init_beam_search(self, - **kwargs, - ): + def init_beam_search( + self, + **kwargs, + ): from funasr.models.uniasr.beam_search import BeamSearchScama from funasr.models.transformer.scorers.ctc import CTCPrefixScorer from funasr.models.transformer.scorers.length_bonus import LengthBonus @@ -810,23 +881,21 @@ class UniASR(torch.nn.Module): decoder = self.decoder2 # 1. Build ASR model scorers = {} - + if self.ctc != None: ctc = CTCPrefixScorer(ctc=self.ctc, eos=self.eos) - scorers.update( - ctc=ctc - ) + scorers.update(ctc=ctc) token_list = kwargs.get("token_list") scorers.update( decoder=decoder, length_bonus=LengthBonus(len(token_list)), ) - + # 3. Build ngram model # ngram is not supported now ngram = None scorers["ngram"] = ngram - + weights = dict( decoder=1.0 - kwargs.get("decoding_ctc_weight", 0.0), ctc=kwargs.get("decoding_ctc_weight", 0.0), @@ -844,17 +913,18 @@ class UniASR(torch.nn.Module): token_list=token_list, pre_beam_score_key=None if self.ctc_weight == 1.0 else "full", ) - + self.beam_search = beam_search - def inference(self, - data_in, - data_lengths=None, - key: list = None, - tokenizer=None, - frontend=None, - **kwargs, - ): + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): decoding_model = kwargs.get("decoding_model", "normal") token_num_relax = kwargs.get("token_num_relax", 5) @@ -868,14 +938,16 @@ class UniASR(torch.nn.Module): decoding_ind = 0 decoding_mode = "model2" # init beamsearch - + if self.beam_search is None: logging.info("enable beam_search") self.init_beam_search(decoding_mode=decoding_mode, **kwargs) self.nbest = kwargs.get("nbest", 1) - + meta_data = {} - if isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank": # fbank + if ( + isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank" + ): # fbank speech, speech_lengths = data_in, data_lengths if len(speech.shape) < 3: speech = speech[None, :, :] @@ -884,17 +956,24 @@ class UniASR(torch.nn.Module): else: # extract fbank feats time1 = time.perf_counter() - audio_sample_list = load_audio_text_image_video(data_in, fs=frontend.fs, audio_fs=kwargs.get("fs", 16000), - data_type=kwargs.get("data_type", "sound"), - tokenizer=tokenizer) + audio_sample_list = load_audio_text_image_video( + data_in, + fs=frontend.fs, + audio_fs=kwargs.get("fs", 16000), + data_type=kwargs.get("data_type", "sound"), + tokenizer=tokenizer, + ) time2 = time.perf_counter() meta_data["load_data"] = f"{time2 - time1:0.3f}" - speech, speech_lengths = extract_fbank(audio_sample_list, data_type=kwargs.get("data_type", "sound"), - frontend=frontend) + speech, speech_lengths = extract_fbank( + audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend + ) time3 = time.perf_counter() meta_data["extract_feat"] = f"{time3 - time2:0.3f}" - meta_data["batch_data_time"] = speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 - + meta_data["batch_data_time"] = ( + speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 + ) + speech = speech.to(device=kwargs["device"]) speech_lengths = speech_lengths.to(device=kwargs["device"]) speech_raw = speech.clone().to(device=kwargs["device"]) @@ -903,10 +982,11 @@ class UniASR(torch.nn.Module): if decoding_mode == "model1": predictor_outs = self.calc_predictor_mask(encoder_out, encoder_out_lens) else: - encoder_out, encoder_out_lens = self.encode2(encoder_out, encoder_out_lens, speech_raw, speech_lengths, ind=decoding_ind) + encoder_out, encoder_out_lens = self.encode2( + encoder_out, encoder_out_lens, speech_raw, speech_lengths, ind=decoding_ind + ) predictor_outs = self.calc_predictor_mask2(encoder_out, encoder_out_lens) - scama_mask = predictor_outs[4] pre_token_length = predictor_outs[1] pre_acoustic_embeds = predictor_outs[0] @@ -914,8 +994,13 @@ class UniASR(torch.nn.Module): minlen = max(0, pre_token_length.sum().item() - token_num_relax) # c. Passed the encoder result and the beam search nbest_hyps = self.beam_search( - x=encoder_out[0], scama_mask=scama_mask, pre_acoustic_embeds=pre_acoustic_embeds, maxlenratio=0.0, - minlenratio=0.0, maxlen=int(maxlen), minlen=int(minlen), + x=encoder_out[0], + scama_mask=scama_mask, + pre_acoustic_embeds=pre_acoustic_embeds, + maxlenratio=0.0, + minlenratio=0.0, + maxlen=int(maxlen), + minlen=int(minlen), ) nbest_hyps = nbest_hyps[: self.nbest] @@ -933,15 +1018,13 @@ class UniASR(torch.nn.Module): # remove blank symbol id, which is assumed to be 0 token_int = list(filter(lambda x: x != 0, token_int)) - # Change integer-ids to tokens token = tokenizer.ids2tokens(token_int) text_postprocessed = tokenizer.tokens2text(token) if not hasattr(tokenizer, "bpemodel"): text_postprocessed, _ = postprocess_utils.sentence_postprocess(token) - result_i = {"key": key[0], "text": text_postprocessed} results.append(result_i) - return results, meta_data \ No newline at end of file + return results, meta_data diff --git a/funasr/models/whisper/model.py b/funasr/models/whisper/model.py index 51fde815d..8e9245af5 100644 --- a/funasr/models/whisper/model.py +++ b/funasr/models/whisper/model.py @@ -38,34 +38,41 @@ class WhisperWarp(nn.Module): dims = kwargs.get("dims", {}) dims = whisper.model.ModelDimensions(**dims) model = whisper.model.Whisper(dims=dims) - + self.model = model - + self.encoder_output_size = self.model.dims.n_audio_state - - def forward(self, ): + + def forward( + self, + ): pass - - def inference(self, - data_in, - data_lengths=None, - key: list = None, - tokenizer=None, - frontend=None, - **kwargs, - ): + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): if kwargs.get("batch_size", 1) > 1: raise NotImplementedError("batch decoding is not implemented") if frontend is None and not hasattr(self, "frontend"): frontend_class = tables.frontend_classes.get("WhisperFrontend") - frontend = frontend_class(n_mels=self.model.dims.n_mels, do_pad_trim=kwargs.get("do_pad_trim", True)) + frontend = frontend_class( + n_mels=self.model.dims.n_mels, do_pad_trim=kwargs.get("do_pad_trim", True) + ) self.frontend = frontend else: frontend = frontend if frontend is not None else self.frontend meta_data = {} - if isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank": # fbank + if ( + isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank" + ): # fbank speech, speech_lengths = data_in, data_lengths if len(speech.shape) < 3: speech = speech[None, :, :] @@ -74,13 +81,18 @@ class WhisperWarp(nn.Module): else: # extract fbank feats time1 = time.perf_counter() - audio_sample_list = load_audio_text_image_video(data_in, fs=frontend.fs if hasattr(frontend, "fs") else 16000, audio_fs=kwargs.get("fs", 16000), - data_type=kwargs.get("data_type", "sound"), - tokenizer=tokenizer) + audio_sample_list = load_audio_text_image_video( + data_in, + fs=frontend.fs if hasattr(frontend, "fs") else 16000, + audio_fs=kwargs.get("fs", 16000), + data_type=kwargs.get("data_type", "sound"), + tokenizer=tokenizer, + ) time2 = time.perf_counter() meta_data["load_data"] = f"{time2 - time1:0.3f}" - speech, speech_lengths = extract_fbank(audio_sample_list, data_type=kwargs.get("data_type", "sound"), - frontend=frontend) + speech, speech_lengths = extract_fbank( + audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend + ) time3 = time.perf_counter() meta_data["extract_feat"] = f"{time3 - time2:0.3f}" frame_shift = frontend.frame_shift if hasattr(frontend, "frame_shift") else 10 @@ -102,6 +114,5 @@ class WhisperWarp(nn.Module): result_i = {"key": key[0], "text": result.text} results.append(result_i) - + return results, meta_data - \ No newline at end of file diff --git a/funasr/models/whisper_lid/decoder.py b/funasr/models/whisper_lid/decoder.py index 4db920546..1c8ab47f8 100644 --- a/funasr/models/whisper_lid/decoder.py +++ b/funasr/models/whisper_lid/decoder.py @@ -29,9 +29,7 @@ class OpenAIWhisperDecoderWarp(nn.Module): super().__init__() assert whisper_model in whisper.available_models() - _model = whisper.load_model( - whisper_model, download_root=download_dir, device="cpu" - ) + _model = whisper.load_model(whisper_model, download_root=download_dir, device="cpu") self.decoders = copy.deepcopy(_model.decoder) attention_dim = self.decoders.token_embedding.embedding_dim @@ -67,10 +65,7 @@ class OpenAIWhisperDecoderWarp(nn.Module): olens: (batch, ) """ tgt, memory = ys_in_pad, hs_pad - tgt = ( - self.decoders.token_embedding(tgt) - + self.decoders.positional_embedding[: tgt.size(1)] - ) + tgt = self.decoders.token_embedding(tgt) + self.decoders.positional_embedding[: tgt.size(1)] tgt = self.dropout(tgt) x = tgt.to(memory.dtype) @@ -81,15 +76,20 @@ class OpenAIWhisperDecoderWarp(nn.Module): memory_mask = None for layer, block in enumerate(self.decoders.blocks): - x = block(x, memory, mask=self.decoders.mask, memory_mask=memory_mask, is_pad_mask=False, is_pad_memory_mask=True) + x = block( + x, + memory, + mask=self.decoders.mask, + memory_mask=memory_mask, + is_pad_mask=False, + is_pad_memory_mask=True, + ) if layer < len(self.decoders.blocks) - 1: x = self.dropout(x) x = self.decoders.ln(x) - x = ( - x @ torch.transpose(self.decoders.token_embedding.weight.to(x.dtype), 0, 1) - ).float() + x = (x @ torch.transpose(self.decoders.token_embedding.weight.to(x.dtype), 0, 1)).float() return x, ys_in_lens @@ -116,10 +116,7 @@ class OpenAIWhisperDecoderWarp(nn.Module): cache implementation is ignored for now for simplicity & correctness """ - x = ( - self.decoders.token_embedding(tgt) - + self.decoders.positional_embedding[: tgt.size(1)] - ) + x = self.decoders.token_embedding(tgt) + self.decoders.positional_embedding[: tgt.size(1)] x = self.dropout(x) x = x.to(memory.dtype) @@ -130,9 +127,7 @@ class OpenAIWhisperDecoderWarp(nn.Module): x = self.decoders.ln(x) y = x[:, -1] - y = ( - y @ torch.transpose(self.decoders.token_embedding.weight.to(x.dtype), 0, 1) - ).float() + y = (y @ torch.transpose(self.decoders.token_embedding.weight.to(x.dtype), 0, 1)).float() y = torch.log_softmax(y, dim=-1) return y, None diff --git a/funasr/models/whisper_lid/encoder.py b/funasr/models/whisper_lid/encoder.py index 7eeb643d6..b6fb28b50 100644 --- a/funasr/models/whisper_lid/encoder.py +++ b/funasr/models/whisper_lid/encoder.py @@ -22,13 +22,13 @@ class OpenAIWhisperEncoderWarp(nn.Module): """ def __init__( - self, - dropout_rate: float = 0.0, - whisper_model: str = "small", - download_dir: str = None, - use_specaug: bool = False, - use_padmask: bool = False, - specaug_conf: Union[dict, None] = None, + self, + dropout_rate: float = 0.0, + whisper_model: str = "small", + download_dir: str = None, + use_specaug: bool = False, + use_padmask: bool = False, + specaug_conf: Union[dict, None] = None, ): super().__init__() @@ -36,9 +36,7 @@ class OpenAIWhisperEncoderWarp(nn.Module): self.dropout = torch.nn.Dropout(dropout_rate) assert whisper_model in whisper.available_models() - _model = whisper.load_model( - whisper_model, download_root=download_dir, device="cpu" - ) + _model = whisper.load_model(whisper_model, download_root=download_dir, device="cpu") self.encoders = copy.deepcopy(_model.encoder) self.encoders.train() @@ -51,9 +49,9 @@ class OpenAIWhisperEncoderWarp(nn.Module): self.use_padmask = use_padmask def whisper_encode( - self, - input: torch.Tensor, - ilens: torch.Tensor = None, + self, + input: torch.Tensor, + ilens: torch.Tensor = None, ) -> torch.Tensor: x = F.gelu(self.encoders.conv1(input)) x = F.gelu(self.encoders.conv2(x)) @@ -69,13 +67,9 @@ class OpenAIWhisperEncoderWarp(nn.Module): if ilens is not None: olens = ( - 1 - + ( - ilens - - self.encoders.conv2.kernel_size[0] - + 2 * self.encoders.conv2.padding[0] - ) - // self.encoders.conv2.stride[0] + 1 + + (ilens - self.encoders.conv2.kernel_size[0] + 2 * self.encoders.conv2.padding[0]) + // self.encoders.conv2.stride[0] ) olens = torch.clamp(olens, max=max_pos) else: @@ -102,10 +96,10 @@ class OpenAIWhisperEncoderWarp(nn.Module): return self.encoders.conv2.weight.shape[0] def forward( - self, - xs_pad: torch.Tensor, - ilens: torch.Tensor, - prev_states: torch.Tensor = None, + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, + prev_states: torch.Tensor = None, ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: feats, feats_lens = xs_pad, ilens diff --git a/funasr/models/whisper_lid/eres2net/ResNet.py b/funasr/models/whisper_lid/eres2net/ResNet.py index 25c79f583..df3b96d24 100644 --- a/funasr/models/whisper_lid/eres2net/ResNet.py +++ b/funasr/models/whisper_lid/eres2net/ResNet.py @@ -23,21 +23,18 @@ class ReLU(nn.Hardtanh): super(ReLU, self).__init__(0, 20, inplace) def __repr__(self): - inplace_str = 'inplace' if self.inplace else '' - return self.__class__.__name__ + ' (' \ - + inplace_str + ')' + inplace_str = "inplace" if self.inplace else "" + return self.__class__.__name__ + " (" + inplace_str + ")" def conv1x1(in_planes, out_planes, stride=1): "1x1 convolution without padding" - return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, - padding=0, bias=False) + return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, padding=0, bias=False) def conv3x3(in_planes, out_planes, stride=1): "3x3 convolution with padding" - return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, - padding=1, bias=False) + return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False) class BasicBlockERes2Net(nn.Module): @@ -64,12 +61,11 @@ class BasicBlockERes2Net(nn.Module): self.shortcut = nn.Sequential() if stride != 1 or in_planes != self.expansion * planes: self.shortcut = nn.Sequential( - nn.Conv2d(in_planes, - self.expansion * planes, - kernel_size=1, - stride=stride, - bias=False), - nn.BatchNorm2d(self.expansion * planes)) + nn.Conv2d( + in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False + ), + nn.BatchNorm2d(self.expansion * planes), + ) self.stride = stride self.width = width self.scale = scale @@ -132,12 +128,11 @@ class BasicBlockERes2Net_diff_AFF(nn.Module): self.shortcut = nn.Sequential() if stride != 1 or in_planes != self.expansion * planes: self.shortcut = nn.Sequential( - nn.Conv2d(in_planes, - self.expansion * planes, - kernel_size=1, - stride=stride, - bias=False), - nn.BatchNorm2d(self.expansion * planes)) + nn.Conv2d( + in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False + ), + nn.BatchNorm2d(self.expansion * planes), + ) self.stride = stride self.width = width self.scale = scale @@ -173,15 +168,17 @@ class BasicBlockERes2Net_diff_AFF(nn.Module): class ERes2Net(nn.Module): - def __init__(self, - block=BasicBlockERes2Net, - block_fuse=BasicBlockERes2Net_diff_AFF, - num_blocks=[3, 4, 6, 3], - m_channels=32, - feat_dim=80, - embedding_size=192, - pooling_func='TSTP', - two_emb_layer=False): + def __init__( + self, + block=BasicBlockERes2Net, + block_fuse=BasicBlockERes2Net_diff_AFF, + num_blocks=[3, 4, 6, 3], + m_channels=32, + feat_dim=80, + embedding_size=192, + pooling_func="TSTP", + two_emb_layer=False, + ): super(ERes2Net, self).__init__() self.in_planes = m_channels self.feat_dim = feat_dim @@ -190,48 +187,32 @@ class ERes2Net(nn.Module): self.two_emb_layer = two_emb_layer self._output_size = embedding_size - self.conv1 = nn.Conv2d(1, - m_channels, - kernel_size=3, - stride=1, - padding=1, - bias=False) + self.conv1 = nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False) self.bn1 = nn.BatchNorm2d(m_channels) - self.layer1 = self._make_layer(block, - m_channels, - num_blocks[0], - stride=1) - self.layer2 = self._make_layer(block, - m_channels * 2, - num_blocks[1], - stride=2) - self.layer3 = self._make_layer(block_fuse, - m_channels * 4, - num_blocks[2], - stride=2) - self.layer4 = self._make_layer(block_fuse, - m_channels * 8, - num_blocks[3], - stride=2) + self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=1) + self.layer2 = self._make_layer(block, m_channels * 2, num_blocks[1], stride=2) + self.layer3 = self._make_layer(block_fuse, m_channels * 4, num_blocks[2], stride=2) + self.layer4 = self._make_layer(block_fuse, m_channels * 8, num_blocks[3], stride=2) # Downsampling module for each layer - self.layer1_downsample = nn.Conv2d(m_channels * 2, m_channels * 4, kernel_size=3, stride=2, padding=1, - bias=False) - self.layer2_downsample = nn.Conv2d(m_channels * 4, m_channels * 8, kernel_size=3, padding=1, stride=2, - bias=False) - self.layer3_downsample = nn.Conv2d(m_channels * 8, m_channels * 16, kernel_size=3, padding=1, stride=2, - bias=False) + self.layer1_downsample = nn.Conv2d( + m_channels * 2, m_channels * 4, kernel_size=3, stride=2, padding=1, bias=False + ) + self.layer2_downsample = nn.Conv2d( + m_channels * 4, m_channels * 8, kernel_size=3, padding=1, stride=2, bias=False + ) + self.layer3_downsample = nn.Conv2d( + m_channels * 8, m_channels * 16, kernel_size=3, padding=1, stride=2, bias=False + ) # Bottom-up fusion module self.fuse_mode12 = AFF(channels=m_channels * 4) self.fuse_mode123 = AFF(channels=m_channels * 8) self.fuse_mode1234 = AFF(channels=m_channels * 16) - self.n_stats = 1 if pooling_func == 'TAP' or pooling_func == "TSDP" else 2 - self.pool = getattr(pooling_layers, pooling_func)( - in_dim=self.stats_dim * block.expansion) - self.seg_1 = nn.Linear(self.stats_dim * block.expansion * self.n_stats, - embedding_size) + self.n_stats = 1 if pooling_func == "TAP" or pooling_func == "TSDP" else 2 + self.pool = getattr(pooling_layers, pooling_func)(in_dim=self.stats_dim * block.expansion) + self.seg_1 = nn.Linear(self.stats_dim * block.expansion * self.n_stats, embedding_size) if self.two_emb_layer: self.seg_bn_1 = nn.BatchNorm1d(embedding_size, affine=False) self.seg_2 = nn.Linear(embedding_size, embedding_size) @@ -301,12 +282,11 @@ class BasicBlockRes2Net(nn.Module): self.shortcut = nn.Sequential() if stride != 1 or in_planes != self.expansion * planes: self.shortcut = nn.Sequential( - nn.Conv2d(in_planes, - self.expansion * planes, - kernel_size=1, - stride=stride, - bias=False), - nn.BatchNorm2d(self.expansion * planes)) + nn.Conv2d( + in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False + ), + nn.BatchNorm2d(self.expansion * planes), + ) self.stride = stride self.width = width self.scale = scale @@ -343,14 +323,16 @@ class BasicBlockRes2Net(nn.Module): class Res2Net(nn.Module): - def __init__(self, - block=BasicBlockRes2Net, - num_blocks=[3, 4, 6, 3], - m_channels=32, - feat_dim=80, - embedding_size=192, - pooling_func='TSTP', - two_emb_layer=False): + def __init__( + self, + block=BasicBlockRes2Net, + num_blocks=[3, 4, 6, 3], + m_channels=32, + feat_dim=80, + embedding_size=192, + pooling_func="TSTP", + two_emb_layer=False, + ): super(Res2Net, self).__init__() self.in_planes = m_channels self.feat_dim = feat_dim @@ -358,35 +340,16 @@ class Res2Net(nn.Module): self.stats_dim = int(feat_dim / 8) * m_channels * 8 self.two_emb_layer = two_emb_layer - self.conv1 = nn.Conv2d(1, - m_channels, - kernel_size=3, - stride=1, - padding=1, - bias=False) + self.conv1 = nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False) self.bn1 = nn.BatchNorm2d(m_channels) - self.layer1 = self._make_layer(block, - m_channels, - num_blocks[0], - stride=1) - self.layer2 = self._make_layer(block, - m_channels * 2, - num_blocks[1], - stride=2) - self.layer3 = self._make_layer(block, - m_channels * 4, - num_blocks[2], - stride=2) - self.layer4 = self._make_layer(block, - m_channels * 8, - num_blocks[3], - stride=2) + self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=1) + self.layer2 = self._make_layer(block, m_channels * 2, num_blocks[1], stride=2) + self.layer3 = self._make_layer(block, m_channels * 4, num_blocks[2], stride=2) + self.layer4 = self._make_layer(block, m_channels * 8, num_blocks[3], stride=2) - self.n_stats = 1 if pooling_func == 'TAP' or pooling_func == "TSDP" else 2 - self.pool = getattr(pooling_layers, pooling_func)( - in_dim=self.stats_dim * block.expansion) - self.seg_1 = nn.Linear(self.stats_dim * block.expansion * self.n_stats, - embedding_size) + self.n_stats = 1 if pooling_func == "TAP" or pooling_func == "TSDP" else 2 + self.pool = getattr(pooling_layers, pooling_func)(in_dim=self.stats_dim * block.expansion) + self.seg_1 = nn.Linear(self.stats_dim * block.expansion * self.n_stats, embedding_size) if self.two_emb_layer: self.seg_bn_1 = nn.BatchNorm1d(embedding_size, affine=False) self.seg_2 = nn.Linear(embedding_size, embedding_size) @@ -422,7 +385,3 @@ class Res2Net(nn.Module): return embed_b else: return embed_a - - - - diff --git a/funasr/models/whisper_lid/eres2net/fusion.py b/funasr/models/whisper_lid/eres2net/fusion.py index 2aff7a721..fbe699efb 100644 --- a/funasr/models/whisper_lid/eres2net/fusion.py +++ b/funasr/models/whisper_lid/eres2net/fusion.py @@ -23,7 +23,6 @@ class AFF(nn.Module): xa = torch.cat((x, ds_y), dim=1) x_att = self.local_att(xa) x_att = 1.0 + torch.tanh(x_att) - xo = torch.mul(x, x_att) + torch.mul(ds_y, 2.0-x_att) + xo = torch.mul(x, x_att) + torch.mul(ds_y, 2.0 - x_att) return xo - diff --git a/funasr/models/whisper_lid/eres2net/pooling_layers.py b/funasr/models/whisper_lid/eres2net/pooling_layers.py index f756ac89d..da23d7b4c 100644 --- a/funasr/models/whisper_lid/eres2net/pooling_layers.py +++ b/funasr/models/whisper_lid/eres2net/pooling_layers.py @@ -57,7 +57,9 @@ class TSTP(nn.Module): count_without_padding = torch.sum(masks, axis=-1) mean_without_padding = sum_without_padding / count_without_padding - var_without_padding = ((x_masked - mean_without_padding.unsqueeze(-1)) ** 2 * masks).sum(-1) / count_without_padding + var_without_padding = ((x_masked - mean_without_padding.unsqueeze(-1)) ** 2 * masks).sum( + -1 + ) / count_without_padding pooling_mean = mean_without_padding pooling_std = torch.sqrt(var_without_padding + 1e-8) @@ -69,8 +71,8 @@ class TSTP(nn.Module): class ASTP(nn.Module): - """ Attentive statistics pooling: Channel- and context-dependent - statistics pooling, first used in ECAPA_TDNN. + """Attentive statistics pooling: Channel- and context-dependent + statistics pooling, first used in ECAPA_TDNN. """ def __init__(self, in_dim, bottleneck_dim=128, global_context_att=False): @@ -81,14 +83,15 @@ class ASTP(nn.Module): # need to transpose inputs. if global_context_att: self.linear1 = nn.Conv1d( - in_dim * 3, bottleneck_dim, - kernel_size=1) # equals W and b in the paper + in_dim * 3, bottleneck_dim, kernel_size=1 + ) # equals W and b in the paper else: self.linear1 = nn.Conv1d( - in_dim, bottleneck_dim, - kernel_size=1) # equals W and b in the paper - self.linear2 = nn.Conv1d(bottleneck_dim, in_dim, - kernel_size=1) # equals V and k in the paper + in_dim, bottleneck_dim, kernel_size=1 + ) # equals W and b in the paper + self.linear2 = nn.Conv1d( + bottleneck_dim, in_dim, kernel_size=1 + ) # equals V and k in the paper def forward(self, x): """ @@ -102,17 +105,15 @@ class ASTP(nn.Module): if self.global_context_att: context_mean = torch.mean(x, dim=-1, keepdim=True).expand_as(x) - context_std = torch.sqrt( - torch.var(x, dim=-1, keepdim=True) + 1e-10).expand_as(x) + context_std = torch.sqrt(torch.var(x, dim=-1, keepdim=True) + 1e-10).expand_as(x) x_in = torch.cat((x, context_mean, context_std), dim=1) else: x_in = x # DON'T use ReLU here! ReLU may be hard to converge. - alpha = torch.tanh( - self.linear1(x_in)) # alpha = F.relu(self.linear1(x_in)) + alpha = torch.tanh(self.linear1(x_in)) # alpha = F.relu(self.linear1(x_in)) alpha = torch.softmax(self.linear2(alpha), dim=2) mean = torch.sum(alpha * x, dim=2) - var = torch.sum(alpha * (x ** 2), dim=2) - mean ** 2 + var = torch.sum(alpha * (x**2), dim=2) - mean**2 std = torch.sqrt(var.clamp(min=1e-10)) return torch.cat([mean, std], dim=1) diff --git a/funasr/models/whisper_lid/eres2net/simple_avg.py b/funasr/models/whisper_lid/eres2net/simple_avg.py index 4fb4c0aa9..9dbe7b398 100644 --- a/funasr/models/whisper_lid/eres2net/simple_avg.py +++ b/funasr/models/whisper_lid/eres2net/simple_avg.py @@ -3,6 +3,7 @@ import torch from funasr.models.encoder.abs_encoder import AbsEncoder from funasr.modules.nets_utils import make_pad_mask + class SimpleAvg(AbsEncoder): def __init__(self, feat_dim): super(SimpleAvg, self).__init__() @@ -14,4 +15,4 @@ class SimpleAvg(AbsEncoder): return avg_x def output_size(self) -> int: - return self.feat_dim \ No newline at end of file + return self.feat_dim diff --git a/funasr/models/whisper_lid/lid_predictor.py b/funasr/models/whisper_lid/lid_predictor.py index 5e042d245..d5b0d37ac 100644 --- a/funasr/models/whisper_lid/lid_predictor.py +++ b/funasr/models/whisper_lid/lid_predictor.py @@ -1,25 +1,31 @@ from funasr.register import tables -from funasr.models.whisper_lid.eres2net.ResNet import ERes2Net, BasicBlockERes2Net, BasicBlockERes2Net_diff_AFF +from funasr.models.whisper_lid.eres2net.ResNet import ( + ERes2Net, + BasicBlockERes2Net, + BasicBlockERes2Net_diff_AFF, +) @tables.register("lid_predictor_classes", "LidPredictor") class LidPredictor(ERes2Net): - def __init__(self, - block=BasicBlockERes2Net, - block_fuse=BasicBlockERes2Net_diff_AFF, - num_blocks=[3, 4, 6, 3], - m_channels=32, - feat_dim=80, - embedding_size=192, - pooling_func='TSTP', - two_emb_layer=False): + def __init__( + self, + block=BasicBlockERes2Net, + block_fuse=BasicBlockERes2Net_diff_AFF, + num_blocks=[3, 4, 6, 3], + m_channels=32, + feat_dim=80, + embedding_size=192, + pooling_func="TSTP", + two_emb_layer=False, + ): super(LidPredictor, self).__init__( - block=block, - block_fuse=block_fuse, - num_blocks=num_blocks, - m_channels=m_channels, - feat_dim=feat_dim, - embedding_size=embedding_size, - pooling_func=pooling_func, - two_emb_layer=two_emb_layer - ) \ No newline at end of file + block=block, + block_fuse=block_fuse, + num_blocks=num_blocks, + m_channels=m_channels, + feat_dim=feat_dim, + embedding_size=embedding_size, + pooling_func=pooling_func, + two_emb_layer=two_emb_layer, + ) diff --git a/funasr/models/whisper_lid/model.py b/funasr/models/whisper_lid/model.py index 6ffb43a6a..0701f6177 100644 --- a/funasr/models/whisper_lid/model.py +++ b/funasr/models/whisper_lid/model.py @@ -22,7 +22,6 @@ from funasr.register import tables class OpenAIWhisperModel(nn.Module): """CTC-attention hybrid Encoder-Decoder model""" - def __init__( self, specaug: str = None, @@ -71,14 +70,12 @@ class OpenAIWhisperModel(nn.Module): decoder_class = tables.decoder_classes.get(decoder) decoder = decoder_class(decoder_conf) if ctc_weight > 0.0: - + if ctc_conf is None: ctc_conf = {} - - ctc = CTC( - odim=vocab_size, encoder_output_size=encoder_output_size, **ctc_conf - ) - + + ctc = CTC(odim=vocab_size, encoder_output_size=encoder_output_size, **ctc_conf) + self.blank_id = blank_id self.sos = sos if sos is not None else vocab_size - 1 self.eos = eos if eos is not None else vocab_size - 1 @@ -102,7 +99,7 @@ class OpenAIWhisperModel(nn.Module): self.decoder = None else: self.decoder = decoder - + self.criterion_att = LabelSmoothingLoss( size=vocab_size, padding_idx=ignore_id, @@ -120,14 +117,14 @@ class OpenAIWhisperModel(nn.Module): self.ctc = None else: self.ctc = ctc - + self.share_embedding = share_embedding if self.share_embedding: self.decoder.embed = None - + self.length_normalized_loss = length_normalized_loss self.beam_search = None - + def forward( self, speech: torch.Tensor, @@ -149,30 +146,30 @@ class OpenAIWhisperModel(nn.Module): text_lengths = text_lengths[:, 0] if len(speech_lengths.size()) > 1: speech_lengths = speech_lengths[:, 0] - + batch_size = speech.shape[0] - + # 1. Encoder encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) intermediate_outs = None if isinstance(encoder_out, tuple): intermediate_outs = encoder_out[1] encoder_out = encoder_out[0] - + loss_att, acc_att, cer_att, wer_att = None, None, None, None loss_ctc, cer_ctc = None, None stats = dict() - + # decoder: CTC branch if self.ctc_weight != 0.0: loss_ctc, cer_ctc = self._calc_ctc_loss( encoder_out, encoder_out_lens, text, text_lengths ) - + # Collect CTC branch stats stats["loss_ctc"] = loss_ctc.detach() if loss_ctc is not None else None stats["cer_ctc"] = cer_ctc - + # Intermediate CTC (optional) loss_interctc = 0.0 if self.interctc_weight != 0.0 and intermediate_outs is not None: @@ -183,25 +180,23 @@ class OpenAIWhisperModel(nn.Module): intermediate_out, encoder_out_lens, text, text_lengths ) loss_interctc = loss_interctc + loss_ic - + # Collect Intermedaite CTC stats stats["loss_interctc_layer{}".format(layer_idx)] = ( loss_ic.detach() if loss_ic is not None else None ) stats["cer_interctc_layer{}".format(layer_idx)] = cer_ic - + loss_interctc = loss_interctc / len(intermediate_outs) - + # calculate whole encoder loss - loss_ctc = ( - 1 - self.interctc_weight - ) * loss_ctc + self.interctc_weight * loss_interctc - + loss_ctc = (1 - self.interctc_weight) * loss_ctc + self.interctc_weight * loss_interctc + # decoder: Attention decoder branch loss_att, acc_att, cer_att, wer_att = self._calc_att_loss( encoder_out, encoder_out_lens, text, text_lengths ) - + # 3. CTC-Att loss definition if self.ctc_weight == 0.0: loss = loss_att @@ -209,25 +204,27 @@ class OpenAIWhisperModel(nn.Module): loss = loss_ctc else: loss = self.ctc_weight * loss_ctc + (1 - self.ctc_weight) * loss_att - + # Collect Attn branch stats stats["loss_att"] = loss_att.detach() if loss_att is not None else None stats["acc"] = acc_att stats["cer"] = cer_att stats["wer"] = wer_att - + # Collect total loss stats stats["loss"] = torch.clone(loss.detach()) - + # force_gatherable: to-device and to-tensor if scalar for DataParallel if self.length_normalized_loss: batch_size = int((text_lengths + 1).sum()) loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) return loss, stats, weight - def encode( - self, speech: torch.Tensor, speech_lengths: torch.Tensor, **kwargs, + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + **kwargs, ) -> Tuple[torch.Tensor, torch.Tensor]: """Frontend + Encoder. Note that this method is used by asr_inference.py Args: @@ -240,30 +237,28 @@ class OpenAIWhisperModel(nn.Module): # Data augmentation if self.specaug is not None and self.training: speech, speech_lengths = self.specaug(speech, speech_lengths) - + # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN if self.normalize is not None: speech, speech_lengths = self.normalize(speech, speech_lengths) - + # Forward encoder # feats: (Batch, Length, Dim) # -> encoder_out: (Batch, Length2, Dim2) if self.encoder.interctc_use_conditioning: - encoder_out, encoder_out_lens, _ = self.encoder( - speech, speech_lengths, ctc=self.ctc - ) + encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths, ctc=self.ctc) else: encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths) intermediate_outs = None if isinstance(encoder_out, tuple): intermediate_outs = encoder_out[1] encoder_out = encoder_out[0] - + if intermediate_outs is not None: return (encoder_out, intermediate_outs), encoder_out_lens - + return encoder_out, encoder_out_lens - + def _calc_att_loss( self, encoder_out: torch.Tensor, @@ -273,12 +268,10 @@ class OpenAIWhisperModel(nn.Module): ): ys_in_pad, ys_out_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id) ys_in_lens = ys_pad_lens + 1 - + # 1. Forward decoder - decoder_out, _ = self.decoder( - encoder_out, encoder_out_lens, ys_in_pad, ys_in_lens - ) - + decoder_out, _ = self.decoder(encoder_out, encoder_out_lens, ys_in_pad, ys_in_lens) + # 2. Compute attention loss loss_att = self.criterion_att(decoder_out, ys_out_pad) acc_att = th_accuracy( @@ -286,16 +279,16 @@ class OpenAIWhisperModel(nn.Module): ys_out_pad, ignore_label=self.ignore_id, ) - + # Compute cer/wer using attention-decoder if self.training or self.error_calculator is None: cer_att, wer_att = None, None else: ys_hat = decoder_out.argmax(dim=-1) cer_att, wer_att = self.error_calculator(ys_hat.cpu(), ys_pad.cpu()) - + return loss_att, acc_att, cer_att, wer_att - + def _calc_ctc_loss( self, encoder_out: torch.Tensor, @@ -305,41 +298,39 @@ class OpenAIWhisperModel(nn.Module): ): # Calc CTC loss loss_ctc = self.ctc(encoder_out, encoder_out_lens, ys_pad, ys_pad_lens) - + # Calc CER using CTC cer_ctc = None if not self.training and self.error_calculator is not None: ys_hat = self.ctc.argmax(encoder_out).data cer_ctc = self.error_calculator(ys_hat.cpu(), ys_pad.cpu(), is_ctc=True) return loss_ctc, cer_ctc - - def init_beam_search(self, - **kwargs, - ): + + def init_beam_search( + self, + **kwargs, + ): from funasr.models.transformer.search import BeamSearch from funasr.models.transformer.scorers.ctc import CTCPrefixScorer from funasr.models.transformer.scorers.length_bonus import LengthBonus - + # 1. Build ASR model scorers = {} - + if self.ctc != None: ctc = CTCPrefixScorer(ctc=self.ctc, eos=self.eos) - scorers.update( - ctc=ctc - ) + scorers.update(ctc=ctc) token_list = kwargs.get("token_list") scorers.update( decoder=self.decoder, length_bonus=LengthBonus(len(token_list)), ) - # 3. Build ngram model # ngram is not supported now ngram = None scorers["ngram"] = ngram - + weights = dict( decoder=1.0 - kwargs.get("decoding_ctc_weight", 0.5), ctc=kwargs.get("decoding_ctc_weight", 0.5), @@ -359,19 +350,20 @@ class OpenAIWhisperModel(nn.Module): ) self.beam_search = beam_search - - def inference(self, - data_in, - data_lengths=None, - key: list=None, - tokenizer=None, - frontend=None, - **kwargs, - ): - + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): + if kwargs.get("batch_size", 1) > 1: raise NotImplementedError("batch decoding is not implemented") - + # init beamsearch if self.beam_search is None: logging.info("enable beam_search") @@ -379,7 +371,9 @@ class OpenAIWhisperModel(nn.Module): self.nbest = kwargs.get("nbest", 1) meta_data = {} - if isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank": # fbank + if ( + isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank" + ): # fbank speech, speech_lengths = data_in, data_lengths if len(speech.shape) < 3: speech = speech[None, :, :] @@ -388,16 +382,23 @@ class OpenAIWhisperModel(nn.Module): else: # extract fbank feats time1 = time.perf_counter() - audio_sample_list = load_audio_text_image_video(data_in, fs=frontend.fs, audio_fs=kwargs.get("fs", 16000), - data_type=kwargs.get("data_type", "sound"), - tokenizer=tokenizer) + audio_sample_list = load_audio_text_image_video( + data_in, + fs=frontend.fs, + audio_fs=kwargs.get("fs", 16000), + data_type=kwargs.get("data_type", "sound"), + tokenizer=tokenizer, + ) time2 = time.perf_counter() meta_data["load_data"] = f"{time2 - time1:0.3f}" - speech, speech_lengths = extract_fbank(audio_sample_list, data_type=kwargs.get("data_type", "sound"), - frontend=frontend) + speech, speech_lengths = extract_fbank( + audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend + ) time3 = time.perf_counter() meta_data["extract_feat"] = f"{time3 - time2:0.3f}" - meta_data["batch_data_time"] = speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 + meta_data["batch_data_time"] = ( + speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 + ) speech = speech.to(device=kwargs["device"]) speech_lengths = speech_lengths.to(device=kwargs["device"]) @@ -405,14 +406,15 @@ class OpenAIWhisperModel(nn.Module): encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) if isinstance(encoder_out, tuple): encoder_out = encoder_out[0] - + # c. Passed the encoder result and the beam search nbest_hyps = self.beam_search( - x=encoder_out[0], maxlenratio=kwargs.get("maxlenratio", 0.0), minlenratio=kwargs.get("minlenratio", 0.0) + x=encoder_out[0], + maxlenratio=kwargs.get("maxlenratio", 0.0), + minlenratio=kwargs.get("minlenratio", 0.0), ) - - nbest_hyps = nbest_hyps[: self.nbest] + nbest_hyps = nbest_hyps[: self.nbest] results = [] b, n, d = encoder_out.size() @@ -424,29 +426,33 @@ class OpenAIWhisperModel(nn.Module): if not hasattr(self, "writer"): self.writer = DatadirWriter(kwargs.get("output_dir")) ibest_writer = self.writer[f"{nbest_idx + 1}best_recog"] - + # remove sos/eos and get results last_pos = -1 if isinstance(hyp.yseq, list): token_int = hyp.yseq[1:last_pos] else: token_int = hyp.yseq[1:last_pos].tolist() - + # remove blank symbol id, which is assumed to be 0 - token_int = list(filter(lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int)) - + token_int = list( + filter( + lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int + ) + ) + # Change integer-ids to tokens token = tokenizer.ids2tokens(token_int) text = tokenizer.tokens2text(token) - + text_postprocessed, _ = postprocess_utils.sentence_postprocess(token) result_i = {"key": key[i], "token": token, "text": text_postprocessed} results.append(result_i) - + if ibest_writer is not None: ibest_writer["token"][key[i]] = " ".join(token) ibest_writer["text"][key[i]] = text_postprocessed - + return results, meta_data @@ -455,18 +461,18 @@ class OpenAIWhisperLIDModel(nn.Module): """WhisperEncoder and EResNet based LID Model""" def __init__( - self, - vocab_size: int, - specaug: str = None, - specaug_conf: dict = None, - encoder: str = None, - encoder_conf: dict = None, - lid_predictor: str = None, - lid_predictor_conf: dict = None, - proj_dim: int = None, - clip_frames: int = None, - random_clip: bool = False, - **kwargs, + self, + vocab_size: int, + specaug: str = None, + specaug_conf: dict = None, + encoder: str = None, + encoder_conf: dict = None, + lid_predictor: str = None, + lid_predictor_conf: dict = None, + proj_dim: int = None, + clip_frames: int = None, + random_clip: bool = False, + **kwargs, ): super().__init__() if specaug is not None: @@ -477,7 +483,7 @@ class OpenAIWhisperLIDModel(nn.Module): lid_predictor_class = tables.lid_predictor_classes.get(lid_predictor) lid_predictor = lid_predictor_class(**lid_predictor_conf) if encoder.output_size() != proj_dim: - self.proj_layer = torch.nn.Linear(encoder.output_size(), proj_dim) + self.proj_layer = torch.nn.Linear(encoder.output_size(), proj_dim) else: self.proj_layer = None self.output_layer = torch.nn.Linear(lid_predictor.output_size(), vocab_size) @@ -498,23 +504,32 @@ class OpenAIWhisperLIDModel(nn.Module): if not hasattr(self.encoder, "interctc_use_conditioning"): self.encoder.interctc_use_conditioning = False - def forward(self, - speech: torch.Tensor, # may be padding - speech_lengths: torch.Tensor, # actual length - lid: torch.Tensor, # lid label, (batch_size, 1) - lid_lengths: torch.Tensor, - ): + def forward( + self, + speech: torch.Tensor, # may be padding + speech_lengths: torch.Tensor, # actual length + lid: torch.Tensor, # lid label, (batch_size, 1) + lid_lengths: torch.Tensor, + ): assert lid.shape[1] == 1 batch_size = speech.shape[0] encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) # re-generate encoder_out if self.clip_frames is None: - reduced_encoder_out = torch.zeros(batch_size, encoder_out_lens.max(), encoder_out.shape[-1]).to(encoder_out.dtype).to(encoder_out.device) + reduced_encoder_out = ( + torch.zeros(batch_size, encoder_out_lens.max(), encoder_out.shape[-1]) + .to(encoder_out.dtype) + .to(encoder_out.device) + ) for i, enc_length in enumerate(encoder_out_lens): reduced_encoder_out[i, :enc_length] = encoder_out[i, :enc_length] else: - reduced_encoder_out = torch.zeros(batch_size, self.clip_frames, encoder_out.shape[-1]).to(encoder_out.dtype).to(encoder_out.device) + reduced_encoder_out = ( + torch.zeros(batch_size, self.clip_frames, encoder_out.shape[-1]) + .to(encoder_out.dtype) + .to(encoder_out.device) + ) if self.random_clip: for i, enc_length in enumerate(encoder_out_lens): if enc_length <= self.clip_frames: @@ -523,7 +538,9 @@ class OpenAIWhisperLIDModel(nn.Module): else: max_start_index = enc_length.item() - self.clip_frames start_index = np.random.randint(0, max_start_index + 1) - reduced_encoder_out[i, :self.clip_frames] = encoder_out[i, start_index:start_index + self.clip_frames] + reduced_encoder_out[i, : self.clip_frames] = encoder_out[ + i, start_index : start_index + self.clip_frames + ] encoder_out_lens[i] = self.clip_frames else: for i, enc_length in enumerate(encoder_out_lens): @@ -548,7 +565,7 @@ class OpenAIWhisperLIDModel(nn.Module): return loss, stats, weight def encode( - self, speech: torch.Tensor, speech_lengths: torch.Tensor + self, speech: torch.Tensor, speech_lengths: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: """Frontend + Encoder. Note that this method is used by asr_inference.py Args: @@ -573,9 +590,7 @@ class OpenAIWhisperLIDModel(nn.Module): # feats: (Batch, Length, Dim) # -> encoder_out: (Batch, Length2, Dim2) if self.encoder.interctc_use_conditioning: - encoder_out, encoder_out_lens, _ = self.encoder( - speech, speech_lengths, ctc=self.ctc - ) + encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths, ctc=self.ctc) else: encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths) intermediate_outs = None @@ -588,20 +603,23 @@ class OpenAIWhisperLIDModel(nn.Module): return encoder_out, encoder_out_lens - def inference(self, - data_in, - data_lengths=None, - key: list = None, - tokenizer=None, - frontend=None, - **kwargs, - ): + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): if kwargs.get("batch_size", 1) > 1: raise NotImplementedError("batch decoding is not implemented") meta_data = {} - if isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank": # fbank + if ( + isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank" + ): # fbank speech, speech_lengths = data_in, data_lengths if len(speech.shape) < 3: speech = speech[None, :, :] @@ -610,39 +628,60 @@ class OpenAIWhisperLIDModel(nn.Module): else: # extract fbank feats time1 = time.perf_counter() - audio_sample_list = load_audio_text_image_video(data_in, fs=frontend.fs, audio_fs=kwargs.get("fs", 16000), - data_type=kwargs.get("data_type", "sound"), - tokenizer=tokenizer) + audio_sample_list = load_audio_text_image_video( + data_in, + fs=frontend.fs, + audio_fs=kwargs.get("fs", 16000), + data_type=kwargs.get("data_type", "sound"), + tokenizer=tokenizer, + ) time2 = time.perf_counter() meta_data["load_data"] = f"{time2 - time1:0.3f}" - speech, speech_lengths = extract_fbank(audio_sample_list, data_type=kwargs.get("data_type", "sound"), - frontend=frontend) + speech, speech_lengths = extract_fbank( + audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend + ) time3 = time.perf_counter() meta_data["extract_feat"] = f"{time3 - time2:0.3f}" - meta_data["batch_data_time"] = speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 + meta_data["batch_data_time"] = ( + speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 + ) speech = speech.to(device=kwargs["device"]) speech_lengths = speech_lengths.to(device=kwargs["device"]) # Encoder - enc, enc_out_lens = self.encode(speech, speech_lengths) + enc, enc_out_lens = self.encode(speech, speech_lengths) inference_clip_length = kwargs.get("inference_clip_length", None) if self.clip_frames is not None: if inference_clip_length is None: - reduced_enc = torch.zeros(enc.shape[0], self.clip_frames, enc.shape[-1]).to(enc.dtype).to(enc.device) + reduced_enc = ( + torch.zeros(enc.shape[0], self.clip_frames, enc.shape[-1]) + .to(enc.dtype) + .to(enc.device) + ) for i, enc_length in enumerate(enc_out_lens): enc_length = self.clip_frames if enc_length >= self.clip_frames else enc_length reduced_enc[i, :enc_length] = enc[i, :enc_length] enc_out_lens[i] = enc_length else: assert inference_clip_length > 0, "inference_clip_length must be larger than 0" - reduced_enc = torch.zeros(enc.shape[0], inference_clip_length, enc.shape[-1]).to(enc.dtype).to(enc.device) + reduced_enc = ( + torch.zeros(enc.shape[0], inference_clip_length, enc.shape[-1]) + .to(enc.dtype) + .to(enc.device) + ) for i, enc_length in enumerate(enc_out_lens): - enc_length = inference_clip_length if enc_length >= inference_clip_length else enc_length + enc_length = ( + inference_clip_length if enc_length >= inference_clip_length else enc_length + ) reduced_enc[i, :enc_length] = enc[i, :enc_length] enc_out_lens[i] = enc_length else: - reduced_enc = torch.zeros(enc.shape[0], enc_out_lens.max(), enc.shape[-1]).to(enc.dtype).to(enc.device) + reduced_enc = ( + torch.zeros(enc.shape[0], enc_out_lens.max(), enc.shape[-1]) + .to(enc.dtype) + .to(enc.device) + ) for i, enc_length in enumerate(enc_out_lens): reduced_enc[i, :enc_length] = enc[i, :enc_length] @@ -662,4 +701,4 @@ class OpenAIWhisperLIDModel(nn.Module): results = [{"key": key[0], "lid": predicted_lid}] - return results, meta_data \ No newline at end of file + return results, meta_data diff --git a/funasr/models/xvector/e2e_sv.py b/funasr/models/xvector/e2e_sv.py index fce0dfea3..56c9e3b10 100644 --- a/funasr/models/xvector/e2e_sv.py +++ b/funasr/models/xvector/e2e_sv.py @@ -43,17 +43,17 @@ class ESPnetSVModel(FunASRModel): """CTC-attention hybrid Encoder-Decoder model""" def __init__( - self, - vocab_size: int, - token_list: Union[Tuple[str, ...], List[str]], - frontend: Optional[AbsFrontend], - specaug: Optional[AbsSpecAug], - normalize: Optional[AbsNormalize], - preencoder: Optional[AbsPreEncoder], - encoder: AbsEncoder, - postencoder: Optional[AbsPostEncoder], - pooling_layer: torch.nn.Module, - decoder: AbsDecoder, + self, + vocab_size: int, + token_list: Union[Tuple[str, ...], List[str]], + frontend: Optional[AbsFrontend], + specaug: Optional[AbsSpecAug], + normalize: Optional[AbsNormalize], + preencoder: Optional[AbsPreEncoder], + encoder: AbsEncoder, + postencoder: Optional[AbsPostEncoder], + pooling_layer: torch.nn.Module, + decoder: AbsDecoder, ): super().__init__() @@ -71,11 +71,11 @@ class ESPnetSVModel(FunASRModel): self.decoder = decoder def forward( - self, - speech: torch.Tensor, - speech_lengths: torch.Tensor, - text: torch.Tensor, - text_lengths: torch.Tensor, + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + text: torch.Tensor, + text_lengths: torch.Tensor, ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]: """Frontend + Encoder + Decoder + Calc loss Args: @@ -87,10 +87,7 @@ class ESPnetSVModel(FunASRModel): assert text_lengths.dim() == 1, text_lengths.shape # Check that batch_size is unified assert ( - speech.shape[0] - == speech_lengths.shape[0] - == text.shape[0] - == text_lengths.shape[0] + speech.shape[0] == speech_lengths.shape[0] == text.shape[0] == text_lengths.shape[0] ), (speech.shape, speech_lengths.shape, text.shape, text_lengths.shape) batch_size = speech.shape[0] @@ -139,9 +136,7 @@ class ESPnetSVModel(FunASRModel): loss_interctc = loss_interctc / len(intermediate_outs) # calculate whole encoder loss - loss_ctc = ( - 1 - self.interctc_weight - ) * loss_ctc + self.interctc_weight * loss_interctc + loss_ctc = (1 - self.interctc_weight) * loss_ctc + self.interctc_weight * loss_interctc if self.use_transducer_decoder: # 2a. Transducer decoder branch @@ -196,11 +191,11 @@ class ESPnetSVModel(FunASRModel): return loss, stats, weight def collect_feats( - self, - speech: torch.Tensor, - speech_lengths: torch.Tensor, - text: torch.Tensor, - text_lengths: torch.Tensor, + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + text: torch.Tensor, + text_lengths: torch.Tensor, ) -> Dict[str, torch.Tensor]: if self.extract_feats_in_collect_stats: feats, feats_lengths = self._extract_feats(speech, speech_lengths) @@ -215,7 +210,7 @@ class ESPnetSVModel(FunASRModel): return {"feats": feats, "feats_lengths": feats_lengths} def encode( - self, speech: torch.Tensor, speech_lengths: torch.Tensor + self, speech: torch.Tensor, speech_lengths: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: """Frontend + Encoder. Note that this method is used by asr_inference.py Args: @@ -244,14 +239,12 @@ class ESPnetSVModel(FunASRModel): # Post-encoder, e.g. NLU if self.postencoder is not None: - encoder_out, encoder_out_lens = self.postencoder( - encoder_out, encoder_out_lens - ) + encoder_out, encoder_out_lens = self.postencoder(encoder_out, encoder_out_lens) return encoder_out, encoder_out_lens def _extract_feats( - self, speech: torch.Tensor, speech_lengths: torch.Tensor + self, speech: torch.Tensor, speech_lengths: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: assert speech_lengths.dim() == 1, speech_lengths.shape @@ -267,4 +260,4 @@ class ESPnetSVModel(FunASRModel): else: # No frontend and no feature extract feats, feats_lengths = speech, speech_lengths - return feats, feats_lengths \ No newline at end of file + return feats, feats_lengths diff --git a/funasr/optimizers/__init__.py b/funasr/optimizers/__init__.py index a1a57a54a..66e2e9d04 100644 --- a/funasr/optimizers/__init__.py +++ b/funasr/optimizers/__init__.py @@ -14,4 +14,4 @@ optim_classes = dict( lbfgs=torch.optim.LBFGS, rmsprop=torch.optim.RMSprop, rprop=torch.optim.Rprop, -) \ No newline at end of file +) diff --git a/funasr/optimizers/fairseq_adam.py b/funasr/optimizers/fairseq_adam.py index 9bdd0f8c8..32fae6a25 100644 --- a/funasr/optimizers/fairseq_adam.py +++ b/funasr/optimizers/fairseq_adam.py @@ -38,13 +38,13 @@ class FairseqAdam(torch.optim.Optimizer): """ def __init__( - self, - params, - lr=1e-3, - adam_betas=(0.9, 0.999), - adam_eps=1e-8, - weight_decay=0, - amsgrad=False, + self, + params, + lr=1e-3, + adam_betas=(0.9, 0.999), + adam_eps=1e-8, + weight_decay=0, + amsgrad=False, ): defaults = dict( lr=lr, betas=adam_betas, eps=adam_eps, weight_decay=weight_decay, amsgrad=amsgrad @@ -104,9 +104,7 @@ class FairseqAdam(torch.optim.Optimizer): state["exp_avg"] = state["exp_avg"].to(p_data_fp32) state["exp_avg_sq"] = state["exp_avg_sq"].to(p_data_fp32) if amsgrad: - state["max_exp_avg_sq"] = state["max_exp_avg_sq"].to( - p_data_fp32 - ) + state["max_exp_avg_sq"] = state["max_exp_avg_sq"].to(p_data_fp32) exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"] if amsgrad: @@ -131,9 +129,7 @@ class FairseqAdam(torch.optim.Optimizer): step_size = group["lr"] * math.sqrt(bias_correction2) / bias_correction1 if group["weight_decay"] != 0: - p_data_fp32.add_( - p_data_fp32, alpha=-group["weight_decay"] * group["lr"] - ) + p_data_fp32.add_(p_data_fp32, alpha=-group["weight_decay"] * group["lr"]) p_data_fp32.addcdiv_(exp_avg, denom, value=-step_size) diff --git a/funasr/register.py b/funasr/register.py index 45e2a858c..384bc5870 100644 --- a/funasr/register.py +++ b/funasr/register.py @@ -3,6 +3,7 @@ import inspect from dataclasses import dataclass import re + @dataclass class RegisterTables: model_classes = {} @@ -24,7 +25,7 @@ class RegisterTables: print("\ntables: \n") fields = vars(self) for classes_key, classes_dict in fields.items(): - + flag = True if key is not None: flag = key in classes_key @@ -37,27 +38,30 @@ class RegisterTables: metas.sort(key=lambda x: x[0]) data = [headers] + metas col_widths = [max(len(str(item)) for item in col) for col in zip(*data)] - - for row in data: - print("| " + " | ".join(str(item).ljust(width) for item, width in zip(row, col_widths)) + " |") - print("\n") + for row in data: + print( + "| " + + " | ".join(str(item).ljust(width) for item, width in zip(row, col_widths)) + + " |" + ) + print("\n") def register(self, register_tables_key: str, key=None): def decorator(target_class): - + if not hasattr(self, register_tables_key): setattr(self, register_tables_key, {}) logging.info("new registry table has been added: {}".format(register_tables_key)) - + registry = getattr(self, register_tables_key) registry_key = key if key is not None else target_class.__name__ - + # assert not registry_key in registry, "(key: {} / class: {}) has been registered already，in {}".format( # registry_key, target_class, register_tables_key) - + registry[registry_key] = target_class - + # meta， headers = ["class name", "register name", "class location"] register_tables_key_meta = register_tables_key + "_meta" if not hasattr(self, register_tables_key_meta): @@ -66,14 +70,18 @@ class RegisterTables: # doc = target_class.__doc__ class_file = inspect.getfile(target_class) class_line = inspect.getsourcelines(target_class)[1] - pattern = r'^.+/funasr/' - class_file = re.sub(pattern, 'funasr/', class_file) + pattern = r"^.+/funasr/" + class_file = re.sub(pattern, "funasr/", class_file) # meata_data = [f"{target_class.__name__}", f"{class_file}:{class_line}"] - meata_data = [f"{registry_key}", f"{target_class.__name__}", f"{class_file}:{class_line}"] + meata_data = [ + f"{registry_key}", + f"{target_class.__name__}", + f"{class_file}:{class_line}", + ] registry_meta[registry_key] = meata_data # print(f"Registering class: {class_file}:{class_line} - {target_class.__name__} as {registry_key}") return target_class - + return decorator @@ -81,4 +89,3 @@ tables = RegisterTables() import funasr - diff --git a/funasr/schedulers/lambdalr_cus.py b/funasr/schedulers/lambdalr_cus.py index 5aad049de..19ad7a8a7 100644 --- a/funasr/schedulers/lambdalr_cus.py +++ b/funasr/schedulers/lambdalr_cus.py @@ -1,7 +1,7 @@ - import torch from torch.optim.lr_scheduler import _LRScheduler + class CustomLambdaLR(_LRScheduler): def __init__(self, optimizer, warmup_steps, last_epoch=-1): self.warmup_steps = warmup_steps @@ -10,13 +10,12 @@ class CustomLambdaLR(_LRScheduler): def get_lr(self): if self.last_epoch < self.warmup_steps: return [ - base_lr * min(self.last_epoch / self.warmup_steps, 1) - for base_lr in self.base_lrs + base_lr * min(self.last_epoch / self.warmup_steps, 1) for base_lr in self.base_lrs ] else: return [base_lr for base_lr in self.base_lrs] - - + + class CustomLambdaLR(_LRScheduler): def __init__(self, optimizer, train_config, last_epoch=-1, verbose=False): self.warmup_steps = train_config.warmup_steps @@ -28,5 +27,7 @@ class CustomLambdaLR(_LRScheduler): if step < self.warmup_steps: lr_scale = step / self.warmup_steps else: - lr_scale = max(0.0, 1 - (step - self.warmup_steps) / (self.total_steps - self.warmup_steps)) + lr_scale = max( + 0.0, 1 - (step - self.warmup_steps) / (self.total_steps - self.warmup_steps) + ) return [base_lr * lr_scale for base_lr in self.base_lrs] diff --git a/funasr/schedulers/noam_lr.py b/funasr/schedulers/noam_lr.py index e08fb635e..b33b4bde6 100644 --- a/funasr/schedulers/noam_lr.py +++ b/funasr/schedulers/noam_lr.py @@ -1,4 +1,5 @@ """Noam learning rate scheduler module.""" + from typing import Union import warnings @@ -56,8 +57,6 @@ class NoamLR(_LRScheduler, AbsBatchStepScheduler): def get_lr(self): step_num = self.last_epoch + 1 return [ - lr - * self.model_size**-0.5 - * min(step_num**-0.5, step_num * self.warmup_steps**-1.5) + lr * self.model_size**-0.5 * min(step_num**-0.5, step_num * self.warmup_steps**-1.5) for lr in self.base_lrs ] diff --git a/funasr/schedulers/tri_stage_scheduler.py b/funasr/schedulers/tri_stage_scheduler.py index c442260c4..7fc13b2ff 100644 --- a/funasr/schedulers/tri_stage_scheduler.py +++ b/funasr/schedulers/tri_stage_scheduler.py @@ -14,12 +14,12 @@ from funasr.schedulers.abs_scheduler import AbsBatchStepScheduler class TriStageLR(_LRScheduler, AbsBatchStepScheduler): def __init__( - self, - optimizer: torch.optim.Optimizer, - last_epoch: int = -1, - phase_ratio: Optional[List[float]] = None, - init_lr_scale: float = 0.01, - final_lr_scale: float = 0.01, + self, + optimizer: torch.optim.Optimizer, + last_epoch: int = -1, + phase_ratio: Optional[List[float]] = None, + init_lr_scale: float = 0.01, + final_lr_scale: float = 0.01, ): self.optimizer = optimizer self.last_epoch = last_epoch @@ -42,9 +42,7 @@ class TriStageLR(_LRScheduler, AbsBatchStepScheduler): self.decay_steps = int(self.max_update * self.phase_ratio[2]) self.warmup_rate = ( - (self.peak_lr - self.init_lr) / self.warmup_steps - if self.warmup_steps != 0 - else 0 + (self.peak_lr - self.init_lr) / self.warmup_steps if self.warmup_steps != 0 else 0 ) self.decay_factor = -math.log(self.final_lr_scale) / self.decay_steps diff --git a/funasr/schedulers/warmup_lr.py b/funasr/schedulers/warmup_lr.py index 95ebacaff..847b53ad7 100644 --- a/funasr/schedulers/warmup_lr.py +++ b/funasr/schedulers/warmup_lr.py @@ -1,4 +1,5 @@ """Warm up learning rate scheduler module.""" + from typing import Union import torch @@ -41,8 +42,6 @@ class WarmupLR(_LRScheduler, AbsBatchStepScheduler): def get_lr(self): step_num = self.last_epoch + 1 return [ - lr - * self.warmup_steps**0.5 - * min(step_num**-0.5, step_num * self.warmup_steps**-1.5) + lr * self.warmup_steps**0.5 * min(step_num**-0.5, step_num * self.warmup_steps**-1.5) for lr in self.base_lrs ] diff --git a/funasr/tokenizer/abs_tokenizer.py b/funasr/tokenizer/abs_tokenizer.py index 136be134b..a629e94ff 100644 --- a/funasr/tokenizer/abs_tokenizer.py +++ b/funasr/tokenizer/abs_tokenizer.py @@ -10,24 +10,26 @@ class AbsTokenizer(ABC): @abstractmethod def text2tokens(self, line: str) -> List[str]: raise NotImplementedError - + @abstractmethod def tokens2text(self, tokens: Iterable[str]) -> str: raise NotImplementedError class BaseTokenizer(ABC): - def __init__(self, token_list: Union[Path, str, Iterable[str]] = None, - unk_symbol: str = "", - **kwargs, - ): - + def __init__( + self, + token_list: Union[Path, str, Iterable[str]] = None, + unk_symbol: str = "", + **kwargs, + ): + if token_list is not None: if isinstance(token_list, (Path, str)) and token_list.endswith(".txt"): token_list = Path(token_list) self.token_list_repr = str(token_list) self.token_list: List[str] = [] - + with token_list.open("r", encoding="utf-8") as f: for idx, line in enumerate(f): line = line.rstrip() @@ -36,11 +38,10 @@ class BaseTokenizer(ABC): token_list = Path(token_list) self.token_list_repr = str(token_list) self.token_list: List[str] = [] - - with open(token_list, 'r', encoding='utf-8') as f: + + with open(token_list, "r", encoding="utf-8") as f: self.token_list = json.load(f) - - + else: self.token_list: List[str] = list(token_list) self.token_list_repr = "" @@ -49,46 +50,44 @@ class BaseTokenizer(ABC): break self.token_list_repr += f"{t}, " self.token_list_repr += f"... (NVocab={(len(self.token_list))})" - + self.token2id: Dict[str, int] = {} for i, t in enumerate(self.token_list): if t in self.token2id: raise RuntimeError(f'Symbol "{t}" is duplicated') self.token2id[t] = i - + self.unk_symbol = unk_symbol if self.unk_symbol not in self.token2id: - raise RuntimeError( - f"Unknown symbol '{unk_symbol}' doesn't exist in the token_list" - ) + raise RuntimeError(f"Unknown symbol '{unk_symbol}' doesn't exist in the token_list") self.unk_id = self.token2id[self.unk_symbol] - + def encode(self, text): tokens = self.text2tokens(text) text_ints = self.tokens2ids(tokens) - + return text_ints - + def decode(self, text_ints): token = self.ids2tokens(text_ints) text = self.tokens2text(token) return text - + def get_num_vocabulary_size(self) -> int: return len(self.token_list) - + def ids2tokens(self, integers: Union[np.ndarray, Iterable[int]]) -> List[str]: if isinstance(integers, np.ndarray) and integers.ndim != 1: raise ValueError(f"Must be 1 dim ndarray, but got {integers.ndim}") return [self.token_list[i] for i in integers] - + def tokens2ids(self, tokens: Iterable[str]) -> List[int]: return [self.token2id.get(i, self.unk_id) for i in tokens] - + @abstractmethod def text2tokens(self, line: str) -> List[str]: raise NotImplementedError - + @abstractmethod def tokens2text(self, tokens: Iterable[str]) -> str: - raise NotImplementedError \ No newline at end of file + raise NotImplementedError diff --git a/funasr/tokenizer/build_tokenizer.py b/funasr/tokenizer/build_tokenizer.py index 9d1cdc3b0..d6080f75d 100644 --- a/funasr/tokenizer/build_tokenizer.py +++ b/funasr/tokenizer/build_tokenizer.py @@ -56,6 +56,4 @@ def build_tokenizer( ) else: - raise ValueError( - f"token_mode must be one of bpe, word, char or phn: " f"{token_type}" - ) + raise ValueError(f"token_mode must be one of bpe, word, char or phn: " f"{token_type}") diff --git a/funasr/tokenizer/char_tokenizer.py b/funasr/tokenizer/char_tokenizer.py index 92c6e677b..805ecd007 100644 --- a/funasr/tokenizer/char_tokenizer.py +++ b/funasr/tokenizer/char_tokenizer.py @@ -8,105 +8,106 @@ import re from funasr.tokenizer.abs_tokenizer import BaseTokenizer from funasr.register import tables + @tables.register("tokenizer_classes", "CharTokenizer") class CharTokenizer(BaseTokenizer): - def __init__( - self, - non_linguistic_symbols: Union[Path, str, Iterable[str]] = None, - space_symbol: str = "", - remove_non_linguistic_symbols: bool = False, - split_with_space: bool = False, - seg_dict: str = None, - **kwargs, - ): - super().__init__(**kwargs) - self.space_symbol = space_symbol - if non_linguistic_symbols is None: - self.non_linguistic_symbols = set() - elif isinstance(non_linguistic_symbols, (Path, str)): - non_linguistic_symbols = Path(non_linguistic_symbols) - try: - with non_linguistic_symbols.open("r", encoding="utf-8") as f: - self.non_linguistic_symbols = set(line.rstrip() for line in f) - except FileNotFoundError: - warnings.warn(f"{non_linguistic_symbols} doesn't exist.") - self.non_linguistic_symbols = set() - else: - self.non_linguistic_symbols = set(non_linguistic_symbols) - self.remove_non_linguistic_symbols = remove_non_linguistic_symbols - self.split_with_space = split_with_space - self.seg_dict = None - seg_dict = seg_dict if seg_dict is not None else kwargs.get("seg_dict_file", None) - if seg_dict is not None: - self.seg_dict = load_seg_dict(seg_dict) - - - def __repr__(self): - return ( - f"{self.__class__.__name__}(" - f'space_symbol="{self.space_symbol}"' - f'non_linguistic_symbols="{self.non_linguistic_symbols}"' - f")" - ) - - def text2tokens(self, line: Union[str, list]) -> List[str]: - - # if self.split_with_space: - - if self.seg_dict is not None: - tokens = line.strip().split(" ") - tokens = seg_tokenize(tokens, self.seg_dict) - else: - tokens = [] - while len(line) != 0: - for w in self.non_linguistic_symbols: - if line.startswith(w): - if not self.remove_non_linguistic_symbols: - tokens.append(line[: len(w)]) - line = line[len(w) :] - break - else: - t = line[0] - if t == " ": - # t = "" - line = line[1:] - continue - tokens.append(t) - line = line[1:] - return tokens - - def tokens2text(self, tokens: Iterable[str]) -> str: - tokens = [t if t != self.space_symbol else " " for t in tokens] - return "".join(tokens) + def __init__( + self, + non_linguistic_symbols: Union[Path, str, Iterable[str]] = None, + space_symbol: str = "", + remove_non_linguistic_symbols: bool = False, + split_with_space: bool = False, + seg_dict: str = None, + **kwargs, + ): + super().__init__(**kwargs) + self.space_symbol = space_symbol + if non_linguistic_symbols is None: + self.non_linguistic_symbols = set() + elif isinstance(non_linguistic_symbols, (Path, str)): + non_linguistic_symbols = Path(non_linguistic_symbols) + try: + with non_linguistic_symbols.open("r", encoding="utf-8") as f: + self.non_linguistic_symbols = set(line.rstrip() for line in f) + except FileNotFoundError: + warnings.warn(f"{non_linguistic_symbols} doesn't exist.") + self.non_linguistic_symbols = set() + else: + self.non_linguistic_symbols = set(non_linguistic_symbols) + self.remove_non_linguistic_symbols = remove_non_linguistic_symbols + self.split_with_space = split_with_space + self.seg_dict = None + seg_dict = seg_dict if seg_dict is not None else kwargs.get("seg_dict_file", None) + if seg_dict is not None: + self.seg_dict = load_seg_dict(seg_dict) + + def __repr__(self): + return ( + f"{self.__class__.__name__}(" + f'space_symbol="{self.space_symbol}"' + f'non_linguistic_symbols="{self.non_linguistic_symbols}"' + f")" + ) + + def text2tokens(self, line: Union[str, list]) -> List[str]: + + # if self.split_with_space: + + if self.seg_dict is not None: + tokens = line.strip().split(" ") + tokens = seg_tokenize(tokens, self.seg_dict) + else: + tokens = [] + while len(line) != 0: + for w in self.non_linguistic_symbols: + if line.startswith(w): + if not self.remove_non_linguistic_symbols: + tokens.append(line[: len(w)]) + line = line[len(w) :] + break + else: + t = line[0] + if t == " ": + # t = "" + line = line[1:] + continue + tokens.append(t) + line = line[1:] + return tokens + + def tokens2text(self, tokens: Iterable[str]) -> str: + tokens = [t if t != self.space_symbol else " " for t in tokens] + return "".join(tokens) def load_seg_dict(seg_dict_file): - seg_dict = {} - assert isinstance(seg_dict_file, str) - with open(seg_dict_file, "r", encoding="utf8") as f: - lines = f.readlines() - for line in lines: - s = line.strip().split() - key = s[0] - value = s[1:] - seg_dict[key] = " ".join(value) - return seg_dict + seg_dict = {} + assert isinstance(seg_dict_file, str) + with open(seg_dict_file, "r", encoding="utf8") as f: + lines = f.readlines() + for line in lines: + s = line.strip().split() + key = s[0] + value = s[1:] + seg_dict[key] = " ".join(value) + return seg_dict + def seg_tokenize(txt, seg_dict): - # pattern = re.compile(r'^[\u4E00-\u9FA50-9]+$') - pattern = re.compile(r"([\u4E00-\u9FA5A-Za-z0-9])") - out_txt = "" - for word in txt: - word = word.lower() - if word in seg_dict: - out_txt += seg_dict[word] + " " - else: - if pattern.match(word): - for char in word: - if char in seg_dict: - out_txt += seg_dict[char] + " " - else: - out_txt += "" + " " - else: - out_txt += "" + " " - return out_txt.strip().split() \ No newline at end of file + # pattern = re.compile(r'^[\u4E00-\u9FA50-9]+$') + pattern = re.compile(r"([\u4E00-\u9FA5A-Za-z0-9])") + out_txt = "" + for word in txt: + word = word.lower() + if word in seg_dict: + out_txt += seg_dict[word] + " " + else: + if pattern.match(word): + for char in word: + if char in seg_dict: + out_txt += seg_dict[char] + " " + else: + out_txt += "" + " " + else: + out_txt += "" + " " + return out_txt.strip().split() diff --git a/funasr/tokenizer/cleaner.py b/funasr/tokenizer/cleaner.py index 1044f504f..882db5c97 100644 --- a/funasr/tokenizer/cleaner.py +++ b/funasr/tokenizer/cleaner.py @@ -1,6 +1,7 @@ from typing import Collection from jaconv import jaconv + # import tacotron_cleaner.cleaners try: diff --git a/funasr/tokenizer/hf_tokenizer.py b/funasr/tokenizer/hf_tokenizer.py index b6043e95b..1ac7ae259 100644 --- a/funasr/tokenizer/hf_tokenizer.py +++ b/funasr/tokenizer/hf_tokenizer.py @@ -1,16 +1,13 @@ - - - from funasr.register import tables + @tables.register("tokenizer_classes", "HuggingfaceTokenizer") def HuggingfaceTokenizer(init_param_path, **kwargs): - try: - from transformers import AutoTokenizer - except: - # print("If you want to use hugging, please `pip install -U transformers`") - pass - tokenizer = AutoTokenizer.from_pretrained(init_param_path) - - return tokenizer + try: + from transformers import AutoTokenizer + except: + # print("If you want to use hugging, please `pip install -U transformers`") + pass + tokenizer = AutoTokenizer.from_pretrained(init_param_path) + return tokenizer diff --git a/funasr/tokenizer/korean_cleaner.py b/funasr/tokenizer/korean_cleaner.py index ee556d42a..eceebe216 100644 --- a/funasr/tokenizer/korean_cleaner.py +++ b/funasr/tokenizer/korean_cleaner.py @@ -19,8 +19,7 @@ class KoreanCleaner: "9": "구", } new_text = "".join( - number_to_kor[char] if char in number_to_kor.keys() else char - for char in text + number_to_kor[char] if char in number_to_kor.keys() else char for char in text ) return new_text @@ -56,9 +55,7 @@ class KoreanCleaner: } new_text = re.sub("[a-z]+", lambda x: str.upper(x.group()), text) new_text = "".join( - upper_alphabet_to_kor[char] - if char in upper_alphabet_to_kor.keys() - else char + upper_alphabet_to_kor[char] if char in upper_alphabet_to_kor.keys() else char for char in new_text ) diff --git a/funasr/tokenizer/phoneme_tokenizer.py b/funasr/tokenizer/phoneme_tokenizer.py index 2820f2111..12b29ec22 100644 --- a/funasr/tokenizer/phoneme_tokenizer.py +++ b/funasr/tokenizer/phoneme_tokenizer.py @@ -238,9 +238,7 @@ class G2pk: """ - def __init__( - self, descritive=False, group_vowels=False, to_syl=False, no_space=False - ): + def __init__(self, descritive=False, group_vowels=False, to_syl=False, no_space=False): self.descritive = descritive self.group_vowels = group_vowels self.to_syl = to_syl diff --git a/funasr/tokenizer/sentencepiece_tokenizer.py b/funasr/tokenizer/sentencepiece_tokenizer.py index 05ca3489a..ff4b3a2ba 100644 --- a/funasr/tokenizer/sentencepiece_tokenizer.py +++ b/funasr/tokenizer/sentencepiece_tokenizer.py @@ -8,11 +8,10 @@ import sentencepiece as spm from funasr.tokenizer.abs_tokenizer import BaseTokenizer from funasr.register import tables + @tables.register("tokenizer_classes", "SentencepiecesTokenizer") class SentencepiecesTokenizer(BaseTokenizer): - def __init__(self, bpemodel: Union[Path, str], - **kwargs - ): + def __init__(self, bpemodel: Union[Path, str], **kwargs): super().__init__(**kwargs) self.bpemodel = str(bpemodel) # NOTE(kamo): @@ -45,4 +44,4 @@ class SentencepiecesTokenizer(BaseTokenizer): def decode(self, line: List[int]): self._build_sentence_piece_processor() - return self.sp.DecodeIds(line) \ No newline at end of file + return self.sp.DecodeIds(line) diff --git a/funasr/tokenizer/token_id_converter.py b/funasr/tokenizer/token_id_converter.py index 1888d7508..c3fc71d5a 100644 --- a/funasr/tokenizer/token_id_converter.py +++ b/funasr/tokenizer/token_id_converter.py @@ -41,9 +41,7 @@ class TokenIDConverter: self.unk_symbol = unk_symbol if self.unk_symbol not in self.token2id: - raise RuntimeError( - f"Unknown symbol '{unk_symbol}' doesn't exist in the token_list" - ) + raise RuntimeError(f"Unknown symbol '{unk_symbol}' doesn't exist in the token_list") self.unk_id = self.token2id[self.unk_symbol] def get_num_vocabulary_size(self) -> int: diff --git a/funasr/tokenizer/whisper_tokenizer.py b/funasr/tokenizer/whisper_tokenizer.py index 0a34d19d1..bd96b832b 100644 --- a/funasr/tokenizer/whisper_tokenizer.py +++ b/funasr/tokenizer/whisper_tokenizer.py @@ -1,46 +1,45 @@ - - from funasr.register import tables + @tables.register("tokenizer_classes", "WhisperTokenizer") def WhisperTokenizer(**kwargs): - try: - from whisper.tokenizer import get_tokenizer - except: - print("Notice: If you want to use whisper, please `pip install -U openai-whisper`") - - language = kwargs.get("language", None) - task = kwargs.get("task", "transcribe") - is_multilingual = kwargs.get("is_multilingual", True) - num_languages = kwargs.get("num_languages", 99) - tokenizer = get_tokenizer( - multilingual=is_multilingual, - num_languages=num_languages, - language=language, - task=task, - ) - - return tokenizer + try: + from whisper.tokenizer import get_tokenizer + except: + print("Notice: If you want to use whisper, please `pip install -U openai-whisper`") + + language = kwargs.get("language", None) + task = kwargs.get("task", "transcribe") + is_multilingual = kwargs.get("is_multilingual", True) + num_languages = kwargs.get("num_languages", 99) + tokenizer = get_tokenizer( + multilingual=is_multilingual, + num_languages=num_languages, + language=language, + task=task, + ) + + return tokenizer @tables.register("tokenizer_classes", "SenseVoiceTokenizer") def SenseVoiceTokenizer(**kwargs): - try: - from funasr.models.sense_voice.whisper_lib.tokenizer import get_tokenizer - except: - print("Notice: If you want to use whisper, please `pip install -U openai-whisper`") - - language = kwargs.get("language", None) - task = kwargs.get("task", None) - is_multilingual = kwargs.get("is_multilingual", True) - num_languages = kwargs.get("num_languages", 8749) - vocab_path = kwargs.get("vocab_path", None) - tokenizer = get_tokenizer( - multilingual=is_multilingual, - num_languages=num_languages, - language=language, - task=task, - vocab_path=vocab_path, - ) - - return tokenizer + try: + from funasr.models.sense_voice.whisper_lib.tokenizer import get_tokenizer + except: + print("Notice: If you want to use whisper, please `pip install -U openai-whisper`") + + language = kwargs.get("language", None) + task = kwargs.get("task", None) + is_multilingual = kwargs.get("is_multilingual", True) + num_languages = kwargs.get("num_languages", 8749) + vocab_path = kwargs.get("vocab_path", None) + tokenizer = get_tokenizer( + multilingual=is_multilingual, + num_languages=num_languages, + language=language, + task=task, + vocab_path=vocab_path, + ) + + return tokenizer diff --git a/funasr/tokenizer/word_tokenizer.py b/funasr/tokenizer/word_tokenizer.py index cbd0673df..c1fd6e039 100644 --- a/funasr/tokenizer/word_tokenizer.py +++ b/funasr/tokenizer/word_tokenizer.py @@ -19,8 +19,7 @@ class WordTokenizer(AbsTokenizer): if not remove_non_linguistic_symbols and non_linguistic_symbols is not None: warnings.warn( - "non_linguistic_symbols is only used " - "when remove_non_linguistic_symbols = True" + "non_linguistic_symbols is only used " "when remove_non_linguistic_symbols = True" ) if non_linguistic_symbols is None: diff --git a/funasr/train_utils/average_nbest_models.py b/funasr/train_utils/average_nbest_models.py index 3413461da..0f0880425 100644 --- a/funasr/train_utils/average_nbest_models.py +++ b/funasr/train_utils/average_nbest_models.py @@ -16,8 +16,7 @@ from collections import OrderedDict from functools import cmp_to_key - -def _get_checkpoint_paths(output_dir: str, last_n: int=5): +def _get_checkpoint_paths(output_dir: str, last_n: int = 5): """ Get the paths of the last 'last_n' checkpoints by parsing filenames in the output directory. @@ -27,7 +26,9 @@ def _get_checkpoint_paths(output_dir: str, last_n: int=5): avg_keep_nbest_models_type = checkpoint["avg_keep_nbest_models_type"] val_step_or_eoch = checkpoint[f"val_{avg_keep_nbest_models_type}_step_or_eoch"] sorted_items = sorted(val_step_or_eoch.items(), key=lambda x: x[1], reverse=True) - sorted_items = sorted_items[:last_n] if avg_keep_nbest_models_type == "acc" else sorted_items[-last_n:] + sorted_items = ( + sorted_items[:last_n] if avg_keep_nbest_models_type == "acc" else sorted_items[-last_n:] + ) checkpoint_paths = [os.path.join(output_dir, key) for key, value in sorted_items[:last_n]] except: print(f"{checkpoint} does not exist, avg the lastet checkpoint.") @@ -36,13 +37,14 @@ def _get_checkpoint_paths(output_dir: str, last_n: int=5): # Filter out checkpoint files and extract epoch numbers checkpoint_files = [f for f in files if f.startswith("model.pt.e")] # Sort files by epoch number in descending order - checkpoint_files.sort(key=lambda x: int(re.search(r'(\d+)', x).group()), reverse=True) + checkpoint_files.sort(key=lambda x: int(re.search(r"(\d+)", x).group()), reverse=True) # Get the last 'last_n' checkpoint paths checkpoint_paths = [os.path.join(output_dir, f) for f in checkpoint_files[:last_n]] return checkpoint_paths + @torch.no_grad() -def average_checkpoints(output_dir: str, last_n: int=5, **kwargs): +def average_checkpoints(output_dir: str, last_n: int = 5, **kwargs): """ Average the last 'last_n' checkpoints' model state_dicts. If a tensor is of type torch.int, perform sum instead of average. @@ -54,7 +56,7 @@ def average_checkpoints(output_dir: str, last_n: int=5, **kwargs): # Load state_dicts from checkpoints for path in checkpoint_paths: if os.path.isfile(path): - state_dicts.append(torch.load(path, map_location='cpu')['state_dict']) + state_dicts.append(torch.load(path, map_location="cpu")["state_dict"]) else: print(f"Checkpoint file {path} not found.") @@ -76,5 +78,5 @@ def average_checkpoints(output_dir: str, last_n: int=5, **kwargs): stacked_tensors = torch.stack(tensors) avg_state_dict[key] = torch.mean(stacked_tensors, dim=0) checkpoint_outpath = os.path.join(output_dir, f"model.pt.avg{last_n}") - torch.save({'state_dict': avg_state_dict}, checkpoint_outpath) - return checkpoint_outpath \ No newline at end of file + torch.save({"state_dict": avg_state_dict}, checkpoint_outpath) + return checkpoint_outpath diff --git a/funasr/train_utils/device_funcs.py b/funasr/train_utils/device_funcs.py index 7919e7d92..dd7fc76a2 100644 --- a/funasr/train_utils/device_funcs.py +++ b/funasr/train_utils/device_funcs.py @@ -8,21 +8,14 @@ import torch def to_device(data, device=None, dtype=None, non_blocking=False, copy=False): """Change the device of object recursively""" if isinstance(data, dict): - return { - k: to_device(v, device, dtype, non_blocking, copy) for k, v in data.items() - } + return {k: to_device(v, device, dtype, non_blocking, copy) for k, v in data.items()} elif dataclasses.is_dataclass(data) and not isinstance(data, type): return type(data)( - *[ - to_device(v, device, dtype, non_blocking, copy) - for v in dataclasses.astuple(data) - ] + *[to_device(v, device, dtype, non_blocking, copy) for v in dataclasses.astuple(data)] ) # maybe namedtuple. I don't know the correct way to judge namedtuple. elif isinstance(data, tuple) and type(data) is not tuple: - return type(data)( - *[to_device(o, device, dtype, non_blocking, copy) for o in data] - ) + return type(data)(*[to_device(o, device, dtype, non_blocking, copy) for o in data]) elif isinstance(data, (list, tuple)): return type(data)(to_device(v, device, dtype, non_blocking, copy) for v in data) elif isinstance(data, np.ndarray): diff --git a/funasr/train_utils/initialize.py b/funasr/train_utils/initialize.py index ee9485048..667585a27 100644 --- a/funasr/train_utils/initialize.py +++ b/funasr/train_utils/initialize.py @@ -19,7 +19,6 @@ def initialize(model: torch.nn.Module, init: str): init: Method of initialization. """ - # weight init for p in model.parameters(): if p.dim() > 1: @@ -40,9 +39,7 @@ def initialize(model: torch.nn.Module, init: str): # reset some modules with default init for m in model.modules(): - if isinstance( - m, (torch.nn.Embedding, torch.nn.LayerNorm, torch.nn.GroupNorm) - ): + if isinstance(m, (torch.nn.Embedding, torch.nn.LayerNorm, torch.nn.GroupNorm)): m.reset_parameters() if hasattr(m, "espnet_initialization_fn"): m.espnet_initialization_fn() @@ -56,4 +53,3 @@ def initialize(model: torch.nn.Module, init: str): model.frontend, "reload_pretrained_parameters", None ): model.frontend.reload_pretrained_parameters() - diff --git a/funasr/train_utils/load_pretrained_model.py b/funasr/train_utils/load_pretrained_model.py index 4f0914ae7..02abfd5e5 100644 --- a/funasr/train_utils/load_pretrained_model.py +++ b/funasr/train_utils/load_pretrained_model.py @@ -9,103 +9,108 @@ import torch.nn import torch.optim import pdb + def filter_state_dict( - dst_state: Dict[str, Union[float, torch.Tensor]], - src_state: Dict[str, Union[float, torch.Tensor]], + dst_state: Dict[str, Union[float, torch.Tensor]], + src_state: Dict[str, Union[float, torch.Tensor]], ): - """Filter name, size mismatch instances between dicts. + """Filter name, size mismatch instances between dicts. - Args: - dst_state: reference state dict for filtering - src_state: target state dict for filtering + Args: + dst_state: reference state dict for filtering + src_state: target state dict for filtering - """ - match_state = {} - for key, value in src_state.items(): - if key in dst_state and (dst_state[key].size() == src_state[key].size()): - match_state[key] = value - else: - if key not in dst_state: - logging.warning( - f"Filter out {key} from pretrained dict" - + " because of name not found in target dict" - ) - else: - logging.warning( - f"Filter out {key} from pretrained dict" - + " because of size mismatch" - + f"({dst_state[key].size()}-{src_state[key].size()})" - ) - return match_state + """ + match_state = {} + for key, value in src_state.items(): + if key in dst_state and (dst_state[key].size() == src_state[key].size()): + match_state[key] = value + else: + if key not in dst_state: + logging.warning( + f"Filter out {key} from pretrained dict" + + " because of name not found in target dict" + ) + else: + logging.warning( + f"Filter out {key} from pretrained dict" + + " because of size mismatch" + + f"({dst_state[key].size()}-{src_state[key].size()})" + ) + return match_state def load_pretrained_model( - path: str, - model: torch.nn.Module, - ignore_init_mismatch: bool=True, - map_location: str = "cpu", - oss_bucket=None, - scope_map=[], - excludes=None, - **kwargs, + path: str, + model: torch.nn.Module, + ignore_init_mismatch: bool = True, + map_location: str = "cpu", + oss_bucket=None, + scope_map=[], + excludes=None, + **kwargs, ): - """Load a model state and set it to the model. + """Load a model state and set it to the model. - Args: - init_param: ::: + Args: + init_param: ::: - Examples: + Examples: - """ - - obj = model - dst_state = obj.state_dict() - - print(f"ckpt: {path}") + """ - if oss_bucket is None: - src_state = torch.load(path, map_location=map_location) - else: - buffer = BytesIO(oss_bucket.get_object(path).read()) - src_state = torch.load(buffer, map_location=map_location) - - src_state = src_state["state_dict"] if "state_dict" in src_state else src_state - src_state = src_state["model_state_dict"] if "model_state_dict" in src_state else src_state - src_state = src_state["model"] if "model" in src_state else src_state - - if isinstance(scope_map, str): - scope_map = scope_map.split(",") - scope_map += ["module.", "None"] - - for k in dst_state.keys(): - - k_src = k + obj = model + dst_state = obj.state_dict() - if scope_map is not None: - src_prefix = "" - dst_prefix = "" - for i in range(0, len(scope_map), 2): - src_prefix = scope_map[i] if scope_map[i].lower() != "none" else "" - dst_prefix = scope_map[i+1] if scope_map[i+1].lower() != "none" else "" - - if dst_prefix == "" and (src_prefix + k) in src_state.keys(): - k_src = src_prefix + k - if not k_src.startswith("module."): - print(f"init param, map: {k} from {k_src} in ckpt") - elif k.startswith(dst_prefix) and k.replace(dst_prefix, src_prefix, 1) in src_state.keys(): - k_src = k.replace(dst_prefix, src_prefix, 1) - if not k_src.startswith("module."): - print(f"init param, map: {k} from {k_src} in ckpt") - - if k_src in src_state.keys(): - if ignore_init_mismatch and dst_state[k].shape != src_state[k_src].shape: - print(f"ignore_init_mismatch:{ignore_init_mismatch}, dst: {k, dst_state[k].shape}, src: {k_src, src_state[k_src].shape}") - else: - dst_state[k] = src_state[k_src] + print(f"ckpt: {path}") + if oss_bucket is None: + src_state = torch.load(path, map_location=map_location) + else: + buffer = BytesIO(oss_bucket.get_object(path).read()) + src_state = torch.load(buffer, map_location=map_location) - else: - print(f"Warning, miss key in ckpt: {k}, mapped: {k_src}") - - flag = obj.load_state_dict(dst_state, strict=True) - # print(flag) + src_state = src_state["state_dict"] if "state_dict" in src_state else src_state + src_state = src_state["model_state_dict"] if "model_state_dict" in src_state else src_state + src_state = src_state["model"] if "model" in src_state else src_state + + if isinstance(scope_map, str): + scope_map = scope_map.split(",") + scope_map += ["module.", "None"] + + for k in dst_state.keys(): + + k_src = k + + if scope_map is not None: + src_prefix = "" + dst_prefix = "" + for i in range(0, len(scope_map), 2): + src_prefix = scope_map[i] if scope_map[i].lower() != "none" else "" + dst_prefix = scope_map[i + 1] if scope_map[i + 1].lower() != "none" else "" + + if dst_prefix == "" and (src_prefix + k) in src_state.keys(): + k_src = src_prefix + k + if not k_src.startswith("module."): + print(f"init param, map: {k} from {k_src} in ckpt") + elif ( + k.startswith(dst_prefix) + and k.replace(dst_prefix, src_prefix, 1) in src_state.keys() + ): + k_src = k.replace(dst_prefix, src_prefix, 1) + if not k_src.startswith("module."): + print(f"init param, map: {k} from {k_src} in ckpt") + + if k_src in src_state.keys(): + if ignore_init_mismatch and dst_state[k].shape != src_state[k_src].shape: + print( + f"ignore_init_mismatch:{ignore_init_mismatch}, dst: {k, dst_state[k].shape}, src: {k_src, src_state[k_src].shape}" + ) + else: + dst_state[k] = src_state[k_src] + + else: + print(f"Warning, miss key in ckpt: {k}, mapped: {k_src}") + + flag = obj.load_state_dict(dst_state, strict=True) + # print(flag) diff --git a/funasr/train_utils/model_summary.py b/funasr/train_utils/model_summary.py index 1001160a9..2aef88a00 100644 --- a/funasr/train_utils/model_summary.py +++ b/funasr/train_utils/model_summary.py @@ -55,9 +55,7 @@ def model_summary(model: torch.nn.Module) -> str: message += "\n\nModel summary:\n" message += f" Class Name: {model.__class__.__name__}\n" message += f" Total Number of model parameters: {tot_params}\n" - message += ( - f" Number of trainable parameters: {num_params} ({percent_trainable}%)\n" - ) + message += f" Number of trainable parameters: {num_params} ({percent_trainable}%)\n" dtype = next(iter(model.parameters())).dtype message += f" Type: {dtype}" diff --git a/funasr/train_utils/recursive_op.py b/funasr/train_utils/recursive_op.py index 286a92daf..1c2a92fd8 100644 --- a/funasr/train_utils/recursive_op.py +++ b/funasr/train_utils/recursive_op.py @@ -1,4 +1,5 @@ """Torch utility module.""" + import torch if torch.distributed.is_available(): diff --git a/funasr/train_utils/trainer.py b/funasr/train_utils/trainer.py index 3ee6885a0..713e17149 100644 --- a/funasr/train_utils/trainer.py +++ b/funasr/train_utils/trainer.py @@ -15,6 +15,7 @@ from funasr.train_utils.recursive_op import recursive_average from funasr.train_utils.average_nbest_models import average_checkpoints from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler + @contextmanager def maybe_autocast(enabled): if enabled: @@ -23,6 +24,7 @@ def maybe_autocast(enabled): else: yield + class Trainer: """ A simple trainer class for training a PyTorch model, saving checkpoints at the end of each epoch, @@ -38,14 +40,16 @@ class Trainer: output_dir (str): Directory where model checkpoints will be saved. resume (str, optional): Path to a checkpoint to resume training from. """ - - def __init__(self, - local_rank, - use_ddp: bool = False, - use_fsdp: bool = False, - use_fp16: bool = False, - output_dir: str="./", - **kwargs): + + def __init__( + self, + local_rank, + use_ddp: bool = False, + use_fsdp: bool = False, + use_fp16: bool = False, + output_dir: str = "./", + **kwargs, + ): """ Initializes the Trainer class with the model, optimizer, scheduler, dataloaders, and other settings. @@ -60,17 +64,17 @@ class Trainer: output_dir (str): The directory where model checkpoints will be saved. Default is './'. resume (str, optional): The file path to a checkpoint to resume training from. """ - + self.output_dir = output_dir if not os.path.exists(self.output_dir): os.makedirs(self.output_dir, exist_ok=True) - self.resume = kwargs.get('resume', True) + self.resume = kwargs.get("resume", True) self.start_epoch = 0 - self.max_epoch = kwargs.get('max_epoch', 100) + self.max_epoch = kwargs.get("max_epoch", 100) self.local_rank = local_rank self.use_ddp = use_ddp self.use_fsdp = use_fsdp - self.device = kwargs.get('device', "cuda") + self.device = kwargs.get("device", "cuda") # self.kwargs = kwargs self.log_interval = kwargs.get("log_interval", 50) self.batch_total = 0 @@ -83,9 +87,7 @@ class Trainer: self.accum_grad = kwargs.get("accum_grad", 1) self.grad_clip = kwargs.get("grad_clip", 10.0) self.grad_clip_type = kwargs.get("grad_clip_type", 2.0) - - - + try: rank = dist.get_rank() world_size = dist.get_world_size() @@ -105,14 +107,16 @@ class Trainer: self.best_step_or_epoch = "" self.val_acc_step_or_eoch = {} self.val_loss_step_or_eoch = {} - - def save_checkpoint(self, epoch, - step=None, - model=None, - optim=None, - scheduler=None, - scaler=None, - ): + + def save_checkpoint( + self, + epoch, + step=None, + model=None, + optim=None, + scheduler=None, + scaler=None, + ): """ Saves a checkpoint containing the model's state, the optimizer's state, and the scheduler's state at the end of the given epoch. This method is @@ -121,15 +125,15 @@ class Trainer: Args: epoch (int): The epoch number at which the checkpoint is being saved. """ - + if self.rank == 0: logging.info(f"Save checkpoint: {epoch}, rank: {self.local_rank}\n") # self.step_or_epoch += 1 state = { - 'epoch': epoch, - 'state_dict': model.state_dict(), - 'optimizer': optim.state_dict(), - 'scheduler': scheduler.state_dict(), + "epoch": epoch, + "state_dict": model.state_dict(), + "optimizer": optim.state_dict(), + "scheduler": scheduler.state_dict(), "saved_ckpts": self.saved_ckpts, "val_acc_step_or_eoch": self.val_acc_step_or_eoch, "val_loss_step_or_eoch": self.val_loss_step_or_eoch, @@ -138,43 +142,59 @@ class Trainer: } if hasattr(model, "module"): state["state_dict"] = model.module.state_dict() - + if scaler: state["scaler_state"] = scaler.state_dict() # Create output directory if it does not exist os.makedirs(self.output_dir, exist_ok=True) if step is None: - ckpt_name = f'model.pt.ep{epoch}' + ckpt_name = f"model.pt.ep{epoch}" else: - ckpt_name = f'model.pt.ep{epoch}.{step}' + ckpt_name = f"model.pt.ep{epoch}.{step}" filename = os.path.join(self.output_dir, ckpt_name) torch.save(state, filename) - - logging.info(f'\nCheckpoint saved to {filename}\n') - latest = Path(os.path.join(self.output_dir, f'model.pt')) + + logging.info(f"\nCheckpoint saved to {filename}\n") + latest = Path(os.path.join(self.output_dir, f"model.pt")) torch.save(state, latest) if self.best_step_or_epoch == "": self.best_step_or_epoch = ckpt_name - + if self.avg_keep_nbest_models_type == "acc": - if self.val_acc_step_or_eoch[ckpt_name] >= self.val_acc_step_or_eoch[self.best_step_or_epoch]: + if ( + self.val_acc_step_or_eoch[ckpt_name] + >= self.val_acc_step_or_eoch[self.best_step_or_epoch] + ): self.best_step_or_epoch = ckpt_name - best_ckpt = Path(os.path.join(self.output_dir, f'model.pt.best')) + best_ckpt = Path(os.path.join(self.output_dir, f"model.pt.best")) torch.save(state, best_ckpt) - logging.info(f"Update best acc: {self.val_acc_step_or_eoch[self.best_step_or_epoch]:.4f}, {best_ckpt}") + logging.info( + f"Update best acc: {self.val_acc_step_or_eoch[self.best_step_or_epoch]:.4f}, {best_ckpt}" + ) else: - logging.info(f"No improvement in acc: {self.val_acc_step_or_eoch[ckpt_name]:.4f} < {self.val_acc_step_or_eoch[self.best_step_or_epoch]:.4f}") + logging.info( + f"No improvement in acc: {self.val_acc_step_or_eoch[ckpt_name]:.4f} < {self.val_acc_step_or_eoch[self.best_step_or_epoch]:.4f}" + ) elif self.avg_keep_nbest_models_type == "loss": - if self.val_loss_step_or_eoch[ckpt_name] <= self.val_loss_step_or_eoch[self.best_step_or_epoch]: + if ( + self.val_loss_step_or_eoch[ckpt_name] + <= self.val_loss_step_or_eoch[self.best_step_or_epoch] + ): self.best_step_or_epoch = ckpt_name - best_ckpt = Path(os.path.join(self.output_dir, f'model.pt.best')) + best_ckpt = Path(os.path.join(self.output_dir, f"model.pt.best")) torch.save(state, best_ckpt) - logging.info(f"Update best loss: {self.val_loss_step_or_eoch[self.best_step_or_epoch]:.4f}, {best_ckpt}") + logging.info( + f"Update best loss: {self.val_loss_step_or_eoch[self.best_step_or_epoch]:.4f}, {best_ckpt}" + ) else: - logging.info(f"No improvement in loss: {self.val_loss_step_or_eoch[ckpt_name]:.4f} > {self.val_loss_step_or_eoch[self.best_step_or_epoch]:.4f}") + logging.info( + f"No improvement in loss: {self.val_loss_step_or_eoch[ckpt_name]:.4f} > {self.val_loss_step_or_eoch[self.best_step_or_epoch]:.4f}" + ) else: print("Undo") - self.saved_ckpts[ckpt_name] = getattr(self, f"val_{self.avg_keep_nbest_models_type}_step_or_eoch")[ckpt_name] + self.saved_ckpts[ckpt_name] = getattr( + self, f"val_{self.avg_keep_nbest_models_type}_step_or_eoch" + )[ckpt_name] if self.keep_nbest_models > 0: if len(self.saved_ckpts) > self.keep_nbest_models: if self.avg_keep_nbest_models_type == "acc": @@ -187,16 +207,17 @@ class Trainer: logging.info(f"Delete: {filename}") if os.path.exists(filename): os.remove(filename) - + if self.use_ddp or self.use_fsdp: dist.barrier() - - def resume_checkpoint(self, - model=None, - optim=None, - scheduler=None, - scaler=None, - ): + + def resume_checkpoint( + self, + model=None, + optim=None, + scheduler=None, + scaler=None, + ): """ Resumes training from a checkpoint at the given file path. Loads the model's state, the optimizer's state, and the scheduler's state. @@ -208,14 +229,14 @@ class Trainer: ckpt = os.path.join(self.output_dir, "model.pt") if os.path.isfile(ckpt): checkpoint = torch.load(ckpt, map_location="cpu") - self.start_epoch = checkpoint['epoch'] + 1 + self.start_epoch = checkpoint["epoch"] + 1 # self.model.load_state_dict(checkpoint['state_dict']) - src_state = checkpoint['state_dict'] + src_state = checkpoint["state_dict"] dst_state = model.state_dict() for k in dst_state.keys(): - if not k.startswith("module.") and "module."+k in src_state.keys(): - k_ddp = "module."+k - elif k.startswith("module.") and "module."+k not in src_state.keys(): + if not k.startswith("module.") and "module." + k in src_state.keys(): + k_ddp = "module." + k + elif k.startswith("module.") and "module." + k not in src_state.keys(): k_ddp = k.replace("module.", "", 1) else: k_ddp = k @@ -223,37 +244,47 @@ class Trainer: dst_state[k] = src_state[k_ddp] else: print(f"Miss key in ckpt: model: {k}, ckpt: {k_ddp}") - + model.load_state_dict(dst_state) - optim.load_state_dict(checkpoint['optimizer']) - scheduler.load_state_dict(checkpoint['scheduler']) - if scaler is not None and 'scaler_state' in checkpoint: - scaler.load_state_dict(checkpoint['scaler_state']) - + optim.load_state_dict(checkpoint["optimizer"]) + scheduler.load_state_dict(checkpoint["scheduler"]) + if scaler is not None and "scaler_state" in checkpoint: + scaler.load_state_dict(checkpoint["scaler_state"]) + self.saved_ckpts = checkpoint["saved_ckpts"] - self.val_acc_step_or_eoch = checkpoint["val_acc_step_or_eoch"] if "val_acc_step_or_eoch" in checkpoint else {} - self.val_loss_step_or_eoch = checkpoint["val_loss_step_or_eoch"] if "val_loss_step_or_eoch" in checkpoint else {} - self.best_step_or_epoch = checkpoint["best_step_or_epoch"] if "best_step_or_epoch" in checkpoint else "" + self.val_acc_step_or_eoch = ( + checkpoint["val_acc_step_or_eoch"] + if "val_acc_step_or_eoch" in checkpoint + else {} + ) + self.val_loss_step_or_eoch = ( + checkpoint["val_loss_step_or_eoch"] + if "val_loss_step_or_eoch" in checkpoint + else {} + ) + self.best_step_or_epoch = ( + checkpoint["best_step_or_epoch"] if "best_step_or_epoch" in checkpoint else "" + ) model.to(self.device) print(f"Checkpoint loaded successfully from '{ckpt}'") else: print(f"No checkpoint found at '{ckpt}', does not resume status!") - + if self.use_ddp or self.use_fsdp: dist.barrier() - - - def train_epoch(self, - model=None, - optim=None, - scheduler=None, - scaler=None, - dataloader_train=None, - dataloader_val=None, - epoch=None, - writer=None, - **kwargs, - ): + + def train_epoch( + self, + model=None, + optim=None, + scheduler=None, + scaler=None, + dataloader_train=None, + dataloader_val=None, + epoch=None, + writer=None, + **kwargs, + ): """ Defines the training process for a single epoch with gradient accumulation. Args: @@ -269,7 +300,7 @@ class Trainer: # Initialize the gradient accumulation optim.zero_grad() speed_stats = {} - + iterator_stop = torch.tensor(0).to(self.device) dataloader_train.batch_sampler.set_epoch(epoch) @@ -293,7 +324,7 @@ class Trainer: time2 = time.perf_counter() with maybe_autocast(self.use_fp16): retval = model(**batch) - + time3 = time.perf_counter() speed_stats["forward_time"] = f"{time3 - time2:0.3f}" loss, stats, weight = retval @@ -306,7 +337,7 @@ class Trainer: if self.use_ddp or self.use_fsdp: dist.all_reduce(weight, op=dist.ReduceOp.SUM) # Now weight is summation over all workers - loss /= weight.sum() # shape:[1] -> shape:[] + loss /= weight.sum() # shape:[1] -> shape:[] # Multiply world_size because DistributedDataParallel # automatically normalizes the gradient by world_size. loss *= self.world_size @@ -318,19 +349,26 @@ class Trainer: loss.backward() time4 = time.perf_counter() speed_stats["backward_time"] = f"{time4 - time3:0.3f}" - - self.train_loss_avg = (self.train_loss_avg*batch_idx + loss.detach().cpu().item())/(batch_idx+1) + + self.train_loss_avg = ( + self.train_loss_avg * batch_idx + loss.detach().cpu().item() + ) / (batch_idx + 1) if "acc" in stats: - self.train_acc_avg = (self.train_acc_avg * batch_idx + stats["acc"].detach().cpu().item()) / (batch_idx + 1) + self.train_acc_avg = ( + self.train_acc_avg * batch_idx + stats["acc"].detach().cpu().item() + ) / (batch_idx + 1) if self.use_ddp or self.use_fsdp: - train_loss_avg = torch.tensor(self.train_loss_avg, dtype=torch.float32).to(self.device) - train_acc_avg = torch.tensor(self.train_acc_avg, dtype=torch.float32).to(self.device) + train_loss_avg = torch.tensor(self.train_loss_avg, dtype=torch.float32).to( + self.device + ) + train_acc_avg = torch.tensor(self.train_acc_avg, dtype=torch.float32).to( + self.device + ) dist.all_reduce(train_loss_avg, op=dist.ReduceOp.SUM) dist.all_reduce(train_acc_avg, op=dist.ReduceOp.SUM) self.train_loss_avg = train_loss_avg.detach().cpu().item() / self.world_size self.train_acc_avg = train_acc_avg.detach().cpu().item() / self.world_size - - + # Perform an optimizer step only after accumulating enough gradients if (batch_idx + 1) % accum_grad == 0: # Perform gradient clipping if it is set @@ -346,7 +384,7 @@ class Trainer: ) optim.zero_grad() # Reset gradients continue - + # Execute an optimization step (update model parameters) if self.use_ddp or self.use_fsdp: dist.barrier() @@ -361,23 +399,25 @@ class Trainer: total_time = f"{time.perf_counter() - time5:0.3f}" time5 = time.perf_counter() speed_stats["optim_time"] = f"{time5 - time4:0.3f}" - + speed_stats["total_time"] = total_time lr = scheduler.get_last_lr()[0] batch_num_epoch = 1 if hasattr(dataloader_train, "__len__"): batch_num_epoch = len(dataloader_train) - self.log(epoch, batch_idx, - batch_num_epoch=batch_num_epoch, - lr=lr, - loss=loss.detach().cpu().item(), - speed_stats=speed_stats, - stats=stats, - writer=writer, - tag="train", - data_split_i=kwargs.get("data_split_i", 0), - data_split_num=kwargs.get("data_split_num", 1), - ) + self.log( + epoch, + batch_idx, + batch_num_epoch=batch_num_epoch, + lr=lr, + loss=loss.detach().cpu().item(), + speed_stats=speed_stats, + stats=stats, + writer=writer, + tag="train", + data_split_i=kwargs.get("data_split_i", 0), + data_split_num=kwargs.get("data_split_num", 1), + ) if (batch_idx + 1) % self.validate_interval == 0: self.validate_epoch( @@ -385,35 +425,41 @@ class Trainer: dataloader_val=dataloader_val, epoch=epoch, writer=writer, - step=batch_idx+1, + step=batch_idx + 1, ) - if (batch_idx+1) % self.save_checkpoint_interval == 0: - self.save_checkpoint(epoch, model=model, optim=optim, scheduler=scheduler, scaler=scaler, step=batch_idx+1) + if (batch_idx + 1) % self.save_checkpoint_interval == 0: + self.save_checkpoint( + epoch, + model=model, + optim=optim, + scheduler=scheduler, + scaler=scaler, + step=batch_idx + 1, + ) time_beg = time.perf_counter() else: if self.use_ddp or self.use_fsdp: iterator_stop.fill_(1) dist.all_reduce(iterator_stop, dist.ReduceOp.SUM) - + if self.use_ddp or self.use_fsdp: dist.barrier() iterator_stop = torch.tensor(0).to(self.device) - - - def validate_epoch(self, - model=None, - dataloader_val=None, - epoch=None, - writer=None, - **kwargs, - ): + def validate_epoch( + self, + model=None, + dataloader_val=None, + epoch=None, + writer=None, + **kwargs, + ): """ Defines the validation process for a single epoch. Should be implemented with the actual model validation steps. - + Args: epoch (int): The current epoch number. """ @@ -421,9 +467,9 @@ class Trainer: dist.barrier() logging.info(f"Validate epoch: {epoch}, rank: {self.local_rank}\n") model.eval() - + with torch.no_grad(): - + speed_stats = {} time5 = time.perf_counter() iterator_stop = torch.tensor(0).to(self.device) @@ -450,7 +496,7 @@ class Trainer: if self.use_ddp or self.use_fsdp: dist.all_reduce(weight, op=dist.ReduceOp.SUM) # Now weight is summation over all workers - loss /= weight.sum() # shape:[1] -> shape:[] + loss /= weight.sum() # shape:[1] -> shape:[] # Multiply world_size because DistributedDataParallel # automatically normalizes the gradient by world_size. loss *= self.world_size @@ -458,12 +504,20 @@ class Trainer: loss = loss time4 = time.perf_counter() - self.val_loss_avg = (self.val_loss_avg*batch_idx + loss.detach().cpu().item())/(batch_idx+1) + self.val_loss_avg = (self.val_loss_avg * batch_idx + loss.detach().cpu().item()) / ( + batch_idx + 1 + ) if "acc" in stats: - self.val_acc_avg = (self.val_acc_avg * batch_idx + stats["acc"].detach().cpu().item()) / (batch_idx + 1) + self.val_acc_avg = ( + self.val_acc_avg * batch_idx + stats["acc"].detach().cpu().item() + ) / (batch_idx + 1) if self.use_ddp or self.use_fsdp: - val_loss_avg = torch.tensor(self.val_loss_avg, dtype=torch.float32).to(self.device) - val_acc_avg = torch.tensor(self.val_acc_avg, dtype=torch.float32).to(self.device) + val_loss_avg = torch.tensor(self.val_loss_avg, dtype=torch.float32).to( + self.device + ) + val_acc_avg = torch.tensor(self.val_acc_avg, dtype=torch.float32).to( + self.device + ) dist.all_reduce(val_loss_avg, op=dist.ReduceOp.SUM) dist.all_reduce(val_acc_avg, op=dist.ReduceOp.SUM) self.val_loss_avg = val_loss_avg.detach().cpu().item() / self.world_size @@ -472,15 +526,17 @@ class Trainer: batch_num_epoch = 1 if hasattr(dataloader_val, "__len__"): batch_num_epoch = len(dataloader_val) - self.log(epoch, batch_idx, - batch_num_epoch=batch_num_epoch, - lr=0.0, - loss=loss.detach().cpu().item(), - speed_stats=speed_stats, - stats=stats, - writer=writer, - tag="val", - ) + self.log( + epoch, + batch_idx, + batch_num_epoch=batch_num_epoch, + lr=0.0, + loss=loss.detach().cpu().item(), + speed_stats=speed_stats, + stats=stats, + writer=writer, + tag="val", + ) else: if self.use_ddp or self.use_fsdp: @@ -488,7 +544,7 @@ class Trainer: dist.all_reduce(iterator_stop, dist.ReduceOp.SUM) if kwargs.get("step", None) is None: - ckpt_name = f'model.pt.ep{epoch}' + ckpt_name = f"model.pt.ep{epoch}" else: ckpt_name = f'model.pt.ep{epoch}.{kwargs.get("step")}' self.val_acc_step_or_eoch[ckpt_name] = self.val_acc_avg @@ -498,34 +554,37 @@ class Trainer: if self.use_ddp or self.use_fsdp: dist.barrier() iterator_stop = torch.tensor(0).to(self.device) - - - def log(self, - epoch=0, - batch_idx=0, - batch_num_epoch=-1, - lr=0.0, - loss=0.0, - speed_stats=None, - stats=None, - writer=None, - tag="train", - data_split_i=0, - data_split_num=1, - **kwargs, - ): - + + def log( + self, + epoch=0, + batch_idx=0, + batch_num_epoch=-1, + lr=0.0, + loss=0.0, + speed_stats=None, + stats=None, + writer=None, + tag="train", + data_split_i=0, + data_split_num=1, + **kwargs, + ): + if (batch_idx + 1) % self.log_interval == 0: - - gpu_info = "GPU, memory: usage: {:.3f} GB, " \ - "peak: {:.3f} GB, " \ - "cache: {:.3f} GB, " \ - "cache_peak: {:.3f} GB".format(torch.cuda.memory_allocated() / 1024 / 1024 / 1024, - torch.cuda.max_memory_allocated() / 1024 / 1024 / 1024, - torch.cuda.memory_reserved() / 1024 / 1024 / 1024, - torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, - ) - + + gpu_info = ( + "GPU, memory: usage: {:.3f} GB, " + "peak: {:.3f} GB, " + "cache: {:.3f} GB, " + "cache_peak: {:.3f} GB".format( + torch.cuda.memory_allocated() / 1024 / 1024 / 1024, + torch.cuda.max_memory_allocated() / 1024 / 1024 / 1024, + torch.cuda.memory_reserved() / 1024 / 1024 / 1024, + torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, + ) + ) + loss_avg_epoch = getattr(self, f"{tag}_loss_avg") acc_avg_epoch = getattr(self, f"{tag}_acc_avg") description = ( @@ -544,23 +603,27 @@ class Trainer: f"{gpu_info}" ) logging.info(description) - + if writer is not None: - writer.add_scalar(f'rank{self.local_rank}_loss/{tag}', loss, self.batch_total) - writer.add_scalar(f'rank{self.local_rank}_lr/{tag}', lr, self.batch_total) - writer.add_scalar(f'rank{self.local_rank}_lr/{tag}', lr, self.batch_total) + writer.add_scalar(f"rank{self.local_rank}_loss/{tag}", loss, self.batch_total) + writer.add_scalar(f"rank{self.local_rank}_lr/{tag}", lr, self.batch_total) + writer.add_scalar(f"rank{self.local_rank}_lr/{tag}", lr, self.batch_total) for key, var in stats.items(): - writer.add_scalar(f'stats_rank{self.local_rank}_{key}/{tag}', var.item(), self.batch_total) + writer.add_scalar( + f"stats_rank{self.local_rank}_{key}/{tag}", var.item(), self.batch_total + ) for key, var in speed_stats.items(): - writer.add_scalar(f'stats_rank{self.local_rank}_{key}/{tag}', eval(var), self.batch_total) - + writer.add_scalar( + f"stats_rank{self.local_rank}_{key}/{tag}", eval(var), self.batch_total + ) + def close(self, writer=None): - + if self.use_ddp or self.use_fsdp: dist.barrier() - + if writer is not None: writer.close() - + if self.use_ddp or self.use_fsdp: - torch.distributed.destroy_process_group() \ No newline at end of file + torch.distributed.destroy_process_group() diff --git a/funasr/utils/datadir_writer.py b/funasr/utils/datadir_writer.py index e555ebcc5..882db7708 100644 --- a/funasr/utils/datadir_writer.py +++ b/funasr/utils/datadir_writer.py @@ -3,7 +3,6 @@ from typing import Union import warnings - class DatadirWriter: """Writer class to create kaldi like data directory. @@ -63,8 +62,7 @@ class DatadirWriter: child.close() if prev_child is not None and prev_child.keys != child.keys: warnings.warn( - f"Ids are mismatching between " - f"{prev_child.path} and {child.path}" + f"Ids are mismatching between " f"{prev_child.path} and {child.path}" ) prev_child = child diff --git a/funasr/utils/export_utils.py b/funasr/utils/export_utils.py index 29502b119..bc7953917 100644 --- a/funasr/utils/export_utils.py +++ b/funasr/utils/export_utils.py @@ -1,68 +1,72 @@ import os import torch -def export_onnx(model, - data_in=None, - quantize: bool = False, - opset_version: int = 14, - **kwargs): - model_scripts = model.export(**kwargs) - export_dir = kwargs.get("output_dir", os.path.dirname(kwargs.get("init_param"))) - os.makedirs(export_dir, exist_ok=True) - - if not isinstance(model_scripts, (list, tuple)): - model_scripts = (model_scripts,) - for m in model_scripts: - m.eval() - _onnx(m, - data_in=data_in, - quantize=quantize, - opset_version=opset_version, - export_dir=export_dir, - **kwargs - ) - print("output dir: {}".format(export_dir)) - - return export_dir - -def _onnx(model, - data_in=None, - quantize: bool = False, - opset_version: int = 14, - export_dir:str = None, - **kwargs): - - dummy_input = model.export_dummy_inputs() - - verbose = kwargs.get("verbose", False) - - export_name = model.export_name() if hasattr(model, "export_name") else "model.onnx" - model_path = os.path.join(export_dir, export_name) - torch.onnx.export( - model, - dummy_input, - model_path, - verbose=verbose, - opset_version=opset_version, - input_names=model.export_input_names(), - output_names=model.export_output_names(), - dynamic_axes=model.export_dynamic_axes() - ) - - if quantize: - from onnxruntime.quantization import QuantType, quantize_dynamic - import onnx - quant_model_path = model_path.replace(".onnx", "_quant.onnx") - if not os.path.exists(quant_model_path): - onnx_model = onnx.load(model_path) - nodes = [n.name for n in onnx_model.graph.node] - nodes_to_exclude = [m for m in nodes if 'output' in m or 'bias_encoder' in m or 'bias_decoder' in m] - quantize_dynamic( - model_input=model_path, - model_output=quant_model_path, - op_types_to_quantize=['MatMul'], - per_channel=True, - reduce_range=False, - weight_type=QuantType.QUInt8, - nodes_to_exclude=nodes_to_exclude, - ) \ No newline at end of file + +def export_onnx(model, data_in=None, quantize: bool = False, opset_version: int = 14, **kwargs): + model_scripts = model.export(**kwargs) + export_dir = kwargs.get("output_dir", os.path.dirname(kwargs.get("init_param"))) + os.makedirs(export_dir, exist_ok=True) + + if not isinstance(model_scripts, (list, tuple)): + model_scripts = (model_scripts,) + for m in model_scripts: + m.eval() + _onnx( + m, + data_in=data_in, + quantize=quantize, + opset_version=opset_version, + export_dir=export_dir, + **kwargs + ) + print("output dir: {}".format(export_dir)) + + return export_dir + + +def _onnx( + model, + data_in=None, + quantize: bool = False, + opset_version: int = 14, + export_dir: str = None, + **kwargs +): + + dummy_input = model.export_dummy_inputs() + + verbose = kwargs.get("verbose", False) + + export_name = model.export_name() if hasattr(model, "export_name") else "model.onnx" + model_path = os.path.join(export_dir, export_name) + torch.onnx.export( + model, + dummy_input, + model_path, + verbose=verbose, + opset_version=opset_version, + input_names=model.export_input_names(), + output_names=model.export_output_names(), + dynamic_axes=model.export_dynamic_axes(), + ) + + if quantize: + from onnxruntime.quantization import QuantType, quantize_dynamic + import onnx + + quant_model_path = model_path.replace(".onnx", "_quant.onnx") + if not os.path.exists(quant_model_path): + onnx_model = onnx.load(model_path) + nodes = [n.name for n in onnx_model.graph.node] + nodes_to_exclude = [ + m for m in nodes if "output" in m or "bias_encoder" in m or "bias_decoder" in m + ] + quantize_dynamic( + model_input=model_path, + model_output=quant_model_path, + op_types_to_quantize=["MatMul"], + per_channel=True, + reduce_range=False, + weight_type=QuantType.QUInt8, + nodes_to_exclude=nodes_to_exclude, + ) diff --git a/funasr/utils/install_model_requirements.py b/funasr/utils/install_model_requirements.py index b67345dc6..d7295db18 100644 --- a/funasr/utils/install_model_requirements.py +++ b/funasr/utils/install_model_requirements.py @@ -1,12 +1,13 @@ import subprocess + def install_requirements(requirements_path): try: result = subprocess.run( - ['pip', 'install', '-r', requirements_path], + ["pip", "install", "-r", requirements_path], stdout=subprocess.PIPE, stderr=subprocess.PIPE, - text=True + text=True, ) # check status @@ -19,10 +20,10 @@ def install_requirements(requirements_path): return False except Exception as e: result = subprocess.run( - ['pip', 'install', '-r', requirements_path], + ["pip", "install", "-r", requirements_path], stdout=subprocess.PIPE, stderr=subprocess.PIPE, - text=True + text=True, ) # check status diff --git a/funasr/utils/load_utils.py b/funasr/utils/load_utils.py index 8ff7115a1..346a00b88 100644 --- a/funasr/utils/load_utils.py +++ b/funasr/utils/load_utils.py @@ -9,6 +9,7 @@ import torchaudio import time import logging from torch.nn.utils.rnn import pad_sequence + try: from funasr.download.file import download_from_url except: @@ -17,39 +18,69 @@ import pdb import subprocess from subprocess import CalledProcessError, run + def is_ffmpeg_installed(): try: - output = subprocess.check_output(['ffmpeg', '-version'], stderr=subprocess.STDOUT) - return 'ffmpeg version' in output.decode('utf-8') + output = subprocess.check_output(["ffmpeg", "-version"], stderr=subprocess.STDOUT) + return "ffmpeg version" in output.decode("utf-8") except (subprocess.CalledProcessError, FileNotFoundError): return False - -use_ffmpeg=False + + +use_ffmpeg = False if is_ffmpeg_installed(): use_ffmpeg = True else: - print("Notice: ffmpeg is not installed. torchaudio is used to load audio\n" - "If you want to use ffmpeg backend to load audio, please install it by:" - "\n\tsudo apt install ffmpeg # ubuntu" - "\n\t# brew install ffmpeg # mac") + print( + "Notice: ffmpeg is not installed. torchaudio is used to load audio\n" + "If you want to use ffmpeg backend to load audio, please install it by:" + "\n\tsudo apt install ffmpeg # ubuntu" + "\n\t# brew install ffmpeg # mac" + ) -def load_audio_text_image_video(data_or_path_or_list, fs: int = 16000, audio_fs: int = 16000, data_type="sound", tokenizer=None, **kwargs): + +def load_audio_text_image_video( + data_or_path_or_list, + fs: int = 16000, + audio_fs: int = 16000, + data_type="sound", + tokenizer=None, + **kwargs, +): if isinstance(data_or_path_or_list, (list, tuple)): if data_type is not None and isinstance(data_type, (list, tuple)): data_types = [data_type] * len(data_or_path_or_list) data_or_path_or_list_ret = [[] for d in data_type] - for i, (data_type_i, data_or_path_or_list_i) in enumerate(zip(data_types, data_or_path_or_list)): - for j, (data_type_j, data_or_path_or_list_j) in enumerate(zip(data_type_i, data_or_path_or_list_i)): - data_or_path_or_list_j = load_audio_text_image_video(data_or_path_or_list_j, fs=fs, audio_fs=audio_fs, data_type=data_type_j, tokenizer=tokenizer, **kwargs) + for i, (data_type_i, data_or_path_or_list_i) in enumerate( + zip(data_types, data_or_path_or_list) + ): + for j, (data_type_j, data_or_path_or_list_j) in enumerate( + zip(data_type_i, data_or_path_or_list_i) + ): + data_or_path_or_list_j = load_audio_text_image_video( + data_or_path_or_list_j, + fs=fs, + audio_fs=audio_fs, + data_type=data_type_j, + tokenizer=tokenizer, + **kwargs, + ) data_or_path_or_list_ret[j].append(data_or_path_or_list_j) return data_or_path_or_list_ret else: - return [load_audio_text_image_video(audio, fs=fs, audio_fs=audio_fs, data_type=data_type, **kwargs) for audio in data_or_path_or_list] - if isinstance(data_or_path_or_list, str) and data_or_path_or_list.startswith('http'): # download url to local file + return [ + load_audio_text_image_video( + audio, fs=fs, audio_fs=audio_fs, data_type=data_type, **kwargs + ) + for audio in data_or_path_or_list + ] + if isinstance(data_or_path_or_list, str) and data_or_path_or_list.startswith( + "http" + ): # download url to local file data_or_path_or_list = download_from_url(data_or_path_or_list) - if isinstance(data_or_path_or_list, str) and os.path.exists(data_or_path_or_list): # local file + if isinstance(data_or_path_or_list, str) and os.path.exists(data_or_path_or_list): # local file if data_type is None or data_type == "sound": # if use_ffmpeg: # data_or_path_or_list = _load_audio_ffmpeg(data_or_path_or_list, sr=fs) @@ -64,14 +95,16 @@ def load_audio_text_image_video(data_or_path_or_list, fs: int = 16000, audio_fs: data_or_path_or_list = data_or_path_or_list.mean(0) except: data_or_path_or_list = _load_audio_ffmpeg(data_or_path_or_list, sr=fs) - data_or_path_or_list = torch.from_numpy(data_or_path_or_list).squeeze() # [n_samples,] + data_or_path_or_list = torch.from_numpy( + data_or_path_or_list + ).squeeze() # [n_samples,] elif data_type == "text" and tokenizer is not None: data_or_path_or_list = tokenizer.encode(data_or_path_or_list) - elif data_type == "image": # undo + elif data_type == "image": # undo pass - elif data_type == "video": # undo + elif data_type == "video": # undo pass - + # if data_in is a file or url, set is_final=True if "cache" in kwargs: kwargs["cache"]["is_final"] = True @@ -81,16 +114,16 @@ def load_audio_text_image_video(data_or_path_or_list, fs: int = 16000, audio_fs: elif isinstance(data_or_path_or_list, np.ndarray): # audio sample point data_or_path_or_list = torch.from_numpy(data_or_path_or_list).squeeze() # [n_samples,] elif isinstance(data_or_path_or_list, str) and data_type == "kaldi_ark": - data_mat = kaldiio.load_mat(data_or_path_or_list) + data_mat = kaldiio.load_mat(data_or_path_or_list) if isinstance(data_mat, tuple): audio_fs, mat = data_mat else: mat = data_mat - if mat.dtype == 'int16' or mat.dtype == 'int32': + if mat.dtype == "int16" or mat.dtype == "int32": mat = mat.astype(np.float64) mat = mat / 32768 - if mat.ndim ==2: - mat = mat[:,0] + if mat.ndim == 2: + mat = mat[:, 0] data_or_path_or_list = mat else: pass @@ -101,30 +134,32 @@ def load_audio_text_image_video(data_or_path_or_list, fs: int = 16000, audio_fs: data_or_path_or_list = resampler(data_or_path_or_list[None, :])[0, :] return data_or_path_or_list + def load_bytes(input): middle_data = np.frombuffer(input, dtype=np.int16) middle_data = np.asarray(middle_data) - if middle_data.dtype.kind not in 'iu': + if middle_data.dtype.kind not in "iu": raise TypeError("'middle_data' must be an array of integers") - dtype = np.dtype('float32') - if dtype.kind != 'f': + dtype = np.dtype("float32") + if dtype.kind != "f": raise TypeError("'dtype' must be a floating point type") - + i = np.iinfo(middle_data.dtype) abs_max = 2 ** (i.bits - 1) offset = i.min + abs_max array = np.frombuffer((middle_data.astype(dtype) - offset) / abs_max, dtype=np.float32) return array -def extract_fbank(data, data_len = None, data_type: str="sound", frontend=None, **kwargs): + +def extract_fbank(data, data_len=None, data_type: str = "sound", frontend=None, **kwargs): if isinstance(data, np.ndarray): data = torch.from_numpy(data) if len(data.shape) < 2: - data = data[None, :] # data: [batch, N] + data = data[None, :] # data: [batch, N] data_len = [data.shape[1]] if data_len is None else data_len elif isinstance(data, torch.Tensor): if len(data.shape) < 2: - data = data[None, :] # data: [batch, N] + data = data[None, :] # data: [batch, N] data_len = [data.shape[1]] if data_len is None else data_len elif isinstance(data, (list, tuple)): data_list, data_len = [], [] @@ -133,14 +168,15 @@ def extract_fbank(data, data_len = None, data_type: str="sound", frontend=None, data_i = torch.from_numpy(data_i) data_list.append(data_i) data_len.append(data_i.shape[0]) - data = pad_sequence(data_list, batch_first=True) # data: [batch, N] + data = pad_sequence(data_list, batch_first=True) # data: [batch, N] data, data_len = frontend(data, data_len, **kwargs) - + if isinstance(data_len, (list, tuple)): data_len = torch.tensor([data_len]) return data.to(torch.float32), data_len.to(torch.int32) + def _load_audio_ffmpeg(file: str, sr: int = 16000): """ Open an audio file and read as mono waveform, resampling as necessary diff --git a/funasr/utils/misc.py b/funasr/utils/misc.py index a853d48b1..5eaa4f871 100644 --- a/funasr/utils/misc.py +++ b/funasr/utils/misc.py @@ -6,50 +6,54 @@ from collections import OrderedDict import numpy as np from omegaconf import DictConfig, OmegaConf + def statistic_model_parameters(model, prefix=None): var_dict = model.state_dict() numel = 0 - for i, key in enumerate(sorted(list([x for x in var_dict.keys() if "num_batches_tracked" not in x]))): + for i, key in enumerate( + sorted(list([x for x in var_dict.keys() if "num_batches_tracked" not in x])) + ): if prefix is None or key.startswith(prefix): numel += var_dict[key].numel() return numel def int2vec(x, vec_dim=8, dtype=np.int32): - b = ('{:0' + str(vec_dim) + 'b}').format(x) + b = ("{:0" + str(vec_dim) + "b}").format(x) # little-endian order: lower bit first - return (np.array(list(b)[::-1]) == '1').astype(dtype) + return (np.array(list(b)[::-1]) == "1").astype(dtype) def seq2arr(seq, vec_dim=8): return np.row_stack([int2vec(int(x), vec_dim) for x in seq]) -def load_scp_as_dict(scp_path, value_type='str', kv_sep=" "): - with io.open(scp_path, 'r', encoding='utf-8') as f: +def load_scp_as_dict(scp_path, value_type="str", kv_sep=" "): + with io.open(scp_path, "r", encoding="utf-8") as f: ret_dict = OrderedDict() for one_line in f.readlines(): one_line = one_line.strip() pos = one_line.find(kv_sep) - key, value = one_line[:pos], one_line[pos + 1:] - if value_type == 'list': - value = value.split(' ') + key, value = one_line[:pos], one_line[pos + 1 :] + if value_type == "list": + value = value.split(" ") ret_dict[key] = value return ret_dict -def load_scp_as_list(scp_path, value_type='str', kv_sep=" "): - with io.open(scp_path, 'r', encoding='utf8') as f: +def load_scp_as_list(scp_path, value_type="str", kv_sep=" "): + with io.open(scp_path, "r", encoding="utf8") as f: ret_dict = [] for one_line in f.readlines(): one_line = one_line.strip() pos = one_line.find(kv_sep) - key, value = one_line[:pos], one_line[pos + 1:] - if value_type == 'list': - value = value.split(' ') + key, value = one_line[:pos], one_line[pos + 1 :] + if value_type == "list": + value = value.split(" ") ret_dict.append((key, value)) return ret_dict + def deep_update(original, update): for key, value in update.items(): if isinstance(value, dict) and key in original: @@ -58,13 +62,12 @@ def deep_update(original, update): deep_update(original[key], value) else: original[key] = value - - + + def prepare_model_dir(**kwargs): - os.makedirs(kwargs.get("output_dir", "./"), exist_ok=True) - + yaml_file = os.path.join(kwargs.get("output_dir", "./"), "config.yaml") OmegaConf.save(config=kwargs, f=yaml_file) print(kwargs) diff --git a/funasr/utils/postprocess_utils.py b/funasr/utils/postprocess_utils.py index efba7551a..34f933c96 100644 --- a/funasr/utils/postprocess_utils.py +++ b/funasr/utils/postprocess_utils.py @@ -6,7 +6,7 @@ from typing import Any, List, Union def isChinese(ch: str): - if '\u4e00' <= ch <= '\u9fff' or '\u0030' <= ch <= '\u0039' or ch == '@': + if "\u4e00" <= ch <= "\u9fff" or "\u0030" <= ch <= "\u0039" or ch == "@": return True return False @@ -14,11 +14,11 @@ def isChinese(ch: str): def isAllChinese(word: Union[List[Any], str]): word_lists = [] for i in word: - cur = i.replace(' ', '') - cur = cur.replace('', '') - cur = cur.replace('', '') - cur = cur.replace('', '') - cur = cur.replace('', '') + cur = i.replace(" ", "") + cur = cur.replace("", "") + cur = cur.replace("", "") + cur = cur.replace("", "") + cur = cur.replace("", "") word_lists.append(cur) if len(word_lists) == 0: @@ -33,11 +33,11 @@ def isAllChinese(word: Union[List[Any], str]): def isAllAlpha(word: Union[List[Any], str]): word_lists = [] for i in word: - cur = i.replace(' ', '') - cur = cur.replace('', '') - cur = cur.replace('', '') - cur = cur.replace('', '') - cur = cur.replace('', '') + cur = i.replace(" ", "") + cur = cur.replace("", "") + cur = cur.replace("", "") + cur = cur.replace("", "") + cur = cur.replace("", "") word_lists.append(cur) if len(word_lists) == 0: @@ -66,12 +66,14 @@ def abbr_dispose(words: List[Any], time_stamp: List[List] = None) -> List[Any]: if num <= last_num: continue - if len(words[num]) == 1 and words[num].encode('utf-8').isalpha(): - if num + 1 < words_size and words[ - num + 1] == ' ' and num + 2 < words_size and len( - words[num + - 2]) == 1 and words[num + - 2].encode('utf-8').isalpha(): + if len(words[num]) == 1 and words[num].encode("utf-8").isalpha(): + if ( + num + 1 < words_size + and words[num + 1] == " " + and num + 2 < words_size + and len(words[num + 2]) == 1 + and words[num + 2].encode("utf-8").isalpha() + ): # found the begin of abbr abbr_begin.append(num) num += 2 @@ -79,11 +81,13 @@ def abbr_dispose(words: List[Any], time_stamp: List[List] = None) -> List[Any]: # to find the end of abbr while True: num += 1 - if num < words_size and words[num] == ' ': + if num < words_size and words[num] == " ": num += 1 - if num < words_size and len( - words[num]) == 1 and words[num].encode( - 'utf-8').isalpha(): + if ( + num < words_size + and len(words[num]) == 1 + and words[num].encode("utf-8").isalpha() + ): abbr_end.pop() abbr_end.append(num) last_num = num @@ -93,11 +97,11 @@ def abbr_dispose(words: List[Any], time_stamp: List[List] = None) -> List[Any]: break for num in range(words_size): - if words[num] == ' ': + if words[num] == " ": ts_nums.append(ts_index) else: ts_nums.append(ts_index) - ts_index += 1 + ts_index += 1 last_num = -1 for num in range(words_size): if num <= last_num: @@ -114,7 +118,7 @@ def abbr_dispose(words: List[Any], time_stamp: List[List] = None) -> List[Any]: last_num = num break else: - if words[num].encode('utf-8').isalpha(): + if words[num].encode("utf-8").isalpha(): abbr_word += words[num].upper() num += 1 word_lists.append(abbr_word) @@ -123,7 +127,7 @@ def abbr_dispose(words: List[Any], time_stamp: List[List] = None) -> List[Any]: ts_lists.append([begin, end]) else: word_lists.append(words[num]) - if time_stamp is not None and words[num] != ' ': + if time_stamp is not None and words[num] != " ": begin = time_stamp[ts_nums[num]][0] end = time_stamp[ts_nums[num]][1] ts_lists.append([begin, end]) @@ -138,18 +142,18 @@ def abbr_dispose(words: List[Any], time_stamp: List[List] = None) -> List[Any]: def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): middle_lists = [] word_lists = [] - word_item = '' + word_item = "" ts_lists = [] # wash words lists for i in words: - word = '' + word = "" if isinstance(i, str): word = i else: - word = i.decode('utf-8') + word = i.decode("utf-8") - if word in ['', '', '', '']: + if word in ["", "", "", ""]: continue else: middle_lists.append(word) @@ -157,7 +161,7 @@ def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): # all chinese characters if isAllChinese(middle_lists): for i, ch in enumerate(middle_lists): - word_lists.append(ch.replace(' ', '')) + word_lists.append(ch.replace(" ", "")) if time_stamp is not None: ts_lists = time_stamp @@ -168,9 +172,9 @@ def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): if ts_flag and time_stamp is not None: begin = time_stamp[i][0] end = time_stamp[i][1] - word = '' - if '@@' in ch: - word = ch.replace('@@', '') + word = "" + if "@@" in ch: + word = ch.replace("@@", "") word_item += word if time_stamp is not None: ts_flag = False @@ -178,8 +182,8 @@ def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): else: word_item += ch word_lists.append(word_item) - word_lists.append(' ') - word_item = '' + word_lists.append(" ") + word_item = "" if time_stamp is not None: ts_flag = True end = time_stamp[i][1] @@ -196,7 +200,7 @@ def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): if ts_flag and time_stamp is not None: begin = time_stamp[i][0] end = time_stamp[i][1] - word = '' + word = "" if isAllChinese(ch): if alpha_blank is True: word_lists.pop() @@ -206,8 +210,8 @@ def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): ts_flag = True ts_lists.append([begin, end]) begin = end - elif '@@' in ch: - word = ch.replace('@@', '') + elif "@@" in ch: + word = ch.replace("@@", "") word_item += word alpha_blank = False if time_stamp is not None: @@ -216,73 +220,74 @@ def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): elif isAllAlpha(ch): word_item += ch word_lists.append(word_item) - word_lists.append(' ') - word_item = '' + word_lists.append(" ") + word_item = "" alpha_blank = True if time_stamp is not None: ts_flag = True - end = time_stamp[i][1] + end = time_stamp[i][1] ts_lists.append([begin, end]) begin = end else: word_lists.append(ch) - if time_stamp is not None: + if time_stamp is not None: word_lists, ts_lists = abbr_dispose(word_lists, ts_lists) real_word_lists = [] for ch in word_lists: - if ch != ' ': + if ch != " ": real_word_lists.append(ch) - sentence = ' '.join(real_word_lists).strip() + sentence = " ".join(real_word_lists).strip() return sentence, ts_lists, real_word_lists else: word_lists = abbr_dispose(word_lists) real_word_lists = [] for ch in word_lists: - if ch != ' ': + if ch != " ": real_word_lists.append(ch) - sentence = ''.join(word_lists).strip() + sentence = "".join(word_lists).strip() return sentence, real_word_lists + def sentence_postprocess_sentencepiece(words): middle_lists = [] word_lists = [] - word_item = '' + word_item = "" # wash words lists for i in words: - word = '' + word = "" if isinstance(i, str): word = i else: - word = i.decode('utf-8') + word = i.decode("utf-8") - if word in ['', '', '', '']: + if word in ["", "", "", ""]: continue else: middle_lists.append(word) # all alpha characters for i, ch in enumerate(middle_lists): - word = '' - if '\u2581' in ch and i == 0: - word_item = '' - word = ch.replace('\u2581', '') + word = "" + if "\u2581" in ch and i == 0: + word_item = "" + word = ch.replace("\u2581", "") word_item += word - elif '\u2581' in ch and i != 0: + elif "\u2581" in ch and i != 0: word_lists.append(word_item) - word_lists.append(' ') - word_item = '' - word = ch.replace('\u2581', '') + word_lists.append(" ") + word_item = "" + word = ch.replace("\u2581", "") word_item += word else: word_item += ch if word_item is not None: word_lists.append(word_item) - #word_lists = abbr_dispose(word_lists) + # word_lists = abbr_dispose(word_lists) real_word_lists = [] for ch in word_lists: - if ch != ' ': + if ch != " ": if ch == "i": ch = ch.replace("i", "I") elif ch == "i'm": @@ -292,5 +297,5 @@ def sentence_postprocess_sentencepiece(words): elif ch == "i'll": ch = ch.replace("i'll", "I'll") real_word_lists.append(ch) - sentence = ''.join(word_lists) - return sentence, real_word_lists \ No newline at end of file + sentence = "".join(word_lists) + return sentence, real_word_lists diff --git a/funasr/utils/speaker_utils.py b/funasr/utils/speaker_utils.py index b769b8577..e470849f8 100644 --- a/funasr/utils/speaker_utils.py +++ b/funasr/utils/speaker_utils.py @@ -19,58 +19,55 @@ def check_audio_list(audio: list): audio_dur = 0 for i in range(len(audio)): seg = audio[i] - assert seg[1] >= seg[0], 'modelscope error: Wrong time stamps.' - assert isinstance(seg[2], np.ndarray), 'modelscope error: Wrong data type.' - assert int(seg[1] * 16000) - int( - seg[0] * 16000 - ) == seg[2].shape[ - 0], 'modelscope error: audio data in list is inconsistent with time length.' + assert seg[1] >= seg[0], "modelscope error: Wrong time stamps." + assert isinstance(seg[2], np.ndarray), "modelscope error: Wrong data type." + assert ( + int(seg[1] * 16000) - int(seg[0] * 16000) == seg[2].shape[0] + ), "modelscope error: audio data in list is inconsistent with time length." if i > 0: - assert seg[0] >= audio[ - i - 1][1], 'modelscope error: Wrong time stamps.' + assert seg[0] >= audio[i - 1][1], "modelscope error: Wrong time stamps." audio_dur += seg[1] - seg[0] return audio_dur # assert audio_dur > 5, 'modelscope error: The effective audio duration is too short.' def sv_preprocess(inputs: Union[np.ndarray, list]): - output = [] - for i in range(len(inputs)): - if isinstance(inputs[i], str): - file_bytes = File.read(inputs[i]) - data, fs = sf.load(io.BytesIO(file_bytes), dtype='float32') - if len(data.shape) == 2: - data = data[:, 0] - data = torch.from_numpy(data).unsqueeze(0) - data = data.squeeze(0) - elif isinstance(inputs[i], np.ndarray): - assert len( - inputs[i].shape - ) == 1, 'modelscope error: Input array should be [N, T]' - data = inputs[i] - if data.dtype in ['int16', 'int32', 'int64']: - data = (data / (1 << 15)).astype('float32') - else: - data = data.astype('float32') - data = torch.from_numpy(data) + output = [] + for i in range(len(inputs)): + if isinstance(inputs[i], str): + file_bytes = File.read(inputs[i]) + data, fs = sf.load(io.BytesIO(file_bytes), dtype="float32") + if len(data.shape) == 2: + data = data[:, 0] + data = torch.from_numpy(data).unsqueeze(0) + data = data.squeeze(0) + elif isinstance(inputs[i], np.ndarray): + assert len(inputs[i].shape) == 1, "modelscope error: Input array should be [N, T]" + data = inputs[i] + if data.dtype in ["int16", "int32", "int64"]: + data = (data / (1 << 15)).astype("float32") else: - raise ValueError( - 'modelscope error: The input type is restricted to audio address and nump array.' - ) - output.append(data) - return output + data = data.astype("float32") + data = torch.from_numpy(data) + else: + raise ValueError( + "modelscope error: The input type is restricted to audio address and nump array." + ) + output.append(data) + return output -def sv_chunk(vad_segments: list, fs = 16000) -> list: +def sv_chunk(vad_segments: list, fs=16000) -> list: config = { - 'seg_dur': 1.5, - 'seg_shift': 0.75, - } + "seg_dur": 1.5, + "seg_shift": 0.75, + } + def seg_chunk(seg_data): seg_st = seg_data[0] data = seg_data[2] - chunk_len = int(config['seg_dur'] * fs) - chunk_shift = int(config['seg_shift'] * fs) + chunk_len = int(config["seg_dur"] * fs) + chunk_shift = int(config["seg_shift"] * fs) last_chunk_ed = 0 seg_res = [] for chunk_st in range(0, data.shape[0], chunk_shift): @@ -81,13 +78,8 @@ def sv_chunk(vad_segments: list, fs = 16000) -> list: chunk_st = max(0, chunk_ed - chunk_len) chunk_data = data[chunk_st:chunk_ed] if chunk_data.shape[0] < chunk_len: - chunk_data = np.pad(chunk_data, - (0, chunk_len - chunk_data.shape[0]), - 'constant') - seg_res.append([ - chunk_st / fs + seg_st, chunk_ed / fs + seg_st, - chunk_data - ]) + chunk_data = np.pad(chunk_data, (0, chunk_len - chunk_data.shape[0]), "constant") + seg_res.append([chunk_st / fs + seg_st, chunk_ed / fs + seg_st, chunk_data]) return seg_res segs = [] @@ -100,16 +92,16 @@ def sv_chunk(vad_segments: list, fs = 16000) -> list: def extract_feature(audio): features = [] for au in audio: - feature = Kaldi.fbank( - au.unsqueeze(0), num_mel_bins=80) + feature = Kaldi.fbank(au.unsqueeze(0), num_mel_bins=80) feature = feature - feature.mean(dim=0, keepdim=True) features.append(feature.unsqueeze(0)) features = torch.cat(features) return features -def postprocess(segments: list, vad_segments: list, - labels: np.ndarray, embeddings: np.ndarray) -> list: +def postprocess( + segments: list, vad_segments: list, labels: np.ndarray, embeddings: np.ndarray +) -> list: assert len(segments) == len(labels) labels = correct_labels(labels) distribute_res = [] @@ -154,16 +146,17 @@ def correct_labels(labels): new_labels.append(id2id[i]) return np.array(new_labels) + def merge_seque(distribute_res): res = [distribute_res[0]] for i in range(1, len(distribute_res)): - if distribute_res[i][2] != res[-1][2] or distribute_res[i][ - 0] > res[-1][1]: + if distribute_res[i][2] != res[-1][2] or distribute_res[i][0] > res[-1][1]: res.append(distribute_res[i]) else: res[-1][1] = distribute_res[i][1] return res + def smooth(res, mindur=1): # short segments are assigned to nearest speakers. for i in range(len(res)): @@ -187,19 +180,18 @@ def smooth(res, mindur=1): def distribute_spk(sentence_list, sd_time_list): sd_sentence_list = [] for d in sentence_list: - sentence_start = d['ts_list'][0][0] - sentence_end = d['ts_list'][-1][1] + sentence_start = d["ts_list"][0][0] + sentence_end = d["ts_list"][-1][1] sentence_spk = 0 max_overlap = 0 for sd_time in sd_time_list: spk_st, spk_ed, spk = sd_time - spk_st = spk_st*1000 - spk_ed = spk_ed*1000 - overlap = max( - min(sentence_end, spk_ed) - max(sentence_start, spk_st), 0) + spk_st = spk_st * 1000 + spk_ed = spk_ed * 1000 + overlap = max(min(sentence_end, spk_ed) - max(sentence_start, spk_st), 0) if overlap > max_overlap: max_overlap = overlap sentence_spk = spk - d['spk'] = sentence_spk + d["spk"] = sentence_spk sd_sentence_list.append(d) return sd_sentence_list diff --git a/funasr/utils/timestamp_tools.py b/funasr/utils/timestamp_tools.py index 232137dba..831d77357 100644 --- a/funasr/utils/timestamp_tools.py +++ b/funasr/utils/timestamp_tools.py @@ -3,6 +3,7 @@ import codecs import logging import argparse import numpy as np + # import edit_distance from itertools import zip_longest @@ -18,20 +19,18 @@ def cif_wo_hidden(alphas, threshold): integrate += alpha list_fires.append(integrate) fire_place = integrate >= threshold - integrate = torch.where(fire_place, - integrate - torch.ones([batch_size], device=alphas.device)*threshold, - integrate) + integrate = torch.where( + fire_place, + integrate - torch.ones([batch_size], device=alphas.device) * threshold, + integrate, + ) fires = torch.stack(list_fires, 1) return fires -def ts_prediction_lfr6_standard(us_alphas, - us_peaks, - char_list, - vad_offset=0.0, - force_time_shift=-1.5, - sil_in_str=True - ): +def ts_prediction_lfr6_standard( + us_alphas, us_peaks, char_list, vad_offset=0.0, force_time_shift=-1.5, sil_in_str=True +): if not len(char_list): return "", [] START_END_THRESHOLD = 5 @@ -41,14 +40,18 @@ def ts_prediction_lfr6_standard(us_alphas, alphas, peaks = us_alphas[0], us_peaks[0] # support inference batch_size=1 only else: alphas, peaks = us_alphas, us_peaks - if char_list[-1] == '': + if char_list[-1] == "": char_list = char_list[:-1] - fire_place = torch.where(peaks>=1.0-1e-4)[0].cpu().numpy() + force_time_shift # total offset + fire_place = ( + torch.where(peaks >= 1.0 - 1e-4)[0].cpu().numpy() + force_time_shift + ) # total offset if len(fire_place) != len(char_list) + 1: - alphas /= (alphas.sum() / (len(char_list) + 1)) + alphas /= alphas.sum() / (len(char_list) + 1) alphas = alphas.unsqueeze(0) - peaks = cif_wo_hidden(alphas, threshold=1.0-1e-4)[0] - fire_place = torch.where(peaks>=1.0-1e-4)[0].cpu().numpy() + force_time_shift # total offset + peaks = cif_wo_hidden(alphas, threshold=1.0 - 1e-4)[0] + fire_place = ( + torch.where(peaks >= 1.0 - 1e-4)[0].cpu().numpy() + force_time_shift + ) # total offset num_frames = peaks.shape[0] timestamp_list = [] new_char_list = [] @@ -59,47 +62,52 @@ def ts_prediction_lfr6_standard(us_alphas, # begin silence if fire_place[0] > START_END_THRESHOLD: # char_list.insert(0, '') - timestamp_list.append([0.0, fire_place[0]*TIME_RATE]) - new_char_list.append('') + timestamp_list.append([0.0, fire_place[0] * TIME_RATE]) + new_char_list.append("") # tokens timestamp - for i in range(len(fire_place)-1): + for i in range(len(fire_place) - 1): new_char_list.append(char_list[i]) - if MAX_TOKEN_DURATION < 0 or fire_place[i+1] - fire_place[i] <= MAX_TOKEN_DURATION: - timestamp_list.append([fire_place[i]*TIME_RATE, fire_place[i+1]*TIME_RATE]) + if MAX_TOKEN_DURATION < 0 or fire_place[i + 1] - fire_place[i] <= MAX_TOKEN_DURATION: + timestamp_list.append([fire_place[i] * TIME_RATE, fire_place[i + 1] * TIME_RATE]) else: # cut the duration to token and sil of the 0-weight frames last long _split = fire_place[i] + MAX_TOKEN_DURATION - timestamp_list.append([fire_place[i]*TIME_RATE, _split*TIME_RATE]) - timestamp_list.append([_split*TIME_RATE, fire_place[i+1]*TIME_RATE]) - new_char_list.append('') + timestamp_list.append([fire_place[i] * TIME_RATE, _split * TIME_RATE]) + timestamp_list.append([_split * TIME_RATE, fire_place[i + 1] * TIME_RATE]) + new_char_list.append("") # tail token and end silence # new_char_list.append(char_list[-1]) if num_frames - fire_place[-1] > START_END_THRESHOLD: _end = (num_frames + fire_place[-1]) * 0.5 - # _end = fire_place[-1] - timestamp_list[-1][1] = _end*TIME_RATE - timestamp_list.append([_end*TIME_RATE, num_frames*TIME_RATE]) + # _end = fire_place[-1] + timestamp_list[-1][1] = _end * TIME_RATE + timestamp_list.append([_end * TIME_RATE, num_frames * TIME_RATE]) new_char_list.append("") else: - timestamp_list[-1][1] = num_frames*TIME_RATE + timestamp_list[-1][1] = num_frames * TIME_RATE if vad_offset: # add offset time in model with vad for i in range(len(timestamp_list)): timestamp_list[i][0] = timestamp_list[i][0] + vad_offset / 1000.0 timestamp_list[i][1] = timestamp_list[i][1] + vad_offset / 1000.0 res_txt = "" for char, timestamp in zip(new_char_list, timestamp_list): - #if char != '': - if not sil_in_str and char == '': continue - res_txt += "{} {} {};".format(char, str(timestamp[0]+0.0005)[:5], str(timestamp[1]+0.0005)[:5]) + # if char != '': + if not sil_in_str and char == "": + continue + res_txt += "{} {} {};".format( + char, str(timestamp[0] + 0.0005)[:5], str(timestamp[1] + 0.0005)[:5] + ) res = [] for char, timestamp in zip(new_char_list, timestamp_list): - if char != '': + if char != "": res.append([int(timestamp[0] * 1000), int(timestamp[1] * 1000)]) return res_txt, res -def timestamp_sentence(punc_id_list, timestamp_postprocessed, text_postprocessed, return_raw_text=False): - punc_list = ['，', '。', '？', '、'] +def timestamp_sentence( + punc_id_list, timestamp_postprocessed, text_postprocessed, return_raw_text=False +): + punc_list = ["，", "。", "？", "、"] res = [] if text_postprocessed is None: return res @@ -111,12 +119,14 @@ def timestamp_sentence(punc_id_list, timestamp_postprocessed, text_postprocessed return res if punc_id_list is None or len(punc_id_list) == 0: - res.append({ - 'text': text_postprocessed.split(), - "start": timestamp_postprocessed[0][0], - "end": timestamp_postprocessed[-1][1], - "timestamp": timestamp_postprocessed, - }) + res.append( + { + "text": text_postprocessed.split(), + "start": timestamp_postprocessed[0][0], + "end": timestamp_postprocessed[-1][1], + "timestamp": timestamp_postprocessed, + } + ) return res if len(punc_id_list) != len(timestamp_postprocessed): logging.warning("length mismatch between punc and timestamp") @@ -126,45 +136,52 @@ def timestamp_sentence(punc_id_list, timestamp_postprocessed, text_postprocessed sentence_start = timestamp_postprocessed[0][0] sentence_end = timestamp_postprocessed[0][1] texts = text_postprocessed.split() - punc_stamp_text_list = list(zip_longest(punc_id_list, timestamp_postprocessed, texts, fillvalue=None)) + punc_stamp_text_list = list( + zip_longest(punc_id_list, timestamp_postprocessed, texts, fillvalue=None) + ) for punc_stamp_text in punc_stamp_text_list: punc_id, timestamp, text = punc_stamp_text # sentence_text += text if text is not None else '' if text is not None: - if 'a' <= text[0] <= 'z' or 'A' <= text[0] <= 'Z': - sentence_text += ' ' + text - elif len(sentence_text) and ('a' <= sentence_text[-1] <= 'z' or 'A' <= sentence_text[-1] <= 'Z'): - sentence_text += ' ' + text + if "a" <= text[0] <= "z" or "A" <= text[0] <= "Z": + sentence_text += " " + text + elif len(sentence_text) and ( + "a" <= sentence_text[-1] <= "z" or "A" <= sentence_text[-1] <= "Z" + ): + sentence_text += " " + text else: sentence_text += text - sentence_text_seg += text + ' ' + sentence_text_seg += text + " " ts_list.append(timestamp) punc_id = int(punc_id) if punc_id is not None else 1 sentence_end = timestamp[1] if timestamp is not None else sentence_end - sentence_text_seg = sentence_text_seg[:-1] if sentence_text_seg[-1] == ' ' else sentence_text_seg + sentence_text_seg = ( + sentence_text_seg[:-1] if sentence_text_seg[-1] == " " else sentence_text_seg + ) if punc_id > 1: sentence_text += punc_list[punc_id - 2] if return_raw_text: - res.append({ - 'text': sentence_text, - "start": sentence_start, - "end": sentence_end, - "timestamp": ts_list, - 'raw_text': sentence_text_seg, - }) + res.append( + { + "text": sentence_text, + "start": sentence_start, + "end": sentence_end, + "timestamp": ts_list, + "raw_text": sentence_text_seg, + } + ) else: - res.append({ - 'text': sentence_text, - "start": sentence_start, - "end": sentence_end, - "timestamp": ts_list, - }) - sentence_text = '' - sentence_text_seg = '' + res.append( + { + "text": sentence_text, + "start": sentence_start, + "end": sentence_end, + "timestamp": ts_list, + } + ) + sentence_text = "" + sentence_text_seg = "" ts_list = [] sentence_start = sentence_end return res - - - diff --git a/funasr/utils/torch_function.py b/funasr/utils/torch_function.py index a078a7ed6..f637bbf82 100644 --- a/funasr/utils/torch_function.py +++ b/funasr/utils/torch_function.py @@ -13,7 +13,7 @@ class MakePadMask(nn.Module): self.mask_pad = torch.Tensor(1 - np.tri(max_seq_len)).type(torch.bool) else: self.mask_pad = torch.Tensor(np.tri(max_seq_len)).type(torch.bool) - + def forward(self, lengths, xs=None, length_dim=-1, maxlen=None): """Make mask tensor containing indices of padded part. This implementation creates the same mask tensor with original make_pad_mask, @@ -25,8 +25,7 @@ class MakePadMask(nn.Module): if xs is not None and len(xs.shape) == 3: if length_dim == 1: - lengths = lengths.unsqueeze(1).expand( - *xs.transpose(1, 2).shape[:2]) + lengths = lengths.unsqueeze(1).expand(*xs.transpose(1, 2).shape[:2]) else: lengths = lengths.unsqueeze(1).expand(*xs.shape[:2]) @@ -44,20 +43,24 @@ class MakePadMask(nn.Module): else: return mask + class sequence_mask(nn.Module): def __init__(self, max_seq_len=512, flip=True): super().__init__() - + def forward(self, lengths, max_seq_len=None, dtype=torch.float32, device=None): if max_seq_len is None: max_seq_len = lengths.max() row_vector = torch.arange(0, max_seq_len, 1).to(lengths.device) matrix = torch.unsqueeze(lengths, dim=-1) mask = row_vector < matrix - + return mask.type(dtype).to(device) if device is not None else mask.type(dtype) -def normalize(input: torch.Tensor, p: float = 2.0, dim: int = 1, out: Optional[torch.Tensor] = None) -> torch.Tensor: + +def normalize( + input: torch.Tensor, p: float = 2.0, dim: int = 1, out: Optional[torch.Tensor] = None +) -> torch.Tensor: if out is None: denom = input.norm(p, dim, keepdim=True).expand_as(input) return input / denom @@ -65,6 +68,7 @@ def normalize(input: torch.Tensor, p: float = 2.0, dim: int = 1, out: Optional[t denom = input.norm(p, dim, keepdim=True).expand_as(input) return torch.div(input, denom, out=out) + def subsequent_mask(size: torch.Tensor): return torch.ones(size, size).tril() @@ -76,5 +80,5 @@ def MakePadMask_test(): print(mask) -if __name__ == '__main__': - MakePadMask_test() \ No newline at end of file +if __name__ == "__main__": + MakePadMask_test() diff --git a/funasr/utils/vad_utils.py b/funasr/utils/vad_utils.py index 6d95e44ab..eba48a972 100644 --- a/funasr/utils/vad_utils.py +++ b/funasr/utils/vad_utils.py @@ -6,30 +6,32 @@ def slice_padding_fbank(speech, speech_lengths, vad_segments): speech_list = [] speech_lengths_list = [] for i, segment in enumerate(vad_segments): - - bed_idx = int(segment[0][0]*16) - end_idx = min(int(segment[0][1]*16), speech_lengths[0]) - speech_i = speech[0, bed_idx: end_idx] - speech_lengths_i = end_idx-bed_idx + + bed_idx = int(segment[0][0] * 16) + end_idx = min(int(segment[0][1] * 16), speech_lengths[0]) + speech_i = speech[0, bed_idx:end_idx] + speech_lengths_i = end_idx - bed_idx speech_list.append(speech_i) speech_lengths_list.append(speech_lengths_i) feats_pad = pad_sequence(speech_list, batch_first=True, padding_value=0.0) speech_lengths_pad = torch.Tensor(speech_lengths_list).int() return feats_pad, speech_lengths_pad + def slice_padding_audio_samples(speech, speech_lengths, vad_segments): speech_list = [] speech_lengths_list = [] for i, segment in enumerate(vad_segments): bed_idx = int(segment[0][0] * 16) end_idx = min(int(segment[0][1] * 16), speech_lengths) - speech_i = speech[bed_idx: end_idx] + speech_i = speech[bed_idx:end_idx] speech_lengths_i = end_idx - bed_idx speech_list.append(speech_i) speech_lengths_list.append(speech_lengths_i) - + return speech_list, speech_lengths_list + def merge_vad(vad_result, max_length=15000): new_result = [] time_step = [t[0] for t in vad_result] + [t[1] for t in vad_result] @@ -37,9 +39,9 @@ def merge_vad(vad_result, max_length=15000): if len(time_step) == 0: return [] bg = 0 - for i in range(len(time_step)-1): + for i in range(len(time_step) - 1): time = time_step[i] - if time_step[i+1] - bg < max_length: + if time_step[i + 1] - bg < max_length: continue if time - bg < max_length * 1.5: new_result.append([bg, time]) @@ -47,8 +49,7 @@ def merge_vad(vad_result, max_length=15000): split_num = int(time - bg) // max_length + 1 spl_l = int(time - bg) // split_num for j in range(split_num): - new_result.append([bg + j*spl_l, bg + (j+1)*spl_l]) + new_result.append([bg + j * spl_l, bg + (j + 1) * spl_l]) bg = time new_result.append([bg, time_step[-1]]) return new_result - \ No newline at end of file diff --git a/runtime/html5/h5Server.py b/runtime/html5/h5Server.py index 96392f800..6c414d5eb 100644 --- a/runtime/html5/h5Server.py +++ b/runtime/html5/h5Server.py @@ -3,63 +3,61 @@ ### Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights ### Reserved. MIT License (https://opensource.org/licenses/MIT) ### -### 2022-2023 by zhaoming,mali aihealthx.com +### 2022-2023 by zhaoming,mali aihealthx.com -from flask import Flask,render_template,request,send_from_directory,jsonify,redirect,url_for -#from gevent.pywsgi import WSGIServer - +from flask import Flask, render_template, request, send_from_directory, jsonify, redirect, url_for + +# from gevent.pywsgi import WSGIServer + import datetime import random import string import time import argparse - -app = Flask(__name__,static_folder='static',static_url_path="/static") -@app.route('/') +app = Flask(__name__, static_folder="static", static_url_path="/static") + + +@app.route("/") def homePage(): - return redirect('/static/index.html') + return redirect("/static/index.html") + - parser = argparse.ArgumentParser() -parser.add_argument("--host", - type=str, - default="0.0.0.0", - required=False, - help="host ip, localhost, 0.0.0.0") -parser.add_argument("--port", - type=int, - default=1337, - required=False, - help="html5 server port") +parser.add_argument( + "--host", type=str, default="0.0.0.0", required=False, help="host ip, localhost, 0.0.0.0" +) +parser.add_argument("--port", type=int, default=1337, required=False, help="html5 server port") -parser.add_argument("--certfile", - type=str, - default="./ssl_key/server.crt", - required=False, - help="certfile for ssl") - -parser.add_argument("--keyfile", - type=str, - default="./ssl_key/server.key", - required=False, - help="keyfile for ssl") +parser.add_argument( + "--certfile", type=str, default="./ssl_key/server.crt", required=False, help="certfile for ssl" +) -if __name__ == '__main__': +parser.add_argument( + "--keyfile", type=str, default="./ssl_key/server.key", required=False, help="keyfile for ssl" +) + +if __name__ == "__main__": args = parser.parse_args() - port=args.port - - #WSGIServer - #ssl = { + port = args.port + + # WSGIServer + # ssl = { # 'certfile': 'server.crt', # 'keyfile': 'server.key' - #} - #httpsServer = WSGIServer(("0.0.0.0",port), app, **ssl) - #httpsServer.serve_forever() - - #flask - print("srv run on ",port) + # } + # httpsServer = WSGIServer(("0.0.0.0",port), app, **ssl) + # httpsServer.serve_forever() - app.run(debug=False,threaded=True,host=args.host,port=port, ssl_context=(args.certfile,args.keyfile)) + # flask + print("srv run on ", port) + + app.run( + debug=False, + threaded=True, + host=args.host, + port=port, + ssl_context=(args.certfile, args.keyfile), + ) diff --git a/runtime/onnxruntime/third_party/kaldi-native-fbank/cmake/cmake_extension.py b/runtime/onnxruntime/third_party/kaldi-native-fbank/cmake/cmake_extension.py index f3ca8744b..25e520927 100644 --- a/runtime/onnxruntime/third_party/kaldi-native-fbank/cmake/cmake_extension.py +++ b/runtime/onnxruntime/third_party/kaldi-native-fbank/cmake/cmake_extension.py @@ -39,7 +39,6 @@ try: # -linux_x86_64.whl self.root_is_pure = False - except ImportError: bdist_wheel = None @@ -84,9 +83,7 @@ class BuildExtension(build_ext): cmake --build {self.build_temp} --target install --config Release -- -m """ print(f"build command is:\n{build_cmd}") - ret = os.system( - f"cmake {cmake_args} -B {self.build_temp} -S {kaldi_native_fbank_dir}" - ) + ret = os.system(f"cmake {cmake_args} -B {self.build_temp} -S {kaldi_native_fbank_dir}") if ret != 0: raise Exception("Failed to configure kaldi_native_fbank") @@ -98,9 +95,7 @@ class BuildExtension(build_ext): else: if make_args == "" and system_make_args == "": print("For fast compilation, run:") - print( - 'export KALDI_NATIVE_FBANK_MAKE_ARGS="-j"; python setup.py install' - ) + print('export KALDI_NATIVE_FBANK_MAKE_ARGS="-j"; python setup.py install') build_cmd = f""" cd {self.build_temp} diff --git a/runtime/python/grpc/grpc_main_client.py b/runtime/python/grpc/grpc_main_client.py index 5c462a707..4606214c8 100644 --- a/runtime/python/grpc/grpc_main_client.py +++ b/runtime/python/grpc/grpc_main_client.py @@ -1,8 +1,8 @@ -''' +""" Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved. MIT License (https://opensource.org/licenses/MIT) 2023 by burkliu(刘柏基) liubaiji@xverse.cn -''' +""" import logging import argparse @@ -13,72 +13,71 @@ import grpc import paraformer_pb2_grpc from paraformer_pb2 import Request, WavFormat, DecodeMode + class GrpcClient: - def __init__(self, wav_path, uri, mode): - self.wav, self.sampling_rate = sf.read(wav_path, dtype='int16') - self.wav_format = WavFormat.pcm - self.audio_chunk_duration = 1000 # ms - self.audio_chunk_size = int(self.sampling_rate * self.audio_chunk_duration / 1000) - self.send_interval = 100 # ms - self.mode = mode + def __init__(self, wav_path, uri, mode): + self.wav, self.sampling_rate = sf.read(wav_path, dtype="int16") + self.wav_format = WavFormat.pcm + self.audio_chunk_duration = 1000 # ms + self.audio_chunk_size = int(self.sampling_rate * self.audio_chunk_duration / 1000) + self.send_interval = 100 # ms + self.mode = mode - # connect to grpc server - channel = grpc.insecure_channel(uri) - self.stub = paraformer_pb2_grpc.ASRStub(channel) - - # start request - for respond in self.stub.Recognize(self.request_iterator()): - logging.info("[receive] mode {}, text {}, is final {}".format( - DecodeMode.Name(respond.mode), respond.text, respond.is_final)) + # connect to grpc server + channel = grpc.insecure_channel(uri) + self.stub = paraformer_pb2_grpc.ASRStub(channel) - def request_iterator(self, mode = DecodeMode.two_pass): - is_first_pack = True - is_final = False - for start in range(0, len(self.wav), self.audio_chunk_size): - request = Request() - audio_chunk = self.wav[start : start + self.audio_chunk_size] + # start request + for respond in self.stub.Recognize(self.request_iterator()): + logging.info( + "[receive] mode {}, text {}, is final {}".format( + DecodeMode.Name(respond.mode), respond.text, respond.is_final + ) + ) - if is_first_pack: - is_first_pack = False - request.sampling_rate = self.sampling_rate - request.mode = self.mode - request.wav_format = self.wav_format - if request.mode == DecodeMode.two_pass or request.mode == DecodeMode.online: - request.chunk_size.extend([5, 10, 5]) + def request_iterator(self, mode=DecodeMode.two_pass): + is_first_pack = True + is_final = False + for start in range(0, len(self.wav), self.audio_chunk_size): + request = Request() + audio_chunk = self.wav[start : start + self.audio_chunk_size] - if start + self.audio_chunk_size >= len(self.wav): - is_final = True - request.is_final = is_final - request.audio_data = audio_chunk.tobytes() - logging.info("[request] audio_data len {}, is final {}".format( - len(request.audio_data), request.is_final)) # int16 = 2bytes - time.sleep(self.send_interval / 1000) - yield request + if is_first_pack: + is_first_pack = False + request.sampling_rate = self.sampling_rate + request.mode = self.mode + request.wav_format = self.wav_format + if request.mode == DecodeMode.two_pass or request.mode == DecodeMode.online: + request.chunk_size.extend([5, 10, 5]) -if __name__ == '__main__': - logging.basicConfig(filename="", format="%(asctime)s %(message)s", level=logging.INFO) - parser = argparse.ArgumentParser() - parser.add_argument("--host", - type=str, - default="127.0.0.1", - required=False, - help="grpc server host ip") - parser.add_argument("--port", - type=int, - default=10100, - required=False, - help="grpc server port") - parser.add_argument("--wav_path", - type=str, - required=True, - help="audio wav path") - args = parser.parse_args() + if start + self.audio_chunk_size >= len(self.wav): + is_final = True + request.is_final = is_final + request.audio_data = audio_chunk.tobytes() + logging.info( + "[request] audio_data len {}, is final {}".format( + len(request.audio_data), request.is_final + ) + ) # int16 = 2bytes + time.sleep(self.send_interval / 1000) + yield request - for mode in [DecodeMode.offline, DecodeMode.online, DecodeMode.two_pass]: - mode_name = DecodeMode.Name(mode) - logging.info("[request] start requesting with mode {}".format(mode_name)) - st = time.time() - uri = '{}:{}'.format(args.host, args.port) - client = GrpcClient(args.wav_path, uri, mode) - logging.info("mode {}, time pass: {}".format(mode_name, time.time() - st)) +if __name__ == "__main__": + logging.basicConfig(filename="", format="%(asctime)s %(message)s", level=logging.INFO) + parser = argparse.ArgumentParser() + parser.add_argument( + "--host", type=str, default="127.0.0.1", required=False, help="grpc server host ip" + ) + parser.add_argument("--port", type=int, default=10100, required=False, help="grpc server port") + parser.add_argument("--wav_path", type=str, required=True, help="audio wav path") + args = parser.parse_args() + + for mode in [DecodeMode.offline, DecodeMode.online, DecodeMode.two_pass]: + mode_name = DecodeMode.Name(mode) + logging.info("[request] start requesting with mode {}".format(mode_name)) + + st = time.time() + uri = "{}:{}".format(args.host, args.port) + client = GrpcClient(args.wav_path, uri, mode) + logging.info("mode {}, time pass: {}".format(mode_name, time.time() - st)) diff --git a/runtime/python/http/client.py b/runtime/python/http/client.py index d6c11eb2b..5024f41df 100644 --- a/runtime/python/http/client.py +++ b/runtime/python/http/client.py @@ -4,21 +4,11 @@ import requests import argparse parser = argparse.ArgumentParser() -parser.add_argument("--host", - type=str, - default="127.0.0.1", - required=False, - help="sever ip") -parser.add_argument("--port", - type=int, - default=8000, - required=False, - help="server port") -parser.add_argument("--audio_path", - type=str, - default='asr_example_zh.wav', - required=False, - help="use audio path") +parser.add_argument("--host", type=str, default="127.0.0.1", required=False, help="sever ip") +parser.add_argument("--port", type=int, default=8000, required=False, help="server port") +parser.add_argument( + "--audio_path", type=str, default="asr_example_zh.wav", required=False, help="use audio path" +) args = parser.parse_args() print("----------- Configuration Arguments -----------") for arg, value in vars(args).items(): @@ -26,9 +16,18 @@ for arg, value in vars(args).items(): print("------------------------------------------------") -url = f'http://{args.host}:{args.port}/recognition' +url = f"http://{args.host}:{args.port}/recognition" headers = {} -files = [('audio', (os.path.basename(args.audio_path), open(args.audio_path, 'rb'), 'application/octet-stream'))] +files = [ + ( + "audio", + ( + os.path.basename(args.audio_path), + open(args.audio_path, "rb"), + "application/octet-stream", + ), + ) +] response = requests.post(url, headers=headers, files=files) print(response.text) diff --git a/runtime/python/http/server.py b/runtime/python/http/server.py index c97ac6385..720d4ce8f 100644 --- a/runtime/python/http/server.py +++ b/runtime/python/http/server.py @@ -15,71 +15,43 @@ logger = get_logger(log_level=logging.INFO) logger.setLevel(logging.INFO) parser = argparse.ArgumentParser() -parser.add_argument("--host", - type=str, - default="0.0.0.0", - required=False, - help="host ip, localhost, 0.0.0.0") -parser.add_argument("--port", - type=int, - default=8000, - required=False, - help="server port") -parser.add_argument("--asr_model", - type=str, - default="paraformer-zh", - help="asr model from https://github.com/alibaba-damo-academy/FunASR?tab=readme-ov-file#model-zoo") -parser.add_argument("--asr_model_revision", - type=str, - default="v2.0.4", - help="") -parser.add_argument("--vad_model", - type=str, - default="fsmn-vad", - help="vad model from https://github.com/alibaba-damo-academy/FunASR?tab=readme-ov-file#model-zoo") -parser.add_argument("--vad_model_revision", - type=str, - default="v2.0.4", - help="") -parser.add_argument("--punc_model", - type=str, - default="ct-punc-c", - help="model from https://github.com/alibaba-damo-academy/FunASR?tab=readme-ov-file#model-zoo") -parser.add_argument("--punc_model_revision", - type=str, - default="v2.0.4", - help="") -parser.add_argument("--ngpu", - type=int, - default=1, - help="0 for cpu, 1 for gpu") -parser.add_argument("--device", - type=str, - default="cuda", - help="cuda, cpu") -parser.add_argument("--ncpu", - type=int, - default=4, - help="cpu cores") -parser.add_argument("--hotword_path", - type=str, - default='hotwords.txt', - help="hot word txt path, only the hot word model works") -parser.add_argument("--certfile", - type=str, - default=None, - required=False, - help="certfile for ssl") -parser.add_argument("--keyfile", - type=str, - default=None, - required=False, - help="keyfile for ssl") -parser.add_argument("--temp_dir", - type=str, - default="temp_dir/", - required=False, - help="temp dir") +parser.add_argument( + "--host", type=str, default="0.0.0.0", required=False, help="host ip, localhost, 0.0.0.0" +) +parser.add_argument("--port", type=int, default=8000, required=False, help="server port") +parser.add_argument( + "--asr_model", + type=str, + default="paraformer-zh", + help="asr model from https://github.com/alibaba-damo-academy/FunASR?tab=readme-ov-file#model-zoo", +) +parser.add_argument("--asr_model_revision", type=str, default="v2.0.4", help="") +parser.add_argument( + "--vad_model", + type=str, + default="fsmn-vad", + help="vad model from https://github.com/alibaba-damo-academy/FunASR?tab=readme-ov-file#model-zoo", +) +parser.add_argument("--vad_model_revision", type=str, default="v2.0.4", help="") +parser.add_argument( + "--punc_model", + type=str, + default="ct-punc-c", + help="model from https://github.com/alibaba-damo-academy/FunASR?tab=readme-ov-file#model-zoo", +) +parser.add_argument("--punc_model_revision", type=str, default="v2.0.4", help="") +parser.add_argument("--ngpu", type=int, default=1, help="0 for cpu, 1 for gpu") +parser.add_argument("--device", type=str, default="cuda", help="cuda, cpu") +parser.add_argument("--ncpu", type=int, default=4, help="cpu cores") +parser.add_argument( + "--hotword_path", + type=str, + default="hotwords.txt", + help="hot word txt path, only the hot word model works", +) +parser.add_argument("--certfile", type=str, default=None, required=False, help="certfile for ssl") +parser.add_argument("--keyfile", type=str, default=None, required=False, help="keyfile for ssl") +parser.add_argument("--temp_dir", type=str, default="temp_dir/", required=False, help="temp dir") args = parser.parse_args() logger.info("----------- Configuration Arguments -----------") for arg, value in vars(args).items(): @@ -90,36 +62,38 @@ os.makedirs(args.temp_dir, exist_ok=True) logger.info("model loading") # load funasr model -model = AutoModel(model=args.asr_model, - model_revision=args.asr_model_revision, - vad_model=args.vad_model, - vad_model_revision=args.vad_model_revision, - punc_model=args.punc_model, - punc_model_revision=args.punc_model_revision, - ngpu=args.ngpu, - ncpu=args.ncpu, - device=args.device, - disable_pbar=True, - disable_log=True) +model = AutoModel( + model=args.asr_model, + model_revision=args.asr_model_revision, + vad_model=args.vad_model, + vad_model_revision=args.vad_model_revision, + punc_model=args.punc_model, + punc_model_revision=args.punc_model_revision, + ngpu=args.ngpu, + ncpu=args.ncpu, + device=args.device, + disable_pbar=True, + disable_log=True, +) logger.info("loaded models!") app = FastAPI(title="FunASR") param_dict = {"sentence_timestamp": True, "batch_size_s": 300} if args.hotword_path is not None and os.path.exists(args.hotword_path): - with open(args.hotword_path, 'r', encoding='utf-8') as f: + with open(args.hotword_path, "r", encoding="utf-8") as f: lines = f.readlines() lines = [line.strip() for line in lines] - hotword = ' '.join(lines) - logger.info(f'热词：{hotword}') - param_dict['hotword'] = hotword + hotword = " ".join(lines) + logger.info(f"热词：{hotword}") + param_dict["hotword"] = hotword @app.post("/recognition") async def api_recognition(audio: UploadFile = File(..., description="audio file")): - suffix = audio.filename.split('.')[-1] - audio_path = f'{args.temp_dir}/{str(uuid.uuid1())}.{suffix}' - async with aiofiles.open(audio_path, 'wb') as out_file: + suffix = audio.filename.split(".")[-1] + audio_path = f"{args.temp_dir}/{str(uuid.uuid1())}.{suffix}" + async with aiofiles.open(audio_path, "wb") as out_file: content = await audio.read() await out_file.write(content) try: @@ -129,7 +103,7 @@ async def api_recognition(audio: UploadFile = File(..., description="audio file" .run(cmd=["ffmpeg", "-nostdin"], capture_stdout=True, capture_stderr=True) ) except Exception as e: - logger.error(f'读取音频文件发生错误，错误信息：{e}') + logger.error(f"读取音频文件发生错误，错误信息：{e}") return {"msg": "读取音频文件发生错误", "code": 1} rec_results = model.generate(input=audio_bytes, is_final=True, **param_dict) # 结果为空 @@ -138,18 +112,22 @@ async def api_recognition(audio: UploadFile = File(..., description="audio file" elif len(rec_results) == 1: # 解析识别结果 rec_result = rec_results[0] - text = rec_result['text'] + text = rec_result["text"] sentences = [] - for sentence in rec_result['sentence_info']: + for sentence in rec_result["sentence_info"]: # 每句话的时间戳 - sentences.append({'text': sentence['text'], 'start': sentence['start'], 'end': sentence['start']}) + sentences.append( + {"text": sentence["text"], "start": sentence["start"], "end": sentence["start"]} + ) ret = {"text": text, "sentences": sentences, "code": 0} - logger.info(f'识别结果：{ret}') + logger.info(f"识别结果：{ret}") return ret else: - logger.info(f'识别结果：{rec_results}') + logger.info(f"识别结果：{rec_results}") return {"msg": "未知错误", "code": -1} -if __name__ == '__main__': - uvicorn.run(app, host=args.host, port=args.port, ssl_keyfile=args.keyfile, ssl_certfile=args.certfile) +if __name__ == "__main__": + uvicorn.run( + app, host=args.host, port=args.port, ssl_keyfile=args.keyfile, ssl_certfile=args.certfile + ) diff --git a/runtime/python/libtorch/demo.py b/runtime/python/libtorch/demo.py index 1a9d9e9ba..1ef9a2024 100644 --- a/runtime/python/libtorch/demo.py +++ b/runtime/python/libtorch/demo.py @@ -1,7 +1,9 @@ from funasr_torch import Paraformer -model_dir = "/nfs/zhifu.gzf/export/damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch" +model_dir = ( + "/nfs/zhifu.gzf/export/damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch" +) model = Paraformer(model_dir, batch_size=1) # cpu # model = Paraformer(model_dir, batch_size=1, device_id=0) # gpu diff --git a/runtime/python/libtorch/funasr_torch/paraformer_bin.py b/runtime/python/libtorch/funasr_torch/paraformer_bin.py index 9954daaa3..68886dfe5 100644 --- a/runtime/python/libtorch/funasr_torch/paraformer_bin.py +++ b/runtime/python/libtorch/funasr_torch/paraformer_bin.py @@ -7,48 +7,46 @@ import copy import librosa import numpy as np -from .utils.utils import (CharTokenizer, Hypothesis, - TokenIDConverter, get_logger, - read_yaml) +from .utils.utils import CharTokenizer, Hypothesis, TokenIDConverter, get_logger, read_yaml from .utils.postprocess_utils import sentence_postprocess from .utils.frontend import WavFrontend from .utils.timestamp_utils import time_stamp_lfr6_onnx + logging = get_logger() import torch -class Paraformer(): - def __init__(self, model_dir: Union[str, Path] = None, - batch_size: int = 1, - device_id: Union[str, int] = "-1", - plot_timestamp_to: str = "", - quantize: bool = False, - intra_op_num_threads: int = 1, - ): +class Paraformer: + def __init__( + self, + model_dir: Union[str, Path] = None, + batch_size: int = 1, + device_id: Union[str, int] = "-1", + plot_timestamp_to: str = "", + quantize: bool = False, + intra_op_num_threads: int = 1, + ): if not Path(model_dir).exists(): - raise FileNotFoundError(f'{model_dir} does not exist.') + raise FileNotFoundError(f"{model_dir} does not exist.") - model_file = os.path.join(model_dir, 'model.torchscripts') + model_file = os.path.join(model_dir, "model.torchscripts") if quantize: - model_file = os.path.join(model_dir, 'model_quant.torchscripts') - config_file = os.path.join(model_dir, 'config.yaml') - cmvn_file = os.path.join(model_dir, 'am.mvn') + model_file = os.path.join(model_dir, "model_quant.torchscripts") + config_file = os.path.join(model_dir, "config.yaml") + cmvn_file = os.path.join(model_dir, "am.mvn") config = read_yaml(config_file) - self.converter = TokenIDConverter(config['token_list']) + self.converter = TokenIDConverter(config["token_list"]) self.tokenizer = CharTokenizer() - self.frontend = WavFrontend( - cmvn_file=cmvn_file, - **config['frontend_conf'] - ) + self.frontend = WavFrontend(cmvn_file=cmvn_file, **config["frontend_conf"]) self.ort_infer = torch.jit.load(model_file) self.batch_size = batch_size self.device_id = device_id self.plot_timestamp_to = plot_timestamp_to - if "predictor_bias" in config['model_conf'].keys(): - self.pred_bias = config['model_conf']['predictor_bias'] + if "predictor_bias" in config["model_conf"].keys(): + self.pred_bias = config["model_conf"]["predictor_bias"] else: self.pred_bias = 0 @@ -57,7 +55,7 @@ class Paraformer(): waveform_nums = len(waveform_list) asr_res = [] for beg_idx in range(0, waveform_nums, self.batch_size): - + end_idx = min(waveform_nums, beg_idx + self.batch_size) feats, feats_len = self.extract_feat(waveform_list[beg_idx:end_idx]) try: @@ -74,51 +72,66 @@ class Paraformer(): else: us_alphas, us_peaks = None, None except: - #logging.warning(traceback.format_exc()) + # logging.warning(traceback.format_exc()) logging.warning("input wav is silence or noise") - preds = [''] + preds = [""] else: preds = self.decode(am_scores, valid_token_lens) if us_peaks is None: for pred in preds: pred = sentence_postprocess(pred) - asr_res.append({'preds': pred}) + asr_res.append({"preds": pred}) else: for pred, us_peaks_ in zip(preds, us_peaks): raw_tokens = pred - timestamp, timestamp_raw = time_stamp_lfr6_onnx(us_peaks_, copy.copy(raw_tokens)) - text_proc, timestamp_proc, _ = sentence_postprocess(raw_tokens, timestamp_raw) + timestamp, timestamp_raw = time_stamp_lfr6_onnx( + us_peaks_, copy.copy(raw_tokens) + ) + text_proc, timestamp_proc, _ = sentence_postprocess( + raw_tokens, timestamp_raw + ) # logging.warning(timestamp) if len(self.plot_timestamp_to): - self.plot_wave_timestamp(waveform_list[0], timestamp, self.plot_timestamp_to) - asr_res.append({'preds': text_proc, 'timestamp': timestamp_proc, "raw_tokens": raw_tokens}) + self.plot_wave_timestamp( + waveform_list[0], timestamp, self.plot_timestamp_to + ) + asr_res.append( + { + "preds": text_proc, + "timestamp": timestamp_proc, + "raw_tokens": raw_tokens, + } + ) return asr_res def plot_wave_timestamp(self, wav, text_timestamp, dest): # TODO: Plot the wav and timestamp results with matplotlib import matplotlib - matplotlib.use('Agg') - matplotlib.rc("font", family='Alibaba PuHuiTi') # set it to a font that your system supports + + matplotlib.use("Agg") + matplotlib.rc( + "font", family="Alibaba PuHuiTi" + ) # set it to a font that your system supports import matplotlib.pyplot as plt + fig, ax1 = plt.subplots(figsize=(11, 3.5), dpi=320) ax2 = ax1.twinx() ax2.set_ylim([0, 2.0]) # plot waveform ax1.set_ylim([-0.3, 0.3]) time = np.arange(wav.shape[0]) / 16000 - ax1.plot(time, wav/wav.max()*0.3, color='gray', alpha=0.4) + ax1.plot(time, wav / wav.max() * 0.3, color="gray", alpha=0.4) # plot lines and text - for (char, start, end) in text_timestamp: - ax1.vlines(start, -0.3, 0.3, ls='--') - ax1.vlines(end, -0.3, 0.3, ls='--') - x_adj = 0.045 if char != '' else 0.12 + for char, start, end in text_timestamp: + ax1.vlines(start, -0.3, 0.3, ls="--") + ax1.vlines(end, -0.3, 0.3, ls="--") + x_adj = 0.045 if char != "" else 0.12 ax1.text((start + end) * 0.5 - x_adj, 0, char) # plt.legend() plotname = "{}/timestamp.png".format(dest) - plt.savefig(plotname, bbox_inches='tight') + plt.savefig(plotname, bbox_inches="tight") - def load_data(self, - wav_content: Union[str, np.ndarray, List[str]], fs: int = None) -> List: + def load_data(self, wav_content: Union[str, np.ndarray, List[str]], fs: int = None) -> List: def load_wav(path: str) -> np.ndarray: waveform, _ = librosa.load(path, sr=fs) return waveform @@ -132,12 +145,9 @@ class Paraformer(): if isinstance(wav_content, list): return [load_wav(path) for path in wav_content] - raise TypeError( - f'The type of {wav_content} is not in [str, np.ndarray, list]') + raise TypeError(f"The type of {wav_content} is not in [str, np.ndarray, list]") - def extract_feat(self, - waveform_list: List[np.ndarray] - ) -> Tuple[np.ndarray, np.ndarray]: + def extract_feat(self, waveform_list: List[np.ndarray]) -> Tuple[np.ndarray, np.ndarray]: feats, feats_len = [], [] for waveform in waveform_list: speech, _ = self.frontend.fbank(waveform) @@ -155,24 +165,23 @@ class Paraformer(): def pad_feats(feats: List[np.ndarray], max_feat_len: int) -> np.ndarray: def pad_feat(feat: np.ndarray, cur_len: int) -> np.ndarray: pad_width = ((0, max_feat_len - cur_len), (0, 0)) - return np.pad(feat, pad_width, 'constant', constant_values=0) + return np.pad(feat, pad_width, "constant", constant_values=0) feat_res = [pad_feat(feat, feat.shape[0]) for feat in feats] feats = np.array(feat_res).astype(np.float32) return feats - def infer(self, feats: np.ndarray, - feats_len: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + def infer(self, feats: np.ndarray, feats_len: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: outputs = self.ort_infer([feats, feats_len]) return outputs def decode(self, am_scores: np.ndarray, token_nums: int) -> List[str]: - return [self.decode_one(am_score, token_num) - for am_score, token_num in zip(am_scores, token_nums)] + return [ + self.decode_one(am_score, token_num) + for am_score, token_num in zip(am_scores, token_nums) + ] - def decode_one(self, - am_score: np.ndarray, - valid_token_num: int) -> List[str]: + def decode_one(self, am_score: np.ndarray, valid_token_num: int) -> List[str]: yseq = am_score.argmax(axis=-1) score = am_score.max(axis=-1) score = np.sum(score, axis=-1) @@ -191,7 +200,6 @@ class Paraformer(): # Change integer-ids to tokens token = self.converter.ids2tokens(token_int) - token = token[:valid_token_num-self.pred_bias] + token = token[: valid_token_num - self.pred_bias] # texts = sentence_postprocess(token) return token - diff --git a/runtime/python/libtorch/funasr_torch/utils/compute_wer.py b/runtime/python/libtorch/funasr_torch/utils/compute_wer.py index 349a3f609..46ec08e5c 100755 --- a/runtime/python/libtorch/funasr_torch/utils/compute_wer.py +++ b/runtime/python/libtorch/funasr_torch/utils/compute_wer.py @@ -2,66 +2,94 @@ import os import numpy as np import sys -def compute_wer(ref_file, - hyp_file, - cer_detail_file): + +def compute_wer(ref_file, hyp_file, cer_detail_file): rst = { - 'Wrd': 0, - 'Corr': 0, - 'Ins': 0, - 'Del': 0, - 'Sub': 0, - 'Snt': 0, - 'Err': 0.0, - 'S.Err': 0.0, - 'wrong_words': 0, - 'wrong_sentences': 0 + "Wrd": 0, + "Corr": 0, + "Ins": 0, + "Del": 0, + "Sub": 0, + "Snt": 0, + "Err": 0.0, + "S.Err": 0.0, + "wrong_words": 0, + "wrong_sentences": 0, } hyp_dict = {} ref_dict = {} - with open(hyp_file, 'r') as hyp_reader: + with open(hyp_file, "r") as hyp_reader: for line in hyp_reader: key = line.strip().split()[0] value = line.strip().split()[1:] hyp_dict[key] = value - with open(ref_file, 'r') as ref_reader: + with open(ref_file, "r") as ref_reader: for line in ref_reader: key = line.strip().split()[0] value = line.strip().split()[1:] ref_dict[key] = value - cer_detail_writer = open(cer_detail_file, 'w') + cer_detail_writer = open(cer_detail_file, "w") for hyp_key in hyp_dict: if hyp_key in ref_dict: - out_item = compute_wer_by_line(hyp_dict[hyp_key], ref_dict[hyp_key]) - rst['Wrd'] += out_item['nwords'] - rst['Corr'] += out_item['cor'] - rst['wrong_words'] += out_item['wrong'] - rst['Ins'] += out_item['ins'] - rst['Del'] += out_item['del'] - rst['Sub'] += out_item['sub'] - rst['Snt'] += 1 - if out_item['wrong'] > 0: - rst['wrong_sentences'] += 1 - cer_detail_writer.write(hyp_key + print_cer_detail(out_item) + '\n') - cer_detail_writer.write("ref:" + '\t' + "".join(ref_dict[hyp_key]) + '\n') - cer_detail_writer.write("hyp:" + '\t' + "".join(hyp_dict[hyp_key]) + '\n') + out_item = compute_wer_by_line(hyp_dict[hyp_key], ref_dict[hyp_key]) + rst["Wrd"] += out_item["nwords"] + rst["Corr"] += out_item["cor"] + rst["wrong_words"] += out_item["wrong"] + rst["Ins"] += out_item["ins"] + rst["Del"] += out_item["del"] + rst["Sub"] += out_item["sub"] + rst["Snt"] += 1 + if out_item["wrong"] > 0: + rst["wrong_sentences"] += 1 + cer_detail_writer.write(hyp_key + print_cer_detail(out_item) + "\n") + cer_detail_writer.write("ref:" + "\t" + "".join(ref_dict[hyp_key]) + "\n") + cer_detail_writer.write("hyp:" + "\t" + "".join(hyp_dict[hyp_key]) + "\n") - if rst['Wrd'] > 0: - rst['Err'] = round(rst['wrong_words'] * 100 / rst['Wrd'], 2) - if rst['Snt'] > 0: - rst['S.Err'] = round(rst['wrong_sentences'] * 100 / rst['Snt'], 2) + if rst["Wrd"] > 0: + rst["Err"] = round(rst["wrong_words"] * 100 / rst["Wrd"], 2) + if rst["Snt"] > 0: + rst["S.Err"] = round(rst["wrong_sentences"] * 100 / rst["Snt"], 2) - cer_detail_writer.write('\n') - cer_detail_writer.write("%WER " + str(rst['Err']) + " [ " + str(rst['wrong_words'])+ " / " + str(rst['Wrd']) + - ", " + str(rst['Ins']) + " ins, " + str(rst['Del']) + " del, " + str(rst['Sub']) + " sub ]" + '\n') - cer_detail_writer.write("%SER " + str(rst['S.Err']) + " [ " + str(rst['wrong_sentences']) + " / " + str(rst['Snt']) + " ]" + '\n') - cer_detail_writer.write("Scored " + str(len(hyp_dict)) + " sentences, " + str(len(hyp_dict) - rst['Snt']) + " not present in hyp." + '\n') + cer_detail_writer.write("\n") + cer_detail_writer.write( + "%WER " + + str(rst["Err"]) + + " [ " + + str(rst["wrong_words"]) + + " / " + + str(rst["Wrd"]) + + ", " + + str(rst["Ins"]) + + " ins, " + + str(rst["Del"]) + + " del, " + + str(rst["Sub"]) + + " sub ]" + + "\n" + ) + cer_detail_writer.write( + "%SER " + + str(rst["S.Err"]) + + " [ " + + str(rst["wrong_sentences"]) + + " / " + + str(rst["Snt"]) + + " ]" + + "\n" + ) + cer_detail_writer.write( + "Scored " + + str(len(hyp_dict)) + + " sentences, " + + str(len(hyp_dict) - rst["Snt"]) + + " not present in hyp." + + "\n" + ) - -def compute_wer_by_line(hyp, - ref): + +def compute_wer_by_line(hyp, ref): hyp = list(map(lambda x: x.lower(), hyp)) ref = list(map(lambda x: x.lower(), ref)) @@ -96,14 +124,7 @@ def compute_wer_by_line(hyp, match_idx = [] i = len_hyp j = len_ref - rst = { - 'nwords': len_ref, - 'cor': 0, - 'wrong': 0, - 'ins': 0, - 'del': 0, - 'sub': 0 - } + rst = {"nwords": len_ref, "cor": 0, "wrong": 0, "ins": 0, "del": 0, "sub": 0} while i >= 0 or j >= 0: i_idx = max(0, i) j_idx = max(0, j) @@ -111,42 +132,57 @@ def compute_wer_by_line(hyp, if ops_matrix[i_idx][j_idx] == 0: # correct if i - 1 >= 0 and j - 1 >= 0: match_idx.append((j - 1, i - 1)) - rst['cor'] += 1 + rst["cor"] += 1 i -= 1 j -= 1 elif ops_matrix[i_idx][j_idx] == 2: # insert i -= 1 - rst['ins'] += 1 + rst["ins"] += 1 elif ops_matrix[i_idx][j_idx] == 3: # delete j -= 1 - rst['del'] += 1 + rst["del"] += 1 elif ops_matrix[i_idx][j_idx] == 1: # substitute i -= 1 j -= 1 - rst['sub'] += 1 + rst["sub"] += 1 if i < 0 and j >= 0: - rst['del'] += 1 + rst["del"] += 1 elif j < 0 and i >= 0: - rst['ins'] += 1 + rst["ins"] += 1 match_idx.reverse() wrong_cnt = cost_matrix[len_hyp][len_ref] - rst['wrong'] = wrong_cnt + rst["wrong"] = wrong_cnt return rst -def print_cer_detail(rst): - return ("(" + "nwords=" + str(rst['nwords']) + ",cor=" + str(rst['cor']) - + ",ins=" + str(rst['ins']) + ",del=" + str(rst['del']) + ",sub=" - + str(rst['sub']) + ") corr:" + '{:.2%}'.format(rst['cor']/rst['nwords']) - + ",cer:" + '{:.2%}'.format(rst['wrong']/rst['nwords'])) -if __name__ == '__main__': +def print_cer_detail(rst): + return ( + "(" + + "nwords=" + + str(rst["nwords"]) + + ",cor=" + + str(rst["cor"]) + + ",ins=" + + str(rst["ins"]) + + ",del=" + + str(rst["del"]) + + ",sub=" + + str(rst["sub"]) + + ") corr:" + + "{:.2%}".format(rst["cor"] / rst["nwords"]) + + ",cer:" + + "{:.2%}".format(rst["wrong"] / rst["nwords"]) + ) + + +if __name__ == "__main__": if len(sys.argv) != 4: print("usage : python compute-wer.py test.ref test.hyp test.wer") sys.exit(0) diff --git a/runtime/python/libtorch/funasr_torch/utils/frontend.py b/runtime/python/libtorch/funasr_torch/utils/frontend.py index fe399551a..743ee5e5f 100644 --- a/runtime/python/libtorch/funasr_torch/utils/frontend.py +++ b/runtime/python/libtorch/funasr_torch/utils/frontend.py @@ -10,22 +10,21 @@ root_dir = Path(__file__).resolve().parent logger_initialized = {} -class WavFrontend(): - """Conventional frontend structure for ASR. - """ +class WavFrontend: + """Conventional frontend structure for ASR.""" def __init__( - self, - cmvn_file: str = None, - fs: int = 16000, - window: str = 'hamming', - n_mels: int = 80, - frame_length: int = 25, - frame_shift: int = 10, - lfr_m: int = 1, - lfr_n: int = 1, - dither: float = 1.0, - **kwargs, + self, + cmvn_file: str = None, + fs: int = 16000, + window: str = "hamming", + n_mels: int = 80, + frame_length: int = 25, + frame_shift: int = 10, + lfr_m: int = 1, + lfr_n: int = 1, + dither: float = 1.0, + **kwargs, ) -> None: opts = knf.FbankOptions() @@ -50,8 +49,7 @@ class WavFrontend(): self.fbank_beg_idx = 0 self.reset_status() - def fbank(self, - waveform: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + def fbank(self, waveform: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: waveform = waveform * (1 << 15) self.fbank_fn = knf.OnlineFbank(self.opts) self.fbank_fn.accept_waveform(self.opts.frame_opts.samp_freq, waveform.tolist()) @@ -63,8 +61,7 @@ class WavFrontend(): feat_len = np.array(mat.shape[0]).astype(np.int32) return feat, feat_len - def fbank_online(self, - waveform: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + def fbank_online(self, waveform: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: waveform = waveform * (1 << 15) # self.fbank_fn = knf.OnlineFbank(self.opts) self.fbank_fn.accept_waveform(self.opts.frame_opts.samp_freq, waveform.tolist()) @@ -102,12 +99,11 @@ class WavFrontend(): T = T + (lfr_m - 1) // 2 for i in range(T_lfr): if lfr_m <= T - i * lfr_n: - LFR_inputs.append( - (inputs[i * lfr_n:i * lfr_n + lfr_m]).reshape(1, -1)) + LFR_inputs.append((inputs[i * lfr_n : i * lfr_n + lfr_m]).reshape(1, -1)) else: # process last LFR frame num_padding = lfr_m - (T - i * lfr_n) - frame = inputs[i * lfr_n:].reshape(-1) + frame = inputs[i * lfr_n :].reshape(-1) for _ in range(num_padding): frame = np.hstack((frame, inputs[-1])) @@ -125,24 +121,26 @@ class WavFrontend(): inputs = (inputs + means) * vars return inputs - def load_cmvn(self,) -> np.ndarray: - with open(self.cmvn_file, 'r', encoding='utf-8') as f: + def load_cmvn( + self, + ) -> np.ndarray: + with open(self.cmvn_file, "r", encoding="utf-8") as f: lines = f.readlines() means_list = [] vars_list = [] for i in range(len(lines)): line_item = lines[i].split() - if line_item[0] == '': + if line_item[0] == "": line_item = lines[i + 1].split() - if line_item[0] == '': - add_shift_line = line_item[3:(len(line_item) - 1)] + if line_item[0] == "": + add_shift_line = line_item[3 : (len(line_item) - 1)] means_list = list(add_shift_line) continue - elif line_item[0] == '': + elif line_item[0] == "": line_item = lines[i + 1].split() - if line_item[0] == '': - rescale_line = line_item[3:(len(line_item) - 1)] + if line_item[0] == "": + rescale_line = line_item[3 : (len(line_item) - 1)] vars_list = list(rescale_line) continue @@ -151,13 +149,14 @@ class WavFrontend(): cmvn = np.array([means, vars]) return cmvn + def load_bytes(input): middle_data = np.frombuffer(input, dtype=np.int16) middle_data = np.asarray(middle_data) - if middle_data.dtype.kind not in 'iu': + if middle_data.dtype.kind not in "iu": raise TypeError("'middle_data' must be an array of integers") - dtype = np.dtype('float32') - if dtype.kind != 'f': + dtype = np.dtype("float32") + if dtype.kind != "f": raise TypeError("'dtype' must be a floating point type") i = np.iinfo(middle_data.dtype) @@ -170,20 +169,25 @@ def load_bytes(input): def test(): path = "/nfs/zhifu.gzf/export/damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch/example/asr_example.wav" import librosa + cmvn_file = "/nfs/zhifu.gzf/export/damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch/am.mvn" config_file = "/nfs/zhifu.gzf/export/damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch/config.yaml" from funasr.runtime.python.onnxruntime.rapid_paraformer.utils.utils import read_yaml + config = read_yaml(config_file) waveform, _ = librosa.load(path, sr=None) frontend = WavFrontend( cmvn_file=cmvn_file, - **config['frontend_conf'], + **config["frontend_conf"], ) - speech, _ = frontend.fbank_online(waveform) #1d, (sample,), numpy - feat, feat_len = frontend.lfr_cmvn(speech) # 2d, (frame, 450), np.float32 -> torch, torch.from_numpy(), dtype, (1, frame, 450) - - frontend.reset_status() # clear cache + speech, _ = frontend.fbank_online(waveform) # 1d, (sample,), numpy + feat, feat_len = frontend.lfr_cmvn( + speech + ) # 2d, (frame, 450), np.float32 -> torch, torch.from_numpy(), dtype, (1, frame, 450) + + frontend.reset_status() # clear cache return feat, feat_len -if __name__ == '__main__': - test() \ No newline at end of file + +if __name__ == "__main__": + test() diff --git a/runtime/python/libtorch/funasr_torch/utils/postprocess_utils.py b/runtime/python/libtorch/funasr_torch/utils/postprocess_utils.py index 575fb90dd..caa6b0057 100644 --- a/runtime/python/libtorch/funasr_torch/utils/postprocess_utils.py +++ b/runtime/python/libtorch/funasr_torch/utils/postprocess_utils.py @@ -6,7 +6,7 @@ from typing import Any, List, Union def isChinese(ch: str): - if '\u4e00' <= ch <= '\u9fff' or '\u0030' <= ch <= '\u0039': + if "\u4e00" <= ch <= "\u9fff" or "\u0030" <= ch <= "\u0039": return True return False @@ -14,9 +14,9 @@ def isChinese(ch: str): def isAllChinese(word: Union[List[Any], str]): word_lists = [] for i in word: - cur = i.replace(' ', '') - cur = cur.replace('', '') - cur = cur.replace('', '') + cur = i.replace(" ", "") + cur = cur.replace("", "") + cur = cur.replace("", "") word_lists.append(cur) if len(word_lists) == 0: @@ -31,9 +31,9 @@ def isAllChinese(word: Union[List[Any], str]): def isAllAlpha(word: Union[List[Any], str]): word_lists = [] for i in word: - cur = i.replace(' ', '') - cur = cur.replace('', '') - cur = cur.replace('', '') + cur = i.replace(" ", "") + cur = cur.replace("", "") + cur = cur.replace("", "") word_lists.append(cur) if len(word_lists) == 0: @@ -62,12 +62,14 @@ def abbr_dispose(words: List[Any], time_stamp: List[List] = None) -> List[Any]: if num <= last_num: continue - if len(words[num]) == 1 and words[num].encode('utf-8').isalpha(): - if num + 1 < words_size and words[ - num + 1] == ' ' and num + 2 < words_size and len( - words[num + - 2]) == 1 and words[num + - 2].encode('utf-8').isalpha(): + if len(words[num]) == 1 and words[num].encode("utf-8").isalpha(): + if ( + num + 1 < words_size + and words[num + 1] == " " + and num + 2 < words_size + and len(words[num + 2]) == 1 + and words[num + 2].encode("utf-8").isalpha() + ): # found the begin of abbr abbr_begin.append(num) num += 2 @@ -75,11 +77,13 @@ def abbr_dispose(words: List[Any], time_stamp: List[List] = None) -> List[Any]: # to find the end of abbr while True: num += 1 - if num < words_size and words[num] == ' ': + if num < words_size and words[num] == " ": num += 1 - if num < words_size and len( - words[num]) == 1 and words[num].encode( - 'utf-8').isalpha(): + if ( + num < words_size + and len(words[num]) == 1 + and words[num].encode("utf-8").isalpha() + ): abbr_end.pop() abbr_end.append(num) last_num = num @@ -89,11 +93,11 @@ def abbr_dispose(words: List[Any], time_stamp: List[List] = None) -> List[Any]: break for num in range(words_size): - if words[num] == ' ': + if words[num] == " ": ts_nums.append(ts_index) else: ts_nums.append(ts_index) - ts_index += 1 + ts_index += 1 last_num = -1 for num in range(words_size): if num <= last_num: @@ -110,7 +114,7 @@ def abbr_dispose(words: List[Any], time_stamp: List[List] = None) -> List[Any]: last_num = num break else: - if words[num].encode('utf-8').isalpha(): + if words[num].encode("utf-8").isalpha(): word_lists.append(words[num].upper()) num += 1 if time_stamp is not None: @@ -118,7 +122,7 @@ def abbr_dispose(words: List[Any], time_stamp: List[List] = None) -> List[Any]: ts_lists.append([begin, end]) else: word_lists.append(words[num]) - if time_stamp is not None and words[num] != ' ': + if time_stamp is not None and words[num] != " ": begin = time_stamp[ts_nums[num]][0] end = time_stamp[ts_nums[num]][1] ts_lists.append([begin, end]) @@ -133,18 +137,18 @@ def abbr_dispose(words: List[Any], time_stamp: List[List] = None) -> List[Any]: def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): middle_lists = [] word_lists = [] - word_item = '' + word_item = "" ts_lists = [] # wash words lists for i in words: - word = '' + word = "" if isinstance(i, str): word = i else: - word = i.decode('utf-8') + word = i.decode("utf-8") - if word in ['', '', '']: + if word in ["", "", ""]: continue else: middle_lists.append(word) @@ -152,7 +156,7 @@ def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): # all chinese characters if isAllChinese(middle_lists): for i, ch in enumerate(middle_lists): - word_lists.append(ch.replace(' ', '')) + word_lists.append(ch.replace(" ", "")) if time_stamp is not None: ts_lists = time_stamp @@ -163,9 +167,9 @@ def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): if ts_flag and time_stamp is not None: begin = time_stamp[i][0] end = time_stamp[i][1] - word = '' - if '@@' in ch: - word = ch.replace('@@', '') + word = "" + if "@@" in ch: + word = ch.replace("@@", "") word_item += word if time_stamp is not None: ts_flag = False @@ -173,8 +177,8 @@ def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): else: word_item += ch word_lists.append(word_item) - word_lists.append(' ') - word_item = '' + word_lists.append(" ") + word_item = "" if time_stamp is not None: ts_flag = True end = time_stamp[i][1] @@ -191,7 +195,7 @@ def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): if ts_flag and time_stamp is not None: begin = time_stamp[i][0] end = time_stamp[i][1] - word = '' + word = "" if isAllChinese(ch): if alpha_blank is True: word_lists.pop() @@ -201,8 +205,8 @@ def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): ts_flag = True ts_lists.append([begin, end]) begin = end - elif '@@' in ch: - word = ch.replace('@@', '') + elif "@@" in ch: + word = ch.replace("@@", "") word_item += word alpha_blank = False if time_stamp is not None: @@ -211,30 +215,30 @@ def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): elif isAllAlpha(ch): word_item += ch word_lists.append(word_item) - word_lists.append(' ') - word_item = '' + word_lists.append(" ") + word_item = "" alpha_blank = True if time_stamp is not None: ts_flag = True - end = time_stamp[i][1] + end = time_stamp[i][1] ts_lists.append([begin, end]) begin = end else: - raise ValueError('invalid character: {}'.format(ch)) + raise ValueError("invalid character: {}".format(ch)) - if time_stamp is not None: + if time_stamp is not None: word_lists, ts_lists = abbr_dispose(word_lists, ts_lists) real_word_lists = [] for ch in word_lists: - if ch != ' ': + if ch != " ": real_word_lists.append(ch) - sentence = ' '.join(real_word_lists).strip() + sentence = " ".join(real_word_lists).strip() return sentence, ts_lists, real_word_lists else: word_lists = abbr_dispose(word_lists) real_word_lists = [] for ch in word_lists: - if ch != ' ': + if ch != " ": real_word_lists.append(ch) - sentence = ''.join(word_lists).strip() + sentence = "".join(word_lists).strip() return sentence, real_word_lists diff --git a/runtime/python/libtorch/funasr_torch/utils/timestamp_utils.py b/runtime/python/libtorch/funasr_torch/utils/timestamp_utils.py index 7d0060cc5..5abbafed8 100644 --- a/runtime/python/libtorch/funasr_torch/utils/timestamp_utils.py +++ b/runtime/python/libtorch/funasr_torch/utils/timestamp_utils.py @@ -3,46 +3,50 @@ import numpy as np def time_stamp_lfr6_onnx(us_cif_peak, char_list, begin_time=0.0, total_offset=-1.5): if not len(char_list): - return '', [] + return "", [] START_END_THRESHOLD = 5 MAX_TOKEN_DURATION = 30 TIME_RATE = 10.0 * 6 / 1000 / 3 # 3 times upsampled cif_peak = us_cif_peak.reshape(-1) num_frames = cif_peak.shape[-1] - if char_list[-1] == '': + if char_list[-1] == "": char_list = char_list[:-1] # char_list = [i for i in text] timestamp_list = [] new_char_list = [] # for bicif model trained with large data, cif2 actually fires when a character starts # so treat the frames between two peaks as the duration of the former token - fire_place = np.where(cif_peak>1.0-1e-4)[0] + total_offset # np format + fire_place = np.where(cif_peak > 1.0 - 1e-4)[0] + total_offset # np format num_peak = len(fire_place) - assert num_peak == len(char_list) + 1 # number of peaks is supposed to be number of tokens + 1 + assert num_peak == len(char_list) + 1 # number of peaks is supposed to be number of tokens + 1 # begin silence if fire_place[0] > START_END_THRESHOLD: # char_list.insert(0, '') - timestamp_list.append([0.0, fire_place[0]*TIME_RATE]) - new_char_list.append('') + timestamp_list.append([0.0, fire_place[0] * TIME_RATE]) + new_char_list.append("") # tokens timestamp - for i in range(len(fire_place)-1): + for i in range(len(fire_place) - 1): new_char_list.append(char_list[i]) - if i == len(fire_place)-2 or MAX_TOKEN_DURATION < 0 or fire_place[i+1] - fire_place[i] < MAX_TOKEN_DURATION: - timestamp_list.append([fire_place[i]*TIME_RATE, fire_place[i+1]*TIME_RATE]) + if ( + i == len(fire_place) - 2 + or MAX_TOKEN_DURATION < 0 + or fire_place[i + 1] - fire_place[i] < MAX_TOKEN_DURATION + ): + timestamp_list.append([fire_place[i] * TIME_RATE, fire_place[i + 1] * TIME_RATE]) else: # cut the duration to token and sil of the 0-weight frames last long _split = fire_place[i] + MAX_TOKEN_DURATION - timestamp_list.append([fire_place[i]*TIME_RATE, _split*TIME_RATE]) - timestamp_list.append([_split*TIME_RATE, fire_place[i+1]*TIME_RATE]) - new_char_list.append('') + timestamp_list.append([fire_place[i] * TIME_RATE, _split * TIME_RATE]) + timestamp_list.append([_split * TIME_RATE, fire_place[i + 1] * TIME_RATE]) + new_char_list.append("") # tail token and end silence if num_frames - fire_place[-1] > START_END_THRESHOLD: _end = (num_frames + fire_place[-1]) / 2 - timestamp_list[-1][1] = _end*TIME_RATE - timestamp_list.append([_end*TIME_RATE, num_frames*TIME_RATE]) + timestamp_list[-1][1] = _end * TIME_RATE + timestamp_list.append([_end * TIME_RATE, num_frames * TIME_RATE]) new_char_list.append("") else: - timestamp_list[-1][1] = num_frames*TIME_RATE + timestamp_list[-1][1] = num_frames * TIME_RATE if begin_time: # add offset time in model with vad for i in range(len(timestamp_list)): timestamp_list[i][0] = timestamp_list[i][0] + begin_time / 1000.0 @@ -53,7 +57,6 @@ def time_stamp_lfr6_onnx(us_cif_peak, char_list, begin_time=0.0, total_offset=-1 res_str += "{} {} {};".format(char, timestamp[0], timestamp[1]) res = [] for char, timestamp in zip(new_char_list, timestamp_list): - if char != '': + if char != "": res.append([int(timestamp[0] * 1000), int(timestamp[1] * 1000)]) return res_str, res - \ No newline at end of file diff --git a/runtime/python/libtorch/funasr_torch/utils/utils.py b/runtime/python/libtorch/funasr_torch/utils/utils.py index 913ddc1e4..f85d4e932 100644 --- a/runtime/python/libtorch/funasr_torch/utils/utils.py +++ b/runtime/python/libtorch/funasr_torch/utils/utils.py @@ -17,24 +17,23 @@ root_dir = Path(__file__).resolve().parent logger_initialized = {} -class TokenIDConverter(): - def __init__(self, token_list: Union[List, str], - ): +class TokenIDConverter: + def __init__( + self, + token_list: Union[List, str], + ): self.token_list = token_list self.unk_symbol = token_list[-1] self.token2id = {v: i for i, v in enumerate(self.token_list)} self.unk_id = self.token2id[self.unk_symbol] - def get_num_vocabulary_size(self) -> int: return len(self.token_list) - def ids2tokens(self, - integers: Union[np.ndarray, Iterable[int]]) -> List[str]: + def ids2tokens(self, integers: Union[np.ndarray, Iterable[int]]) -> List[str]: if isinstance(integers, np.ndarray) and integers.ndim != 1: - raise TokenIDConverterError( - f"Must be 1 dim ndarray, but got {integers.ndim}") + raise TokenIDConverterError(f"Must be 1 dim ndarray, but got {integers.ndim}") return [self.token_list[i] for i in integers] def tokens2ids(self, tokens: Iterable[str]) -> List[int]: @@ -42,7 +41,7 @@ class TokenIDConverter(): return [self.token2id.get(i, self.unk_id) for i in tokens] -class CharTokenizer(): +class CharTokenizer: def __init__( self, symbol_value: Union[Path, str, Iterable[str]] = None, @@ -77,7 +76,7 @@ class CharTokenizer(): if line.startswith(w): if not self.remove_non_linguistic_symbols: tokens.append(line[: len(w)]) - line = line[len(w):] + line = line[len(w) :] break else: t = line[0] @@ -100,7 +99,6 @@ class CharTokenizer(): ) - class Hypothesis(NamedTuple): """Hypothesis data type.""" @@ -120,15 +118,15 @@ class Hypothesis(NamedTuple): def read_yaml(yaml_path: Union[str, Path]) -> Dict: if not Path(yaml_path).exists(): - raise FileExistsError(f'The {yaml_path} does not exist.') + raise FileExistsError(f"The {yaml_path} does not exist.") - with open(str(yaml_path), 'rb') as f: + with open(str(yaml_path), "rb") as f: data = yaml.load(f, Loader=yaml.Loader) return data @functools.lru_cache() -def get_logger(name='funasr_torch'): +def get_logger(name="funasr_torch"): """Initialize and get a logger by name. If the logger has not been initialized, this method will initialize the logger by adding one or two handlers, otherwise the initialized logger will @@ -148,8 +146,8 @@ def get_logger(name='funasr_torch'): return logger formatter = logging.Formatter( - '[%(asctime)s] %(name)s %(levelname)s: %(message)s', - datefmt="%Y/%m/%d %H:%M:%S") + "[%(asctime)s] %(name)s %(levelname)s: %(message)s", datefmt="%Y/%m/%d %H:%M:%S" + ) sh = logging.StreamHandler() sh.setFormatter(formatter) diff --git a/runtime/python/libtorch/setup.py b/runtime/python/libtorch/setup.py index 4b20c0b2d..58df4fe20 100644 --- a/runtime/python/libtorch/setup.py +++ b/runtime/python/libtorch/setup.py @@ -3,19 +3,19 @@ from pathlib import Path import setuptools from setuptools import find_packages + def get_readme(): root_dir = Path(__file__).resolve().parent - readme_path = str(root_dir / 'README.md') + readme_path = str(root_dir / "README.md") print(readme_path) - with open(readme_path, 'r', encoding='utf-8') as f: + with open(readme_path, "r", encoding="utf-8") as f: readme = f.read() return readme - setuptools.setup( - name='funasr_torch', - version='0.0.4', + name="funasr_torch", + version="0.0.4", platforms="Any", url="https://github.com/alibaba-damo-academy/FunASR.git", author="Speech Lab of DAMO Academy, Alibaba Group", @@ -23,24 +23,24 @@ setuptools.setup( description="FunASR: A Fundamental End-to-End Speech Recognition Toolkit", license="The MIT License", long_description=get_readme(), - long_description_content_type='text/markdown', + long_description_content_type="text/markdown", include_package_data=True, - install_requires=["librosa", - "onnxruntime>=1.7.0", - "scipy", - "numpy>=1.19.3", - "kaldi-native-fbank", - "PyYAML>=5.1.2", - "torch-quant >= 0.4.0"], - packages=find_packages(include=["torch_paraformer*"]), - keywords=[ - 'funasr, paraformer, funasr_torch' + install_requires=[ + "librosa", + "onnxruntime>=1.7.0", + "scipy", + "numpy>=1.19.3", + "kaldi-native-fbank", + "PyYAML>=5.1.2", + "torch-quant >= 0.4.0", ], + packages=find_packages(include=["torch_paraformer*"]), + keywords=["funasr, paraformer, funasr_torch"], classifiers=[ - 'Programming Language :: Python :: 3.6', - 'Programming Language :: Python :: 3.7', - 'Programming Language :: Python :: 3.8', - 'Programming Language :: Python :: 3.9', - 'Programming Language :: Python :: 3.10', + "Programming Language :: Python :: 3.6", + "Programming Language :: Python :: 3.7", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", ], ) diff --git a/runtime/python/onnxruntime/demo_contextual_paraformer.py b/runtime/python/onnxruntime/demo_contextual_paraformer.py index d8aee235d..04b0ce9a1 100644 --- a/runtime/python/onnxruntime/demo_contextual_paraformer.py +++ b/runtime/python/onnxruntime/demo_contextual_paraformer.py @@ -4,9 +4,9 @@ from pathlib import Path model_dir = "damo/speech_paraformer-large-contextual_asr_nat-zh-cn-16k-common-vocab8404" model = ContextualParaformer(model_dir, batch_size=1) -wav_path = ['{}/.cache/modelscope/hub/{}/example/asr_example.wav'.format(Path.home(), model_dir)] +wav_path = ["{}/.cache/modelscope/hub/{}/example/asr_example.wav".format(Path.home(), model_dir)] -hotwords = '你的热词 魔搭' +hotwords = "你的热词 魔搭" result = model(wav_path, hotwords) print(result) diff --git a/runtime/python/onnxruntime/demo_paraformer_offline.py b/runtime/python/onnxruntime/demo_paraformer_offline.py index c6fc79c66..3fe6d0a7e 100644 --- a/runtime/python/onnxruntime/demo_paraformer_offline.py +++ b/runtime/python/onnxruntime/demo_paraformer_offline.py @@ -9,7 +9,7 @@ model = Paraformer(model_dir, batch_size=1, quantize=False) # when using paraformer-large-vad-punc model, you can set plot_timestamp_to="./xx.png" to get figure of alignment besides timestamps # model = Paraformer(model_dir, batch_size=1, plot_timestamp_to="test.png") -wav_path = ['{}/.cache/modelscope/hub/{}/example/asr_example.wav'.format(Path.home(), model_dir)] +wav_path = ["{}/.cache/modelscope/hub/{}/example/asr_example.wav".format(Path.home(), model_dir)] result = model(wav_path) print(result) diff --git a/runtime/python/onnxruntime/demo_paraformer_online.py b/runtime/python/onnxruntime/demo_paraformer_online.py index 210f3eaf8..bb72b41a1 100644 --- a/runtime/python/onnxruntime/demo_paraformer_online.py +++ b/runtime/python/onnxruntime/demo_paraformer_online.py @@ -3,17 +3,19 @@ from funasr_onnx.paraformer_online_bin import Paraformer from pathlib import Path model_dir = "damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-online" -wav_path = ['{}/.cache/modelscope/hub/{}/example/asr_example.wav'.format(Path.home(), model_dir)] +wav_path = ["{}/.cache/modelscope/hub/{}/example/asr_example.wav".format(Path.home(), model_dir)] chunk_size = [5, 10, 5] -model = Paraformer(model_dir, batch_size=1, quantize=True, chunk_size=chunk_size, intra_op_num_threads=4) # only support batch_size = 1 +model = Paraformer( + model_dir, batch_size=1, quantize=True, chunk_size=chunk_size, intra_op_num_threads=4 +) # only support batch_size = 1 ##online asr speech, sample_rate = soundfile.read(wav_path) speech_length = speech.shape[0] sample_offset = 0 step = chunk_size[1] * 960 -param_dict = {'cache': dict()} +param_dict = {"cache": dict()} final_result = "" for sample_offset in range(0, speech_length, min(step, speech_length - sample_offset)): if sample_offset + step >= speech_length - 1: @@ -21,10 +23,9 @@ for sample_offset in range(0, speech_length, min(step, speech_length - sample_of is_final = True else: is_final = False - param_dict['is_final'] = is_final - rec_result = model(audio_in=speech[sample_offset: sample_offset + step], - param_dict=param_dict) + param_dict["is_final"] = is_final + rec_result = model(audio_in=speech[sample_offset : sample_offset + step], param_dict=param_dict) if len(rec_result) > 0: - final_result += rec_result[0]["preds"][0] + final_result += rec_result[0]["preds"][0] print(rec_result) print(final_result) diff --git a/runtime/python/onnxruntime/demo_punc_offline.py b/runtime/python/onnxruntime/demo_punc_offline.py index c56b484a5..206aca2ea 100644 --- a/runtime/python/onnxruntime/demo_punc_offline.py +++ b/runtime/python/onnxruntime/demo_punc_offline.py @@ -1,9 +1,9 @@ from funasr_onnx import CT_Transformer -#model_dir = "damo/punc_ct-transformer_zh-cn-common-vocab272727-pytorch" +# model_dir = "damo/punc_ct-transformer_zh-cn-common-vocab272727-pytorch" model_dir = "damo/punc_ct-transformer_cn-en-common-vocab471067-large" model = CT_Transformer(model_dir) -text_in="跨境河流是养育沿岸人民的生命之源长期以来为帮助下游地区防灾减灾中方技术人员在上游地区极为恶劣的自然条件下克服巨大困难甚至冒着生命危险向印方提供汛期水文资料处理紧急事件中方重视印方在跨境河流问题上的关切愿意进一步完善双方联合工作机制凡是中方能做的我们都会去做而且会做得更好我请印度朋友们放心中国在上游的任何开发利用都会经过科学规划和论证兼顾上下游的利益" +text_in = "跨境河流是养育沿岸人民的生命之源长期以来为帮助下游地区防灾减灾中方技术人员在上游地区极为恶劣的自然条件下克服巨大困难甚至冒着生命危险向印方提供汛期水文资料处理紧急事件中方重视印方在跨境河流问题上的关切愿意进一步完善双方联合工作机制凡是中方能做的我们都会去做而且会做得更好我请印度朋友们放心中国在上游的任何开发利用都会经过科学规划和论证兼顾上下游的利益" result = model(text_in) print(result[0]) diff --git a/runtime/python/onnxruntime/demo_punc_online.py b/runtime/python/onnxruntime/demo_punc_online.py index f69866288..36019fb99 100644 --- a/runtime/python/onnxruntime/demo_punc_online.py +++ b/runtime/python/onnxruntime/demo_punc_online.py @@ -3,10 +3,10 @@ from funasr_onnx import CT_Transformer_VadRealtime model_dir = "damo/punc_ct-transformer_zh-cn-common-vad_realtime-vocab272727" model = CT_Transformer_VadRealtime(model_dir) -text_in = "跨境河流是养育沿岸|人民的生命之源长期以来为帮助下游地区防灾减灾中方技术人员|在上游地区极为恶劣的自然条件下克服巨大困难甚至冒着生命危险|向印方提供汛期水文资料处理紧急事件中方重视印方在跨境河流>问题上的关切|愿意进一步完善双方联合工作机制|凡是|中方能做的我们|都会去做而且会做得更好我请印度朋友们放心中国在上游的|任何开发利用都会经过科学|规划和论证兼顾上下游的利益" +text_in = "跨境河流是养育沿岸|人民的生命之源长期以来为帮助下游地区防灾减灾中方技术人员|在上游地区极为恶劣的自然条件下克服巨大困难甚至冒着生命危险|向印方提供汛期水文资料处理紧急事件中方重视印方在跨境河流>问题上的关切|愿意进一步完善双方联合工作机制|凡是|中方能做的我们|都会去做而且会做得更好我请印度朋友们放心中国在上游的|任何开发利用都会经过科学|规划和论证兼顾上下游的利益" vads = text_in.split("|") -rec_result_all="" +rec_result_all = "" param_dict = {"cache": []} for vad in vads: result = model(vad, param_dict=param_dict) diff --git a/runtime/python/onnxruntime/demo_seaco_paraformer.py b/runtime/python/onnxruntime/demo_seaco_paraformer.py index 29cff3e14..f03719b30 100644 --- a/runtime/python/onnxruntime/demo_seaco_paraformer.py +++ b/runtime/python/onnxruntime/demo_seaco_paraformer.py @@ -4,9 +4,9 @@ from pathlib import Path model_dir = "iic/speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch" model = SeacoParaformer(model_dir, batch_size=1) -wav_path = ['{}/.cache/modelscope/hub/{}/example/asr_example.wav'.format(Path.home(), model_dir)] +wav_path = ["{}/.cache/modelscope/hub/{}/example/asr_example.wav".format(Path.home(), model_dir)] -hotwords = '你的热词 魔搭' +hotwords = "你的热词 魔搭" result = model(wav_path, hotwords) print(result) diff --git a/runtime/python/onnxruntime/demo_vad_offline.py b/runtime/python/onnxruntime/demo_vad_offline.py index 555436bf7..79c984904 100644 --- a/runtime/python/onnxruntime/demo_vad_offline.py +++ b/runtime/python/onnxruntime/demo_vad_offline.py @@ -2,7 +2,9 @@ from funasr_onnx import Fsmn_vad from pathlib import Path model_dir = "damo/speech_fsmn_vad_zh-cn-16k-common-pytorch" -wav_path = '{}/.cache/modelscope/hub/damo/speech_fsmn_vad_zh-cn-16k-common-pytorch/example/vad_example.wav'.format(Path.home()) +wav_path = "{}/.cache/modelscope/hub/damo/speech_fsmn_vad_zh-cn-16k-common-pytorch/example/vad_example.wav".format( + Path.home() +) model = Fsmn_vad(model_dir) diff --git a/runtime/python/onnxruntime/demo_vad_online.py b/runtime/python/onnxruntime/demo_vad_online.py index a80f58a55..b3e580be2 100644 --- a/runtime/python/onnxruntime/demo_vad_online.py +++ b/runtime/python/onnxruntime/demo_vad_online.py @@ -3,7 +3,9 @@ import soundfile from pathlib import Path model_dir = "damo/speech_fsmn_vad_zh-cn-16k-common-pytorch" -wav_path = '{}/.cache/modelscope/hub/damo/speech_fsmn_vad_zh-cn-16k-common-pytorch/example/vad_example.wav'.format(Path.home()) +wav_path = "{}/.cache/modelscope/hub/damo/speech_fsmn_vad_zh-cn-16k-common-pytorch/example/vad_example.wav".format( + Path.home() +) model = Fsmn_vad_online(model_dir) @@ -14,15 +16,16 @@ speech_length = speech.shape[0] # sample_offset = 0 step = 1600 -param_dict = {'in_cache': []} +param_dict = {"in_cache": []} for sample_offset in range(0, speech_length, min(step, speech_length - sample_offset)): if sample_offset + step >= speech_length - 1: step = speech_length - sample_offset is_final = True else: is_final = False - param_dict['is_final'] = is_final - segments_result = model(audio_in=speech[sample_offset: sample_offset + step], - param_dict=param_dict) + param_dict["is_final"] = is_final + segments_result = model( + audio_in=speech[sample_offset : sample_offset + step], param_dict=param_dict + ) if segments_result: - print(segments_result) \ No newline at end of file + print(segments_result) diff --git a/runtime/python/onnxruntime/funasr_client_http.py b/runtime/python/onnxruntime/funasr_client_http.py index 1ddd62ca1..bd0d9d669 100644 --- a/runtime/python/onnxruntime/funasr_client_http.py +++ b/runtime/python/onnxruntime/funasr_client_http.py @@ -3,16 +3,25 @@ import requests import threading -with open('A2_0.wav','rb') as f: - test_wav_bytes=f.read() -url='http://127.0.0.1:8888/api/asr' -def send_post(i,url,wav_bytes): - r1=requests.post(url,json={'wav_base64':str(base64.b64encode(wav_bytes), 'utf-8')}) - print('线程:',i,r1.json()) +with open("A2_0.wav", "rb") as f: + test_wav_bytes = f.read() +url = "http://127.0.0.1:8888/api/asr" + + +def send_post(i, url, wav_bytes): + r1 = requests.post(url, json={"wav_base64": str(base64.b64encode(wav_bytes), "utf-8")}) + print("线程:", i, r1.json()) for i in range(100): - t = threading.Thread(target=send_post, args=(i,url,test_wav_bytes,)) + t = threading.Thread( + target=send_post, + args=( + i, + url, + test_wav_bytes, + ), + ) t.start() - #t.join() -print('完成测试') \ No newline at end of file + # t.join() +print("完成测试") diff --git a/runtime/python/onnxruntime/funasr_onnx/paraformer_bin.py b/runtime/python/onnxruntime/funasr_onnx/paraformer_bin.py index ee216098f..2269273ec 100644 --- a/runtime/python/onnxruntime/funasr_onnx/paraformer_bin.py +++ b/runtime/python/onnxruntime/funasr_onnx/paraformer_bin.py @@ -11,11 +11,16 @@ import copy import librosa import numpy as np -from .utils.utils import (CharTokenizer, Hypothesis, ONNXRuntimeError, - OrtInferSession, TokenIDConverter, get_logger, - read_yaml) -from .utils.postprocess_utils import (sentence_postprocess, - sentence_postprocess_sentencepiece) +from .utils.utils import ( + CharTokenizer, + Hypothesis, + ONNXRuntimeError, + OrtInferSession, + TokenIDConverter, + get_logger, + read_yaml, +) +from .utils.postprocess_utils import sentence_postprocess, sentence_postprocess_sentencepiece from .utils.frontend import WavFrontend from .utils.timestamp_utils import time_stamp_lfr6_onnx from .utils.utils import pad_list @@ -23,72 +28,70 @@ from .utils.utils import pad_list logging = get_logger() -class Paraformer(): +class Paraformer: """ Author: Speech Lab of DAMO Academy, Alibaba Group Paraformer: Fast and Accurate Parallel Transformer for Non-autoregressive End-to-End Speech Recognition https://arxiv.org/abs/2206.08317 """ - def __init__(self, model_dir: Union[str, Path] = None, - batch_size: int = 1, - device_id: Union[str, int] = "-1", - plot_timestamp_to: str = "", - quantize: bool = False, - intra_op_num_threads: int = 4, - cache_dir: str = None, - **kwargs - ): + + def __init__( + self, + model_dir: Union[str, Path] = None, + batch_size: int = 1, + device_id: Union[str, int] = "-1", + plot_timestamp_to: str = "", + quantize: bool = False, + intra_op_num_threads: int = 4, + cache_dir: str = None, + **kwargs, + ): if not Path(model_dir).exists(): try: from modelscope.hub.snapshot_download import snapshot_download except: - raise "You are exporting model from modelscope, please install modelscope and try it again. To install modelscope, you could:\n" \ - "\npip3 install -U modelscope\n" \ - "For the users in China, you could install with the command:\n" \ - "\npip3 install -U modelscope -i https://mirror.sjtu.edu.cn/pypi/web/simple" + raise "You are exporting model from modelscope, please install modelscope and try it again. To install modelscope, you could:\n" "\npip3 install -U modelscope\n" "For the users in China, you could install with the command:\n" "\npip3 install -U modelscope -i https://mirror.sjtu.edu.cn/pypi/web/simple" try: model_dir = snapshot_download(model_dir, cache_dir=cache_dir) except: - raise "model_dir must be model_name in modelscope or local path downloaded from modelscope, but is {}".format(model_dir) - - model_file = os.path.join(model_dir, 'model.onnx') + raise "model_dir must be model_name in modelscope or local path downloaded from modelscope, but is {}".format( + model_dir + ) + + model_file = os.path.join(model_dir, "model.onnx") if quantize: - model_file = os.path.join(model_dir, 'model_quant.onnx') + model_file = os.path.join(model_dir, "model_quant.onnx") if not os.path.exists(model_file): print(".onnx is not exist, begin to export onnx") try: from funasr import AutoModel except: - raise "You are exporting onnx, please install funasr and try it again. To install funasr, you could:\n" \ - "\npip3 install -U funasr\n" \ - "For the users in China, you could install with the command:\n" \ - "\npip3 install -U funasr -i https://mirror.sjtu.edu.cn/pypi/web/simple" + raise "You are exporting onnx, please install funasr and try it again. To install funasr, you could:\n" "\npip3 install -U funasr\n" "For the users in China, you could install with the command:\n" "\npip3 install -U funasr -i https://mirror.sjtu.edu.cn/pypi/web/simple" model = AutoModel(model=model_dir) model_dir = model.export(type="onnx", quantize=quantize, **kwargs) - - config_file = os.path.join(model_dir, 'config.yaml') - cmvn_file = os.path.join(model_dir, 'am.mvn') + + config_file = os.path.join(model_dir, "config.yaml") + cmvn_file = os.path.join(model_dir, "am.mvn") config = read_yaml(config_file) - token_list = os.path.join(model_dir, 'tokens.json') - with open(token_list, 'r', encoding='utf-8') as f: + token_list = os.path.join(model_dir, "tokens.json") + with open(token_list, "r", encoding="utf-8") as f: token_list = json.load(f) self.converter = TokenIDConverter(token_list) self.tokenizer = CharTokenizer() - self.frontend = WavFrontend( - cmvn_file=cmvn_file, - **config['frontend_conf'] + self.frontend = WavFrontend(cmvn_file=cmvn_file, **config["frontend_conf"]) + self.ort_infer = OrtInferSession( + model_file, device_id, intra_op_num_threads=intra_op_num_threads ) - self.ort_infer = OrtInferSession(model_file, device_id, intra_op_num_threads=intra_op_num_threads) self.batch_size = batch_size self.plot_timestamp_to = plot_timestamp_to - if "predictor_bias" in config['model_conf'].keys(): - self.pred_bias = config['model_conf']['predictor_bias'] + if "predictor_bias" in config["model_conf"].keys(): + self.pred_bias = config["model_conf"]["predictor_bias"] else: self.pred_bias = 0 if "lang" in config: - self.language = config['lang'] + self.language = config["lang"] else: self.language = None @@ -97,7 +100,7 @@ class Paraformer(): waveform_nums = len(waveform_list) asr_res = [] for beg_idx in range(0, waveform_nums, self.batch_size): - + end_idx = min(waveform_nums, beg_idx + self.batch_size) feats, feats_len = self.extract_feat(waveform_list[beg_idx:end_idx]) try: @@ -109,9 +112,9 @@ class Paraformer(): else: us_alphas, us_peaks = None, None except ONNXRuntimeError: - #logging.warning(traceback.format_exc()) + # logging.warning(traceback.format_exc()) logging.warning("input wav is silence or noise") - preds = [''] + preds = [""] else: preds = self.decode(am_scores, valid_token_lens) if us_peaks is None: @@ -120,43 +123,58 @@ class Paraformer(): pred = sentence_postprocess_sentencepiece(pred) else: pred = sentence_postprocess(pred) - asr_res.append({'preds': pred}) + asr_res.append({"preds": pred}) else: for pred, us_peaks_ in zip(preds, us_peaks): raw_tokens = pred - timestamp, timestamp_raw = time_stamp_lfr6_onnx(us_peaks_, copy.copy(raw_tokens)) - text_proc, timestamp_proc, _ = sentence_postprocess(raw_tokens, timestamp_raw) + timestamp, timestamp_raw = time_stamp_lfr6_onnx( + us_peaks_, copy.copy(raw_tokens) + ) + text_proc, timestamp_proc, _ = sentence_postprocess( + raw_tokens, timestamp_raw + ) # logging.warning(timestamp) if len(self.plot_timestamp_to): - self.plot_wave_timestamp(waveform_list[0], timestamp, self.plot_timestamp_to) - asr_res.append({'preds': text_proc, 'timestamp': timestamp_proc, "raw_tokens": raw_tokens}) + self.plot_wave_timestamp( + waveform_list[0], timestamp, self.plot_timestamp_to + ) + asr_res.append( + { + "preds": text_proc, + "timestamp": timestamp_proc, + "raw_tokens": raw_tokens, + } + ) return asr_res def plot_wave_timestamp(self, wav, text_timestamp, dest): # TODO: Plot the wav and timestamp results with matplotlib import matplotlib - matplotlib.use('Agg') - matplotlib.rc("font", family='Alibaba PuHuiTi') # set it to a font that your system supports + + matplotlib.use("Agg") + matplotlib.rc( + "font", family="Alibaba PuHuiTi" + ) # set it to a font that your system supports import matplotlib.pyplot as plt + fig, ax1 = plt.subplots(figsize=(11, 3.5), dpi=320) ax2 = ax1.twinx() ax2.set_ylim([0, 2.0]) # plot waveform ax1.set_ylim([-0.3, 0.3]) time = np.arange(wav.shape[0]) / 16000 - ax1.plot(time, wav/wav.max()*0.3, color='gray', alpha=0.4) + ax1.plot(time, wav / wav.max() * 0.3, color="gray", alpha=0.4) # plot lines and text - for (char, start, end) in text_timestamp: - ax1.vlines(start, -0.3, 0.3, ls='--') - ax1.vlines(end, -0.3, 0.3, ls='--') - x_adj = 0.045 if char != '' else 0.12 + for char, start, end in text_timestamp: + ax1.vlines(start, -0.3, 0.3, ls="--") + ax1.vlines(end, -0.3, 0.3, ls="--") + x_adj = 0.045 if char != "" else 0.12 ax1.text((start + end) * 0.5 - x_adj, 0, char) # plt.legend() plotname = "{}/timestamp.png".format(dest) - plt.savefig(plotname, bbox_inches='tight') + plt.savefig(plotname, bbox_inches="tight") - def load_data(self, - wav_content: Union[str, np.ndarray, List[str]], fs: int = None) -> List: + def load_data(self, wav_content: Union[str, np.ndarray, List[str]], fs: int = None) -> List: def load_wav(path: str) -> np.ndarray: waveform, _ = librosa.load(path, sr=fs) return waveform @@ -170,12 +188,9 @@ class Paraformer(): if isinstance(wav_content, list): return [load_wav(path) for path in wav_content] - raise TypeError( - f'The type of {wav_content} is not in [str, np.ndarray, list]') + raise TypeError(f"The type of {wav_content} is not in [str, np.ndarray, list]") - def extract_feat(self, - waveform_list: List[np.ndarray] - ) -> Tuple[np.ndarray, np.ndarray]: + def extract_feat(self, waveform_list: List[np.ndarray]) -> Tuple[np.ndarray, np.ndarray]: feats, feats_len = [], [] for waveform in waveform_list: speech, _ = self.frontend.fbank(waveform) @@ -191,24 +206,23 @@ class Paraformer(): def pad_feats(feats: List[np.ndarray], max_feat_len: int) -> np.ndarray: def pad_feat(feat: np.ndarray, cur_len: int) -> np.ndarray: pad_width = ((0, max_feat_len - cur_len), (0, 0)) - return np.pad(feat, pad_width, 'constant', constant_values=0) + return np.pad(feat, pad_width, "constant", constant_values=0) feat_res = [pad_feat(feat, feat.shape[0]) for feat in feats] feats = np.array(feat_res).astype(np.float32) return feats - def infer(self, feats: np.ndarray, - feats_len: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + def infer(self, feats: np.ndarray, feats_len: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: outputs = self.ort_infer([feats, feats_len]) return outputs def decode(self, am_scores: np.ndarray, token_nums: int) -> List[str]: - return [self.decode_one(am_score, token_num) - for am_score, token_num in zip(am_scores, token_nums)] + return [ + self.decode_one(am_score, token_num) + for am_score, token_num in zip(am_scores, token_nums) + ] - def decode_one(self, - am_score: np.ndarray, - valid_token_num: int) -> List[str]: + def decode_one(self, am_score: np.ndarray, valid_token_num: int) -> List[str]: yseq = am_score.argmax(axis=-1) score = am_score.max(axis=-1) score = np.sum(score, axis=-1) @@ -227,7 +241,7 @@ class Paraformer(): # Change integer-ids to tokens token = self.converter.ids2tokens(token_int) - token = token[:valid_token_num-self.pred_bias] + token = token[: valid_token_num - self.pred_bias] # texts = sentence_postprocess(token) return token @@ -238,81 +252,80 @@ class ContextualParaformer(Paraformer): Paraformer: Fast and Accurate Parallel Transformer for Non-autoregressive End-to-End Speech Recognition https://arxiv.org/abs/2206.08317 """ - def __init__(self, model_dir: Union[str, Path] = None, - batch_size: int = 1, - device_id: Union[str, int] = "-1", - plot_timestamp_to: str = "", - quantize: bool = False, - intra_op_num_threads: int = 4, - cache_dir: str = None, - **kwargs - ): + + def __init__( + self, + model_dir: Union[str, Path] = None, + batch_size: int = 1, + device_id: Union[str, int] = "-1", + plot_timestamp_to: str = "", + quantize: bool = False, + intra_op_num_threads: int = 4, + cache_dir: str = None, + **kwargs, + ): if not Path(model_dir).exists(): try: from modelscope.hub.snapshot_download import snapshot_download except: - raise "You are exporting model from modelscope, please install modelscope and try it again. To install modelscope, you could:\n" \ - "\npip3 install -U modelscope\n" \ - "For the users in China, you could install with the command:\n" \ - "\npip3 install -U modelscope -i https://mirror.sjtu.edu.cn/pypi/web/simple" + raise "You are exporting model from modelscope, please install modelscope and try it again. To install modelscope, you could:\n" "\npip3 install -U modelscope\n" "For the users in China, you could install with the command:\n" "\npip3 install -U modelscope -i https://mirror.sjtu.edu.cn/pypi/web/simple" try: model_dir = snapshot_download(model_dir, cache_dir=cache_dir) except: - raise "model_dir must be model_name in modelscope or local path downloaded from modelscope, but is {}".format(model_dir) - + raise "model_dir must be model_name in modelscope or local path downloaded from modelscope, but is {}".format( + model_dir + ) + if quantize: - model_bb_file = os.path.join(model_dir, 'model_quant.onnx') - model_eb_file = os.path.join(model_dir, 'model_eb_quant.onnx') + model_bb_file = os.path.join(model_dir, "model_quant.onnx") + model_eb_file = os.path.join(model_dir, "model_eb_quant.onnx") else: - model_bb_file = os.path.join(model_dir, 'model.onnx') - model_eb_file = os.path.join(model_dir, 'model_eb.onnx') + model_bb_file = os.path.join(model_dir, "model.onnx") + model_eb_file = os.path.join(model_dir, "model_eb.onnx") if not (os.path.exists(model_eb_file) and os.path.exists(model_bb_file)): print(".onnx is not exist, begin to export onnx") try: from funasr import AutoModel except: - raise "You are exporting onnx, please install funasr and try it again. To install funasr, you could:\n" \ - "\npip3 install -U funasr\n" \ - "For the users in China, you could install with the command:\n" \ - "\npip3 install -U funasr -i https://mirror.sjtu.edu.cn/pypi/web/simple" + raise "You are exporting onnx, please install funasr and try it again. To install funasr, you could:\n" "\npip3 install -U funasr\n" "For the users in China, you could install with the command:\n" "\npip3 install -U funasr -i https://mirror.sjtu.edu.cn/pypi/web/simple" model = AutoModel(model=model_dir) model_dir = model.export(type="onnx", quantize=quantize, **kwargs) - - config_file = os.path.join(model_dir, 'config.yaml') - cmvn_file = os.path.join(model_dir, 'am.mvn') + + config_file = os.path.join(model_dir, "config.yaml") + cmvn_file = os.path.join(model_dir, "am.mvn") config = read_yaml(config_file) - token_list = os.path.join(model_dir, 'tokens.json') - with open(token_list, 'r', encoding='utf-8') as f: + token_list = os.path.join(model_dir, "tokens.json") + with open(token_list, "r", encoding="utf-8") as f: token_list = json.load(f) - + # revert token_list into vocab dict self.vocab = {} for i, token in enumerate(token_list): - self.vocab[token] = i + self.vocab[token] = i self.converter = TokenIDConverter(token_list) self.tokenizer = CharTokenizer() - self.frontend = WavFrontend( - cmvn_file=cmvn_file, - **config['frontend_conf'] + self.frontend = WavFrontend(cmvn_file=cmvn_file, **config["frontend_conf"]) + self.ort_infer_bb = OrtInferSession( + model_bb_file, device_id, intra_op_num_threads=intra_op_num_threads + ) + self.ort_infer_eb = OrtInferSession( + model_eb_file, device_id, intra_op_num_threads=intra_op_num_threads ) - self.ort_infer_bb = OrtInferSession(model_bb_file, device_id, intra_op_num_threads=intra_op_num_threads) - self.ort_infer_eb = OrtInferSession(model_eb_file, device_id, intra_op_num_threads=intra_op_num_threads) self.batch_size = batch_size self.plot_timestamp_to = plot_timestamp_to - if "predictor_bias" in config['model_conf'].keys(): - self.pred_bias = config['model_conf']['predictor_bias'] + if "predictor_bias" in config["model_conf"].keys(): + self.pred_bias = config["model_conf"]["predictor_bias"] else: self.pred_bias = 0 - def __call__(self, - wav_content: Union[str, np.ndarray, List[str]], - hotwords: str, - **kwargs) -> List: + def __call__( + self, wav_content: Union[str, np.ndarray, List[str]], hotwords: str, **kwargs + ) -> List: # make hotword list hotwords, hotwords_length = self.proc_hotword(hotwords) # import pdb; pdb.set_trace() @@ -333,14 +346,14 @@ class ContextualParaformer(Paraformer): outputs = self.bb_infer(feats, feats_len, bias_embed) am_scores, valid_token_lens = outputs[0], outputs[1] except ONNXRuntimeError: - #logging.warning(traceback.format_exc()) + # logging.warning(traceback.format_exc()) logging.warning("input wav is silence or noise") - preds = [''] + preds = [""] else: preds = self.decode(am_scores, valid_token_lens) for pred in preds: pred = sentence_postprocess(pred) - asr_res.append({'preds': pred}) + asr_res.append({"preds": pred}) return asr_res def proc_hotword(self, hotwords): @@ -348,16 +361,20 @@ class ContextualParaformer(Paraformer): hotwords_length = [len(i) - 1 for i in hotwords] hotwords_length.append(0) hotwords_length = np.array(hotwords_length) + # hotwords.append('') def word_map(word): hotwords = [] for c in word: if c not in self.vocab.keys(): hotwords.append(8403) - logging.warning("oov character {} found in hotword {}, replaced by ".format(c, word)) + logging.warning( + "oov character {} found in hotword {}, replaced by ".format(c, word) + ) else: hotwords.append(self.vocab[c]) return np.array(hotwords) + hotword_int = [word_map(i) for i in hotwords] # import pdb; pdb.set_trace() hotword_int.append(np.array([1])) @@ -365,8 +382,9 @@ class ContextualParaformer(Paraformer): # import pdb; pdb.set_trace() return hotwords, hotwords_length - def bb_infer(self, feats: np.ndarray, - feats_len: np.ndarray, bias_embed) -> Tuple[np.ndarray, np.ndarray]: + def bb_infer( + self, feats: np.ndarray, feats_len: np.ndarray, bias_embed + ) -> Tuple[np.ndarray, np.ndarray]: outputs = self.ort_infer_bb([feats, feats_len, bias_embed]) return outputs @@ -375,12 +393,12 @@ class ContextualParaformer(Paraformer): return outputs def decode(self, am_scores: np.ndarray, token_nums: int) -> List[str]: - return [self.decode_one(am_score, token_num) - for am_score, token_num in zip(am_scores, token_nums)] + return [ + self.decode_one(am_score, token_num) + for am_score, token_num in zip(am_scores, token_nums) + ] - def decode_one(self, - am_score: np.ndarray, - valid_token_num: int) -> List[str]: + def decode_one(self, am_score: np.ndarray, valid_token_num: int) -> List[str]: yseq = am_score.argmax(axis=-1) score = am_score.max(axis=-1) score = np.sum(score, axis=-1) @@ -399,12 +417,12 @@ class ContextualParaformer(Paraformer): # Change integer-ids to tokens token = self.converter.ids2tokens(token_int) - token = token[:valid_token_num-self.pred_bias] + token = token[: valid_token_num - self.pred_bias] # texts = sentence_postprocess(token) return token - - + + class SeacoParaformer(ContextualParaformer): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - # no difference with contextual_paraformer in method of calling onnx models \ No newline at end of file + # no difference with contextual_paraformer in method of calling onnx models diff --git a/runtime/python/onnxruntime/funasr_onnx/paraformer_online_bin.py b/runtime/python/onnxruntime/funasr_onnx/paraformer_online_bin.py index 05563aa79..9b68b2f05 100644 --- a/runtime/python/onnxruntime/funasr_onnx/paraformer_online_bin.py +++ b/runtime/python/onnxruntime/funasr_onnx/paraformer_online_bin.py @@ -8,75 +8,78 @@ import copy import librosa import numpy as np -from .utils.utils import (CharTokenizer, Hypothesis, ONNXRuntimeError, - OrtInferSession, TokenIDConverter, get_logger, - read_yaml) +from .utils.utils import ( + CharTokenizer, + Hypothesis, + ONNXRuntimeError, + OrtInferSession, + TokenIDConverter, + get_logger, + read_yaml, +) from .utils.postprocess_utils import sentence_postprocess from .utils.frontend import WavFrontendOnline, SinusoidalPositionEncoderOnline logging = get_logger() -class Paraformer(): - def __init__(self, model_dir: Union[str, Path] = None, - batch_size: int = 1, - chunk_size: List = [5, 10, 5], - device_id: Union[str, int] = "-1", - quantize: bool = False, - intra_op_num_threads: int = 4, - cache_dir: str = None, - **kwargs - ): +class Paraformer: + def __init__( + self, + model_dir: Union[str, Path] = None, + batch_size: int = 1, + chunk_size: List = [5, 10, 5], + device_id: Union[str, int] = "-1", + quantize: bool = False, + intra_op_num_threads: int = 4, + cache_dir: str = None, + **kwargs, + ): if not Path(model_dir).exists(): try: from modelscope.hub.snapshot_download import snapshot_download except: - raise "You are exporting model from modelscope, please install modelscope and try it again. To install modelscope, you could:\n" \ - "\npip3 install -U modelscope\n" \ - "For the users in China, you could install with the command:\n" \ - "\npip3 install -U modelscope -i https://mirror.sjtu.edu.cn/pypi/web/simple" + raise "You are exporting model from modelscope, please install modelscope and try it again. To install modelscope, you could:\n" "\npip3 install -U modelscope\n" "For the users in China, you could install with the command:\n" "\npip3 install -U modelscope -i https://mirror.sjtu.edu.cn/pypi/web/simple" try: model_dir = snapshot_download(model_dir, cache_dir=cache_dir) except: - raise "model_dir must be model_name in modelscope or local path downloaded from modelscope, but is {}".format(model_dir) - - encoder_model_file = os.path.join(model_dir, 'model.onnx') - decoder_model_file = os.path.join(model_dir, 'decoder.onnx') + raise "model_dir must be model_name in modelscope or local path downloaded from modelscope, but is {}".format( + model_dir + ) + + encoder_model_file = os.path.join(model_dir, "model.onnx") + decoder_model_file = os.path.join(model_dir, "decoder.onnx") if quantize: - encoder_model_file = os.path.join(model_dir, 'model_quant.onnx') - decoder_model_file = os.path.join(model_dir, 'decoder_quant.onnx') + encoder_model_file = os.path.join(model_dir, "model_quant.onnx") + decoder_model_file = os.path.join(model_dir, "decoder_quant.onnx") if not os.path.exists(encoder_model_file) or not os.path.exists(decoder_model_file): print(".onnx is not exist, begin to export onnx") try: from funasr import AutoModel except: - raise "You are exporting onnx, please install funasr and try it again. To install funasr, you could:\n" \ - "\npip3 install -U funasr\n" \ - "For the users in China, you could install with the command:\n" \ - "\npip3 install -U funasr -i https://mirror.sjtu.edu.cn/pypi/web/simple" + raise "You are exporting onnx, please install funasr and try it again. To install funasr, you could:\n" "\npip3 install -U funasr\n" "For the users in China, you could install with the command:\n" "\npip3 install -U funasr -i https://mirror.sjtu.edu.cn/pypi/web/simple" model = AutoModel(model=model_dir) model_dir = model.export(type="onnx", quantize=quantize, **kwargs) - config_file = os.path.join(model_dir, 'config.yaml') - cmvn_file = os.path.join(model_dir, 'am.mvn') + config_file = os.path.join(model_dir, "config.yaml") + cmvn_file = os.path.join(model_dir, "am.mvn") config = read_yaml(config_file) - token_list = os.path.join(model_dir, 'tokens.json') - with open(token_list, 'r', encoding='utf-8') as f: + token_list = os.path.join(model_dir, "tokens.json") + with open(token_list, "r", encoding="utf-8") as f: token_list = json.load(f) self.converter = TokenIDConverter(token_list) self.tokenizer = CharTokenizer() - self.frontend = WavFrontendOnline( - cmvn_file=cmvn_file, - **config['frontend_conf'] - ) + self.frontend = WavFrontendOnline(cmvn_file=cmvn_file, **config["frontend_conf"]) self.pe = SinusoidalPositionEncoderOnline() - self.ort_encoder_infer = OrtInferSession(encoder_model_file, device_id, - intra_op_num_threads=intra_op_num_threads) - self.ort_decoder_infer = OrtInferSession(decoder_model_file, device_id, - intra_op_num_threads=intra_op_num_threads) + self.ort_encoder_infer = OrtInferSession( + encoder_model_file, device_id, intra_op_num_threads=intra_op_num_threads + ) + self.ort_decoder_infer = OrtInferSession( + decoder_model_file, device_id, intra_op_num_threads=intra_op_num_threads + ) self.batch_size = batch_size self.chunk_size = chunk_size self.encoder_output_size = config["encoder_conf"]["output_size"] @@ -95,7 +98,9 @@ class Paraformer(): cache["cif_alphas"] = np.zeros((batch_size, 1)).astype(np.float32) cache["chunk_size"] = self.chunk_size cache["last_chunk"] = False - cache["feats"] = np.zeros((batch_size, self.chunk_size[0] + self.chunk_size[2], self.feats_dims)).astype(np.float32) + cache["feats"] = np.zeros( + (batch_size, self.chunk_size[0] + self.chunk_size[2], self.feats_dims) + ).astype(np.float32) cache["decoder_fsmn"] = [] for i in range(self.fsmn_layer): fsmn_cache = np.zeros((batch_size, self.fsmn_dims, self.fsmn_lorder)).astype(np.float32) @@ -108,31 +113,31 @@ class Paraformer(): # process last chunk overlap_feats = np.concatenate((cache["feats"], feats), axis=1) if cache["is_final"]: - cache["feats"] = overlap_feats[:, -self.chunk_size[0]:, :] + cache["feats"] = overlap_feats[:, -self.chunk_size[0] :, :] if not cache["last_chunk"]: - padding_length = sum(self.chunk_size) - overlap_feats.shape[1] - overlap_feats = np.pad(overlap_feats, ((0, 0), (0, padding_length), (0, 0))) + padding_length = sum(self.chunk_size) - overlap_feats.shape[1] + overlap_feats = np.pad(overlap_feats, ((0, 0), (0, padding_length), (0, 0))) else: - cache["feats"] = overlap_feats[:, -(self.chunk_size[0] + self.chunk_size[2]):, :] + cache["feats"] = overlap_feats[:, -(self.chunk_size[0] + self.chunk_size[2]) :, :] return overlap_feats def __call__(self, audio_in: np.ndarray, **kwargs): waveforms = np.expand_dims(audio_in, axis=0) - param_dict = kwargs.get('param_dict', dict()) - is_final = param_dict.get('is_final', False) - cache = param_dict.get('cache', dict()) + param_dict = kwargs.get("param_dict", dict()) + is_final = param_dict.get("is_final", False) + cache = param_dict.get("cache", dict()) asr_res = [] - + if waveforms.shape[1] < 16 * 60 and is_final and len(cache) > 0: cache["last_chunk"] = True feats = cache["feats"] feats_len = np.array([feats.shape[1]]).astype(np.int32) asr_res = self.infer(feats, feats_len, cache) return asr_res - + feats, feats_len = self.extract_feat(waveforms, is_final) if feats.shape[1] != 0: - feats *= self.encoder_output_size ** 0.5 + feats *= self.encoder_output_size**0.5 cache = self.prepare_cache(cache) cache["is_final"] = is_final @@ -145,16 +150,19 @@ class Paraformer(): feats = self.add_overlap_chunk(feats, cache) else: # first chunk - feats_chunk1 = self.add_overlap_chunk(feats[:, :self.chunk_size[1], :], cache) + feats_chunk1 = self.add_overlap_chunk(feats[:, : self.chunk_size[1], :], cache) feats_len = np.array([feats_chunk1.shape[1]]).astype(np.int32) asr_res_chunk1 = self.infer(feats_chunk1, feats_len, cache) # last chunk cache["last_chunk"] = True - feats_chunk2 = self.add_overlap_chunk(feats[:, -(feats.shape[1] + self.chunk_size[2] - self.chunk_size[1]):, :], cache) + feats_chunk2 = self.add_overlap_chunk( + feats[:, -(feats.shape[1] + self.chunk_size[2] - self.chunk_size[1]) :, :], + cache, + ) feats_len = np.array([feats_chunk2.shape[1]]).astype(np.int32) asr_res_chunk2 = self.infer(feats_chunk2, feats_len, cache) - + asr_res_chunk = asr_res_chunk1 + asr_res_chunk2 res = {} for pred in asr_res_chunk: @@ -188,17 +196,18 @@ class Paraformer(): dec_input.extend(cache["decoder_fsmn"]) dec_output = self.ort_decoder_infer(dec_input) logits, sample_ids, cache["decoder_fsmn"] = dec_output[0], dec_output[1], dec_output[2:] - cache["decoder_fsmn"] = [item[:, :, -self.fsmn_lorder:] for item in cache["decoder_fsmn"]] + cache["decoder_fsmn"] = [ + item[:, :, -self.fsmn_lorder :] for item in cache["decoder_fsmn"] + ] preds = self.decode(logits, acoustic_embeds_len) for pred in preds: pred = sentence_postprocess(pred) - asr_res.append({'preds': pred}) + asr_res.append({"preds": pred}) return asr_res - def load_data(self, - wav_content: Union[str, np.ndarray, List[str]], fs: int = None) -> List: + def load_data(self, wav_content: Union[str, np.ndarray, List[str]], fs: int = None) -> List: def load_wav(path: str) -> np.ndarray: waveform, _ = librosa.load(path, sr=fs) return waveform @@ -212,12 +221,11 @@ class Paraformer(): if isinstance(wav_content, list): return [load_wav(path) for path in wav_content] - raise TypeError( - f'The type of {wav_content} is not in [str, np.ndarray, list]') + raise TypeError(f"The type of {wav_content} is not in [str, np.ndarray, list]") - def extract_feat(self, - waveforms: np.ndarray, is_final: bool = False - ) -> Tuple[np.ndarray, np.ndarray]: + def extract_feat( + self, waveforms: np.ndarray, is_final: bool = False + ) -> Tuple[np.ndarray, np.ndarray]: waveforms_lens = np.zeros(waveforms.shape[0]).astype(np.int32) for idx, waveform in enumerate(waveforms): waveforms_lens[idx] = waveform.shape[-1] @@ -226,12 +234,12 @@ class Paraformer(): return feats.astype(np.float32), feats_len.astype(np.int32) def decode(self, am_scores: np.ndarray, token_nums: int) -> List[str]: - return [self.decode_one(am_score, token_num) - for am_score, token_num in zip(am_scores, token_nums)] + return [ + self.decode_one(am_score, token_num) + for am_score, token_num in zip(am_scores, token_nums) + ] - def decode_one(self, - am_score: np.ndarray, - valid_token_num: int) -> List[str]: + def decode_one(self, am_score: np.ndarray, valid_token_num: int) -> List[str]: yseq = am_score.argmax(axis=-1) score = am_score.max(axis=-1) score = np.sum(score, axis=-1) @@ -261,15 +269,15 @@ class Paraformer(): list_frames = [] cache_alphas = [] cache_hiddens = [] - alphas[:, :self.chunk_size[0]] = 0.0 - alphas[:, sum(self.chunk_size[:2]):] = 0.0 + alphas[:, : self.chunk_size[0]] = 0.0 + alphas[:, sum(self.chunk_size[:2]) :] = 0.0 if cache is not None and "cif_alphas" in cache and "cif_hidden" in cache: hidden = np.concatenate((cache["cif_hidden"], hidden), axis=1) alphas = np.concatenate((cache["cif_alphas"], alphas), axis=1) if cache is not None and "last_chunk" in cache and cache["last_chunk"]: tail_hidden = np.zeros((batch_size, 1, hidden_size)).astype(np.float32) tail_alphas = np.array([[self.tail_threshold]]).astype(np.float32) - tail_alphas =np.tile(tail_alphas, (batch_size, 1)) + tail_alphas = np.tile(tail_alphas, (batch_size, 1)) hidden = np.concatenate((hidden, tail_hidden), axis=1) alphas = np.concatenate((alphas, tail_alphas), axis=1) @@ -317,5 +325,6 @@ class Paraformer(): cache["cif_hidden"] = np.stack(cache_hiddens, axis=0) cache["cif_hidden"] = np.expand_dims(cache["cif_hidden"], axis=0) - return np.stack(list_ls, axis=0).astype(np.float32), np.stack(token_length, axis=0).astype(np.int32) - + return np.stack(list_ls, axis=0).astype(np.float32), np.stack(token_length, axis=0).astype( + np.int32 + ) diff --git a/runtime/python/onnxruntime/funasr_onnx/punc_bin.py b/runtime/python/onnxruntime/funasr_onnx/punc_bin.py index db45baaaf..6208c09e0 100644 --- a/runtime/python/onnxruntime/funasr_onnx/punc_bin.py +++ b/runtime/python/onnxruntime/funasr_onnx/punc_bin.py @@ -7,68 +7,72 @@ from pathlib import Path from typing import List, Union, Tuple import numpy as np import json -from .utils.utils import (ONNXRuntimeError, - OrtInferSession, get_logger, - read_yaml) -from .utils.utils import (TokenIDConverter, split_to_mini_sentence,code_mix_split_words,code_mix_split_words_jieba) +from .utils.utils import ONNXRuntimeError, OrtInferSession, get_logger, read_yaml +from .utils.utils import ( + TokenIDConverter, + split_to_mini_sentence, + code_mix_split_words, + code_mix_split_words_jieba, +) + logging = get_logger() -class CT_Transformer(): +class CT_Transformer: """ Author: Speech Lab of DAMO Academy, Alibaba Group CT-Transformer: Controllable time-delay transformer for real-time punctuation prediction and disfluency detection https://arxiv.org/pdf/2003.01309.pdf """ - def __init__(self, model_dir: Union[str, Path] = None, - batch_size: int = 1, - device_id: Union[str, int] = "-1", - quantize: bool = False, - intra_op_num_threads: int = 4, - cache_dir: str = None, - **kwargs - ): - + + def __init__( + self, + model_dir: Union[str, Path] = None, + batch_size: int = 1, + device_id: Union[str, int] = "-1", + quantize: bool = False, + intra_op_num_threads: int = 4, + cache_dir: str = None, + **kwargs + ): + if not Path(model_dir).exists(): try: from modelscope.hub.snapshot_download import snapshot_download except: - raise "You are exporting model from modelscope, please install modelscope and try it again. To install modelscope, you could:\n" \ - "\npip3 install -U modelscope\n" \ - "For the users in China, you could install with the command:\n" \ - "\npip3 install -U modelscope -i https://mirror.sjtu.edu.cn/pypi/web/simple" + raise "You are exporting model from modelscope, please install modelscope and try it again. To install modelscope, you could:\n" "\npip3 install -U modelscope\n" "For the users in China, you could install with the command:\n" "\npip3 install -U modelscope -i https://mirror.sjtu.edu.cn/pypi/web/simple" try: model_dir = snapshot_download(model_dir, cache_dir=cache_dir) except: raise "model_dir must be model_name in modelscope or local path downloaded from modelscope, but is {}".format( - model_dir) - - model_file = os.path.join(model_dir, 'model.onnx') + model_dir + ) + + model_file = os.path.join(model_dir, "model.onnx") if quantize: - model_file = os.path.join(model_dir, 'model_quant.onnx') + model_file = os.path.join(model_dir, "model_quant.onnx") if not os.path.exists(model_file): print(".onnx is not exist, begin to export onnx") try: from funasr import AutoModel except: - raise "You are exporting onnx, please install funasr and try it again. To install funasr, you could:\n" \ - "\npip3 install -U funasr\n" \ - "For the users in China, you could install with the command:\n" \ - "\npip3 install -U funasr -i https://mirror.sjtu.edu.cn/pypi/web/simple" + raise "You are exporting onnx, please install funasr and try it again. To install funasr, you could:\n" "\npip3 install -U funasr\n" "For the users in China, you could install with the command:\n" "\npip3 install -U funasr -i https://mirror.sjtu.edu.cn/pypi/web/simple" model = AutoModel(model=model_dir) model_dir = model.export(type="onnx", quantize=quantize, **kwargs) - - config_file = os.path.join(model_dir, 'config.yaml') + + config_file = os.path.join(model_dir, "config.yaml") config = read_yaml(config_file) - token_list = os.path.join(model_dir, 'tokens.json') - with open(token_list, 'r', encoding='utf-8') as f: + token_list = os.path.join(model_dir, "tokens.json") + with open(token_list, "r", encoding="utf-8") as f: token_list = json.load(f) self.converter = TokenIDConverter(token_list) - self.ort_infer = OrtInferSession(model_file, device_id, intra_op_num_threads=intra_op_num_threads) + self.ort_infer = OrtInferSession( + model_file, device_id, intra_op_num_threads=intra_op_num_threads + ) self.batch_size = 1 - self.punc_list = config["model_conf"]['punc_list'] + self.punc_list = config["model_conf"]["punc_list"] self.period = 0 for i in range(len(self.punc_list)): if self.punc_list[i] == ",": @@ -77,7 +81,7 @@ class CT_Transformer(): self.punc_list[i] = "？" elif self.punc_list[i] == "。": self.period = i - self.jieba_usr_dict_path = os.path.join(model_dir, 'jieba_usr_dict') + self.jieba_usr_dict_path = os.path.join(model_dir, "jieba_usr_dict") if os.path.exists(self.jieba_usr_dict_path): self.seg_jieba = True self.code_mix_split_words_jieba = code_mix_split_words_jieba(self.jieba_usr_dict_path) @@ -102,15 +106,15 @@ class CT_Transformer(): mini_sentence = mini_sentences[mini_sentence_i] mini_sentence_id = mini_sentences_id[mini_sentence_i] mini_sentence = cache_sent + mini_sentence - mini_sentence_id = np.array(cache_sent_id + mini_sentence_id, dtype='int32') + mini_sentence_id = np.array(cache_sent_id + mini_sentence_id, dtype="int32") data = { - "text": mini_sentence_id[None,:], - "text_lengths": np.array([len(mini_sentence_id)], dtype='int32'), + "text": mini_sentence_id[None, :], + "text_lengths": np.array([len(mini_sentence_id)], dtype="int32"), } try: - outputs = self.infer(data['text'], data['text_lengths']) + outputs = self.infer(data["text"], data["text_lengths"]) y = outputs[0] - punctuations = np.argmax(y,axis=-1)[0] + punctuations = np.argmax(y, axis=-1)[0] assert punctuations.size == len(mini_sentence) except ONNXRuntimeError: logging.warning("error") @@ -120,26 +124,36 @@ class CT_Transformer(): sentenceEnd = -1 last_comma_index = -1 for i in range(len(punctuations) - 2, 1, -1): - if self.punc_list[punctuations[i]] == "。" or self.punc_list[punctuations[i]] == "？": + if ( + self.punc_list[punctuations[i]] == "。" + or self.punc_list[punctuations[i]] == "？" + ): sentenceEnd = i break if last_comma_index < 0 and self.punc_list[punctuations[i]] == "，": last_comma_index = i - if sentenceEnd < 0 and len(mini_sentence) > cache_pop_trigger_limit and last_comma_index >= 0: + if ( + sentenceEnd < 0 + and len(mini_sentence) > cache_pop_trigger_limit + and last_comma_index >= 0 + ): # The sentence it too long, cut off at a comma. sentenceEnd = last_comma_index punctuations[sentenceEnd] = self.period - cache_sent = mini_sentence[sentenceEnd + 1:] - cache_sent_id = mini_sentence_id[sentenceEnd + 1:].tolist() - mini_sentence = mini_sentence[0:sentenceEnd + 1] - punctuations = punctuations[0:sentenceEnd + 1] + cache_sent = mini_sentence[sentenceEnd + 1 :] + cache_sent_id = mini_sentence_id[sentenceEnd + 1 :].tolist() + mini_sentence = mini_sentence[0 : sentenceEnd + 1] + punctuations = punctuations[0 : sentenceEnd + 1] new_mini_sentence_punc += [int(x) for x in punctuations] words_with_punc = [] for i in range(len(mini_sentence)): if i > 0: - if len(mini_sentence[i][0].encode()) == 1 and len(mini_sentence[i - 1][0].encode()) == 1: + if ( + len(mini_sentence[i][0].encode()) == 1 + and len(mini_sentence[i - 1][0].encode()) == 1 + ): mini_sentence[i] = " " + mini_sentence[i] words_with_punc.append(mini_sentence[i]) if self.punc_list[punctuations[i]] != "_": @@ -157,8 +171,7 @@ class CT_Transformer(): new_mini_sentence_punc_out = new_mini_sentence_punc[:-1] + [self.period] return new_mini_sentence_out, new_mini_sentence_punc_out - def infer(self, feats: np.ndarray, - feats_len: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + def infer(self, feats: np.ndarray, feats_len: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: outputs = self.ort_infer([feats, feats_len]) return outputs @@ -169,8 +182,8 @@ class CT_Transformer_VadRealtime(CT_Transformer): CT-Transformer: Controllable time-delay transformer for real-time punctuation prediction and disfluency detection https://arxiv.org/pdf/2003.01309.pdf """ - def __init__(self, *args, **kwargs - ): + + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) def __call__(self, text: str, param_dict: map, split_size=20): @@ -191,7 +204,7 @@ class CT_Transformer_VadRealtime(CT_Transformer): assert len(mini_sentences) == len(mini_sentences_id) cache_sent = [] - cache_sent_id = np.array([], dtype='int32') + cache_sent_id = np.array([], dtype="int32") sentence_punc_list = [] sentence_words_list = [] cache_pop_trigger_limit = 200 @@ -200,19 +213,23 @@ class CT_Transformer_VadRealtime(CT_Transformer): mini_sentence = mini_sentences[mini_sentence_i] mini_sentence_id = mini_sentences_id[mini_sentence_i] mini_sentence = cache_sent + mini_sentence - mini_sentence_id = np.concatenate((cache_sent_id, mini_sentence_id), axis=0,dtype='int32') + mini_sentence_id = np.concatenate( + (cache_sent_id, mini_sentence_id), axis=0, dtype="int32" + ) text_length = len(mini_sentence_id) vad_mask = self.vad_mask(text_length, len(cache))[None, None, :, :].astype(np.float32) data = { - "input": mini_sentence_id[None,:], - "text_lengths": np.array([text_length], dtype='int32'), + "input": mini_sentence_id[None, :], + "text_lengths": np.array([text_length], dtype="int32"), "vad_mask": vad_mask, - "sub_masks": vad_mask + "sub_masks": vad_mask, } try: - outputs = self.infer(data['input'], data['text_lengths'], data['vad_mask'], data["sub_masks"]) + outputs = self.infer( + data["input"], data["text_lengths"], data["vad_mask"], data["sub_masks"] + ) y = outputs[0] - punctuations = np.argmax(y,axis=-1)[0] + punctuations = np.argmax(y, axis=-1)[0] assert punctuations.size == len(mini_sentence) except ONNXRuntimeError: logging.warning("error") @@ -222,20 +239,27 @@ class CT_Transformer_VadRealtime(CT_Transformer): sentenceEnd = -1 last_comma_index = -1 for i in range(len(punctuations) - 2, 1, -1): - if self.punc_list[punctuations[i]] == "。" or self.punc_list[punctuations[i]] == "？": + if ( + self.punc_list[punctuations[i]] == "。" + or self.punc_list[punctuations[i]] == "？" + ): sentenceEnd = i break if last_comma_index < 0 and self.punc_list[punctuations[i]] == "，": last_comma_index = i - if sentenceEnd < 0 and len(mini_sentence) > cache_pop_trigger_limit and last_comma_index >= 0: + if ( + sentenceEnd < 0 + and len(mini_sentence) > cache_pop_trigger_limit + and last_comma_index >= 0 + ): # The sentence it too long, cut off at a comma. sentenceEnd = last_comma_index punctuations[sentenceEnd] = self.period - cache_sent = mini_sentence[sentenceEnd + 1:] - cache_sent_id = mini_sentence_id[sentenceEnd + 1:] - mini_sentence = mini_sentence[0:sentenceEnd + 1] - punctuations = punctuations[0:sentenceEnd + 1] + cache_sent = mini_sentence[sentenceEnd + 1 :] + cache_sent_id = mini_sentence_id[sentenceEnd + 1 :] + mini_sentence = mini_sentence[0 : sentenceEnd + 1] + punctuations = punctuations[0 : sentenceEnd + 1] punctuations_np = [int(x) for x in punctuations] new_mini_sentence_punc += punctuations_np @@ -247,7 +271,10 @@ class CT_Transformer_VadRealtime(CT_Transformer): sentence_punc_list_out = [] for i in range(0, len(sentence_words_list)): if i > 0: - if len(sentence_words_list[i][0].encode()) == 1 and len(sentence_words_list[i - 1][-1].encode()) == 1: + if ( + len(sentence_words_list[i][0].encode()) == 1 + and len(sentence_words_list[i - 1][-1].encode()) == 1 + ): sentence_words_list[i] = " " + sentence_words_list[i] if skip_num < len(cache): skip_num += 1 @@ -264,7 +291,7 @@ class CT_Transformer_VadRealtime(CT_Transformer): if sentence_punc_list[i] == "。" or sentence_punc_list[i] == "？": sentenceEnd = i break - cache_out = sentence_words_list[sentenceEnd + 1:] + cache_out = sentence_words_list[sentenceEnd + 1 :] if sentence_out[-1] in self.punc_list: sentence_out = sentence_out[:-1] sentence_punc_list_out[-1] = "_" @@ -282,15 +309,12 @@ class CT_Transformer_VadRealtime(CT_Transformer): ret = np.ones((size, size), dtype=dtype) if vad_pos <= 0 or vad_pos >= size: return ret - sub_corner = np.zeros( - (vad_pos - 1, size - vad_pos), dtype=dtype) - ret[0:vad_pos - 1, vad_pos:] = sub_corner + sub_corner = np.zeros((vad_pos - 1, size - vad_pos), dtype=dtype) + ret[0 : vad_pos - 1, vad_pos:] = sub_corner return ret - def infer(self, feats: np.ndarray, - feats_len: np.ndarray, - vad_mask: np.ndarray, - sub_masks: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + def infer( + self, feats: np.ndarray, feats_len: np.ndarray, vad_mask: np.ndarray, sub_masks: np.ndarray + ) -> Tuple[np.ndarray, np.ndarray]: outputs = self.ort_infer([feats, feats_len, vad_mask, sub_masks]) return outputs - diff --git a/runtime/python/onnxruntime/funasr_onnx/utils/e2e_vad.py b/runtime/python/onnxruntime/funasr_onnx/utils/e2e_vad.py index 3cda80da4..41e357ee7 100644 --- a/runtime/python/onnxruntime/funasr_onnx/utils/e2e_vad.py +++ b/runtime/python/onnxruntime/funasr_onnx/utils/e2e_vad.py @@ -8,6 +8,7 @@ from typing import List, Tuple, Dict, Any import math import numpy as np + class VadStateMachine(Enum): kVadInStateStartPointNotDetected = 1 kVadInStateInSpeechSegment = 2 @@ -37,36 +38,36 @@ class VadDetectMode(Enum): class VADXOptions: def __init__( - self, - sample_rate: int = 16000, - detect_mode: int = VadDetectMode.kVadMutipleUtteranceDetectMode.value, - snr_mode: int = 0, - max_end_silence_time: int = 800, - max_start_silence_time: int = 3000, - do_start_point_detection: bool = True, - do_end_point_detection: bool = True, - window_size_ms: int = 200, - sil_to_speech_time_thres: int = 150, - speech_to_sil_time_thres: int = 150, - speech_2_noise_ratio: float = 1.0, - do_extend: int = 1, - lookback_time_start_point: int = 200, - lookahead_time_end_point: int = 100, - max_single_segment_time: int = 60000, - nn_eval_block_size: int = 8, - dcd_block_size: int = 4, - snr_thres: int = -100.0, - noise_frame_num_used_for_snr: int = 100, - decibel_thres: int = -100.0, - speech_noise_thres: float = 0.6, - fe_prior_thres: float = 1e-4, - silence_pdf_num: int = 1, - sil_pdf_ids: List[int] = [0], - speech_noise_thresh_low: float = -0.1, - speech_noise_thresh_high: float = 0.3, - output_frame_probs: bool = False, - frame_in_ms: int = 10, - frame_length_ms: int = 25, + self, + sample_rate: int = 16000, + detect_mode: int = VadDetectMode.kVadMutipleUtteranceDetectMode.value, + snr_mode: int = 0, + max_end_silence_time: int = 800, + max_start_silence_time: int = 3000, + do_start_point_detection: bool = True, + do_end_point_detection: bool = True, + window_size_ms: int = 200, + sil_to_speech_time_thres: int = 150, + speech_to_sil_time_thres: int = 150, + speech_2_noise_ratio: float = 1.0, + do_extend: int = 1, + lookback_time_start_point: int = 200, + lookahead_time_end_point: int = 100, + max_single_segment_time: int = 60000, + nn_eval_block_size: int = 8, + dcd_block_size: int = 4, + snr_thres: int = -100.0, + noise_frame_num_used_for_snr: int = 100, + decibel_thres: int = -100.0, + speech_noise_thres: float = 0.6, + fe_prior_thres: float = 1e-4, + silence_pdf_num: int = 1, + sil_pdf_ids: List[int] = [0], + speech_noise_thresh_low: float = -0.1, + speech_noise_thresh_high: float = 0.3, + output_frame_probs: bool = False, + frame_in_ms: int = 10, + frame_length_ms: int = 25, ): self.sample_rate = sample_rate self.detect_mode = detect_mode @@ -127,8 +128,13 @@ class E2EVadFrameProb(object): class WindowDetector(object): - def __init__(self, window_size_ms: int, sil_to_speech_time: int, - speech_to_sil_time: int, frame_size_ms: int): + def __init__( + self, + window_size_ms: int, + sil_to_speech_time: int, + speech_to_sil_time: int, + frame_size_ms: int, + ): self.window_size_ms = window_size_ms self.sil_to_speech_time = sil_to_speech_time self.speech_to_sil_time = speech_to_sil_time @@ -174,11 +180,17 @@ class WindowDetector(object): self.win_state[self.cur_win_pos] = cur_frame_state self.cur_win_pos = (self.cur_win_pos + 1) % self.win_size_frame - if self.pre_frame_state == FrameState.kFrameStateSil and self.win_sum >= self.sil_to_speech_frmcnt_thres: + if ( + self.pre_frame_state == FrameState.kFrameStateSil + and self.win_sum >= self.sil_to_speech_frmcnt_thres + ): self.pre_frame_state = FrameState.kFrameStateSpeech return AudioChangeState.kChangeStateSil2Speech - if self.pre_frame_state == FrameState.kFrameStateSpeech and self.win_sum <= self.speech_to_sil_frmcnt_thres: + if ( + self.pre_frame_state == FrameState.kFrameStateSpeech + and self.win_sum <= self.speech_to_sil_frmcnt_thres + ): self.pre_frame_state = FrameState.kFrameStateSil return AudioChangeState.kChangeStateSpeech2Sil @@ -192,19 +204,22 @@ class WindowDetector(object): return int(self.frame_size_ms) -class E2EVadModel(): +class E2EVadModel: """ Author: Speech Lab of DAMO Academy, Alibaba Group Deep-FSMN for Large Vocabulary Continuous Speech Recognition https://arxiv.org/abs/1803.05030 """ + def __init__(self, vad_post_args: Dict[str, Any]): super(E2EVadModel, self).__init__() self.vad_opts = VADXOptions(**vad_post_args) - self.windows_detector = WindowDetector(self.vad_opts.window_size_ms, - self.vad_opts.sil_to_speech_time_thres, - self.vad_opts.speech_to_sil_time_thres, - self.vad_opts.frame_in_ms) + self.windows_detector = WindowDetector( + self.vad_opts.window_size_ms, + self.vad_opts.sil_to_speech_time_thres, + self.vad_opts.speech_to_sil_time_thres, + self.vad_opts.frame_in_ms, + ) # self.encoder = encoder # init variables self.is_final = False @@ -226,7 +241,9 @@ class E2EVadModel(): self.output_data_buf = [] self.output_data_buf_offset = 0 self.frame_probs = [] - self.max_end_sil_frame_cnt_thresh = self.vad_opts.max_end_silence_time - self.vad_opts.speech_to_sil_time_thres + self.max_end_sil_frame_cnt_thresh = ( + self.vad_opts.max_end_silence_time - self.vad_opts.speech_to_sil_time_thres + ) self.speech_noise_thres = self.vad_opts.speech_noise_thres self.scores = None self.idx_pre_chunk = 0 @@ -257,7 +274,9 @@ class E2EVadModel(): self.output_data_buf = [] self.output_data_buf_offset = 0 self.frame_probs = [] - self.max_end_sil_frame_cnt_thresh = self.vad_opts.max_end_silence_time - self.vad_opts.speech_to_sil_time_thres + self.max_end_sil_frame_cnt_thresh = ( + self.vad_opts.max_end_silence_time - self.vad_opts.speech_to_sil_time_thres + ) self.speech_noise_thres = self.vad_opts.speech_noise_thres self.scores = None self.idx_pre_chunk = 0 @@ -287,36 +306,58 @@ class E2EVadModel(): self.data_buf_size = self.data_buf_all_size else: self.data_buf_all_size += len(self.waveform[0]) - for offset in range(0, self.waveform.shape[1] - frame_sample_length + 1, frame_shift_length): + for offset in range( + 0, self.waveform.shape[1] - frame_sample_length + 1, frame_shift_length + ): self.decibel.append( - 10 * math.log10(np.square((self.waveform[0][offset: offset + frame_sample_length])).sum() + \ - 0.000001)) + 10 + * math.log10( + np.square((self.waveform[0][offset : offset + frame_sample_length])).sum() + + 0.000001 + ) + ) def ComputeScores(self, scores: np.ndarray) -> None: # scores = self.encoder(feats, in_cache) # return B * T * D self.vad_opts.nn_eval_block_size = scores.shape[1] self.frm_cnt += scores.shape[1] # count total frames - self.scores=scores + self.scores = scores def PopDataBufTillFrame(self, frame_idx: int) -> None: # need check again while self.data_buf_start_frame < frame_idx: - if self.data_buf_size >= int(self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000): + if self.data_buf_size >= int( + self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000 + ): self.data_buf_start_frame += 1 - self.data_buf_size = self.data_buf_all_size-self.data_buf_start_frame * int( - self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000) + self.data_buf_size = self.data_buf_all_size - self.data_buf_start_frame * int( + self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000 + ) - def PopDataToOutputBuf(self, start_frm: int, frm_cnt: int, first_frm_is_start_point: bool, - last_frm_is_end_point: bool, end_point_is_sent_end: bool) -> None: + def PopDataToOutputBuf( + self, + start_frm: int, + frm_cnt: int, + first_frm_is_start_point: bool, + last_frm_is_end_point: bool, + end_point_is_sent_end: bool, + ) -> None: self.PopDataBufTillFrame(start_frm) - expected_sample_number = int(frm_cnt * self.vad_opts.sample_rate * self.vad_opts.frame_in_ms / 1000) + expected_sample_number = int( + frm_cnt * self.vad_opts.sample_rate * self.vad_opts.frame_in_ms / 1000 + ) if last_frm_is_end_point: - extra_sample = max(0, int(self.vad_opts.frame_length_ms * self.vad_opts.sample_rate / 1000 - \ - self.vad_opts.sample_rate * self.vad_opts.frame_in_ms / 1000)) + extra_sample = max( + 0, + int( + self.vad_opts.frame_length_ms * self.vad_opts.sample_rate / 1000 + - self.vad_opts.sample_rate * self.vad_opts.frame_in_ms / 1000 + ), + ) expected_sample_number += int(extra_sample) if end_point_is_sent_end: expected_sample_number = max(expected_sample_number, self.data_buf_size) if self.data_buf_size < expected_sample_number: - print('error in calling pop data_buf\n') + print("error in calling pop data_buf\n") if len(self.output_data_buf) == 0 or first_frm_is_start_point: self.output_data_buf.append(E2EVadSpeechBufWithDoa()) @@ -326,15 +367,17 @@ class E2EVadModel(): self.output_data_buf[-1].doa = 0 cur_seg = self.output_data_buf[-1] if cur_seg.end_ms != start_frm * self.vad_opts.frame_in_ms: - print('warning\n') + print("warning\n") out_pos = len(cur_seg.buffer) # cur_seg.buff现在没做任何操作 data_to_pop = 0 if end_point_is_sent_end: data_to_pop = expected_sample_number else: - data_to_pop = int(frm_cnt * self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000) + data_to_pop = int( + frm_cnt * self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000 + ) if data_to_pop > self.data_buf_size: - print('VAD data_to_pop is bigger than self.data_buf_size!!!\n') + print("VAD data_to_pop is bigger than self.data_buf_size!!!\n") data_to_pop = self.data_buf_size expected_sample_number = self.data_buf_size @@ -346,7 +389,7 @@ class E2EVadModel(): # cur_seg.buffer[out_pos++] = data_buf_.back() out_pos += 1 if cur_seg.end_ms != start_frm * self.vad_opts.frame_in_ms: - print('Something wrong with the VAD algorithm\n') + print("Something wrong with the VAD algorithm\n") self.data_buf_start_frame += frm_cnt cur_seg.end_ms = (start_frm + frm_cnt) * self.vad_opts.frame_in_ms if first_frm_is_start_point: @@ -369,11 +412,14 @@ class E2EVadModel(): if self.vad_opts.do_start_point_detection: pass if self.confirmed_start_frame != -1: - print('not reset vad properly\n') + print("not reset vad properly\n") else: self.confirmed_start_frame = start_frame - if not fake_result and self.vad_state_machine == VadStateMachine.kVadInStateStartPointNotDetected: + if ( + not fake_result + and self.vad_state_machine == VadStateMachine.kVadInStateStartPointNotDetected + ): self.PopDataToOutputBuf(self.confirmed_start_frame, 1, True, False, False) def OnVoiceEnd(self, end_frame: int, fake_result: bool, is_last_frame: bool) -> None: @@ -382,7 +428,7 @@ class E2EVadModel(): if self.vad_opts.do_end_point_detection: pass if self.confirmed_end_frame != -1: - print('not reset vad properly\n') + print("not reset vad properly\n") else: self.confirmed_end_frame = end_frame if not fake_result: @@ -419,7 +465,10 @@ class E2EVadModel(): assert len(self.sil_pdf_ids) == self.vad_opts.silence_pdf_num if len(self.sil_pdf_ids) > 0: assert len(self.scores) == 1 # 只支持batch_size = 1的测试 - sil_pdf_scores = [self.scores[0][t - self.idx_pre_chunk][sil_pdf_id] for sil_pdf_id in self.sil_pdf_ids] + sil_pdf_scores = [ + self.scores[0][t - self.idx_pre_chunk][sil_pdf_id] + for sil_pdf_id in self.sil_pdf_ids + ] sum_score = sum(sil_pdf_scores) noise_prob = math.log(sum_score) * self.vad_opts.speech_2_noise_ratio total_score = 1.0 @@ -442,15 +491,21 @@ class E2EVadModel(): if self.noise_average_decibel < -99.9: self.noise_average_decibel = cur_decibel else: - self.noise_average_decibel = (cur_decibel + self.noise_average_decibel * ( - self.vad_opts.noise_frame_num_used_for_snr - - 1)) / self.vad_opts.noise_frame_num_used_for_snr + self.noise_average_decibel = ( + cur_decibel + + self.noise_average_decibel * (self.vad_opts.noise_frame_num_used_for_snr - 1) + ) / self.vad_opts.noise_frame_num_used_for_snr return frame_state - def __call__(self, score: np.ndarray, waveform: np.ndarray, - is_final: bool = False, max_end_sil: int = 800, online: bool = False - ): + def __call__( + self, + score: np.ndarray, + waveform: np.ndarray, + is_final: bool = False, + max_end_sil: int = 800, + online: bool = False, + ): self.max_end_sil_frame_cnt_thresh = max_end_sil - self.vad_opts.speech_to_sil_time_thres self.waveform = waveform # compute decibel for each frame self.ComputeDecibel() @@ -478,8 +533,10 @@ class E2EVadModel(): end_ms = -1 self.next_seg = False else: - if not is_final and (not self.output_data_buf[i].contain_seg_start_point or not self.output_data_buf[ - i].contain_seg_end_point): + if not is_final and ( + not self.output_data_buf[i].contain_seg_start_point + or not self.output_data_buf[i].contain_seg_end_point + ): continue start_ms = self.output_data_buf[i].start_ms end_ms = self.output_data_buf[i].end_ms @@ -517,7 +574,9 @@ class E2EVadModel(): return 0 - def DetectOneFrame(self, cur_frm_state: FrameState, cur_frm_idx: int, is_final_frame: bool) -> None: + def DetectOneFrame( + self, cur_frm_state: FrameState, cur_frm_idx: int, is_final_frame: bool + ) -> None: tmp_cur_frm_state = FrameState.kFrameStateInvalid if cur_frm_state == FrameState.kFrameStateSpeech: if math.fabs(1.0) > self.vad_opts.fe_prior_thres: @@ -534,7 +593,9 @@ class E2EVadModel(): self.pre_end_silence_detected = False start_frame = 0 if self.vad_state_machine == VadStateMachine.kVadInStateStartPointNotDetected: - start_frame = max(self.data_buf_start_frame, cur_frm_idx - self.LatencyFrmNumAtStartPoint()) + start_frame = max( + self.data_buf_start_frame, cur_frm_idx - self.LatencyFrmNumAtStartPoint() + ) self.OnVoiceStart(start_frame) self.vad_state_machine = VadStateMachine.kVadInStateInSpeechSegment for t in range(start_frame + 1, cur_frm_idx + 1): @@ -542,8 +603,10 @@ class E2EVadModel(): elif self.vad_state_machine == VadStateMachine.kVadInStateInSpeechSegment: for t in range(self.latest_confirmed_speech_frame + 1, cur_frm_idx): self.OnVoiceDetected(t) - if cur_frm_idx - self.confirmed_start_frame + 1 > \ - self.vad_opts.max_single_segment_time / frm_shift_in_ms: + if ( + cur_frm_idx - self.confirmed_start_frame + 1 + > self.vad_opts.max_single_segment_time / frm_shift_in_ms + ): self.OnVoiceEnd(cur_frm_idx, False, False) self.vad_state_machine = VadStateMachine.kVadInStateEndPointDetected elif not is_final_frame: @@ -557,8 +620,10 @@ class E2EVadModel(): if self.vad_state_machine == VadStateMachine.kVadInStateStartPointNotDetected: pass elif self.vad_state_machine == VadStateMachine.kVadInStateInSpeechSegment: - if cur_frm_idx - self.confirmed_start_frame + 1 > \ - self.vad_opts.max_single_segment_time / frm_shift_in_ms: + if ( + cur_frm_idx - self.confirmed_start_frame + 1 + > self.vad_opts.max_single_segment_time / frm_shift_in_ms + ): self.OnVoiceEnd(cur_frm_idx, False, False) self.vad_state_machine = VadStateMachine.kVadInStateEndPointDetected elif not is_final_frame: @@ -570,8 +635,10 @@ class E2EVadModel(): elif AudioChangeState.kChangeStateSpeech2Speech == state_change: self.continous_silence_frame_count = 0 if self.vad_state_machine == VadStateMachine.kVadInStateInSpeechSegment: - if cur_frm_idx - self.confirmed_start_frame + 1 > \ - self.vad_opts.max_single_segment_time / frm_shift_in_ms: + if ( + cur_frm_idx - self.confirmed_start_frame + 1 + > self.vad_opts.max_single_segment_time / frm_shift_in_ms + ): self.max_time_out = True self.OnVoiceEnd(cur_frm_idx, False, False) self.vad_state_machine = VadStateMachine.kVadInStateEndPointDetected @@ -585,40 +652,53 @@ class E2EVadModel(): self.continous_silence_frame_count += 1 if self.vad_state_machine == VadStateMachine.kVadInStateStartPointNotDetected: # silence timeout, return zero length decision - if ((self.vad_opts.detect_mode == VadDetectMode.kVadSingleUtteranceDetectMode.value) and ( - self.continous_silence_frame_count * frm_shift_in_ms > self.vad_opts.max_start_silence_time)) \ - or (is_final_frame and self.number_end_time_detected == 0): + if ( + (self.vad_opts.detect_mode == VadDetectMode.kVadSingleUtteranceDetectMode.value) + and ( + self.continous_silence_frame_count * frm_shift_in_ms + > self.vad_opts.max_start_silence_time + ) + ) or (is_final_frame and self.number_end_time_detected == 0): for t in range(self.lastest_confirmed_silence_frame + 1, cur_frm_idx): self.OnSilenceDetected(t) self.OnVoiceStart(0, True) - self.OnVoiceEnd(0, True, False); + self.OnVoiceEnd(0, True, False) self.vad_state_machine = VadStateMachine.kVadInStateEndPointDetected else: if cur_frm_idx >= self.LatencyFrmNumAtStartPoint(): self.OnSilenceDetected(cur_frm_idx - self.LatencyFrmNumAtStartPoint()) elif self.vad_state_machine == VadStateMachine.kVadInStateInSpeechSegment: - if self.continous_silence_frame_count * frm_shift_in_ms >= self.max_end_sil_frame_cnt_thresh: + if ( + self.continous_silence_frame_count * frm_shift_in_ms + >= self.max_end_sil_frame_cnt_thresh + ): lookback_frame = int(self.max_end_sil_frame_cnt_thresh / frm_shift_in_ms) if self.vad_opts.do_extend: - lookback_frame -= int(self.vad_opts.lookahead_time_end_point / frm_shift_in_ms) + lookback_frame -= int( + self.vad_opts.lookahead_time_end_point / frm_shift_in_ms + ) lookback_frame -= 1 lookback_frame = max(0, lookback_frame) self.OnVoiceEnd(cur_frm_idx - lookback_frame, False, False) self.vad_state_machine = VadStateMachine.kVadInStateEndPointDetected - elif cur_frm_idx - self.confirmed_start_frame + 1 > \ - self.vad_opts.max_single_segment_time / frm_shift_in_ms: + elif ( + cur_frm_idx - self.confirmed_start_frame + 1 + > self.vad_opts.max_single_segment_time / frm_shift_in_ms + ): self.OnVoiceEnd(cur_frm_idx, False, False) self.vad_state_machine = VadStateMachine.kVadInStateEndPointDetected elif self.vad_opts.do_extend and not is_final_frame: if self.continous_silence_frame_count <= int( - self.vad_opts.lookahead_time_end_point / frm_shift_in_ms): + self.vad_opts.lookahead_time_end_point / frm_shift_in_ms + ): self.OnVoiceDetected(cur_frm_idx) else: self.MaybeOnVoiceEndIfLastFrame(is_final_frame, cur_frm_idx) else: pass - if self.vad_state_machine == VadStateMachine.kVadInStateEndPointDetected and \ - self.vad_opts.detect_mode == VadDetectMode.kVadMutipleUtteranceDetectMode.value: + if ( + self.vad_state_machine == VadStateMachine.kVadInStateEndPointDetected + and self.vad_opts.detect_mode == VadDetectMode.kVadMutipleUtteranceDetectMode.value + ): self.ResetDetection() - diff --git a/runtime/python/onnxruntime/funasr_onnx/utils/frontend.py b/runtime/python/onnxruntime/funasr_onnx/utils/frontend.py index 295e7b519..7b38f8b76 100644 --- a/runtime/python/onnxruntime/funasr_onnx/utils/frontend.py +++ b/runtime/python/onnxruntime/funasr_onnx/utils/frontend.py @@ -11,22 +11,21 @@ root_dir = Path(__file__).resolve().parent logger_initialized = {} -class WavFrontend(): - """Conventional frontend structure for ASR. - """ +class WavFrontend: + """Conventional frontend structure for ASR.""" def __init__( - self, - cmvn_file: str = None, - fs: int = 16000, - window: str = 'hamming', - n_mels: int = 80, - frame_length: int = 25, - frame_shift: int = 10, - lfr_m: int = 1, - lfr_n: int = 1, - dither: float = 1.0, - **kwargs, + self, + cmvn_file: str = None, + fs: int = 16000, + window: str = "hamming", + n_mels: int = 80, + frame_length: int = 25, + frame_shift: int = 10, + lfr_m: int = 1, + lfr_n: int = 1, + dither: float = 1.0, + **kwargs, ) -> None: opts = knf.FbankOptions() @@ -51,8 +50,7 @@ class WavFrontend(): self.fbank_beg_idx = 0 self.reset_status() - def fbank(self, - waveform: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + def fbank(self, waveform: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: waveform = waveform * (1 << 15) self.fbank_fn = knf.OnlineFbank(self.opts) self.fbank_fn.accept_waveform(self.opts.frame_opts.samp_freq, waveform.tolist()) @@ -64,8 +62,7 @@ class WavFrontend(): feat_len = np.array(mat.shape[0]).astype(np.int32) return feat, feat_len - def fbank_online(self, - waveform: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + def fbank_online(self, waveform: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: waveform = waveform * (1 << 15) # self.fbank_fn = knf.OnlineFbank(self.opts) self.fbank_fn.accept_waveform(self.opts.frame_opts.samp_freq, waveform.tolist()) @@ -103,12 +100,11 @@ class WavFrontend(): T = T + (lfr_m - 1) // 2 for i in range(T_lfr): if lfr_m <= T - i * lfr_n: - LFR_inputs.append( - (inputs[i * lfr_n:i * lfr_n + lfr_m]).reshape(1, -1)) + LFR_inputs.append((inputs[i * lfr_n : i * lfr_n + lfr_m]).reshape(1, -1)) else: # process last LFR frame num_padding = lfr_m - (T - i * lfr_n) - frame = inputs[i * lfr_n:].reshape(-1) + frame = inputs[i * lfr_n :].reshape(-1) for _ in range(num_padding): frame = np.hstack((frame, inputs[-1])) @@ -126,24 +122,26 @@ class WavFrontend(): inputs = (inputs + means) * vars return inputs - def load_cmvn(self,) -> np.ndarray: - with open(self.cmvn_file, 'r', encoding='utf-8') as f: + def load_cmvn( + self, + ) -> np.ndarray: + with open(self.cmvn_file, "r", encoding="utf-8") as f: lines = f.readlines() means_list = [] vars_list = [] for i in range(len(lines)): line_item = lines[i].split() - if line_item[0] == '': + if line_item[0] == "": line_item = lines[i + 1].split() - if line_item[0] == '': - add_shift_line = line_item[3:(len(line_item) - 1)] + if line_item[0] == "": + add_shift_line = line_item[3 : (len(line_item) - 1)] means_list = list(add_shift_line) continue - elif line_item[0] == '': + elif line_item[0] == "": line_item = lines[i + 1].split() - if line_item[0] == '': - rescale_line = line_item[3:(len(line_item) - 1)] + if line_item[0] == "": + rescale_line = line_item[3 : (len(line_item) - 1)] vars_list = list(rescale_line) continue @@ -158,8 +156,12 @@ class WavFrontendOnline(WavFrontend): super().__init__(**kwargs) # self.fbank_fn = knf.OnlineFbank(self.opts) # add variables - self.frame_sample_length = int(self.opts.frame_opts.frame_length_ms * self.opts.frame_opts.samp_freq / 1000) - self.frame_shift_sample_length = int(self.opts.frame_opts.frame_shift_ms * self.opts.frame_opts.samp_freq / 1000) + self.frame_sample_length = int( + self.opts.frame_opts.frame_length_ms * self.opts.frame_opts.samp_freq / 1000 + ) + self.frame_shift_sample_length = int( + self.opts.frame_opts.frame_shift_ms * self.opts.frame_opts.samp_freq / 1000 + ) self.waveform = None self.reserve_waveforms = None self.input_cache = None @@ -167,23 +169,26 @@ class WavFrontendOnline(WavFrontend): @staticmethod # inputs has catted the cache - def apply_lfr(inputs: np.ndarray, lfr_m: int, lfr_n: int, is_final: bool = False) -> Tuple[ - np.ndarray, np.ndarray, int]: + def apply_lfr( + inputs: np.ndarray, lfr_m: int, lfr_n: int, is_final: bool = False + ) -> Tuple[np.ndarray, np.ndarray, int]: """ Apply lfr with data """ LFR_inputs = [] T = inputs.shape[0] # include the right context - T_lfr = int(np.ceil((T - (lfr_m - 1) // 2) / lfr_n)) # minus the right context: (lfr_m - 1) // 2 + T_lfr = int( + np.ceil((T - (lfr_m - 1) // 2) / lfr_n) + ) # minus the right context: (lfr_m - 1) // 2 splice_idx = T_lfr for i in range(T_lfr): if lfr_m <= T - i * lfr_n: - LFR_inputs.append((inputs[i * lfr_n:i * lfr_n + lfr_m]).reshape(1, -1)) + LFR_inputs.append((inputs[i * lfr_n : i * lfr_n + lfr_m]).reshape(1, -1)) else: # process last LFR frame if is_final: num_padding = lfr_m - (T - i * lfr_n) - frame = (inputs[i * lfr_n:]).reshape(-1) + frame = (inputs[i * lfr_n :]).reshape(-1) for _ in range(num_padding): frame = np.hstack((frame, inputs[-1])) LFR_inputs.append(frame) @@ -197,24 +202,27 @@ class WavFrontendOnline(WavFrontend): return LFR_outputs.astype(np.float32), lfr_splice_cache, splice_idx @staticmethod - def compute_frame_num(sample_length: int, frame_sample_length: int, frame_shift_sample_length: int) -> int: + def compute_frame_num( + sample_length: int, frame_sample_length: int, frame_shift_sample_length: int + ) -> int: frame_num = int((sample_length - frame_sample_length) / frame_shift_sample_length + 1) return frame_num if frame_num >= 1 and sample_length >= frame_sample_length else 0 - def fbank( - self, - input: np.ndarray, - input_lengths: np.ndarray + self, input: np.ndarray, input_lengths: np.ndarray ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: self.fbank_fn = knf.OnlineFbank(self.opts) batch_size = input.shape[0] if self.input_cache is None: self.input_cache = np.empty((batch_size, 0), dtype=np.float32) input = np.concatenate((self.input_cache, input), axis=1) - frame_num = self.compute_frame_num(input.shape[-1], self.frame_sample_length, self.frame_shift_sample_length) + frame_num = self.compute_frame_num( + input.shape[-1], self.frame_sample_length, self.frame_shift_sample_length + ) # update self.in_cache - self.input_cache = input[:, -(input.shape[-1] - frame_num * self.frame_shift_sample_length):] + self.input_cache = input[ + :, -(input.shape[-1] - frame_num * self.frame_shift_sample_length) : + ] waveforms = np.empty(0, dtype=np.float32) feats_pad = np.empty(0, dtype=np.float32) feats_lens = np.empty(0, dtype=np.int32) @@ -225,9 +233,15 @@ class WavFrontendOnline(WavFrontend): for i in range(batch_size): waveform = input[i] waveforms.append( - waveform[:((frame_num - 1) * self.frame_shift_sample_length + self.frame_sample_length)]) + waveform[ + : ( + (frame_num - 1) * self.frame_shift_sample_length + + self.frame_sample_length + ) + ] + ) waveform = waveform * (1 << 15) - + self.fbank_fn.accept_waveform(self.opts.frame_opts.samp_freq, waveform.tolist()) frames = self.fbank_fn.num_frames_ready mat = np.empty([frames, self.opts.mel_opts.num_bins]) @@ -249,22 +263,20 @@ class WavFrontendOnline(WavFrontend): return self.fbanks, self.fbanks_lens def lfr_cmvn( - self, - input: np.ndarray, - input_lengths: np.ndarray, - is_final: bool = False + self, input: np.ndarray, input_lengths: np.ndarray, is_final: bool = False ) -> Tuple[np.ndarray, np.ndarray, List[int]]: batch_size = input.shape[0] feats = [] feats_lens = [] lfr_splice_frame_idxs = [] for i in range(batch_size): - mat = input[i, :input_lengths[i], :] + mat = input[i, : input_lengths[i], :] lfr_splice_frame_idx = -1 if self.lfr_m != 1 or self.lfr_n != 1: # update self.lfr_splice_cache in self.apply_lfr - mat, self.lfr_splice_cache[i], lfr_splice_frame_idx = self.apply_lfr(mat, self.lfr_m, self.lfr_n, - is_final) + mat, self.lfr_splice_cache[i], lfr_splice_frame_idx = self.apply_lfr( + mat, self.lfr_m, self.lfr_n, is_final + ) if self.cmvn_file is not None: mat = self.apply_cmvn(mat) feat_length = mat.shape[0] @@ -276,47 +288,70 @@ class WavFrontendOnline(WavFrontend): feats_pad = np.array(feats) return feats_pad, feats_lens, lfr_splice_frame_idxs - def extract_fbank( - self, input: np.ndarray, input_lengths: np.ndarray, is_final: bool = False + self, input: np.ndarray, input_lengths: np.ndarray, is_final: bool = False ) -> Tuple[np.ndarray, np.ndarray]: batch_size = input.shape[0] - assert batch_size == 1, 'we support to extract feature online only when the batch size is equal to 1 now' + assert ( + batch_size == 1 + ), "we support to extract feature online only when the batch size is equal to 1 now" waveforms, feats, feats_lengths = self.fbank(input, input_lengths) # input shape: B T D if feats.shape[0]: - self.waveforms = waveforms if self.reserve_waveforms is None else np.concatenate( - (self.reserve_waveforms, waveforms), axis=1) + self.waveforms = ( + waveforms + if self.reserve_waveforms is None + else np.concatenate((self.reserve_waveforms, waveforms), axis=1) + ) if not self.lfr_splice_cache: for i in range(batch_size): - self.lfr_splice_cache.append(np.expand_dims(feats[i][0, :], axis=0).repeat((self.lfr_m - 1) // 2, axis=0)) - + self.lfr_splice_cache.append( + np.expand_dims(feats[i][0, :], axis=0).repeat((self.lfr_m - 1) // 2, axis=0) + ) + if feats_lengths[0] + self.lfr_splice_cache[0].shape[0] >= self.lfr_m: lfr_splice_cache_np = np.stack(self.lfr_splice_cache) # B T D feats = np.concatenate((lfr_splice_cache_np, feats), axis=1) feats_lengths += lfr_splice_cache_np[0].shape[0] frame_from_waveforms = int( - (self.waveforms.shape[1] - self.frame_sample_length) / self.frame_shift_sample_length + 1) + (self.waveforms.shape[1] - self.frame_sample_length) + / self.frame_shift_sample_length + + 1 + ) minus_frame = (self.lfr_m - 1) // 2 if self.reserve_waveforms is None else 0 - feats, feats_lengths, lfr_splice_frame_idxs = self.lfr_cmvn(feats, feats_lengths, is_final) + feats, feats_lengths, lfr_splice_frame_idxs = self.lfr_cmvn( + feats, feats_lengths, is_final + ) if self.lfr_m == 1: self.reserve_waveforms = None else: reserve_frame_idx = lfr_splice_frame_idxs[0] - minus_frame # print('reserve_frame_idx: ' + str(reserve_frame_idx)) # print('frame_frame: ' + str(frame_from_waveforms)) - self.reserve_waveforms = self.waveforms[:, reserve_frame_idx * self.frame_shift_sample_length:frame_from_waveforms * self.frame_shift_sample_length] - sample_length = (frame_from_waveforms - 1) * self.frame_shift_sample_length + self.frame_sample_length + self.reserve_waveforms = self.waveforms[ + :, + reserve_frame_idx + * self.frame_shift_sample_length : frame_from_waveforms + * self.frame_shift_sample_length, + ] + sample_length = ( + frame_from_waveforms - 1 + ) * self.frame_shift_sample_length + self.frame_sample_length self.waveforms = self.waveforms[:, :sample_length] else: # update self.reserve_waveforms and self.lfr_splice_cache - self.reserve_waveforms = self.waveforms[:, - :-(self.frame_sample_length - self.frame_shift_sample_length)] + self.reserve_waveforms = self.waveforms[ + :, : -(self.frame_sample_length - self.frame_shift_sample_length) + ] for i in range(batch_size): - self.lfr_splice_cache[i] = np.concatenate((self.lfr_splice_cache[i], feats[i]), axis=0) + self.lfr_splice_cache[i] = np.concatenate( + (self.lfr_splice_cache[i], feats[i]), axis=0 + ) return np.empty(0, dtype=np.float32), feats_lengths else: if is_final: - self.waveforms = waveforms if self.reserve_waveforms is None else self.reserve_waveforms + self.waveforms = ( + waveforms if self.reserve_waveforms is None else self.reserve_waveforms + ) feats = np.stack(self.lfr_splice_cache) feats_lengths = np.zeros(batch_size, dtype=np.int32) + feats.shape[1] feats, feats_lengths, _ = self.lfr_cmvn(feats, feats_lengths, is_final) @@ -333,13 +368,14 @@ class WavFrontendOnline(WavFrontend): self.input_cache = None self.lfr_splice_cache = [] + def load_bytes(input): middle_data = np.frombuffer(input, dtype=np.int16) middle_data = np.asarray(middle_data) - if middle_data.dtype.kind not in 'iu': + if middle_data.dtype.kind not in "iu": raise TypeError("'middle_data' must be an array of integers") - dtype = np.dtype('float32') - if dtype.kind != 'f': + dtype = np.dtype("float32") + if dtype.kind != "f": raise TypeError("'dtype' must be a floating point type") i = np.iinfo(middle_data.dtype) @@ -349,9 +385,8 @@ def load_bytes(input): return array -class SinusoidalPositionEncoderOnline(): - '''Streaming Positional encoding. - ''' +class SinusoidalPositionEncoderOnline: + """Streaming Positional encoding.""" def encode(self, positions: np.ndarray = None, depth: int = None, dtype: np.dtype = np.float32): batch_size = positions.shape[0] @@ -365,29 +400,34 @@ class SinusoidalPositionEncoderOnline(): def forward(self, x, start_idx=0): batch_size, timesteps, input_dim = x.shape - positions = np.arange(1, timesteps+1+start_idx)[None, :] + positions = np.arange(1, timesteps + 1 + start_idx)[None, :] position_encoding = self.encode(positions, input_dim, x.dtype) - return x + position_encoding[:, start_idx: start_idx + timesteps] + return x + position_encoding[:, start_idx : start_idx + timesteps] def test(): path = "/nfs/zhifu.gzf/export/damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch/example/asr_example.wav" import librosa + cmvn_file = "/nfs/zhifu.gzf/export/damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch/am.mvn" config_file = "/nfs/zhifu.gzf/export/damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch/config.yaml" from funasr.runtime.python.onnxruntime.rapid_paraformer.utils.utils import read_yaml + config = read_yaml(config_file) waveform, _ = librosa.load(path, sr=None) frontend = WavFrontend( cmvn_file=cmvn_file, - **config['frontend_conf'], + **config["frontend_conf"], ) - speech, _ = frontend.fbank_online(waveform) #1d, (sample,), numpy - feat, feat_len = frontend.lfr_cmvn(speech) # 2d, (frame, 450), np.float32 -> torch, torch.from_numpy(), dtype, (1, frame, 450) - - frontend.reset_status() # clear cache + speech, _ = frontend.fbank_online(waveform) # 1d, (sample,), numpy + feat, feat_len = frontend.lfr_cmvn( + speech + ) # 2d, (frame, 450), np.float32 -> torch, torch.from_numpy(), dtype, (1, frame, 450) + + frontend.reset_status() # clear cache return feat, feat_len -if __name__ == '__main__': + +if __name__ == "__main__": test() diff --git a/runtime/python/onnxruntime/funasr_onnx/utils/postprocess_utils.py b/runtime/python/onnxruntime/funasr_onnx/utils/postprocess_utils.py index 14d6c7687..da8065e56 100644 --- a/runtime/python/onnxruntime/funasr_onnx/utils/postprocess_utils.py +++ b/runtime/python/onnxruntime/funasr_onnx/utils/postprocess_utils.py @@ -8,7 +8,7 @@ from typing import Any, List, Union def isChinese(ch: str): - if '\u4e00' <= ch <= '\u9fff' or '\u0030' <= ch <= '\u0039': + if "\u4e00" <= ch <= "\u9fff" or "\u0030" <= ch <= "\u0039": return True return False @@ -16,9 +16,9 @@ def isChinese(ch: str): def isAllChinese(word: Union[List[Any], str]): word_lists = [] for i in word: - cur = i.replace(' ', '') - cur = cur.replace('', '') - cur = cur.replace('', '') + cur = i.replace(" ", "") + cur = cur.replace("", "") + cur = cur.replace("", "") word_lists.append(cur) if len(word_lists) == 0: @@ -33,9 +33,9 @@ def isAllChinese(word: Union[List[Any], str]): def isAllAlpha(word: Union[List[Any], str]): word_lists = [] for i in word: - cur = i.replace(' ', '') - cur = cur.replace('', '') - cur = cur.replace('', '') + cur = i.replace(" ", "") + cur = cur.replace("", "") + cur = cur.replace("", "") word_lists.append(cur) if len(word_lists) == 0: @@ -64,12 +64,14 @@ def abbr_dispose(words: List[Any], time_stamp: List[List] = None) -> List[Any]: if num <= last_num: continue - if len(words[num]) == 1 and words[num].encode('utf-8').isalpha(): - if num + 1 < words_size and words[ - num + 1] == ' ' and num + 2 < words_size and len( - words[num + - 2]) == 1 and words[num + - 2].encode('utf-8').isalpha(): + if len(words[num]) == 1 and words[num].encode("utf-8").isalpha(): + if ( + num + 1 < words_size + and words[num + 1] == " " + and num + 2 < words_size + and len(words[num + 2]) == 1 + and words[num + 2].encode("utf-8").isalpha() + ): # found the begin of abbr abbr_begin.append(num) num += 2 @@ -77,11 +79,13 @@ def abbr_dispose(words: List[Any], time_stamp: List[List] = None) -> List[Any]: # to find the end of abbr while True: num += 1 - if num < words_size and words[num] == ' ': + if num < words_size and words[num] == " ": num += 1 - if num < words_size and len( - words[num]) == 1 and words[num].encode( - 'utf-8').isalpha(): + if ( + num < words_size + and len(words[num]) == 1 + and words[num].encode("utf-8").isalpha() + ): abbr_end.pop() abbr_end.append(num) last_num = num @@ -91,11 +95,11 @@ def abbr_dispose(words: List[Any], time_stamp: List[List] = None) -> List[Any]: break for num in range(words_size): - if words[num] == ' ': + if words[num] == " ": ts_nums.append(ts_index) else: ts_nums.append(ts_index) - ts_index += 1 + ts_index += 1 last_num = -1 for num in range(words_size): if num <= last_num: @@ -112,7 +116,7 @@ def abbr_dispose(words: List[Any], time_stamp: List[List] = None) -> List[Any]: last_num = num break else: - if words[num].encode('utf-8').isalpha(): + if words[num].encode("utf-8").isalpha(): word_lists.append(words[num].upper()) num += 1 if time_stamp is not None: @@ -120,7 +124,7 @@ def abbr_dispose(words: List[Any], time_stamp: List[List] = None) -> List[Any]: ts_lists.append([begin, end]) else: word_lists.append(words[num]) - if time_stamp is not None and words[num] != ' ': + if time_stamp is not None and words[num] != " ": begin = time_stamp[ts_nums[num]][0] end = time_stamp[ts_nums[num]][1] ts_lists.append([begin, end]) @@ -135,18 +139,18 @@ def abbr_dispose(words: List[Any], time_stamp: List[List] = None) -> List[Any]: def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): middle_lists = [] word_lists = [] - word_item = '' + word_item = "" ts_lists = [] # wash words lists for i in words: - word = '' + word = "" if isinstance(i, str): word = i else: - word = i.decode('utf-8') + word = i.decode("utf-8") - if word in ['', '', '']: + if word in ["", "", ""]: continue else: middle_lists.append(word) @@ -154,7 +158,7 @@ def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): # all chinese characters if isAllChinese(middle_lists): for i, ch in enumerate(middle_lists): - word_lists.append(ch.replace(' ', '')) + word_lists.append(ch.replace(" ", "")) if time_stamp is not None: ts_lists = time_stamp @@ -165,9 +169,9 @@ def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): if ts_flag and time_stamp is not None: begin = time_stamp[i][0] end = time_stamp[i][1] - word = '' - if '@@' in ch: - word = ch.replace('@@', '') + word = "" + if "@@" in ch: + word = ch.replace("@@", "") word_item += word if time_stamp is not None: ts_flag = False @@ -175,8 +179,8 @@ def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): else: word_item += ch word_lists.append(word_item) - word_lists.append(' ') - word_item = '' + word_lists.append(" ") + word_item = "" if time_stamp is not None: ts_flag = True end = time_stamp[i][1] @@ -193,7 +197,7 @@ def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): if ts_flag and time_stamp is not None: begin = time_stamp[i][0] end = time_stamp[i][1] - word = '' + word = "" if isAllChinese(ch): if alpha_blank is True: word_lists.pop() @@ -203,8 +207,8 @@ def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): ts_flag = True ts_lists.append([begin, end]) begin = end - elif '@@' in ch: - word = ch.replace('@@', '') + elif "@@" in ch: + word = ch.replace("@@", "") word_item += word alpha_blank = False if time_stamp is not None: @@ -213,73 +217,74 @@ def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): elif isAllAlpha(ch): word_item += ch word_lists.append(word_item) - word_lists.append(' ') - word_item = '' + word_lists.append(" ") + word_item = "" alpha_blank = True if time_stamp is not None: ts_flag = True - end = time_stamp[i][1] + end = time_stamp[i][1] ts_lists.append([begin, end]) begin = end else: - raise ValueError('invalid character: {}'.format(ch)) + raise ValueError("invalid character: {}".format(ch)) - if time_stamp is not None: + if time_stamp is not None: word_lists, ts_lists = abbr_dispose(word_lists, ts_lists) real_word_lists = [] for ch in word_lists: - if ch != ' ': + if ch != " ": real_word_lists.append(ch) - sentence = ' '.join(real_word_lists).strip() + sentence = " ".join(real_word_lists).strip() return sentence, ts_lists, real_word_lists else: word_lists = abbr_dispose(word_lists) real_word_lists = [] for ch in word_lists: - if ch != ' ': + if ch != " ": real_word_lists.append(ch) - sentence = ''.join(word_lists).strip() + sentence = "".join(word_lists).strip() return sentence, real_word_lists + def sentence_postprocess_sentencepiece(words): middle_lists = [] word_lists = [] - word_item = '' + word_item = "" # wash words lists for i in words: - word = '' + word = "" if isinstance(i, str): word = i else: - word = i.decode('utf-8') + word = i.decode("utf-8") - if word in ['', '', '', '']: + if word in ["", "", "", ""]: continue else: middle_lists.append(word) # all alpha characters for i, ch in enumerate(middle_lists): - word = '' - if '\u2581' in ch and i == 0: - word_item = '' - word = ch.replace('\u2581', '') + word = "" + if "\u2581" in ch and i == 0: + word_item = "" + word = ch.replace("\u2581", "") word_item += word - elif '\u2581' in ch and i != 0: + elif "\u2581" in ch and i != 0: word_lists.append(word_item) - word_lists.append(' ') - word_item = '' - word = ch.replace('\u2581', '') + word_lists.append(" ") + word_item = "" + word = ch.replace("\u2581", "") word_item += word else: word_item += ch if word_item is not None: word_lists.append(word_item) - #word_lists = abbr_dispose(word_lists) + # word_lists = abbr_dispose(word_lists) real_word_lists = [] for ch in word_lists: - if ch != ' ': + if ch != " ": if ch == "i": ch = ch.replace("i", "I") elif ch == "i'm": @@ -289,5 +294,5 @@ def sentence_postprocess_sentencepiece(words): elif ch == "i'll": ch = ch.replace("i'll", "I'll") real_word_lists.append(ch) - sentence = ''.join(word_lists) - return sentence, real_word_lists \ No newline at end of file + sentence = "".join(word_lists) + return sentence, real_word_lists diff --git a/runtime/python/onnxruntime/funasr_onnx/utils/timestamp_utils.py b/runtime/python/onnxruntime/funasr_onnx/utils/timestamp_utils.py index a9ec8cc1b..c8ba5a7e0 100644 --- a/runtime/python/onnxruntime/funasr_onnx/utils/timestamp_utils.py +++ b/runtime/python/onnxruntime/funasr_onnx/utils/timestamp_utils.py @@ -13,40 +13,44 @@ def time_stamp_lfr6_onnx(us_cif_peak, char_list, begin_time=0.0, total_offset=-1 TIME_RATE = 10.0 * 6 / 1000 / 3 # 3 times upsampled cif_peak = us_cif_peak.reshape(-1) num_frames = cif_peak.shape[-1] - if char_list[-1] == '': + if char_list[-1] == "": char_list = char_list[:-1] # char_list = [i for i in text] timestamp_list = [] new_char_list = [] # for bicif model trained with large data, cif2 actually fires when a character starts # so treat the frames between two peaks as the duration of the former token - fire_place = np.where(cif_peak>1.0-1e-4)[0] + total_offset # np format + fire_place = np.where(cif_peak > 1.0 - 1e-4)[0] + total_offset # np format num_peak = len(fire_place) - assert num_peak == len(char_list) + 1 # number of peaks is supposed to be number of tokens + 1 + assert num_peak == len(char_list) + 1 # number of peaks is supposed to be number of tokens + 1 # begin silence if fire_place[0] > START_END_THRESHOLD: # char_list.insert(0, '') - timestamp_list.append([0.0, fire_place[0]*TIME_RATE]) - new_char_list.append('') + timestamp_list.append([0.0, fire_place[0] * TIME_RATE]) + new_char_list.append("") # tokens timestamp - for i in range(len(fire_place)-1): + for i in range(len(fire_place) - 1): new_char_list.append(char_list[i]) - if i == len(fire_place)-2 or MAX_TOKEN_DURATION < 0 or fire_place[i+1] - fire_place[i] < MAX_TOKEN_DURATION: - timestamp_list.append([fire_place[i]*TIME_RATE, fire_place[i+1]*TIME_RATE]) + if ( + i == len(fire_place) - 2 + or MAX_TOKEN_DURATION < 0 + or fire_place[i + 1] - fire_place[i] < MAX_TOKEN_DURATION + ): + timestamp_list.append([fire_place[i] * TIME_RATE, fire_place[i + 1] * TIME_RATE]) else: # cut the duration to token and sil of the 0-weight frames last long _split = fire_place[i] + MAX_TOKEN_DURATION - timestamp_list.append([fire_place[i]*TIME_RATE, _split*TIME_RATE]) - timestamp_list.append([_split*TIME_RATE, fire_place[i+1]*TIME_RATE]) - new_char_list.append('') + timestamp_list.append([fire_place[i] * TIME_RATE, _split * TIME_RATE]) + timestamp_list.append([_split * TIME_RATE, fire_place[i + 1] * TIME_RATE]) + new_char_list.append("") # tail token and end silence if num_frames - fire_place[-1] > START_END_THRESHOLD: _end = (num_frames + fire_place[-1]) / 2 - timestamp_list[-1][1] = _end*TIME_RATE - timestamp_list.append([_end*TIME_RATE, num_frames*TIME_RATE]) + timestamp_list[-1][1] = _end * TIME_RATE + timestamp_list.append([_end * TIME_RATE, num_frames * TIME_RATE]) new_char_list.append("") else: - timestamp_list[-1][1] = num_frames*TIME_RATE + timestamp_list[-1][1] = num_frames * TIME_RATE if begin_time: # add offset time in model with vad for i in range(len(timestamp_list)): timestamp_list[i][0] = timestamp_list[i][0] + begin_time / 1000.0 @@ -57,7 +61,6 @@ def time_stamp_lfr6_onnx(us_cif_peak, char_list, begin_time=0.0, total_offset=-1 res_str += "{} {} {};".format(char, timestamp[0], timestamp[1]) res = [] for char, timestamp in zip(new_char_list, timestamp_list): - if char != '': + if char != "": res.append([int(timestamp[0] * 1000), int(timestamp[1] * 1000)]) return res_str, res - \ No newline at end of file diff --git a/runtime/python/onnxruntime/funasr_onnx/utils/utils.py b/runtime/python/onnxruntime/funasr_onnx/utils/utils.py index 260a85eb7..c39d433bf 100644 --- a/runtime/python/onnxruntime/funasr_onnx/utils/utils.py +++ b/runtime/python/onnxruntime/funasr_onnx/utils/utils.py @@ -8,9 +8,15 @@ from typing import Any, Dict, Iterable, List, NamedTuple, Set, Tuple, Union import re import numpy as np import yaml + try: - from onnxruntime import (GraphOptimizationLevel, InferenceSession, - SessionOptions, get_available_providers, get_device) + from onnxruntime import ( + GraphOptimizationLevel, + InferenceSession, + SessionOptions, + get_available_providers, + get_device, + ) except: print("please pip3 install onnxruntime") import jieba @@ -33,7 +39,8 @@ def pad_list(xs, pad_value, max_len=None): return pad -''' + +""" def make_pad_mask(lengths, xs=None, length_dim=-1, maxlen=None): if length_dim == 0: raise ValueError("length_dim cannot be 0: {}".format(length_dim)) @@ -66,26 +73,26 @@ def make_pad_mask(lengths, xs=None, length_dim=-1, maxlen=None): ) mask = mask[ind].expand_as(xs).to(xs.device) return mask -''' +""" -class TokenIDConverter(): - def __init__(self, token_list: Union[List, str], - ): + +class TokenIDConverter: + def __init__( + self, + token_list: Union[List, str], + ): self.token_list = token_list self.unk_symbol = token_list[-1] self.token2id = {v: i for i, v in enumerate(self.token_list)} self.unk_id = self.token2id[self.unk_symbol] - def get_num_vocabulary_size(self) -> int: return len(self.token_list) - def ids2tokens(self, - integers: Union[np.ndarray, Iterable[int]]) -> List[str]: + def ids2tokens(self, integers: Union[np.ndarray, Iterable[int]]) -> List[str]: if isinstance(integers, np.ndarray) and integers.ndim != 1: - raise TokenIDConverterError( - f"Must be 1 dim ndarray, but got {integers.ndim}") + raise TokenIDConverterError(f"Must be 1 dim ndarray, but got {integers.ndim}") return [self.token_list[i] for i in integers] def tokens2ids(self, tokens: Iterable[str]) -> List[int]: @@ -93,7 +100,7 @@ class TokenIDConverter(): return [self.token2id.get(i, self.unk_id) for i in tokens] -class CharTokenizer(): +class CharTokenizer: def __init__( self, symbol_value: Union[Path, str, Iterable[str]] = None, @@ -128,7 +135,7 @@ class CharTokenizer(): if line.startswith(w): if not self.remove_non_linguistic_symbols: tokens.append(line[: len(w)]) - line = line[len(w):] + line = line[len(w) :] break else: t = line[0] @@ -151,7 +158,6 @@ class CharTokenizer(): ) - class Hypothesis(NamedTuple): """Hypothesis data type.""" @@ -177,7 +183,7 @@ class ONNXRuntimeError(Exception): pass -class OrtInferSession(): +class OrtInferSession: def __init__(self, model_file, device_id=-1, intra_op_num_threads=4): device_id = str(device_id) sess_opt = SessionOptions() @@ -186,54 +192,56 @@ class OrtInferSession(): sess_opt.enable_cpu_mem_arena = False sess_opt.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL - cuda_ep = 'CUDAExecutionProvider' + cuda_ep = "CUDAExecutionProvider" cuda_provider_options = { "device_id": device_id, "arena_extend_strategy": "kNextPowerOfTwo", "cudnn_conv_algo_search": "EXHAUSTIVE", "do_copy_in_default_stream": "true", } - cpu_ep = 'CPUExecutionProvider' + cpu_ep = "CPUExecutionProvider" cpu_provider_options = { "arena_extend_strategy": "kSameAsRequested", } EP_list = [] - if device_id != "-1" and get_device() == 'GPU' \ - and cuda_ep in get_available_providers(): + if device_id != "-1" and get_device() == "GPU" and cuda_ep in get_available_providers(): EP_list = [(cuda_ep, cuda_provider_options)] EP_list.append((cpu_ep, cpu_provider_options)) self._verify_model(model_file) - self.session = InferenceSession(model_file, - sess_options=sess_opt, - providers=EP_list) + self.session = InferenceSession(model_file, sess_options=sess_opt, providers=EP_list) if device_id != "-1" and cuda_ep not in self.session.get_providers(): - warnings.warn(f'{cuda_ep} is not avaiable for current env, the inference part is automatically shifted to be executed under {cpu_ep}.\n' - 'Please ensure the installed onnxruntime-gpu version matches your cuda and cudnn version, ' - 'you can check their relations from the offical web site: ' - 'https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html', - RuntimeWarning) + warnings.warn( + f"{cuda_ep} is not avaiable for current env, the inference part is automatically shifted to be executed under {cpu_ep}.\n" + "Please ensure the installed onnxruntime-gpu version matches your cuda and cudnn version, " + "you can check their relations from the offical web site: " + "https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html", + RuntimeWarning, + ) - def __call__(self, - input_content: List[Union[np.ndarray, np.ndarray]]) -> np.ndarray: + def __call__(self, input_content: List[Union[np.ndarray, np.ndarray]]) -> np.ndarray: input_dict = dict(zip(self.get_input_names(), input_content)) try: return self.session.run(self.get_output_names(), input_dict) except Exception as e: - raise ONNXRuntimeError('ONNXRuntime inferece failed.') from e + raise ONNXRuntimeError("ONNXRuntime inferece failed.") from e - def get_input_names(self, ): + def get_input_names( + self, + ): return [v.name for v in self.session.get_inputs()] - def get_output_names(self,): + def get_output_names( + self, + ): return [v.name for v in self.session.get_outputs()] - def get_character_list(self, key: str = 'character'): + def get_character_list(self, key: str = "character"): return self.meta_dict[key].splitlines() - def have_key(self, key: str = 'character') -> bool: + def have_key(self, key: str = "character") -> bool: self.meta_dict = self.session.get_modelmeta().custom_metadata_map if key in self.meta_dict.keys(): return True @@ -243,9 +251,10 @@ class OrtInferSession(): def _verify_model(model_path): model_path = Path(model_path) if not model_path.exists(): - raise FileNotFoundError(f'{model_path} does not exists.') + raise FileNotFoundError(f"{model_path} does not exists.") if not model_path.is_file(): - raise FileExistsError(f'{model_path} is not a file.') + raise FileExistsError(f"{model_path} is not a file.") + def split_to_mini_sentence(words: list, word_limit: int = 20): assert word_limit > 1 @@ -255,11 +264,12 @@ def split_to_mini_sentence(words: list, word_limit: int = 20): length = len(words) sentence_len = length // word_limit for i in range(sentence_len): - sentences.append(words[i * word_limit:(i + 1) * word_limit]) + sentences.append(words[i * word_limit : (i + 1) * word_limit]) if length % word_limit > 0: - sentences.append(words[sentence_len * word_limit:]) + sentences.append(words[sentence_len * word_limit :]) return sentences + def code_mix_split_words(text: str): words = [] segs = text.split() @@ -280,23 +290,26 @@ def code_mix_split_words(text: str): words.append(current_word) return words -def isEnglish(text:str): - if re.search('^[a-zA-Z\']+$', text): + +def isEnglish(text: str): + if re.search("^[a-zA-Z']+$", text): return True else: return False + def join_chinese_and_english(input_list): - line = '' + line = "" for token in input_list: if isEnglish(token): - line = line + ' ' + token + line = line + " " + token else: line = line + token line = line.strip() return line + def code_mix_split_words_jieba(seg_dict_file: str): jieba.load_userdict(seg_dict_file) @@ -307,48 +320,50 @@ def code_mix_split_words_jieba(seg_dict_file: str): token_list_tmp = [] language_flag = None for token in input_list: - if isEnglish(token) and language_flag == 'Chinese': + if isEnglish(token) and language_flag == "Chinese": token_list_all.append(token_list_tmp) - langauge_list.append('Chinese') + langauge_list.append("Chinese") token_list_tmp = [] - elif not isEnglish(token) and language_flag == 'English': + elif not isEnglish(token) and language_flag == "English": token_list_all.append(token_list_tmp) - langauge_list.append('English') + langauge_list.append("English") token_list_tmp = [] - + token_list_tmp.append(token) - + if isEnglish(token): - language_flag = 'English' + language_flag = "English" else: - language_flag = 'Chinese' - + language_flag = "Chinese" + if token_list_tmp: token_list_all.append(token_list_tmp) langauge_list.append(language_flag) - + result_list = [] for token_list_tmp, language_flag in zip(token_list_all, langauge_list): - if language_flag == 'English': + if language_flag == "English": result_list.extend(token_list_tmp) else: seg_list = jieba.cut(join_chinese_and_english(token_list_tmp), HMM=False) result_list.extend(seg_list) - + return result_list + return _fn + def read_yaml(yaml_path: Union[str, Path]) -> Dict: if not Path(yaml_path).exists(): - raise FileExistsError(f'The {yaml_path} does not exist.') + raise FileExistsError(f"The {yaml_path} does not exist.") - with open(str(yaml_path), 'rb') as f: + with open(str(yaml_path), "rb") as f: data = yaml.load(f, Loader=yaml.Loader) return data @functools.lru_cache() -def get_logger(name='funasr_onnx'): +def get_logger(name="funasr_onnx"): """Initialize and get a logger by name. If the logger has not been initialized, this method will initialize the logger by adding one or two handlers, otherwise the initialized logger will @@ -368,8 +383,8 @@ def get_logger(name='funasr_onnx'): return logger formatter = logging.Formatter( - '[%(asctime)s] %(name)s %(levelname)s: %(message)s', - datefmt="%Y/%m/%d %H:%M:%S") + "[%(asctime)s] %(name)s %(levelname)s: %(message)s", datefmt="%Y/%m/%d %H:%M:%S" + ) sh = logging.StreamHandler() sh.setFormatter(formatter) diff --git a/runtime/python/onnxruntime/funasr_onnx/vad_bin.py b/runtime/python/onnxruntime/funasr_onnx/vad_bin.py index a2f443ae9..c195bb3b3 100644 --- a/runtime/python/onnxruntime/funasr_onnx/vad_bin.py +++ b/runtime/python/onnxruntime/funasr_onnx/vad_bin.py @@ -10,323 +10,321 @@ import copy import librosa import numpy as np -from .utils.utils import (ONNXRuntimeError, - OrtInferSession, get_logger, - read_yaml) +from .utils.utils import ONNXRuntimeError, OrtInferSession, get_logger, read_yaml from .utils.frontend import WavFrontend, WavFrontendOnline from .utils.e2e_vad import E2EVadModel logging = get_logger() -class Fsmn_vad(): - """ - Author: Speech Lab of DAMO Academy, Alibaba Group - Deep-FSMN for Large Vocabulary Continuous Speech Recognition - https://arxiv.org/abs/1803.05030 - """ - def __init__(self, model_dir: Union[str, Path] = None, - batch_size: int = 1, - device_id: Union[str, int] = "-1", - quantize: bool = False, - intra_op_num_threads: int = 4, - max_end_sil: int = None, - cache_dir: str = None, - **kwargs - ): - - if not Path(model_dir).exists(): - try: - from modelscope.hub.snapshot_download import snapshot_download - except: - raise "You are exporting model from modelscope, please install modelscope and try it again. To install modelscope, you could:\n" \ - "\npip3 install -U modelscope\n" \ - "For the users in China, you could install with the command:\n" \ - "\npip3 install -U modelscope -i https://mirror.sjtu.edu.cn/pypi/web/simple" - try: - model_dir = snapshot_download(model_dir, cache_dir=cache_dir) - except: - raise "model_dir must be model_name in modelscope or local path downloaded from modelscope, but is {}".format( - model_dir) - - model_file = os.path.join(model_dir, 'model.onnx') - if quantize: - model_file = os.path.join(model_dir, 'model_quant.onnx') - if not os.path.exists(model_file): - print(".onnx is not exist, begin to export onnx") - try: - from funasr import AutoModel - except: - raise "You are exporting onnx, please install funasr and try it again. To install funasr, you could:\n" \ - "\npip3 install -U funasr\n" \ - "For the users in China, you could install with the command:\n" \ - "\npip3 install -U funasr -i https://mirror.sjtu.edu.cn/pypi/web/simple" - - model = AutoModel(model=model_dir) - model_dir = model.export(type="onnx", quantize=quantize, **kwargs) - config_file = os.path.join(model_dir, 'config.yaml') - cmvn_file = os.path.join(model_dir, 'am.mvn') - config = read_yaml(config_file) - - self.frontend = WavFrontend( - cmvn_file=cmvn_file, - **config['frontend_conf'] - ) - self.ort_infer = OrtInferSession(model_file, device_id, intra_op_num_threads=intra_op_num_threads) - self.batch_size = batch_size - self.vad_scorer = E2EVadModel(config["model_conf"]) - self.max_end_sil = max_end_sil if max_end_sil is not None else config["model_conf"]["max_end_silence_time"] - self.encoder_conf = config["encoder_conf"] - - def prepare_cache(self, in_cache: list = []): - if len(in_cache) > 0: - return in_cache - fsmn_layers = self.encoder_conf["fsmn_layers"] - proj_dim = self.encoder_conf["proj_dim"] - lorder = self.encoder_conf["lorder"] - for i in range(fsmn_layers): - cache = np.zeros((1, proj_dim, lorder-1, 1)).astype(np.float32) - in_cache.append(cache) - return in_cache - - - def __call__(self, audio_in: Union[str, np.ndarray, List[str]], **kwargs) -> List: - waveform_list = self.load_data(audio_in, self.frontend.opts.frame_opts.samp_freq) - waveform_nums = len(waveform_list) - is_final = kwargs.get('kwargs', False) +class Fsmn_vad: + """ + Author: Speech Lab of DAMO Academy, Alibaba Group + Deep-FSMN for Large Vocabulary Continuous Speech Recognition + https://arxiv.org/abs/1803.05030 + """ - segments = [[]] * self.batch_size - for beg_idx in range(0, waveform_nums, self.batch_size): - - end_idx = min(waveform_nums, beg_idx + self.batch_size) - waveform = waveform_list[beg_idx:end_idx] - feats, feats_len = self.extract_feat(waveform) - waveform = np.array(waveform) - param_dict = kwargs.get('param_dict', dict()) - in_cache = param_dict.get('in_cache', list()) - in_cache = self.prepare_cache(in_cache) - try: - t_offset = 0 - step = int(min(feats_len.max(), 6000)) - for t_offset in range(0, int(feats_len), min(step, feats_len - t_offset)): - if t_offset + step >= feats_len - 1: - step = feats_len - t_offset - is_final = True - else: - is_final = False - feats_package = feats[:, t_offset:int(t_offset + step), :] - waveform_package = waveform[:, t_offset * 160:min(waveform.shape[-1], (int(t_offset + step) - 1) * 160 + 400)] + def __init__( + self, + model_dir: Union[str, Path] = None, + batch_size: int = 1, + device_id: Union[str, int] = "-1", + quantize: bool = False, + intra_op_num_threads: int = 4, + max_end_sil: int = None, + cache_dir: str = None, + **kwargs, + ): - inputs = [feats_package] - # inputs = [feats] - inputs.extend(in_cache) - scores, out_caches = self.infer(inputs) - in_cache = out_caches - segments_part = self.vad_scorer(scores, waveform_package, is_final=is_final, max_end_sil=self.max_end_sil, online=False) - # segments = self.vad_scorer(scores, waveform[0][None, :], is_final=is_final, max_end_sil=self.max_end_sil) + if not Path(model_dir).exists(): + try: + from modelscope.hub.snapshot_download import snapshot_download + except: + raise "You are exporting model from modelscope, please install modelscope and try it again. To install modelscope, you could:\n" "\npip3 install -U modelscope\n" "For the users in China, you could install with the command:\n" "\npip3 install -U modelscope -i https://mirror.sjtu.edu.cn/pypi/web/simple" + try: + model_dir = snapshot_download(model_dir, cache_dir=cache_dir) + except: + raise "model_dir must be model_name in modelscope or local path downloaded from modelscope, but is {}".format( + model_dir + ) - if segments_part: - for batch_num in range(0, self.batch_size): - segments[batch_num] += segments_part[batch_num] - - except ONNXRuntimeError: - # logging.warning(traceback.format_exc()) - logging.warning("input wav is silence or noise") - segments = '' - - return segments + model_file = os.path.join(model_dir, "model.onnx") + if quantize: + model_file = os.path.join(model_dir, "model_quant.onnx") + if not os.path.exists(model_file): + print(".onnx is not exist, begin to export onnx") + try: + from funasr import AutoModel + except: + raise "You are exporting onnx, please install funasr and try it again. To install funasr, you could:\n" "\npip3 install -U funasr\n" "For the users in China, you could install with the command:\n" "\npip3 install -U funasr -i https://mirror.sjtu.edu.cn/pypi/web/simple" - def load_data(self, - wav_content: Union[str, np.ndarray, List[str]], fs: int = None) -> List: - def load_wav(path: str) -> np.ndarray: - waveform, _ = librosa.load(path, sr=fs) - return waveform - - if isinstance(wav_content, np.ndarray): - return [wav_content] - - if isinstance(wav_content, str): - return [load_wav(wav_content)] - - if isinstance(wav_content, list): - return [load_wav(path) for path in wav_content] - - raise TypeError( - f'The type of {wav_content} is not in [str, np.ndarray, list]') - - def extract_feat(self, - waveform_list: List[np.ndarray] - ) -> Tuple[np.ndarray, np.ndarray]: - feats, feats_len = [], [] - for waveform in waveform_list: - speech, _ = self.frontend.fbank(waveform) - feat, feat_len = self.frontend.lfr_cmvn(speech) - feats.append(feat) - feats_len.append(feat_len) - - feats = self.pad_feats(feats, np.max(feats_len)) - feats_len = np.array(feats_len).astype(np.int32) - return feats, feats_len - - @staticmethod - def pad_feats(feats: List[np.ndarray], max_feat_len: int) -> np.ndarray: - def pad_feat(feat: np.ndarray, cur_len: int) -> np.ndarray: - pad_width = ((0, max_feat_len - cur_len), (0, 0)) - return np.pad(feat, pad_width, 'constant', constant_values=0) - - feat_res = [pad_feat(feat, feat.shape[0]) for feat in feats] - feats = np.array(feat_res).astype(np.float32) - return feats - - def infer(self, feats: List) -> Tuple[np.ndarray, np.ndarray]: - - outputs = self.ort_infer(feats) - scores, out_caches = outputs[0], outputs[1:] - return scores, out_caches + model = AutoModel(model=model_dir) + model_dir = model.export(type="onnx", quantize=quantize, **kwargs) + config_file = os.path.join(model_dir, "config.yaml") + cmvn_file = os.path.join(model_dir, "am.mvn") + config = read_yaml(config_file) + + self.frontend = WavFrontend(cmvn_file=cmvn_file, **config["frontend_conf"]) + self.ort_infer = OrtInferSession( + model_file, device_id, intra_op_num_threads=intra_op_num_threads + ) + self.batch_size = batch_size + self.vad_scorer = E2EVadModel(config["model_conf"]) + self.max_end_sil = ( + max_end_sil if max_end_sil is not None else config["model_conf"]["max_end_silence_time"] + ) + self.encoder_conf = config["encoder_conf"] + + def prepare_cache(self, in_cache: list = []): + if len(in_cache) > 0: + return in_cache + fsmn_layers = self.encoder_conf["fsmn_layers"] + proj_dim = self.encoder_conf["proj_dim"] + lorder = self.encoder_conf["lorder"] + for i in range(fsmn_layers): + cache = np.zeros((1, proj_dim, lorder - 1, 1)).astype(np.float32) + in_cache.append(cache) + return in_cache + + def __call__(self, audio_in: Union[str, np.ndarray, List[str]], **kwargs) -> List: + waveform_list = self.load_data(audio_in, self.frontend.opts.frame_opts.samp_freq) + waveform_nums = len(waveform_list) + is_final = kwargs.get("kwargs", False) + + segments = [[]] * self.batch_size + for beg_idx in range(0, waveform_nums, self.batch_size): + + end_idx = min(waveform_nums, beg_idx + self.batch_size) + waveform = waveform_list[beg_idx:end_idx] + feats, feats_len = self.extract_feat(waveform) + waveform = np.array(waveform) + param_dict = kwargs.get("param_dict", dict()) + in_cache = param_dict.get("in_cache", list()) + in_cache = self.prepare_cache(in_cache) + try: + t_offset = 0 + step = int(min(feats_len.max(), 6000)) + for t_offset in range(0, int(feats_len), min(step, feats_len - t_offset)): + if t_offset + step >= feats_len - 1: + step = feats_len - t_offset + is_final = True + else: + is_final = False + feats_package = feats[:, t_offset : int(t_offset + step), :] + waveform_package = waveform[ + :, + t_offset + * 160 : min(waveform.shape[-1], (int(t_offset + step) - 1) * 160 + 400), + ] + + inputs = [feats_package] + # inputs = [feats] + inputs.extend(in_cache) + scores, out_caches = self.infer(inputs) + in_cache = out_caches + segments_part = self.vad_scorer( + scores, + waveform_package, + is_final=is_final, + max_end_sil=self.max_end_sil, + online=False, + ) + # segments = self.vad_scorer(scores, waveform[0][None, :], is_final=is_final, max_end_sil=self.max_end_sil) + + if segments_part: + for batch_num in range(0, self.batch_size): + segments[batch_num] += segments_part[batch_num] + + except ONNXRuntimeError: + # logging.warning(traceback.format_exc()) + logging.warning("input wav is silence or noise") + segments = "" + + return segments + + def load_data(self, wav_content: Union[str, np.ndarray, List[str]], fs: int = None) -> List: + def load_wav(path: str) -> np.ndarray: + waveform, _ = librosa.load(path, sr=fs) + return waveform + + if isinstance(wav_content, np.ndarray): + return [wav_content] + + if isinstance(wav_content, str): + return [load_wav(wav_content)] + + if isinstance(wav_content, list): + return [load_wav(path) for path in wav_content] + + raise TypeError(f"The type of {wav_content} is not in [str, np.ndarray, list]") + + def extract_feat(self, waveform_list: List[np.ndarray]) -> Tuple[np.ndarray, np.ndarray]: + feats, feats_len = [], [] + for waveform in waveform_list: + speech, _ = self.frontend.fbank(waveform) + feat, feat_len = self.frontend.lfr_cmvn(speech) + feats.append(feat) + feats_len.append(feat_len) + + feats = self.pad_feats(feats, np.max(feats_len)) + feats_len = np.array(feats_len).astype(np.int32) + return feats, feats_len + + @staticmethod + def pad_feats(feats: List[np.ndarray], max_feat_len: int) -> np.ndarray: + def pad_feat(feat: np.ndarray, cur_len: int) -> np.ndarray: + pad_width = ((0, max_feat_len - cur_len), (0, 0)) + return np.pad(feat, pad_width, "constant", constant_values=0) + + feat_res = [pad_feat(feat, feat.shape[0]) for feat in feats] + feats = np.array(feat_res).astype(np.float32) + return feats + + def infer(self, feats: List) -> Tuple[np.ndarray, np.ndarray]: + + outputs = self.ort_infer(feats) + scores, out_caches = outputs[0], outputs[1:] + return scores, out_caches -class Fsmn_vad_online(): - """ - Author: Speech Lab of DAMO Academy, Alibaba Group - Deep-FSMN for Large Vocabulary Continuous Speech Recognition - https://arxiv.org/abs/1803.05030 - """ - def __init__(self, model_dir: Union[str, Path] = None, - batch_size: int = 1, - device_id: Union[str, int] = "-1", - quantize: bool = False, - intra_op_num_threads: int = 4, - max_end_sil: int = None, - cache_dir: str = None, - **kwargs - ): - if not Path(model_dir).exists(): - try: - from modelscope.hub.snapshot_download import snapshot_download - except: - raise "You are exporting model from modelscope, please install modelscope and try it again. To install modelscope, you could:\n" \ - "\npip3 install -U modelscope\n" \ - "For the users in China, you could install with the command:\n" \ - "\npip3 install -U modelscope -i https://mirror.sjtu.edu.cn/pypi/web/simple" - try: - model_dir = snapshot_download(model_dir, cache_dir=cache_dir) - except: - raise "model_dir must be model_name in modelscope or local path downloaded from modelscope, but is {}".format( - model_dir) - - model_file = os.path.join(model_dir, 'model.onnx') - if quantize: - model_file = os.path.join(model_dir, 'model_quant.onnx') - if not os.path.exists(model_file): - print(".onnx is not exist, begin to export onnx") - try: - from funasr import AutoModel - except: - raise "You are exporting onnx, please install funasr and try it again. To install funasr, you could:\n" \ - "\npip3 install -U funasr\n" \ - "For the users in China, you could install with the command:\n" \ - "\npip3 install -U funasr -i https://mirror.sjtu.edu.cn/pypi/web/simple" - - model = AutoModel(model=model_dir) - model_dir = model.export(type="onnx", quantize=quantize, **kwargs) - - config_file = os.path.join(model_dir, 'config.yaml') - cmvn_file = os.path.join(model_dir, 'am.mvn') - config = read_yaml(config_file) - - self.frontend = WavFrontendOnline( - cmvn_file=cmvn_file, - **config['frontend_conf'] - ) - self.ort_infer = OrtInferSession(model_file, device_id, intra_op_num_threads=intra_op_num_threads) - self.batch_size = batch_size - self.vad_scorer = E2EVadModel(config["model_conf"]) - self.max_end_sil = max_end_sil if max_end_sil is not None else config["model_conf"]["max_end_silence_time"] - self.encoder_conf = config["encoder_conf"] - - def prepare_cache(self, in_cache: list = []): - if len(in_cache) > 0: - return in_cache - fsmn_layers = self.encoder_conf["fsmn_layers"] - proj_dim = self.encoder_conf["proj_dim"] - lorder = self.encoder_conf["lorder"] - for i in range(fsmn_layers): - cache = np.zeros((1, proj_dim, lorder - 1, 1)).astype(np.float32) - in_cache.append(cache) - return in_cache - - def __call__(self, audio_in: np.ndarray, **kwargs) -> List: - waveforms = np.expand_dims(audio_in, axis=0) - - param_dict = kwargs.get('param_dict', dict()) - is_final = param_dict.get('is_final', False) - feats, feats_len = self.extract_feat(waveforms, is_final) - segments = [] - if feats.size != 0: - in_cache = param_dict.get('in_cache', list()) - in_cache = self.prepare_cache(in_cache) - try: - inputs = [feats] - inputs.extend(in_cache) - scores, out_caches = self.infer(inputs) - param_dict['in_cache'] = out_caches - waveforms = self.frontend.get_waveforms() - segments = self.vad_scorer(scores, waveforms, is_final=is_final, max_end_sil=self.max_end_sil, - online=True) - - - except ONNXRuntimeError: - # logging.warning(traceback.format_exc()) - logging.warning("input wav is silence or noise") - segments = [] - return segments - - def load_data(self, - wav_content: Union[str, np.ndarray, List[str]], fs: int = None) -> List: - def load_wav(path: str) -> np.ndarray: - waveform, _ = librosa.load(path, sr=fs) - return waveform - - if isinstance(wav_content, np.ndarray): - return [wav_content] - - if isinstance(wav_content, str): - return [load_wav(wav_content)] - - if isinstance(wav_content, list): - return [load_wav(path) for path in wav_content] - - raise TypeError( - f'The type of {wav_content} is not in [str, np.ndarray, list]') - - def extract_feat(self, - waveforms: np.ndarray, is_final: bool = False - ) -> Tuple[np.ndarray, np.ndarray]: - waveforms_lens = np.zeros(waveforms.shape[0]).astype(np.int32) - for idx, waveform in enumerate(waveforms): - waveforms_lens[idx] = waveform.shape[-1] - - feats, feats_len = self.frontend.extract_fbank(waveforms, waveforms_lens, is_final) - # feats.append(feat) - # feats_len.append(feat_len) - - # feats = self.pad_feats(feats, np.max(feats_len)) - # feats_len = np.array(feats_len).astype(np.int32) - return feats.astype(np.float32), feats_len.astype(np.int32) - - @staticmethod - def pad_feats(feats: List[np.ndarray], max_feat_len: int) -> np.ndarray: - def pad_feat(feat: np.ndarray, cur_len: int) -> np.ndarray: - pad_width = ((0, max_feat_len - cur_len), (0, 0)) - return np.pad(feat, pad_width, 'constant', constant_values=0) - - feat_res = [pad_feat(feat, feat.shape[0]) for feat in feats] - feats = np.array(feat_res).astype(np.float32) - return feats - - def infer(self, feats: List) -> Tuple[np.ndarray, np.ndarray]: - - outputs = self.ort_infer(feats) - scores, out_caches = outputs[0], outputs[1:] - return scores, out_caches +class Fsmn_vad_online: + """ + Author: Speech Lab of DAMO Academy, Alibaba Group + Deep-FSMN for Large Vocabulary Continuous Speech Recognition + https://arxiv.org/abs/1803.05030 + """ + def __init__( + self, + model_dir: Union[str, Path] = None, + batch_size: int = 1, + device_id: Union[str, int] = "-1", + quantize: bool = False, + intra_op_num_threads: int = 4, + max_end_sil: int = None, + cache_dir: str = None, + **kwargs, + ): + if not Path(model_dir).exists(): + try: + from modelscope.hub.snapshot_download import snapshot_download + except: + raise "You are exporting model from modelscope, please install modelscope and try it again. To install modelscope, you could:\n" "\npip3 install -U modelscope\n" "For the users in China, you could install with the command:\n" "\npip3 install -U modelscope -i https://mirror.sjtu.edu.cn/pypi/web/simple" + try: + model_dir = snapshot_download(model_dir, cache_dir=cache_dir) + except: + raise "model_dir must be model_name in modelscope or local path downloaded from modelscope, but is {}".format( + model_dir + ) + + model_file = os.path.join(model_dir, "model.onnx") + if quantize: + model_file = os.path.join(model_dir, "model_quant.onnx") + if not os.path.exists(model_file): + print(".onnx is not exist, begin to export onnx") + try: + from funasr import AutoModel + except: + raise "You are exporting onnx, please install funasr and try it again. To install funasr, you could:\n" "\npip3 install -U funasr\n" "For the users in China, you could install with the command:\n" "\npip3 install -U funasr -i https://mirror.sjtu.edu.cn/pypi/web/simple" + + model = AutoModel(model=model_dir) + model_dir = model.export(type="onnx", quantize=quantize, **kwargs) + + config_file = os.path.join(model_dir, "config.yaml") + cmvn_file = os.path.join(model_dir, "am.mvn") + config = read_yaml(config_file) + + self.frontend = WavFrontendOnline(cmvn_file=cmvn_file, **config["frontend_conf"]) + self.ort_infer = OrtInferSession( + model_file, device_id, intra_op_num_threads=intra_op_num_threads + ) + self.batch_size = batch_size + self.vad_scorer = E2EVadModel(config["model_conf"]) + self.max_end_sil = ( + max_end_sil if max_end_sil is not None else config["model_conf"]["max_end_silence_time"] + ) + self.encoder_conf = config["encoder_conf"] + + def prepare_cache(self, in_cache: list = []): + if len(in_cache) > 0: + return in_cache + fsmn_layers = self.encoder_conf["fsmn_layers"] + proj_dim = self.encoder_conf["proj_dim"] + lorder = self.encoder_conf["lorder"] + for i in range(fsmn_layers): + cache = np.zeros((1, proj_dim, lorder - 1, 1)).astype(np.float32) + in_cache.append(cache) + return in_cache + + def __call__(self, audio_in: np.ndarray, **kwargs) -> List: + waveforms = np.expand_dims(audio_in, axis=0) + + param_dict = kwargs.get("param_dict", dict()) + is_final = param_dict.get("is_final", False) + feats, feats_len = self.extract_feat(waveforms, is_final) + segments = [] + if feats.size != 0: + in_cache = param_dict.get("in_cache", list()) + in_cache = self.prepare_cache(in_cache) + try: + inputs = [feats] + inputs.extend(in_cache) + scores, out_caches = self.infer(inputs) + param_dict["in_cache"] = out_caches + waveforms = self.frontend.get_waveforms() + segments = self.vad_scorer( + scores, waveforms, is_final=is_final, max_end_sil=self.max_end_sil, online=True + ) + + except ONNXRuntimeError: + # logging.warning(traceback.format_exc()) + logging.warning("input wav is silence or noise") + segments = [] + return segments + + def load_data(self, wav_content: Union[str, np.ndarray, List[str]], fs: int = None) -> List: + def load_wav(path: str) -> np.ndarray: + waveform, _ = librosa.load(path, sr=fs) + return waveform + + if isinstance(wav_content, np.ndarray): + return [wav_content] + + if isinstance(wav_content, str): + return [load_wav(wav_content)] + + if isinstance(wav_content, list): + return [load_wav(path) for path in wav_content] + + raise TypeError(f"The type of {wav_content} is not in [str, np.ndarray, list]") + + def extract_feat( + self, waveforms: np.ndarray, is_final: bool = False + ) -> Tuple[np.ndarray, np.ndarray]: + waveforms_lens = np.zeros(waveforms.shape[0]).astype(np.int32) + for idx, waveform in enumerate(waveforms): + waveforms_lens[idx] = waveform.shape[-1] + + feats, feats_len = self.frontend.extract_fbank(waveforms, waveforms_lens, is_final) + # feats.append(feat) + # feats_len.append(feat_len) + + # feats = self.pad_feats(feats, np.max(feats_len)) + # feats_len = np.array(feats_len).astype(np.int32) + return feats.astype(np.float32), feats_len.astype(np.int32) + + @staticmethod + def pad_feats(feats: List[np.ndarray], max_feat_len: int) -> np.ndarray: + def pad_feat(feat: np.ndarray, cur_len: int) -> np.ndarray: + pad_width = ((0, max_feat_len - cur_len), (0, 0)) + return np.pad(feat, pad_width, "constant", constant_values=0) + + feat_res = [pad_feat(feat, feat.shape[0]) for feat in feats] + feats = np.array(feat_res).astype(np.float32) + return feats + + def infer(self, feats: List) -> Tuple[np.ndarray, np.ndarray]: + + outputs = self.ort_infer(feats) + scores, out_caches = outputs[0], outputs[1:] + return scores, out_caches diff --git a/runtime/python/onnxruntime/funasr_server_http.py b/runtime/python/onnxruntime/funasr_server_http.py index e92f3022d..a6107229e 100644 --- a/runtime/python/onnxruntime/funasr_server_http.py +++ b/runtime/python/onnxruntime/funasr_server_http.py @@ -4,35 +4,40 @@ import io import soundfile as sf import uvicorn from fastapi import FastAPI, Body + app = FastAPI() from funasr_onnx import Paraformer + model_dir = "damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-onnx" model = Paraformer(model_dir, batch_size=1, quantize=True) + async def recognition_onnx(waveform): - result=model(waveform)[0]["preds"][0] + result = model(waveform)[0]["preds"][0] return result + + @app.post("/api/asr") async def asr(item: dict = Body(...)): try: - audio_bytes= base64.b64decode(bytes(item['wav_base64'], 'utf-8')) + audio_bytes = base64.b64decode(bytes(item["wav_base64"], "utf-8")) waveform, _ = sf.read(io.BytesIO(audio_bytes)) - result=await recognition_onnx(waveform) + result = await recognition_onnx(waveform) ret = {"results": result, "code": 0} except: - print('请求出错，这里是处理出错的') - ret = {"results": '', "code": 1} + print("请求出错，这里是处理出错的") + ret = {"results": "", "code": 1} return ret -if __name__ == '__main__': - parser = argparse.ArgumentParser(description='API Service') - parser.add_argument('--listen', default='0.0.0.0', type=str, help='the network to listen') - parser.add_argument('--port', default=8888, type=int, help='the port to listen') +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="API Service") + parser.add_argument("--listen", default="0.0.0.0", type=str, help="the network to listen") + parser.add_argument("--port", default=8888, type=int, help="the port to listen") args = parser.parse_args() - print('start...') - print('server on:',args) - - uvicorn.run(app, host=args.listen, port=args.port) \ No newline at end of file + print("start...") + print("server on:", args) + + uvicorn.run(app, host=args.listen, port=args.port) diff --git a/runtime/python/onnxruntime/setup.py b/runtime/python/onnxruntime/setup.py index 39d57e7d5..b2977cedb 100644 --- a/runtime/python/onnxruntime/setup.py +++ b/runtime/python/onnxruntime/setup.py @@ -5,15 +5,15 @@ import setuptools def get_readme(): root_dir = Path(__file__).resolve().parent - readme_path = str(root_dir / 'README.md') + readme_path = str(root_dir / "README.md") print(readme_path) - with open(readme_path, 'r', encoding='utf-8') as f: + with open(readme_path, "r", encoding="utf-8") as f: readme = f.read() return readme -MODULE_NAME = 'funasr_onnx' -VERSION_NUM = '0.3.1' +MODULE_NAME = "funasr_onnx" +VERSION_NUM = "0.3.1" setuptools.setup( name=MODULE_NAME, @@ -23,27 +23,26 @@ setuptools.setup( author="Speech Lab of DAMO Academy, Alibaba Group", author_email="funasr@list.alibaba-inc.com", description="FunASR: A Fundamental End-to-End Speech Recognition Toolkit", - license='MIT', + license="MIT", long_description=get_readme(), - long_description_content_type='text/markdown', + long_description_content_type="text/markdown", include_package_data=True, - install_requires=["librosa", - "onnxruntime>=1.7.0", - "scipy", - "numpy>=1.19.3", - "kaldi-native-fbank", - "PyYAML>=5.1.2", - "onnx" - ], - packages=[MODULE_NAME, f'{MODULE_NAME}.utils'], - keywords=[ - 'funasr,asr' + install_requires=[ + "librosa", + "onnxruntime>=1.7.0", + "scipy", + "numpy>=1.19.3", + "kaldi-native-fbank", + "PyYAML>=5.1.2", + "onnx", ], + packages=[MODULE_NAME, f"{MODULE_NAME}.utils"], + keywords=["funasr,asr"], classifiers=[ - 'Programming Language :: Python :: 3.6', - 'Programming Language :: Python :: 3.7', - 'Programming Language :: Python :: 3.8', - 'Programming Language :: Python :: 3.9', - 'Programming Language :: Python :: 3.10', + "Programming Language :: Python :: 3.6", + "Programming Language :: Python :: 3.7", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", ], ) diff --git a/runtime/python/utils/compute_wer.py b/runtime/python/utils/compute_wer.py index 349a3f609..46ec08e5c 100755 --- a/runtime/python/utils/compute_wer.py +++ b/runtime/python/utils/compute_wer.py @@ -2,66 +2,94 @@ import os import numpy as np import sys -def compute_wer(ref_file, - hyp_file, - cer_detail_file): + +def compute_wer(ref_file, hyp_file, cer_detail_file): rst = { - 'Wrd': 0, - 'Corr': 0, - 'Ins': 0, - 'Del': 0, - 'Sub': 0, - 'Snt': 0, - 'Err': 0.0, - 'S.Err': 0.0, - 'wrong_words': 0, - 'wrong_sentences': 0 + "Wrd": 0, + "Corr": 0, + "Ins": 0, + "Del": 0, + "Sub": 0, + "Snt": 0, + "Err": 0.0, + "S.Err": 0.0, + "wrong_words": 0, + "wrong_sentences": 0, } hyp_dict = {} ref_dict = {} - with open(hyp_file, 'r') as hyp_reader: + with open(hyp_file, "r") as hyp_reader: for line in hyp_reader: key = line.strip().split()[0] value = line.strip().split()[1:] hyp_dict[key] = value - with open(ref_file, 'r') as ref_reader: + with open(ref_file, "r") as ref_reader: for line in ref_reader: key = line.strip().split()[0] value = line.strip().split()[1:] ref_dict[key] = value - cer_detail_writer = open(cer_detail_file, 'w') + cer_detail_writer = open(cer_detail_file, "w") for hyp_key in hyp_dict: if hyp_key in ref_dict: - out_item = compute_wer_by_line(hyp_dict[hyp_key], ref_dict[hyp_key]) - rst['Wrd'] += out_item['nwords'] - rst['Corr'] += out_item['cor'] - rst['wrong_words'] += out_item['wrong'] - rst['Ins'] += out_item['ins'] - rst['Del'] += out_item['del'] - rst['Sub'] += out_item['sub'] - rst['Snt'] += 1 - if out_item['wrong'] > 0: - rst['wrong_sentences'] += 1 - cer_detail_writer.write(hyp_key + print_cer_detail(out_item) + '\n') - cer_detail_writer.write("ref:" + '\t' + "".join(ref_dict[hyp_key]) + '\n') - cer_detail_writer.write("hyp:" + '\t' + "".join(hyp_dict[hyp_key]) + '\n') + out_item = compute_wer_by_line(hyp_dict[hyp_key], ref_dict[hyp_key]) + rst["Wrd"] += out_item["nwords"] + rst["Corr"] += out_item["cor"] + rst["wrong_words"] += out_item["wrong"] + rst["Ins"] += out_item["ins"] + rst["Del"] += out_item["del"] + rst["Sub"] += out_item["sub"] + rst["Snt"] += 1 + if out_item["wrong"] > 0: + rst["wrong_sentences"] += 1 + cer_detail_writer.write(hyp_key + print_cer_detail(out_item) + "\n") + cer_detail_writer.write("ref:" + "\t" + "".join(ref_dict[hyp_key]) + "\n") + cer_detail_writer.write("hyp:" + "\t" + "".join(hyp_dict[hyp_key]) + "\n") - if rst['Wrd'] > 0: - rst['Err'] = round(rst['wrong_words'] * 100 / rst['Wrd'], 2) - if rst['Snt'] > 0: - rst['S.Err'] = round(rst['wrong_sentences'] * 100 / rst['Snt'], 2) + if rst["Wrd"] > 0: + rst["Err"] = round(rst["wrong_words"] * 100 / rst["Wrd"], 2) + if rst["Snt"] > 0: + rst["S.Err"] = round(rst["wrong_sentences"] * 100 / rst["Snt"], 2) - cer_detail_writer.write('\n') - cer_detail_writer.write("%WER " + str(rst['Err']) + " [ " + str(rst['wrong_words'])+ " / " + str(rst['Wrd']) + - ", " + str(rst['Ins']) + " ins, " + str(rst['Del']) + " del, " + str(rst['Sub']) + " sub ]" + '\n') - cer_detail_writer.write("%SER " + str(rst['S.Err']) + " [ " + str(rst['wrong_sentences']) + " / " + str(rst['Snt']) + " ]" + '\n') - cer_detail_writer.write("Scored " + str(len(hyp_dict)) + " sentences, " + str(len(hyp_dict) - rst['Snt']) + " not present in hyp." + '\n') + cer_detail_writer.write("\n") + cer_detail_writer.write( + "%WER " + + str(rst["Err"]) + + " [ " + + str(rst["wrong_words"]) + + " / " + + str(rst["Wrd"]) + + ", " + + str(rst["Ins"]) + + " ins, " + + str(rst["Del"]) + + " del, " + + str(rst["Sub"]) + + " sub ]" + + "\n" + ) + cer_detail_writer.write( + "%SER " + + str(rst["S.Err"]) + + " [ " + + str(rst["wrong_sentences"]) + + " / " + + str(rst["Snt"]) + + " ]" + + "\n" + ) + cer_detail_writer.write( + "Scored " + + str(len(hyp_dict)) + + " sentences, " + + str(len(hyp_dict) - rst["Snt"]) + + " not present in hyp." + + "\n" + ) - -def compute_wer_by_line(hyp, - ref): + +def compute_wer_by_line(hyp, ref): hyp = list(map(lambda x: x.lower(), hyp)) ref = list(map(lambda x: x.lower(), ref)) @@ -96,14 +124,7 @@ def compute_wer_by_line(hyp, match_idx = [] i = len_hyp j = len_ref - rst = { - 'nwords': len_ref, - 'cor': 0, - 'wrong': 0, - 'ins': 0, - 'del': 0, - 'sub': 0 - } + rst = {"nwords": len_ref, "cor": 0, "wrong": 0, "ins": 0, "del": 0, "sub": 0} while i >= 0 or j >= 0: i_idx = max(0, i) j_idx = max(0, j) @@ -111,42 +132,57 @@ def compute_wer_by_line(hyp, if ops_matrix[i_idx][j_idx] == 0: # correct if i - 1 >= 0 and j - 1 >= 0: match_idx.append((j - 1, i - 1)) - rst['cor'] += 1 + rst["cor"] += 1 i -= 1 j -= 1 elif ops_matrix[i_idx][j_idx] == 2: # insert i -= 1 - rst['ins'] += 1 + rst["ins"] += 1 elif ops_matrix[i_idx][j_idx] == 3: # delete j -= 1 - rst['del'] += 1 + rst["del"] += 1 elif ops_matrix[i_idx][j_idx] == 1: # substitute i -= 1 j -= 1 - rst['sub'] += 1 + rst["sub"] += 1 if i < 0 and j >= 0: - rst['del'] += 1 + rst["del"] += 1 elif j < 0 and i >= 0: - rst['ins'] += 1 + rst["ins"] += 1 match_idx.reverse() wrong_cnt = cost_matrix[len_hyp][len_ref] - rst['wrong'] = wrong_cnt + rst["wrong"] = wrong_cnt return rst -def print_cer_detail(rst): - return ("(" + "nwords=" + str(rst['nwords']) + ",cor=" + str(rst['cor']) - + ",ins=" + str(rst['ins']) + ",del=" + str(rst['del']) + ",sub=" - + str(rst['sub']) + ") corr:" + '{:.2%}'.format(rst['cor']/rst['nwords']) - + ",cer:" + '{:.2%}'.format(rst['wrong']/rst['nwords'])) -if __name__ == '__main__': +def print_cer_detail(rst): + return ( + "(" + + "nwords=" + + str(rst["nwords"]) + + ",cor=" + + str(rst["cor"]) + + ",ins=" + + str(rst["ins"]) + + ",del=" + + str(rst["del"]) + + ",sub=" + + str(rst["sub"]) + + ") corr:" + + "{:.2%}".format(rst["cor"] / rst["nwords"]) + + ",cer:" + + "{:.2%}".format(rst["wrong"] / rst["nwords"]) + ) + + +if __name__ == "__main__": if len(sys.argv) != 4: print("usage : python compute-wer.py test.ref test.hyp test.wer") sys.exit(0) diff --git a/runtime/python/utils/proce_text.py b/runtime/python/utils/proce_text.py index 9e517a4e1..d03febd96 100755 --- a/runtime/python/utils/proce_text.py +++ b/runtime/python/utils/proce_text.py @@ -1,4 +1,3 @@ - import sys import re @@ -7,25 +6,25 @@ out_f = sys.argv[2] with open(in_f, "r", encoding="utf-8") as f: - lines = f.readlines() + lines = f.readlines() with open(out_f, "w", encoding="utf-8") as f: - for line in lines: - outs = line.strip().split(" ", 1) - if len(outs) == 2: - idx, text = outs - text = re.sub("", "", text) - text = re.sub("", "", text) - text = re.sub("@@", "", text) - text = re.sub("@", "", text) - text = re.sub("", "", text) - text = re.sub(" ", "", text) - text = text.lower() - else: - idx = outs[0] - text = " " + for line in lines: + outs = line.strip().split(" ", 1) + if len(outs) == 2: + idx, text = outs + text = re.sub("", "", text) + text = re.sub("", "", text) + text = re.sub("@@", "", text) + text = re.sub("@", "", text) + text = re.sub("", "", text) + text = re.sub(" ", "", text) + text = text.lower() + else: + idx = outs[0] + text = " " - text = [x for x in text] - text = " ".join(text) - out = "{} {}\n".format(idx, text) - f.write(out) + text = [x for x in text] + text = " ".join(text) + out = "{} {}\n".format(idx, text) + f.write(out) diff --git a/runtime/python/utils/test_cer.py b/runtime/python/utils/test_cer.py index d795d338b..58285c248 100644 --- a/runtime/python/utils/test_cer.py +++ b/runtime/python/utils/test_cer.py @@ -5,44 +5,52 @@ import librosa from funasr.utils.types import str2bool import argparse + parser = argparse.ArgumentParser() -parser.add_argument('--model_dir', type=str, required=True) -parser.add_argument('--backend', type=str, default='onnx', help='["onnx", "torch"]') -parser.add_argument('--wav_file', type=str, default=None, help='amp fallback number') -parser.add_argument('--quantize', type=str2bool, default=False, help='quantized model') -parser.add_argument('--intra_op_num_threads', type=int, default=1, help='intra_op_num_threads for onnx') -parser.add_argument('--output_dir', type=str, default=None, help='amp fallback number') +parser.add_argument("--model_dir", type=str, required=True) +parser.add_argument("--backend", type=str, default="onnx", help='["onnx", "torch"]') +parser.add_argument("--wav_file", type=str, default=None, help="amp fallback number") +parser.add_argument("--quantize", type=str2bool, default=False, help="quantized model") +parser.add_argument( + "--intra_op_num_threads", type=int, default=1, help="intra_op_num_threads for onnx" +) +parser.add_argument("--output_dir", type=str, default=None, help="amp fallback number") args = parser.parse_args() from funasr.runtime.python.libtorch.funasr_torch import Paraformer + if args.backend == "onnx": from funasr.runtime.python.onnxruntime.funasr_onnx import Paraformer - -model = Paraformer(args.model_dir, batch_size=1, quantize=args.quantize, intra_op_num_threads=args.intra_op_num_threads) -wav_file_f = open(args.wav_file, 'r') +model = Paraformer( + args.model_dir, + batch_size=1, + quantize=args.quantize, + intra_op_num_threads=args.intra_op_num_threads, +) + +wav_file_f = open(args.wav_file, "r") wav_files = wav_file_f.readlines() output_dir = args.output_dir if not os.path.exists(output_dir): os.makedirs(output_dir) -if os.name == 'nt': # Windows - newline = '\r\n' -else: # Linux Mac - newline = '\n' +if os.name == "nt": # Windows + newline = "\r\n" +else: # Linux Mac + newline = "\n" text_f = open(os.path.join(output_dir, "text"), "w", newline=newline) token_f = open(os.path.join(output_dir, "token"), "w", newline=newline) for i, wav_path_i in enumerate(wav_files): wav_name, wav_path = wav_path_i.strip().split() result = model(wav_path) - text_i = "{} {}\n".format(wav_name, result[0]['preds'][0]) - token_i = "{} {}\n".format(wav_name, result[0]['preds'][1]) + text_i = "{} {}\n".format(wav_name, result[0]["preds"][0]) + token_i = "{} {}\n".format(wav_name, result[0]["preds"][1]) text_f.write(text_i) text_f.flush() token_f.write(token_i) token_f.flush() text_f.close() token_f.close() - diff --git a/runtime/python/utils/test_rtf.py b/runtime/python/utils/test_rtf.py index 3fe96a3b2..112025d6d 100644 --- a/runtime/python/utils/test_rtf.py +++ b/runtime/python/utils/test_rtf.py @@ -1,55 +1,79 @@ - import time import sys import librosa from funasr.utils.types import str2bool import argparse + parser = argparse.ArgumentParser() -parser.add_argument('--model_dir', type=str, required=True) -parser.add_argument('--backend', type=str, default='onnx', help='["onnx", "torch"]') -parser.add_argument('--wav_file', type=str, default=None, help='amp fallback number') -parser.add_argument('--quantize', type=str2bool, default=False, help='quantized model') -parser.add_argument('--intra_op_num_threads', type=int, default=1, help='intra_op_num_threads for onnx') +parser.add_argument("--model_dir", type=str, required=True) +parser.add_argument("--backend", type=str, default="onnx", help='["onnx", "torch"]') +parser.add_argument("--wav_file", type=str, default=None, help="amp fallback number") +parser.add_argument("--quantize", type=str2bool, default=False, help="quantized model") +parser.add_argument( + "--intra_op_num_threads", type=int, default=1, help="intra_op_num_threads for onnx" +) args = parser.parse_args() from funasr.runtime.python.libtorch.funasr_torch import Paraformer + if args.backend == "onnx": from funasr.runtime.python.onnxruntime.funasr_onnx import Paraformer - -model = Paraformer(args.model_dir, batch_size=1, quantize=args.quantize, intra_op_num_threads=args.intra_op_num_threads) -wav_file_f = open(args.wav_file, 'r') +model = Paraformer( + args.model_dir, + batch_size=1, + quantize=args.quantize, + intra_op_num_threads=args.intra_op_num_threads, +) + +wav_file_f = open(args.wav_file, "r") wav_files = wav_file_f.readlines() # warm-up total = 0.0 num = 30 -wav_path = wav_files[0].split("\t")[1].strip() if "\t" in wav_files[0] else wav_files[0].split(" ")[1].strip() +wav_path = ( + wav_files[0].split("\t")[1].strip() + if "\t" in wav_files[0] + else wav_files[0].split(" ")[1].strip() +) for i in range(num): beg_time = time.time() result = model(wav_path) end_time = time.time() - duration = end_time-beg_time + duration = end_time - beg_time total += duration print(result) - print("num: {}, time, {}, avg: {}, rtf: {}".format(len(wav_path), duration, total/(i+1), (total/(i+1))/5.53)) + print( + "num: {}, time, {}, avg: {}, rtf: {}".format( + len(wav_path), duration, total / (i + 1), (total / (i + 1)) / 5.53 + ) + ) # infer time beg_time = time.time() for i, wav_path_i in enumerate(wav_files): - wav_path = wav_path_i.split("\t")[1].strip() if "\t" in wav_path_i else wav_path_i.split(" ")[1].strip() + wav_path = ( + wav_path_i.split("\t")[1].strip() + if "\t" in wav_path_i + else wav_path_i.split(" ")[1].strip() + ) result = model(wav_path) end_time = time.time() -duration = (end_time-beg_time)*1000 +duration = (end_time - beg_time) * 1000 print("total_time_comput_ms: {}".format(int(duration))) duration_time = 0.0 for i, wav_path_i in enumerate(wav_files): - wav_path = wav_path_i.split("\t")[1].strip() if "\t" in wav_path_i else wav_path_i.split(" ")[1].strip() + wav_path = ( + wav_path_i.split("\t")[1].strip() + if "\t" in wav_path_i + else wav_path_i.split(" ")[1].strip() + ) waveform, _ = librosa.load(wav_path, sr=16000) - duration_time += len(waveform)/16.0 + duration_time += len(waveform) / 16.0 print("total_time_wav_ms: {}".format(int(duration_time))) -print("total_rtf: {:.5}".format(duration/duration_time)) \ No newline at end of file +print("total_rtf: {:.5}".format(duration / duration_time)) diff --git a/runtime/python/utils/test_rtf_gpu.py b/runtime/python/utils/test_rtf_gpu.py index 02d5ac685..38d429401 100644 --- a/runtime/python/utils/test_rtf_gpu.py +++ b/runtime/python/utils/test_rtf_gpu.py @@ -1,58 +1,82 @@ - import time import sys import librosa from funasr.utils.types import str2bool import argparse + parser = argparse.ArgumentParser() -parser.add_argument('--model_dir', type=str, required=True) -parser.add_argument('--backend', type=str, default='onnx', help='["onnx", "torch"]') -parser.add_argument('--wav_file', type=str, default=None, help='amp fallback number') -parser.add_argument('--quantize', type=str2bool, default=False, help='quantized model') -parser.add_argument('--intra_op_num_threads', type=int, default=1, help='intra_op_num_threads for onnx') -parser.add_argument('--batch_size', type=int, default=1, help='batch_size for onnx') +parser.add_argument("--model_dir", type=str, required=True) +parser.add_argument("--backend", type=str, default="onnx", help='["onnx", "torch"]') +parser.add_argument("--wav_file", type=str, default=None, help="amp fallback number") +parser.add_argument("--quantize", type=str2bool, default=False, help="quantized model") +parser.add_argument( + "--intra_op_num_threads", type=int, default=1, help="intra_op_num_threads for onnx" +) +parser.add_argument("--batch_size", type=int, default=1, help="batch_size for onnx") args = parser.parse_args() from funasr.runtime.python.libtorch.funasr_torch import Paraformer + if args.backend == "onnx": from funasr.runtime.python.onnxruntime.funasr_onnx import Paraformer - -model = Paraformer(args.model_dir, batch_size=args.batch_size, quantize=args.quantize, intra_op_num_threads=args.intra_op_num_threads) -wav_file_f = open(args.wav_file, 'r') +model = Paraformer( + args.model_dir, + batch_size=args.batch_size, + quantize=args.quantize, + intra_op_num_threads=args.intra_op_num_threads, +) + +wav_file_f = open(args.wav_file, "r") wav_files = wav_file_f.readlines() # warm-up total = 0.0 num = 30 -wav_path = wav_files[0].split("\t")[1].strip() if "\t" in wav_files[0] else wav_files[0].split(" ")[1].strip() +wav_path = ( + wav_files[0].split("\t")[1].strip() + if "\t" in wav_files[0] + else wav_files[0].split(" ")[1].strip() +) for i in range(num): beg_time = time.time() result = model(wav_path) end_time = time.time() - duration = end_time-beg_time + duration = end_time - beg_time total += duration print(result) - print("num: {}, time, {}, avg: {}, rtf: {}".format(len(wav_path), duration, total/(i+1), (total/(i+1))/5.53)) + print( + "num: {}, time, {}, avg: {}, rtf: {}".format( + len(wav_path), duration, total / (i + 1), (total / (i + 1)) / 5.53 + ) + ) # infer time wav_path = [] beg_time = time.time() for i, wav_path_i in enumerate(wav_files): - wav_path_i = wav_path_i.split("\t")[1].strip() if "\t" in wav_path_i else wav_path_i.split(" ")[1].strip() + wav_path_i = ( + wav_path_i.split("\t")[1].strip() + if "\t" in wav_path_i + else wav_path_i.split(" ")[1].strip() + ) wav_path += [wav_path_i] result = model(wav_path) end_time = time.time() -duration = (end_time-beg_time)*1000 +duration = (end_time - beg_time) * 1000 print("total_time_comput_ms: {}".format(int(duration))) duration_time = 0.0 for i, wav_path_i in enumerate(wav_files): - wav_path = wav_path_i.split("\t")[1].strip() if "\t" in wav_path_i else wav_path_i.split(" ")[1].strip() + wav_path = ( + wav_path_i.split("\t")[1].strip() + if "\t" in wav_path_i + else wav_path_i.split(" ")[1].strip() + ) waveform, _ = librosa.load(wav_path, sr=16000) - duration_time += len(waveform)/16.0 + duration_time += len(waveform) / 16.0 print("total_time_wav_ms: {}".format(int(duration_time))) -print("total_rtf: {:.5}".format(duration/duration_time)) \ No newline at end of file +print("total_rtf: {:.5}".format(duration / duration_time)) diff --git a/runtime/python/websocket/funasr_client_api.py b/runtime/python/websocket/funasr_client_api.py index 093ed5186..3fa242241 100644 --- a/runtime/python/websocket/funasr_client_api.py +++ b/runtime/python/websocket/funasr_client_api.py @@ -1,9 +1,10 @@ -''' +""" Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved. MIT License (https://opensource.org/licenses/MIT) 2022-2023 by zhaomingwork@qq.com -''' +""" + # pip install websocket-client import ssl from websocket import ABNF @@ -14,138 +15,146 @@ import traceback import json import time import numpy as np + + # class for recognizer in websocket -class Funasr_websocket_recognizer(): - ''' +class Funasr_websocket_recognizer: + """ python asr recognizer lib - ''' - def __init__(self, host="127.0.0.1", - port="30035", - is_ssl=True, - chunk_size="0, 10, 5", - chunk_interval=10, - mode="offline", - wav_name="default"): - ''' - host: server host ip - port: server port - is_ssl: True for wss protocal, False for ws - ''' - try: - if is_ssl == True: - ssl_context = ssl.SSLContext() - ssl_context.check_hostname = False - ssl_context.verify_mode = ssl.CERT_NONE - uri = "wss://{}:{}".format(host, port) - ssl_opt={"cert_reqs": ssl.CERT_NONE} - else: - uri = "ws://{}:{}".format(host, port) - ssl_context = None - ssl_opt=None - self.host = host - self.port = port - - self.msg_queue = Queue() # used for recognized result text + """ - print("connect to url",uri) - self.websocket=create_connection(uri, ssl=ssl_context, sslopt=ssl_opt) - - self.thread_msg = threading.Thread(target=Funasr_websocket_recognizer.thread_rec_msg, args=(self,)) - self.thread_msg.start() - chunk_size = [int(x) for x in chunk_size.split(",")] - stride = int(60 * chunk_size[1] / chunk_interval / 1000 * 16000 * 2) - chunk_num = (len(audio_bytes) - 1) // stride + 1 - - message = json.dumps({"mode": mode, - "chunk_size": chunk_size, - "encoder_chunk_look_back": 4, - "decoder_chunk_look_back": 1, - "chunk_interval": chunk_interval, - "wav_name": wav_name, - "is_speaking": True}) - - self.websocket.send(message) - - print("send json",message) - - except Exception as e: + def __init__( + self, + host="127.0.0.1", + port="30035", + is_ssl=True, + chunk_size="0, 10, 5", + chunk_interval=10, + mode="offline", + wav_name="default", + ): + """ + host: server host ip + port: server port + is_ssl: True for wss protocal, False for ws + """ + try: + if is_ssl == True: + ssl_context = ssl.SSLContext() + ssl_context.check_hostname = False + ssl_context.verify_mode = ssl.CERT_NONE + uri = "wss://{}:{}".format(host, port) + ssl_opt = {"cert_reqs": ssl.CERT_NONE} + else: + uri = "ws://{}:{}".format(host, port) + ssl_context = None + ssl_opt = None + self.host = host + self.port = port + + self.msg_queue = Queue() # used for recognized result text + + print("connect to url", uri) + self.websocket = create_connection(uri, ssl=ssl_context, sslopt=ssl_opt) + + self.thread_msg = threading.Thread( + target=Funasr_websocket_recognizer.thread_rec_msg, args=(self,) + ) + self.thread_msg.start() + chunk_size = [int(x) for x in chunk_size.split(",")] + stride = int(60 * chunk_size[1] / chunk_interval / 1000 * 16000 * 2) + chunk_num = (len(audio_bytes) - 1) // stride + 1 + + message = json.dumps( + { + "mode": mode, + "chunk_size": chunk_size, + "encoder_chunk_look_back": 4, + "decoder_chunk_look_back": 1, + "chunk_interval": chunk_interval, + "wav_name": wav_name, + "is_speaking": True, + } + ) + + self.websocket.send(message) + + print("send json", message) + + except Exception as e: print("Exception:", e) traceback.print_exc() - + # threads for rev msg def thread_rec_msg(self): try: - while(True): - msg=self.websocket.recv() - if msg is None or len(msg) == 0: - continue - msg = json.loads(msg) - - self.msg_queue.put(msg) + while True: + msg = self.websocket.recv() + if msg is None or len(msg) == 0: + continue + msg = json.loads(msg) + + self.msg_queue.put(msg) except Exception as e: print("client closed") - + # feed data to asr engine, wait_time means waiting for result until time out def feed_chunk(self, chunk, wait_time=0.01): try: - self.websocket.send(chunk, ABNF.OPCODE_BINARY) + self.websocket.send(chunk, ABNF.OPCODE_BINARY) # loop to check if there is a message, timeout in 0.01s - while(True): - msg = self.msg_queue.get(timeout=wait_time) - if self.msg_queue.empty(): - break - + while True: + msg = self.msg_queue.get(timeout=wait_time) + if self.msg_queue.empty(): + break + return msg except: return "" - - def close(self,timeout=1): + + def close(self, timeout=1): message = json.dumps({"is_speaking": False}) self.websocket.send(message) # sleep for timeout seconds to wait for result time.sleep(timeout) - msg="" - while(not self.msg_queue.empty()): + msg = "" + while not self.msg_queue.empty(): msg = self.msg_queue.get() - + self.websocket.close() # only resturn the last msg return msg - -if __name__ == '__main__': - - print('example for Funasr_websocket_recognizer') + + +if __name__ == "__main__": + + print("example for Funasr_websocket_recognizer") import wave + wav_path = "/Users/zhifu/Downloads/modelscope_models/speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch/example/asr_example.wav" with wave.open(wav_path, "rb") as wav_file: - params = wav_file.getparams() - frames = wav_file.readframes(wav_file.getnframes()) - audio_bytes = bytes(frames) - - + params = wav_file.getparams() + frames = wav_file.readframes(wav_file.getnframes()) + audio_bytes = bytes(frames) + stride = int(60 * 10 / 10 / 1000 * 16000 * 2) chunk_num = (len(audio_bytes) - 1) // stride + 1 - # create an recognizer - rcg = Funasr_websocket_recognizer(host="127.0.0.1", - port="10095", - is_ssl=True, - mode="2pass", - chunk_size="0,10,5") + # create an recognizer + rcg = Funasr_websocket_recognizer( + host="127.0.0.1", port="10095", is_ssl=True, mode="2pass", chunk_size="0,10,5" + ) # loop to send chunk for i in range(chunk_num): - beg = i * stride - data = audio_bytes[beg:beg + stride] - - text = rcg.feed_chunk(data,wait_time=0.02) - if len(text)>0: - print("text",text) - time.sleep(0.05) - + beg = i * stride + data = audio_bytes[beg : beg + stride] + + text = rcg.feed_chunk(data, wait_time=0.02) + if len(text) > 0: + print("text", text) + time.sleep(0.05) + # get last message text = rcg.close(timeout=3) - print("text",text) - - - + print("text", text) diff --git a/runtime/python/websocket/funasr_wss_client.py b/runtime/python/websocket/funasr_wss_client.py index b30964a92..dfea7da7e 100644 --- a/runtime/python/websocket/funasr_wss_client.py +++ b/runtime/python/websocket/funasr_wss_client.py @@ -3,11 +3,13 @@ import os import time import websockets, ssl import asyncio + # import threading import argparse import json import traceback from multiprocessing import Process + # from funasr.fileio.datadir_writer import DatadirWriter import logging @@ -15,72 +17,34 @@ import logging logging.basicConfig(level=logging.ERROR) parser = argparse.ArgumentParser() -parser.add_argument("--host", - type=str, - default="localhost", - required=False, - help="host ip, localhost, 0.0.0.0") -parser.add_argument("--port", - type=int, - default=10095, - required=False, - help="grpc server port") -parser.add_argument("--chunk_size", - type=str, - default="5, 10, 5", - help="chunk") -parser.add_argument("--encoder_chunk_look_back", - type=int, - default=4, - help="chunk") -parser.add_argument("--decoder_chunk_look_back", - type=int, - default=0, - help="chunk") -parser.add_argument("--chunk_interval", - type=int, - default=10, - help="chunk") -parser.add_argument("--hotword", - type=str, - default="", - help="hotword file path, one hotword perline (e.g.:阿里巴巴 20)") -parser.add_argument("--audio_in", - type=str, - default=None, - help="audio_in") -parser.add_argument("--audio_fs", - type=int, - default=16000, - help="audio_fs") -parser.add_argument("--send_without_sleep", - action="store_true", - default=True, - help="if audio_in is set, send_without_sleep") -parser.add_argument("--thread_num", - type=int, - default=1, - help="thread_num") -parser.add_argument("--words_max_print", - type=int, - default=10000, - help="chunk") -parser.add_argument("--output_dir", - type=str, - default=None, - help="output_dir") -parser.add_argument("--ssl", - type=int, - default=1, - help="1 for ssl connect, 0 for no ssl") -parser.add_argument("--use_itn", - type=int, - default=1, - help="1 for using itn, 0 for not itn") -parser.add_argument("--mode", - type=str, - default="2pass", - help="offline, online, 2pass") +parser.add_argument( + "--host", type=str, default="localhost", required=False, help="host ip, localhost, 0.0.0.0" +) +parser.add_argument("--port", type=int, default=10095, required=False, help="grpc server port") +parser.add_argument("--chunk_size", type=str, default="5, 10, 5", help="chunk") +parser.add_argument("--encoder_chunk_look_back", type=int, default=4, help="chunk") +parser.add_argument("--decoder_chunk_look_back", type=int, default=0, help="chunk") +parser.add_argument("--chunk_interval", type=int, default=10, help="chunk") +parser.add_argument( + "--hotword", + type=str, + default="", + help="hotword file path, one hotword perline (e.g.:阿里巴巴 20)", +) +parser.add_argument("--audio_in", type=str, default=None, help="audio_in") +parser.add_argument("--audio_fs", type=int, default=16000, help="audio_fs") +parser.add_argument( + "--send_without_sleep", + action="store_true", + default=True, + help="if audio_in is set, send_without_sleep", +) +parser.add_argument("--thread_num", type=int, default=1, help="thread_num") +parser.add_argument("--words_max_print", type=int, default=10000, help="chunk") +parser.add_argument("--output_dir", type=str, default=None, help="output_dir") +parser.add_argument("--ssl", type=int, default=1, help="1 for ssl connect, 0 for no ssl") +parser.add_argument("--use_itn", type=int, default=1, help="1 for using itn, 0 for not itn") +parser.add_argument("--mode", type=str, default="2pass", help="offline, online, 2pass") args = parser.parse_args() args.chunk_size = [int(x) for x in args.chunk_size.split(",")] @@ -89,12 +53,12 @@ print(args) from queue import Queue voices = Queue() -offline_msg_done=False +offline_msg_done = False if args.output_dir is not None: # if os.path.exists(args.output_dir): # os.remove(args.output_dir) - + if not os.path.exists(args.output_dir): os.makedirs(args.output_dir) @@ -102,6 +66,7 @@ if args.output_dir is not None: async def record_microphone(): is_finished = False import pyaudio + # print("2") global voices FORMAT = pyaudio.paInt16 @@ -112,11 +77,9 @@ async def record_microphone(): p = pyaudio.PyAudio() - stream = p.open(format=FORMAT, - channels=CHANNELS, - rate=RATE, - input=True, - frames_per_buffer=CHUNK) + stream = p.open( + format=FORMAT, channels=CHANNELS, rate=RATE, input=True, frames_per_buffer=CHUNK + ) # hotwords fst_dict = {} hotword_msg = "" @@ -139,27 +102,31 @@ async def record_microphone(): use_itn = True if args.use_itn == 0: - use_itn=False - - message = json.dumps({"mode": args.mode, - "chunk_size": args.chunk_size, - "chunk_interval": args.chunk_interval, - "encoder_chunk_look_back": args.encoder_chunk_look_back, - "decoder_chunk_look_back": args.decoder_chunk_look_back, - "wav_name": "microphone", - "is_speaking": True, - "hotwords": hotword_msg, - "itn": use_itn, - }) - #voices.put(message) + use_itn = False + + message = json.dumps( + { + "mode": args.mode, + "chunk_size": args.chunk_size, + "chunk_interval": args.chunk_interval, + "encoder_chunk_look_back": args.encoder_chunk_look_back, + "decoder_chunk_look_back": args.decoder_chunk_look_back, + "wav_name": "microphone", + "is_speaking": True, + "hotwords": hotword_msg, + "itn": use_itn, + } + ) + # voices.put(message) await websocket.send(message) while True: data = stream.read(CHUNK) message = data - #voices.put(message) + # voices.put(message) await websocket.send(message) await asyncio.sleep(0.005) + async def record_from_scp(chunk_begin, chunk_size): global voices is_finished = False @@ -188,33 +155,34 @@ async def record_from_scp(chunk_begin, chunk_size): hotword_msg = json.dumps(fst_dict) else: hotword_msg = args.hotword - print (hotword_msg) + print(hotword_msg) sample_rate = args.audio_fs wav_format = "pcm" - use_itn=True + use_itn = True if args.use_itn == 0: - use_itn=False - + use_itn = False + if chunk_size > 0: - wavs = wavs[chunk_begin:chunk_begin + chunk_size] + wavs = wavs[chunk_begin : chunk_begin + chunk_size] for wav in wavs: wav_splits = wav.strip().split() - + wav_name = wav_splits[0] if len(wav_splits) > 1 else "demo" wav_path = wav_splits[1] if len(wav_splits) > 1 else wav_splits[0] - if not len(wav_path.strip())>0: - continue + if not len(wav_path.strip()) > 0: + continue if wav_path.endswith(".pcm"): with open(wav_path, "rb") as f: audio_bytes = f.read() elif wav_path.endswith(".wav"): import wave + with wave.open(wav_path, "rb") as wav_file: params = wav_file.getparams() sample_rate = wav_file.getframerate() frames = wav_file.readframes(wav_file.getnframes()) - audio_bytes = bytes(frames) + audio_bytes = bytes(frames) else: wav_format = "others" with open(wav_path, "rb") as f: @@ -225,88 +193,97 @@ async def record_from_scp(chunk_begin, chunk_size): # print(stride) # send first time - message = json.dumps({"mode": args.mode, - "chunk_size": args.chunk_size, - "chunk_interval": args.chunk_interval, - "encoder_chunk_look_back": args.encoder_chunk_look_back, - "decoder_chunk_look_back": args.decoder_chunk_look_back, - "audio_fs":sample_rate, - "wav_name": wav_name, - "wav_format": wav_format, - "is_speaking": True, - "hotwords": hotword_msg, - "itn": use_itn}) + message = json.dumps( + { + "mode": args.mode, + "chunk_size": args.chunk_size, + "chunk_interval": args.chunk_interval, + "encoder_chunk_look_back": args.encoder_chunk_look_back, + "decoder_chunk_look_back": args.decoder_chunk_look_back, + "audio_fs": sample_rate, + "wav_name": wav_name, + "wav_format": wav_format, + "is_speaking": True, + "hotwords": hotword_msg, + "itn": use_itn, + } + ) - #voices.put(message) + # voices.put(message) await websocket.send(message) is_speaking = True for i in range(chunk_num): beg = i * stride - data = audio_bytes[beg:beg + stride] + data = audio_bytes[beg : beg + stride] message = data - #voices.put(message) + # voices.put(message) await websocket.send(message) if i == chunk_num - 1: is_speaking = False message = json.dumps({"is_speaking": is_speaking}) - #voices.put(message) + # voices.put(message) await websocket.send(message) - - sleep_duration = 0.001 if args.mode == "offline" else 60 * args.chunk_size[1] / args.chunk_interval / 1000 - + + sleep_duration = ( + 0.001 + if args.mode == "offline" + else 60 * args.chunk_size[1] / args.chunk_interval / 1000 + ) + await asyncio.sleep(sleep_duration) - - if not args.mode=="offline": + + if not args.mode == "offline": await asyncio.sleep(2) # offline model need to wait for message recved - - if args.mode=="offline": - global offline_msg_done - while not offline_msg_done: - await asyncio.sleep(1) - + + if args.mode == "offline": + global offline_msg_done + while not offline_msg_done: + await asyncio.sleep(1) + await websocket.close() - async def message(id): - global websocket,voices,offline_msg_done + global websocket, voices, offline_msg_done text_print = "" text_print_2pass_online = "" text_print_2pass_offline = "" if args.output_dir is not None: - ibest_writer = open(os.path.join(args.output_dir, "text.{}".format(id)), "a", encoding="utf-8") + ibest_writer = open( + os.path.join(args.output_dir, "text.{}".format(id)), "a", encoding="utf-8" + ) else: ibest_writer = None try: - while True: - + while True: + meg = await websocket.recv() meg = json.loads(meg) wav_name = meg.get("wav_name", "demo") text = meg["text"] - timestamp="" + timestamp = "" offline_msg_done = meg.get("is_final", False) if "timestamp" in meg: timestamp = meg["timestamp"] if ibest_writer is not None: - if timestamp !="": + if timestamp != "": text_write_line = "{}\t{}\t{}\n".format(wav_name, text, timestamp) else: text_write_line = "{}\t{}\n".format(wav_name, text) ibest_writer.write(text_write_line) - if 'mode' not in meg: + if "mode" not in meg: continue if meg["mode"] == "online": text_print += "{}".format(text) - text_print = text_print[-args.words_max_print:] - os.system('clear') + text_print = text_print[-args.words_max_print :] + os.system("clear") print("\rpid" + str(id) + ": " + text_print) elif meg["mode"] == "offline": - if timestamp !="": + if timestamp != "": text_print += "{} timestamp: {}".format(text, timestamp) else: text_print += "{}".format(text) @@ -323,58 +300,59 @@ async def message(id): text_print_2pass_online = "" text_print = text_print_2pass_offline + "{}".format(text) text_print_2pass_offline += "{}".format(text) - text_print = text_print[-args.words_max_print:] - os.system('clear') + text_print = text_print[-args.words_max_print :] + os.system("clear") print("\rpid" + str(id) + ": " + text_print) # offline_msg_done=True except Exception as e: - print("Exception:", e) - #traceback.print_exc() - #await websocket.close() - - + print("Exception:", e) + # traceback.print_exc() + # await websocket.close() async def ws_client(id, chunk_begin, chunk_size): - if args.audio_in is None: - chunk_begin=0 - chunk_size=1 - global websocket,voices,offline_msg_done - - for i in range(chunk_begin,chunk_begin+chunk_size): - offline_msg_done=False - voices = Queue() - if args.ssl == 1: - ssl_context = ssl.SSLContext() - ssl_context.check_hostname = False - ssl_context.verify_mode = ssl.CERT_NONE - uri = "wss://{}:{}".format(args.host, args.port) - else: - uri = "ws://{}:{}".format(args.host, args.port) - ssl_context = None - print("connect to", uri) - async with websockets.connect(uri, subprotocols=["binary"], ping_interval=None, ssl=ssl_context) as websocket: - if args.audio_in is not None: - task = asyncio.create_task(record_from_scp(i, 1)) + if args.audio_in is None: + chunk_begin = 0 + chunk_size = 1 + global websocket, voices, offline_msg_done + + for i in range(chunk_begin, chunk_begin + chunk_size): + offline_msg_done = False + voices = Queue() + if args.ssl == 1: + ssl_context = ssl.SSLContext() + ssl_context.check_hostname = False + ssl_context.verify_mode = ssl.CERT_NONE + uri = "wss://{}:{}".format(args.host, args.port) else: - task = asyncio.create_task(record_microphone()) - task3 = asyncio.create_task(message(str(id)+"_"+str(i))) #processid+fileid - await asyncio.gather(task, task3) - exit(0) - + uri = "ws://{}:{}".format(args.host, args.port) + ssl_context = None + print("connect to", uri) + async with websockets.connect( + uri, subprotocols=["binary"], ping_interval=None, ssl=ssl_context + ) as websocket: + if args.audio_in is not None: + task = asyncio.create_task(record_from_scp(i, 1)) + else: + task = asyncio.create_task(record_microphone()) + task3 = asyncio.create_task(message(str(id) + "_" + str(i))) # processid+fileid + await asyncio.gather(task, task3) + exit(0) + def one_thread(id, chunk_begin, chunk_size): asyncio.get_event_loop().run_until_complete(ws_client(id, chunk_begin, chunk_size)) asyncio.get_event_loop().run_forever() -if __name__ == '__main__': + +if __name__ == "__main__": # for microphone if args.audio_in is None: p = Process(target=one_thread, args=(0, 0, 0)) p.start() p.join() - print('end') + print("end") else: # calculate the number of wavs for each preocess if args.audio_in.endswith(".scp"): @@ -388,7 +366,6 @@ if __name__ == '__main__': wav_path = wav_splits[1] if len(wav_splits) > 1 else wav_splits[0] audio_type = os.path.splitext(wav_path)[-1].lower() - total_len = len(wavs) if total_len >= args.thread_num: chunk_size = int(total_len / args.thread_num) @@ -413,4 +390,4 @@ if __name__ == '__main__': for i in process_list: p.join() - print('end') + print("end") diff --git a/runtime/python/websocket/funasr_wss_server.py b/runtime/python/websocket/funasr_wss_server.py index 015d87b80..9ee30513f 100644 --- a/runtime/python/websocket/funasr_wss_server.py +++ b/runtime/python/websocket/funasr_wss_server.py @@ -10,71 +10,56 @@ import ssl parser = argparse.ArgumentParser() -parser.add_argument("--host", - type=str, - default="0.0.0.0", - required=False, - help="host ip, localhost, 0.0.0.0") -parser.add_argument("--port", - type=int, - default=10095, - required=False, - help="grpc server port") -parser.add_argument("--asr_model", - type=str, - default="iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch", - help="model from modelscope") -parser.add_argument("--asr_model_revision", - type=str, - default="v2.0.4", - help="") -parser.add_argument("--asr_model_online", - type=str, - default="iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-online", - help="model from modelscope") -parser.add_argument("--asr_model_online_revision", - type=str, - default="v2.0.4", - help="") -parser.add_argument("--vad_model", - type=str, - default="iic/speech_fsmn_vad_zh-cn-16k-common-pytorch", - help="model from modelscope") -parser.add_argument("--vad_model_revision", - type=str, - default="v2.0.4", - help="") -parser.add_argument("--punc_model", - type=str, - default="iic/punc_ct-transformer_zh-cn-common-vad_realtime-vocab272727", - help="model from modelscope") -parser.add_argument("--punc_model_revision", - type=str, - default="v2.0.4", - help="") -parser.add_argument("--ngpu", - type=int, - default=1, - help="0 for cpu, 1 for gpu") -parser.add_argument("--device", - type=str, - default="cuda", - help="cuda, cpu") -parser.add_argument("--ncpu", - type=int, - default=4, - help="cpu cores") -parser.add_argument("--certfile", - type=str, - default="../../ssl_key/server.crt", - required=False, - help="certfile for ssl") +parser.add_argument( + "--host", type=str, default="0.0.0.0", required=False, help="host ip, localhost, 0.0.0.0" +) +parser.add_argument("--port", type=int, default=10095, required=False, help="grpc server port") +parser.add_argument( + "--asr_model", + type=str, + default="iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch", + help="model from modelscope", +) +parser.add_argument("--asr_model_revision", type=str, default="v2.0.4", help="") +parser.add_argument( + "--asr_model_online", + type=str, + default="iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-online", + help="model from modelscope", +) +parser.add_argument("--asr_model_online_revision", type=str, default="v2.0.4", help="") +parser.add_argument( + "--vad_model", + type=str, + default="iic/speech_fsmn_vad_zh-cn-16k-common-pytorch", + help="model from modelscope", +) +parser.add_argument("--vad_model_revision", type=str, default="v2.0.4", help="") +parser.add_argument( + "--punc_model", + type=str, + default="iic/punc_ct-transformer_zh-cn-common-vad_realtime-vocab272727", + help="model from modelscope", +) +parser.add_argument("--punc_model_revision", type=str, default="v2.0.4", help="") +parser.add_argument("--ngpu", type=int, default=1, help="0 for cpu, 1 for gpu") +parser.add_argument("--device", type=str, default="cuda", help="cuda, cpu") +parser.add_argument("--ncpu", type=int, default=4, help="cpu cores") +parser.add_argument( + "--certfile", + type=str, + default="../../ssl_key/server.crt", + required=False, + help="certfile for ssl", +) -parser.add_argument("--keyfile", - type=str, - default="../../ssl_key/server.key", - required=False, - help="keyfile for ssl") +parser.add_argument( + "--keyfile", + type=str, + default="../../ssl_key/server.key", + required=False, + help="keyfile for ssl", +) args = parser.parse_args() @@ -84,232 +69,266 @@ print("model loading") from funasr import AutoModel # asr -model_asr = AutoModel(model=args.asr_model, - model_revision=args.asr_model_revision, - ngpu=args.ngpu, - ncpu=args.ncpu, - device=args.device, - disable_pbar=True, - disable_log=True, - ) +model_asr = AutoModel( + model=args.asr_model, + model_revision=args.asr_model_revision, + ngpu=args.ngpu, + ncpu=args.ncpu, + device=args.device, + disable_pbar=True, + disable_log=True, +) # asr -model_asr_streaming = AutoModel(model=args.asr_model_online, - model_revision=args.asr_model_online_revision, - ngpu=args.ngpu, - ncpu=args.ncpu, - device=args.device, - disable_pbar=True, - disable_log=True, - ) +model_asr_streaming = AutoModel( + model=args.asr_model_online, + model_revision=args.asr_model_online_revision, + ngpu=args.ngpu, + ncpu=args.ncpu, + device=args.device, + disable_pbar=True, + disable_log=True, +) # vad -model_vad = AutoModel(model=args.vad_model, - model_revision=args.vad_model_revision, - ngpu=args.ngpu, - ncpu=args.ncpu, - device=args.device, - disable_pbar=True, - disable_log=True, - # chunk_size=60, - ) +model_vad = AutoModel( + model=args.vad_model, + model_revision=args.vad_model_revision, + ngpu=args.ngpu, + ncpu=args.ncpu, + device=args.device, + disable_pbar=True, + disable_log=True, + # chunk_size=60, +) if args.punc_model != "": - model_punc = AutoModel(model=args.punc_model, - model_revision=args.punc_model_revision, - ngpu=args.ngpu, - ncpu=args.ncpu, - device=args.device, - disable_pbar=True, - disable_log=True, - ) + model_punc = AutoModel( + model=args.punc_model, + model_revision=args.punc_model_revision, + ngpu=args.ngpu, + ncpu=args.ncpu, + device=args.device, + disable_pbar=True, + disable_log=True, + ) else: - model_punc = None - + model_punc = None print("model loaded! only support one client at the same time now!!!!") -async def ws_reset(websocket): - print("ws reset now, total num is ",len(websocket_users)) - websocket.status_dict_asr_online["cache"] = {} - websocket.status_dict_asr_online["is_final"] = True - websocket.status_dict_vad["cache"] = {} - websocket.status_dict_vad["is_final"] = True - websocket.status_dict_punc["cache"] = {} - - await websocket.close() +async def ws_reset(websocket): + print("ws reset now, total num is ", len(websocket_users)) + + websocket.status_dict_asr_online["cache"] = {} + websocket.status_dict_asr_online["is_final"] = True + websocket.status_dict_vad["cache"] = {} + websocket.status_dict_vad["is_final"] = True + websocket.status_dict_punc["cache"] = {} + + await websocket.close() async def clear_websocket(): - for websocket in websocket_users: - await ws_reset(websocket) - websocket_users.clear() - + for websocket in websocket_users: + await ws_reset(websocket) + websocket_users.clear() async def ws_serve(websocket, path): - frames = [] - frames_asr = [] - frames_asr_online = [] - global websocket_users - # await clear_websocket() - websocket_users.add(websocket) - websocket.status_dict_asr = {} - websocket.status_dict_asr_online = {"cache": {}, "is_final": False} - websocket.status_dict_vad = {'cache': {}, "is_final": False} - websocket.status_dict_punc = {'cache': {}} - websocket.chunk_interval = 10 - websocket.vad_pre_idx = 0 - speech_start = False - speech_end_i = -1 - websocket.wav_name = "microphone" - websocket.mode = "2pass" - print("new user connected", flush=True) - - try: - async for message in websocket: - if isinstance(message, str): - messagejson = json.loads(message) - - if "is_speaking" in messagejson: - websocket.is_speaking = messagejson["is_speaking"] - websocket.status_dict_asr_online["is_final"] = not websocket.is_speaking - if "chunk_interval" in messagejson: - websocket.chunk_interval = messagejson["chunk_interval"] - if "wav_name" in messagejson: - websocket.wav_name = messagejson.get("wav_name") - if "chunk_size" in messagejson: - chunk_size = messagejson["chunk_size"] - if isinstance(chunk_size, str): - chunk_size = chunk_size.split(',') - websocket.status_dict_asr_online["chunk_size"] = [int(x) for x in chunk_size] - if "encoder_chunk_look_back" in messagejson: - websocket.status_dict_asr_online["encoder_chunk_look_back"] = messagejson["encoder_chunk_look_back"] - if "decoder_chunk_look_back" in messagejson: - websocket.status_dict_asr_online["decoder_chunk_look_back"] = messagejson["decoder_chunk_look_back"] - if "hotword" in messagejson: - websocket.status_dict_asr["hotword"] = messagejson["hotword"] - if "mode" in messagejson: - websocket.mode = messagejson["mode"] - - websocket.status_dict_vad["chunk_size"] = int(websocket.status_dict_asr_online["chunk_size"][1]*60/websocket.chunk_interval) - if len(frames_asr_online) > 0 or len(frames_asr) > 0 or not isinstance(message, str): - if not isinstance(message, str): - frames.append(message) - duration_ms = len(message)//32 - websocket.vad_pre_idx += duration_ms - - # asr online - frames_asr_online.append(message) - websocket.status_dict_asr_online["is_final"] = speech_end_i != -1 - if len(frames_asr_online) % websocket.chunk_interval == 0 or websocket.status_dict_asr_online["is_final"]: - if websocket.mode == "2pass" or websocket.mode == "online": - audio_in = b"".join(frames_asr_online) - try: - await async_asr_online(websocket, audio_in) - except: - print(f"error in asr streaming, {websocket.status_dict_asr_online}") - frames_asr_online = [] - if speech_start: - frames_asr.append(message) - # vad online - try: - speech_start_i, speech_end_i = await async_vad(websocket, message) - except: - print("error in vad") - if speech_start_i != -1: - speech_start = True - beg_bias = (websocket.vad_pre_idx-speech_start_i)//duration_ms - frames_pre = frames[-beg_bias:] - frames_asr = [] - frames_asr.extend(frames_pre) - # asr punc offline - if speech_end_i != -1 or not websocket.is_speaking: - # print("vad end point") - if websocket.mode == "2pass" or websocket.mode == "offline": - audio_in = b"".join(frames_asr) - try: - await async_asr(websocket, audio_in) - except: - print("error in asr offline") - frames_asr = [] - speech_start = False - frames_asr_online = [] - websocket.status_dict_asr_online["cache"] = {} - if not websocket.is_speaking: - websocket.vad_pre_idx = 0 - frames = [] - websocket.status_dict_vad["cache"] = {} - else: - frames = frames[-20:] - - - except websockets.ConnectionClosed: - print("ConnectionClosed...", websocket_users,flush=True) - await ws_reset(websocket) - websocket_users.remove(websocket) - except websockets.InvalidState: - print("InvalidState...") - except Exception as e: - print("Exception:", e) + frames = [] + frames_asr = [] + frames_asr_online = [] + global websocket_users + # await clear_websocket() + websocket_users.add(websocket) + websocket.status_dict_asr = {} + websocket.status_dict_asr_online = {"cache": {}, "is_final": False} + websocket.status_dict_vad = {"cache": {}, "is_final": False} + websocket.status_dict_punc = {"cache": {}} + websocket.chunk_interval = 10 + websocket.vad_pre_idx = 0 + speech_start = False + speech_end_i = -1 + websocket.wav_name = "microphone" + websocket.mode = "2pass" + print("new user connected", flush=True) + + try: + async for message in websocket: + if isinstance(message, str): + messagejson = json.loads(message) + + if "is_speaking" in messagejson: + websocket.is_speaking = messagejson["is_speaking"] + websocket.status_dict_asr_online["is_final"] = not websocket.is_speaking + if "chunk_interval" in messagejson: + websocket.chunk_interval = messagejson["chunk_interval"] + if "wav_name" in messagejson: + websocket.wav_name = messagejson.get("wav_name") + if "chunk_size" in messagejson: + chunk_size = messagejson["chunk_size"] + if isinstance(chunk_size, str): + chunk_size = chunk_size.split(",") + websocket.status_dict_asr_online["chunk_size"] = [int(x) for x in chunk_size] + if "encoder_chunk_look_back" in messagejson: + websocket.status_dict_asr_online["encoder_chunk_look_back"] = messagejson[ + "encoder_chunk_look_back" + ] + if "decoder_chunk_look_back" in messagejson: + websocket.status_dict_asr_online["decoder_chunk_look_back"] = messagejson[ + "decoder_chunk_look_back" + ] + if "hotword" in messagejson: + websocket.status_dict_asr["hotword"] = messagejson["hotword"] + if "mode" in messagejson: + websocket.mode = messagejson["mode"] + + websocket.status_dict_vad["chunk_size"] = int( + websocket.status_dict_asr_online["chunk_size"][1] * 60 / websocket.chunk_interval + ) + if len(frames_asr_online) > 0 or len(frames_asr) > 0 or not isinstance(message, str): + if not isinstance(message, str): + frames.append(message) + duration_ms = len(message) // 32 + websocket.vad_pre_idx += duration_ms + + # asr online + frames_asr_online.append(message) + websocket.status_dict_asr_online["is_final"] = speech_end_i != -1 + if ( + len(frames_asr_online) % websocket.chunk_interval == 0 + or websocket.status_dict_asr_online["is_final"] + ): + if websocket.mode == "2pass" or websocket.mode == "online": + audio_in = b"".join(frames_asr_online) + try: + await async_asr_online(websocket, audio_in) + except: + print(f"error in asr streaming, {websocket.status_dict_asr_online}") + frames_asr_online = [] + if speech_start: + frames_asr.append(message) + # vad online + try: + speech_start_i, speech_end_i = await async_vad(websocket, message) + except: + print("error in vad") + if speech_start_i != -1: + speech_start = True + beg_bias = (websocket.vad_pre_idx - speech_start_i) // duration_ms + frames_pre = frames[-beg_bias:] + frames_asr = [] + frames_asr.extend(frames_pre) + # asr punc offline + if speech_end_i != -1 or not websocket.is_speaking: + # print("vad end point") + if websocket.mode == "2pass" or websocket.mode == "offline": + audio_in = b"".join(frames_asr) + try: + await async_asr(websocket, audio_in) + except: + print("error in asr offline") + frames_asr = [] + speech_start = False + frames_asr_online = [] + websocket.status_dict_asr_online["cache"] = {} + if not websocket.is_speaking: + websocket.vad_pre_idx = 0 + frames = [] + websocket.status_dict_vad["cache"] = {} + else: + frames = frames[-20:] + + except websockets.ConnectionClosed: + print("ConnectionClosed...", websocket_users, flush=True) + await ws_reset(websocket) + websocket_users.remove(websocket) + except websockets.InvalidState: + print("InvalidState...") + except Exception as e: + print("Exception:", e) async def async_vad(websocket, audio_in): - - segments_result = model_vad.generate(input=audio_in, **websocket.status_dict_vad)[0]["value"] - # print(segments_result) - - speech_start = -1 - speech_end = -1 - - if len(segments_result) == 0 or len(segments_result) > 1: - return speech_start, speech_end - if segments_result[0][0] != -1: - speech_start = segments_result[0][0] - if segments_result[0][1] != -1: - speech_end = segments_result[0][1] - return speech_start, speech_end + + segments_result = model_vad.generate(input=audio_in, **websocket.status_dict_vad)[0]["value"] + # print(segments_result) + + speech_start = -1 + speech_end = -1 + + if len(segments_result) == 0 or len(segments_result) > 1: + return speech_start, speech_end + if segments_result[0][0] != -1: + speech_start = segments_result[0][0] + if segments_result[0][1] != -1: + speech_end = segments_result[0][1] + return speech_start, speech_end async def async_asr(websocket, audio_in): - if len(audio_in) > 0: - # print(len(audio_in)) - rec_result = model_asr.generate(input=audio_in, **websocket.status_dict_asr)[0] - # print("offline_asr, ", rec_result) - if model_punc is not None and len(rec_result["text"])>0: - # print("offline, before punc", rec_result, "cache", websocket.status_dict_punc) - rec_result = model_punc.generate(input=rec_result['text'], **websocket.status_dict_punc)[0] - # print("offline, after punc", rec_result) - if len(rec_result["text"])>0: - # print("offline", rec_result) - mode = "2pass-offline" if "2pass" in websocket.mode else websocket.mode - message = json.dumps({"mode": mode, "text": rec_result["text"], "wav_name": websocket.wav_name,"is_final":websocket.is_speaking}) - await websocket.send(message) + if len(audio_in) > 0: + # print(len(audio_in)) + rec_result = model_asr.generate(input=audio_in, **websocket.status_dict_asr)[0] + # print("offline_asr, ", rec_result) + if model_punc is not None and len(rec_result["text"]) > 0: + # print("offline, before punc", rec_result, "cache", websocket.status_dict_punc) + rec_result = model_punc.generate( + input=rec_result["text"], **websocket.status_dict_punc + )[0] + # print("offline, after punc", rec_result) + if len(rec_result["text"]) > 0: + # print("offline", rec_result) + mode = "2pass-offline" if "2pass" in websocket.mode else websocket.mode + message = json.dumps( + { + "mode": mode, + "text": rec_result["text"], + "wav_name": websocket.wav_name, + "is_final": websocket.is_speaking, + } + ) + await websocket.send(message) async def async_asr_online(websocket, audio_in): - if len(audio_in) > 0: - # print(websocket.status_dict_asr_online.get("is_final", False)) - rec_result = model_asr_streaming.generate(input=audio_in, **websocket.status_dict_asr_online)[0] - # print("online, ", rec_result) - if websocket.mode == "2pass" and websocket.status_dict_asr_online.get("is_final", False): - return - # websocket.status_dict_asr_online["cache"] = dict() - if len(rec_result["text"]): - mode = "2pass-online" if "2pass" in websocket.mode else websocket.mode - message = json.dumps({"mode": mode, "text": rec_result["text"], "wav_name": websocket.wav_name,"is_final":websocket.is_speaking}) - await websocket.send(message) + if len(audio_in) > 0: + # print(websocket.status_dict_asr_online.get("is_final", False)) + rec_result = model_asr_streaming.generate( + input=audio_in, **websocket.status_dict_asr_online + )[0] + # print("online, ", rec_result) + if websocket.mode == "2pass" and websocket.status_dict_asr_online.get("is_final", False): + return + # websocket.status_dict_asr_online["cache"] = dict() + if len(rec_result["text"]): + mode = "2pass-online" if "2pass" in websocket.mode else websocket.mode + message = json.dumps( + { + "mode": mode, + "text": rec_result["text"], + "wav_name": websocket.wav_name, + "is_final": websocket.is_speaking, + } + ) + await websocket.send(message) -if len(args.certfile)>0: - ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) - - # Generate with Lets Encrypt, copied to this location, chown to current user and 400 permissions - ssl_cert = args.certfile - ssl_key = args.keyfile - - ssl_context.load_cert_chain(ssl_cert, keyfile=ssl_key) - start_server = websockets.serve(ws_serve, args.host, args.port, subprotocols=["binary"], ping_interval=None,ssl=ssl_context) + +if len(args.certfile) > 0: + ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + + # Generate with Lets Encrypt, copied to this location, chown to current user and 400 permissions + ssl_cert = args.certfile + ssl_key = args.keyfile + + ssl_context.load_cert_chain(ssl_cert, keyfile=ssl_key) + start_server = websockets.serve( + ws_serve, args.host, args.port, subprotocols=["binary"], ping_interval=None, ssl=ssl_context + ) else: - start_server = websockets.serve(ws_serve, args.host, args.port, subprotocols=["binary"], ping_interval=None) + start_server = websockets.serve( + ws_serve, args.host, args.port, subprotocols=["binary"], ping_interval=None + ) asyncio.get_event_loop().run_until_complete(start_server) asyncio.get_event_loop().run_forever() diff --git a/runtime/tools/fst/ctc_token_fst.py b/runtime/tools/fst/ctc_token_fst.py index e4cdabd4b..5f06e5217 100755 --- a/runtime/tools/fst/ctc_token_fst.py +++ b/runtime/tools/fst/ctc_token_fst.py @@ -4,21 +4,20 @@ import sys -with open(sys.argv[1], 'r') as fread: - print('0 0 ') +with open(sys.argv[1], "r") as fread: + print("0 0 ") nodeX = 1 for entry in fread.readlines(): - entry = entry.replace('\n', '').strip() - fields = entry.split(' ') + entry = entry.replace("\n", "").strip() + fields = entry.split(" ") phone = fields[0] - if phone == '' or phone == '': + if phone == "" or phone == "": continue - if '#' in phone: - print(str(0) + ' ' + str(0) + ' ' + '' + ' ' + phone) + if "#" in phone: + print(str(0) + " " + str(0) + " " + "" + " " + phone) else: - print(str(0) + ' ' + str(nodeX) + ' ' + phone + ' ' + phone) - print(str(nodeX) + ' ' + str(nodeX) + ' ' + phone + ' ') - print(str(nodeX) + ' ' + str(0) + ' ' + ' ') + print(str(0) + " " + str(nodeX) + " " + phone + " " + phone) + print(str(nodeX) + " " + str(nodeX) + " " + phone + " ") + print(str(nodeX) + " " + str(0) + " " + " ") nodeX += 1 - print('0') - + print("0") diff --git a/runtime/tools/fst/generate_lexicon.py b/runtime/tools/fst/generate_lexicon.py index a5a47ec6e..f8c744340 100755 --- a/runtime/tools/fst/generate_lexicon.py +++ b/runtime/tools/fst/generate_lexicon.py @@ -8,18 +8,19 @@ import sys # sys.argv[3]: lexicon file for corpus.dict lex_dict = {} -with open(sys.argv[2], 'r', encoding='utf8') as fin: +with open(sys.argv[2], "r", encoding="utf8") as fin: for line in fin: - words = line.strip().split('\t') + words = line.strip().split("\t") if len(words) != 2: continue lex_dict[words[0]] = words[1] -with open(sys.argv[1], 'r', encoding='utf8') as fin, \ - open(sys.argv[3], 'w', encoding='utf8') as fout: +with open(sys.argv[1], "r", encoding="utf8") as fin, open( + sys.argv[3], "w", encoding="utf8" +) as fout: for line in fin: word = line.strip() - if word == '' or word == '': + if word == "" or word == "": continue word_lex = "" if word in lex_dict: @@ -29,7 +30,6 @@ with open(sys.argv[1], 'r', encoding='utf8') as fin, \ if word[i] in lex_dict: word_lex += " " + lex_dict[word[i]] else: - word_lex += " " - - fout.write('{}\t{}\n'.format(word, word_lex.strip())) + word_lex += " " + fout.write("{}\t{}\n".format(word, word_lex.strip())) diff --git a/runtime/triton_gpu/client/decode_manifest_triton.py b/runtime/triton_gpu/client/decode_manifest_triton.py index 3a8d57fed..482c715e0 100644 --- a/runtime/triton_gpu/client/decode_manifest_triton.py +++ b/runtime/triton_gpu/client/decode_manifest_triton.py @@ -78,9 +78,7 @@ DEFAULT_MANIFEST_FILENAME = "/mnt/samsung-t7/yuekai/aishell-test-dev-manifests/d def get_args(): - parser = argparse.ArgumentParser( - formatter_class=argparse.ArgumentDefaultsHelpFormatter - ) + parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument( "--server-addr", @@ -225,9 +223,7 @@ async def send( lengths = np.array([[len(waveform)]], dtype=np.int32) inputs = [ - protocol_client.InferInput( - "WAV", samples.shape, np_to_triton_dtype(samples.dtype) - ), + protocol_client.InferInput("WAV", samples.shape, np_to_triton_dtype(samples.dtype)), protocol_client.InferInput( "WAV_LENS", lengths.shape, np_to_triton_dtype(lengths.dtype) ), @@ -360,9 +356,7 @@ async def send_streaming( decoding_results = b" ".join(decoding_results).decode("utf-8") else: # For wenet - decoding_results = response.as_numpy("TRANSCRIPTS")[0].decode( - "utf-8" - ) + decoding_results = response.as_numpy("TRANSCRIPTS")[0].decode("utf-8") chunk_end = time.time() - chunk_start latency_data.append((chunk_end, chunk_len / sample_rate)) @@ -406,15 +400,11 @@ async def main(): if args.streaming or args.simulate_streaming: frame_shift_ms = 10 frame_length_ms = 25 - add_frames = math.ceil( - (frame_length_ms - frame_shift_ms) / frame_shift_ms - ) + add_frames = math.ceil((frame_length_ms - frame_shift_ms) / frame_shift_ms) # decode_window_length: input sequence length of streaming encoder if args.context > 0: # decode window length calculation for wenet - decode_window_length = ( - args.chunk_size - 1 - ) * args.subsampling + args.context + decode_window_length = (args.chunk_size - 1) * args.subsampling + args.context else: # decode window length calculation for icefall decode_window_length = ( @@ -437,10 +427,7 @@ async def main(): compute_cer=compute_cer, model_name=args.model_name, first_chunk_in_secs=first_chunk_ms / 1000, - other_chunk_in_secs=args.chunk_size - * args.subsampling - * frame_shift_ms - / 1000, + other_chunk_in_secs=args.chunk_size * args.subsampling * frame_shift_ms / 1000, task_index=i, ) ) @@ -455,10 +442,7 @@ async def main(): compute_cer=compute_cer, model_name=args.model_name, first_chunk_in_secs=first_chunk_ms / 1000, - other_chunk_in_secs=args.chunk_size - * args.subsampling - * frame_shift_ms - / 1000, + other_chunk_in_secs=args.chunk_size * args.subsampling * frame_shift_ms / 1000, task_index=i, simulate_mode=True, ) @@ -496,15 +480,10 @@ async def main(): s = f"RTF: {rtf:.4f}\n" s += f"total_duration: {total_duration:.3f} seconds\n" s += f"({total_duration/3600:.2f} hours)\n" - s += ( - f"processing time: {elapsed:.3f} seconds " - f"({elapsed/3600:.2f} hours)\n" - ) + s += f"processing time: {elapsed:.3f} seconds " f"({elapsed/3600:.2f} hours)\n" if args.streaming or args.simulate_streaming: - latency_list = [ - chunk_end for (chunk_end, chunk_duration) in latency_data - ] + latency_list = [chunk_end for (chunk_end, chunk_duration) in latency_data] latency_ms = sum(latency_list) / float(len(latency_list)) * 1000.0 latency_variance = np.var(latency_list, dtype=np.float64) * 1000.0 s += f"latency_variance: {latency_variance:.2f}\n" @@ -530,9 +509,7 @@ async def main(): print(f.readline()) # Detailed errors if args.stats_file: - stats = await triton_client.get_inference_statistics( - model_name="", as_json=True - ) + stats = await triton_client.get_inference_statistics(model_name="", as_json=True) with open(args.stats_file, "w") as f: json.dump(stats, f) diff --git a/runtime/triton_gpu/client/decode_manifest_triton_wo_cuts.py b/runtime/triton_gpu/client/decode_manifest_triton_wo_cuts.py index ad121c60d..379b28df0 100644 --- a/runtime/triton_gpu/client/decode_manifest_triton_wo_cuts.py +++ b/runtime/triton_gpu/client/decode_manifest_triton_wo_cuts.py @@ -60,14 +60,12 @@ from tritonclient.utils import np_to_triton_dtype from icefall.utils import store_transcripts, write_error_stats DEFAULT_MANIFEST_FILENAME = "./aishell_test.txt" # noqa -DEFAULT_ROOT = './' -DEFAULT_ROOT = '/mfs/songtao/researchcode/FunASR/data/' +DEFAULT_ROOT = "./" +DEFAULT_ROOT = "/mfs/songtao/researchcode/FunASR/data/" def get_args(): - parser = argparse.ArgumentParser( - formatter_class=argparse.ArgumentDefaultsHelpFormatter - ) + parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument( "--server-addr", @@ -185,7 +183,7 @@ def load_manifest(fp): with open(fp) as f: for i, dp in enumerate(f.readlines()): dp = eval(dp) - dp['id'] = i + dp["id"] = i data.append(dp) return data @@ -195,19 +193,19 @@ def split_dps(dps, num_tasks): # import pdb;pdb.set_trace() assert len(dps) > num_tasks - one_task_num = len(dps)//num_tasks + one_task_num = len(dps) // num_tasks for i in range(0, len(dps), one_task_num): - if i+one_task_num >= len(dps): + if i + one_task_num >= len(dps): for k, j in enumerate(range(i, len(dps))): dps_splited[k].append(dps[j]) else: - dps_splited.append(dps[i:i+one_task_num]) + dps_splited.append(dps[i : i + one_task_num]) return dps_splited def load_audio(path): audio = AudioSegment.from_wav(path).set_frame_rate(16000).set_channels(1) - audiop_np = np.array(audio.get_array_of_samples())/32768.0 + audiop_np = np.array(audio.get_array_of_samples()) / 32768.0 return audiop_np.astype(np.float32), audio.duration_seconds @@ -227,16 +225,14 @@ async def send( if i % log_interval == 0: print(f"{name}: {i}/{len(dps)}") - waveform, duration = load_audio( - os.path.join(DEFAULT_ROOT, dp['audio_filepath'])) + waveform, duration = load_audio(os.path.join(DEFAULT_ROOT, dp["audio_filepath"])) sample_rate = 16000 # padding to nearset 10 seconds samples = np.zeros( ( 1, - 10 * sample_rate * - (int(len(waveform) / sample_rate // 10) + 1), + 10 * sample_rate * (int(len(waveform) / sample_rate // 10) + 1), ), dtype=np.float32, ) @@ -245,9 +241,7 @@ async def send( lengths = np.array([[len(waveform)]], dtype=np.int32) inputs = [ - protocol_client.InferInput( - "WAV", samples.shape, np_to_triton_dtype(samples.dtype) - ), + protocol_client.InferInput("WAV", samples.shape, np_to_triton_dtype(samples.dtype)), protocol_client.InferInput( "WAV_LENS", lengths.shape, np_to_triton_dtype(lengths.dtype) ), @@ -271,16 +265,16 @@ async def send( total_duration += duration if compute_cer: - ref = dp['text'].split() + ref = dp["text"].split() hyp = decoding_results.split() ref = list("".join(ref)) hyp = list("".join(hyp)) - results.append((dp['id'], ref, hyp)) + results.append((dp["id"], ref, hyp)) else: results.append( ( - dp['id'], - dp['text'].split(), + dp["id"], + dp["text"].split(), decoding_results.split(), ) ) # noqa @@ -309,7 +303,7 @@ async def send_streaming( if i % log_interval == 0: print(f"{name}: {i}/{len(dps)}") - waveform, duration = load_audio(dp['audio_filepath']) + waveform, duration = load_audio(dp["audio_filepath"]) sample_rate = 16000 wav_segs = [] @@ -318,10 +312,10 @@ async def send_streaming( while j < len(waveform): if j == 0: stride = int(first_chunk_in_secs * sample_rate) - wav_segs.append(waveform[j: j + stride]) + wav_segs.append(waveform[j : j + stride]) else: stride = int(other_chunk_in_secs * sample_rate) - wav_segs.append(waveform[j: j + stride]) + wav_segs.append(waveform[j : j + stride]) j += len(wav_segs[-1]) sequence_id = task_index + 10086 @@ -380,25 +374,23 @@ async def send_streaming( decoding_results = b" ".join(decoding_results).decode("utf-8") else: # For wenet - decoding_results = response.as_numpy("TRANSCRIPTS")[0].decode( - "utf-8" - ) + decoding_results = response.as_numpy("TRANSCRIPTS")[0].decode("utf-8") chunk_end = time.time() - chunk_start latency_data.append((chunk_end, chunk_len / sample_rate)) total_duration += duration if compute_cer: - ref = dp['text'].split() + ref = dp["text"].split() hyp = decoding_results.split() ref = list("".join(ref)) hyp = list("".join(hyp)) - results.append((dp['id'], ref, hyp)) + results.append((dp["id"], ref, hyp)) else: results.append( ( - dp['id'], - dp['text'].split(), + dp["id"], + dp["text"].split(), decoding_results.split(), ) ) # noqa @@ -426,15 +418,11 @@ async def main(): if args.streaming or args.simulate_streaming: frame_shift_ms = 10 frame_length_ms = 25 - add_frames = math.ceil( - (frame_length_ms - frame_shift_ms) / frame_shift_ms - ) + add_frames = math.ceil((frame_length_ms - frame_shift_ms) / frame_shift_ms) # decode_window_length: input sequence length of streaming encoder if args.context > 0: # decode window length calculation for wenet - decode_window_length = ( - args.chunk_size - 1 - ) * args.subsampling + args.context + decode_window_length = (args.chunk_size - 1) * args.subsampling + args.context else: # decode window length calculation for icefall decode_window_length = ( @@ -457,10 +445,7 @@ async def main(): compute_cer=compute_cer, model_name=args.model_name, first_chunk_in_secs=first_chunk_ms / 1000, - other_chunk_in_secs=args.chunk_size - * args.subsampling - * frame_shift_ms - / 1000, + other_chunk_in_secs=args.chunk_size * args.subsampling * frame_shift_ms / 1000, task_index=i, ) ) @@ -475,10 +460,7 @@ async def main(): compute_cer=compute_cer, model_name=args.model_name, first_chunk_in_secs=first_chunk_ms / 1000, - other_chunk_in_secs=args.chunk_size - * args.subsampling - * frame_shift_ms - / 1000, + other_chunk_in_secs=args.chunk_size * args.subsampling * frame_shift_ms / 1000, task_index=i, simulate_mode=True, ) @@ -516,15 +498,10 @@ async def main(): s = f"RTF: {rtf:.4f}\n" s += f"total_duration: {total_duration:.3f} seconds\n" s += f"({total_duration/3600:.2f} hours)\n" - s += ( - f"processing time: {elapsed:.3f} seconds " - f"({elapsed/3600:.2f} hours)\n" - ) + s += f"processing time: {elapsed:.3f} seconds " f"({elapsed/3600:.2f} hours)\n" if args.streaming or args.simulate_streaming: - latency_list = [ - chunk_end for (chunk_end, chunk_duration) in latency_data - ] + latency_list = [chunk_end for (chunk_end, chunk_duration) in latency_data] latency_ms = sum(latency_list) / float(len(latency_list)) * 1000.0 latency_variance = np.var(latency_list, dtype=np.float64) * 1000.0 s += f"latency_variance: {latency_variance:.2f}\n" @@ -550,9 +527,7 @@ async def main(): print(f.readline()) # Detailed errors if args.stats_file: - stats = await triton_client.get_inference_statistics( - model_name="", as_json=True - ) + stats = await triton_client.get_inference_statistics(model_name="", as_json=True) with open(args.stats_file, "w") as f: json.dump(stats, f) diff --git a/runtime/triton_gpu/client/speech_client.py b/runtime/triton_gpu/client/speech_client.py index 1bafbe246..15482a760 100644 --- a/runtime/triton_gpu/client/speech_client.py +++ b/runtime/triton_gpu/client/speech_client.py @@ -71,9 +71,7 @@ class StreamingSpeechClient(object): # the exact first chunk length frames # since the subsampling will look ahead several frames first_chunk_length = (chunk_size - 1) * subsampling + context - add_frames = math.ceil( - (frame_length_ms - frame_shift_ms) / frame_shift_ms - ) + add_frames = math.ceil((frame_length_ms - frame_shift_ms) / frame_shift_ms) first_chunk_ms = (first_chunk_length + add_frames) * frame_shift_ms other_chunk_ms = chunk_size * subsampling * frame_shift_ms self.first_chunk_in_secs = first_chunk_ms / 1000 diff --git a/runtime/triton_gpu/model_repo_paraformer_large_offline/feature_extractor/1/model.py b/runtime/triton_gpu/model_repo_paraformer_large_offline/feature_extractor/1/model.py index c556dafc5..e6eb3b935 100644 --- a/runtime/triton_gpu/model_repo_paraformer_large_offline/feature_extractor/1/model.py +++ b/runtime/triton_gpu/model_repo_paraformer_large_offline/feature_extractor/1/model.py @@ -25,8 +25,10 @@ import json import yaml from typing import Any, Dict, Iterable, List, NamedTuple, Set, Tuple, Union + class LFR(torch.nn.Module): - """Batch LFR: https://github.com/Mddct/devil-asr/blob/main/patch/lfr.py """ + """Batch LFR: https://github.com/Mddct/devil-asr/blob/main/patch/lfr.py""" + def __init__(self, m: int = 7, n: int = 6) -> None: """ Actually, this implements stacking frames and skipping frames. @@ -42,8 +44,9 @@ class LFR(torch.nn.Module): self.left_padding_nums = math.ceil((self.m - 1) // 2) - def forward(self, input_tensor: torch.Tensor, - input_lens: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + def forward( + self, input_tensor: torch.Tensor, input_lens: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: B, _, D = input_tensor.size() n_lfr = torch.ceil(input_lens / self.n) @@ -70,44 +73,43 @@ class LFR(torch.nn.Module): # stack input_tensor = torch.cat([head_frames, input_tensor, tail_frames], dim=1) - index = torch.arange(T_all_max, - device=input_tensor.device, - dtype=input_lens.dtype).unsqueeze(0).repeat(B, 1) # [B, T_all_max] - index_mask = (index < - (self.left_padding_nums + input_lens).unsqueeze(1) - ) #[B, T_all_max] + index = ( + torch.arange(T_all_max, device=input_tensor.device, dtype=input_lens.dtype) + .unsqueeze(0) + .repeat(B, 1) + ) # [B, T_all_max] + index_mask = index < (self.left_padding_nums + input_lens).unsqueeze(1) # [B, T_all_max] - tail_index_mask = torch.logical_not( - index >= (T_all.unsqueeze(1))) & index_mask - tail = torch.ones(T_all_max, - dtype=input_lens.dtype, - device=input_tensor.device).unsqueeze(0).repeat(B, 1) * ( - T_all_max - 1) # [B, T_all_max] - indices = torch.where(torch.logical_or(index_mask, tail_index_mask), - index, tail) + tail_index_mask = torch.logical_not(index >= (T_all.unsqueeze(1))) & index_mask + tail = torch.ones(T_all_max, dtype=input_lens.dtype, device=input_tensor.device).unsqueeze( + 0 + ).repeat(B, 1) * ( + T_all_max - 1 + ) # [B, T_all_max] + indices = torch.where(torch.logical_or(index_mask, tail_index_mask), index, tail) input_tensor = torch.gather(input_tensor, 1, indices.unsqueeze(2).repeat(1, 1, D)) input_tensor = input_tensor.unfold(1, self.m, step=self.n).transpose(2, 3) return input_tensor.reshape(B, -1, D * self.m), new_len -class WavFrontend(): - """Conventional frontend structure for ASR. - """ + +class WavFrontend: + """Conventional frontend structure for ASR.""" def __init__( - self, - cmvn_file: str = None, - fs: int = 16000, - window: str = 'hamming', - n_mels: int = 80, - frame_length: int = 25, - frame_shift: int = 10, - filter_length_min: int = -1, - filter_length_max: float = -1, - lfr_m: int = 7, - lfr_n: int = 6, - dither: float = 1.0 + self, + cmvn_file: str = None, + fs: int = 16000, + window: str = "hamming", + n_mels: int = 80, + frame_length: int = 25, + frame_shift: int = 10, + filter_length_min: int = -1, + filter_length_max: float = -1, + lfr_m: int = 7, + lfr_n: int = 6, + dither: float = 1.0, ) -> None: self.fs = fs @@ -133,31 +135,33 @@ class WavFrontend(): batch, frame, dim = inputs.shape means = np.tile(self.cmvn[0:1, :dim], (frame, 1)) vars = np.tile(self.cmvn[1:2, :dim], (frame, 1)) - + means = torch.from_numpy(means).to(inputs.device) vars = torch.from_numpy(vars).to(inputs.device) # print(inputs.shape, means.shape, vars.shape) inputs = (inputs + means) * vars return inputs - def load_cmvn(self,) -> np.ndarray: - with open(self.cmvn_file, 'r', encoding='utf-8') as f: + def load_cmvn( + self, + ) -> np.ndarray: + with open(self.cmvn_file, "r", encoding="utf-8") as f: lines = f.readlines() means_list = [] vars_list = [] for i in range(len(lines)): line_item = lines[i].split() - if line_item[0] == '': + if line_item[0] == "": line_item = lines[i + 1].split() - if line_item[0] == '': - add_shift_line = line_item[3:(len(line_item) - 1)] + if line_item[0] == "": + add_shift_line = line_item[3 : (len(line_item) - 1)] means_list = list(add_shift_line) continue - elif line_item[0] == '': + elif line_item[0] == "": line_item = lines[i + 1].split() - if line_item[0] == '': - rescale_line = line_item[3:(len(line_item) - 1)] + if line_item[0] == "": + rescale_line = line_item[3 : (len(line_item) - 1)] vars_list = list(rescale_line) continue @@ -197,16 +201,14 @@ class TritonPythonModel: * model_version: Model version * model_name: Model name """ - self.model_config = model_config = json.loads(args['model_config']) + self.model_config = model_config = json.loads(args["model_config"]) self.max_batch_size = max(model_config["max_batch_size"], 1) self.device = "cuda" # Get OUTPUT0 configuration - output0_config = pb_utils.get_output_config_by_name( - model_config, "speech") + output0_config = pb_utils.get_output_config_by_name(model_config, "speech") # Convert Triton types to numpy types - output0_dtype = pb_utils.triton_string_to_numpy( - output0_config['data_type']) + output0_dtype = pb_utils.triton_string_to_numpy(output0_config["data_type"]) if output0_dtype == np.float32: self.output0_dtype = torch.float32 @@ -214,42 +216,36 @@ class TritonPythonModel: self.output0_dtype = torch.float16 # Get OUTPUT1 configuration - output1_config = pb_utils.get_output_config_by_name( - model_config, "speech_lengths") + output1_config = pb_utils.get_output_config_by_name(model_config, "speech_lengths") # Convert Triton types to numpy types - self.output1_dtype = pb_utils.triton_string_to_numpy( - output1_config['data_type']) + self.output1_dtype = pb_utils.triton_string_to_numpy(output1_config["data_type"]) - params = self.model_config['parameters'] + params = self.model_config["parameters"] for li in params.items(): key, value = li value = value["string_value"] if key == "config_path": - with open(str(value), 'rb') as f: + with open(str(value), "rb") as f: config = yaml.load(f, Loader=yaml.Loader) if key == "cmvn_path": cmvn_path = str(value) opts = kaldifeat.FbankOptions() - opts.frame_opts.dither = 1.0 # TODO: 0.0 or 1.0 - opts.frame_opts.window_type = config['frontend_conf']['window'] - opts.mel_opts.num_bins = int(config['frontend_conf']['n_mels']) - opts.frame_opts.frame_shift_ms = float(config['frontend_conf']['frame_shift']) - opts.frame_opts.frame_length_ms = float(config['frontend_conf']['frame_length']) - opts.frame_opts.samp_freq = int(config['frontend_conf']['fs']) + opts.frame_opts.dither = 1.0 # TODO: 0.0 or 1.0 + opts.frame_opts.window_type = config["frontend_conf"]["window"] + opts.mel_opts.num_bins = int(config["frontend_conf"]["n_mels"]) + opts.frame_opts.frame_shift_ms = float(config["frontend_conf"]["frame_shift"]) + opts.frame_opts.frame_length_ms = float(config["frontend_conf"]["frame_length"]) + opts.frame_opts.samp_freq = int(config["frontend_conf"]["fs"]) opts.device = torch.device(self.device) self.opts = opts self.feature_extractor = Fbank(self.opts) self.feature_size = opts.mel_opts.num_bins - self.frontend = WavFrontend( - cmvn_file=cmvn_path, - **config['frontend_conf']) + self.frontend = WavFrontend(cmvn_file=cmvn_path, **config["frontend_conf"]) - def extract_feat(self, - waveform_list: List[np.ndarray] - ) -> Tuple[np.ndarray, np.ndarray]: + def extract_feat(self, waveform_list: List[np.ndarray]) -> Tuple[np.ndarray, np.ndarray]: feats, feats_len = [], [] wavs = [] for waveform in waveform_list: @@ -258,18 +254,21 @@ class TritonPythonModel: features = self.feature_extractor(wavs) features_len = [feature.shape[0] for feature in features] - speech = torch.zeros((len(features), max(features_len), self.opts.mel_opts.num_bins), - dtype=self.output0_dtype, device=self.device) + speech = torch.zeros( + (len(features), max(features_len), self.opts.mel_opts.num_bins), + dtype=self.output0_dtype, + device=self.device, + ) for i, feature in enumerate(features): - speech[i,:int(features_len[i])] = feature - speech_lens = torch.tensor(features_len,dtype=torch.int64).to(self.device) - + speech[i, : int(features_len[i])] = feature + speech_lens = torch.tensor(features_len, dtype=torch.int64).to(self.device) + feats, feats_len = self.frontend.lfr(speech, speech_lens) feats_len = feats_len.type(torch.int32) - + feats = self.frontend.apply_cmvn_batch(feats) feats = feats.type(self.output0_dtype) - + return feats, feats_len def execute(self, requests): @@ -294,23 +293,22 @@ class TritonPythonModel: batch_len = [] responses = [] for request in requests: - + input0 = pb_utils.get_input_tensor_by_name(request, "wav") input1 = pb_utils.get_input_tensor_by_name(request, "wav_lens") - cur_b_wav = input0.as_numpy() * (1 << 15) # b x -1 + cur_b_wav = input0.as_numpy() * (1 << 15) # b x -1 total_waves.append(cur_b_wav) features, feats_len = self.extract_feat(total_waves) for i in range(features.shape[0]): - speech = features[i:i+1][:int(feats_len[i].cpu())] + speech = features[i : i + 1][: int(feats_len[i].cpu())] speech_lengths = feats_len[i].unsqueeze(0).unsqueeze(0) speech, speech_lengths = speech.cpu(), speech_lengths.cpu() out0 = pb_utils.Tensor.from_dlpack("speech", to_dlpack(speech)) - out1 = pb_utils.Tensor.from_dlpack("speech_lengths", - to_dlpack(speech_lengths)) + out1 = pb_utils.Tensor.from_dlpack("speech_lengths", to_dlpack(speech_lengths)) inference_response = pb_utils.InferenceResponse(output_tensors=[out0, out1]) responses.append(inference_response) return responses diff --git a/runtime/triton_gpu/model_repo_paraformer_large_offline/scoring/1/model.py b/runtime/triton_gpu/model_repo_paraformer_large_offline/scoring/1/model.py index ef6278dc8..d582e50ff 100644 --- a/runtime/triton_gpu/model_repo_paraformer_large_offline/scoring/1/model.py +++ b/runtime/triton_gpu/model_repo_paraformer_large_offline/scoring/1/model.py @@ -23,6 +23,7 @@ import json import os import yaml + class TritonPythonModel: """Your Python model must use the same class name. Every Python model that is created must have "TritonPythonModel" as the class name. @@ -44,20 +45,18 @@ class TritonPythonModel: * model_version: Model version * model_name: Model name """ - self.model_config = model_config = json.loads(args['model_config']) + self.model_config = model_config = json.loads(args["model_config"]) self.max_batch_size = max(model_config["max_batch_size"], 1) # # Get OUTPUT0 configuration - output0_config = pb_utils.get_output_config_by_name( - model_config, "OUTPUT0") + output0_config = pb_utils.get_output_config_by_name(model_config, "OUTPUT0") # # Convert Triton types to numpy types - self.out0_dtype = pb_utils.triton_string_to_numpy( - output0_config['data_type']) - - self.init_vocab(self.model_config['parameters']) + self.out0_dtype = pb_utils.triton_string_to_numpy(output0_config["data_type"]) + + self.init_vocab(self.model_config["parameters"]) def init_vocab(self, parameters): - blank_id=0 + blank_id = 0 for li in parameters.items(): key, value = li value = value["string_value"] @@ -67,16 +66,16 @@ class TritonPythonModel: lm_path = value elif key == "vocabulary": self.vocab_dict = self.load_vocab(value) - if key == 'ignore_id': + if key == "ignore_id": ignore_id = int(value) def load_vocab(self, vocab_file): """ load lang_char.txt """ - with open(str(vocab_file), 'rb') as f: + with open(str(vocab_file), "rb") as f: config = yaml.load(f, Loader=yaml.Loader) - return config['token_list'] + return config["token_list"] def execute(self, requests): """`execute` must be implemented in every Python model. `execute` @@ -110,31 +109,40 @@ class TritonPythonModel: # Perform inference on the request and append it to responses list... in_0 = pb_utils.get_input_tensor_by_name(request, "logits") in_1 = pb_utils.get_input_tensor_by_name(request, "token_num") - + logits, token_num = from_dlpack(in_0.to_dlpack()), from_dlpack(in_1.to_dlpack()).cpu() max_token_num = max(max_token_num, token_num) assert logits.shape[0] == 1 logits_list.append(logits) token_num_list.append(token_num) - total_seq +=1 + total_seq += 1 - logits_batch = torch.zeros(len(logits_list), max_token_num, len(self.vocab_dict), dtype=torch.float32, device=logits.device) + logits_batch = torch.zeros( + len(logits_list), + max_token_num, + len(self.vocab_dict), + dtype=torch.float32, + device=logits.device, + ) token_num_batch = torch.zeros(len(logits_list)) for i, (logits, token_num) in enumerate(zip(logits_list, token_num_list)): - logits_batch[i][:int(token_num)] = logits[0][:int(token_num)] + logits_batch[i][: int(token_num)] = logits[0][: int(token_num)] token_num_batch[i] = token_num yseq_batch = logits_batch.argmax(axis=-1).tolist() token_int_batch = [list(filter(lambda x: x not in (0, 2), yseq)) for yseq in yseq_batch] tokens_batch = [[self.vocab_dict[i] for i in token_int] for token_int in token_int_batch] - - hyps = [u"".join([t if t != "" else " " for t in tokens]).encode('utf-8') for tokens in tokens_batch] + + hyps = [ + "".join([t if t != "" else " " for t in tokens]).encode("utf-8") + for tokens in tokens_batch + ] responses = [] for i in range(total_seq): - sents = np.array(hyps[i:i+1]) + sents = np.array(hyps[i : i + 1]) out0 = pb_utils.Tensor("OUTPUT0", sents.astype(self.out0_dtype)) inference_response = pb_utils.InferenceResponse(output_tensors=[out0]) responses.append(inference_response) @@ -146,4 +154,4 @@ class TritonPythonModel: Implementing `finalize` function is optional. This function allows the model to perform any necessary clean ups before exit. """ - print('Cleaning up...') + print("Cleaning up...") diff --git a/runtime/triton_gpu/model_repo_paraformer_large_online/cif_search/1/model.py b/runtime/triton_gpu/model_repo_paraformer_large_online/cif_search/1/model.py index 96ad8214a..aeecd3301 100755 --- a/runtime/triton_gpu/model_repo_paraformer_large_online/cif_search/1/model.py +++ b/runtime/triton_gpu/model_repo_paraformer_large_online/cif_search/1/model.py @@ -23,10 +23,14 @@ class LimitedDict(OrderedDict): class CIFSearch: """CIFSearch: https://github.com/alibaba-damo-academy/FunASR/blob/main/runtime/python/onnxruntime/funasr_onnx - /paraformer_online_bin.py """ + /paraformer_online_bin.py""" + def __init__(self): - self.cache = {"cif_hidden": np.zeros((1, 1, 512)).astype(np.float32), - "cif_alphas": np.zeros((1, 1)).astype(np.float32), "last_chunk": False} + self.cache = { + "cif_hidden": np.zeros((1, 1, 512)).astype(np.float32), + "cif_alphas": np.zeros((1, 1)).astype(np.float32), + "last_chunk": False, + } self.chunk_size = [5, 10, 5] self.tail_threshold = 0.45 self.cif_threshold = 1.0 @@ -38,8 +42,8 @@ class CIFSearch: list_frames = [] cache_alphas = [] cache_hiddens = [] - alphas[:, :self.chunk_size[0]] = 0.0 - alphas[:, sum(self.chunk_size[:2]):] = 0.0 + alphas[:, : self.chunk_size[0]] = 0.0 + alphas[:, sum(self.chunk_size[:2]) :] = 0.0 if self.cache is not None and "cif_alphas" in self.cache and "cif_hidden" in self.cache: hidden = np.concatenate((self.cache["cif_hidden"], hidden), axis=1) @@ -95,7 +99,9 @@ class CIFSearch: self.cache["cif_hidden"] = np.stack(cache_hiddens, axis=0) self.cache["cif_hidden"] = np.expand_dims(self.cache["cif_hidden"], axis=0) - return np.stack(list_ls, axis=0).astype(np.float32), np.stack(token_length, axis=0).astype(np.int32) + return np.stack(list_ls, axis=0).astype(np.float32), np.stack(token_length, axis=0).astype( + np.int32 + ) class TritonPythonModel: @@ -119,18 +125,16 @@ class TritonPythonModel: * model_version: Model version * model_name: Model name """ - self.model_config = model_config = json.loads(args['model_config']) + self.model_config = model_config = json.loads(args["model_config"]) self.max_batch_size = max(model_config["max_batch_size"], 1) # # Get OUTPUT0 configuration - output0_config = pb_utils.get_output_config_by_name( - model_config, "transcripts") + output0_config = pb_utils.get_output_config_by_name(model_config, "transcripts") # # Convert Triton types to numpy types - self.out0_dtype = pb_utils.triton_string_to_numpy( - output0_config['data_type']) + self.out0_dtype = pb_utils.triton_string_to_numpy(output0_config["data_type"]) + + self.init_vocab(self.model_config["parameters"]) - self.init_vocab(self.model_config['parameters']) - self.cif_search_cache = LimitedDict(1024) self.start = LimitedDict(1024) @@ -142,9 +146,9 @@ class TritonPythonModel: self.vocab_dict = self.load_vocab(value) def load_vocab(self, vocab_file): - with open(str(vocab_file), 'rb') as f: + with open(str(vocab_file), "rb") as f: config = yaml.load(f, Loader=yaml.Loader) - return config['token_list'] + return config["token_list"] async def execute(self, requests): """`execute` must be implemented in every Python model. `execute` @@ -187,7 +191,7 @@ class TritonPythonModel: in_start = pb_utils.get_input_tensor_by_name(request, "START") start = in_start.as_numpy()[0][0] - + in_corrid = pb_utils.get_input_tensor_by_name(request, "CORRID") corrid = in_corrid.as_numpy()[0][0] @@ -202,19 +206,21 @@ class TritonPythonModel: self.start[corrid] = 1 if end: self.cif_search_cache[corrid].cache["last_chunk"] = True - + acoustic, acoustic_len = self.cif_search_cache[corrid].infer(hidden, alphas) - batch_result[corrid] = '' + batch_result[corrid] = "" if acoustic.shape[1] == 0: - continue + continue else: qualified_corrid.append(corrid) input_tensor0 = pb_utils.Tensor("enc", hidden) input_tensor1 = pb_utils.Tensor("enc_len", np.array([hidden_len], dtype=np.int32)) input_tensor2 = pb_utils.Tensor("acoustic_embeds", acoustic) - input_tensor3 = pb_utils.Tensor("acoustic_embeds_len", np.array([acoustic_len], dtype=np.int32)) + input_tensor3 = pb_utils.Tensor( + "acoustic_embeds_len", np.array([acoustic_len], dtype=np.int32) + ) input_tensors = [input_tensor0, input_tensor1, input_tensor2, input_tensor3] - + if self.start[corrid] and end: flag = 3 elif end: @@ -225,12 +231,12 @@ class TritonPythonModel: else: flag = 0 inference_request = pb_utils.InferenceRequest( - model_name='decoder', - requested_output_names=['sample_ids'], + model_name="decoder", + requested_output_names=["sample_ids"], inputs=input_tensors, - request_id='', + request_id="", correlation_id=corrid, - flags=flag + flags=flag, ) inference_response_awaits.append(inference_request.async_exec()) @@ -240,9 +246,9 @@ class TritonPythonModel: if inference_response.has_error(): raise pb_utils.TritonModelException(inference_response.error().message()) else: - sample_ids = pb_utils.get_output_tensor_by_name(inference_response, 'sample_ids') + sample_ids = pb_utils.get_output_tensor_by_name(inference_response, "sample_ids") token_ids = from_dlpack(sample_ids.to_dlpack()).cpu().numpy()[0] - + # Change integer-ids to tokens tokens = [self.vocab_dict[token_id] for token_id in token_ids] batch_result[index_corrid] = "".join(tokens) @@ -252,7 +258,7 @@ class TritonPythonModel: out0 = pb_utils.Tensor("transcripts", sent.astype(self.out0_dtype)) inference_response = pb_utils.InferenceResponse(output_tensors=[out0]) responses.append(inference_response) - + if batch_end[i]: del self.cif_search_cache[index_corrid] del self.start[index_corrid] @@ -264,5 +270,4 @@ class TritonPythonModel: Implementing `finalize` function is optional. This function allows the model to perform any necessary clean ups before exit. """ - print('Cleaning up...') - + print("Cleaning up...") diff --git a/runtime/triton_gpu/model_repo_paraformer_large_online/feature_extractor/1/model.py b/runtime/triton_gpu/model_repo_paraformer_large_online/feature_extractor/1/model.py index 7d81a986f..716a07546 100755 --- a/runtime/triton_gpu/model_repo_paraformer_large_online/feature_extractor/1/model.py +++ b/runtime/triton_gpu/model_repo_paraformer_large_online/feature_extractor/1/model.py @@ -36,7 +36,7 @@ class Fbank(torch.nn.Module): class Feat(object): - def __init__(self, seqid, offset_ms, sample_rate, frame_stride, device='cpu'): + def __init__(self, seqid, offset_ms, sample_rate, frame_stride, device="cpu"): self.seqid = seqid self.sample_rate = sample_rate self.wav = torch.tensor([], device=device) @@ -45,14 +45,14 @@ class Feat(object): self.frame_stride = int(frame_stride) self.device = device self.lfr_m = 7 - + def add_wavs(self, wav: torch.tensor): wav = wav.to(self.device) self.wav = torch.cat((self.wav, wav), axis=0) - + def get_seg_wav(self): seg = self.wav[:] - self.wav = self.wav[-self.offset:] + self.wav = self.wav[-self.offset :] return seg def add_frames(self, frames: torch.tensor): @@ -60,14 +60,13 @@ class Feat(object): frames: seq_len x feat_sz """ if self.frames is None: - self.frames = torch.cat((frames[0, :].repeat((self.lfr_m - 1) // 2, 1), - frames), axis=0) + self.frames = torch.cat((frames[0, :].repeat((self.lfr_m - 1) // 2, 1), frames), axis=0) else: self.frames = torch.cat([self.frames, frames], axis=0) - + def get_frames(self, num_frames: int): - seg = self.frames[0: num_frames] - self.frames = self.frames[self.frame_stride:] + seg = self.frames[0:num_frames] + self.frames = self.frames[self.frame_stride :] return seg @@ -91,7 +90,7 @@ class TritonPythonModel: * model_version: Model version * model_name: Model name """ - self.model_config = model_config = json.loads(args['model_config']) + self.model_config = model_config = json.loads(args["model_config"]) self.max_batch_size = max(model_config["max_batch_size"], 1) if "GPU" in model_config["instance_group"][0]["kind"]: @@ -100,35 +99,33 @@ class TritonPythonModel: self.device = "cpu" # Get OUTPUT0 configuration - output0_config = pb_utils.get_output_config_by_name( - model_config, "speech") + output0_config = pb_utils.get_output_config_by_name(model_config, "speech") # Convert Triton types to numpy types - self.output0_dtype = pb_utils.triton_string_to_numpy( - output0_config['data_type']) + self.output0_dtype = pb_utils.triton_string_to_numpy(output0_config["data_type"]) if self.output0_dtype == np.float32: self.dtype = torch.float32 else: self.dtype = torch.float16 - self.feature_size = output0_config['dims'][-1] - self.decoding_window = output0_config['dims'][-2] + self.feature_size = output0_config["dims"][-1] + self.decoding_window = output0_config["dims"][-2] - params = self.model_config['parameters'] + params = self.model_config["parameters"] for li in params.items(): key, value = li value = value["string_value"] if key == "config_path": - with open(str(value), 'rb') as f: + with open(str(value), "rb") as f: config = yaml.load(f, Loader=yaml.Loader) opts = kaldifeat.FbankOptions() opts.frame_opts.dither = 0.0 - opts.frame_opts.window_type = config['frontend_conf']['window'] - opts.mel_opts.num_bins = int(config['frontend_conf']['n_mels']) - opts.frame_opts.frame_shift_ms = float(config['frontend_conf']['frame_shift']) - opts.frame_opts.frame_length_ms = float(config['frontend_conf']['frame_length']) - opts.frame_opts.samp_freq = int(config['frontend_conf']['fs']) + opts.frame_opts.window_type = config["frontend_conf"]["window"] + opts.mel_opts.num_bins = int(config["frontend_conf"]["n_mels"]) + opts.frame_opts.frame_shift_ms = float(config["frontend_conf"]["frame_shift"]) + opts.frame_opts.frame_length_ms = float(config["frontend_conf"]["frame_length"]) + opts.frame_opts.samp_freq = int(config["frontend_conf"]["fs"]) opts.device = torch.device(self.device) self.opts = opts self.feature_extractor = Fbank(self.opts) @@ -177,8 +174,7 @@ class TritonPythonModel: # wav_len = from_dlpack(input1.to_dlpack())[0] wav_len = len(wav) if wav_len < self.chunk_size: - temp = torch.zeros(self.chunk_size, dtype=torch.float32, - device=self.device) + temp = torch.zeros(self.chunk_size, dtype=torch.float32, device=self.device) temp[0:wav_len] = wav[:] wav = temp @@ -192,10 +188,9 @@ class TritonPythonModel: end = in_end.as_numpy()[0][0] if start: - self.seq_feat[corrid] = Feat(corrid, self.offset_ms, - self.sample_rate, - self.frame_stride, - self.device) + self.seq_feat[corrid] = Feat( + corrid, self.offset_ms, self.sample_rate, self.frame_stride, self.device + ) if ready: self.seq_feat[corrid].add_wavs(wav) diff --git a/runtime/triton_gpu/model_repo_paraformer_large_online/lfr_cmvn_pe/export_lfr_cmvn_pe_onnx.py b/runtime/triton_gpu/model_repo_paraformer_large_online/lfr_cmvn_pe/export_lfr_cmvn_pe_onnx.py index c3328fd69..92a76d808 100755 --- a/runtime/triton_gpu/model_repo_paraformer_large_online/lfr_cmvn_pe/export_lfr_cmvn_pe_onnx.py +++ b/runtime/triton_gpu/model_repo_paraformer_large_online/lfr_cmvn_pe/export_lfr_cmvn_pe_onnx.py @@ -8,14 +8,16 @@ import torch.nn.functional as F class LFR_CMVN_PE(torch.nn.Module): - def __init__(self, - mean: torch.Tensor, - istd: torch.Tensor, - m: int = 7, - n: int = 6, - max_len: int = 5000, - encoder_input_size: int = 560, - encoder_output_size: int = 512): + def __init__( + self, + mean: torch.Tensor, + istd: torch.Tensor, + m: int = 7, + n: int = 6, + max_len: int = 5000, + encoder_input_size: int = 560, + encoder_output_size: int = 512, + ): super().__init__() # LRF @@ -34,12 +36,14 @@ class LFR_CMVN_PE(torch.nn.Module): self.encoder_output_size = encoder_output_size self.max_len = max_len self.pe = torch.zeros(self.max_len, self.encoder_input_size) - position = torch.arange(0, self.max_len, - dtype=torch.float32).unsqueeze(1) + position = torch.arange(0, self.max_len, dtype=torch.float32).unsqueeze(1) div_term = torch.exp( - torch.arange((self.encoder_input_size/2), dtype=torch.float32) * - -(math.log(10000.0) / (self.encoder_input_size/2-1))) - self.pe[:, 0::1] = torch.cat((torch.sin(position * div_term), torch.cos(position * div_term)), dim=1) + torch.arange((self.encoder_input_size / 2), dtype=torch.float32) + * -(math.log(10000.0) / (self.encoder_input_size / 2 - 1)) + ) + self.pe[:, 0::1] = torch.cat( + (torch.sin(position * div_term), torch.cos(position * div_term)), dim=1 + ) def forward(self, x, cache, offset): """ @@ -54,9 +58,9 @@ class LFR_CMVN_PE(torch.nn.Module): x = x.view(B, -1, D * self.m) x = (x + self.mean) * self.istd - x = x * (self.encoder_output_size ** 0.5) + x = x * (self.encoder_output_size**0.5) - index = offset + torch.arange(1, x.size(1)+1).to(dtype=torch.int32) + index = offset + torch.arange(1, x.size(1) + 1).to(dtype=torch.int32) pos_emb = F.embedding(index, self.pe) # B X T X d_model r_cache = x + pos_emb @@ -68,23 +72,23 @@ class LFR_CMVN_PE(torch.nn.Module): def load_cmvn(cmvn_file): - with open(cmvn_file, 'r', encoding='utf-8') as f: + with open(cmvn_file, "r", encoding="utf-8") as f: lines = f.readlines() means_list = [] vars_list = [] for i in range(len(lines)): line_item = lines[i].split() - if line_item[0] == '': + if line_item[0] == "": line_item = lines[i + 1].split() - if line_item[0] == '': - add_shift_line = line_item[3:(len(line_item) - 1)] + if line_item[0] == "": + add_shift_line = line_item[3 : (len(line_item) - 1)] means_list = list(add_shift_line) continue - elif line_item[0] == '': + elif line_item[0] == "": line_item = lines[i + 1].split() - if line_item[0] == '': - rescale_line = line_item[3:(len(line_item) - 1)] + if line_item[0] == "": + rescale_line = line_item[3 : (len(line_item) - 1)] vars_list = list(rescale_line) continue @@ -94,6 +98,7 @@ def load_cmvn(cmvn_file): vars = torch.from_numpy(vars) return means, vars + if __name__ == "__main__": means, vars = load_cmvn("am.mvn") means = torch.tile(means, (10, 1)) @@ -102,30 +107,36 @@ if __name__ == "__main__": model = LFR_CMVN_PE(means, vars) model.eval() - all_names = ['chunk_xs', 'cache', 'offset', 'chunk_xs_out', 'chunk_xs_out_len', 'r_cache', 'r_offset'] + all_names = [ + "chunk_xs", + "cache", + "offset", + "chunk_xs_out", + "chunk_xs_out_len", + "r_cache", + "r_offset", + ] dynamic_axes = {} for name in all_names: - dynamic_axes[name] = {0: 'B'} + dynamic_axes[name] = {0: "B"} input_data1 = torch.randn(4, 61, 80).to(torch.float32) input_data2 = torch.randn(4, 10, 560).to(torch.float32) input_data3 = torch.randn(4, 1).to(torch.int32) onnx_path = "./1/lfr_cmvn_pe.onnx" - torch.onnx.export(model, - (input_data1, input_data2, input_data3), - onnx_path, - export_params=True, - opset_version=11, - do_constant_folding=True, - input_names=['chunk_xs', 'cache', 'offset'], - output_names=['chunk_xs_out', 'chunk_xs_out_len', 'r_cache', 'r_offset'], - dynamic_axes=dynamic_axes, - verbose=False - ) + torch.onnx.export( + model, + (input_data1, input_data2, input_data3), + onnx_path, + export_params=True, + opset_version=11, + do_constant_folding=True, + input_names=["chunk_xs", "cache", "offset"], + output_names=["chunk_xs_out", "chunk_xs_out_len", "r_cache", "r_offset"], + dynamic_axes=dynamic_axes, + verbose=False, + ) print("export to onnx model succeed!") - - - diff --git a/setup.py b/setup.py index f9abf6498..dffe59482 100644 --- a/setup.py +++ b/setup.py @@ -18,7 +18,7 @@ requirements = { "kaldiio>=2.17.0", "torch_complex", # "nltk>=3.4.5", - "sentencepiece", # train + "sentencepiece", # train "jieba", "rotary_embedding_torch", # "ffmpeg-python", @@ -53,7 +53,7 @@ requirements = { "torch_optimizer", "fairscale", "transformers", - "openai-whisper" + "openai-whisper", ], "setup": [ "numpy", @@ -81,9 +81,9 @@ requirements = { "recommonmark>=0.4.0", "nbsphinx>=0.4.2", "sphinx-markdown-tables>=0.0.12", - "configargparse>=1.2.1" + "configargparse>=1.2.1", ], - "llm":[ + "llm": [ "transformers>=4.32.0", "accelerate", "tiktoken", @@ -102,9 +102,7 @@ requirements["test"].extend(requirements["train"]) install_requires = requirements["install"] setup_requires = requirements["setup"] tests_require = requirements["test"] -extras_require = { - k: v for k, v in requirements.items() if k not in ["install", "setup"] -} +extras_require = {k: v for k, v in requirements.items() if k not in ["install", "setup"]} dirname = os.path.dirname(__file__) version_file = os.path.join(dirname, "funasr", "version.txt") @@ -139,13 +137,15 @@ setup( "License :: OSI Approved :: Apache Software License", "Topic :: Software Development :: Libraries :: Python Modules", ], - entry_points={"console_scripts": [ - "funasr = funasr.bin.inference:main_hydra", - "funasr-train = funasr.bin.train:main_hydra", - "funasr-export = funasr.bin.export:main_hydra", - "scp2jsonl = funasr.datasets.audio_datasets.scp2jsonl:main_hydra", - "jsonl2scp = funasr.datasets.audio_datasets.jsonl2scp:main_hydra", - "funasr-scp2jsonl = funasr.datasets.audio_datasets.scp2jsonl:main_hydra", - "funasr-jsonl2scp = funasr.datasets.audio_datasets.jsonl2scp:main_hydra", - ]}, + entry_points={ + "console_scripts": [ + "funasr = funasr.bin.inference:main_hydra", + "funasr-train = funasr.bin.train:main_hydra", + "funasr-export = funasr.bin.export:main_hydra", + "scp2jsonl = funasr.datasets.audio_datasets.scp2jsonl:main_hydra", + "jsonl2scp = funasr.datasets.audio_datasets.jsonl2scp:main_hydra", + "funasr-scp2jsonl = funasr.datasets.audio_datasets.scp2jsonl:main_hydra", + "funasr-jsonl2scp = funasr.datasets.audio_datasets.jsonl2scp:main_hydra", + ] + }, ) diff --git a/tests/run_test.py b/tests/run_test.py index 4aa075e55..9ff648bc1 100644 --- a/tests/run_test.py +++ b/tests/run_test.py @@ -19,7 +19,7 @@ def gather_test_cases(test_dir, pattern, list_tests): for case in case_list: test_case = unittest.defaultTestLoader.discover(start_dir=test_dir, pattern=case) test_suite.addTest(test_case) - if hasattr(test_case, '__iter__'): + if hasattr(test_case, "__iter__"): for subcase in test_case: if list_tests: print(subcase) @@ -40,12 +40,12 @@ def main(args): sys.exit(len(result.errors)) -if __name__ == '__main__': - parser = argparse.ArgumentParser('test runner') - parser.add_argument('--list_tests', action='store_true', help='list all tests') - parser.add_argument('--pattern', default='test_*.py', help='test file pattern') - parser.add_argument('--test_dir', default='tests', help='directory to be tested') - parser.add_argument('--disable_profile', action='store_true', help='disable profiling') +if __name__ == "__main__": + parser = argparse.ArgumentParser("test runner") + parser.add_argument("--list_tests", action="store_true", help="list all tests") + parser.add_argument("--pattern", default="test_*.py", help="test file pattern") + parser.add_argument("--test_dir", default="tests", help="directory to be tested") + parser.add_argument("--disable_profile", action="store_true", help="disable profiling") args = parser.parse_args() - print(f'working dir: {os.getcwd()}') - main(args) \ No newline at end of file + print(f"working dir: {os.getcwd()}") + main(args) diff --git a/tests/test_asr_inference_pipeline.py b/tests/test_asr_inference_pipeline.py index f68f29b62..ba436bbde 100644 --- a/tests/test_asr_inference_pipeline.py +++ b/tests/test_asr_inference_pipeline.py @@ -11,22 +11,27 @@ class TestConformerInferencePipelines(unittest.TestCase): def test_funasr_path(self): import funasr import os + logger.info("run_dir:{0} ; funasr_path: {1}".format(os.getcwd(), funasr.__file__)) def test_aishell1(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_conformer_asr_nat-zh-cn-16k-aishell1-vocab4234-pytorch') + model="damo/speech_conformer_asr_nat-zh-cn-16k-aishell1-vocab4234-pytorch", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav') + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav" + ) logger.info("asr inference result: {0}".format(rec_result)) def test_aishell2(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_conformer_asr_nat-zh-cn-16k-aishell2-vocab5212-pytorch') + model="damo/speech_conformer_asr_nat-zh-cn-16k-aishell2-vocab5212-pytorch", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav') + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav" + ) logger.info("asr inference result: {0}".format(rec_result)) @@ -34,23 +39,28 @@ class TestData2vecInferencePipelines(unittest.TestCase): def test_funasr_path(self): import funasr import os + logger.info("run_dir:{0} ; funasr_path: {1}".format(os.getcwd(), funasr.__file__)) def test_transformer(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_data2vec_pretrain-zh-cn-aishell2-16k-pytorch') + model="damo/speech_data2vec_pretrain-zh-cn-aishell2-16k-pytorch", + ) rec_result = inference_pipeline( - audio_in='https://modelscope.oss-cn-beijing.aliyuncs.com/test/audios/asr_example.wav') + audio_in="https://modelscope.oss-cn-beijing.aliyuncs.com/test/audios/asr_example.wav" + ) logger.info("asr inference result: {0}".format(rec_result)) assert rec_result["text"] == "每一天都要快乐喔" def test_paraformer(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_data2vec_pretrain-paraformer-zh-cn-aishell2-16k') + model="damo/speech_data2vec_pretrain-paraformer-zh-cn-aishell2-16k", + ) rec_result = inference_pipeline( - audio_in='https://modelscope.oss-cn-beijing.aliyuncs.com/test/audios/asr_example.wav') + audio_in="https://modelscope.oss-cn-beijing.aliyuncs.com/test/audios/asr_example.wav" + ) logger.info("asr inference result: {0}".format(rec_result)) assert rec_result["text"] == "每一天都要快乐喔" @@ -59,15 +69,18 @@ class TestMfccaInferencePipelines(unittest.TestCase): def test_funasr_path(self): import funasr import os + logger.info("run_dir:{0} ; funasr_path: {1}".format(os.getcwd(), funasr.__file__)) def test_alimeeting(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='NPU-ASLP/speech_mfcca_asr-zh-cn-16k-alimeeting-vocab4950', - model_revision='v3.0.0') + model="NPU-ASLP/speech_mfcca_asr-zh-cn-16k-alimeeting-vocab4950", + model_revision="v3.0.0", + ) rec_result = inference_pipeline( - audio_in='https://pre.modelscope.cn/api/v1/models/NPU-ASLP/speech_mfcca_asr-zh-cn-16k-alimeeting-vocab4950/repo?Revision=master&FilePath=example/asr_example_mc.wav') + audio_in="https://pre.modelscope.cn/api/v1/models/NPU-ASLP/speech_mfcca_asr-zh-cn-16k-alimeeting-vocab4950/repo?Revision=master&FilePath=example/asr_example_mc.wav" + ) logger.info("asr inference result: {0}".format(rec_result)) @@ -75,103 +88,124 @@ class TestParaformerInferencePipelines(unittest.TestCase): def test_funasr_path(self): import funasr import os + logger.info("run_dir:{0} ; funasr_path: {1}".format(os.getcwd(), funasr.__file__)) def test_paraformer_large_contextual_common(self): param_dict = dict() - param_dict['hotword'] = "https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/hotword.txt" + param_dict["hotword"] = ( + "https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/hotword.txt" + ) inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_paraformer-large-contextual_asr_nat-zh-cn-16k-common-vocab8404', - param_dict=param_dict) + model="damo/speech_paraformer-large-contextual_asr_nat-zh-cn-16k-common-vocab8404", + param_dict=param_dict, + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_hotword.wav') + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_hotword.wav" + ) logger.info("asr inference result: {0}".format(rec_result)) assert rec_result["text"] == "国务院发展研究中心市场经济研究所副所长邓郁松认为" def test_paraformer_large_aishell1(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_paraformer-large_asr_nat-zh-cn-16k-aishell1-vocab8404-pytorch') + model="damo/speech_paraformer-large_asr_nat-zh-cn-16k-aishell1-vocab8404-pytorch", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav') + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav" + ) logger.info("asr inference result: {0}".format(rec_result)) assert rec_result["text"] == "欢迎大家来体验达摩院推出的语音识别模型" def test_paraformer_large_aishell2(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_paraformer-large_asr_nat-zh-cn-16k-aishell2-vocab8404-pytorch') + model="damo/speech_paraformer-large_asr_nat-zh-cn-16k-aishell2-vocab8404-pytorch", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav') + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav" + ) logger.info("asr inference result: {0}".format(rec_result)) assert rec_result["text"] == "欢迎大家来体验达摩院推出的语音识别模型" def test_paraformer_large_common(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch') + model="damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav') + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav" + ) logger.info("asr inference result: {0}".format(rec_result)) assert rec_result["text"] == "欢迎大家来体验达摩院推出的语音识别模型" def test_paraformer_large_online_common(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-online', - model_revision='v1.0.6', + model="damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-online", + model_revision="v1.0.6", update_model=False, - mode="paraformer_fake_streaming" + mode="paraformer_fake_streaming", ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav') + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav" + ) logger.info("asr inference result: {0}".format(rec_result)) assert rec_result["text"] == "欢迎大家来体验达摩院推出的语音识别模型" def test_paraformer_online_common(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_paraformer_asr_nat-zh-cn-16k-common-vocab8404-online', - model_revision='v1.0.6', + model="damo/speech_paraformer_asr_nat-zh-cn-16k-common-vocab8404-online", + model_revision="v1.0.6", update_model=False, - mode="paraformer_fake_streaming" + mode="paraformer_fake_streaming", ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav') + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav" + ) logger.info("asr inference result: {0}".format(rec_result)) assert rec_result["text"] == "欢迎大家来体验达摩院推出的语音识别模型" def test_paraformer_tiny_commandword(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_paraformer-tiny-commandword_asr_nat-zh-cn-16k-vocab544-pytorch') + model="damo/speech_paraformer-tiny-commandword_asr_nat-zh-cn-16k-vocab544-pytorch", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh_command.wav') + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh_command.wav" + ) logger.info("asr inference result: {0}".format(rec_result)) def test_paraformer_8k(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_paraformer_asr_nat-zh-cn-8k-common-vocab8358-tensorflow1') + model="damo/speech_paraformer_asr_nat-zh-cn-8k-common-vocab8358-tensorflow1", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_8K.wav') + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_8K.wav" + ) logger.info("asr inference result: {0}".format(rec_result)) def test_paraformer_aishell1(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_paraformer_asr_nat-zh-cn-16k-aishell1-vocab4234-pytorch') + model="damo/speech_paraformer_asr_nat-zh-cn-16k-aishell1-vocab4234-pytorch", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav') + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav" + ) logger.info("asr inference result: {0}".format(rec_result)) def test_paraformer_aishell2(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_paraformer_asr_nat-zh-cn-16k-aishell2-vocab5212-pytorch') + model="damo/speech_paraformer_asr_nat-zh-cn-16k-aishell2-vocab5212-pytorch", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav') + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav" + ) logger.info("asr inference result: {0}".format(rec_result)) @@ -179,22 +213,27 @@ class TestParaformerBertInferencePipelines(unittest.TestCase): def test_funasr_path(self): import funasr import os + logger.info("run_dir:{0} ; funasr_path: {1}".format(os.getcwd(), funasr.__file__)) def test_aishell1(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_paraformerbert_asr_nat-zh-cn-16k-aishell1-vocab4234-pytorch') + model="damo/speech_paraformerbert_asr_nat-zh-cn-16k-aishell1-vocab4234-pytorch", + ) rec_result = inference_pipeline( - audio_in='https://modelscope.oss-cn-beijing.aliyuncs.com/test/audios/asr_example.wav') + audio_in="https://modelscope.oss-cn-beijing.aliyuncs.com/test/audios/asr_example.wav" + ) logger.info("asr inference result: {0}".format(rec_result)) def test_aishell2(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_paraformerbert_asr_nat-zh-cn-16k-aishell2-vocab5212-pytorch') + model="damo/speech_paraformerbert_asr_nat-zh-cn-16k-aishell2-vocab5212-pytorch", + ) rec_result = inference_pipeline( - audio_in='https://modelscope.oss-cn-beijing.aliyuncs.com/test/audios/asr_example.wav') + audio_in="https://modelscope.oss-cn-beijing.aliyuncs.com/test/audios/asr_example.wav" + ) logger.info("asr inference result: {0}".format(rec_result)) @@ -202,304 +241,370 @@ class TestUniasrInferencePipelines(unittest.TestCase): def test_funasr_path(self): import funasr import os + logger.info("run_dir:{0} ; funasr_path: {1}".format(os.getcwd(), funasr.__file__)) def test_uniasr_2pass_cantonese_chs_16k_common_offline(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-cantonese-CHS-16k-common-vocab1468-tensorflow1-online') + model="damo/speech_UniASR_asr_2pass-cantonese-CHS-16k-common-vocab1468-tensorflow1-online", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_cantonese-CHS.wav', - param_dict={"decoding_model": "offline"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_cantonese-CHS.wav", + param_dict={"decoding_model": "offline"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_cantonese_chs_16k_common_online(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-cantonese-CHS-16k-common-vocab1468-tensorflow1-online') + model="damo/speech_UniASR_asr_2pass-cantonese-CHS-16k-common-vocab1468-tensorflow1-online", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_cantonese-CHS.wav', - param_dict={"decoding_model": "normal"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_cantonese-CHS.wav", + param_dict={"decoding_model": "normal"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_cn_dialect_offline(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-cn-dialect-16k-vocab8358-tensorflow1-offline') + model="damo/speech_UniASR_asr_2pass-cn-dialect-16k-vocab8358-tensorflow1-offline", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav') + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav" + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_cn_dialect_online(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-cn-dialect-16k-vocab8358-tensorflow1-online') + model="damo/speech_UniASR_asr_2pass-cn-dialect-16k-vocab8358-tensorflow1-online", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav') + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav" + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_de_common_offline(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-de-16k-common-vocab3690-tensorflow1-offline') + model="damo/speech_UniASR_asr_2pass-de-16k-common-vocab3690-tensorflow1-offline", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_de.wav', - param_dict={"decoding_model": "offline"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_de.wav", + param_dict={"decoding_model": "offline"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_de_common_online(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-de-16k-common-vocab3690-tensorflow1-online') + model="damo/speech_UniASR_asr_2pass-de-16k-common-vocab3690-tensorflow1-online", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_de.wav', - param_dict={"decoding_model": "normal"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_de.wav", + param_dict={"decoding_model": "normal"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_en_common_offline(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-en-16k-common-vocab1080-tensorflow1-offline') + model="damo/speech_UniASR_asr_2pass-en-16k-common-vocab1080-tensorflow1-offline", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_en.wav', - param_dict={"decoding_model": "offline"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_en.wav", + param_dict={"decoding_model": "offline"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_en_common_online(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-en-16k-common-vocab1080-tensorflow1-online') + model="damo/speech_UniASR_asr_2pass-en-16k-common-vocab1080-tensorflow1-online", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_en.wav', - param_dict={"decoding_model": "normal"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_en.wav", + param_dict={"decoding_model": "normal"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_es_common_offline(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-es-16k-common-vocab3445-tensorflow1-offline') + model="damo/speech_UniASR_asr_2pass-es-16k-common-vocab3445-tensorflow1-offline", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_es.wav', - param_dict={"decoding_model": "offline"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_es.wav", + param_dict={"decoding_model": "offline"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_es_common_online(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-es-16k-common-vocab3445-tensorflow1-online') + model="damo/speech_UniASR_asr_2pass-es-16k-common-vocab3445-tensorflow1-online", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_es.wav', - param_dict={"decoding_model": "normal"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_es.wav", + param_dict={"decoding_model": "normal"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_fa_common_offline(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-fa-16k-common-vocab1257-pytorch-offline') + model="damo/speech_UniASR_asr_2pass-fa-16k-common-vocab1257-pytorch-offline", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_fa.wav', - param_dict={"decoding_model": "offline"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_fa.wav", + param_dict={"decoding_model": "offline"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_fa_common_online(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-fa-16k-common-vocab1257-pytorch-online') + model="damo/speech_UniASR_asr_2pass-fa-16k-common-vocab1257-pytorch-online", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_fa.wav', - param_dict={"decoding_model": "normal"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_fa.wav", + param_dict={"decoding_model": "normal"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_fr_common_offline(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-fr-16k-common-vocab3472-tensorflow1-offline') + model="damo/speech_UniASR_asr_2pass-fr-16k-common-vocab3472-tensorflow1-offline", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_fr.wav', - param_dict={"decoding_model": "offline"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_fr.wav", + param_dict={"decoding_model": "offline"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_fr_common_online(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-fr-16k-common-vocab3472-tensorflow1-online') + model="damo/speech_UniASR_asr_2pass-fr-16k-common-vocab3472-tensorflow1-online", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_fr.wav', - param_dict={"decoding_model": "normal"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_fr.wav", + param_dict={"decoding_model": "normal"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_id_common_offline(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-id-16k-common-vocab1067-tensorflow1-online') + model="damo/speech_UniASR_asr_2pass-id-16k-common-vocab1067-tensorflow1-online", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_id.wav', - param_dict={"decoding_model": "offline"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_id.wav", + param_dict={"decoding_model": "offline"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_id_common_online(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-id-16k-common-vocab1067-tensorflow1-online') + model="damo/speech_UniASR_asr_2pass-id-16k-common-vocab1067-tensorflow1-online", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_id.wav', - param_dict={"decoding_model": "normal"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_id.wav", + param_dict={"decoding_model": "normal"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_ja_common_offline(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-ja-16k-common-vocab93-tensorflow1-offline') + model="damo/speech_UniASR_asr_2pass-ja-16k-common-vocab93-tensorflow1-offline", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_ja.wav', - param_dict={"decoding_model": "offline"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_ja.wav", + param_dict={"decoding_model": "offline"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_ja_common_online(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-ja-16k-common-vocab93-tensorflow1-online') + model="damo/speech_UniASR_asr_2pass-ja-16k-common-vocab93-tensorflow1-online", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_ja.wav', - param_dict={"decoding_model": "normal"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_ja.wav", + param_dict={"decoding_model": "normal"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_ko_common_offline(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-ko-16k-common-vocab6400-tensorflow1-offline') + model="damo/speech_UniASR_asr_2pass-ko-16k-common-vocab6400-tensorflow1-offline", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_ko.wav', - param_dict={"decoding_model": "offline"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_ko.wav", + param_dict={"decoding_model": "offline"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_ko_common_online(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-ko-16k-common-vocab6400-tensorflow1-online') + model="damo/speech_UniASR_asr_2pass-ko-16k-common-vocab6400-tensorflow1-online", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_ko.wav', - param_dict={"decoding_model": "normal"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_ko.wav", + param_dict={"decoding_model": "normal"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_minnan_common_offline(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-minnan-16k-common-vocab3825') + model="damo/speech_UniASR_asr_2pass-minnan-16k-common-vocab3825", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav', - param_dict={"decoding_model": "offline"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav", + param_dict={"decoding_model": "offline"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_pt_common_offline(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-pt-16k-common-vocab1617-tensorflow1-offline') + model="damo/speech_UniASR_asr_2pass-pt-16k-common-vocab1617-tensorflow1-offline", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_pt.wav', - param_dict={"decoding_model": "offline"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_pt.wav", + param_dict={"decoding_model": "offline"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_pt_common_online(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-pt-16k-common-vocab1617-tensorflow1-online') + model="damo/speech_UniASR_asr_2pass-pt-16k-common-vocab1617-tensorflow1-online", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_pt.wav', - param_dict={"decoding_model": "normal"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_pt.wav", + param_dict={"decoding_model": "normal"}, + ) logger.info("asr inference result: {0}".format(rec_result)) - + def test_uniasr_2pass_ru_common_offline(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-ru-16k-common-vocab1664-tensorflow1-offline') + model="damo/speech_UniASR_asr_2pass-ru-16k-common-vocab1664-tensorflow1-offline", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_ru.wav', - param_dict={"decoding_model": "offline"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_ru.wav", + param_dict={"decoding_model": "offline"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_ru_common_online(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-ru-16k-common-vocab1664-tensorflow1-online') + model="damo/speech_UniASR_asr_2pass-ru-16k-common-vocab1664-tensorflow1-online", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_ru.wav', - param_dict={"decoding_model": "normal"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_ru.wav", + param_dict={"decoding_model": "normal"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_vi_common_offline(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-vi-16k-common-vocab1001-pytorch-offline') + model="damo/speech_UniASR_asr_2pass-vi-16k-common-vocab1001-pytorch-offline", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_vi.wav', - param_dict={"decoding_model": "offline"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_vi.wav", + param_dict={"decoding_model": "offline"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_vi_common_online(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-vi-16k-common-vocab1001-pytorch-online') + model="damo/speech_UniASR_asr_2pass-vi-16k-common-vocab1001-pytorch-online", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_vi.wav', - param_dict={"decoding_model": "normal"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_vi.wav", + param_dict={"decoding_model": "normal"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_zhcn_8k_common_vocab3445_offline(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-zh-cn-8k-common-vocab3445-pytorch-offline') + model="damo/speech_UniASR_asr_2pass-zh-cn-8k-common-vocab3445-pytorch-offline", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav', - param_dict={"decoding_model": "offline"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav", + param_dict={"decoding_model": "offline"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_zhcn_8k_common_vocab3445_online(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-zh-cn-8k-common-vocab3445-pytorch-online') + model="damo/speech_UniASR_asr_2pass-zh-cn-8k-common-vocab3445-pytorch-online", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav', - param_dict={"decoding_model": "normal"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav", + param_dict={"decoding_model": "normal"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_zhcn_8k_common_vocab8358_offline(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-zh-cn-8k-common-vocab8358-tensorflow1-offline') + model="damo/speech_UniASR_asr_2pass-zh-cn-8k-common-vocab8358-tensorflow1-offline", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav', - param_dict={"decoding_model": "offline"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav", + param_dict={"decoding_model": "offline"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_zhcn_8k_common_vocab8358_online(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-zh-cn-8k-common-vocab8358-tensorflow1-online') + model="damo/speech_UniASR_asr_2pass-zh-cn-8k-common-vocab8358-tensorflow1-online", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav', - param_dict={"decoding_model": "normal"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav", + param_dict={"decoding_model": "normal"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_zhcn_16k_common_vocab8358_offline(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-zh-cn-16k-common-vocab8358-tensorflow1-offline') + model="damo/speech_UniASR_asr_2pass-zh-cn-16k-common-vocab8358-tensorflow1-offline", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav', - param_dict={"decoding_model": "offline"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav", + param_dict={"decoding_model": "offline"}, + ) logger.info("asr inference result: {0}".format(rec_result)) def test_uniasr_2pass_zhcn_16k_common_vocab8358_online(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_UniASR_asr_2pass-zh-cn-16k-common-vocab8358-tensorflow1-online') + model="damo/speech_UniASR_asr_2pass-zh-cn-16k-common-vocab8358-tensorflow1-online", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav', - param_dict={"decoding_model": "normal"}) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav", + param_dict={"decoding_model": "normal"}, + ) logger.info("asr inference result: {0}".format(rec_result)) - -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/test_asr_vad_punc_inference_pipeline.py b/tests/test_asr_vad_punc_inference_pipeline.py index f86f23dab..e8767674f 100644 --- a/tests/test_asr_vad_punc_inference_pipeline.py +++ b/tests/test_asr_vad_punc_inference_pipeline.py @@ -6,28 +6,31 @@ from modelscope.utils.logger import get_logger logger = get_logger() + class TestParaformerInferencePipelines(unittest.TestCase): def test_funasr_path(self): import funasr import os + logger.info("run_dir:{0} ; funasr_path: {1}".format(os.getcwd(), funasr.__file__)) def test_inference_pipeline(self): inference_pipeline = pipeline( task=Tasks.auto_speech_recognition, - model='damo/speech_paraformer-large-vad-punc_asr_nat-zh-cn-16k-common-vocab8404-pytorch', + model="damo/speech_paraformer-large-vad-punc_asr_nat-zh-cn-16k-common-vocab8404-pytorch", model_revision="v1.2.1", - vad_model='damo/speech_fsmn_vad_zh-cn-16k-common-pytorch', + vad_model="damo/speech_fsmn_vad_zh-cn-16k-common-pytorch", vad_model_revision="v1.1.8", - punc_model='damo/punc_ct-transformer_zh-cn-common-vocab272727-pytorch', + punc_model="damo/punc_ct-transformer_zh-cn-common-vocab272727-pytorch", punc_model_revision="v1.1.6", ngpu=1, ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav') + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav" + ) logger.info("asr_vad_punc inference result: {0}".format(rec_result)) assert rec_result["text"] == "欢迎大家来体验达摩院推出的语音识别模型。" -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/test_lm_pipeline.py b/tests/test_lm_pipeline.py index 3a5ec57a3..f67c9a57f 100644 --- a/tests/test_lm_pipeline.py +++ b/tests/test_lm_pipeline.py @@ -6,20 +6,22 @@ from modelscope.utils.logger import get_logger logger = get_logger() + class TestTransformerInferencePipelines(unittest.TestCase): def test_funasr_path(self): import funasr import os + logger.info("run_dir:{0} ; funasr_path: {1}".format(os.getcwd(), funasr.__file__)) def test_inference_pipeline(self): inference_pipeline = pipeline( task=Tasks.language_score_prediction, - model='damo/speech_transformer_lm_zh-cn-common-vocab8404-pytorch', + model="damo/speech_transformer_lm_zh-cn-common-vocab8404-pytorch", ) rec_result = inference_pipeline(text_in="hello 大 家 好 呀") logger.info("lm inference result: {0}".format(rec_result)) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/test_punctuation_pipeline.py b/tests/test_punctuation_pipeline.py index e582042bb..8ae87d170 100644 --- a/tests/test_punctuation_pipeline.py +++ b/tests/test_punctuation_pipeline.py @@ -6,16 +6,18 @@ from modelscope.utils.logger import get_logger logger = get_logger() + class TestTransformerInferencePipelines(unittest.TestCase): def test_funasr_path(self): import funasr import os + logger.info("run_dir:{0} ; funasr_path: {1}".format(os.getcwd(), funasr.__file__)) def test_inference_pipeline(self): inference_pipeline = pipeline( task=Tasks.punctuation, - model='damo/punc_ct-transformer_zh-cn-common-vocab272727-pytorch', + model="damo/punc_ct-transformer_zh-cn-common-vocab272727-pytorch", model_revision="v1.1.7", ) inputs = "./egs_modelscope/punctuation/punc_ct-transformer_zh-cn-common-vocab272727-pytorch/data/punc_example.txt" @@ -25,7 +27,7 @@ class TestTransformerInferencePipelines(unittest.TestCase): def test_vadrealtime_inference_pipeline(self): inference_pipeline = pipeline( task=Tasks.punctuation, - model='damo/punc_ct-transformer_zh-cn-common-vad_realtime-vocab272727', + model="damo/punc_ct-transformer_zh-cn-common-vad_realtime-vocab272727", ) inputs = "跨境河流是养育沿岸|人民的生命之源长期以来为帮助下游地区防灾减灾中方技术人员|在上游地区极为恶劣的自然条件下克服巨大困难甚至冒着生命危险|向印方提供汛期水文资料处理紧急事件中方重视印方在跨境河流问题上的关切|愿意进一步完善双方联合工作机制|凡是|中方能做的我们|都会去做而且会做得更好我请印度朋友们放心中国在上游的|任何开发利用都会经过科学|规划和论证兼顾上下游的利益" vads = inputs.split("|") @@ -37,5 +39,5 @@ class TestTransformerInferencePipelines(unittest.TestCase): logger.info("punctuation inference result: {0}".format(rec_result_all)) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/test_sv_inference_pipeline.py b/tests/test_sv_inference_pipeline.py index c4e427efc..76754ab21 100644 --- a/tests/test_sv_inference_pipeline.py +++ b/tests/test_sv_inference_pipeline.py @@ -12,27 +12,37 @@ class TestXVectorInferencePipelines(unittest.TestCase): def test_funasr_path(self): import funasr import os + logger.info("run_dir:{0} ; funasr_path: {1}".format(os.getcwd(), funasr.__file__)) def test_inference_pipeline(self): inference_sv_pipline = pipeline( task=Tasks.speaker_verification, - model='damo/speech_xvector_sv-zh-cn-cnceleb-16k-spk3465-pytorch' + model="damo/speech_xvector_sv-zh-cn-cnceleb-16k-spk3465-pytorch", ) # the same speaker - rec_result = inference_sv_pipline(audio_in=( - 'https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/sv_example_enroll.wav', - 'https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/sv_example_same.wav')) - assert abs(rec_result["scores"][0]-0.85) < 0.1 and abs(rec_result["scores"][1]-0.14) < 0.1 - logger.info(f"Similarity {rec_result['scores']}") - - # different speaker - rec_result = inference_sv_pipline(audio_in=( - 'https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/sv_example_enroll.wav', - 'https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/sv_example_different.wav')) - assert abs(rec_result["scores"][0]-0.0) < 0.1 and abs(rec_result["scores"][1]-1.0) < 0.1 + rec_result = inference_sv_pipline( + audio_in=( + "https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/sv_example_enroll.wav", + "https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/sv_example_same.wav", + ) + ) + assert ( + abs(rec_result["scores"][0] - 0.85) < 0.1 and abs(rec_result["scores"][1] - 0.14) < 0.1 + ) logger.info(f"Similarity {rec_result['scores']}") -if __name__ == '__main__': - unittest.main() \ No newline at end of file + # different speaker + rec_result = inference_sv_pipline( + audio_in=( + "https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/sv_example_enroll.wav", + "https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/sv_example_different.wav", + ) + ) + assert abs(rec_result["scores"][0] - 0.0) < 0.1 and abs(rec_result["scores"][1] - 1.0) < 0.1 + logger.info(f"Similarity {rec_result['scores']}") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_tp_pipeline.py b/tests/test_tp_pipeline.py index 07084f253..c4853ab13 100644 --- a/tests/test_tp_pipeline.py +++ b/tests/test_tp_pipeline.py @@ -6,25 +6,53 @@ from modelscope.utils.logger import get_logger logger = get_logger() + class TestTimestampPredictionPipelines(unittest.TestCase): def test_funasr_path(self): import funasr import os + logger.info("run_dir:{0} ; funasr_path: {1}".format(os.getcwd(), funasr.__file__)) def test_inference_pipeline(self): inference_pipeline = pipeline( task=Tasks.speech_timestamp, - model='damo/speech_timestamp_prediction-v1-16k-offline', - model_revision='v1.1.0') + model="damo/speech_timestamp_prediction-v1-16k-offline", + model_revision="v1.1.0", + ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_timestamps.wav', - text_in='一 个 东 太 平 洋 国 家 为 什 么 跑 到 西 太 平 洋 来 了 呢',) + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_timestamps.wav", + text_in="一 个 东 太 平 洋 国 家 为 什 么 跑 到 西 太 平 洋 来 了 呢", + ) print(rec_result) logger.info("punctuation inference result: {0}".format(rec_result)) - assert rec_result=={'text': ' 0.000 0.380;一 0.380 0.560;个 0.560 0.800;东 0.800 0.980;太 0.980 1.140;平 1.140 1.260;洋 1.260 1.440;国 1.440 1.680;家 1.680 1.920; 1.920 2.040;为 2.040 2.200;什 2.200 2.320;么 2.320 2.500;跑 2.500 2.680;到 2.680 2.860;西 2.860 3.040;太 3.040 3.200;平 3.200 3.380;洋 3.380 3.500;来 3.500 3.640;了 3.640 3.800;呢 3.800 4.150; 4.150 4.440;', 'timestamp': [[380, 560], [560, 800], [800, 980], [980, 1140], [1140, 1260], [1260, 1440], [1440, 1680], [1680, 1920], [2040, 2200], [2200, 2320], [2320, 2500], [2500, 2680], [2680, 2860], [2860, 3040], [3040, 3200], [3200, 3380], [3380, 3500], [3500, 3640], [3640, 3800], [3800, 4150]]} + assert rec_result == { + "text": " 0.000 0.380;一 0.380 0.560;个 0.560 0.800;东 0.800 0.980;太 0.980 1.140;平 1.140 1.260;洋 1.260 1.440;国 1.440 1.680;家 1.680 1.920; 1.920 2.040;为 2.040 2.200;什 2.200 2.320;么 2.320 2.500;跑 2.500 2.680;到 2.680 2.860;西 2.860 3.040;太 3.040 3.200;平 3.200 3.380;洋 3.380 3.500;来 3.500 3.640;了 3.640 3.800;呢 3.800 4.150; 4.150 4.440;", + "timestamp": [ + [380, 560], + [560, 800], + [800, 980], + [980, 1140], + [1140, 1260], + [1260, 1440], + [1440, 1680], + [1680, 1920], + [2040, 2200], + [2200, 2320], + [2320, 2500], + [2500, 2680], + [2680, 2860], + [2860, 3040], + [3040, 3200], + [3200, 3380], + [3380, 3500], + [3500, 3640], + [3640, 3800], + [3800, 4150], + ], + } -if __name__ == '__main__': - unittest.main() \ No newline at end of file +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_vad_inference_pipeline.py b/tests/test_vad_inference_pipeline.py index 50b8db3b5..7fb0ce484 100644 --- a/tests/test_vad_inference_pipeline.py +++ b/tests/test_vad_inference_pipeline.py @@ -11,6 +11,7 @@ class TestFSMNInferencePipelines(unittest.TestCase): def test_funasr_path(self): import funasr import os + logger.info("run_dir:{0} ; funasr_path: {1}".format(os.getcwd(), funasr.__file__)) def test_8k(self): @@ -19,15 +20,41 @@ class TestFSMNInferencePipelines(unittest.TestCase): model="damo/speech_fsmn_vad_zh-cn-8k-common", ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/vad_example_8k.wav') + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/vad_example_8k.wav" + ) logger.info("vad inference result: {0}".format(rec_result)) - assert rec_result["text"] == [[0, 1960], [2870, 6730], [7960, 10180], [12140, 14830], [15740, 19400], - [20220, 24230], [25540, 27290], [30070, 30970], [32070, 34280], [35990, 37050], - [39400, 41020], [41810, 47320], [48120, 52150], [53560, 58310], [59290, 62210], - [63110, 66420], [67300, 68280], [69670, 71770], [73100, 75550], [76850, 78500], - [79380, 83280], [85000, 92320], [93560, 94110], [94990, 95620], [96940, 97590], - [98400, 100530], [101600, 104890], [108780, 110900], [112020, 113460], - [114210, 115030]] + assert rec_result["text"] == [ + [0, 1960], + [2870, 6730], + [7960, 10180], + [12140, 14830], + [15740, 19400], + [20220, 24230], + [25540, 27290], + [30070, 30970], + [32070, 34280], + [35990, 37050], + [39400, 41020], + [41810, 47320], + [48120, 52150], + [53560, 58310], + [59290, 62210], + [63110, 66420], + [67300, 68280], + [69670, 71770], + [73100, 75550], + [76850, 78500], + [79380, 83280], + [85000, 92320], + [93560, 94110], + [94990, 95620], + [96940, 97590], + [98400, 100530], + [101600, 104890], + [108780, 110900], + [112020, 113460], + [114210, 115030], + ] def test_16k(self): inference_pipeline = pipeline( @@ -35,13 +62,24 @@ class TestFSMNInferencePipelines(unittest.TestCase): model="damo/speech_fsmn_vad_zh-cn-16k-common-pytorch", ) rec_result = inference_pipeline( - audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/vad_example.wav') + audio_in="https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/vad_example.wav" + ) logger.info("vad inference result: {0}".format(rec_result)) - assert rec_result["text"] == [[70, 2340], [2620, 6200], [6480, 23670], [23950, 26250], [26780, 28990], - [29950, 31430], [31750, 37600], [38210, 46900], [47310, 49630], [49910, 56460], - [56740, 59540], [59820, 70450]] + assert rec_result["text"] == [ + [70, 2340], + [2620, 6200], + [6480, 23670], + [23950, 26250], + [26780, 28990], + [29950, 31430], + [31750, 37600], + [38210, 46900], + [47310, 49630], + [49910, 56460], + [56740, 59540], + [59820, 70450], + ] -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() -