Как создать новый столбец с плотными векторами в таблице Pyspark с помощью UDF Pandas? - PullRequest
0 голосов
/ 20 сентября 2019

Моя таблица хранится в pyspark в блоках данных.Таблица имеет два столбца id и text.Я пытаюсь получить плотный вектор для столбца text.У меня есть модель ML для генерации плотного представления текста в новый столбец с именем dense_embedding.Модель генерирует пустой массив для представления входного текста.работай так model.encode(text_input).Я хочу использовать эту модель для генерации плотного представления всего текста для столбца text.

Вот что я сделал:

from pyspark.sql.functions import pandas_udf, PandasUDFType
from pyspark.ml.linalg import Vectors
from pyspark.sql.types import *
import pandas as pd

# Use pandas_udf to define a Pandas UDF
@pandas_udf('???', PandasUDFType.SCALAR)
# Input/output are text and dense vector
def embedding(v):
    return Vectors.dense(model.encode([v]))

small.withColumn('dense_embedding', embedding(small.text))

Я не уверен, какой тип данных я должен поместить в функцию pandas_udf?это правильно, чтобы преобразовать dens_vector как то, что я сделал?

...