sparknlp.annotator.DistilBertEmbeddings¶

class sparknlp.annotator.DistilBertEmbeddings(classname='com.johnsnowlabs.nlp.embeddings.DistilBertEmbeddings', java_model=None)[source]¶

Bases: sparknlp.common.AnnotatorModel, sparknlp.common.HasEmbeddingsProperties, sparknlp.common.HasCaseSensitiveProperties, sparknlp.common.HasStorageRef, sparknlp.common.HasBatchedAnnotate

DistilBERT is a small, fast, cheap and light Transformer model trained by distilling BERT base. It has 40% less parameters than bert-base-uncased, runs 60% faster while preserving over 95% of BERT’s performances as measured on the GLUE language understanding benchmark.

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

>>> embeddings = DistilBertEmbeddings.pretrained() \
...     .setInputCols(["document", "token"]) \
...     .setOutputCol("embeddings")

The default model is "distilbert_base_cased", if no name is provided. For available pretrained models please see the Models Hub.

For extended examples of usage, see the Spark NLP Workshop. Models from the HuggingFace 🤗 Transformers library are also compatible with Spark NLP 🚀. To see which models are compatible and how to import them see Import Transformers into Spark NLP 🚀.

Input Annotation types	Output Annotation type
`DOCUMENT, TOKEN`	`WORD_EMBEDDINGS`

Parameters

batchSize: Size of every batch, by default 8
dimension: Number of embedding dimensions, by default 768
caseSensitive: Whether to ignore case in tokens for embeddings matching, by default False
maxSentenceLength: Max sentence length to process, by default 128
configProtoBytes: ConfigProto from tensorflow, serialized into byte array.

Notes

DistilBERT doesn’t have token_type_ids, you don’t need to indicate which token belongs to which segment. Just separate your segments with the separation token tokenizer.sep_token (or [SEP]).
DistilBERT doesn’t have options to select the input positions (position_ids input). This could be added if necessary though, just let us know if you need this option.

References

The DistilBERT model was proposed in the paper DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter.

Paper Abstract:

As Transfer Learning from large-scale pre-trained models becomes more prevalent in Natural Language Processing (NLP), operating these large models in on-the- edge and/or under constrained computational training or inference budgets remains challenging. In this work, we propose a method to pre-train a smaller general-purpose language representation model, called DistilBERT, which can then be fine-tuned with good performances on a wide range of tasks like its larger counterparts. While most prior work investigated the use of distillation for building task-specific models, we leverage knowledge distillation during the pretraining phase and show that it is possible to reduce the size of a BERT model by 40%, while retaining 97% of its language understanding capabilities and being 60% faster. To leverage the inductive biases learned by larger models during pretraining, we introduce a triple loss combining language modeling, distillation and cosine-distance losses. Our smaller, faster and lighter model is cheaper to pre-train and we demonstrate its capabilities for on-device computations in a proof-of-concept experiment and a comparative on-device study.

Examples

>>> import sparknlp
>>> from sparknlp.base import *
>>> from sparknlp.annotator import *
>>> from pyspark.ml import Pipeline
>>> documentAssembler = DocumentAssembler() \
...     .setInputCol("text") \
...     .setOutputCol("document")
>>> tokenizer = Tokenizer() \
...     .setInputCols(["document"]) \
...     .setOutputCol("token")
>>> embeddings = DistilBertEmbeddings.pretrained() \
...     .setInputCols(["document", "token"]) \
...     .setOutputCol("embeddings") \
...     .setCaseSensitive(True)
>>> embeddingsFinisher = EmbeddingsFinisher() \
...     .setInputCols(["embeddings"]) \
...     .setOutputCols("finished_embeddings") \
...     .setOutputAsVector(True) \
...     .setCleanAnnotations(False)
>>> pipeline = Pipeline() \
...     .setStages([
...       documentAssembler,
...       tokenizer,
...       embeddings,
...       embeddingsFinisher
...     ])
>>> data = spark.createDataFrame([["This is a sentence."]]).toDF("text")
>>> result = pipeline.fit(data).transform(data)
>>> result.selectExpr("explode(finished_embeddings) as result").show(5, 80)
+--------------------------------------------------------------------------------+
|                                                                          result|
+--------------------------------------------------------------------------------+
|[0.1127224713563919,-0.1982710212469101,0.5360898375511169,-0.272536993026733...|
|[0.35534414649009705,0.13215228915214539,0.40981462597846985,0.14036104083061...|
|[0.328085333108902,-0.06269335001707077,-0.017595693469047546,-0.024373905733...|
|[0.15617232024669647,0.2967822253704071,0.22324979305267334,-0.04568954557180...|
|[0.45411425828933716,0.01173491682857275,0.190129816532135,0.1178255230188369...|
+--------------------------------------------------------------------------------+

Methods

`__init__`([classname, java_model])	Initialize this instance with a Java model object.
`clear`(param)	Clears a param from the param map if it has been explicitly set.
`copy`([extra])	Creates a copy of this instance with the same uid and some extra params.
`explainParam`(param)	Explains a single param and returns its name, doc, and optional default value and user-supplied value in a string.
`explainParams`()	Returns the documentation of all params with their optionally default values and user-supplied values.
`extractParamMap`([extra])	Extracts the embedded default param values and user-supplied values, and then merges them with extra values from input into a flat param map, where the latter value is used if there exist conflicts, i.e., with ordering: default param values < user-supplied values < extra.
`getBatchSize`()	Gets current batch size.
`getCaseSensitive`()	Gets whether to ignore case in tokens for embeddings matching.
`getDimension`()	Gets embeddings dimension.
`getInputCols`()	Gets current column names of input annotations.
`getLazyAnnotator`()	Gets whether Annotator should be evaluated lazily in a RecursivePipeline.
`getOrDefault`(param)	Gets the value of a param in the user-supplied param map or its default value.
`getOutputCol`()	Gets output column name of annotations.
`getParam`(paramName)	Gets a param by its name.
`getParamValue`(paramName)	Gets the value of a parameter.
`getStorageRef`()	Gets unique reference name for identification.
`hasDefault`(param)	Checks whether a param has a default value.
`hasParam`(paramName)	Tests whether this instance contains a param with a given (string) name.
`isDefined`(param)	Checks whether a param is explicitly set by user or has a default value.
`isSet`(param)	Checks whether a param is explicitly set by user.
`load`(path)	Reads an ML instance from the input path, a shortcut of read().load(path).
`loadSavedModel`(folder, spark_session)	Loads a locally saved model.
`pretrained`([name, lang, remote_loc])	Downloads and loads a pretrained model.
`read`()	Returns an MLReader instance for this class.
`save`(path)	Save this ML instance to the given path, a shortcut of 'write().save(path)'.
`set`(param, value)	Sets a parameter in the embedded param map.
`setBatchSize`(v)	Sets batch size.
`setCaseSensitive`(value)	Sets whether to ignore case in tokens for embeddings matching.
`setConfigProtoBytes`(b)	Sets configProto from tensorflow, serialized into byte array.
`setDimension`(value)	Sets embeddings dimension.
`setInputCols`(*value)	Sets column names of input annotations.
`setLazyAnnotator`(value)	Sets whether Annotator should be evaluated lazily in a RecursivePipeline.
`setMaxSentenceLength`(value)	Sets max sentence length to process.
`setOutputCol`(value)	Sets output column name of annotations.
`setParamValue`(paramName)	Sets the value of a parameter.
`setParams`()
`setStorageRef`(value)	Sets unique reference name for identification.
`transform`(dataset[, params])	Transforms the input dataset with optional parameters.
`write`()	Returns an MLWriter instance for this ML instance.

Attributes

`batchSize`
`caseSensitive`
`configProtoBytes`
`dimension`
`getter_attrs`
`inputCols`
`lazyAnnotator`
`maxSentenceLength`
`name`
`outputCol`
`params`	Returns all params ordered by name.
`storageRef`

clear(param)¶: Clears a param from the param map if it has been explicitly set.

copy(extra=None)¶

Creates a copy of this instance with the same uid and some extra params. This implementation first calls Params.copy and then make a copy of the companion Java pipeline component with extra params. So both the Python wrapper and the Java pipeline component get copied.

Parameters: extra – Extra parameters to copy to the new instance
Returns: Copy of this instance

explainParam(param)¶: Explains a single param and returns its name, doc, and optional default value and user-supplied value in a string.

explainParams()¶: Returns the documentation of all params with their optionally default values and user-supplied values.

extractParamMap(extra=None)¶

Extracts the embedded default param values and user-supplied values, and then merges them with extra values from input into a flat param map, where the latter value is used if there exist conflicts, i.e., with ordering: default param values < user-supplied values < extra.

Parameters: extra – extra param values
Returns: merged param map

getBatchSize()¶

Gets current batch size.

Returns

int: Current batch size

getCaseSensitive()¶

Gets whether to ignore case in tokens for embeddings matching.

Returns

bool: Whether to ignore case in tokens for embeddings matching

getDimension()¶: Gets embeddings dimension.

getInputCols()¶: Gets current column names of input annotations.

getLazyAnnotator()¶: Gets whether Annotator should be evaluated lazily in a RecursivePipeline.

getOrDefault(param)¶: Gets the value of a param in the user-supplied param map or its default value. Raises an error if neither is set.

getOutputCol()¶: Gets output column name of annotations.

getParam(paramName)¶: Gets a param by its name.

getParamValue(paramName)¶

Gets the value of a parameter.

Parameters

paramNamestr: Name of the parameter

getStorageRef()¶

Gets unique reference name for identification.

Returns

str: Unique reference name for identification

hasDefault(param)¶: Checks whether a param has a default value.

hasParam(paramName)¶: Tests whether this instance contains a param with a given (string) name.

isDefined(param)¶: Checks whether a param is explicitly set by user or has a default value.

isSet(param)¶: Checks whether a param is explicitly set by user.

classmethod load(path)¶: Reads an ML instance from the input path, a shortcut of read().load(path).

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

DistilBertEmbeddings: The restored model

property params¶: Returns all params ordered by name. The default implementation uses dir() to get all attributes of type Param.

static pretrained(name='distilbert_base_cased', lang='en', remote_loc=None)[source]¶

Downloads and loads a pretrained model.

Parameters

namestr, optional: Name of the pretrained model, by default “distilbert_base_cased”
langstr, optional: Language of the pretrained model, by default “en”
remote_locstr, optional: Optional remote address of the resource, by default None. Will use Spark NLPs repositories otherwise.

Returns

DistilBertEmbeddings: The restored model

classmethod read()¶: Returns an MLReader instance for this class.

save(path)¶: Save this ML instance to the given path, a shortcut of ‘write().save(path)’.

set(param, value)¶: Sets a parameter in the embedded param map.

setBatchSize(v)¶

Sets batch size.

Parameters

vint: Batch size

setCaseSensitive(value)¶

Sets whether to ignore case in tokens for embeddings matching.

Parameters

valuebool: Whether to ignore case in tokens for embeddings matching

setConfigProtoBytes(b)[source]¶

Sets configProto from tensorflow, serialized into byte array.

Parameters

bList[str]: ConfigProto from tensorflow, serialized into byte array

setDimension(value)¶

Sets embeddings dimension.

Parameters

valueint: Embeddings dimension

setInputCols(*value)¶

Sets column names of input annotations.

Parameters

*valuestr: Input columns for the annotator

setLazyAnnotator(value)¶

Sets whether Annotator should be evaluated lazily in a RecursivePipeline.

Parameters

valuebool: Whether Annotator should be evaluated lazily in a RecursivePipeline

setMaxSentenceLength(value)[source]¶

Sets max sentence length to process.

Parameters

valueint: Max sentence length to process

setOutputCol(value)¶

Sets output column name of annotations.

Parameters

valuestr: Name of output column

setParamValue(paramName)¶

Sets the value of a parameter.

Parameters

paramNamestr: Name of the parameter

setStorageRef(value)¶

Sets unique reference name for identification.

Parameters

valuestr: Unique reference name for identification

transform(dataset, params=None)¶

Transforms the input dataset with optional parameters.

Parameters

dataset – input dataset, which is an instance of pyspark.sql.DataFrame
params – an optional param map that overrides embedded params.

Returns

transformed dataset

New in version 1.3.0.

uid¶: A unique id for the object.

write()¶: Returns an MLWriter instance for this ML instance.

sparknlp.annotator.DependencyParserModel

sparknlp.annotator.DistilBertForTokenClassification