将关键字提取模型应用于包含文本的Spark DataFrame列

Question

想象一下，我有一个Spark DataFrame行星，其中包含2列-id和描述：

我想应用不同的关键字提取模型，并为每个< code>id从< code>description列的相应文本中提取关键字。

在几乎所有实现示例中，模型都接受文本作为字符串。
但是，例如，我如何实现keyBERT模型（查看此处的示例）处理描述列而不是简单的文本字符串？

在输出中，我需要使用模型找到的前n个关键字的新列来扩展我的原始数据框：

我将不胜感激任何帮助！

Answer 1

您可以使用RDD.map将函数应用于rdd中的每个元素。

我太懒了，无法真正安装keyBERT，所以我将制作自己的extract_keywords函数来展示一个示例。

rdd = sc.parallelize(enumerate(("Mercury is the closest planet to the Sun. The smallest planet in the Solar System, Mercury has no natural satellites. The dominant geological features are impact craters or basins with ejecta blankets, the remains of early volcanic activity including magma flows, and lobed ridges or rupes that were probably produced by a period of contraction early in the planet's history.", "Venus is close in size to Earth and has a thick silicate mantle around an iron core and evidence of internal geological activity. Venus is the hottest planet mainly due to the amount of greenhouse gases in the atmosphere. Venus has no magnetic field that would prevent depletion of its substantial atmosphere, which suggests that its atmosphere is being replenished by volcanic eruptions."), start=1))

from heapq import nlargest

def extract_keywords(s):
    return nlargest(4, s.split(), key=len)

rdd = rdd.map(lambda x: (x[0], x[1], extract_keywords(x[1])))

print( rdd.collect() )
# [(1, "Mercury is ...", ['satellites.', 'contraction', 'geological', 'blankets,']),
#  (2, 'Venus is ...', ['atmosphere.', 'substantial', 'atmosphere,', 'replenished'])]

我不熟悉keyBERT，但从他们展示的示例来看，我认为以下两个片段之一可能是合适的：

# EITHER USE SAME MODEL FOR EVERYONE
kw_model = KeyBERT()
rdd = rdd.map(lambda x: (x[0], x[1], kw_model.extract_keywords(x[1])))

# OR USE NEW MODEL FOR EACH ELEMENT
rdd = rdd.map(lambda x: (x[0], x[1], KeyBERT().extract_keywords(x[1])))