Spaces:
Sleeping
Sleeping
File size: 7,282 Bytes
db36668 0271b59 db36668 8496047 db36668 0271b59 db36668 0271b59 db36668 0271b59 db36668 0271b59 db36668 0271b59 db36668 0271b59 8496047 0271b59 8496047 0271b59 db36668 0271b59 8496047 0271b59 db36668 0271b59 db36668 0271b59 db36668 0271b59 db36668 0271b59 db36668 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 |
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 spaces
import torch
from datasets import Dataset, Audio
from dataspeech import rate_apply, pitch_apply, snr_apply, squim_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
@spaces.GPU(duration=60)
def analyze(
batch: List[Dict[str, Any]],
audio_column_name: str, text_column_name: str,
cache: Optional[Dict[str, List[Any]]] = None,
) -> List[List[Any]]:
cache = {} if cache is None else cache
# TODO: add speaker and gender to app
speaker_id_column_name = "speaker_id"
gender_column_name = "gender"
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
squim_dataset = tmp_dataset.map(
squim_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,},
)
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"])
enriched_dataset = enriched_dataset.add_column("stoi", squim_dataset["stoi"]).add_column("si-sdr", squim_dataset["sdr"]).add_column("pesq", squim_dataset["pesq"])
## 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", "stoi", "pesq", "si-sdr"]: # 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
return batch
def run_dataspeech(
rows: Iterable[Row], audio_column_name: str, text_column_name: str
) -> Iterable[Any]:
cache: Dict[str, List[Any]] = {}
for batch in batched(rows, BATCH_SIZE):
yield analyze(
batch=batch,
audio_column_name=audio_column_name,
text_column_name=text_column_name,
cache=cache,
)
|