Options
- Mark as New
- Bookmark
- Subscribe
- Mute
- Subscribe to RSS Feed
- Permalink
- Report Inappropriate Content
06-19-2024 12:11 AM
Hmm interesting, then it's something else.
The below code works for me on a Shared access mode cluster. (I don't know what your input dataset looks like):
df = spark.sql("SELECT '1' as content")
from fastembed import TextEmbedding, SparseTextEmbedding
from pyspark.sql.pandas.functions import pandas_udf, PandasUDFType
from pyspark.sql.types import StructType, StructField, StringType, ArrayType, FloatType, IntegerType
import pandas as pd
from pyspark.sql.functions import col
@pandas_udf(ArrayType(FloatType()))
def generate_dense_embeddings(contents: pd.Series) -> pd.Series:
small_embedding_model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5", cache_dir="/tmp/local_cache/")
dense_embeddings_list = small_embedding_model.embed(contents)
return pd.Series(list(dense_embeddings_list))
df=df.limit(50)
df.show(10)
embeddings = df.withColumn("dense_vector", generate_dense_embeddings(col("content")))
embeddings.show(10)Are you sure your cluster setup is sufficient enough for what you're trying to achieve?