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,7 +3,33 @@ from collections import namedtuple
 from typing import Dict, List, NamedTuple, Optional, Union
 
 import torch
-from torchaudio._torchaudio_decoder import (
+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,
@@ -17,8 +43,8 @@ from torchaudio._torchaudio_decoder import (
         _SmearingMode,
         _Trie,
         _ZeroLM,
-)
-from torchaudio.utils import download_asset
+    )
+
 
 __all__ = ["CTCHypothesis", "CTCDecoder", "ctc_decoder", "download_pretrained_files"]