[Wav2Vec2] PyCTCDecode Integration to support language model boosted decoding

docs/source/model_doc/wav2vec2.rst

@@ -67,20 +67,13 @@ Wav2Vec2Processor
     :members: __call__, pad, from_pretrained, save_pretrained, batch_decode, decode, as_target_processor
 
 
-Wav2Vec2 specific outputs
+Wav2Vec2ProcessorWithLM
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-.. autoclass:: transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2BaseModelOutput
-    :members: 
-
-.. autoclass:: transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForPreTrainingOutput
-    :members: 
+.. autoclass:: transformers.Wav2Vec2ProcessorWithLM
+    :members: __call__, pad, from_pretrained, save_pretrained, batch_decode, decode, as_target_processor
 
-.. autoclass:: transformers.models.wav2vec2.modeling_flax_wav2vec2.FlaxWav2Vec2BaseModelOutput
-    :members: 
 
-.. autoclass:: transformers.models.wav2vec2.modeling_flax_wav2vec2.FlaxWav2Vec2ForPreTrainingOutput
-    :members: 
 
 
 Wav2Vec2Model
@@ -143,3 +136,19 @@ FlaxWav2Vec2ForPreTraining
 
 .. autoclass:: transformers.FlaxWav2Vec2ForPreTraining
     :members: __call__
+
+Wav2Vec2 specific outputs
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2BaseModelOutput
+    :members: 
+
+.. autoclass:: transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForPreTrainingOutput
+    :members: 
+
+.. autoclass:: transformers.models.wav2vec2.modeling_flax_wav2vec2.FlaxWav2Vec2BaseModelOutput
+    :members: 
+
+.. autoclass:: transformers.models.wav2vec2.modeling_flax_wav2vec2.FlaxWav2Vec2ForPreTrainingOutput
+    :members: 
+

setup.py

@@ -151,6 +151,7 @@
     "tokenizers>=0.10.1,<0.11",
     "torch>=1.0,<1.10",
     "torchaudio",
+    "pyctcdecode>=0.2.0",
     "tqdm>=4.27",
     "unidic>=1.0.2",
     "unidic_lite>=1.0.7",
@@ -256,7 +257,7 @@ def run(self):
 extras["serving"] = deps_list("pydantic", "uvicorn", "fastapi", "starlette")
 extras["audio"] = deps_list("librosa")
 extras["speech"] = deps_list("torchaudio") + extras["audio"]  # `pip install ".[speech]"` is deprecated and `pip install ".[torch-speech]"` should be used instead
-extras["torch-speech"] = deps_list("torchaudio") + extras["audio"]
+extras["torch-speech"] = deps_list("torchaudio", "pyctcdecode") + extras["audio"]
 extras["tf-speech"] = extras["audio"]
 extras["flax-speech"] = extras["audio"]
 extras["vision"] = deps_list("Pillow")

src/transformers/__init__.py

@@ -298,6 +298,7 @@
         "Wav2Vec2CTCTokenizer",
         "Wav2Vec2FeatureExtractor",
         "Wav2Vec2Processor",
+        "Wav2Vec2ProcessorWithLM",
         "Wav2Vec2Tokenizer",
     ],
     "models.xlm": ["XLM_PRETRAINED_CONFIG_ARCHIVE_MAP", "XLMConfig", "XLMTokenizer"],
@@ -2191,6 +2192,7 @@
         Wav2Vec2CTCTokenizer,
         Wav2Vec2FeatureExtractor,
         Wav2Vec2Processor,
+        Wav2Vec2ProcessorWithLM,
         Wav2Vec2Tokenizer,
     )
     from .models.xlm import XLM_PRETRAINED_CONFIG_ARCHIVE_MAP, XLMConfig, XLMTokenizer

src/transformers/dependency_versions_table.py

@@ -69,6 +69,7 @@
     "tokenizers": "tokenizers>=0.10.1,<0.11",
     "torch": "torch>=1.0,<1.10",
     "torchaudio": "torchaudio",
+    "pyctcdecode": "pyctcdecode>=0.2.0",
     "tqdm": "tqdm>=4.27",
     "unidic": "unidic>=1.0.2",
     "unidic_lite": "unidic_lite>=1.0.7",

src/transformers/file_utils.py

@@ -220,6 +220,14 @@
     _torchaudio_available = False
 
 
+_pyctcdecode_available = importlib.util.find_spec("pyctcdecode") is not None
+try:
+    _pyctcdecode_version = importlib_metadata.version("pyctcdecode")
+    logger.debug(f"Successfully imported pyctcdecode version {_pyctcdecode_version}")
+except importlib_metadata.PackageNotFoundError:
+    _pyctcdecode_available = False
+
+
 torch_cache_home = os.getenv("TORCH_HOME", os.path.join(os.getenv("XDG_CACHE_HOME", "~/.cache"), "torch"))
 old_default_cache_path = os.path.join(torch_cache_home, "transformers")
 # New default cache, shared with the Datasets library
@@ -294,6 +302,10 @@ def is_torch_available():
     return _torch_available
 
 
+def is_pyctcdecode_available():
+    return _pyctcdecode_available
+
+
 def is_torch_cuda_available():
     if is_torch_available():
         import torch
@@ -650,6 +662,12 @@ def wrapper(*args, **kwargs):
 `pip install pytesseract`
 """
 
+# docstyle-ignore
+PYCTCDECODE_IMPORT_ERROR = """
+{0} requires the pyctcdecode library but it was not found in your environment. You can install it with pip:
+`pip install pyctcdecode`
+"""
+
 
 BACKENDS_MAPPING = OrderedDict(
     [
@@ -659,6 +677,7 @@ def wrapper(*args, **kwargs):
         ("flax", (is_flax_available, FLAX_IMPORT_ERROR)),
         ("pandas", (is_pandas_available, PANDAS_IMPORT_ERROR)),
         ("protobuf", (is_protobuf_available, PROTOBUF_IMPORT_ERROR)),
+        ("pyctcdecode", (is_pyctcdecode_available, PYCTCDECODE_IMPORT_ERROR)),
         ("pytesseract", (is_pytesseract_available, PYTESSERACT_IMPORT_ERROR)),
         ("scatter", (is_scatter_available, SCATTER_IMPORT_ERROR)),
         ("sentencepiece", (is_sentencepiece_available, SENTENCEPIECE_IMPORT_ERROR)),

src/transformers/models/wav2vec2/__init__.py

@@ -24,6 +24,7 @@
     "configuration_wav2vec2": ["WAV_2_VEC_2_PRETRAINED_CONFIG_ARCHIVE_MAP", "Wav2Vec2Config"],
     "feature_extraction_wav2vec2": ["Wav2Vec2FeatureExtractor"],
     "processing_wav2vec2": ["Wav2Vec2Processor"],
+    "processing_wav2vec2_with_lm": ["Wav2Vec2ProcessorWithLM"],
     "tokenization_wav2vec2": ["Wav2Vec2CTCTokenizer", "Wav2Vec2Tokenizer"],
 }
 
@@ -59,6 +60,7 @@
     from .configuration_wav2vec2 import WAV_2_VEC_2_PRETRAINED_CONFIG_ARCHIVE_MAP, Wav2Vec2Config
     from .feature_extraction_wav2vec2 import Wav2Vec2FeatureExtractor
     from .processing_wav2vec2 import Wav2Vec2Processor
+    from .processing_wav2vec2_with_lm import Wav2Vec2ProcessorWithLM
     from .tokenization_wav2vec2 import Wav2Vec2CTCTokenizer, Wav2Vec2Tokenizer
 
     if is_torch_available():