more return_timestamps=True

ylacombe · Sep 12, 2024 · a1217e8 · a1217e8
1 parent 1fd3d96
commit a1217e8
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/models/whisper/test_modeling_whisper.py b/tests/models/whisper/test_modeling_whisper.py
@@ -2385,7 +2385,7 @@ def test_tiny_token_timestamp_generation_longform(self):
         )
 
         inputs = inputs.to(torch_device)
-        generate_outputs = model.generate(**inputs, return_segments=True, return_token_timestamps=True)
+        generate_outputs = model.generate(**inputs, return_segments=True, return_token_timestamps=True, return_timestamps=True)
 
         token_timestamps_shape = [
             [segment["token_timestamps"].shape for segment in segment_list]
@@ -3264,6 +3264,7 @@ def test_whisper_empty_longform(self):
             "num_beams": 5,
             "language": "fr",
             "task": "transcribe",
+            "return_timestamps": True,
         }
 
         torch.manual_seed(0)