Interview transcription using WhisperX model, Part 1.

device = "cuda"
batch_size = 2
compute_type = "int8" # "float32"
model_dir = "./model/"
HF_TOKEN = os.getenv("HUGGINGFACE_TOKEN")

print(f"{elapsed()} 1. Loading model and audio")

model = whisperx.load_model("turbo", device, compute_type=compute_type, download_root=model_dir)
audio = whisperx.load_audio(audio_file)

print(f"{elapsed()} 2. Load Whisper model and transcribe")
transcription_results = model.transcribe(audio, batch_size=batch_size)

print(f"{elapsed()} 3. Align transcription")

model_a, metadata = whisperx.load_align_model(language_code=transcription_results["language"], device=device)
alignment_results = whisperx.align(transcription_results["segments"], model_a, metadata, audio, device, return_char_alignments=False)

print(f"{elapsed()} 4. Speaker diarization")
diarize_model = whisperx.DiarizationPipeline(use_auth_token=HF_TOKEN, device=device)
diarization_results = diarize_model(audio)

print(f"{elapsed()} 5. Assigning words to speakers")
result = whisperx.assign_word_speakers(diarization_results, alignment_results)
pprint.pprint(result["segments"][:4], depth=2)

print(f"{elapsed()} 6. Process and format the output")

structured_output = []
current_speaker = None
current_start = None
current_end = None
current_text = []

def format_time(seconds):
    return str(timedelta(seconds=int(seconds)))

for segment in result["segments"]:

    speaker = segment.get("speaker", "UNKNOWN_SPEAKER")
    start = segment.get("start", "unknown")
    end = segment.get("end", "unknown")
    text = segment.get("text", "")

    if speaker != current_speaker:
        if current_speaker is not None:
            structured_output.append([current_speaker, format_time(current_start), format_time(current_end), " ".join(current_text)])

        current_speaker = speaker
        current_start = start
        current_text = [text]
    else:
        current_text.append(text)

    current_end = end

if current_speaker is not None:
    structured_output.append([current_speaker, format_time(current_start), format_time(current_end), " ".join(current_text)])

print(f"{elapsed()} 7. Save to CSV")

with open(final_output_path, mode="w", newline="", encoding="utf-8") as file:
    writer = csv.writer(file)
    writer.writerow(["Speaker", "Start", "End", "Speech"])
    writer.writerows(structured_output)

print(f"{elapsed()} Transcription saved to {final_output_path}")