riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-07-31 02:02:21 +06:00
Code Issues Actions 23 Packages Projects Releases Wiki Activity

[Wav2Vec2ProcessorWithLM] improve decoder downlaod (#15040)

Browse Source
This commit is contained in:
Patrick von Platen 2022-01-11 11:59:38 +01:00 committed by GitHub
parent 6ea6266625
commit efb35a4107
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 22 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
src/transformers/models/wav2vec2_with_lm/processing_wav2vec2_with_lm.py
View File

@ -166,7 +166,14 @@ class Wav2Vec2ProcessorWithLM:
# BeamSearchDecoderCTC has no auto class
kwargs.pop("_from_auto", None)
decoder = BeamSearchDecoderCTC.load_from_hf_hub(pretrained_model_name_or_path, **kwargs)
# make sure that only relevant filenames are downloaded
language_model_filenames = os.path.join(BeamSearchDecoderCTC._LANGUAGE_MODEL_SERIALIZED_DIRECTORY, "*")
alphabet_filename = BeamSearchDecoderCTC._ALPHABET_SERIALIZED_FILENAME
allow_regex = [language_model_filenames, alphabet_filename]
decoder = BeamSearchDecoderCTC.load_from_hf_hub(
pretrained_model_name_or_path, allow_regex=allow_regex, **kwargs
)
# set language model attributes
for attribute in ["alpha", "beta", "unk_score_offset", "score_boundary"]:

14
tests/test_processor_wav2vec2_with_lm.py
View File

@ -18,6 +18,7 @@ import shutil
import tempfile
import unittest
from multiprocessing import Pool
from pathlib import Path
import numpy as np
@ -234,3 +235,16 @@ class Wav2Vec2ProcessorWithLMTest(unittest.TestCase):
self.assertListEqual(decoded_decoder, decoded_processor)
self.assertListEqual(["<s> </s> </s>", "<s> <s> </s>"], decoded_processor)
def test_decoder_download_ignores_files(self):
processor = Wav2Vec2ProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm")
language_model = processor.decoder.model_container[processor.decoder._model_key]
path_to_cached_dir = Path(language_model._kenlm_model.path.decode("utf-8")).parent.parent.absolute()
downloaded_decoder_files = os.listdir(path_to_cached_dir)
# test that only decoder relevant files from
# https://huggingface.co/hf-internal-testing/processor_with_lm/tree/main
# are downloaded and none of the rest (e.g. README.md, ...)
self.assertListEqual(downloaded_decoder_files, ["alphabet.json", "language_model"])