将 MUSE 从 TensorFlow 转换为 PyTorch

本仓库包含使用 PyTorch 调用 TF Hub 上的 MUSE（多语言通用句子编码器）Transformer 模型的代码。

[!IMPORTANT] PyTorch 模型不仅可以用于推理，还可以用于额外的训练和微调！

了解更多项目信息：GitHub。

使用方法

该模型可直接通过 torch 在 HF Models 中获取（目前不支持 transformers 库的原生调用）。

模型初始化和使用代码：

import torch
from functools import partial
from architecture import MUSE
from tokenizer import get_tokenizer, tokenize

PATH_TO_PT_MODEL = "model.pt"
PATH_TO_TF_MODEL = "universal-sentence-encoder-multilingual-large-3"

tokenizer = get_tokenizer(PATH_TO_TF_MODEL)
tokenize = partial(tokenize, tokenizer=tokenizer)

model_torch = MUSE(
    num_embeddings=128010,
    embedding_dim=512,
    d_model=512,
    num_heads=8,
)
model_torch.load_state_dict(
    torch.load(PATH_TO_PT_MODEL)
)

sentence = "Hello, world!"
res = model_torch(tokenize(sentence))

[!NOTE] 目前，分词使用的是原始 TF Hub 模型的检查点，因此在上面的代码中需要加载它。