Whisper JAX é uma implementação do modelo Whisper da OpenAI em JAX, capaz de acelerar a transcrição de fala em até 70 vezes em TPUs. Esta ferramenta torna a transcrição de áudio por IA extremamente rápida e simples, ideal para desenvolvedores, pesquisadores e usuários comuns.

Por que escolher o Whisper JAX?

É atualmente a implementação mais rápida do Whisper, superando a versão PyTorch da OpenAI em até 70 vezes em TPUs. Compatível com CPU, GPU e TPU, oferece alta velocidade em diferentes ambientes de hardware.

• Velocidade impressionante: até 70x mais rápido que o Whisper original
• Suporte a múltiplos hardwares: CPU, GPU e TPU
• Processamento em lote: processamento paralelo de blocos de áudio, eficiência até 10x maior
• Precisão mantida: perda inferior a 1% na acurácia mesmo com ganho de velocidade

Instalação simplificada

Certifique-se de ter o JAX instalado e execute o comando abaixo para instalar o Whisper JAX:

pip install git+https://gitcode.com/gh_mirrors/wh/whisper-jax.git

Para atualizra para a versão mais recente:

pip install --upgrade --no-deps --force-reinstall git+https://gitcode.com/gh_mirrors/wh/whisper-jax.git

Primeiros passos em 5 minutos

1. Importar e criar o pipeline

from whisper_jax import FlaxWhisperPipeline

transcriber = FlaxWhisperPipeline("openai/whisper-large-v2")

2. Compilação inicial (apenas uma vez)

A primeira execução realiza compilação JIT, que pode ser mais lenta, mas é necessária apenas uma vez:

# Compilação JIT - feita uma única vez
resultado = transcriber("audio.mp3")

3. Transcrições subsequentes (ultrarrápidas)

Após a compilação, as chamadas usam funções em cache e são extremamente rápidas:

resultado = transcriber("audio.mp3")

Configurações avançadas

Aceleração com meia precisão

from whisper_jax import FlaxWhisperPipeline
import jax.numpy as jnp

transcriber = FlaxWhisperPipeline("openai/whisper-large-v2", dtype=jnp.bfloat16)

Otimização com processamento em lote

transcriber = FlaxWhisperPipeline("openai/whisper-large-v2", batch_size=16)

Tradução de fala

resultado = transcriber("audio.mp3", task="translate")

Obter timestamps

saida = transcriber("audio.mp3", task="transcribe", return_timestamps=True)
texto = saida["text"]
blocos = saida["chunks"]

Exemplo de configuração completa

from whisper_jax import FlaxWhisperPipeline
import jax.numpy as jnp

transcriber = FlaxWhisperPipeline(
    "openai/whisper-large-v2", 
    dtype=jnp.bfloat16, 
    batch_size=16
)

saida = transcriber("audio.mp3", task="transcribe", return_timestamps=True)

Comparativo de desempenho

Criar seu próprio endpoint de transcrição

Clone o repositório e instale as dependências:

git clone https://gitcode.com/gh_mirrors/wh/whisper-jax
cd whisper-jax
pip install -e .["endpoint"]

Inicie o aplicativo Gradio:

python app/app.py

Envie requisições de transcrição via cliente:

from gradio_client import Client

API_URL = "http://SEU-IP-DISPOSITIVO:7860/"
client = Client(API_URL)

def transcrever_audio(caminho_audio):
    texto, tempo = client.predict(
        caminho_audio,
        "transcribe",
        False,
        api_name="/predict_1"
    )
    return texto

resultado = transcrever_audio("seu_audio.mp3")

Estrutura do projeto

• whisper_jax/ – Módulo principal
  • modeling_flax_whisper.py – Implementação do modelo Flax Whisper
  • pipeline.py – Classe principal FlaxWhisperPipeline
  • partitioner.py – Implementação do particionador
  • layers.py – Definição das camadas neurais
  • train_state.py – Gerenciamento do estado de treinamento
• app/ – Arquivos da aplicação
  • app.py – Arquivo principal do Gradio
  • run_app.sh – Script de inciialização
  • monitor.sh – Script de monitoramento
• benchmarks/ – Scripts de teste de desempenho

Dicas e melhores práticas

Escolha do tamanho do modelo

Whisper JAX suporta todos os tamanhos do Whisper:

• tiny (39M) – rápido e leve, para tempo real
• base (74M) – equilíbrio entre velocidade e precisão
• small (244M) – recomendado para produção
• medium (769M) – alta precisão
• large-v2 (1550M) – melhor precisão

Otimização de memória

• Use meia precisão (float16/bfloat16) para reduzir consumo
• Ajuste batch_size conforme o hardware
• Para áudios longos, utilize processamento em blocos

Tratamento de erros

try:
    resultado = transcriber("audio.mp3")
except Exception as erro:
    print(f"Falha na transcrição: {erro}")