speech_to_text.py

import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
from transformers.utils import is_flash_attn_2_available
import torchaudio
import numpy as np
import re

def axelify(text):
    pattern = r'\ba\w*(x|ks)\w*la\b'
    corrected_text = re.sub(pattern, "Axela", text, flags=re.IGNORECASE)
    corrected_text = corrected_text.lower().replace("excellent", "")
    return corrected_text

def load_speech_model(device="cpu"):
    """
    Load and return the Whisper model and processor.
    
    Args:
        device: Device to load the model on ("cpu" or "cuda")
    
    Returns:
        Tuple of (model, processor, device)
    """
    model_id = "distil-whisper/distil-small.en"
    torch_dtype = torch.float32
    device = torch.device(device)

    model = AutoModelForSpeechSeq2Seq.from_pretrained(
        model_id,
        dtype=torch_dtype,
        low_cpu_mem_usage=True,
        use_safetensors=True
    ).to(device)
    model.eval()

    processor = AutoProcessor.from_pretrained(model_id)
    
    return model, processor, device

def speech_to_text(audio_input, model, processor, device, sample_rate=16000):
    """
    Convert speech to text using Whisper.
    
    Args:
        audio_input: Either a file path (str) or numpy array of audio samples (int16)
        model: Pre-loaded Whisper model
        processor: Pre-loaded Whisper processor
        device: Device the model is on
        sample_rate: Sample rate in Hz (default: 16000Hz). Only used for numpy arrays.
    
    Returns:
        Transcribed and axelified text
    """

    # Handle both file path and numpy array inputs
    if isinstance(audio_input, str):
        # Load from file
        audio, sr = torchaudio.load(audio_input, format="wav")
    elif isinstance(audio_input, np.ndarray):
        # Convert numpy array to torch tensor
        # Assuming input is int16, normalize to float32 in range [-1, 1]
        audio = torch.from_numpy(audio_input.astype(np.float32) / 32768.0)
        if audio.dim() == 1:
            audio = audio.unsqueeze(0)  # Add channel dimension if needed
        sr = sample_rate
    else:
        raise ValueError("audio_input must be either a file path (str) or numpy array")

    # Whisper expects 16kHz mono
    if sr != 16000:
        resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=16000)
        audio = resampler(audio)
    if audio.shape[0] > 1:
        audio = torch.mean(audio, dim=0, keepdim=True)  # Convert to mono

    # Preprocess
    inputs = processor(audio.squeeze(), return_tensors="pt").to(device)

    # Inference
    with torch.no_grad():
        generated_ids = model.generate(
            inputs["input_features"],
            max_new_tokens=128
        )
        transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]

    return axelify(transcription)