如何将n-gram组合成Spark中的一个词汇?

时间:2016-08-08 23:29:19

标签: python apache-spark nlp pyspark apache-spark-ml

想知道是否有内置Spark功能将1,2,n-gram功能组合到一个词汇表中。在n=2中设置NGram,然后调用CountVectorizer会导致只包含2克的字典。我真正想要的是将所有频繁的1克,2克等组合成一个字典用于我的语料库。

1 个答案:

答案 0 :(得分:12)

您可以训练单独的NGramCountVectorizer模型,并使用VectorAssembler进行合并。

from pyspark.ml.feature import NGram, CountVectorizer, VectorAssembler
from pyspark.ml import Pipeline


def build_ngrams(inputCol="tokens", n=3):

    ngrams = [
        NGram(n=i, inputCol="tokens", outputCol="{0}_grams".format(i))
        for i in range(1, n + 1)
    ]

    vectorizers = [
        CountVectorizer(inputCol="{0}_grams".format(i),
            outputCol="{0}_counts".format(i))
        for i in range(1, n + 1)
    ]

    assembler = [VectorAssembler(
        inputCols=["{0}_counts".format(i) for i in range(1, n + 1)],
        outputCol="features"
    )]

    return Pipeline(stages=ngrams + vectorizers + assembler)

使用示例:

df = spark.createDataFrame([
  (1, ["a", "b", "c", "d"]),
  (2, ["d", "e", "d"])
], ("id", "tokens"))

build_ngrams().fit(df).transform(df)