Whisper JAX:让你的语音转文字功能速度快70倍!

RayRay
Whisper JAXJAX音频转录GPU批处理热门

项目简介

该存储库包含 OpenAI 的 Whisper 模型的优化 JAX代码,主要基于 Hugging Face Transformers Whisper 实现。与 OpenAI的 PyTorch 代码相比,Whisper JAX 的运行速度快了70 倍以上,使其成为可用的最快的 Whisper 运行。

JAX 代码在 CPU、GPU 和 TPU 上兼容,并且可以独立运行或作为推理端点。

安装

Whisper JAX 使用 Python 3.9 和 JAX 版本 0.4.5 进行了测试。安装假定您的设备上已安装最新版本的 JAX 包。您可以使用官方 JAX 安装指南来执行此操作:

https://github.com/google/jax#installation

一旦安装了适当版本的 JAX,就可以通过 pip 安装 Whisper JAX

pip install git+https://github.com/sanchit-gandhi/whisper-jax.git

要将 Whisper JAX 包更新到最新版本,只需运行:

pip install --upgrade --no-deps --force-reinstall git+https://github.com/sanchit-gandhi/whisper-jax.git

管道使用

运行 Whisper JAX 的推荐方式是通过FlaxWhisperPipline抽象类。此类处理所有必要的预处理和后处理,并包装生成方法以实现跨加速器设备的数据并行性。

Whisper JAX 利用 JAX 的pmap功能实现跨 GPU/TPU 设备的数据并行性。该函数在第一次调用时进行即时 (JIT)编译。此后,该函数将被缓存,使其能够以超快的速度运行:

from whisper_jax import FlaxWhisperPipline # instantiate pipeline pipeline = FlaxWhisperPipline("openai/whisper-large-v2") # JIT compile the forward call - slow, but we only do once text = pipeline("audio.mp3") # used cached function thereafter - super fast!! text = pipeline("audio.mp3")

半精度

通过在实例化管道时传递 dtype 参数,可以以半精度运行模型计算。通过以半精度存储中间张量,这将大大加快计算速度。模型权重的精度没有变化。

对于大多数 GPU,dtype 应设置为jnp.float16. 对于 A100 GPU 或 TPU,dtype 应设置为jnp.bfloat16:

from whisper_jax import FlaxWhisperPipline import jax.numpy as jnp # instantiate pipeline in bfloat16 pipeline = FlaxWhisperPipline("openai/whisper-large-v2", dtype=jnp.bfloat16)

批处理

Whisper JAX 还提供跨加速器设备批量处理单个音频输入的选项。音频首先被分成 30 秒的片段,然后将片段分派到模型进行并行转录。所得到的转录在边界处缝合在一起以给出单个、统一的转录。实际上,如果选择的批处理大小足够大,则与顺序转录音频样本相比,批处理提供了 10 倍的加速,并且 WER 1的损失不到 1%。

要启用批处理,请batch_size在实例化管道时传递参数:

from whisper_jax import FlaxWhisperPipline # instantiate pipeline with batching pipeline = FlaxWhisperPipline("openai/whisper-large-v2", batch_size=16)

任务

默认情况下,管道以所说的语言转录音频文件。对于语音翻译,请将参数设置 task为"translate":

# translate text = pipeline("audio.mp3", task="translate")

时间戳

FlaxWhisperPipline还支持时间戳预测。请注意,启用时间戳将需要对前向调用进行第二次 JIT 编译,包括时间戳输出:

# transcribe and return timestamps outputs = pipeline("audio.mp3", task="transcribe", return_timestamps=True) text = outputs["text"] # transcription chunks = outputs["chunks"] # transcription + timestamps

高级用法

更高级的用户可能希望探索不同的并行化技术。Whisper JAX 代码构建在T5x 代码库之上,这意味着它可以使用 T5x 分区约定的模型、激活和数据并行性来运行。要使用 T5x 分区,必须定义逻辑轴规则和模型分区数量。更多详细信息,用户可参考官方T5x分区指南:https://github.com/google-research/t5x/blob/main/docs/usage/partitioning.md

基准测试

我们将 Whisper JAX 与官方OpenAI 实现和🤗 Transformers 实现进行比较。我们在长度不断增加的音频样本上对模型进行基准测试,并报告 10 次重复运行的平均推理时间(以秒为单位)。对于所有三个系统,我们将预加载的音频文件传递给模型并测量前向传递的时间。将加载音频文件的任务留给系统会增加所有基准时间的相等偏移,因此加载和转录音频文件的实际时间将高于报告的数字。

OpenAI 和 Transformers 都在 GPU 上的 PyTorch 中运行。Whisper JAX 在 GPU 和 TPU 上的 JAX 中运行。OpenAI 按照说话的顺序依次转录音频。Transformers 和 Whisper JAX 都使用批处理算法,其中音频块被一起批处理并并行转录(请参阅批处理部分)。

表 :长度增加的音频文件的平均推理时间(以秒为单位)。GPU设备是单个A100 40GB GPU。TPU 设备是单个 TPU v4-8。

benchmark

项目链接

https://github.com/sanchit-gandhi/whisper-jax

编辑推荐精选

咔片PPT

咔片PPT

AI助力,做PPT更简单!

咔片是一款轻量化在线演示设计工具,借助 AI 技术,实现从内容生成到智能设计的一站式 PPT 制作服务。支持多种文档格式导入生成 PPT,提供海量模板、智能美化、素材替换等功能,适用于销售、教师、学生等各类人群,能高效制作出高品质 PPT,满足不同场景演示需求。

讯飞绘文

讯飞绘文

选题、配图、成文,一站式创作,让内容运营更高效

讯飞绘文,一个AI集成平台,支持写作、选题、配图、排版和发布。高效生成适用于各类媒体的定制内容,加速品牌传播,提升内容营销效果。

AI助手热门AI工具AI创作AI辅助写作讯飞绘文内容运营个性化文章多平台分发
材料星

材料星

专业的AI公文写作平台,公文写作神器

AI 材料星,专业的 AI 公文写作辅助平台,为体制内工作人员提供高效的公文写作解决方案。拥有海量公文文库、9 大核心 AI 功能,支持 30 + 文稿类型生成,助力快速完成领导讲话、工作总结、述职报告等材料,提升办公效率,是体制打工人的得力写作神器。

openai-agents-python

openai-agents-python

OpenAI Agents SDK,助力开发者便捷使用 OpenAI 相关功能。

openai-agents-python 是 OpenAI 推出的一款强大 Python SDK,它为开发者提供了与 OpenAI 模型交互的高效工具,支持工具调用、结果处理、追踪等功能,涵盖多种应用场景,如研究助手、财务研究等,能显著提升开发效率,让开发者更轻松地利用 OpenAI 的技术优势。

Hunyuan3D-2

Hunyuan3D-2

高分辨率纹理 3D 资产生成

Hunyuan3D-2 是腾讯开发的用于 3D 资产生成的强大工具,支持从文本描述、单张图片或多视角图片生成 3D 模型,具备快速形状生成能力,可生成带纹理的高质量 3D 模型,适用于多个领域,为 3D 创作提供了高效解决方案。

3FS

3FS

一个具备存储、管理和客户端操作等多种功能的分布式文件系统相关项目。

3FS 是一个功能强大的分布式文件系统项目,涵盖了存储引擎、元数据管理、客户端工具等多个模块。它支持多种文件操作,如创建文件和目录、设置布局等,同时具备高效的事件循环、节点选择和协程池管理等特性。适用于需要大规模数据存储和管理的场景,能够提高系统的性能和可靠性,是分布式存储领域的优质解决方案。

TRELLIS

TRELLIS

用于可扩展和多功能 3D 生成的结构化 3D 潜在表示

TRELLIS 是一个专注于 3D 生成的项目,它利用结构化 3D 潜在表示技术,实现了可扩展且多功能的 3D 生成。项目提供了多种 3D 生成的方法和工具,包括文本到 3D、图像到 3D 等,并且支持多种输出格式,如 3D 高斯、辐射场和网格等。通过 TRELLIS,用户可以根据文本描述或图像输入快速生成高质量的 3D 资产,适用于游戏开发、动画制作、虚拟现实等多个领域。

ai-agents-for-beginners

ai-agents-for-beginners

10 节课教你开启构建 AI 代理所需的一切知识

AI Agents for Beginners 是一个专为初学者打造的课程项目,提供 10 节课程,涵盖构建 AI 代理的必备知识,支持多种语言,包含规划设计、工具使用、多代理等丰富内容,助您快速入门 AI 代理领域。

AEE

AEE

AI Excel全自动制表工具

AEE 在线 AI 全自动 Excel 编辑器,提供智能录入、自动公式、数据整理、图表生成等功能,高效处理 Excel 任务,提升办公效率。支持自动高亮数据、批量计算、不规则数据录入,适用于企业、教育、金融等多场景。

UI-TARS-desktop

UI-TARS-desktop

基于 UI-TARS 视觉语言模型的桌面应用,可通过自然语言控制计算机进行多模态操作。

UI-TARS-desktop 是一款功能强大的桌面应用,基于 UI-TARS(视觉语言模型)构建。它具备自然语言控制、截图与视觉识别、精确的鼠标键盘控制等功能,支持跨平台使用(Windows/MacOS),能提供实时反馈和状态显示,且数据完全本地处理,保障隐私安全。该应用集成了多种大语言模型和搜索方式,还可进行文件系统操作。适用于需要智能交互和自动化任务的场景,如信息检索、文件管理等。其提供了详细的文档,包括快速启动、部署、贡献指南和 SDK 使用说明等,方便开发者使用和扩展。

下拉加载更多