src/AI_Models/wave2vec_inference.py

import torch
from transformers import AutoModelForCTC, AutoProcessor, Wav2Vec2Processor, Wav2Vec2ForCTC
import onnxruntime as rt
import numpy as np
import librosa
import warnings
import os
warnings.filterwarnings("ignore")


class Wave2Vec2Inference:
    def __init__(self, model_name, use_gpu=True):
        # Auto-detect device
        if use_gpu:
            if torch.backends.mps.is_available():
                self.device = "mps"
            elif torch.cuda.is_available():
                self.device = "cuda"
            else:
                self.device = "cpu"
        else:
            self.device = "cpu"
        
        print(f"Using device: {self.device}")
        
        # Load model and processor
        self.processor = AutoProcessor.from_pretrained(model_name)
        self.model = AutoModelForCTC.from_pretrained(model_name)
        self.model.to(self.device)
        self.model.eval()
        
        # Disable gradients for inference
        torch.set_grad_enabled(False)

    def buffer_to_text(self, audio_buffer):
        if len(audio_buffer) == 0:
            return ""

        # Convert to tensor
        if isinstance(audio_buffer, np.ndarray):
            audio_tensor = torch.from_numpy(audio_buffer).float()
        else:
            audio_tensor = torch.tensor(audio_buffer, dtype=torch.float32)

        # Process audio
        inputs = self.processor(
            audio_tensor,
            sampling_rate=16_000,
            return_tensors="pt",
            padding=True,
        )

        # Move to device
        input_values = inputs.input_values.to(self.device)
        attention_mask = inputs.attention_mask.to(self.device) if "attention_mask" in inputs else None

        # Inference
        with torch.no_grad():
            if attention_mask is not None:
                logits = self.model(input_values, attention_mask=attention_mask).logits
            else:
                logits = self.model(input_values).logits

        # Decode
        predicted_ids = torch.argmax(logits, dim=-1)
        if self.device != "cpu":
            predicted_ids = predicted_ids.cpu()
        
        transcription = self.processor.batch_decode(predicted_ids)[0]
        return transcription.lower().strip()

    def file_to_text(self, filename):
        try:
            audio_input, _ = librosa.load(filename, sr=16000, dtype=np.float32)
            return self.buffer_to_text(audio_input)
        except Exception as e:
            print(f"Error loading audio file {filename}: {e}")
            return ""


class Wave2Vec2ONNXInference:
    def __init__(self, model_name, onnx_path, use_gpu=True):
        self.processor = Wav2Vec2Processor.from_pretrained(model_name)
        
        # Setup ONNX Runtime
        options = rt.SessionOptions()
        options.graph_optimization_level = rt.GraphOptimizationLevel.ORT_ENABLE_ALL
        
        # Choose providers based on GPU availability
        providers = []
        if use_gpu and rt.get_available_providers():
            if 'CUDAExecutionProvider' in rt.get_available_providers():
                providers.append('CUDAExecutionProvider')
        providers.append('CPUExecutionProvider')
        
        self.model = rt.InferenceSession(onnx_path, options, providers=providers)
        self.input_name = self.model.get_inputs()[0].name
        print(f"ONNX model loaded with providers: {self.model.get_providers()}")

    def buffer_to_text(self, audio_buffer):
        if len(audio_buffer) == 0:
            return ""

        # Convert to tensor
        if isinstance(audio_buffer, np.ndarray):
            audio_tensor = torch.from_numpy(audio_buffer).float()
        else:
            audio_tensor = torch.tensor(audio_buffer, dtype=torch.float32)

        # Process audio
        inputs = self.processor(
            audio_tensor,
            sampling_rate=16_000,
            return_tensors="np",
            padding=True,
        )

        # ONNX inference
        input_values = inputs.input_values.astype(np.float32)
        onnx_outputs = self.model.run(None, {self.input_name: input_values})[0]
        
        # Decode
        prediction = np.argmax(onnx_outputs, axis=-1)
        transcription = self.processor.decode(prediction.squeeze().tolist())
        return transcription.lower().strip()

    def file_to_text(self, filename):
        try:
            audio_input, _ = librosa.load(filename, sr=16000, dtype=np.float32)
            return self.buffer_to_text(audio_input)
        except Exception as e:
            print(f"Error loading audio file {filename}: {e}")
            return ""


def convert_to_onnx(model_id_or_path, onnx_model_name):
    """Convert PyTorch model to ONNX format"""
    print(f"Converting {model_id_or_path} to ONNX...")
    model = Wav2Vec2ForCTC.from_pretrained(model_id_or_path)
    model.eval()
    
    # Create dummy input
    audio_len = 250000
    dummy_input = torch.randn(1, audio_len, requires_grad=True)

    torch.onnx.export(
        model,
        dummy_input,
        onnx_model_name,
        export_params=True,
        opset_version=14,
        do_constant_folding=True,
        input_names=["input"],
        output_names=["output"],
        dynamic_axes={
            "input": {1: "audio_len"},
            "output": {1: "audio_len"},
        },
    )
    print(f"ONNX model saved to: {onnx_model_name}")


def quantize_onnx_model(onnx_model_path, quantized_model_path):
    """Quantize ONNX model for faster inference"""
    print("Starting quantization...")
    from onnxruntime.quantization import quantize_dynamic, QuantType

    quantize_dynamic(
        onnx_model_path, 
        quantized_model_path, 
        weight_type=QuantType.QUInt8
    )
    print(f"Quantized model saved to: {quantized_model_path}")


def export_to_onnx(model_name, quantize=False):
    """
    Export model to ONNX format with optional quantization
    
    Args:
        model_name: HuggingFace model name
        quantize: Whether to also create quantized version
        
    Returns:
        tuple: (onnx_path, quantized_path or None)
    """
    onnx_filename = f"{model_name.split('/')[-1]}.onnx"
    convert_to_onnx(model_name, onnx_filename)
    
    quantized_path = None
    if quantize:
        quantized_path = onnx_filename.replace('.onnx', '.quantized.onnx')
        quantize_onnx_model(onnx_filename, quantized_path)
    
    return onnx_filename, quantized_path


def create_inference(model_name, use_onnx=False, onnx_path=None, use_gpu=True, use_onnx_quantize=False):
    """
    Create optimized inference instance
    
    Args:
        model_name: HuggingFace model name
        use_onnx: Whether to use ONNX runtime
        onnx_path: Path to ONNX model file
        use_gpu: Whether to use GPU if available
        use_onnx_quantize: Whether to use quantized ONNX model
        
    Returns:
        Inference instance
    """
    if use_onnx:
        if not onnx_path or not os.path.exists(onnx_path):
            # Convert to ONNX if path not provided or doesn't exist
            onnx_filename = f"{model_name.split('/')[-1]}.onnx"
            convert_to_onnx(model_name, onnx_filename)
            onnx_path = onnx_filename
            
        if use_onnx_quantize:
            quantized_path = onnx_path.replace('.onnx', '.quantized.onnx')
            if not os.path.exists(quantized_path):
                quantize_onnx_model(onnx_path, quantized_path)
            onnx_path = quantized_path
            
        print(f"Using ONNX model: {onnx_path}")
        return Wave2Vec2ONNXInference(model_name, onnx_path, use_gpu)
    else:
        print("Using PyTorch model")
        return Wave2Vec2Inference(model_name, use_gpu)


if __name__ == "__main__":
    import time
    
    model_name = "facebook/wav2vec2-large-960h-lv60-self"
    test_file = "test.wav"
    
    if not os.path.exists(test_file):
        print(f"Test file {test_file} not found. Please provide a valid audio file.")
        exit(1)
    
    # Test different configurations
    configs = [
        {"use_onnx": False, "use_gpu": True},
        {"use_onnx": True, "use_gpu": True, "use_onnx_quantize": False},
        {"use_onnx": True, "use_gpu": True, "use_onnx_quantize": True},
    ]
    
    for config in configs:
        print(f"\n=== Testing config: {config} ===")
        
        # Create inference instance
        asr = create_inference(model_name, **config)
        
        # Warm up
        asr.file_to_text(test_file)
        
        # Test performance
        times = []
        for i in range(5):
            start_time = time.time()
            text = asr.file_to_text(test_file)
            end_time = time.time()
            execution_time = end_time - start_time
            times.append(execution_time)
            print(f"Run {i+1}: {execution_time:.3f}s - {text[:50]}...")
        
        avg_time = sum(times) / len(times)
        print(f"Average time: {avg_time:.3f}s")