跳到主要内容

MLflow 和 OpenAI Whisper 入门

下载此笔记本

在本教程中，了解 OpenAI Whisper（一个 ASR 系统）与 MLflow 的集成。

本教程将教你什么​

  • 使用 Whisper 模型建立一个音频转录 **pipeline**。
  • 使用 MLflow **记录**和管理 Whisper 模型。
  • 推断和理解 Whisper 模型的 **签名**。
  • **加载**并与 MLflow 中存储的 Whisper 模型进行交互。
  • 利用 MLflow 的 **pyfunc** 进行 Whisper 模型服务和转录任务。

什么是 Whisper？​

Whisper 是由 OpenAI 开发的一个多功能 ASR 模型，经过训练可实现高精度语音转文本转换。它因在多种口音和环境下的训练而脱颖而出，可通过 Transformers 库轻松使用。

为什么选择 MLflow 配合 Whisper？​

将 MLflow 与 Whisper 集成可以增强 ASR 模型管理。

  • 实验跟踪：有助于跟踪模型配置和性能，以获得最佳结果。
  • 模型管理：集中管理 Whisper 模型的不同版本，提高组织性和可访问性。
  • 可复现性：通过跟踪重现模型行为所需的所有组件，确保转录的一致性。
  • 部署：简化了 Whisper 模型在各种生产环境中的部署，确保高效应用。

对 Whisper 感兴趣？想了解更多关于 Whisper 为 ASR 领域带来的转录能力方面的重大突破，您可以阅读白皮书，并在 OpenAI 的研究网站上了解更多关于活跃开发和进展。

准备好提升您的语音转文本能力了吗？让我们一起探索使用 MLflow 和 Whisper 进行自动语音识别！

python
# Disable tokenizers warnings when constructing pipelines
%env TOKENIZERS_PARALLELISM=false

import warnings

# Disable a few less-than-useful UserWarnings from setuptools and pydantic
warnings.filterwarnings("ignore", category=UserWarning)
env: TOKENIZERS_PARALLELISM=false

环境设置和音频数据获取​

使用 Whisper 进行转录的初步步骤：获取 音频并设置 MLflow。

在开始使用 OpenAI Whisper 进行音频转录过程之前，需要完成一些准备步骤，以确保一切就绪，从而获得顺畅高效的转录体验。

音频获取​

第一步是获取要处理的音频文件。在本教程中，我们使用 NASA 提供的一个公开可用的音频文件。这个示例音频提供了一个实际的例子来演示 Whisper 的转录能力。

模型和 Pipeline 初始化​

我们从 Transformers 库加载 Whisper 模型及其分词器和特征提取器。这些组件对于处理音频数据并将其转换为 Whisper 模型可以理解和转录的格式至关重要。接下来，我们使用 Whisper 模型创建一个转录 pipeline。这个 pipeline 简化了将音频数据输入模型并获取转录的过程。

MLflow 环境设置​

除了模型和音频数据设置之外，我们还初始化了 MLflow 环境。MLflow 用于跟踪和管理我们的实验，提供了一种有组织的方式来记录转录过程和结果。

以下代码块涵盖了这些初始设置步骤，为我们使用 Whisper 模型进行音频转录任务奠定了基础。

python
import requests
import transformers

import mlflow

# Acquire an audio file that is in the public domain
resp = requests.get(
  "https://www.nasa.gov/wp-content/uploads/2015/01/590325main_ringtone_kennedy_WeChoose.mp3"
)
resp.raise_for_status()
audio = resp.content

# Set the task that our pipeline implementation will be using
task = "automatic-speech-recognition"

# Define the model instance
architecture = "openai/whisper-large-v3"

# Load the components and necessary configuration for Whisper ASR from the Hugging Face Hub
model = transformers.WhisperForConditionalGeneration.from_pretrained(architecture)
tokenizer = transformers.WhisperTokenizer.from_pretrained(architecture)
feature_extractor = transformers.WhisperFeatureExtractor.from_pretrained(architecture)
model.generation_config.alignment_heads = [[2, 2], [3, 0], [3, 2], [3, 3], [3, 4], [3, 5]]

# Instantiate our pipeline for ASR using the Whisper model
audio_transcription_pipeline = transformers.pipeline(
  task=task, model=model, tokenizer=tokenizer, feature_extractor=feature_extractor
)
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.

格式化转录输出​

在本节中，我们介绍了一个仅用于增强此 Jupyter notebook 演示中转录输出可读性的实用函数。请注意，此函数仅用于演示目的，不应包含在生产代码中，也不应用于本教程以外的任何其他用途。

format_transcription 函数接收一段长转录文本，并通过将其分割成句子并插入换行符来格式化它。这使得在 notebook 环境中打印输出时更容易阅读。

python
def format_transcription(transcription):
  """
  Function for formatting a long string by splitting into sentences and adding newlines.
  """
  # Split the transcription into sentences, ensuring we don't split on abbreviations or initials
  sentences = [
      sentence.strip() + ("." if not sentence.endswith(".") else "")
      for sentence in transcription.split(". ")
      if sentence
  ]

  # Join the sentences with a newline character
  return "
".join(sentences)

执行转录 Pipeline​

使用 Whisper pipeline 执行音频转录并查看输出。

设置好 Whisper 模型和音频转录 pipeline 后，下一步是处理音频文件以提取其转录。本教程的这一部分至关重要，因为它演示了 Whisper 模型在将口语转换为书面文本方面的实际应用。

转录过程​

下面的代码块将音频文件输入 pipeline，然后生成转录。先前定义的 format_transcription 函数通过使用句子分割和换行符来格式化输出，从而提高可读性。

预保存测试的重要性​

在 MLflow 中保存模型之前测试转录 pipeline 至关重要。此步骤可验证模型是否按预期工作，确保准确性和可靠性。这种验证可以避免部署后出现问题，并确认模型与训练数据的一致性。它还提供了一个基准，用于与模型从 MLflow 加载后的输出进行比较，从而确保性能一致。

执行以下代码以转录音频，并评估 Whisper 模型提供的转录的质量和准确性。

python
# Verify that our pipeline is capable of processing an audio file and transcribing it
transcription = audio_transcription_pipeline(audio)

print(format_transcription(transcription["text"]))
We choose to go to the moon in this decade and do the other things.
Not because they are easy, but because they are hard.
3, 2, 1, 0.
All engines running.
Liftoff.
We have a liftoff.
32 minutes past the hour.
Liftoff on Apollo 11.

模型签名和配置​

为 Whisper 生成模型签名，以了解其输入和输出数据需求。

模型签名对于定义 Whisper 模型输入和输出的架构至关重要，它明确了预期的数据类型和结构。此步骤可确保模型正确处理输入并输出结构化数据。

处理不同的音频格式​

虽然默认签名涵盖了二进制音频数据，但 transformers flavor 支持多种格式，包括 numpy 数组和基于 URL 的输入。这种灵活性允许 Whisper 从各种来源进行转录，尽管此处未演示基于 URL 的转录。

模型配置​

设置模型配置涉及音频处理的参数，例如分块和步长长度。这些设置可以根据不同的转录需求进行调整，以提高 Whisper 在特定场景下的性能。

运行下一个代码块以推断模型签名并配置关键参数，使 Whisper 的功能与您项目的需求保持一致。

python
# Specify parameters and their defaults that we would like to be exposed for manipulation during inference time
model_config = {
  "chunk_length_s": 20,
  "stride_length_s": [5, 3],
}

# Define the model signature by using the input and output of our pipeline, as well as specifying our inference parameters that will allow for those parameters to
# be overridden at inference time.
signature = mlflow.models.infer_signature(
  audio,
  mlflow.transformers.generate_signature_output(audio_transcription_pipeline, audio),
  params=model_config,
)

# Visualize the signature
signature
inputs: 
[binary]
outputs: 
[string]
params: 
['chunk_length_s': long (default: 20), 'stride_length_s': long (default: [5, 3]) (shape: (-1,))]

创建实验​

我们创建一个新的 MLflow 实验，以便我们要将模型记录到的运行不会记录到默认实验，而是具有其自己的上下文相关条目。

python
# If you are running this tutorial in local mode, leave the next line commented out.
# Otherwise, uncomment the following line and set your tracking uri to your local or remote tracking server.

# mlflow.set_tracking_uri("http://127.0.0.1:8080")

mlflow.set_experiment("Whisper Transcription ASR")
<Experiment: artifact_location='file:///Users/benjamin.wilson/repos/mlflow-fork/mlflow/docs/source/llms/transformers/tutorials/audio-transcription/mlruns/864092483920291025', creation_time=1701294423466, experiment_id='864092483920291025', last_update_time=1701294423466, lifecycle_stage='active', name='Whisper Transcription ASR', tags={}>

使用 MLflow 记录模型​

了解如何使用 MLflow 记录 Whisper 模型及其配置。

在 MLflow 中记录 Whisper 模型是捕获模型重现、共享和部署的关键信息的重要步骤。此过程包括

模型记录的关键组件​

  • 模型信息：包括模型、其签名和输入示例。
  • 模型配置：为模型设置的任何特定参数，例如分块长度或步长长度。

使用 MLflow 的 log_model 函数​

此函数在 MLflow 运行中使用，用于记录模型及其配置。它确保记录了模型使用的所有必要组件。

执行下一个单元格中的代码将在当前的 MLflow 实验中记录 Whisper 模型。这包括将模型存储在指定的 artifact 路径中，并记录在推理过程中将应用的默认配置。

python
# Log the pipeline
with mlflow.start_run():
  model_info = mlflow.transformers.log_model(
      transformers_model=audio_transcription_pipeline,
      name="whisper_transcriber",
      signature=signature,
      input_example=audio,
      model_config=model_config,
      # Since MLflow 2.11.0, you can save the model in 'reference-only' mode to reduce storage usage by not saving
      # the base model weights but only the reference to the HuggingFace model hub. To enable this, uncomment the
      # following line:
      # save_pretrained=False,
  )

加载和使用模型 Pipeline​

探索如何从 MLflow 加载和使用 Whisper 模型 pipeline。

在 MLflow 中记录 Whisper 模型后，下一步是加载和使用它进行推理。此过程可确保我们记录的模型按预期运行，并可有效地用于音频转录等任务。

加载模型​

模型使用 MLflow 的 load_model 函数以其原生格式加载。此步骤验证模型在 MLflow 中记录后是否可以无缝检索和使用。

使用加载的模型​

加载后，模型即可进行推理。我们通过将 MP3 音频文件传递给模型并获取其转录来演示这一点。此测试是模型记录后功能的实际演示。

此步骤是更复杂的部署场景之前的验证形式。确保模型在其原生格式中正确运行有助于故障排除，并简化部署过程，特别是对于像 Whisper 这样庞大而复杂的模型。

python
# Load the pipeline in its native format
loaded_transcriber = mlflow.transformers.load_model(model_uri=model_info.model_uri)

# Perform transcription with the native pipeline implementation
transcription = loaded_transcriber(audio)

print(f"
Whisper native output transcription:
{format_transcription(transcription['text'])}")
2023/11/30 12:51:43 INFO mlflow.transformers: 'runs:/f7503a09d20f4fb481544968b5ed28dd/whisper_transcriber' resolved as 'file:///Users/benjamin.wilson/repos/mlflow-fork/mlflow/docs/source/llms/transformers/tutorials/audio-transcription/mlruns/864092483920291025/f7503a09d20f4fb481544968b5ed28dd/artifacts/whisper_transcriber'
Loading checkpoint shards:   0%|          | 0/13 [00:00<?, ?it/s]
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
Whisper native output transcription:
We choose to go to the moon in this decade and do the other things.
Not because they are easy, but because they are hard.
3, 2, 1, 0.
All engines running.
Liftoff.
We have a liftoff.
32 minutes past the hour.
Liftoff on Apollo 11.

使用 Pyfunc Flavor 进行推理​

了解 MLflow 的 pyfunc flavor 如何促进灵活的模型部署。

MLflow 的 pyfunc flavor 为模型推理提供了一个通用接口，在各种机器学习框架和部署环境中提供灵活性。此功能对于部署原始框架可能不可用或需要更灵活接口的模型非常有用。

使用 Pyfunc 加载和预测​

以下代码说明了如何将 Whisper 模型加载为 pyfunc 并用于预测。此方法突出了 MLflow 在各种场景下适应和部署模型的能力。

输出格式注意事项​

请注意使用 pyfunc 与原生格式相比，输出格式的差异。pyfunc 输出符合标准的 pyfunc 输出签名，通常表示为 List[str] 类型，与更广泛的 MLflow 模型输出标准保持一致。

python
# Load the saved transcription pipeline as a generic python function
pyfunc_transcriber = mlflow.pyfunc.load_model(model_uri=model_info.model_uri)

# Ensure that the pyfunc wrapper is capable of transcribing passed-in audio
pyfunc_transcription = pyfunc_transcriber.predict([audio])

# Note: the pyfunc return type if `return_timestamps` is set is a JSON encoded string.
print(f"
Pyfunc output transcription:
{format_transcription(pyfunc_transcription[0])}")
Loading checkpoint shards:   0%|          | 0/13 [00:00<?, ?it/s]
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
2023/11/30 12:52:02 WARNING mlflow.transformers: params provided to the `predict` method will override the inference configuration saved with the model. If the params provided are not valid for the pipeline, MlflowException will be raised.
Pyfunc output transcription:
We choose to go to the moon in this decade and do the other things.
Not because they are easy, but because they are hard.
3, 2, 1, 0.
All engines running.
Liftoff.
We have a liftoff.
32 minutes past the hour.
Liftoff on Apollo 11.

教程总结​

在本教程中，我们探讨了如何

  • 使用 OpenAI Whisper 模型设置音频转录 pipeline。
  • 格式化和准备音频数据进行转录。
  • 使用 MLflow 记录、加载和使用模型，同时利用原生和 pyfunc flavors 进行推理。
  • 格式化输出以便于阅读并在 Jupyter Notebook 环境中使用。

我们已经看到了使用 MLflow 管理机器学习生命周期的好处，包括实验跟踪、模型版本控制、可复现性和部署。通过将 MLflow 与 Transformers 库集成，我们简化了使用最先进的 NLP 模型的工作流程，使得跟踪、管理和部署尖端 NLP 应用程序更加容易。