jhj0517 commited on
Commit
0e9733d
·
1 Parent(s): 5f8203e

Add diarization parameter

Browse files
Files changed (1) hide show
  1. tests/test_transcription.py +7 -5
tests/test_transcription.py CHANGED
@@ -10,17 +10,18 @@ import os
10
 
11
 
12
  @pytest.mark.parametrize(
13
- "whisper_type,vad_filter,bgm_separation",
14
  [
15
- ("whisper", False, False),
16
- ("faster-whisper", False, False),
17
- ("insanely_fast_whisper", False, False)
18
  ]
19
  )
20
  def test_transcribe(
21
  whisper_type: str,
22
  vad_filter: bool,
23
  bgm_separation: bool,
 
24
  ):
25
  audio_path_dir = os.path.join(WEBUI_DIR, "tests")
26
  audio_path = os.path.join(audio_path_dir, "jfk.wav")
@@ -36,7 +37,8 @@ def test_transcribe(
36
  model_size=TEST_WHISPER_MODEL,
37
  vad_filter=vad_filter,
38
  is_bgm_separate=bgm_separation,
39
- uvr_enable_offload=True
 
40
  ).as_list()
41
 
42
  subtitle_str, file_path = whisper_inferencer.transcribe_file(
 
10
 
11
 
12
  @pytest.mark.parametrize(
13
+ "whisper_type,vad_filter,bgm_separation,diarization",
14
  [
15
+ ("whisper", False, False, False),
16
+ ("faster-whisper", False, False, False),
17
+ ("insanely_fast_whisper", False, False, False)
18
  ]
19
  )
20
  def test_transcribe(
21
  whisper_type: str,
22
  vad_filter: bool,
23
  bgm_separation: bool,
24
+ diarization: bool,
25
  ):
26
  audio_path_dir = os.path.join(WEBUI_DIR, "tests")
27
  audio_path = os.path.join(audio_path_dir, "jfk.wav")
 
37
  model_size=TEST_WHISPER_MODEL,
38
  vad_filter=vad_filter,
39
  is_bgm_separate=bgm_separation,
40
+ uvr_enable_offload=True,
41
+ is_diarize=diarization,
42
  ).as_list()
43
 
44
  subtitle_str, file_path = whisper_inferencer.transcribe_file(