interview/transcription_engine.py

import os
import time
import threading
import logging
from .models import TranscriptionProviderConfig

log = logging.getLogger(__name__)

# --- Global caches para los modelos cargados en memoria ---
# (Esto evita cargar los modelos desde cero en cada peticion)
_whisper_model = None
_whisper_lock = threading.Lock()

_nemotron_model = None
_nemotron_lock = threading.Lock()

def get_whisper_model(model_size="base"):
    global _whisper_model
    if _whisper_model is None:
        try:
            import whisper
            with _whisper_lock:
                if _whisper_model is None:
                    log.info(f"Cargando modelo Whisper ({model_size}) en memoria...")
                    _whisper_model = whisper.load_model(model_size)
        except ImportError:
            log.error("La librería 'whisper' no está instalada. Ejecuta: pip install openai-whisper")
            return None
    return _whisper_model

def get_nemotron_model(model_name="nvidia/nemotron-3.5-asr-streaming-0.6b"):
    global _nemotron_model
    if _nemotron_model is None:
        try:
            import torch
            # Fix PyTorch 2.4+ strict weights_only loading breaking NeMo
            _orig_load = torch.load
            def patched_load(*args, **kwargs):
                kwargs['weights_only'] = False
                return _orig_load(*args, **kwargs)
            torch.load = patched_load

            import nemo.collections.asr as nemo_asr
            with _nemotron_lock:
                if _nemotron_model is None:
                    log.info(f"Cargando modelo Nemotron ({model_name}) en memoria...")
                    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
                    _nemotron_model = nemo_asr.models.ASRModel.from_pretrained(model_name, map_location=device)
            
            # Restore original load
            torch.load = _orig_load
        except ImportError:
            log.error("La librería 'nemo_toolkit' no está instalada.")
            return None
        except Exception as e:
            log.error(f"Error cargando Nemotron: {e}")
            if 'torch' in locals() and hasattr(torch, 'load') and '_orig_load' in locals():
                torch.load = _orig_load
            return None
    return _nemotron_model

def transcribe_audio_file(audio_path, provider_type="whisper_local"):
    """
    Toma la ruta de un archivo de audio y el tipo de proveedor de transcripcion.
    Retorna (texto_transcrito, tiempo_de_ejecucion_ms, error_msg)
    """
    start_time = time.time()
    transcribed_text = ""
    error_msg = None

    try:
        config = TranscriptionProviderConfig.objects.filter(provider_type=provider_type, is_active=True).first()
        
        if provider_type == "whisper_local":
            model_size = config.model if config and config.model else "base"
            model = get_whisper_model(model_size)
            if not model:
                raise RuntimeError("Motor Whisper Local no disponible (faltan dependencias o falló la carga).")
            
            with _whisper_lock:
                result = model.transcribe(audio_path, language="es", fp16=False)
                transcribed_text = result["text"]

        elif provider_type == "whisper_api":
            try:
                from openai import OpenAI
            except ImportError:
                raise RuntimeError("La librería 'openai' no está instalada.")
                
            if not config or not config.api_key:
                raise RuntimeError("API Key no configurada para Whisper API.")
                
            client = OpenAI(api_key=config.api_key)
            with open(audio_path, "rb") as audio_file:
                transcription = client.audio.transcriptions.create(
                    model="whisper-1", 
                    file=audio_file,
                    language="es"
                )
                transcribed_text = transcription.text

        elif provider_type == "nemotron":
            model_name = config.model if config and config.model else "nvidia/nemotron-3.5-asr-streaming-0.6b"
            model = get_nemotron_model(model_name)
            if not model:
                log.warning("Nemotron falló. Haciendo fallback automático a whisper_local.")
                return transcribe_audio_file(audio_path, provider_type="whisper_local")
            
            with _nemotron_lock:
                import json
                import tempfile
                
                import json
                import tempfile

                # ponytail: get duration, NeMo's Lhotse adapter crashes without it
                try:
                    import librosa
                    duration = float(librosa.get_duration(path=audio_path))
                except Exception:
                    try:
                        import soundfile as sf
                        duration = float(sf.info(audio_path).duration)
                    except Exception:
                        duration = 10000.0 # safe fallback if both fail

                # ponytail: write the dict to a manifest file, since NeMo doesn't accept dicts natively in audio[]
                with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json') as f:
                    manifest_data = {
                        "audio_filepath": audio_path,
                        "duration": duration,
                        "target_lang": "es",
                        "lang": "es", # LazyNeMoIterator uses 'lang' by default to set supervision.language!
                        "text": ""
                    }
                    f.write(json.dumps(manifest_data) + '\n')
                    manifest_path = f.name

                try:
                    # ponytail: pass as positional arg without kwargs to avoid API signature mismatches
                    results = model.transcribe([manifest_path])
                    if isinstance(results, tuple):
                        results = results[0]
                    if isinstance(results, list) and len(results) > 0:
                        if isinstance(results[0], tuple):
                            transcribed_text = results[0][0]
                        else:
                            transcribed_text = results[0]
                        
                        # ponytail: NeMo can return Hypothesis objects instead of strings! Extract the text.
                        if hasattr(transcribed_text, 'text'):
                            transcribed_text = transcribed_text.text
                        elif not isinstance(transcribed_text, str):
                            transcribed_text = str(transcribed_text)
                except Exception as e:
                    raise e
                finally:
                    import os
                    if os.path.exists(manifest_path):
                        os.remove(manifest_path)

        else:
            raise ValueError(f"Proveedor de transcripción '{provider_type}' no soportado en el backend.")

    except Exception as e:
        log.exception(f"Error en transcribe_audio_file ({provider_type})")
        error_msg = str(e)

    end_time = time.time()
    execution_time_ms = int((end_time - start_time) * 1000)

    return transcribed_text, execution_time_ms, error_msg