add: note

note.txt

@@ -0,0 +1,2 @@
+- tạo range cho các item trong dataset
+- chạy parallel 

phoneme_eval.py

@@ -2,7 +2,7 @@ import pandas as pd
 from utils.load_model import run_hubert_base, run_whisper, run_model, run_timit, run_wavlm_large_phoneme, run_gruut
 from utils.audio_process import calculate_error_rate, load_audio
 from utils.cmu_process import clean_cmu, cmu_to_ipa, text_to_phoneme
-from constants import DATASETS
+from constants import DATASETS, FINAL_SIZE
 from datasets import load_dataset, Audio
 import argparse
 
@@ -12,8 +12,8 @@ MODEL_RUNNERS = {
     "Whisper": run_whisper,
     "HuBERT fine-tuned": run_model,
     "Timit": run_timit,
-    "speech31/wavlm-large-english-phoneme": run_wavlm_large_phoneme,
-    "bookbot/wav2vec2-ljspeech-gruut": run_gruut,
+    "WavLM": run_wavlm_large_phoneme,
+    "LJSpeech Gruut": run_gruut,
 }
 
 def set_output(model, pre_pho, ref_pho, duration, per, score):
@@ -42,23 +42,32 @@ def get_output(model, wav, reference_phoneme):
 
 def benchmark_all(example):
     """
-    Run all models on a single dataset example.
+    Run all models on a single dataset example in parallel.
     """
     # Load waveform manually to avoid datasets' torchcodec dependency
-    wav = load_audio(example["audio"]) 
-    reference_phoneme = example["phonetic"]
+    wav = load_audio(example["audio"])
+    reference_phoneme = example["phonetic"] 
     reference_phoneme = cmu_to_ipa(clean_cmu(reference_phoneme))
 
-    # Run all models
-    results = [
-        get_output("HuBERT-Base", wav, reference_phoneme),
-        get_output("Whisper", wav, reference_phoneme),
-        get_output("HuBERT fine-tuned", wav, reference_phoneme),
-        get_output("Timit", wav, reference_phoneme),
-        get_output("WavLM", wav, reference_phoneme),
-        get_output("LJSpeech Gruut", wav, reference_phoneme),
+    # Run all models in parallel using ThreadPoolExecutor
+    from concurrent.futures import ThreadPoolExecutor
+    
+    models = [
+        "HuBERT-Base",
+        "Whisper", 
+        "HuBERT fine-tuned",
+        "Timit",
+        "WavLM",
+        "LJSpeech Gruut"
     ]
 
+    with ThreadPoolExecutor(max_workers=len(models)) as executor:
+        futures = [
+            executor.submit(get_output, model, wav, reference_phoneme)
+            for model in models
+        ]
+        results = [future.result() for future in futures]
+
     return pd.DataFrame(results)
 
 def benchmark_dataset(dataset):
@@ -127,124 +136,112 @@ def load_dataset_with_limits(dataset_config, max_samples=None, use_streaming=Fal
         print(f"[warn] skip dataset {dataset_config['name']}: {e}")
         return None
 
-def main():
-    # Parse command line arguments
+def parse_cli_args():
+    """
+    Parse and return CLI arguments for the evaluation script.
+    """
     parser = argparse.ArgumentParser(description='Phoneme Detection Evaluation')
     parser.add_argument('--max-samples', type=int, default=None, 
                          help='Override max_samples for all datasets')
     parser.add_argument('--dataset', type=str, default=None,
                          help='Process only specific dataset (by name)')
-    args = parser.parse_args()
-    
-    per_model_results = {}
+    return parser.parse_args()
 
-    for dataset_config in DATASETS:
-        # Skip dataset if specific dataset is requested and this isn't it
-        if args.dataset and args.dataset not in dataset_config["name"]:
-            continue
-            
-        # Override max_samples if provided via command line
-        max_samples = args.max_samples if args.max_samples is not None else dataset_config.get("max_samples")
-        use_streaming = dataset_config.get("use_streaming", False)
-        
-        # Load dataset with limits
-        dataset = load_dataset_with_limits(
-            dataset_config, 
-            max_samples=max_samples,
-            use_streaming=use_streaming
-        )
-        
-        if dataset is None:
-            continue
+def cast_audio_column_safely(dataset):
+    """
+    Ensure the dataset's 'audio' column is set to non-decoding Audio.
+    """
+    try:
+        dataset = dataset.cast_column("audio", Audio(decode=False))
+    except Exception:
+        pass
+    return dataset
 
-        try:
-            dataset = dataset.cast_column("audio", Audio(decode=False))
-        except Exception:
-            pass
+def prepare_dataset_for_evaluation(dataset, dataset_config, max_samples):
+    """
+    Normalize, deduplicate, and filter dataset examples for evaluation.
+    Handles both streaming and non-streaming datasets.
+    Returns a finalized small dataset suitable for benchmarking.
+    """
+    field = dataset_config["field"]
+    use_streaming = dataset_config.get("use_streaming", False)
 
-        field = dataset_config["field"]
-        
-        # Handle streaming datasets differently
-        if use_streaming:
-            print("Processing streaming dataset...")
-            valid_samples = []
-            
-            # Set a reasonable limit for streaming (max 100 samples)
-            streaming_limit = max(max_samples or 100, 100)
-            
-            for example in dataset:
-                # Convert text to phonemes if needed
-                if field == "text":
-                    phonetic_text = text_to_phoneme(example[field])
-                    example = {**example, "phonetic": phonetic_text}
-                    current_field = "phonetic"
-                else:
-                    current_field = field
-                
-                # Check if valid
-                if current_field in example:
-                    phoneme_tokens = example[current_field].split()
-                    if len(phoneme_tokens) >= 10:
-                        valid_samples.append(example)
-                        # Stop when we reach the streaming limit
-                        if len(valid_samples) >= streaming_limit:
-                            break
-            
-            print(f"Found {len(valid_samples)} valid samples")
-            if len(valid_samples) == 0:
-                print("No valid samples found, skipping dataset")
-                continue
-                
-            # Convert to regular dataset for processing
-            from datasets import Dataset
-            dataset_final = Dataset.from_list(valid_samples)
-            field = "phonetic" if field == "text" else field
-        else:
-            # Regular dataset processing
+    if use_streaming:
+        print("Processing streaming dataset...")
+        valid_samples = []
+
+        streaming_limit = min(max_samples, FINAL_SIZE)
+
+        for example in dataset:
             if field == "text":
-                dataset = dataset.map(lambda x: {"phonetic": text_to_phoneme(x[field])})
-                field = "phonetic"
+                phonetic_text = text_to_phoneme(example[field])
+                example = {**example, "phonetic": phonetic_text}
+                current_field = "phonetic"
+            else:
+                current_field = field
 
-            unique_texts = dataset.unique(field)
-            print("Unique phonetic strings (", dataset_config["name"], "):", len(unique_texts))
+            if current_field in example:
+                phoneme_tokens = example[current_field].split()
+                if len(phoneme_tokens) >= 10:
+                    valid_samples.append(example)
+                    if len(valid_samples) >= streaming_limit:
+                        break
 
-            dataset_unique = dataset.filter(lambda x: x[field] in unique_texts)
+        print(f"Found {len(valid_samples)} valid samples")
+        if len(valid_samples) == 0:
+            print("No valid samples found, skipping dataset")
+            return None
 
-            def is_valid(example):
-                phoneme_tokens = example[field].split()
-                return len(phoneme_tokens) >= 10
+        from datasets import Dataset
+        dataset_final = Dataset.from_list(valid_samples)
+        return dataset_final
+    else:
+        if field == "text":
+            dataset = dataset.map(lambda x: {"phonetic": text_to_phoneme(x[field])})
+            field = "phonetic"
 
-            dataset_filtered = dataset_unique.filter(is_valid)
-            # Use smaller final size for evaluation
-            final_size = min(100, len(dataset_filtered))
-            dataset_final = dataset_filtered.shuffle(seed=42).select(range(final_size))
+        unique_texts = dataset.unique(field)
+        print("Unique phonetic strings (", dataset_config["name"], "):", len(unique_texts))
 
-        print(dataset_final)
-        print("Final size:", len(dataset_final))
+        dataset_unique = dataset.filter(lambda x: x[field] in unique_texts)
 
-        # Limit to 10 samples for benchmarking
-        benchmark_size = min(10, len(dataset_final))
-        full_results, avg_stats = benchmark_dataset(dataset_final.select(range(benchmark_size)))
-        print("Average Statistic per model (", dataset_config["name"], "):")
-        print(avg_stats)
+        def is_valid(example):
+            phoneme_tokens = example[field].split()
+            return len(phoneme_tokens) >= 10
 
-        # Use dataset name as key (extract the actual name part)
-        dataset_key = dataset_config["name"].split("/")[-1]  # Get the last part after the slash
-        for _, row in avg_stats.iterrows():
-            model_name = str(row["model"]).replace(" ", "-")
-            per = float(row["Average PER"]) if row["Average PER"] is not None else None
-            avg_dur = float(row["Average Duration (s)"]) if row["Average Duration (s)"] is not None else None
+        dataset_filtered = dataset_unique.filter(is_valid)
+        final_size = min(FINAL_SIZE, len(dataset_filtered))
+        dataset_final = dataset_filtered.shuffle(seed=42).select(range(final_size))
+        return dataset_final
 
-            if model_name not in per_model_results:
-                per_model_results[model_name] = {}
-            per_model_results[model_name][dataset_key] = {"per": per, "avg_duration": avg_dur}
+def evaluate_dataset(dataset_final):
+    """
+    Run benchmarking on a capped subset of the dataset and return both
+    the full per-example results and the aggregated stats per model.
+    """
+    benchmark_size = min(FINAL_SIZE, len(dataset_final))
+    return benchmark_dataset(dataset_final.select(range(benchmark_size)))
 
-    # Save results for leaderboard consumption (one JSON per model)
+def update_aggregates(per_model_results, avg_stats, dataset_name):
+    """
+    Update the aggregate dictionary per model with results from one dataset.
+    """
+    dataset_key = dataset_name.split("/")[-1]
+    for _, row in avg_stats.iterrows():
+        model_name = str(row["model"]).replace(" ", "-")
+        per = float(row["Average PER"]) if row["Average PER"] is not None else None
+        avg_dur = float(row["Average Duration (s)"]) if row["Average Duration (s)"] is not None else None
+
+        if model_name not in per_model_results:
+            per_model_results[model_name] = {}
+        per_model_results[model_name][dataset_key] = {"per": per, "avg_duration": avg_dur}
+
+def save_leaderboard_results(per_model_results, results_dir="eval-results"):
+    """
+    Persist one JSON file per model for the leaderboard app to consume.
+    """
     import json, os, time
-    # results_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "eval-results")
-    results_dir = os.path.join("eval-results")
     os.makedirs(results_dir, exist_ok=True)
-
     timestamp = int(time.time())
     for model_name, task_results in per_model_results.items():
         org_model = f"{model_name}"
@@ -261,6 +258,50 @@ def main():
             json.dump(payload, f, ensure_ascii=False, indent=2)
         print(f"Saved leaderboard result: {out_path}")
 
+def process_single_dataset(dataset_config, args, per_model_results):
+    """
+    Load, normalize, evaluate a single dataset and update aggregates.
+    """
+    if args.dataset and args.dataset not in dataset_config["name"]:
+        return
+
+    max_samples = args.max_samples if args.max_samples is not None else dataset_config.get("max_samples")
+    use_streaming = dataset_config.get("use_streaming", False)
+
+    dataset = load_dataset_with_limits(
+        dataset_config,
+        max_samples=max_samples,
+        use_streaming=use_streaming
+    )
+
+    if dataset is None:
+        return
+
+    dataset = cast_audio_column_safely(dataset)
+
+    dataset_final = prepare_dataset_for_evaluation(dataset, dataset_config, max_samples)
+    if dataset_final is None:
+        return
+
+    print(dataset_final)
+    print("Final size:", len(dataset_final))
+
+    full_results, avg_stats = evaluate_dataset(dataset_final)
+    print("Average Statistic per model (", dataset_config["name"], "):")
+    print(avg_stats)
+
+    update_aggregates(per_model_results, avg_stats, dataset_config["name"])
+
+def main():
+    args = parse_cli_args()
+    
+    per_model_results = {}
+
+    for dataset_config in DATASETS:
+        process_single_dataset(dataset_config, args, per_model_results)
+
+    save_leaderboard_results(per_model_results)
+
 
 if __name__ == "__main__":
     main()

utils/load_model.py

@@ -9,7 +9,6 @@ from transformers import (
 from .cmu_process import text_to_phoneme, cmu_to_ipa, clean_cmu
 
 from dotenv import load_dotenv
-import torch.backends.cudnn as cudnn
 
 # Load environment variables from .env file
 load_dotenv()
@@ -18,10 +17,6 @@ load_dotenv()
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print("Using device:", device)
 
-# Enable faster cudnn autotuner for variable input lengths
-if device.type == "cuda":
-    cudnn.benchmark = True
-
 # === Helper: move all tensors to model device ===
 def to_device(batch, device):
     if isinstance(batch, dict):
@@ -66,16 +61,9 @@ wavlm_model = AutoModelForCTC.from_pretrained("speech31/wavlm-large-english-phon
 def run_hubert_base(wav):
     start = time.time()
     inputs = base_proc(wav, sampling_rate=16000, return_tensors="pt", padding=True).input_values
-    if device.type == "cuda":
-        try:
-            inputs = inputs.pin_memory()
-        except Exception:
-            pass
-        inputs = inputs.to(device, non_blocking=True)
-    else:
-        inputs = inputs.to(device)
+    inputs = inputs.to(device)
     
-    with torch.inference_mode():
+    with torch.no_grad():
         logits = base_model(inputs).logits
     ids = torch.argmax(logits, dim=-1)
     text = base_proc.batch_decode(ids)[0]
@@ -87,43 +75,14 @@ def run_whisper(wav):
     start = time.time()
 
     inputs = whisper_proc(wav, sampling_rate=16000, return_tensors="pt")
-    input_features = inputs.input_features
-    if device.type == "cuda":
-        try:
-            input_features = input_features.pin_memory()
-        except Exception:
-            pass
-        input_features = input_features.to(device, non_blocking=True)
-    else:
-        input_features = input_features.to(device)
+    input_features = inputs.input_features.to(device)
     attention_mask = inputs.get("attention_mask", None)
     gen_kwargs = {"language": "en"}
     if attention_mask is not None:
-        if device.type == "cuda":
-            try:
-                attention_mask = attention_mask.pin_memory()
-            except Exception:
-                pass
-            gen_kwargs["attention_mask"] = attention_mask.to(device, non_blocking=True)
-        else:
-            gen_kwargs["attention_mask"] = attention_mask.to(device)
-
-    # Force English transcription and use greedy decoding with short max tokens for speed
-    try:
-        forced_ids = whisper_proc.get_decoder_prompt_ids(language="en", task="transcribe")
-    except Exception:
-        forced_ids = None
-
-    with torch.inference_mode():
-        pred_ids = whisper_model.generate(
-            input_features,
-            forced_decoder_ids=forced_ids,
-            do_sample=False,
-            num_beams=1,
-            max_new_tokens=64,
-            use_cache=True,
-            **gen_kwargs,
-        )
+        gen_kwargs["attention_mask"] = attention_mask.to(device)
+
+    with torch.no_grad():
+        pred_ids = whisper_model.generate(input_features, **gen_kwargs)
 
     text = whisper_proc.batch_decode(pred_ids, skip_special_tokens=True)[0]
     phonemes = text_to_phoneme(text)
@@ -134,18 +93,10 @@ def run_model(wav):
     start = time.time()
 
     # Prepare input (BatchEncoding supports .to(device))
-    inputs = proc(wav, sampling_rate=16000, return_tensors="pt")
-    if device.type == "cuda":
-        try:
-            inputs = inputs.pin_memory()
-        except Exception:
-            pass
-        inputs = inputs.to(device, non_blocking=True)
-    else:
-        inputs = inputs.to(device)
+    inputs = proc(wav, sampling_rate=16000, return_tensors="pt").to(device)
 
     # Forward pass
-    with torch.inference_mode():
+    with torch.no_grad():
         logits = model(**inputs).logits
 
     # Greedy decode
@@ -159,17 +110,10 @@ def run_timit(wav):
     start = time.time()
     # Read and process the input
     inputs = timit_proc(wav, sampling_rate=16_000, return_tensors="pt", padding=True)
-    if device.type == "cuda":
-        try:
-            inputs = inputs.pin_memory()
-        except Exception:
-            pass
-        inputs = inputs.to(device, non_blocking=True)
-    else:
-        inputs = inputs.to(device)
+    inputs = inputs.to(device)
     
     # Forward pass
-    with torch.inference_mode():
+    with torch.no_grad():
         logits = timit_model(inputs.input_values, attention_mask=inputs.attention_mask).logits
     
     # Decode id into string
@@ -189,18 +133,10 @@ def run_gruut(wav):
         sampling_rate=16000, 
         return_tensors="pt", 
         padding=True
-    )
-    if device.type == "cuda":
-        try:
-            inputs = inputs.pin_memory()
-        except Exception:
-            pass
-        inputs = inputs.to(device, non_blocking=True)
-    else:
-        inputs = inputs.to(device)
+    ).to(device)
 
     # Forward pass
-    with torch.inference_mode():
+    with torch.no_grad():
         logits = gruut_model(**inputs).logits
 
     # Greedy decode → IPA phonemes
@@ -219,21 +155,13 @@ def run_wavlm_large_phoneme(wav):
         sampling_rate=16000,
         return_tensors="pt",
         padding=True
-    )
-    if device.type == "cuda":
-        try:
-            inputs = inputs.pin_memory()
-        except Exception:
-            pass
-        inputs = inputs.to(device, non_blocking=True)
-    else:
-        inputs = inputs.to(device)
+    ).to(device)
 
     input_values = inputs.input_values
     attention_mask = inputs.get("attention_mask", None)
 
     # Forward pass
-    with torch.inference_mode():
+    with torch.no_grad():
         logits = wavlm_model(input_values, attention_mask=attention_mask).logits
 
     # Greedy decode → phoneme tokens