tests/test_simple.py

#!/usr/bin/env python3
"""
Teste Simples - Pipeline LLaMA-Omni2
=====================================
Teste com áudio real em português
"""

import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import torch
import numpy as np
import time
from pipelines.llama_omni2_oficial import LLaMAOmni2Correct
from gtts import gTTS
import tempfile
import os
import soundfile as sf

print("\n" + "="*60)
print("🎤 TESTE SIMPLES - LLaMA-Omni2")
print("="*60)

# Verificar GPU
if not torch.cuda.is_available():
    print("❌ GPU não disponível!")
    exit()

print(f"✅ GPU: {torch.cuda.get_device_name(0)}")
print(f"💾 Memória: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")

# Criar áudio REAL da pergunta
pergunta = "Qual é a capital do Brasil?"
print(f"\n🎤 Criando áudio REAL da pergunta: '{pergunta}'")

# Gerar áudio com gTTS
tts = gTTS(text=pergunta, lang="pt", slow=False)
with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as f:
    tts.save(f.name)
    temp_mp3 = f.name

# Converter para 16kHz
data, sr = sf.read(temp_mp3)
if sr != 16000:
    import librosa
    data = librosa.resample(data, orig_sr=sr, target_sr=16000)

os.remove(temp_mp3)
audio = data.astype(np.float32)

print(f"   ✅ Áudio criado: {len(audio)/16000:.1f}s")

# Carregar modelo na GPU
print("\n📦 Carregando modelo na GPU...")
inicio = time.time()
model = LLaMAOmni2Correct(device="cuda")
print(f"⏱️ Tempo de carga: {time.time() - inicio:.1f}s")

# Warmup
print("\n🔥 Warmup...")
warmup_audio = np.random.randn(16000).astype(np.float32) * 0.01
model.process(warmup_audio)

# Teste real com áudio da pergunta
print("\n⚡ Processando pergunta REAL:")
print(f"   🎤 PERGUNTA: '{pergunta}'")
print("   ⏳ Processando...")

inicio = time.time()
resposta, audio_resposta = model.process(audio)
tempo = time.time() - inicio

print("\n" + "="*60)
print("📊 RESULTADO:")
print("="*60)
print(f"❓ PERGUNTA: {pergunta}")
print(f"💬 RESPOSTA: {resposta if resposta else '(vazio)'}")
print(f"⏱️ TEMPO GPU: {tempo:.2f}s")

# Verificar coerência
if resposta:
    resposta_lower = resposta.lower()
    if any(x in resposta_lower for x in ["brasília", "brasilia", "capital", "brazil"]):
        print("✅ RESPOSTA COERENTE!")
    else:
        print("⚠️ Resposta não menciona Brasília")
        
print("="*60)