如何基于TF-Hub预训练模型计算WMD并转换为spaCy/word2vec格式？

嗨，我来帮你梳理几个可行的解决方案，不用纠结于格式转换，或者也能轻松完成格式适配～

其实WMD的核心是词嵌入的距离计算，不一定非要把TF Hub模型转成spaCy/word2vec格式。我们可以直接利用USE生成词嵌入，再结合Gensim的WMD工具来计算：

import tensorflow as tf
import tensorflow_hub as hub
from gensim.models import KeyedVectors

# 加载USE模型
embed = hub.load("https://tfhub.dev/google/universal-sentence-encoder/4")

# 生成指定词汇的USE嵌入，并封装成Gensim可识别的格式
def build_use_keyed_vectors(vocab_list):
    # 把每个词作为独立句子传入USE，得到嵌入
    embeddings = embed(vocab_list).numpy()
    # 构建Gensim的KeyedVectors对象
    kv = KeyedVectors(vector_size=512)  # USE的嵌入维度是512
    kv.add_vectors(vocab_list, embeddings)
    return kv

# 示例句子
sent1 = "cat sat on the mat".split()
sent2 = "dog lay on the rug".split()

# 收集所有需要计算的词汇
all_words = list(set(sent1 + sent2))
# 生成嵌入表
use_kv = build_use_keyed_vectors(all_words)

# 计算WMD
wmd_distance = use_kv.wmdistance(sent1, sent2)
print(f"Word Mover's Distance: {wmd_distance}")

这个方法跳过了格式转换步骤，直接利用USE的能力，效率更高，也能保证嵌入的一致性。

如果你确实需要把USE嵌入保存成word2vec的.bin或.vec格式，可以按以下步骤操作：

import tensorflow as tf
import tensorflow_hub as hub
from gensim.models import KeyedVectors

# 加载USE模型
embed = hub.load("https://tfhub.dev/google/universal-sentence-encoder/4")

# 替换成你自己的词汇表（可以从语料库提取或手动定义）
custom_vocab = ["cat", "dog", "sat", "lay", "mat", "rug", "the", "on"]

# 生成词汇对应的嵌入
embeddings = embed(custom_vocab).numpy()

# 构建KeyedVectors对象
kv = KeyedVectors(vector_size=512)
kv.add_vectors(custom_vocab, embeddings)

# 保存成文本格式的.vec文件
kv.save_word2vec_format("use_embeddings.vec", binary=False)
# 保存成二进制格式的.bin文件
kv.save_word2vec_format("use_embeddings.bin", binary=True)

# 后续可以直接用Gensim加载
# loaded_kv = KeyedVectors.load_word2vec_format("use_embeddings.bin", binary=True)

如果你想在spaCy生态中使用USE嵌入，有两种方式：

方式A：转成spaCy兼容的向量格式 #

1. 先按方案2生成use_embeddings.vec文件
2. 用spaCy的命令行工具将其转换为spaCy的二进制向量格式：

python -m spacy init vectors en use_embeddings.vec ./use_spacy_vectors --name use_vectors

3. 在代码中加载使用：

import spacy

# 加载基础spaCy模型，替换其向量表
nlp = spacy.load("en_core_web_sm")
nlp.vocab.vectors.from_disk("./use_spacy_vectors/vectors")

# 测试：获取token的USE嵌入
doc = nlp("cat sat on the mat")
print(doc[0].vector.shape)  # 输出(512,)，和USE维度一致

方式B：自定义spaCy管道组件动态生成嵌入 #

如果不想预先生成所有词汇的嵌入，可以写一个自定义组件，在处理文本时实时调用USE获取词嵌入：

import spacy
import tensorflow as tf
import tensorflow_hub as hub

# 加载USE模型
embed = hub.load("https://tfhub.dev/google/universal-sentence-encoder/4")

def use_embedding_component(doc):
    # 提取所有token的文本
    token_texts = [token.text for token in doc]
    # 生成嵌入
    embeddings = embed(token_texts).numpy()
    # 为每个token设置USE嵌入
    for token, emb in zip(doc, embeddings):
        token.set_vector(emb)
    return doc

# 加载spaCy基础模型并添加自定义组件
nlp = spacy.load("en_core_web_sm")
nlp.add_pipe("use_embedding", after="tagger", func=use_embedding_component)

# 测试
doc = nlp("cat sat on the mat")
print(doc[0].vector)  # 输出该词的USE嵌入向量

内容的提问来源于stack exchange，提问作者davislf2