Allow user-managed Pool in Wav2Vec2ProcessorWithLM.batch_decode (#18351)
* [Wav2Vec2] Allow user-managed Pool in Wav2Vec2ProcessorWithLM.batch_decode * [Wav2Vec2] Add user-managed LM's pool tests and usage examples * Improve styling Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com> * [Wav2Vec2] Fix hyperlink references Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
This commit is contained in:
committed by
GitHub
parent
bf0e094142
commit
af150e4a1c
@@ -18,6 +18,7 @@ import copy
|
||||
import glob
|
||||
import inspect
|
||||
import math
|
||||
import multiprocessing
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
@@ -26,7 +27,7 @@ from datasets import load_dataset
|
||||
|
||||
from huggingface_hub import snapshot_download
|
||||
from transformers import Wav2Vec2Config, is_tf_available
|
||||
from transformers.testing_utils import is_flaky, require_librosa, require_pyctcdecode, require_tf, slow
|
||||
from transformers.testing_utils import CaptureLogger, is_flaky, require_librosa, require_pyctcdecode, require_tf, slow
|
||||
from transformers.utils import is_librosa_available, is_pyctcdecode_available
|
||||
|
||||
from ...test_configuration_common import ConfigTester
|
||||
@@ -42,6 +43,7 @@ if is_tf_available():
|
||||
|
||||
if is_pyctcdecode_available():
|
||||
from transformers import Wav2Vec2ProcessorWithLM
|
||||
from transformers.models.wav2vec2_with_lm import processing_wav2vec2_with_lm
|
||||
|
||||
|
||||
if is_librosa_available():
|
||||
@@ -590,3 +592,56 @@ class TFWav2Vec2ModelIntegrationTest(unittest.TestCase):
|
||||
transcription = processor.batch_decode(logits.numpy()).text
|
||||
|
||||
self.assertEqual(transcription[0], "el libro ha sido escrito por cervantes")
|
||||
|
||||
@require_pyctcdecode
|
||||
@require_librosa
|
||||
def test_wav2vec2_with_lm_pool(self):
|
||||
downloaded_folder = snapshot_download("patrickvonplaten/common_voice_es_sample")
|
||||
file_path = glob.glob(downloaded_folder + "/*")[0]
|
||||
sample = librosa.load(file_path, sr=16_000)[0]
|
||||
|
||||
model = TFWav2Vec2ForCTC.from_pretrained("patrickvonplaten/wav2vec2-large-xlsr-53-spanish-with-lm")
|
||||
processor = Wav2Vec2ProcessorWithLM.from_pretrained("patrickvonplaten/wav2vec2-large-xlsr-53-spanish-with-lm")
|
||||
|
||||
input_values = processor(sample, return_tensors="tf").input_values
|
||||
|
||||
logits = model(input_values).logits
|
||||
|
||||
# test user-managed pool
|
||||
with multiprocessing.get_context("fork").Pool(2) as pool:
|
||||
transcription = processor.batch_decode(logits.numpy(), pool).text
|
||||
|
||||
self.assertEqual(transcription[0], "el libro ha sido escrito por cervantes")
|
||||
|
||||
# user-managed pool + num_processes should trigger a warning
|
||||
with CaptureLogger(processing_wav2vec2_with_lm.logger) as cl, multiprocessing.get_context("fork").Pool(
|
||||
2
|
||||
) as pool:
|
||||
transcription = processor.batch_decode(logits.numpy(), pool, num_processes=2).text
|
||||
|
||||
self.assertIn("num_process", cl.out)
|
||||
self.assertIn("it will be ignored", cl.out)
|
||||
|
||||
self.assertEqual(transcription[0], "el libro ha sido escrito por cervantes")
|
||||
|
||||
@require_pyctcdecode
|
||||
@require_librosa
|
||||
def test_wav2vec2_with_lm_invalid_pool(self):
|
||||
downloaded_folder = snapshot_download("patrickvonplaten/common_voice_es_sample")
|
||||
file_path = glob.glob(downloaded_folder + "/*")[0]
|
||||
sample = librosa.load(file_path, sr=16_000)[0]
|
||||
|
||||
model = TFWav2Vec2ForCTC.from_pretrained("patrickvonplaten/wav2vec2-large-xlsr-53-spanish-with-lm")
|
||||
processor = Wav2Vec2ProcessorWithLM.from_pretrained("patrickvonplaten/wav2vec2-large-xlsr-53-spanish-with-lm")
|
||||
|
||||
input_values = processor(sample, return_tensors="tf").input_values
|
||||
|
||||
logits = model(input_values).logits
|
||||
|
||||
# change default start method, which should trigger a warning if different than fork
|
||||
multiprocessing.set_start_method("spawn")
|
||||
with CaptureLogger(processing_wav2vec2_with_lm.logger) as cl:
|
||||
transcription = processor.batch_decode(logits.numpy()).text
|
||||
|
||||
self.assertIn("Falling back to sequential decoding.", cl.out)
|
||||
self.assertEqual(transcription[0], "el libro ha sido escrito por cervantes")
|
||||
|
||||
Reference in New Issue
Block a user