analyze.py

from itertools import count, islice
from typing import Any, Iterable, Literal, Optional, TypeVar, Union, overload, Dict, List, Tuple
from collections import defaultdict
import json

import torch

from datasets import Dataset, Audio

from dataspeech import rate_apply, pitch_apply, snr_apply
from metadata_to_text import bins_to_text, speaker_level_relative_to_gender

Row = Dict[str, Any]
T = TypeVar("T")
BATCH_SIZE = 20


@overload
def batched(it: Iterable[T], n: int) -> Iterable[List[T]]:
    ...


@overload
def batched(it: Iterable[T], n: int, with_indices: Literal[False]) -> Iterable[List[T]]:
    ...

@overload
def batched(it: Iterable[T], n: int, with_indices: Literal[True]) -> Iterable[Tuple[List[int], List[T]]]:
    ...


def batched(
    it: Iterable[T], n: int, with_indices: bool = False
) -> Union[Iterable[List[T]], Iterable[Tuple[List[int], List[T]]]]:
    it, indices = iter(it), count()
    while batch := list(islice(it, n)):
        yield (list(islice(indices, len(batch))), batch) if with_indices else batch


def analyze(
    batch: List[Dict[str, Any]],
    cache: Optional[Dict[str, List[Any]]] = None,
) -> List[List[Any]]:
    cache = {} if cache is None else cache            
    return batch


def run_dataspeech(
    rows: Iterable[Row], audio_column_name: str, text_column_name: str
) -> Iterable[Any]:
    cache: Dict[str, List[Any]] = {}

    # TODO: add speaker and gender to app
    speaker_id_column_name = "speaker_id"
    gender_column_name = "gender"

    for batch in batched(rows, BATCH_SIZE):
        tmp_dict = defaultdict(list)
        for sample in batch:
            for key in sample:
                if key in [audio_column_name, text_column_name, speaker_id_column_name, gender_column_name]:
                    tmp_dict[key].append(sample[key]) if key != audio_column_name else tmp_dict[key].append(sample[key][0]["src"])

        tmp_dataset = Dataset.from_dict(tmp_dict).cast_column(audio_column_name, Audio())
        
        
        ## 1. Extract continous tags
        pitch_dataset = tmp_dataset.map(
            pitch_apply,
            batched=True,
            batch_size=BATCH_SIZE,
            with_rank=True if torch.cuda.device_count()>0 else False,
            num_proc=torch.cuda.device_count(),
            remove_columns=[audio_column_name], # tricks to avoid rewritting audio
            fn_kwargs={"audio_column_name": audio_column_name, "penn_batch_size": 4096},
        )

        snr_dataset = tmp_dataset.map(
            snr_apply,
            batched=True,
            batch_size=BATCH_SIZE,
            with_rank=True if torch.cuda.device_count()>0 else False,
            num_proc=torch.cuda.device_count(),
            remove_columns=[audio_column_name], # tricks to avoid rewritting audio
            fn_kwargs={"audio_column_name": audio_column_name},
        )
        
        rate_dataset = tmp_dataset.map(
            rate_apply,
            with_rank=False,
            num_proc=1,
            remove_columns=[audio_column_name], # tricks to avoid rewritting audio
            fn_kwargs={"audio_column_name": audio_column_name, "text_column_name": text_column_name},
        )
    
        enriched_dataset = pitch_dataset.add_column("snr", snr_dataset["snr"]).add_column("c50", snr_dataset["c50"])
        enriched_dataset = enriched_dataset.add_column("speaking_rate", rate_dataset["speaking_rate"]).add_column("phonemes", rate_dataset["phonemes"])
        
        
        ## 2. Map continuous tags to text tags
        
        text_bins_dict = {}
        with open("./create_dataset_app/v01_text_bins.json") as json_file:
            text_bins_dict = json.load(json_file)
                
        bin_edges_dict = {}
        with open("./create_dataset_app/v01_bin_edges.json") as json_file:
            bin_edges_dict = json.load(json_file)

        speaker_level_pitch_bins = text_bins_dict.get("speaker_level_pitch_bins")
        speaker_rate_bins = text_bins_dict.get("speaker_rate_bins")
        snr_bins = text_bins_dict.get("snr_bins")
        reverberation_bins = text_bins_dict.get("reverberation_bins")
        utterance_level_std = text_bins_dict.get("utterance_level_std")
        
        enriched_dataset = [enriched_dataset]
        if "gender" in batch[0] and "speaker_id" in batch[0]:
            bin_edges = None
            if "pitch_bins_male" in bin_edges_dict and "pitch_bins_female" in bin_edges_dict:
                bin_edges = {"male": bin_edges_dict["pitch_bins_male"], "female": bin_edges_dict["pitch_bins_female"]}

            enriched_dataset, _ = speaker_level_relative_to_gender(enriched_dataset, speaker_level_pitch_bins, "speaker_id", "gender", "utterance_pitch_mean", "pitch", batch_size=20, num_workers=1, std_tolerance=None, save_dir=None, only_save_plot=False, bin_edges=bin_edges)
            
        enriched_dataset, _ = bins_to_text(enriched_dataset, speaker_rate_bins, "speaking_rate", "speaking_rate", batch_size=20, num_workers=1, leading_split_for_bins=None, std_tolerance=None, save_dir=None, only_save_plot=False, bin_edges=bin_edges_dict.get("speaking_rate",None))
        enriched_dataset, _ = bins_to_text(enriched_dataset, snr_bins, "snr", "noise", batch_size=20, num_workers=1, leading_split_for_bins=None, std_tolerance=None, save_dir=None, only_save_plot=False, bin_edges=bin_edges_dict.get("noise",None), lower_range=None)
        enriched_dataset, _ = bins_to_text(enriched_dataset, reverberation_bins, "c50", "reverberation", batch_size=20, num_workers=1, leading_split_for_bins=None, std_tolerance=None, save_dir=None, only_save_plot=False, bin_edges=bin_edges_dict.get("reverberation",None))
        enriched_dataset, _ = bins_to_text(enriched_dataset, utterance_level_std, "utterance_pitch_std", "speech_monotony", batch_size=20, num_workers=1, leading_split_for_bins=None, std_tolerance=None, save_dir=None, only_save_plot=False, bin_edges=bin_edges_dict.get("speech_monotony",None))


        enriched_dataset = enriched_dataset[0]
    
        for i,sample in enumerate(batch):
            new_sample = {}
            new_sample[audio_column_name] = f"<audio src='{sample[audio_column_name][0]['src']}' controls></audio>"
            for col in ["speaking_rate", "reverberation", "noise", "speech_monotony", "c50", "snr",]: # phonemes, speaking_rate, utterance_pitch_std, utterance_pitch_mean
                new_sample[col] = enriched_dataset[col][i]
            if "gender" in batch[0] and "speaker_id" in batch[0]:
                new_sample["pitch"] = enriched_dataset["pitch"][i]
                new_sample[gender_column_name] = sample[col]
                new_sample[speaker_id_column_name] = sample[col]
            
            new_sample[text_column_name] = sample[text_column_name]   
            batch[i] = new_sample
            
        yield analyze(
            batch=batch,
            cache=cache,
        )