Whisper v3

whisperv3.py

import os
import csv
import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
from tqdm import tqdm


def transcribe_audio_files(input_dir, output_csv):
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
    model_id = "openai/whisper-large-v3"

    model = AutoModelForSpeechSeq2Seq.from_pretrained(
        model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
    )
    model.to(device)
    processor = AutoProcessor.from_pretrained(model_id)

    pipe = pipeline(
        "automatic-speech-recognition",
        model=model,
        tokenizer=processor.tokenizer,
        feature_extractor=processor.feature_extractor,
        torch_dtype=torch_dtype,
        device=device,
        generate_kwargs={"language": "<|vi|>", "task": "transcribe"},
    )

    audio_files = [
        f
        for f in os.listdir(input_dir)
        if f.endswith((".mp3", ".wav", ".flac", ".ogg"))
    ]

    with open(output_csv, "w", newline="", encoding="utf-8") as csvfile:
        csvwriter = csv.writer(csvfile)
        csvwriter.writerow(["File", "Transcription"])

        for audio_file in tqdm(audio_files, desc="Transcribing files"):
            file_path = os.path.join(input_dir, audio_file)
            try:
                result = pipe(file_path)
                transcription = result["text"]
                csvwriter.writerow([file_path, transcription])
            except Exception as e:
                print(f"Error processing {file_path}: {str(e)}")
                csvwriter.writerow([file_path, f"Error: {str(e)}"])

    print(f"Transcription complete. Results saved to {output_csv}")


if __name__ == "__main__":
    input_directory = "tv360/"
    output_csv_file = "whisper-l-v3.csv"
    transcribe_audio_files(input_directory, output_csv_file)