tests/test_llama_omni_api.py

#!/usr/bin/env python3
"""
Teste completo para o LLaMA-Omni2-0.5B no Hugging Face
Este script pode:
1. Transcrever áudio localmente e enviar para o modelo
2. Enviar texto diretamente para o modelo
3. Facilita o teste manual com interface web
4. Testar a API diretamente de modo programático
"""

import os
import sys
import time
import argparse
import requests
import subprocess
import webbrowser
from pathlib import Path
from gradio_client import Client

# Configurações padrão
DEFAULT_API_URL = "https://marcosremar2-llama-omni.hf.space"
DEFAULT_OUTPUT_DIR = "./output"
MODEL_NAME = "LLaMA-Omni2-0.5B"

def transcribe_audio_locally(audio_file_path):
    """
    Transcreve áudio localmente usando whisper se disponível
    Caso contrário, retorna uma mensagem padrão
    """
    try:
        # Tenta usar whisper CLI se disponível
        result = subprocess.run(
            ["whisper", audio_file_path, "--model", "tiny", "--output_format", "txt"],
            capture_output=True,
            text=True,
            check=True
        )
        transcript_file = f"{os.path.splitext(audio_file_path)[0]}.txt"
        if os.path.exists(transcript_file):
            with open(transcript_file, "r") as f:
                transcript = f.read().strip()
            print(f"Transcrição: {transcript}")
            return transcript
    except (subprocess.CalledProcessError, FileNotFoundError) as e:
        print(f"Whisper não disponível ou erro: {e}")
    
    # Mensagem padrão
    print("Usando mensagem de teste padrão, já que whisper não está disponível")
    return f"Olá, estou testando o modelo {MODEL_NAME}. Você pode me responder em português?"

def check_url_accessibility(url):
    """Verifica se a URL é acessível"""
    try:
        response = requests.get(url, timeout=10)
        if response.status_code == 200:
            return True
        else:
            print(f"URL retornou código {response.status_code}")
            return False
    except Exception as e:
        print(f"Erro ao acessar URL: {e}")
        return False

def save_text_to_file(text, output_dir, filename="text.txt"):
    """Salva texto em arquivo para fácil cópia"""
    os.makedirs(output_dir, exist_ok=True)
    filepath = os.path.join(output_dir, filename)
    
    with open(filepath, "w") as f:
        f.write(text)
    
    print(f"Texto salvo em: {filepath}")
    return filepath

def test_api_programmatically(api_url, text_input, output_dir=DEFAULT_OUTPUT_DIR):
    """
    Testa a API do modelo programaticamente enviando um texto
    e salvando a resposta
    """
    output_path = os.path.join(output_dir, f"response_{int(time.time())}.txt")
    os.makedirs(output_dir, exist_ok=True)
    
    print(f"Testando API em: {api_url}")
    print(f"Texto de entrada: {text_input[:50]}..." if len(text_input) > 50 else f"Texto de entrada: {text_input}")
    
    try:
        # Conecta ao app Gradio com timeout aumentado
        client = Client(
            api_url,
            httpx_kwargs={"timeout": 300.0}  # 5 minutos de timeout
        )
        
        print("Conectado à API com sucesso")
        
        # Lista os endpoints disponíveis
        print("Endpoints disponíveis:")
        client.view_api()
        
        # Envia o prompt para o modelo
        print(f"\nUsando endpoint de geração de texto (/lambda_1)...")
        print(f"Enviando prompt: '{text_input[:50]}...'")
        job = client.submit(
            text_input,
            MODEL_NAME,
            api_name="/lambda_1"
        )
        
        print("Requisição enviada, aguardando resposta...")
        result = job.result()
        print(f"Resposta recebida (tamanho: {len(str(result))} caracteres)")
        
        # Salva a resposta em arquivo
        with open(output_path, "w") as f:
            f.write(str(result))
        
        print(f"Resposta salva em: {output_path}")
        
        # Tenta obter informações do modelo
        try:
            print("\nConsultando informações do modelo...")
            model_info = client.submit(api_name="/lambda").result()
            print(f"Informações do modelo: {model_info}")
        except Exception as model_error:
            print(f"Erro ao obter informações do modelo: {str(model_error)}")
        
        return True, result
        
    except Exception as e:
        print(f"Erro durante requisição à API: {str(e)}")
        print("Isso pode ocorrer porque o Space está dormindo e precisa de tempo para iniciar.")
        print("Tente acessar o Space diretamente primeiro: " + api_url)
        print(f"\nNota: Esta API é para o modelo {MODEL_NAME} e não processa áudio diretamente.")
        print("Para trabalhar com áudio, você precisaria primeiro transcrever o áudio usando Whisper,")
        print("e então enviar o texto transcrito para esta API.")
        return False, None

def test_manual_interface(api_url, text_input, output_dir=DEFAULT_OUTPUT_DIR):
    """
    Prepara o teste manual do modelo via interface web:
    1. Salva o texto em arquivo para fácil cópia
    2. Abre a interface web para teste manual
    """
    # Verifica se a URL é acessível
    print(f"Verificando acessibilidade de {api_url}...")
    if not check_url_accessibility(api_url):
        print(f"Aviso: {api_url} não está acessível. Teste manual pode não ser possível.")
    
    # Salva o texto em arquivo para fácil cópia
    transcript_file = save_text_to_file(text_input, output_dir, "transcription.txt")
    
    # Instruções para teste manual
    print("\n" + "=" * 50)
    print(f"INSTRUÇÕES PARA TESTE MANUAL DO {MODEL_NAME}")
    print("=" * 50)
    print(f"1. O texto foi salvo em: {transcript_file}")
    print(f"2. Abrindo {api_url} no navegador...")
    print("3. Copie o texto do arquivo salvo e cole no campo 'Input Text'")
    print("4. Clique no botão 'Generate'")
    print("5. Quando receber a resposta, copie e salve para seus registros")
    print("=" * 50 + "\n")
    
    # Abre a URL no navegador padrão
    try:
        webbrowser.open(api_url)
        return True
    except Exception as e:
        print(f"Erro ao abrir navegador: {e}")
        print(f"Por favor, visite manualmente: {api_url}")
        return False

def main():
    parser = argparse.ArgumentParser(description=f"Teste para {MODEL_NAME} no Hugging Face")
    parser.add_argument("--api-url", type=str, default=DEFAULT_API_URL,
                        help=f"URL da interface Gradio (padrão: {DEFAULT_API_URL})")
    parser.add_argument("--audio-file", type=str, default="test.mp3",
                        help="Caminho para o arquivo de áudio a ser transcrito localmente (opcional)")
    parser.add_argument("--text", type=str, default=None,
                        help="Texto para usar diretamente (em vez de transcrever áudio)")
    parser.add_argument("--output-dir", type=str, default=DEFAULT_OUTPUT_DIR,
                        help="Diretório para salvar a transcrição e respostas")
    parser.add_argument("--mode", type=str, choices=["api", "manual", "both"], default="both",
                        help="Modo de teste: api (programático), manual (navegador) ou both (ambos)")
    args = parser.parse_args()
    
    # Converte caminhos relativos para absolutos
    if args.audio_file and not os.path.isabs(args.audio_file):
        if not os.path.exists(args.audio_file):
            script_dir = os.path.dirname(os.path.abspath(__file__))
            args.audio_file = os.path.join(script_dir, args.audio_file)
    
    if args.output_dir and not os.path.isabs(args.output_dir):
        script_dir = os.path.dirname(os.path.abspath(__file__))
        args.output_dir = os.path.join(script_dir, args.output_dir)
    
    # Obtém texto de entrada da transcrição ou do parâmetro
    input_text = args.text
    if not input_text and args.audio_file:
        if os.path.exists(args.audio_file):
            input_text = transcribe_audio_locally(args.audio_file)
        else:
            print(f"Arquivo de áudio não encontrado: {args.audio_file}")
            input_text = f"Olá, estou testando o modelo {MODEL_NAME}. Você pode me responder em português?"
    if not input_text:
        input_text = f"Olá, estou testando o modelo {MODEL_NAME}. Você pode me responder em português?"
    
    print(f"Texto de entrada: {input_text}")
    
    # Executa os testes conforme o modo selecionado
    success = True
    if args.mode in ["api", "both"]:
        api_success, _ = test_api_programmatically(args.api_url, input_text, args.output_dir)
        success = success and api_success
    
    if args.mode in ["manual", "both"]:
        manual_success = test_manual_interface(args.api_url, input_text, args.output_dir)
        success = success and manual_success
    
    # Sai com código apropriado
    sys.exit(0 if success else 1)

if __name__ == "__main__":
    main()