src/eval/evaluate_nptel_pure.py

import os
import sys
import json
import argparse
import numpy as np
import librosa
import torch
from tqdm import tqdm
import Levenshtein
import nltk

# Ensure the parent directory and current directory are on sys.path for local imports
# Add project root to sys.path
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from transformers import Wav2Vec2Processor
from src.models.phoneme_embedder import Wav2Vec2PhonemeEmbedder
from src.g2p.g2p_utils import G2PManager
from src.utils.audio_utils import AudioPreprocessor

def parse_args():
    parser = argparse.ArgumentParser(description="Evaluate Phoneme Embedder on Local NPTEL-pure Dataset")
    parser.add_argument("--dataset_dir", default="sample_dataset/nptel-pure", help="Path to the local NPTEL-pure dataset")
    parser.add_argument("--model_dir", default="MihirRPatil/nptel-asr-phoneme-v2", help="Hugging Face repo ID or path to local model checkpoints")
    parser.add_argument("--processor_dir", default="models/processor_dir", help="Path to local Wav2Vec2 processor directory")
    parser.add_argument("--dict_path", default="src/g2p/output_v2_detailed.dict", help="Path to G2P mapping dictionary")
    parser.add_argument("--limit", type=int, default=None, help="Limit number of samples to process (for quick testing)")
    parser.add_argument("--transcript_mode", choices=["corrected", "original", "metadata"], default="corrected",
                        help="Primary transcript source: human-corrected (corrected), original-ASR (original), or JSON-metadata (metadata)")
    return parser.parse_args()

def get_transcript(dataset_dir, file_hash, transcript_mode="corrected"):
    """
    Retrieves the transcript for a given file hash.
    Tries different modes and falls back sequentially.
    """
    corrected_path = os.path.join(dataset_dir, "corrected_txt", f"{file_hash}.txt")
    original_path = os.path.join(dataset_dir, "original_txt", f"{file_hash}.txt")
    metadata_path = os.path.join(dataset_dir, "metadata", f"{file_hash}.json")

    # Tiered search based on transcript_mode configuration
    if transcript_mode == "corrected":
        search_order = [corrected_path, original_path, metadata_path]
    elif transcript_mode == "original":
        search_order = [original_path, corrected_path, metadata_path]
    else:
        search_order = [metadata_path, corrected_path, original_path]

    for path in search_order:
        if not os.path.exists(path):
            continue
        try:
            if path.endswith(".txt"):
                with open(path, "r", encoding="utf-8") as f:
                    text = f.read().strip()
                    if text:
                        return text
            elif path.endswith(".json"):
                with open(path, "r", encoding="utf-8") as f:
                    data = json.load(f)
                    text = data.get("original_phrase", "").strip()
                    if text:
                        return text
        except Exception:
            pass
    return None

def load_processor_and_model(model_dir, processor_dir, device):
    """
    Loads Wav2Vec2Processor and Wav2Vec2PhonemeEmbedder with local path checks and HF fallbacks.
    """
    # 1. Load Processor
    processor = None
    if os.path.exists(processor_dir):
        print(f"Loading processor from local directory: {processor_dir}...")
        try:
            processor = Wav2Vec2Processor.from_pretrained(processor_dir)
        except Exception as e:
            print(f"⚠️ Failed to load processor from local directory {processor_dir}: {e}")
    
    # Fallback to loading processor from the model source
    if processor is None:
        print(f"Trying to load processor from model source: {model_dir}...")
        try:
            processor = Wav2Vec2Processor.from_pretrained(model_dir)
        except Exception as e:
            print(f"⚠️ Failed to load processor from {model_dir}: {e}. Falling back to facebook/wav2vec2-xlsr-53...")
            processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-xlsr-53")

    # 2. Load Model
    local_weights = None
    if os.path.exists(model_dir):
        checkpoints = sorted(
            [d for d in os.listdir(model_dir) if d.startswith("checkpoint-")],
            key=lambda x: int(x.split("-")[1]) if "-" in x else 0
        )
        if checkpoints:
            local_weights = os.path.join(model_dir, checkpoints[-1])
            print(f"✅ Found latest local checkpoint at: {local_weights}")
        elif os.path.exists(os.path.join(model_dir, "model.safetensors")):
            local_weights = model_dir
            print(f"✅ Found model weights at root of model_dir: {local_weights}")

    if local_weights:
        print(f"🚀 Loading pre-trained state from local path: {local_weights}...")
        model = Wav2Vec2PhonemeEmbedder.from_pretrained(local_weights)
    else:
        print(f"🚀 Loading model directly from HF Hub: {model_dir}...")
        model = Wav2Vec2PhonemeEmbedder.from_pretrained(model_dir)
        
    model = model.to(device)
    model.eval()
    return processor, model

def main():
    args = parse_args()
    
    # Download required NLTK resources
    print("Checking NLTK resources...")
    for res in ['averaged_perceptron_tagger', 'averaged_perceptron_tagger_eng', 'cmudict', 'punkt', 'punkt_tab']:
        try:
            nltk.download(res, quiet=True)
        except Exception:
            pass

    # Directory sanity checks
    wav_dir = os.path.join(args.dataset_dir, "wav")
    if not os.path.exists(wav_dir):
        print(f"❌ Error: WAV directory '{wav_dir}' not found.")
        return

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"🖥️ Using device: {device}")

    # Load Model and Processor
    processor, model = load_processor_and_model(args.model_dir, args.processor_dir, device)
    pad_token_id = processor.tokenizer.pad_token_id or 0

    # Build token to phoneme mapping for legible output logs
    try:
        vocab = processor.tokenizer.get_vocab()
        id2phoneme = {v: k for k, v in vocab.items()}
    except Exception:
        # Fallback if vocab cannot be extracted
        vocab_json_path = os.path.join(args.model_dir, "vocab.json")
        if os.path.exists(vocab_json_path):
            with open(vocab_json_path, "r", encoding="utf-8") as f:
                v_dict = json.load(f)
            id2phoneme = {v: k for k, v in v_dict.items()}
        else:
            id2phoneme = {}

    # Initialize utilities
    preprocessor = AudioPreprocessor(sr=16000)
    
    # Check if local dict path exists, if not fall back to None to let G2PManager autodetect in its subfolders
    dict_path = args.dict_path
    if not os.path.exists(dict_path):
        print(f"ℹ️ Local G2P dictionary path '{dict_path}' not found. Letting G2PManager auto-detect output_full.dict...")
        dict_path = None
    g2p = G2PManager(dict_path=dict_path)

    # Get all WAV files
    wav_files = sorted([f for f in os.listdir(wav_dir) if f.endswith(".wav")])
    if args.limit:
        wav_files = wav_files[:args.limit]

    print(f"\n📊 Evaluating NPTEL-pure dataset ({len(wav_files)} samples)...")

    per_scores = []
    skipped = 0
    error_count = 0
    max_error_prints = 5

    for filename in tqdm(wav_files):
        file_hash = os.path.splitext(filename)[0]
        audio_path = os.path.join(wav_dir, filename)
        
        # Retrieve transcript with fallbacks
        transcript = get_transcript(args.dataset_dir, file_hash, args.transcript_mode)

        if not transcript:
            if error_count < max_error_prints:
                print(f"⚠️  Skipped {file_hash}: Transcript not found in corrected, original, or metadata files.")
                error_count += 1
            skipped += 1
            continue

        try:
            # 1. Load Audio
            audio_array, sr = librosa.load(audio_path, sr=None)
            if sr != 16000:
                audio_array = librosa.resample(audio_array, orig_sr=sr, target_sr=16000)

            # 2. Preprocess Audio (Spectral Subtraction + Silero VAD)
            clean_audio = preprocessor.preprocess(audio_array)
            if len(clean_audio) == 0:
                if error_count < max_error_prints:
                    print(f"⚠️  Skipped {file_hash}: VAD trimmed audio to 0 samples.")
                    error_count += 1
                skipped += 1
                continue

            # 3. Extract Audio Features
            input_values = processor(clean_audio, sampling_rate=16000).input_values[0]
            input_tensor = torch.tensor(input_values, dtype=torch.float32).unsqueeze(0).to(device)

            # 4. G2P conversion of target transcript
            target_phonemes = g2p.convert_sentence(transcript)
            if len(target_phonemes) == 0:
                if error_count < max_error_prints:
                    print(f"⚠️  Skipped {file_hash}: G2P conversion resulted in empty phonemes list.")
                    error_count += 1
                skipped += 1
                continue
                
            target_ids = processor.tokenizer.convert_tokens_to_ids(target_phonemes)
            clean_ref = [rid for rid in target_ids if rid >= 0 and rid != pad_token_id]

            if not clean_ref:
                if error_count < max_error_prints:
                    print(f"⚠️  Skipped {file_hash}: Clean tokenized reference is empty.")
                    error_count += 1
                skipped += 1
                continue

            # 5. Model Inference
            with torch.no_grad():
                outputs = model(input_tensor)
                logits = outputs["logits"] if isinstance(outputs, dict) else outputs.logits
                pred_ids = torch.argmax(logits, dim=-1)[0].cpu().numpy().tolist()

            # 6. Collapse duplicate predictions (CTC decoding)
            collapsed_pred = []
            prev = None
            for pid in pred_ids:
                if pid == prev or pid == pad_token_id:
                    prev = pid
                    continue
                prev = pid
                collapsed_pred.append(pid)

            # 7. Compute Phoneme Error Rate (PER)
            dist = Levenshtein.distance(clean_ref, collapsed_pred)
            max_len = max(len(clean_ref), len(collapsed_pred), 1)
            per = dist / max_len
            per_scores.append(per)

            # 8. Detailed Comparison Log for the first few samples
            if len(per_scores) <= 3:
                ref_phonemes_str = " ".join([id2phoneme.get(i, f"[{i}]") for i in clean_ref])
                pred_phonemes_str = " ".join([id2phoneme.get(i, f"[{i}]") for i in collapsed_pred])
                print(f"\n--- Detailed Log: Sample {len(per_scores)} ({file_hash}) ---")
                print(f"Transcript: {transcript}")
                print(f"Ref Phonemes: {ref_phonemes_str}")
                print(f"Hyp Phonemes: {pred_phonemes_str}")
                print(f"PER: {per:.2%}")
                print("-" * 50)

        except Exception as e:
            if error_count < max_error_prints:
                print(f"⚠️  Error processing sample {file_hash}: {e}")
                error_count += 1
            skipped += 1
            continue

    # Report Final Statistics
    if per_scores:
        mean_per = np.mean(per_scores)
        median_per = np.median(per_scores)
        std_per = np.std(per_scores)
        total_processed = len(per_scores)
        
        print("\n" + "="*50)
        print("          NPTEL-PURE EVALUATION REPORT")
        print("="*50)
        print(f"Total Files Scanned: {len(wav_files)}")
        print(f"Successfully Processed: {total_processed}")
        print(f"Skipped / Failed:     {skipped}")
        print("-"*50)
        print(f"Mean PER:    {mean_per:.2%}")
        print(f"Median PER:  {median_per:.2%}")
        print(f"Std Dev PER: {std_per:.2%}")
        print("="*50)
    else:
        print("\n❌ Error: Failed to evaluate any samples.")

if __name__ == "__main__":
    main()