j_apply_vad_to_speech_profiles.py

import os
import threading

from typing import Any, cast


from dotenv import load_dotenv
from pydub import AudioSegment

from utils.stt.vad import apply_vad_for_speech_profile, VADEmptyError

load_dotenv('../../.env')
os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = '../../' + os.getenv('GOOGLE_APPLICATION_CREDENTIALS', '')

from database._client import get_users_uid
from utils.other.storage import get_profile_audio_if_exists, upload_profile_audio


def execute():
    os.makedirs('_temp', exist_ok=True)
    uids = get_users_uid()
    print('execute for', len(uids))

    def single(uid: str) -> None:
        file_path = get_profile_audio_if_exists(uid)
        if not file_path:
            return

        try:
            apply_vad_for_speech_profile(file_path)
        except VADEmptyError:
            print('VAD empty for', uid)
            return

        aseg = cast(
            Any, AudioSegment.from_wav(file_path)
        )  # pyright: ignore[reportUnknownMemberType]  # pydub has no type stubs
        if aseg.duration_seconds < 5 or aseg.duration_seconds > 180:
            print('Invalid duration for', uid)
            return
        upload_profile_audio(file_path, uid)

    threads: list[threading.Thread] = []
    for i, uid in enumerate(uids):
        threads.append(threading.Thread(target=single, args=(uid,)))

    count = 20
    chunks = [threads[i : i + count] for i in range(0, len(threads), count)]
    for i, chunk in enumerate(chunks):
        [thread.start() for thread in chunk]
        [thread.join() for thread in chunk]
        print('Chunk', i + 1, 'completed')