How to use Spacy with Pyspark?

Viewed 923

I'm trying to use Spacy with Pyspark. Since Spacy models are not serializable, they cannot be broadcasted. The solution I have found is to load the models on worker-side. But loading models for each task can be time consuming. The solution I have found is to use a global variable in which I store the loaded model. And since "spark.python.worker.reuse" is True by default, the model will be loaded only once in each worker.

accu = spark.sparkContext.accumulator(0)
spmodel = None
def get_sp():
    global spmodel
    if not spmodel:
        accu.add(1)
        spmodel = spacy.load("ja_core_news_sm")
    return spmodel

def compute_ner(s_text:pd.Series) -> pd.DataFrame:
    nlp = get_sp()
    docs = nlp.pipe(s_text)
    ents=[
         [(ent.text, ent.label_) for ent in doc.ents]
         for doc in docs
         ]
    preds = []
    # Some Instructions .. 
    return pd.DataFrame(preds)
udf_calculate = F.pandas_udf(compute_ner, returnType=return_types)
ner_df = df.repartition(16).withColumn('ner', udf_calculate('text'))

I have used an accumulator to see how many times the model is loaded. I have tested the application on local mode with 8 dedicated cores. The DataFrame I have used is divided on 16 partitions. I was surprised to discover that the accumulator value is 16. So spmodel is always None. I'm a little bit confused, can someone explain what's happening, or maybe I am missing something.

0 Answers
Related