MLflow 和 OpenAI Whisper 入门

下载此笔记本

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

本教程内容

• 使用 Whisper 模型建立一个音频转录管道。
• 使用 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 进行自动语音识别！

# 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 的转录功能。

模型和管道初始化

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

MLflow 环境设置

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

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

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 函数接收一个长的转录文本字符串，通过将其分割成句子并插入换行符来对其进行格式化。这样可以提高在笔记本环境中打印输出时的可读性。

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)

执行转录管道

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

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

转录过程

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

预保存测试的重要性

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

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

# 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 的功能与您项目的需求保持一致。

# 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 实验，以便我们要将模型记录到的运行不会记录到默认实验，而是具有其自己的上下文相关条目。

# 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 模型。这包括将模型存储在指定的工件路径中，并记录推理过程中将应用的默认配置。

# 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,
  )

加载和使用模型管道

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

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

加载模型

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

使用加载的模型

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

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

# 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 模型输出标准保持一致。

# 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 模型设置音频转录管道。
• 格式化和准备音频数据以进行转录。
• 使用 MLflow 记录、加载和使用模型，利用原生和 pyfunc flavor 进行推理。
• 格式化输出，以提高可读性并在 Jupyter Notebook 环境中使用。

我们已经看到了 MLflow 在管理机器学习生命周期（包括实验跟踪、模型版本控制、可复现性和部署）方面的优势。通过将 MLflow 与 Transformers 库集成，我们简化了使用最先进的 NLP 模型的过程，从而更轻松地跟踪、管理和部署最先进的 NLP 应用。