Apply keyword extraction model to a Spark DataFrame column that contains text

user_10797482 0
Viewed 191

Imagine I have a Spark DataFrame planet that contains 2 columns - id and description:

id description
1 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.
2 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.

I want to apply different keyword extraction models and for each id extract keywords from corresponding text in the description column.

In almost all implementation examples, the models accept text as a string.
But how can I implement, for example, keyBERT model (look at the example here) processing description column instead of a simple text string?

In the output, I need to get my original dataframe extended with a new column with top n keywords found by the model:

id description keywords
1 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. Mercury, Solar System, early volcanic activity, geological features, magma flows
2 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. Venus, internal geological activity, thick silicate mantle, atmosphere, iron core

I would appreciate any help!

pyspark apache-spark-sql nltk nlp keyword
1 Answers

You can use RDD.map to apply a function to every element in your rdd.

I'm too lazy to actually install keyBERT, so I'll make up my own extract_keywords function to show an example.

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'])]

I'm not familiar with keyBERT, but from the example they show, I think one of the following two snippets could be appropriate:

# 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])))
user_30807230
Related