Spaces:
Runtime error
Runtime error
| # Copyright 2022 Xiaomi Corp. (authors: Fangjun Kuang) | |
| # | |
| # See LICENSE for clarification regarding multiple authors | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| from huggingface_hub import hf_hub_download | |
| from functools import lru_cache | |
| from offline_asr import OfflineAsr | |
| sample_rate = 16000 | |
| def get_pretrained_model(repo_id: str) -> OfflineAsr: | |
| if repo_id in chinese_models: | |
| return chinese_models[repo_id](repo_id) | |
| elif repo_id in english_models: | |
| return english_models[repo_id](repo_id) | |
| elif repo_id in chinese_english_mixed_models: | |
| return chinese_english_mixed_models[repo_id](repo_id) | |
| else: | |
| raise ValueError(f"Unsupported repo_id: {repo_id}") | |
| def _get_nn_model_filename( | |
| repo_id: str, | |
| filename: str, | |
| subfolder: str = "exp", | |
| ) -> str: | |
| nn_model_filename = hf_hub_download( | |
| repo_id=repo_id, | |
| filename=filename, | |
| subfolder=subfolder, | |
| ) | |
| return nn_model_filename | |
| def _get_bpe_model_filename( | |
| repo_id: str, | |
| filename: str = "bpe.model", | |
| subfolder: str = "data/lang_bpe_500", | |
| ) -> str: | |
| bpe_model_filename = hf_hub_download( | |
| repo_id=repo_id, | |
| filename=filename, | |
| subfolder=subfolder, | |
| ) | |
| return bpe_model_filename | |
| def _get_token_filename( | |
| repo_id: str, | |
| filename: str = "tokens.txt", | |
| subfolder: str = "data/lang_char", | |
| ) -> str: | |
| token_filename = hf_hub_download( | |
| repo_id=repo_id, | |
| filename=filename, | |
| subfolder=subfolder, | |
| ) | |
| return token_filename | |
| def _get_aishell2_pretrained_model(repo_id: str) -> OfflineAsr: | |
| assert repo_id in [ | |
| # context-size 1 | |
| "yuekai/icefall-asr-aishell2-pruned-transducer-stateless5-A-2022-07-12", # noqa | |
| # context-size 2 | |
| "yuekai/icefall-asr-aishell2-pruned-transducer-stateless5-B-2022-07-12", # noqa | |
| ] | |
| nn_model_filename = _get_nn_model_filename( | |
| repo_id=repo_id, | |
| filename="cpu_jit.pt", | |
| ) | |
| token_filename = _get_token_filename(repo_id=repo_id) | |
| return OfflineAsr( | |
| nn_model_filename=nn_model_filename, | |
| bpe_model_filename=None, | |
| token_filename=token_filename, | |
| sample_rate=sample_rate, | |
| device="cpu", | |
| ) | |
| def _get_gigaspeech_pre_trained_model(repo_id: str) -> OfflineAsr: | |
| assert repo_id in [ | |
| "wgb14/icefall-asr-gigaspeech-pruned-transducer-stateless2", | |
| ] | |
| nn_model_filename = _get_nn_model_filename( | |
| # It is converted from https://huggingface.co/wgb14/icefall-asr-gigaspeech-pruned-transducer-stateless2 # noqa | |
| repo_id="csukuangfj/icefall-asr-gigaspeech-pruned-transducer-stateless2", # noqa | |
| filename="cpu_jit-epoch-29-avg-11-torch-1.10.0.pt", | |
| ) | |
| bpe_model_filename = _get_bpe_model_filename(repo_id=repo_id) | |
| return OfflineAsr( | |
| nn_model_filename=nn_model_filename, | |
| bpe_model_filename=bpe_model_filename, | |
| token_filename=None, | |
| sample_rate=sample_rate, | |
| device="cpu", | |
| ) | |
| def _get_librispeech_pre_trained_model(repo_id: str) -> OfflineAsr: | |
| assert repo_id in [ | |
| "csukuangfj/icefall-asr-librispeech-pruned-transducer-stateless3-2022-05-13", # noqa | |
| ] | |
| nn_model_filename = _get_nn_model_filename( | |
| repo_id=repo_id, | |
| filename="cpu_jit.pt", | |
| ) | |
| bpe_model_filename = _get_bpe_model_filename(repo_id=repo_id) | |
| return OfflineAsr( | |
| nn_model_filename=nn_model_filename, | |
| bpe_model_filename=bpe_model_filename, | |
| token_filename=None, | |
| sample_rate=sample_rate, | |
| device="cpu", | |
| ) | |
| def _get_wenetspeech_pre_trained_model(repo_id: str): | |
| assert repo_id in [ | |
| "luomingshuang/icefall_asr_wenetspeech_pruned_transducer_stateless2", | |
| ] | |
| nn_model_filename = _get_nn_model_filename( | |
| repo_id=repo_id, | |
| filename="cpu_jit_epoch_10_avg_2_torch_1.7.1.pt", | |
| ) | |
| token_filename = _get_token_filename(repo_id=repo_id) | |
| return OfflineAsr( | |
| nn_model_filename=nn_model_filename, | |
| bpe_model_filename=None, | |
| token_filename=token_filename, | |
| sample_rate=sample_rate, | |
| device="cpu", | |
| ) | |
| def _get_tal_csasr_pre_trained_model(repo_id: str): | |
| assert repo_id in [ | |
| "luomingshuang/icefall_asr_tal-csasr_pruned_transducer_stateless5", | |
| ] | |
| nn_model_filename = _get_nn_model_filename( | |
| repo_id=repo_id, | |
| filename="cpu_jit.pt", | |
| ) | |
| token_filename = _get_token_filename(repo_id=repo_id) | |
| return OfflineAsr( | |
| nn_model_filename=nn_model_filename, | |
| bpe_model_filename=None, | |
| token_filename=token_filename, | |
| sample_rate=sample_rate, | |
| device="cpu", | |
| ) | |
| chinese_models = { | |
| "yuekai/icefall-asr-aishell2-pruned-transducer-stateless5-A-2022-07-12": _get_aishell2_pretrained_model, # noqa | |
| "yuekai/icefall-asr-aishell2-pruned-transducer-stateless5-B-2022-07-12": _get_aishell2_pretrained_model, # noqa | |
| "luomingshuang/icefall_asr_wenetspeech_pruned_transducer_stateless2": _get_wenetspeech_pre_trained_model, # noqa | |
| } | |
| english_models = { | |
| "wgb14/icefall-asr-gigaspeech-pruned-transducer-stateless2": _get_gigaspeech_pre_trained_model, # noqa | |
| "csukuangfj/icefall-asr-librispeech-pruned-transducer-stateless3-2022-05-13": _get_librispeech_pre_trained_model, # noqa | |
| } | |
| chinese_english_mixed_models = { | |
| "luomingshuang/icefall_asr_tal-csasr_pruned_transducer_stateless5": _get_tal_csasr_pre_trained_model, # noqa | |
| } | |
| all_models = { | |
| **chinese_models, | |
| **english_models, | |
| **chinese_english_mixed_models, | |
| } | |
| language_to_models = { | |
| "Chinese": sorted(chinese_models.keys()), | |
| "English": sorted(english_models.keys()), | |
| "Chinese+English": sorted(chinese_english_mixed_models.keys()), | |
| } | |