""" IndicWav2Vec Validator ====================== Uses AI4Bharat's IndicWav2Vec model for forced alignment and ASR. https://github.com/AI4Bharat/indicwav2vec This model is pre-trained on 40 Indian languages and can be used for: - Speech recognition (ASR) - Forced alignment via CTC segmentation """ import os import time from typing import Optional, List, Dict, Any from pathlib import Path from .base import ( BaseValidator, ValidationResult, WordAlignment, ValidatorStatus, normalize_language_code ) class IndicWav2VecValidator(BaseValidator): """ Validator using AI4Bharat's IndicWav2Vec model. Provides: - ASR transcription for comparison - CTC-based forced alignment with word timestamps """ name = "indicwav2vec" description = "AI4Bharat IndicWav2Vec - Multilingual Indian ASR & Alignment" # Model variants available MODEL_VARIANTS = { "base": "ai4bharat/indicwav2vec-hindi", # Fallback "large": "ai4bharat/indicwav2vec_v1_bengali", # Example "multilingual": "facebook/wav2vec2-large-xlsr-53", # Multilingual fallback } # Language-specific fine-tuned models (if available) LANGUAGE_MODELS = { "hi": "ai4bharat/indicwav2vec-hindi", "bn": "ai4bharat/indicwav2vec_v1_bengali", "te": "facebook/wav2vec2-large-xlsr-53", # Use XLSR for Telugu "ta": "facebook/wav2vec2-large-xlsr-53", "kn": "facebook/wav2vec2-large-xlsr-53", "ml": "facebook/wav2vec2-large-xlsr-53", "gu": "facebook/wav2vec2-large-xlsr-53", "mr": "facebook/wav2vec2-large-xlsr-53", } def __init__( self, enabled: bool = True, model_variant: str = "multilingual", device: str = "auto", **kwargs ): """ Initialize IndicWav2Vec validator. Args: enabled: Whether validator is active model_variant: Which model variant to use device: "cuda", "cpu", or "auto" """ super().__init__(enabled=enabled, **kwargs) self.model_variant = model_variant self.device_preference = device self.processor = None self.model = None self.device = None def setup(self) -> bool: """Load the Wav2Vec2 model and processor.""" try: import torch from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor # Determine device if self.device_preference == "auto": self.device = "cuda" if torch.cuda.is_available() else "cpu" else: self.device = self.device_preference print(f"[{self.name}] Loading model on {self.device}...") # Use XLSR multilingual as default (good coverage) model_name = "facebook/wav2vec2-large-xlsr-53" self.processor = Wav2Vec2Processor.from_pretrained(model_name) self.model = Wav2Vec2ForCTC.from_pretrained(model_name) self.model.to(self.device) self.model.eval() print(f"[{self.name}] Model loaded successfully") return True except ImportError as e: print(f"[{self.name}] Missing dependency: {e}") print(f"[{self.name}] Install with: pip install transformers torch torchaudio") return False except Exception as e: print(f"[{self.name}] Setup error: {e}") return False def _load_audio(self, audio_path: str, target_sr: int = 16000): """Load and resample audio to target sample rate.""" import torchaudio waveform, sample_rate = torchaudio.load(audio_path) # Resample if needed if sample_rate != target_sr: resampler = torchaudio.transforms.Resample(sample_rate, target_sr) waveform = resampler(waveform) # Convert to mono if stereo if waveform.shape[0] > 1: waveform = waveform.mean(dim=0, keepdim=True) return waveform.squeeze(), target_sr def _get_word_alignments_ctc( self, logits, transcription: str, audio_duration: float ) -> List[WordAlignment]: """ Extract word-level alignments from CTC logits. This is an approximation based on CTC blank predictions. """ import torch # Get predicted token IDs predicted_ids = torch.argmax(logits, dim=-1).squeeze() # CTC blank token is usually 0 blank_id = self.processor.tokenizer.pad_token_id or 0 # Find non-blank segments words = transcription.split() if not words: return [] # Simple proportional alignment based on character count # (More sophisticated CTC segmentation would use ctc-segmentation library) total_chars = sum(len(w) for w in words) if total_chars == 0: return [] alignments = [] current_time = 0.0 for word in words: word_duration = (len(word) / total_chars) * audio_duration alignments.append(WordAlignment( word=word, start_time=current_time, end_time=current_time + word_duration, confidence=None # CTC doesn't give per-word confidence easily )) current_time += word_duration return alignments def validate( self, audio_path: str, reference_text: Optional[str] = None, language: str = "te" ) -> ValidationResult: """ Process audio through IndicWav2Vec. Args: audio_path: Path to audio file reference_text: Optional reference transcription language: Language code Returns: ValidationResult with ASR output and alignments """ if not self.enabled: return ValidationResult( validator_name=self.name, audio_path=audio_path, success=False, error_message="Validator disabled" ) if not self.ensure_setup(): return ValidationResult( validator_name=self.name, audio_path=audio_path, success=False, error_message="Model not loaded" ) start_time = time.time() lang_code = normalize_language_code(language) try: import torch # Load audio waveform, sr = self._load_audio(audio_path) audio_duration = len(waveform) / sr # Process through model inputs = self.processor( waveform.numpy(), sampling_rate=sr, return_tensors="pt", padding=True ) input_values = inputs.input_values.to(self.device) with torch.no_grad(): outputs = self.model(input_values) logits = outputs.logits # Decode transcription predicted_ids = torch.argmax(logits, dim=-1) transcription = self.processor.batch_decode(predicted_ids)[0] # Get alignments (basic CTC-based) alignments = self._get_word_alignments_ctc( logits.cpu(), transcription, audio_duration ) # Calculate confidence from logits (softmax probability of predicted tokens) probs = torch.softmax(logits, dim=-1) max_probs = probs.max(dim=-1).values overall_confidence = max_probs.mean().item() processing_time = time.time() - start_time return ValidationResult( validator_name=self.name, audio_path=audio_path, success=True, transcription=transcription, word_alignments=alignments, overall_confidence=overall_confidence, processing_time_sec=processing_time, audio_duration_sec=audio_duration, raw_output={ "language": lang_code, "model": "wav2vec2-large-xlsr-53" } ) except Exception as e: return ValidationResult( validator_name=self.name, audio_path=audio_path, success=False, error_message=str(e), processing_time_sec=time.time() - start_time ) def cleanup(self): """Release GPU memory.""" if self.model is not None: del self.model self.model = None if self.processor is not None: del self.processor self.processor = None import torch if torch.cuda.is_available(): torch.cuda.empty_cache()