Spark环境下生成类Gensim风格Bigrams的技术方案问询

要在Spark中生成和Gensim一致的全局高频Bigrams，核心是先基于全量语料筛选出符合频率/PMI阈值的双词组合，再将这些组合合并回原文本，而不是每条文本单独生成所有二元组。以下分两种方案实现：

方案一：Spark NLP实现（优先推荐） #

步骤1：全局统计并筛选高频Bigrams #

首先生成所有可能的双词组合，统计全局出现频率，再通过min_count和PMI（点互信息）筛选出符合条件的Bigrams（对齐Gensim的Phrases逻辑）：

from sparknlp.annotator import NGramGenerator
from sparknlp.base import DocumentAssembler, Finisher
from pyspark.sql.functions import explode, col, log, split
from pyspark.sql.types import StringType

# 假设你的预处理DataFrame名为`process`，包含`lemmatized`列（Array[String]类型的分词结果）
# 1. 将分词转换为Spark NLP的Document格式，生成所有可能的Bigrams
document_assembler = DocumentAssembler() \
    .setInputCol("lemmatized") \
    .setOutputCol("document")

ngram_generator = NGramGenerator() \
    .setInputCols(["document"]) \
    .setOutputCol("ngrams") \
    .setN(2) \
    .setDelimiter("_") \
    .setEnableCumulative(False)

finisher = Finisher() \
    .setInputCols(["ngrams"]) \
    .setOutputCols(["all_bigrams"]) \
    .setOutputAsArray(True)

# 生成包含所有Bigrams的DataFrame
bigram_df = document_assembler.transform(process)
bigram_df = ngram_generator.transform(bigram_df)
bigram_df = finisher.transform(bigram_df)

# 2. 统计每个Bigram的全局出现次数，过滤掉低频组合（对应Gensim的min_count=5）
bigram_counts = bigram_df.select(explode(col("all_bigrams")).alias("bigram")) \
    .groupBy("bigram") \
    .count() \
    .filter(col("count") >= 5)

# 3. 计算PMI（点互信息）筛选有意义的Bigrams（对应Gensim的threshold）
# 先统计单个词的全局频率
word_counts = process.select(explode(col("lemmatized")).alias("word")) \
    .groupBy("word") \
    .count() \
    .withColumnRenamed("count", "word_count")

# 拆分Bigram为两个独立词
bigram_split = bigram_counts.withColumn("words", split(col("bigram"), "_")) \
    .withColumn("word1", col("words")[0]) \
    .withColumn("word2", col("words")[1]) \
    .drop("words")

# 计算PMI：PMI = log( (Bigram出现次数 * 总词数) / (词1出现次数 * 词2出现次数) )
total_words = word_counts.agg(sum("word_count")).first()[0]

bigram_pmi = bigram_split.join(word_counts, bigram_split.word1 == word_counts.word, "left") \
    .withColumnRenamed("word_count", "count1") \
    .drop("word") \
    .join(word_counts, bigram_split.word2 == word_counts.word, "left") \
    .withColumnRenamed("word_count", "count2") \
    .drop("word") \
    .withColumn("pmi", log( (col("count") * total_words) / (col("count1") * col("count2")) )) \
    .filter(col("pmi") >= 4.0)  # 阈值可根据语料调整，Gensim默认threshold=100对应PMI约4-5

# 得到最终要合并的Bigrams列表
selected_bigrams = bigram_pmi.select("bigram").rdd.map(lambda x: x[0]).collect()

步骤2：将筛选出的Bigrams合并回原文本 #

使用广播变量传递Bigrams集合，通过UDF高效合并原分词中的连续双词：

from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType

# 将Bigrams集合转为广播变量，避免重复加载到每个Task
bigrams_broadcast = spark.sparkContext.broadcast(set(selected_bigrams))

def merge_bigrams(tokens):
    merged = []
    i = 0
    token_len = len(tokens)
    while i < token_len:
        # 检查当前词和下一个词是否组成目标Bigram
        if i < token_len - 1 and f"{tokens[i]}_{tokens[i+1]}" in bigrams_broadcast.value:
            merged.append(f"{tokens[i]}_{tokens[i+1]}")
            i += 2
        else:
            merged.append(tokens[i])
            i += 1
    return merged

# 注册UDF并生成处理后的列
merge_udf = udf(merge_bigrams, ArrayType(StringType()))
tweets_with_bigrams = process.withColumn("lemmatized_with_bigrams", merge_udf(col("lemmatized")))

方案二：Spark MLlib实现 #

如果无法使用Spark NLP，可通过MLlib的NGram组件实现相同逻辑：

步骤1：统计筛选Bigrams #

from pyspark.ml.feature import NGram
from pyspark.sql.functions import explode, col, split, log

# 生成所有可能的Bigrams
ngram = NGram(n=2, inputCol="lemmatized", outputCol="all_bigrams")
bigram_df = ngram.transform(process)

# 统计频率并过滤低频组合
bigram_counts = bigram_df.select(explode(col("all_bigrams")).alias("bigram")) \
    .groupBy("bigram") \
    .count() \
    .filter(col("count") >= 5)

# 后续PMI计算和筛选逻辑与Spark NLP方案完全一致，最终得到selected_bigrams列表

步骤2：合并Bigrams到原文本 #

使用和Spark NLP方案相同的广播变量+UDF逻辑即可，代码完全复用。

为什么你之前的方法无效？ #

• Spark NLP的NGramGenerator：仅生成单条文本内的所有连续双词，没有基于全量语料筛选高频组合，因此得到的是所有可能的二元组，而非有意义的Bigrams。
• UDF调用Gensim Phrases：按行单独训练Phrases模型，仅利用当前行的分词数据，完全没有用到全局语料，生成的Bigrams不具备通用性，不符合需求。

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