Skip to content

Commit

Permalink
Add w2v2bert to pipeline (#28585)
Browse files Browse the repository at this point in the history
* generalize asr pipeline to fbank models

* change w2v2 pipeline output

* Update test_pipelines_automatic_speech_recognition.py
  • Loading branch information
ylacombe authored Jan 19, 2024
1 parent b2748a6 commit 268fc1f
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 2 deletions.
7 changes: 5 additions & 2 deletions src/transformers/pipelines/automatic_speech_recognition.py
Original file line number Diff line number Diff line change
Expand Up @@ -517,8 +517,11 @@ def _forward(self, model_inputs, return_timestamps=False, generate_kwargs=None):
out["stride"] = stride

else:
input_values = model_inputs.pop("input_values")
outputs = self.model(input_values=input_values, attention_mask=attention_mask)
inputs = {
self.model.main_input_name: model_inputs.pop(self.model.main_input_name),
"attention_mask": attention_mask,
}
outputs = self.model(**inputs)
logits = outputs.logits

if self.type == "ctc_with_lm":
Expand Down
17 changes: 17 additions & 0 deletions tests/pipelines/test_pipelines_automatic_speech_recognition.py
Original file line number Diff line number Diff line change
Expand Up @@ -298,6 +298,23 @@ def test_torch_large(self):
output = speech_recognizer(filename)
self.assertEqual(output, {"text": "A MAN SAID TO THE UNIVERSE SIR I EXIST"})

@require_torch
@slow
def test_torch_large_with_input_features(self):
speech_recognizer = pipeline(
task="automatic-speech-recognition",
model="hf-audio/wav2vec2-bert-CV16-en",
framework="pt",
)
waveform = np.tile(np.arange(1000, dtype=np.float32), 34)
output = speech_recognizer(waveform)
self.assertEqual(output, {"text": ""})

ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation").sort("id")
filename = ds[40]["file"]
output = speech_recognizer(filename)
self.assertEqual(output, {"text": "a man said to the universe sir i exist"})

@slow
@require_torch
@slow
Expand Down

0 comments on commit 268fc1f

Please sign in to comment.