sparknlp.annotator.embeddings.bge_m3_embeddings#

Classes#

class BGEM3Embeddings(classname='com.johnsnowlabs.nlp.embeddings.BGEM3Embeddings', java_model=None)[source]#

Sentence embeddings using BGE-M3.

BGE-M3 is a versatile multilingual embedding model from BAAI built on the xlm-roberta-large backbone. Unlike the English dense-only BGE models exposed through BGEEmbeddings, BGE-M3 supports up to 8192 tokens, over 100 languages, and produces both:

• a dense embedding (in Annotation.embeddings), and

• a sparse / lexical {token: weight} map (in Annotation.metadata under the sparse_weights key when setReturnSparseEmbeddings() is enabled).

Both outputs are emitted from a single SENTENCE_EMBEDDINGS output column.

Pretrained models can be loaded with pretrained() of the companion object:

>>> embeddings = BGEM3Embeddings.pretrained() \
...     .setInputCols(["document"]) \
...     .setOutputCol("bge_m3_embeddings")

The default model is "bge_m3", if no name is provided.

For available pretrained models please see the Models Hub.

Input Annotation types	Output Annotation type
DOCUMENT	SENTENCE_EMBEDDINGS

References

BGE M3-Embedding: Multi-Lingual, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation

BGE Github Repository

Parameters:

batchSize

Size of every batch, by default 8

dimension

Number of embedding dimensions, by default 1024

caseSensitive

Whether to ignore case in tokens for embeddings matching, by default True

maxSentenceLength

Max sentence length to process, by default 512 (up to 8192)

returnSparseEmbeddings

Whether to compute the sparse lexical embeddings

Examples

>>> import sparknlp
>>> from sparknlp.base import *
>>> from sparknlp.annotator import *
>>> from pyspark.ml import Pipeline
>>> documentAssembler = DocumentAssembler() \
...     .setInputCol("text") \
...     .setOutputCol("document")
>>> embeddings = BGEM3Embeddings.pretrained() \
...     .setInputCols(["document"]) \
...     .setOutputCol("bge_m3")
>>> embeddingsFinisher = EmbeddingsFinisher() \
...     .setInputCols(["bge_m3"]) \
...     .setOutputCols("finished_embeddings") \
...     .setOutputAsVector(True)
>>> pipeline = Pipeline().setStages([
...     documentAssembler,
...     embeddings,
...     embeddingsFinisher
... ])
>>> data = spark.createDataFrame([["El BGE-M3 admite recuperación densa y dispersa."]]).toDF("text")
>>> result = pipeline.fit(data).transform(data)
>>> result.selectExpr("explode(finished_embeddings) as result").show(1, 80)

name = 'BGEM3Embeddings'[source]#

inputAnnotatorTypes[source]#

outputAnnotatorType = 'sentence_embeddings'[source]#

max_length_limit = 8192[source]#

returnSparseEmbeddings[source]#

setReturnSparseEmbeddings(value)[source]#

Sets whether to compute the sparse lexical embeddings

Parameters:

valuebool

Whether to return the sparse lexical embeddings

getReturnSparseEmbeddings()[source]#

Gets whether the sparse lexical embeddings are computed.

Returns:

bool

Whether the sparse lexical embeddings are returned

static loadSavedModel(folder, spark_session)[source]#

Loads a locally saved model.

Parameters:

folderstr

Folder of the saved model

spark_sessionpyspark.sql.SparkSession

The current SparkSession

Returns:

BGEM3Embeddings

The restored model

static pretrained(name='bge_m3', lang='xx', remote_loc=None)[source]#

Downloads and loads a pretrained model.

Parameters:

namestr, optional

Name of the pretrained model, by default “bge_m3”

langstr, optional

Language of the pretrained model, by default “xx”

remote_locstr, optional

Optional remote address of the resource, by default None. Will use Spark NLPs repositories otherwise.

Returns:

BGEM3Embeddings

The restored model