Моя таблица хранится в 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 как то, что я сделал?