Switch to flashlight decoder from upstream (#2557)

Summary:
Pull Request resolved: https://github.com/pytorch/audio/pull/2557

Allow the use of flahslight-decoder from upstream

Reviewed By: carolineechen

Differential Revision: D37983846

fbshipit-source-id: edb1b701bd18718b3b10cf51cc63d3924d4cc073

torchaudio/models/decoder/__init__.py

@@ -7,20 +7,8 @@ _LAZILY_IMPORTED = [
 ]
 
 
-def _init_extension():
-    import torchaudio
-
-    torchaudio._extension._load_lib("libtorchaudio_decoder")
-
-    global _INITIALIZED
-    _INITIALIZED = True
-
-
 def __getattr__(name: str):
     if name in _LAZILY_IMPORTED:
-        if not _INITIALIZED:
-            _init_extension()
-
         try:
             from . import _ctc_decoder
         except AttributeError as err:

torchaudio/models/decoder/_ctc_decoder.py

@@ -3,23 +3,49 @@ from collections import namedtuple
 from typing import Dict, List, NamedTuple, Optional, Union
 
 import torch
-from torchaudio._torchaudio_decoder import (
-    _create_word_dict,
-    _CriterionType,
-    _Dictionary,
-    _KenLM,
-    _LexiconDecoder,
-    _LexiconDecoderOptions,
-    _LexiconFreeDecoder,
-    _LexiconFreeDecoderOptions,
-    _LM,
-    _load_words,
-    _SmearingMode,
-    _Trie,
-    _ZeroLM,
-)
+import torchaudio
 from torchaudio.utils import download_asset
 
+try:
+    # We prioritize the version from upstream flashlight here.
+    # This will allow applications that use the upstream flashlight
+    # alongside torchaudio.
+    from flashlight.lib.text.decoder import (
+        CriterionType as _CriterionType,
+        KenLM as _KenLM,
+        LexiconDecoder as _LexiconDecoder,
+        LexiconDecoderOptions as _LexiconDecoderOptions,
+        LexiconFreeDecoder as _LexiconFreeDecoder,
+        LexiconFreeDecoderOptions as _LexiconFreeDecoderOptions,
+        LM as _LM,
+        SmearingMode as _SmearingMode,
+        Trie as _Trie,
+        ZeroLM as _ZeroLM,
+    )
+    from flashlight.lib.text.dictionary import (
+        create_word_dict as _create_word_dict,
+        Dictionary as _Dictionary,
+        load_words as _load_words,
+    )
+except Exception:
+    torchaudio._extension._load_lib("libtorchaudio_decoder")
+    from torchaudio._torchaudio_decoder import (
+        _create_word_dict,
+        _CriterionType,
+        _Dictionary,
+        _KenLM,
+        _LexiconDecoder,
+        _LexiconDecoderOptions,
+        _LexiconFreeDecoder,
+        _LexiconFreeDecoderOptions,
+        _LM,
+        _load_words,
+        _SmearingMode,
+        _Trie,
+        _ZeroLM,
+    )
+
+
 __all__ = ["CTCHypothesis", "CTCDecoder", "ctc_decoder", "download_pretrained_files"]