sparknlp_jsl.annotator.embeddings.entity_chunk_embeddings#

Module Contents#

Classes#

EntityChunkEmbeddings

Weighted average embeddings of multiple named entities chunk annotations.

class EntityChunkEmbeddings(classname='com.johnsnowlabs.nlp.annotators.embeddings.EntityChunkEmbeddings', java_model=None)#

Bases: sparknlp.annotator.BertSentenceEmbeddings, sparknlp_jsl.common.HasEngine

Weighted average embeddings of multiple named entities chunk annotations.

Entity Chunk Embeddings uses BERT Sentence embeddings to compute a weighted average vector represention of related entity chunks. The input the model consists of chunks of recognized named entities. One or more entities are selected as target entities and for each of them a list of related entities is specified (if empty, all other entities are assumed to be related).

The model looks for chunks of the target entities and then tries to pair each target entity (e.g. DRUG) with other related entities (e.g. DOSAGE, STRENGTH, FORM, etc). The criterion for pairing a target entity with another related entity is that they appear in the same sentence and the maximal syntactic distance is below a predefined threshold.

The relationship between target and related entities is one-to-many, meaning that if there multiple instances of the same target entity (e.g.) within a sentence, the model will map a related entity (e.g. DOSAGE) to at most one of the instances of the target entity. For example, if there is a sentence “The patient was given 125 mg of paracetamol and metformin”, the model will pair “125 mg” to “paracetamol”, but not to “metformin”.

The output of the model is an average embeddings of the chunks of each of the target entities and their related entities. It is possible to specify a particular weight for each entity type.

An entity can be defined both as target a entity and as a related entity for some other target entity. For example, we may want to compute the embeddings of SYMPTOMs and their related entities, as well as the embeddings of DRUGs and their related entities, one of each is also SYMPTOM. In such cases, it is possible to use the TARGET_ENTITY:RELATED_ENTITY notation to specify the weight of an related entity (e.g. “DRUG:SYMPTOM” to set the weight of SYMPTOM when it appears as an related entity to target entity DRUG). The relative weights of entities for particular entity chunk embeddings are available in the annotations metadata.

This model is a subclass of BertSentenceEmbeddings and shares all parameters with it. It can load any pretrained BertSentenceEmbeddings model.

The default model is “sbiobert_base_cased_mli” from “clinical/models”. Other available models can be found at Models Hub.

Input Annotation types

Output Annotation type

DEPENDENCY, CHUNK

SENTENCE_EMBEDDINGS

Parameters:
  • targetEntities (dict) –

    Target entities and their related entities. A target entity mapped to an empty list of related entities means all other entities are assumed to be related to it. The default value is an empty dictionary, meaning no target entities are specified.

    Entity names are case insensitive.

  • entityWeights (dict) –

    Relative weights of entities. By default the dictionary is empty and all entities have equal weights. If it is non-empty and some entity is not in it, then its weight is set to 0.

    The notation TARGET_ENTITY:RELATED_ENTITY can be used to specify the weight of a entity which is related to specific target entity (e.g. “DRUG:SYMPTOM” -> 0.3f).

    Entity names are case insensitive.

  • maxSyntacticDistance (int) – Maximal syntactic distance between related entities. Default value is 2.

Examples

>>> import sparknlp
>>> from sparknlp.base import *
>>> from sparknlp_jsl.common import *
>>> from sparknlp.annotator import *
>>> from sparknlp.training import *
>>> import sparknlp_jsl
>>> from sparknlp_jsl.base import *
>>> from sparknlp_jsl.annotator import *
>>> from pyspark.ml import Pipeline
>>> documenter = DocumentAssembler()\
...     .setInputCol("text")\
...     .setOutputCol("documents")
>>> sentence_detector = SentenceDetector() \
...     .setInputCols("documents") \
...     .setOutputCol("sentences")
>>> tokenizer = Tokenizer() \
...     .setInputCols("sentences") \
...     .setOutputCol("tokens")
>>> embeddings = WordEmbeddingsModel() \
...     .pretrained("embeddings_clinical", "en", "clinical/models")\
...     .setInputCols(["sentences", "tokens"])\
...     .setOutputCol("embeddings")
>>> ner_model = MedicalNerModel()\
...     .pretrained("ner_posology_large", "en", "clinical/models")\
...     .setInputCols(["sentences", "tokens", "embeddings"])\
...     .setOutputCol("ner")
>>> ner_converter = NerConverterInternal()\
...     .setInputCols("sentences", "tokens", "ner")\
...     .setOutputCol("ner_chunks")
>>> pos_tager = PerceptronModel()\
...     .pretrained("pos_clinical", "en", "clinical/models")\
...     .setInputCols("sentences", "tokens")\
...     .setOutputCol("pos_tags")
>>> dependency_parser = DependencyParserModel()\
...     .pretrained("dependency_conllu", "en")\
...     .setInputCols(["sentences", "pos_tags", "tokens"])\
...     .setOutputCol("dependencies")
>>> drug_chunk_embeddings = EntityChunkEmbeddings()\
...     .pretrained("sbiobert_base_cased_mli","en","clinical/models")\
...     .setInputCols(["ner_chunks", "dependencies"])\
...     .setOutputCol("drug_chunk_embeddings")\
...     .setMaxSyntacticDistance(3)\
...     .setTargetEntities({"DRUG": []})
...     .setEntityWeights({"DRUG": 0.8, "STRENGTH": 0.2, "DOSAGE": 0.2, "FORM": 0.5})
>>> sampleData = "The parient was given metformin 125 mg, 250 mg of coumadin and then one pill paracetamol"
>>> data = SparkContextForTest.spark.createDataFrame([[sampleData]]).toDF("text")
>>> pipeline = Pipeline().setStages([
...     documenter,
...     sentence_detector,
...     tokenizer,
...     embeddings,
...     ner_model,
...     ner_converter,
...     pos_tager,
...     dependency_parser,
...     drug_chunk_embeddings])
>>> results = pipeline.fit(data).transform(data)
>>> results = results \
...     .selectExpr("explode(drug_chunk_embeddings) AS drug_chunk") \
...     .selectExpr("drug_chunk.result", "slice(drug_chunk.embeddings, 1, 5) AS drug_embedding") \
...     .cache()
>>> results.show(truncate=False)
+-----------------------------+-----------------------------------------------------------------+
|                       result|                                                  drug_embedding"|
+-----------------------------+-----------------------------------------------------------------+
|metformin 125 mg             |[-0.267413, 0.07614058, -0.5620966, 0.83838946, 0.8911504]       |
|250 mg coumadin              |[0.22319649, -0.07094894, -0.6885556, 0.79176235, 0.82672405]    |
|one pill paracetamol         |[-0.10939768, -0.29242, -0.3574444, 0.3981813, 0.79609615]       |
+-----------------------------+-----------------------------------------------------------------+
batchSize#
caseSensitive#
configProtoBytes#
dimension#
entityWeights#
getter_attrs = []#
inputAnnotatorTypes#
inputCols#
isLong#
lazyAnnotator#
maxSyntacticDistance#
name = 'EntityChunkEmbeddings'#
optionalInputAnnotatorTypes = []#
outputAnnotatorType#
outputCol#
storageRef#
targetEntities#
uid#
clear(param: pyspark.ml.param.Param) None#

Clears a param from the param map if it has been explicitly set.

copy(extra: pyspark.ml._typing.ParamMap | None = None) JP#

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 (dict, optional) – Extra parameters to copy to the new instance

Returns:

Copy of this instance

Return type:

JavaParams

explainParam(param: str | Param) str#

Explains a single param and returns its name, doc, and optional default value and user-supplied value in a string.

explainParams() str#

Returns the documentation of all params with their optionally default values and user-supplied values.

extractParamMap(extra: pyspark.ml._typing.ParamMap | None = None) pyspark.ml._typing.ParamMap#

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 (dict, optional) – extra param values

Returns:

merged param map

Return type:

dict

getBatchSize()#

Gets current batch size.

Returns:

Current batch size

Return type:

int

getCaseSensitive()#

Gets whether to ignore case in tokens for embeddings matching.

Returns:

Whether to ignore case in tokens for embeddings matching

Return type:

bool

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: str) Any#
getOrDefault(param: Param[T]) T

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: str) Param#

Gets a param by its name.

getParamValue(paramName)#

Gets the value of a parameter.

Parameters:

paramName (str) – Name of the parameter

getStorageRef()#

Gets unique reference name for identification.

Returns:

Unique reference name for identification

Return type:

str

hasDefault(param: str | Param[Any]) bool#

Checks whether a param has a default value.

hasParam(paramName: str) bool#

Tests whether this instance contains a param with a given (string) name.

inputColsValidation(value)#
isDefined(param: str | Param[Any]) bool#

Checks whether a param is explicitly set by user or has a default value.

isSet(param: str | Param[Any]) bool#

Checks whether a param is explicitly set by user.

classmethod load(path: str) RL#

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

static loadSavedModel(folder, spark_session, use_openvino=False)#

Loads a locally saved model.

Parameters:
  • folder (str) – Folder of the saved model

  • spark_session (pyspark.sql.SparkSession) – The current SparkSession

  • use_openvino (bool) – Use OpenVINO backend

Returns:

The restored model

Return type:

BertSentenceEmbeddings

static pretrained(name='sbiobert_base_cased_mli', lang='en', remote_loc='clinical/models')#

Downloads and loads a pretrained model.

Parameters:
  • name (str, optional) – Name of the pretrained model, by default “sbiobert_base_cased_mli”

  • lang (str, optional) – Language of the pretrained model, by default “en”

  • remote_loc (str, optional) – Optional remote address of the resource, by default None. Will use Spark NLPs repositories otherwise.

Returns:

The restored model

Return type:

EntityChunkEmbeddings

classmethod read()#

Returns an MLReader instance for this class.

save(path: str) None#

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

set(param: Param, value: Any) None#

Sets a parameter in the embedded param map.

setBatchSize(v)#

Sets batch size.

Parameters:

v (int) – Batch size

setCaseSensitive(value)#

Sets whether to ignore case in tokens for embeddings matching.

Parameters:

value (bool) – Whether to ignore case in tokens for embeddings matching

setConfigProtoBytes(b)#

Sets configProto from tensorflow, serialized into byte array.

Parameters:

b (List[int]) – ConfigProto from tensorflow, serialized into byte array

setDimension(value)#

Sets embeddings dimension.

Parameters:

value (int) – Embeddings dimension

setEntityWeights(weights: dict = {})#

Sets the relative weights of the embeddings of specific entities.

By default the dictionary is empty and all entities have equal weights.

If non-empty and some entity is not in it, then its weight is set to 0.

Parameters:

weights (: dict[str, float]) – Dictionary with the relative weighs of entities. The notation TARGET_ENTITY:RELATED_ENTITY can be used to specify the weight of a entity which is related to specific target entity (e.g. “DRUG:SYMPTOM”: 0.3). Entity names are case-insensitive.

setInputCols(*value)#

Sets column names of input annotations.

Parameters:

*value (List[str]) – Input columns for the annotator

setIsLong(value)#

Sets whether to use Long type instead of Int type for inputs buffer.

Some Bert models require Long instead of Int.

Parameters:

value (bool) – Whether to use Long type instead of Int type for inputs buffer

setLazyAnnotator(value)#

Sets whether Annotator should be evaluated lazily in a RecursivePipeline.

Parameters:

value (bool) – Whether Annotator should be evaluated lazily in a RecursivePipeline

setMaxSyntacticDistance(distance: int = 2)#

Sets the maximal syntactic distance between related entities. Default value is 2.

Parameters:

distance (int) – Maximal syntactic distance

setOutputCol(value)#

Sets output column name of annotations.

Parameters:

value (str) – Name of output column

setParamValue(paramName)#

Sets the value of a parameter.

Parameters:

paramName (str) – Name of the parameter

setParams()#
setStorageRef(value)#

Sets unique reference name for identification.

Parameters:

value (str) – Unique reference name for identification

setTargetEntities(entities: dict = {})#

Sets the target entities and maps them to their related entities.

A target entity with an empty list of related entities means all other entities are assumed to be related to it.

Parameters:

entities (dict[str, list[str]]) – Dictionary with target and related entities (TARGET: [RELATED1, RELATED2,…]). If the list of related entities is empty, then all non-target entities are considered. Entity names are case insensitive.

transform(dataset: pyspark.sql.dataframe.DataFrame, params: pyspark.ml._typing.ParamMap | None = None) pyspark.sql.dataframe.DataFrame#

Transforms the input dataset with optional parameters.

New in version 1.3.0.

Parameters:
  • dataset (pyspark.sql.DataFrame) – input dataset

  • params (dict, optional) – an optional param map that overrides embedded params.

Returns:

transformed dataset

Return type:

pyspark.sql.DataFrame

write() JavaMLWriter#

Returns an MLWriter instance for this ML instance.