在所有 Spark 工作节点上初始化 gensim 对象

Question

I have a function and a UDF which I have created,

    def test(string):
        path_index=SparkFiles.get("corpus_final_production.index")
        path_dictionary=SparkFiles.get('dictionary_production.gensim')
        path_corpus=SparkFiles.get("corpus_final_production")
        dictionary = corpora.Dictionary.load(path_dictionary)
        corpus = corpora.MmCorpus(path_corpus)
        tf_idf = gensim.models.TfidfModel(corpus)
        index_tmpfile = get_tmpfile(path_index)
        sims = gensim.similarities.Similarity(index_tmpfile,tf_idf[corpus],num_features=len(dictionary))
        query_doc=word_tokenize(string.lower())
        query_doc_bow=dictionary.doc2bow(query_doc)
        query_doc_tf_idf=tf_idf[query_doc_bow]
        sum_of_sims=np.sum(sims[query_doc_tf_idf], dtype=np.float32)
        max_sims=np.amax(sims[query_doc_tf_idf])
        max_count=np.count_nonzero(sims[query_doc_tf_idf] >= max_sims-0.05)
        max_sims_origin=file_docs[np.argmax(sims[query_doc_tf_idf])]
        return max_sims_origin

   test_udf = udf(lambda x: test(x),StringType())
   df_new = garuda.withColumn('max_sim_origin', test_udf(garuda.text))

It is working fine but as you see I am applying a rowwise action to the pyspark dataframe. For every row, the dictionary corpus and sims get generated with the indexes which takes close to 6 minutes for every row.

Is there a way for me to initialize the dictionary, corpus and index files before onto every worker node instead of calling it in the UDF.

I am new to spark so every little help will be appreciated

I have added all the dictionary and corpus files, as it is pregenerated using sc.addFile()

Answer 1

You may try the following:

1. Creating your gensim components once
2. Broadcasting reusable gensim components to every worker
3. Modifying your UDF to utilize broadcasted components
4. Apply your modified UDF test

# Step 1
path_index=SparkFiles.get("corpus_final_production.index")
path_dictionary=SparkFiles.get('dictionary_production.gensim')
path_corpus=SparkFiles.get("corpus_final_production")
dictionary = corpora.Dictionary.load(path_dictionary)
corpus = corpora.MmCorpus(path_corpus)
tf_idf = gensim.models.TfidfModel(corpus)
index_tmpfile = get_tmpfile(path_index)
sims = gensim.similarities.Similarity(index_tmpfile,tf_idf[corpus],num_features=len(dictionary))

# Step 2
dictionaryBC = sparkSession.sparkContext.broadcast(dictionary)
tfidfBC = sparkSession.sparkContext.broadcast(tfidf)
simsBC = sparkSession.sparkContext.broadcast(sims)

# Step 3
def test(string):
    dictionary = dictionaryBC.value
    tfidf = tfidfBC.value
    sims = simsBC.value

    query_doc=word_tokenize(string.lower())
    query_doc_bow=dictionary.doc2bow(query_doc)
    query_doc_tf_idf=tf_idf[query_doc_bow]
    sum_of_sims=np.sum(sims[query_doc_tf_idf], dtype=np.float32)
    max_sims=np.amax(sims[query_doc_tf_idf])
    max_count=np.count_nonzero(sims[query_doc_tf_idf] >= max_sims-0.05)
    max_sims_origin=file_docs[np.argmax(sims[query_doc_tf_idf])]
    
    return max_sims_origin

test_udf = udf(lambda x: test(x),StringType())

# Step 4
df_new = garuda.withColumn('max_sim_origin', test_udf(garuda.text))

Let me know if this works for you