wav2vec2 fixes (#40341)

* Changed datasets to avoid a datasets error

* Changed back split to test
This commit is contained in:
Rémi Ouazan
2025-08-22 11:32:29 +02:00
committed by GitHub
parent d7fe3111ff
commit e018b77c89

View File

@ -97,7 +97,7 @@ def _test_wav2vec2_with_lm_invalid_pool(in_queue, out_queue, timeout):
try:
_ = in_queue.get(timeout=timeout)
ds = load_dataset("mozilla-foundation/common_voice_11_0", "es", split="test", streaming=True)
ds = load_dataset("fixie-ai/common_voice_17_0", "es", split="test", streaming=True)
sample = next(iter(ds))
resampled_audio = torchaudio.functional.resample(
@ -118,8 +118,9 @@ def _test_wav2vec2_with_lm_invalid_pool(in_queue, out_queue, timeout):
with CaptureLogger(pyctcdecode.decoder.logger) as cl, multiprocessing.get_context("spawn").Pool(1) as pool:
transcription = processor.batch_decode(logits.cpu().numpy(), pool).text
expected_transcription = "el resto de los equipos se mantienen en su sede"
unittest.TestCase().assertIn("Falling back to sequential decoding.", cl.out)
unittest.TestCase().assertEqual(transcription[0], "habitan aguas poco profundas y rocosas")
unittest.TestCase().assertEqual(transcription[0], expected_transcription)
# force batch_decode to internally create a spawn pool, which should trigger a warning if different than fork
multiprocessing.set_start_method("spawn", force=True)
@ -127,7 +128,7 @@ def _test_wav2vec2_with_lm_invalid_pool(in_queue, out_queue, timeout):
transcription = processor.batch_decode(logits.cpu().numpy()).text
unittest.TestCase().assertIn("Falling back to sequential decoding.", cl.out)
unittest.TestCase().assertEqual(transcription[0], "habitan aguas poco profundas y rocosas")
unittest.TestCase().assertEqual(transcription[0], expected_transcription)
except Exception:
error = f"{traceback.format_exc()}"
@ -1828,7 +1829,7 @@ class Wav2Vec2ModelIntegrationTest(unittest.TestCase):
@require_pyctcdecode
@require_torchaudio
def test_wav2vec2_with_lm(self):
ds = load_dataset("mozilla-foundation/common_voice_11_0", "es", split="test", streaming=True)
ds = load_dataset("fixie-ai/common_voice_17_0", "es", split="test", streaming=True)
sample = next(iter(ds))
resampled_audio = torchaudio.functional.resample(
@ -1847,12 +1848,13 @@ class Wav2Vec2ModelIntegrationTest(unittest.TestCase):
transcription = processor.batch_decode(logits.cpu().numpy()).text
self.assertEqual(transcription[0], "habitan aguas poco profundas y rocosas")
expected_transcription = "el resto de los equipos se mantienen en su sede"
self.assertEqual(transcription[0], expected_transcription)
@require_pyctcdecode
@require_torchaudio
def test_wav2vec2_with_lm_pool(self):
ds = load_dataset("mozilla-foundation/common_voice_11_0", "es", split="test", streaming=True)
ds = load_dataset("fixie-ai/common_voice_17_0", "es", split="test", streaming=True)
sample = next(iter(ds))
resampled_audio = torchaudio.functional.resample(
@ -1873,7 +1875,8 @@ class Wav2Vec2ModelIntegrationTest(unittest.TestCase):
with multiprocessing.get_context("fork").Pool(2) as pool:
transcription = processor.batch_decode(logits.cpu().numpy(), pool).text
self.assertEqual(transcription[0], "habitan aguas poco profundas y rocosas")
expected_transcription = "el resto de los equipos se mantienen en su sede"
self.assertEqual(transcription[0], expected_transcription)
# user-managed pool + num_processes should trigger a warning
with (
@ -1885,7 +1888,7 @@ class Wav2Vec2ModelIntegrationTest(unittest.TestCase):
self.assertIn("num_process", cl.out)
self.assertIn("it will be ignored", cl.out)
self.assertEqual(transcription[0], "habitan aguas poco profundas y rocosas")
self.assertEqual(transcription[0], expected_transcription)
@require_pyctcdecode
@require_torchaudio
@ -1951,7 +1954,7 @@ class Wav2Vec2ModelIntegrationTest(unittest.TestCase):
LANG_MAP = {"it": "ita", "es": "spa", "fr": "fra", "en": "eng"}
def run_model(lang):
ds = load_dataset("mozilla-foundation/common_voice_11_0", lang, split="test", streaming=True)
ds = load_dataset("fixie-ai/common_voice_17_0", lang, split="test", streaming=True)
sample = next(iter(ds))
wav2vec2_lang = LANG_MAP[lang]
@ -1976,9 +1979,9 @@ class Wav2Vec2ModelIntegrationTest(unittest.TestCase):
return transcription
TRANSCRIPTIONS = {
"it": "il libro ha suscitato molte polemiche a causa dei suoi contenuti",
"es": "habitan aguas poco profundas y rocosas",
"fr": "ce dernier est volé tout au long de l'histoire romaine",
"it": "viaggiarono in italia dove lavorarono con hamilton",
"es": "el resto de los equipos se mantienen en su sede",
"fr": "il a obtenu son batchelor of lows",
"en": "joe keton disapproved of films and buster also had reservations about the media",
}