bge_small_en_v1.5 model from BAAI

Description

Pretrained BGEEmbeddings model, adapted from Hugging Face and curated to provide scalability and production-readiness using Spark NLP.bge_small_en_v1.5 is a English model originally trained by BAAI.

Download Copy S3 URI

How to use

 
documentAssembler = DocumentAssembler() \
      .setInputCol("text") \
      .setOutputCol("document")
    
embeds = BGEEmbeddings.pretrained("bge_small_en_v1.5","en") \
      .setInputCols(["document"]) \
      .setOutputCol("embeddings")       
        
pipeline = Pipeline().setStages([documentAssembler, embeds])
data = spark.createDataFrame([["I love spark-nlp"]]).toDF("text")
pipelineModel = pipeline.fit(data)
pipelineDF = pipelineModel.transform(data)


val documentAssembler = new DocumentAssembler() 
    .setInputCol("text") 
    .setOutputCol("document")
    
val embeds = BGEEmbeddings.pretrained("bge_small_en_v1.5","en") 
    .setInputCols(Array("document")) 
    .setOutputCol("embeddings")

val pipeline = new Pipeline().setStages(Array(documentAssembler, embeds))
val data = Seq("I love spark-nlp").toDF("text")
val pipelineModel = pipeline.fit(data)
val pipelineDF = pipelineModel.transform(data)

Model Information

Model Name: bge_small_en_v1.5
Compatibility: Spark NLP 5.5.1+
License: Open Source
Edition: Official
Input Labels: [document]
Output Labels: [bge]
Language: en
Size: 78.9 MB

References

https://huggingface.co/BAAI/bge-small-en-v1.5