模型管理

为 Firefox 准备模型

可用于 Firefox 的模型应该具有不同量化级别的 ONNX 权重。

为了确保我们与 Transformers.js 兼容,我们使用该项目提供的转换脚本,该脚本会检查模型架构是否有效并已过测试。

为此,请按照以下步骤操作

  • 确保您的模型已在 Hugging Face 上发布,并使用 PyTorch 或 SafeTensor 权重。

  • 克隆 https://github.com/xenova/transformers.js 并检出 v3 分支。

  • 进入 scripts/ 目录。

  • 在那里创建一个虚拟环境并从本地 requirements.txt 文件安装依赖项。

然后您可以运行

python convert.py --model_id organizationId/modelId --quantize --modes fp16 q8 q4 --task the-inference-task

您将在 models/organizationId/modelId 中获得一个新目录,其中包含 onnx 目录和其他文件。将所有内容上传到 Hugging Face。

恭喜!您已拥有一个与 Firefox 兼容的模型。您现在可以在 about:inference 中尝试它。

请注意,对于具有两个文件的编码器-解码器模型,您可能需要将 decoder_model_quantized.onnx 重命名为 decoder_model_merged_quantized.onnx,并对 fp16、q4 版本进行类似更改。您无需重命名编码器模型。

生命周期

当 Firefox 使用模型时,它将

  1. 读取存储在远程设置中的元数据。

  2. 从我们的中心下载模型文件。

  3. 将文件存储在 IndexDB 中。

1. 远程设置

我们在远程设置中拥有两个集合。

  • ml-onnx-runtime:提供运行推理平台所需的所有 WASM 文件。

  • ml-inference-options:为每个 taskId 提供运行选项列表,例如 modelId

运行推理 API 将根据需要下载 WASM 文件,然后查看 ml-inference-options 中是否存在该任务的条目,以获取选项。这使我们能够为每个任务设置默认运行选项。

这也是我们如何在不更改 Firefox 代码的情况下更新模型的方法:在远程设置中为模型设置新的版本将触发用户的重新下载。

ml-inference-options 中的记录由 featureId 唯一标识。未提供时,将回退到 taskName。此集合将提供该功能所需的所有选项。

例如,PDF.js 图像转文本记录为

{
"featureId": "pdfjs-alt-text"
"dtype":"q8",
"modelId":"mozilla/distilvit",
"taskName":"image-to-text",
"processorId":"mozilla/distilvit",
"tokenizerId":"mozilla/distilvit",
"modelRevision":"v0.5.0",
"processorRevision":"v0.5.0"
}

如果您在 Firefox 中添加新的推理调用,请在 FEATURES 中创建一个新的唯一 featureId,并在 ml-inference-options 中添加一个包含任务设置的记录。

通过这样做,您将能够使用此简单的调用创建引擎

const engine = await createEngine({featureId: "pdfjs-alt-text"});

2. 模型中心

我们的模型中心遵循与 Hugging Face 相同的结构,每个模型的文件都在唯一的 URL 下。

https://model-hub.mozilla.org/<organization>/<model>/<revision>/<path>

其中:- organizationname 是模型 ID。例如“ mozilla/distivit” - revision 是分支或版本 - path 是文件的路径。

从中心下载的模型文件存储在 IndexDB 中,因此用户无需再次下载它们。

模型文件

模型包含几个文件,例如其配置、分词器、训练元数据和权重。

以下是您将遇到的最常见文件。

1. 模型权重

  • pytorch_model.bin:包含 PyTorch 模型的模型权重。它是一个序列化文件,保存神经网络的参数。

  • tf_model.h5:TensorFlow 版本的模型权重。

  • flax_model.msgpack:对于使用 Flax 框架构建的模型,此文件以 JAX 和 Flax 使用的格式包含模型权重。

  • onnx:包含不同量化级别的 ONNX 权重文件的子目录。 **它们是我们平台使用的文件**

2. 模型配置

config.json 文件包含模型架构所需的所有配置,例如层数、隐藏单元、注意力头、激活函数等等。这允许 Hugging Face 库完全按照定义重建模型。

3. 分词器文件

  • vocab.txtvocab.json:词汇文件,将标记(单词、子词或字符)映射到 ID。不同的分词器(BERT、GPT-2 等)将具有不同的格式。

  • tokenizer.json:存储完整的分词器配置和映射。

  • tokenizer_config.json:此文件包含分词器使用的特定设置,例如它是否区分大小写或它使用的特殊标记(例如 [CLS]、[SEP] 等)。

4. 预处理文件

  • special_tokens_map.json:将特殊标记(如填充、CLS、SEP 等)映射到分词器使用的标记 ID。

  • added_tokens.json:如果在原始词汇之外添加了任何其他标记(如自定义标记或特定领域标记),则将它们存储在此文件中。

5. 训练元数据

  • training_args.bin:包含训练期间使用的参数,例如学习率、批大小和其他超参数。此文件允许更轻松地复制训练过程。

  • trainer_state.json:捕获训练器的状态,例如时期信息和优化器状态,这对于恢复训练很有用。

  • optimizer.pt:存储 PyTorch 模型的优化器状态,允许从中断处恢复训练。

6. 模型卡片

README.mdmodel_card.json。模型卡片提供有关模型的文档,包括其预期用途、训练数据、性能指标、道德考虑因素和任何限制的详细信息。这可以是 README.md 或结构化为 model_card.json

7. 分词和特征提取文件

  • merges.txt:对于字节对编码 (BPE) 分词器,此文件包含用于将单词拆分为子词的合并操作。

  • preprocessor_config.json:包含在传递到模型之前应用于输入的任何预处理或特征提取步骤的配置详细信息。

版本控制

revision 字段用于确定应从中心下载哪个版本的模型。您可以先提供 main 分支的服务,但是一旦发布模型,您应该开始对其进行版本控制。

我们使用的 version 方案比较宽松。它可以是 main 或遵循扩展 semver 的版本。

[v]MAJOR.MINOR[.PATCH][.(alpha|beta|pre|post|rc|)NUMBER]

我们不提供任何排序函数。

示例

  • v1.0

  • v2.3.4

  • 1.2.1

  • 1.0.0-beta1

  • 1.0.0.alpha2

  • 1.0.0.rc1

要对模型进行版本控制,您可以在 Hugging Face 上使用 git tag v1.0 && git push –tags 推送标签,并在 GCP 存储桶中创建一个新目录,并将模型文件复制到该目录。