如何从Spark NLP的WordEmbeddingsModel中获取词汇表以构建嵌入矩阵

问题背景 #

使用Spark NLP构建包含WordEmbeddingsModel的Pipeline后，Python API层面没有直接提供getVocab方法，需要借助底层Java实现来获取词汇表。

解决方案 #

训练完成的PipelineModel通过stages属性保存了所有流水线阶段，先从中定位到WordEmbeddingsModel实例，再通过其Java对象调用词汇表相关方法：

# 从训练好的PipelineModel中提取WordEmbeddingsModel阶段
embeddings_stage = [stage for stage in model.stages if isinstance(stage, WordEmbeddingsModel)][0]

# 通过Java对象获取词汇表，自动转为Python列表
vocab_list = embeddings_stage._java_obj.getVocab()

# 若需要获取对应嵌入向量矩阵，可调用getEmbeddings方法
embeddings_java_matrix = embeddings_stage._java_obj.getEmbeddings()

# 转换为NumPy数组用于后续处理
import numpy as np
embeddings_np_matrix = np.array(embeddings_java_matrix)

补充说明 #

• Spark NLP的Python注解器是Java实现的上层封装，不少底层功能需要通过_java_obj访问原生Java方法
• getVocab()返回的Java集合会自动转换为Python列表，可直接用于后续词汇映射逻辑
• getEmbeddings()返回的矩阵可转换为NumPy数组，快速构建符合需求的嵌入矩阵

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