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.