src/utils.py

# utils.py
import opencc
import os
from pathlib import Path
import sys
from typing import Optional
import multiprocessing

# Detect logical cores (vCPUs available to the container)
# On HF Spaces free tier, cpu_count() reports 16 but only 2 are actually available
detected_cpus = multiprocessing.cpu_count()
if os.environ.get('SPACE_ID'):
    # HF Spaces free tier limitation
    num_vcpus = min(detected_cpus, 2)
else:
    num_vcpus = detected_cpus

model_names = {
    "tiny English":"tiny", 
    "tiny Arabic":"tiny-ar", 
    "tiny Chinese":"tiny-zh", 
    "tiny Japanese":"tiny-ja", 
    "tiny Korean":"tiny-ko", 
    "tiny Ukrainian":"tiny-uk", 
    "tiny Vietnamese":"tiny-vi",
    "base English":"base", 
    "base Spanish":"base-es"
}

# Using only the two specified sherpa-onnx models from Hugging Face
sensevoice_models = {
    "SenseVoice Small (2024)": "csukuangfj/sherpa-onnx-sense-voice-zh-en-ja-ko-yue-2024-07-17",
    "SenseVoice Small (2025 int8)": "csukuangfj/sherpa-onnx-sense-voice-zh-en-ja-ko-yue-int8-2025-09-09",
}

# X-ASR zipformer transducer (offline) models, loaded via
# sherpa_onnx.OfflineRecognizer.from_transducer.
xasr_models = {
    "X-ASR Zipformer zh-en punct (int8)": "csukuangfj2/sherpa-onnx-x-asr-zipformer-transducer-zh-en-punct-int8-2026-06-03",
}

# Qwen3-ASR (offline) models, loaded via
# sherpa_onnx.OfflineRecognizer.from_qwen3_asr.
qwen3_models = {
    "Qwen3-ASR 0.6B (int8)": "csukuangfj2/sherpa-onnx-qwen3-asr-0.6B-int8-2026-03-25",
}

available_gguf_llms = {
    "Gemma-3-1B": ("bartowski/google_gemma-3-1b-it-qat-GGUF", "google_gemma-3-1b-it-qat-Q4_0.gguf"),
    "Gemma-3-270M": ("bartowski/google_gemma-3-270m-it-qat-GGUF", "google_gemma-3-270m-it-qat-Q8_0.gguf"),
    "Gemma-3-3N-E2B": ("unsloth/gemma-3n-E2B-it-GGUF", "gemma-3n-E2B-it-Q4_0.gguf"),
    "Gemma-3-3N-E4B": ("unsloth/gemma-3n-E4B-it-GGUF", "gemma-3n-E4B-it-Q4_0.gguf"),
    "Gemma-4-E2B": ("unsloth/gemma-4-E2B-it-GGUF", "gemma-4-E2B-it-Q4_0.gguf"),
    "Gemma-4-E4B": ("unsloth/gemma-4-E4B-it-GGUF", "gemma-4-E4B-it-Q4_0.gguf"),
}

s2tw_converter = opencc.OpenCC('s2twp')

def get_writable_model_dir():
    """Get appropriate model directory for HF Spaces"""
    # Check for HF Spaces environment
    if os.environ.get('SPACE_ID'):
        # Use HF Spaces cache directory
        cache_dir = Path('/tmp/models')
    else:
        # Use standard cache directory
        cache_dir = Path.home() / ".cache" / "speech_assistant" / "models"
    
    # Ensure directory exists
    cache_dir.mkdir(parents=True, exist_ok=True)
    return cache_dir

def download_sensevoice_model(model_name: str) -> Path:
    """Download SenseVoice model from Hugging Face using official tools"""
    
    try:
        from huggingface_hub import snapshot_download
        from huggingface_hub.utils import HFValidationError
    except ImportError:
        raise ImportError("Please install huggingface_hub: pip install huggingface_hub")
    
    # Use model_name directly as repo_id
    repo_id = model_name
    model_cache_dir = get_writable_model_dir()
    local_dir = model_cache_dir / model_name.replace("/", "--")
    
    # Check if model already exists
    model_file = "model.int8.onnx" if "int8" in model_name else "model.onnx"
    model_file_path = local_dir / model_file
    tokens_file_path = local_dir / "tokens.txt"
    
    if model_file_path.exists() and tokens_file_path.exists():
        print(f"Model {model_name} already exists, skipping download")
        return local_dir
    
    # Remove existing incomplete model directory
    if local_dir.exists():
        import shutil
        print(f"Removing incomplete model directory: {local_dir}")
        shutil.rmtree(local_dir)
    
    print(f"Downloading {model_name} from Hugging Face")
    print("This may take several minutes depending on your connection...")
    
    try:
        # Use HF's snapshot_download for reliable download
        snapshot_download(
            repo_id=repo_id,
            local_dir=str(local_dir),
            resume_download=True,  # Resume if interrupted
            max_workers=4,  # Parallel downloads
        )
        
        print(f"Model {model_name} downloaded successfully!")
        return local_dir
        
    except HFValidationError as e:
        print(f"Hugging Face validation error: {e}")
        raise
    except Exception as e:
        print(f"Download failed: {str(e)}")
        # Clean up partial download
        if local_dir.exists():
            import shutil
            shutil.rmtree(local_dir)
        raise e

def load_sensevoice_model(model_name: str = "csukuangfj/sherpa-onnx-sense-voice-zh-en-ja-ko-yue-2024-07-17"):
    """Load SenseVoice ONNX model from Hugging Face"""
    try:
        # Try to import sherpa-onnx
        import sherpa_onnx
        
        print(f"Loading model: {model_name}")
        
        # Download model if not exists
        model_path = download_sensevoice_model(model_name)
        
        # Determine which model file to use
        model_file = "model.int8.onnx" if "int8" in model_name else "model.onnx"
        model_file_path = model_path / model_file
        
        print("Initializing recognizer...")
        
        # Initialize recognizer with proper settings
        recognizer = sherpa_onnx.OfflineRecognizer.from_sense_voice(
            model=str(model_file_path),
            tokens=str(model_path / "tokens.txt"),
            use_itn=True,  # Enable inverse text normalization
            language="auto"  # Auto-detect language
        )
        
        print("Model loaded successfully!")
        return recognizer
    except Exception as e:
        print(f"Failed to load SenseVoice model: {e}")
        # Try to force redownload on next attempt
        model_cache_dir = get_writable_model_dir()
        model_dir = model_cache_dir / model_name.replace("/", "--")
        if model_dir.exists():
            import shutil
            print(f"Removing model directory for redownload: {model_dir}")
            shutil.rmtree(model_dir)
        raise e


def download_sherpa_repo(repo_id: str, sentinel_files) -> Path:
    """Download a sherpa-onnx model repo from Hugging Face into the cache.

    ``sentinel_files`` are paths (relative to the repo root) whose presence means
    the model is already fully downloaded, so the download can be skipped.
    Returns the local directory containing the model files.
    """
    try:
        from huggingface_hub import snapshot_download
    except ImportError:
        raise ImportError("Please install huggingface_hub: pip install huggingface_hub")

    model_cache_dir = get_writable_model_dir()
    local_dir = model_cache_dir / repo_id.replace("/", "--")

    if all((local_dir / f).exists() for f in sentinel_files):
        print(f"Model {repo_id} already exists, skipping download")
        return local_dir

    if local_dir.exists():
        import shutil
        print(f"Removing incomplete model directory: {local_dir}")
        shutil.rmtree(local_dir)

    print(f"Downloading {repo_id} from Hugging Face (this may take a few minutes)...")
    snapshot_download(repo_id=repo_id, local_dir=str(local_dir), max_workers=4)
    print(f"Model {repo_id} downloaded successfully!")
    return local_dir


def load_xasr_model(model_name: str):
    """Load an X-ASR zipformer transducer model (offline) via sherpa-onnx."""
    import sherpa_onnx

    print(f"Loading X-ASR model: {model_name}")
    model_dir = download_sherpa_repo(
        model_name,
        sentinel_files=(
            "encoder-epoch-99-avg-1.int8.onnx",
            "decoder-epoch-99-avg-1.onnx",
            "joiner-epoch-99-avg-1.int8.onnx",
            "tokens.txt",
        ),
    )
    recognizer = sherpa_onnx.OfflineRecognizer.from_transducer(
        encoder=str(model_dir / "encoder-epoch-99-avg-1.int8.onnx"),
        decoder=str(model_dir / "decoder-epoch-99-avg-1.onnx"),
        joiner=str(model_dir / "joiner-epoch-99-avg-1.int8.onnx"),
        tokens=str(model_dir / "tokens.txt"),
        num_threads=num_vcpus,
        sample_rate=16000,
        feature_dim=80,
        decoding_method="greedy_search",
    )
    print("X-ASR model loaded successfully!")
    return recognizer


def load_qwen3_model(model_name: str):
    """Load a Qwen3-ASR model (offline) via sherpa-onnx."""
    import sherpa_onnx

    print(f"Loading Qwen3-ASR model: {model_name}")
    model_dir = download_sherpa_repo(
        model_name,
        sentinel_files=(
            "conv_frontend.onnx",
            "encoder.int8.onnx",
            "decoder.int8.onnx",
            "tokenizer/vocab.json",
        ),
    )
    recognizer = sherpa_onnx.OfflineRecognizer.from_qwen3_asr(
        conv_frontend=str(model_dir / "conv_frontend.onnx"),
        encoder=str(model_dir / "encoder.int8.onnx"),
        decoder=str(model_dir / "decoder.int8.onnx"),
        tokenizer=str(model_dir / "tokenizer"),
        num_threads=num_vcpus,
    )
    print("Qwen3-ASR model loaded successfully!")
    return recognizer