使用前面估算器的属性作为ML管道中的参数

Question

我正在使用Pyspark的Pipelines库中的ML来预处理文本并计算所有令牌的 TF-IDF 值。我还创建了一个自定义Transformer，它为每个文本片段返回具有最高TF-IDF值的5个令牌。主要代码如下：

%pyspark
tokenizer = RegexTokenizer(inputCol="text", outputCol="tokenized", pattern="\\W")
remover = StopWordsRemover(inputCol="tokenized", outputCol="filtered")
count_vectorizer = CountVectorizer(inputCol="filtered", outputCol="count", vocabSize=pow(2,10))
idf = IDF(inputCol="count", outputCol="TF-IDF")
normalizer = Normalizer(inputCol="TF-IDF", outputCol="normalized", p=2.0)
top_token_extractor = TopTokenExtractor(inputCol="normalized", outputCol="topTokens", vocabulary=model.stages[2].vocabulary) # !!! does not work

pipeline = Pipeline(stages=[tokenizer, remover, count_vectorizer, idf, normalizer, top_token_extractor])

model = pipeline.fit(df)

以下是TopTokenExtractor的实施：

%pyspark
from pyspark import keyword_only
from pyspark.ml.pipeline import Transformer 
from pyspark.ml.param.shared import HasInputCol, HasOutputCol, Param 
from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, StringType

class TopTokenExtractor(Transformer, HasInputCol, HasOutputCol): 

    @keyword_only
    def __init__(self, inputCol=None, outputCol=None, vocabulary=None): 
        super(TopTokenExtractor, self).__init__() 
        self.vocabulary = Param(self, "vocabulary", "")
        self._setDefault(vocabulary=set())
        kwargs = self._input_kwargs 
        self.setParams(**kwargs) 

    @keyword_only
    def setParams(self, inputCol=None, outputCol=None, vocabulary=None): 
        kwargs = self._input_kwargs 
        return self._set(**kwargs) 

    def setVocabulary(self, value):
        self._paramMap[self.vocabulary] = value
        return self

    def getVocabulary(self):
        return self.getOrDefault(self.vocabulary)

    def _transform(self, dataset): 
        out_col = self.getOutputCol() 
        in_col = dataset[self.getInputCol()]
        vocabulary = self.getVocabulary()

        def f(s): 
            token_tuples = sorted(list(zip(s.indices, s.values)), key=lambda x: x[1], reverse=True)
            top_tokens = list()
            for i in range(0, min(5, len(token_tuples))):
                top_tokens.append(vocabulary[token_tuples[i][0]])
            return top_tokens

        t = ArrayType(StringType()) 
        return dataset.withColumn(out_col, udf(f, t)(in_col)) 

问题是，为了返回令牌列表而不是索引，我需要将CountVectorizer中的词汇表作为参数传递给TopTokenExtractor。调用pipeline.fit(df)后，model.stages[2].vocabulary可以访问词汇表，但我无法弄清楚如何在管道过程中将其作为参数传递。这有可能吗？

作为一种解决方法，我可能会将管道分成两部分，但如果可能的话，我更愿意拥有一个管道。