lib/model.py

#!/usr/bin/python
# -*- coding: utf-8 -*-

import gc
from time import time
import torch
import whisperx as wx

from .config import (
    DEVICE,
    COMPUTE_TYPE,
    BATCH_SIZE,
)
# -->> Tunables <<---------------------


# -->> Definitions <<------------------


# -->> API <<--------------------------


def transcribe_audio(audio_file, audio_path, transcript_folder_path):
    # Transcribe the audio
    print("Starting transcription...")
    print("Loading model...")
    time_1 = time()
    model = wx.load_model(
        "large-v2", device=DEVICE, compute_type=COMPUTE_TYPE, language="en"
    )
    time_2 = time()
    print("Loading audio...")
    time_3 = time()
    audio = wx.load_audio(audio_path)
    time_4 = time()
    print("Transcribing...")
    time_5 = time()
    result = model.transcribe(audio, batch_size=BATCH_SIZE)
    time_6 = time()
    print("Transcription complete!")

    print("\nTime Report:   ")
    print("Loading model:   ", round(time_2 - time_1, 2), " [s]")
    print("Loading audio:   ", round(time_4 - time_3, 2), " [s]")
    print("Transcribing:    ", round(time_6 - time_5, 2), " [s]")
    print("Total:           ", round(time_6 - time_1, 2), " [s]")

    # Save the transcript to a file
    text = "\n ".join([i["text"] for i in result["segments"]])

    # Free memory
    gc.collect()
    torch.cuda.empty_cache()
    del model

    return text


# -->> Execute <<----------------------


# -->> Export <<-----------------------