PySpark + sentence-transformers - UDF cannot pickle '_thread.RLock' object

Agudolive 0
Viewed 42

I would like to use PySpark and the library sentence-transformers to speed up the embedding of a corpus of documents.

From what I understand I can't just call :

embedd_text = udf(lambda x: model.encode(x).tolist(), ArrayType(ArrayType(FloatType())))

Because the model can't be pickled and be passed to the workers.

That's why I put the code that way:

class text_embedder:
    @property
    def model(self):
        return SentenceTransformer('dangvantuan/sentence-camembert-large')
        
    def embedd_text(self, x):
        return self.model.encode(x).tolist()
    
    def return_embedd_text_column(self, df):
        embedd_text_udf = udf(self.embedd_text, returnType=ArrayType(ArrayType(FloatType())))
        return df.withColumn("text_embedded", embedd_text_udf(data.text))

That way, the model should be created within the UDF call and distributed to the workers.

The problem is now I'm facing TypeError: cannot pickle '_thread.RLock' object and I don't know where the problem comes from.

My question is: is it possible to do what I try to do, if yes, is it a good idea to do it like that performance-wise.

Thanks in advance for your help

sentence-transformers pyspark apache-spark python
1 Answers

In order to use the model in the udf each executor node needs to have the model available.
you can use Broadcast in order to send the model to each executor.
Then in the udf, instead of calling model you will need to call the broadcast_variable.value which is the model.

from pyspark.sql.functions import *
from pyspark.sql.types import *


df = spark.createDataFrame([
    (["This framework generates embeddings for each input sentence"],),
    (["Sentences are passed as a list of string"],),
    (["The quick brown fox jumps over the lazy dog"],),
], schema=["sentence"])

broadcast_model = spark.sparkContext.broadcast(model)

embedd_text = udf(lambda x: broadcast_model.value.encode(x).tolist(), ArrayType(ArrayType(FloatType())))

df.withColumn("embedded", embedd_text("sentence")).show()

+--------------------+--------------------+
|            sentence|            embedded|
+--------------------+--------------------+
|[This framework g...|[[-0.013717361, -...|
|[Sentences are pa...|[[0.06387022, 0.0...|
|[The quick brown ...|[[0.035496805, 0....|
+--------------------+--------------------+
walking0
Related