class SnowFlakeEmbeddings extends AnnotatorModel[SnowFlakeEmbeddings] with HasBatchedAnnotate[SnowFlakeEmbeddings] with WriteTensorflowModel with WriteOnnxModel with WriteOpenvinoModel with HasEmbeddingsProperties with HasStorageRef with HasCaseSensitiveProperties with HasEngine
Sentence embeddings using SnowFlake.
snowflake-arctic-embed is a suite of text embedding models that focuses on creating high-quality retrieval models optimized for performance.
Pretrained models can be loaded with pretrained
of the companion object:
val embeddings = SnowFlakeEmbeddings.pretrained() .setInputCols("document") .setOutputCol("snowflake_embeddings")
The default model is "snowflake_artic_m"
, if no name is provided.
For available pretrained models please see the Models Hub.
For extended examples of usage, see SnowFlakeEmbeddingsTestSpec.
Sources :
Arctic-Embed: Scalable, Efficient, and Accurate Text Embedding Models
Paper abstract
The models are trained by leveraging existing open-source text representation models, such as bert-base-uncased, and are trained in a multi-stage pipeline to optimize their retrieval performance. First, the models are trained with large batches of query-document pairs where negatives are derived in-batch—pretraining leverages about 400m samples of a mix of public datasets and proprietary web search data. Following pretraining models are further optimized with long training on a smaller dataset (about 1m samples) of triplets of query, positive document, and negative document derived from hard harmful mining. Mining of the negatives and data curation is crucial to retrieval accuracy. A detailed technical report will be available shortly.
Example
import spark.implicits._ import com.johnsnowlabs.nlp.base.DocumentAssembler import com.johnsnowlabs.nlp.annotators.Tokenizer import com.johnsnowlabs.nlp.embeddings.SnowFlakeEmbeddings import com.johnsnowlabs.nlp.EmbeddingsFinisher import org.apache.spark.ml.Pipeline val documentAssembler = new DocumentAssembler() .setInputCol("text") .setOutputCol("document") val embeddings = SnowFlakeEmbeddings.pretrained() .setInputCols("document") .setOutputCol("snowflake_embeddings") val embeddingsFinisher = new EmbeddingsFinisher() .setInputCols("snowflake_embeddings") .setOutputCols("finished_embeddings") .setOutputAsVector(true) val pipeline = new Pipeline().setStages(Array( documentAssembler, embeddings, embeddingsFinisher )) val data = Seq("hello world", "hello moon").toDF("text") val result = pipeline.fit(data).transform(data) result.selectExpr("explode(finished_embeddings) as result").show(5, 80) --------------------+ finished_embeddings| --------------------+ [[-0.45763275, 0....| [[-0.43076283, 0....| --------------------+
- See also
Annotators Main Page for a list of transformer based embeddings
- Grouped
- Alphabetic
- By Inheritance
- SnowFlakeEmbeddings
- HasEngine
- HasCaseSensitiveProperties
- HasStorageRef
- HasEmbeddingsProperties
- HasProtectedParams
- WriteOpenvinoModel
- WriteOnnxModel
- WriteTensorflowModel
- HasBatchedAnnotate
- AnnotatorModel
- CanBeLazy
- RawAnnotator
- HasOutputAnnotationCol
- HasInputAnnotationCols
- HasOutputAnnotatorType
- ParamsAndFeaturesWritable
- HasFeatures
- DefaultParamsWritable
- MLWritable
- Model
- Transformer
- PipelineStage
- Logging
- Params
- Serializable
- Serializable
- Identifiable
- AnyRef
- Any
- Hide All
- Show All
- Public
- All
Parameters
A list of (hyper-)parameter keys this annotator can take. Users can set and get the parameter values through setters and getters, respectively.
-
val
batchSize: IntParam
Size of every batch (Default depends on model).
Size of every batch (Default depends on model).
- Definition Classes
- HasBatchedAnnotate
-
val
caseSensitive: BooleanParam
Whether to ignore case in index lookups (Default depends on model)
Whether to ignore case in index lookups (Default depends on model)
- Definition Classes
- HasCaseSensitiveProperties
-
val
configProtoBytes: IntArrayParam
ConfigProto from tensorflow, serialized into byte array.
ConfigProto from tensorflow, serialized into byte array. Get with
config_proto.SerializeToString()
-
val
dimension: ProtectedParam[Int]
Number of embedding dimensions (Default depends on model)
Number of embedding dimensions (Default depends on model)
- Definition Classes
- HasEmbeddingsProperties
-
val
engine: Param[String]
This param is set internally once via loadSavedModel.
This param is set internally once via loadSavedModel. That's why there is no setter
- Definition Classes
- HasEngine
-
val
maxSentenceLength: IntParam
Max sentence length to process (Default:
128
) -
val
poolingStrategy: Param[String]
Pooling strategy to use for sentence embeddings.
Pooling strategy to use for sentence embeddings.
Available pooling strategies for sentence embeddings are:
"cls"
: leading[CLS]
token"cls_avg"
: leading[CLS]
token + mean of all other tokens"last"
: embeddings of the last token in the sequence"avg"
: mean of all tokens"max"
: max of all embedding values for the token sequence"all"
: return all token embeddings"int"
: An integer number, which represents the index of the token to use as the embedding
-
val
signatures: MapFeature[String, String]
It contains TF model signatures for the laded saved model
-
val
storageRef: Param[String]
Unique identifier for storage (Default:
this.uid
)Unique identifier for storage (Default:
this.uid
)- Definition Classes
- HasStorageRef
-
val
vocabulary: MapFeature[String, Int]
Vocabulary used to encode the words to ids with WordPieceEncoder
Members
-
implicit
class
ProtectedParam[T] extends Param[T]
- Definition Classes
- HasProtectedParams
-
type
AnnotatorType = String
- Definition Classes
- HasOutputAnnotatorType
-
def
batchAnnotate(batchedAnnotations: Seq[Array[Annotation]]): Seq[Seq[Annotation]]
takes a document and annotations and produces new annotations of this annotator's annotation type
takes a document and annotations and produces new annotations of this annotator's annotation type
- batchedAnnotations
Annotations that correspond to inputAnnotationCols generated by previous annotators if any
- returns
any number of annotations processed for every input annotation. Not necessary one to one relationship
- Definition Classes
- SnowFlakeEmbeddings → HasBatchedAnnotate
-
def
batchProcess(rows: Iterator[_]): Iterator[Row]
- Definition Classes
- HasBatchedAnnotate
-
final
def
clear(param: Param[_]): SnowFlakeEmbeddings.this.type
- Definition Classes
- Params
-
def
copy(extra: ParamMap): SnowFlakeEmbeddings
requirement for annotators copies
requirement for annotators copies
- Definition Classes
- RawAnnotator → Model → Transformer → PipelineStage → Params
-
def
createDatabaseConnection(database: Name): RocksDBConnection
- Definition Classes
- HasStorageRef
-
def
explainParam(param: Param[_]): String
- Definition Classes
- Params
-
def
explainParams(): String
- Definition Classes
- Params
-
final
def
extractParamMap(): ParamMap
- Definition Classes
- Params
-
final
def
extractParamMap(extra: ParamMap): ParamMap
- Definition Classes
- Params
-
val
features: ArrayBuffer[Feature[_, _, _]]
- Definition Classes
- HasFeatures
-
final
def
get[T](param: Param[T]): Option[T]
- Definition Classes
- Params
-
final
def
getDefault[T](param: Param[T]): Option[T]
- Definition Classes
- Params
-
def
getInputCols: Array[String]
- returns
input annotations columns currently used
- Definition Classes
- HasInputAnnotationCols
-
def
getLazyAnnotator: Boolean
- Definition Classes
- CanBeLazy
-
final
def
getOrDefault[T](param: Param[T]): T
- Definition Classes
- Params
-
final
def
getOutputCol: String
Gets annotation column name going to generate
Gets annotation column name going to generate
- Definition Classes
- HasOutputAnnotationCol
-
def
getParam(paramName: String): Param[Any]
- Definition Classes
- Params
- def getPoolingStrategy: String
-
def
getStorageRef: String
- Definition Classes
- HasStorageRef
-
final
def
hasDefault[T](param: Param[T]): Boolean
- Definition Classes
- Params
-
def
hasParam(paramName: String): Boolean
- Definition Classes
- Params
-
def
hasParent: Boolean
- Definition Classes
- Model
-
val
inputAnnotatorTypes: Array[String]
Annotator reference id.
Annotator reference id. Used to identify elements in metadata or to refer to this annotator type
- Definition Classes
- SnowFlakeEmbeddings → HasInputAnnotationCols
-
final
def
isDefined(param: Param[_]): Boolean
- Definition Classes
- Params
-
final
def
isSet(param: Param[_]): Boolean
- Definition Classes
- Params
-
val
lazyAnnotator: BooleanParam
- Definition Classes
- CanBeLazy
-
def
onWrite(path: String, spark: SparkSession): Unit
- Definition Classes
- SnowFlakeEmbeddings → ParamsAndFeaturesWritable
-
val
optionalInputAnnotatorTypes: Array[String]
- Definition Classes
- HasInputAnnotationCols
-
val
outputAnnotatorType: AnnotatorType
- Definition Classes
- SnowFlakeEmbeddings → HasOutputAnnotatorType
-
lazy val
params: Array[Param[_]]
- Definition Classes
- Params
-
var
parent: Estimator[SnowFlakeEmbeddings]
- Definition Classes
- Model
-
def
save(path: String): Unit
- Definition Classes
- MLWritable
- Annotations
- @Since( "1.6.0" ) @throws( ... )
- def sentenceStartTokenId: Int
-
def
set[T](param: ProtectedParam[T], value: T): SnowFlakeEmbeddings.this.type
Sets the value for a protected Param.
Sets the value for a protected Param.
If the parameter was already set, it will not be set again. Default values do not count as a set value and can be overridden.
- T
Type of the parameter
- param
Protected parameter to set
- value
Value for the parameter
- returns
This object
- Definition Classes
- HasProtectedParams
-
final
def
set[T](param: Param[T], value: T): SnowFlakeEmbeddings.this.type
- Definition Classes
- Params
-
final
def
setInputCols(value: String*): SnowFlakeEmbeddings.this.type
- Definition Classes
- HasInputAnnotationCols
-
def
setInputCols(value: Array[String]): SnowFlakeEmbeddings.this.type
Overrides required annotators column if different than default
Overrides required annotators column if different than default
- Definition Classes
- HasInputAnnotationCols
-
def
setLazyAnnotator(value: Boolean): SnowFlakeEmbeddings.this.type
- Definition Classes
- CanBeLazy
-
final
def
setOutputCol(value: String): SnowFlakeEmbeddings.this.type
Overrides annotation column name when transforming
Overrides annotation column name when transforming
- Definition Classes
- HasOutputAnnotationCol
-
def
setParent(parent: Estimator[SnowFlakeEmbeddings]): SnowFlakeEmbeddings
- Definition Classes
- Model
-
def
setStorageRef(value: String): SnowFlakeEmbeddings.this.type
- Definition Classes
- HasStorageRef
-
def
toString(): String
- Definition Classes
- Identifiable → AnyRef → Any
- def tokenize(sentences: Seq[Annotation]): Seq[WordpieceTokenizedSentence]
-
final
def
transform(dataset: Dataset[_]): DataFrame
Given requirements are met, this applies ML transformation within a Pipeline or stand-alone Output annotation will be generated as a new column, previous annotations are still available separately metadata is built at schema level to record annotations structural information outside its content
Given requirements are met, this applies ML transformation within a Pipeline or stand-alone Output annotation will be generated as a new column, previous annotations are still available separately metadata is built at schema level to record annotations structural information outside its content
- dataset
Dataset[Row]
- Definition Classes
- AnnotatorModel → Transformer
-
def
transform(dataset: Dataset[_], paramMap: ParamMap): DataFrame
- Definition Classes
- Transformer
- Annotations
- @Since( "2.0.0" )
-
def
transform(dataset: Dataset[_], firstParamPair: ParamPair[_], otherParamPairs: ParamPair[_]*): DataFrame
- Definition Classes
- Transformer
- Annotations
- @Since( "2.0.0" ) @varargs()
-
final
def
transformSchema(schema: StructType): StructType
requirement for pipeline transformation validation.
requirement for pipeline transformation validation. It is called on fit()
- Definition Classes
- RawAnnotator → PipelineStage
-
val
uid: String
- Definition Classes
- SnowFlakeEmbeddings → Identifiable
-
def
validateStorageRef(dataset: Dataset[_], inputCols: Array[String], annotatorType: String): Unit
- Definition Classes
- HasStorageRef
-
def
write: MLWriter
- Definition Classes
- ParamsAndFeaturesWritable → DefaultParamsWritable → MLWritable
-
def
writeOnnxModel(path: String, spark: SparkSession, onnxWrapper: OnnxWrapper, suffix: String, fileName: String): Unit
- Definition Classes
- WriteOnnxModel
-
def
writeOnnxModels(path: String, spark: SparkSession, onnxWrappersWithNames: Seq[(OnnxWrapper, String)], suffix: String): Unit
- Definition Classes
- WriteOnnxModel
-
def
writeOpenvinoModel(path: String, spark: SparkSession, openvinoWrapper: OpenvinoWrapper, suffix: String, fileName: String): Unit
- Definition Classes
- WriteOpenvinoModel
-
def
writeOpenvinoModels(path: String, spark: SparkSession, ovWrappersWithNames: Seq[(OpenvinoWrapper, String)], suffix: String): Unit
- Definition Classes
- WriteOpenvinoModel
-
def
writeTensorflowHub(path: String, tfPath: String, spark: SparkSession, suffix: String = "_use"): Unit
- Definition Classes
- WriteTensorflowModel
-
def
writeTensorflowModel(path: String, spark: SparkSession, tensorflow: TensorflowWrapper, suffix: String, filename: String, configProtoBytes: Option[Array[Byte]] = None): Unit
- Definition Classes
- WriteTensorflowModel
-
def
writeTensorflowModelV2(path: String, spark: SparkSession, tensorflow: TensorflowWrapper, suffix: String, filename: String, configProtoBytes: Option[Array[Byte]] = None, savedSignatures: Option[Map[String, String]] = None): Unit
- Definition Classes
- WriteTensorflowModel
Parameter setters
- def sentenceEndTokenId: Int
-
def
setBatchSize(size: Int): SnowFlakeEmbeddings.this.type
Size of every batch.
Size of every batch.
- Definition Classes
- HasBatchedAnnotate
-
def
setCaseSensitive(value: Boolean): SnowFlakeEmbeddings.this.type
Whether to lowercase tokens or not
Whether to lowercase tokens or not
- Definition Classes
- SnowFlakeEmbeddings → HasCaseSensitiveProperties
- def setConfigProtoBytes(bytes: Array[Int]): SnowFlakeEmbeddings.this.type
-
def
setDimension(value: Int): SnowFlakeEmbeddings.this.type
Set Embeddings dimensions for the BERT model Only possible to set this when the first time is saved dimension is not changeable, it comes from BERT config file
Set Embeddings dimensions for the BERT model Only possible to set this when the first time is saved dimension is not changeable, it comes from BERT config file
- Definition Classes
- SnowFlakeEmbeddings → HasEmbeddingsProperties
- def setMaxSentenceLength(value: Int): SnowFlakeEmbeddings.this.type
- def setModelIfNotSet(spark: SparkSession, tensorflowWrapper: Option[TensorflowWrapper], onnxWrapper: Option[OnnxWrapper], openvinoWrapper: Option[OpenvinoWrapper]): SnowFlakeEmbeddings
-
def
setPoolingStrategy(value: String): SnowFlakeEmbeddings.this.type
Pooling strategy to use for sentence embeddings.
Pooling strategy to use for sentence embeddings.
Available pooling strategies for sentence embeddings are:
"cls"
: leading[CLS]
token"cls_avg"
: leading[CLS]
token + mean of all other tokens"last"
: embeddings of the last token in the sequence"avg"
: mean of all tokens"max"
: max of all embedding features of the entire token sequence"int"
: An integer number, which represents the index of the token to use as the embedding
- def setSignatures(value: Map[String, String]): SnowFlakeEmbeddings.this.type
- def setVocabulary(value: Map[String, Int]): SnowFlakeEmbeddings.this.type
Parameter getters
-
def
getBatchSize: Int
Size of every batch.
Size of every batch.
- Definition Classes
- HasBatchedAnnotate
-
def
getCaseSensitive: Boolean
- Definition Classes
- HasCaseSensitiveProperties
- def getConfigProtoBytes: Option[Array[Byte]]
-
def
getDimension: Int
- Definition Classes
- HasEmbeddingsProperties
-
def
getEngine: String
- Definition Classes
- HasEngine
- def getMaxSentenceLength: Int
- def getModelIfNotSet: SnowFlake
- def getSignatures: Option[Map[String, String]]