test_models.py

import os
import io
import time
import sys
import numpy as np
from PIL import Image
import scipy.io.wavfile as wavfile

# Add parent directory to path if needed (though running from health-screener is recommended)
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

try:
    import local_mode
    import client
except ImportError as e:
    print(f"Import Error: {e}")
    print("Please run this script from the 'health-screener' directory.")
    sys.exit(1)


def generate_dummy_audio() -> bytes:
    """Generate 1 second of silent 16kHz audio as WAV bytes."""
    sr = 16000
    y = np.zeros(sr, dtype=np.int16)
    out_buf = io.BytesIO()
    wavfile.write(out_buf, sr, y)
    return out_buf.getvalue()


def generate_dummy_image() -> Image.Image:
    """Generate a simple 128x128 placeholder image."""
    return Image.new("RGB", (128, 128), color=(34, 197, 94))  # Premium green color


def print_section(title: str):
    print("\n" + "=" * 60)
    print(f" TEST: {title}")
    print("=" * 60)


def run_tests():
    print("============================================================")
    print("            ASHA Health Screener - Diagnostic Tool          ")
    print("============================================================")

    # 0. System Information Check
    print(f"Python Version : {sys.version.split()[0]}")
    import torch

    print(f"PyTorch Version: {torch.__version__}")
    cuda_available = torch.cuda.is_available()
    print(
        f"CUDA (GPU)     : {'Available (Using GPU)' if cuda_available else 'Not Available (Using CPU)'}"
    )
    if cuda_available:
        print(f"GPU Device     : {torch.cuda.get_device_name(0)}")
    print(f"CPU Threads    : {os.cpu_count()}")
    print("============================================================\n")

    # Set local mode environment variable to force local tests
    os.environ["LOCAL_MODE"] = "true"

    # --- TEST 1: Whisper Hindi Speech-to-Text ---
    print_section("ASR (Hindi Speech-to-Text) - Whisper Tiny")
    try:
        audio_bytes = generate_dummy_audio()
        print("Initializing local Whisper (tiny)...")
        start_time = time.time()
        # Run local transcription
        text = local_mode.local_transcribe(audio_bytes)
        duration = time.time() - start_time
        print(f"Result text: '{text}'")
        print(f"STATUS     : SUCCESS (Time taken: {duration:.2f}s)")
    except Exception as e:
        print("STATUS     : FAILED")
        print(f"Error      : {e}")

    # --- TEST 2: SmolVLM Vision-to-Text ---
    print_section("VLM (Image Observation) - SmolVLM-256M")
    try:
        img = generate_dummy_image()
        prompt = "Describe the color of this image."
        print("Initializing local SmolVLM...")
        start_time = time.time()
        description = local_mode.local_vision(img, prompt)
        duration = time.time() - start_time
        print(f"Result description: '{description}'")
        print(f"STATUS            : SUCCESS (Time taken: {duration:.2f}s)")
    except Exception as e:
        print("STATUS            : FAILED")
        print(f"Error             : {e}")

    # --- TEST 3: Nemotron LLM Triage ---
    print_section("LLM (Clinical Triage) - nvidia/Nemotron-Mini-4B-Instruct")
    try:
        system_prompt = 'Always respond with a valid JSON: {"status": "ok"}'
        user_message = "Perform diagnostic self-test."

        print("Initializing local Nemotron LLM...")
        start_time = time.time()
        response = local_mode.local_inference(
            system_prompt, user_message, max_tokens=100
        )
        duration = time.time() - start_time
        print(f"Result response:\n{response}")
        print(f"STATUS         : SUCCESS (Time taken: {duration:.2f}s)")
    except Exception as e:
        print("STATUS         : FAILED")
        print(f"Error          : {e}")

    # --- TEST 4: Hindi TTS (MMS-TTS) ---
    print_section("TTS (Hindi Speech Synthesis) - Facebook MMS-TTS")
    try:
        test_text = "नमस्ते, यह एक परीक्षण संदेश है।"
        print("Initializing local Hindi Vits Model...")
        start_time = time.time()
        audio_out = client.synthesize_hindi(test_text)
        duration = time.time() - start_time
        audio_size = len(audio_out) if audio_out else 0
        print(f"Audio Output size: {audio_size} bytes")
        if audio_size > 0:
            print(f"STATUS           : SUCCESS (Time taken: {duration:.2f}s)")
        else:
            print("STATUS           : FAILED (Returned empty audio)")
    except Exception as e:
        print("STATUS           : FAILED")
        print(f"Error            : {e}")

    # --- TEST 5: English TTS (Kokoro-v1.0) ---
    print_section("TTS (English Speech Synthesis) - Kokoro")
    try:
        test_text = "Hello, this is a diagnostic self-test."
        print("Initializing local Kokoro pipeline...")
        start_time = time.time()
        audio_out = client.synthesize_english(test_text)
        duration = time.time() - start_time
        audio_size = len(audio_out) if audio_out else 0
        print(f"Audio Output size: {audio_size} bytes")
        if audio_size > 0:
            print(f"STATUS           : SUCCESS (Time taken: {duration:.2f}s)")
        else:
            print("STATUS           : FAILED (Returned empty audio)")
    except Exception as e:
        print("STATUS           : FAILED")
        print(f"Error            : {e}")

    # --- TEST 6: Modal Cloud Connection ---
    print_section("Modal Cloud Integration (Online Mode)")
    # Temporarily restore online mode to check connection
    os.environ["LOCAL_MODE"] = "false"
    try:
        import modal

        print("Checking Modal setup/token configuration...")
        token_id = os.environ.get("MODAL_TOKEN_ID", "")
        token_secret = os.environ.get("MODAL_TOKEN_SECRET", "")
        print(f"MODAL_TOKEN_ID set     : {'Yes' if token_id else 'No'}")
        print(f"MODAL_TOKEN_SECRET set : {'Yes' if token_secret else 'No'}")

        # Test look up to check server availability
        print("Checking if village-health-screener app is deployed on Modal...")
        modal.Function.lookup("village-health-screener", "run_transcription")
        print("STATUS                 : Deployed (Modal Backend is active and reachable!)")
    except ImportError:
        print("STATUS                 : FAILED (Modal SDK is not installed)")
    except Exception as e:
        print("STATUS                 : OFFLINE/UNREACHABLE")
        print(f"Detail/Warning         : {e}")

    print("\n" + "=" * 60)
    print("                  DIAGNOSTICS COMPLETED                     ")
    print("=" * 60)


if __name__ == "__main__":
    run_tests()