Update analyze.py

analyze.py

@@ -2,6 +2,7 @@ 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
 
@@ -36,113 +37,117 @@ def batched(
     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            
-    return batch
-
-
-def run_dataspeech(
-    rows: Iterable[Row], audio_column_name: str, text_column_name: str
-) -> Iterable[Any]:
-    cache: Dict[str, 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"
 
-    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},
-        )
+    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"])
 
-        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},
-        )
+    tmp_dataset = Dataset.from_dict(tmp_dict).cast_column(audio_column_name, Audio())
     
-        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")
+    
+    ## 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 = [enriched_dataset]
+    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]:
-            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"]}
+            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
 
-            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))
 
+def run_dataspeech(
+    rows: Iterable[Row], audio_column_name: str, text_column_name: str
+) -> Iterable[Any]:
+    cache: Dict[str, List[Any]] = {}
 
-        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
-            
+
+    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,
         )