372 lines
9.4 KiB
Python
372 lines
9.4 KiB
Python
"""
|
|
# Handles transcribing audio provided through --input-audio
|
|
"""
|
|
|
|
import os
|
|
import json
|
|
import argparse
|
|
|
|
import torch
|
|
import torchaudio
|
|
|
|
"""
|
|
try:
|
|
import whisperx
|
|
except Exception as e:
|
|
whisperx = None
|
|
print(f"Error while querying for whisperx: {str(e)}")
|
|
pass
|
|
"""
|
|
|
|
from transformers import pipeline
|
|
|
|
from functools import cache
|
|
from tqdm.auto import tqdm
|
|
from pathlib import Path
|
|
|
|
from ..utils import coerce_dtype
|
|
|
|
def pad(num, zeroes):
|
|
return str(num).zfill(zeroes+1)
|
|
|
|
def process_items( items, stride=0, stride_offset=0 ):
|
|
items = sorted( items )
|
|
return items if stride == 0 else [ item for i, item in enumerate( items ) if (i+stride_offset) % stride == 0 ]
|
|
|
|
# major cringe but should automatically unload models when loading a different one
|
|
_cached_models = {
|
|
"model": (None, None),
|
|
"diarization": (None, None),
|
|
"align": (None, None),
|
|
}
|
|
# yes I can write a decorator to do this
|
|
def _load_model(model_name="openai/whisper-large-v3", device="cuda", dtype="float16", language="auto", backend="auto", attention="sdpa"):
|
|
cache_key = f'{model_name}:{device}:{dtype}:{language}'
|
|
if _cached_models["model"][0] == cache_key:
|
|
return _cached_models["model"][1]
|
|
|
|
del _cached_models["model"]
|
|
|
|
if not isinstance( dtype, str ):
|
|
if dtype == torch.float32:
|
|
dtype = "float32"
|
|
elif dtype == torch.float16:
|
|
dtype = "float16"
|
|
elif dtype == torch.bfloat16:
|
|
dtype = "bfloat16"
|
|
|
|
# doesnt support it for some reason
|
|
if dtype == "bfloat16":
|
|
dtype = "float16"
|
|
|
|
kwargs = {}
|
|
kwargs["compute_type"] = dtype
|
|
kwargs["task"] = "transcribe"
|
|
kwargs["device"] = device
|
|
|
|
if language != "auto":
|
|
kwargs["language"] = language
|
|
|
|
if backend == "auto" and whisperx is not None:
|
|
backend = "whisperx"
|
|
|
|
if backend == "whisperx":
|
|
model_name = model_name.replace("openai/whisper-", "")
|
|
model = whisperx.load_model(model_name, **kwargs)
|
|
else:
|
|
model = pipeline(
|
|
"automatic-speech-recognition",
|
|
model=model_name,
|
|
torch_dtype=coerce_dtype(dtype),
|
|
device=device,
|
|
model_kwargs={"attn_implementation": attention},
|
|
)
|
|
|
|
_cached_models["model"] = (cache_key, model)
|
|
return model
|
|
|
|
def _load_diarization_model(device="cuda", backend="auto"):
|
|
cache_key = f'{device}'
|
|
|
|
if _cached_models["diarization"][0] == cache_key:
|
|
return _cached_models["diarization"][1]
|
|
del _cached_models["diarization"]
|
|
|
|
if backend == "auto" and whisperx is not None:
|
|
backend = "whisperx"
|
|
|
|
if backend == "whisperx":
|
|
model = whisperx.DiarizationPipeline(device=device)
|
|
else:
|
|
model = None # to do later
|
|
|
|
_cached_models["diarization"] = (cache_key, model)
|
|
return model
|
|
|
|
def _load_align_model(language, device="cuda", backend="auto"):
|
|
cache_key = f'{language}:{device}'
|
|
|
|
if _cached_models["align"][0] == cache_key:
|
|
return _cached_models["align"][1]
|
|
del _cached_models["align"]
|
|
|
|
if backend == "auto" and whisperx is not None:
|
|
backend = "whisperx"
|
|
|
|
if backend == "whisperx":
|
|
model = whisperx.load_align_model(language_code=language, device=device)
|
|
else:
|
|
model = None # to do later
|
|
|
|
_cached_models["align"] = (cache_key, model)
|
|
return model
|
|
|
|
# yes I can just do a for-loop
|
|
def unload_model():
|
|
del _cached_models["model"]
|
|
del _cached_models["diarization"]
|
|
del _cached_models["align"]
|
|
|
|
_cached_models["model"] = (None, None)
|
|
_cached_models["diarization"] = (None, None)
|
|
_cached_models["align"] = (None, None)
|
|
|
|
def transcribe(
|
|
audio,
|
|
language = "auto",
|
|
diarize = False,
|
|
batch_size = 16,
|
|
verbose=False,
|
|
align=True,
|
|
**model_kwargs,
|
|
):
|
|
metadata = {
|
|
"segments": [],
|
|
"language": "",
|
|
"text": "",
|
|
"start": 0,
|
|
"end": 0,
|
|
}
|
|
|
|
# load requested models
|
|
model_kwargs["backend"] = "automatic-speech-recognition"
|
|
device = model_kwargs.get("device", "cuda")
|
|
model = _load_model(language=language, **model_kwargs)
|
|
|
|
result = model(
|
|
str(audio),
|
|
chunk_length_s=30,
|
|
batch_size=batch_size,
|
|
generate_kwargs={"task": "transcribe", "language": None if language == "auto" else language},
|
|
return_timestamps="word" if align else False,
|
|
return_language=True,
|
|
)
|
|
|
|
start = 0
|
|
end = 0
|
|
segments = []
|
|
for segment in result["chunks"]:
|
|
text = segment["text"]
|
|
|
|
if "timestamp" in segment:
|
|
s, e = segment["timestamp"]
|
|
start = min( start, s )
|
|
end = max( end, e )
|
|
else:
|
|
s, e = None, None
|
|
|
|
if language == "auto":
|
|
language = segment["language"]
|
|
|
|
segments.append({
|
|
"start": s,
|
|
"end": e,
|
|
"text": text,
|
|
})
|
|
|
|
if language != "auto":
|
|
metadata["language"] = language
|
|
|
|
metadata["segments"] = segments
|
|
metadata["text"] = result["text"].strip()
|
|
metadata["start"] = start
|
|
metadata["end"] = end
|
|
|
|
return metadata
|
|
|
|
# for backwards compat since it also handles some other things for me
|
|
"""
|
|
def transcribe_whisperx(
|
|
audio,
|
|
language = "auto",
|
|
diarize = False,
|
|
batch_size = 16,
|
|
verbose=False,
|
|
align=True,
|
|
**model_kwargs,
|
|
):
|
|
metadata = {
|
|
"segments": [],
|
|
"language": "",
|
|
"text": "",
|
|
"start": 0,
|
|
"end": 0,
|
|
}
|
|
|
|
# load requested models
|
|
device = model_kwargs.get("device", "cuda")
|
|
model = _load_model(language=language, **model_kwargs)
|
|
diarize_model = _load_diarization_model(device=device) if diarize else None
|
|
|
|
# audio is a path, load it
|
|
if isinstance(audio, str) or isinstance(audio, Path):
|
|
#audio = load_audio(audio)
|
|
audio = whisperx.load_audio(audio)
|
|
|
|
result = model.transcribe(audio, batch_size=batch_size)
|
|
|
|
if language == "auto":
|
|
language = result["language"]
|
|
|
|
if align:
|
|
align_model, align_model_metadata = _load_align_model(language=language, device=device)
|
|
result = whisperx.align(result["segments"], align_model, align_model_metadata, audio, device, return_char_alignments=False)
|
|
|
|
if diarize_model is not None:
|
|
diarize_segments = diarize_model(audio)
|
|
result = whisperx.assign_word_speakers(diarize_segments, result)
|
|
|
|
text = []
|
|
start = 0
|
|
end = 0
|
|
for segment in result["segments"]:
|
|
text.append( segment["text"] )
|
|
start = min( start, segment["start"] )
|
|
end = max( end, segment["end"] )
|
|
|
|
metadata["language"] = language
|
|
metadata["segments"] = result["segments"]
|
|
metadata["text"] = " ".join(text).strip()
|
|
metadata["start"] = start
|
|
metadata["end"] = end
|
|
|
|
return metadata
|
|
"""
|
|
|
|
def transcribe_batch(
|
|
input_audio = "voices",
|
|
input_voice = None,
|
|
output_metadata = "training/metadata",
|
|
model_name = "openai/whisper-large-v3",
|
|
|
|
skip_existing = True,
|
|
diarize = False,
|
|
|
|
stride = 0,
|
|
stride_offset = 0,
|
|
|
|
batch_size = 16,
|
|
device = "cuda",
|
|
dtype = "float16",
|
|
):
|
|
# to-do: make this also prepared from args
|
|
language_map = {} # k = group, v = language
|
|
|
|
ignore_groups = [] # skip these groups
|
|
ignore_speakers = [] # skip these speakers
|
|
|
|
only_groups = [] # only process these groups
|
|
only_speakers = [] # only process these speakers
|
|
|
|
if input_voice is not None:
|
|
only_speakers = [input_voice]
|
|
|
|
for dataset_name in os.listdir(f'./{input_audio}/'):
|
|
if not os.path.isdir(f'./{input_audio}/{dataset_name}/'):
|
|
continue
|
|
|
|
if group_name in ignore_groups:
|
|
continue
|
|
if only_groups and group_name not in only_groups:
|
|
continue
|
|
|
|
for speaker_id in tqdm(process_items(os.listdir(f'./{input_audio}/{dataset_name}/')), desc="Processing speaker"):
|
|
if not os.path.isdir(f'./{input_audio}/{dataset_name}/{speaker_id}'):
|
|
continue
|
|
|
|
if speaker_id in ignore_speakers:
|
|
continue
|
|
if only_speakers and speaker_id not in only_speakers:
|
|
continue
|
|
|
|
outpath = Path(f'./{output_metadata}/{dataset_name}/{speaker_id}/whisper.json')
|
|
|
|
if outpath.exists():
|
|
metadata = json.loads(open(outpath, 'r', encoding='utf-8').read())
|
|
else:
|
|
os.makedirs(f'./{output_metadata}/{dataset_name}/{speaker_id}/', exist_ok=True)
|
|
metadata = {}
|
|
|
|
for filename in tqdm(os.listdir(f'./{input_audio}/{dataset_name}/{speaker_id}/'), desc=f"Processing speaker: {speaker_id}"):
|
|
if skip_existing and filename in metadata:
|
|
continue
|
|
|
|
if ".json" in filename:
|
|
continue
|
|
|
|
inpath = f'./{input_audio}/{dataset_name}/{speaker_id}/{filename}'
|
|
|
|
if os.path.isdir(inpath):
|
|
continue
|
|
|
|
metadata[filename] = transcribe( inpath, model_name=model_name, diarize=diarize, device=device, dtype=dtype )
|
|
|
|
open(outpath, 'w', encoding='utf-8').write(json.dumps(metadata))
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
|
|
parser.add_argument("--input-audio", type=str, default="voices")
|
|
parser.add_argument("--input-voice", type=str, default=None)
|
|
parser.add_argument("--output-metadata", type=str, default="training/metadata")
|
|
|
|
parser.add_argument("--model-name", type=str, default="openai/whisper-large-v3")
|
|
parser.add_argument("--skip-existing", action="store_true")
|
|
parser.add_argument("--diarize", action="store_true")
|
|
parser.add_argument("--batch-size", type=int, default=16)
|
|
parser.add_argument("--stride", type=int, default=0)
|
|
parser.add_argument("--stride-offset", type=int, default=0)
|
|
|
|
parser.add_argument("--device", type=str, default="cuda")
|
|
parser.add_argument("--dtype", type=str, default="bfloat16")
|
|
parser.add_argument("--amp", action="store_true")
|
|
# parser.add_argument("--raise-exceptions", action="store_true")
|
|
|
|
args = parser.parse_args()
|
|
|
|
# do some assumption magic
|
|
# to-do: find a nice way to spawn multiple processes where tqdm plays nicely
|
|
if args.device.isnumeric():
|
|
args.stride = torch.cuda.device_count()
|
|
args.stride_offset = int(args.device)
|
|
args.device = f'cuda:{args.device}'
|
|
|
|
transcribe_batch(
|
|
input_audio = args.input_audio,
|
|
input_voice = args.input_voice,
|
|
output_metadata = args.output_metadata,
|
|
model_name = args.model_name,
|
|
|
|
skip_existing = args.skip_existing,
|
|
diarize = args.diarize,
|
|
|
|
stride = args.stride,
|
|
stride_offset = args.stride_offset,
|
|
|
|
batch_size = args.batch_size,
|
|
device = args.device,
|
|
dtype = args.dtype,
|
|
)
|
|
|
|
if __name__ == "__main__":
|
|
main() |