Создание модели преобразователя предложений в Spark Mllib

Question

Я использовал предварительно обученную модель из библиотеки преобразователей предложений, чтобы проверить сходство между двумя предложениями. Теперь мне нужно, чтобы эта конкретная модель была реализована с помощью spark mllib. Какие-либо предложения? Я очень ценю любую помощь, которую вы можете предоставить.

https://www.sbert.net/ https://spark.apache.org/mllib/

Answer 1

Один подход, который я нашел работающим, заключается в использовании пользовательской функции Pandas , которая кодирует текст и возвращает вложение. Затем этот столбец внедрения можно использовать с MLlib.

      import pandas as pd
import pyspark.sql.functions as F
from pyspark.sql.types import ArrayType, DoubleType, StringType
from sentence_transformers import SentenceTransformer

# import sbert model
model = SentenceTransformer("all-MiniLM-L6-v2")

# sentences to encode
sentences = [
    "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.",
]

# create spark df with sentences
data = spark.createDataFrame(sentences, StringType(), ["sentences"])
data.show()

# create a pandas udf that will encode the text and return an array of doubles
@F.pandas_udf(returnType=ArrayType(DoubleType()))
def encode(x: pd.Series) -> pd.Series:
    return pd.Series(model.encode(x).tolist())

# apply udf and show 
data.withColumn("embedding", encode("value")).show()

выход

      +--------------------+
|               value|
+--------------------+
|This framework ge...|
|Sentences are pas...|
|The quick brown f...|
+--------------------+

+--------------------+--------------------+
|               value|           embedding|
+--------------------+--------------------+
|This framework ge...|[-0.0137173617258...|
|Sentences are pas...|[0.05645250156521...|
|The quick brown f...|[0.04393352568149...|
+--------------------+--------------------+