Spaces:
Build error
Build error
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 | |
def batched(it: Iterable[T], n: int) -> Iterable[List[T]]: | |
... | |
def batched(it: Iterable[T], n: int, with_indices: Literal[False]) -> Iterable[List[T]]: | |
... | |
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("./v01_text_bins.json") as json_file: | |
text_bins_dict = json.load(json_file) | |
bin_edges_dict = {} | |
with open("./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, | |
) | |