4992350528

Build Type

github

Committed by GitHub

Commit Message

Merge 424c7ff18 into ef7906c5e

Pull Request Pull Request #13797: SPARKNLP-835: ProtectedParam and ProtectedFeature

Run Details

24 of 24 new or added lines in 6 files covered. (100.0%)

8643 of 13129 relevant lines covered (65.83%)

0.66 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

69.23

/src/main/scala/com/johnsnowlabs/nlp/annotators/ner/crf/NerCrfModel.scala

/*
 * Copyright 2017-2022 John Snow Labs
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.johnsnowlabs.nlp.annotators.ner.crf

import com.johnsnowlabs.ml.crf.{FbCalculator, LinearChainCrfModel}
import com.johnsnowlabs.nlp.AnnotatorType._
import com.johnsnowlabs.nlp._
import com.johnsnowlabs.nlp.annotators.common.Annotated.{NerTaggedSentence, PosTaggedSentence}
import com.johnsnowlabs.nlp.annotators.common._
import com.johnsnowlabs.nlp.serialization.{MapFeature, StructFeature}
import com.johnsnowlabs.storage.HasStorageRef
import org.apache.spark.ml.param.{BooleanParam, StringArrayParam}
import org.apache.spark.ml.util._
import org.apache.spark.sql.Dataset

import scala.collection.Map

/** Extracts Named Entities based on a CRF Model.
  *
  * This Named Entity recognition annotator allows for a generic model to be trained by utilizing
  * a CRF machine learning algorithm. The data should have columns of type `DOCUMENT, TOKEN, POS,
  * WORD_EMBEDDINGS`. These can be extracted with for example
  *   - a [[com.johnsnowlabs.nlp.annotators.sbd.pragmatic.SentenceDetector SentenceDetector]],
  *   - a [[com.johnsnowlabs.nlp.annotators.Tokenizer Tokenizer]] and
  *   - a [[com.johnsnowlabs.nlp.annotators.pos.perceptron.PerceptronModel PerceptronModel]].
  *
  * This is the instantiated model of the [[NerCrfApproach]]. For training your own model, please
  * see the documentation of that class.
  *
  * Pretrained models can be loaded with `pretrained` of the companion object:
  * {{{
  * val nerTagger = NerCrfModel.pretrained()
  *   .setInputCols("sentence", "token", "word_embeddings", "pos")
  *   .setOutputCol("ner"
  * }}}
  * The default model is `"ner_crf"`, if no name is provided. For available pretrained models
  * please see the [[https://sparknlp.org/models?task=Named+Entity+Recognition Models Hub]].
  *
  * For extended examples of usage, see the
  * [[https://github.com/JohnSnowLabs/spark-nlp/blob/master/examples/python/training/english/crf-ner/ner_dl_crf.ipynb Examples]].
  *
  * ==Example==
  * {{{
  * import spark.implicits._
  * import com.johnsnowlabs.nlp.base.DocumentAssembler
  * import com.johnsnowlabs.nlp.annotators.Tokenizer
  * import com.johnsnowlabs.nlp.annotators.sbd.pragmatic.SentenceDetector
  * import com.johnsnowlabs.nlp.embeddings.WordEmbeddingsModel
  * import com.johnsnowlabs.nlp.annotators.pos.perceptron.PerceptronModel
  * import com.johnsnowlabs.nlp.annotators.ner.crf.NerCrfModel
  * import org.apache.spark.ml.Pipeline
  *
  * // First extract the prerequisites for the NerCrfModel
  * val documentAssembler = new DocumentAssembler()
  *   .setInputCol("text")
  *   .setOutputCol("document")
  *
  * val sentence = new SentenceDetector()
  *   .setInputCols("document")
  *   .setOutputCol("sentence")
  *
  * val tokenizer = new Tokenizer()
  *   .setInputCols("sentence")
  *   .setOutputCol("token")
  *
  * val embeddings = WordEmbeddingsModel.pretrained()
  *   .setInputCols("sentence", "token")
  *   .setOutputCol("word_embeddings")
  *
  * val posTagger = PerceptronModel.pretrained()
  *   .setInputCols("sentence", "token")
  *   .setOutputCol("pos")
  *
  * // Then NER can be extracted
  * val nerTagger = NerCrfModel.pretrained()
  *   .setInputCols("sentence", "token", "word_embeddings", "pos")
  *   .setOutputCol("ner")
  *
  * val pipeline = new Pipeline().setStages(Array(
  *   documentAssembler,
  *   sentence,
  *   tokenizer,
  *   embeddings,
  *   posTagger,
  *   nerTagger
  * ))
  *
  * val data = Seq("U.N. official Ekeus heads for Baghdad.").toDF("text")
  * val result = pipeline.fit(data).transform(data)
  *
  * result.select("ner.result").show(false)
  * +------------------------------------+
  * |result                              |
  * +------------------------------------+
  * |[I-ORG, O, O, I-PER, O, O, I-LOC, O]|
  * +------------------------------------+
  * }}}
  *
  * @see
  *   [[com.johnsnowlabs.nlp.annotators.ner.dl.NerDLModel NerDLModel]] for a deep learning based
  *   approach
  * @see
  *   [[com.johnsnowlabs.nlp.annotators.ner.NerConverter NerConverter]] to further process the
  *   results
  * @param uid
  *   required uid for storing annotator to disk
  * @groupname anno Annotator types
  * @groupdesc anno
  *   Required input and expected output annotator types
  * @groupname Ungrouped Members
  * @groupname param Parameters
  * @groupname setParam Parameter setters
  * @groupname getParam Parameter getters
  * @groupname Ungrouped Members
  * @groupprio param  1
  * @groupprio anno  2
  * @groupprio Ungrouped 3
  * @groupprio setParam  4
  * @groupprio getParam  5
  * @groupdesc param
  *   A list of (hyper-)parameter keys this annotator can take. Users can set and get the
  *   parameter values through setters and getters, respectively.
  */
class NerCrfModel(override val uid: String)
    extends AnnotatorModel[NerCrfModel]
    with HasSimpleAnnotate[NerCrfModel]
    with HasStorageRef {

  def this() = this(Identifiable.randomUID("NER"))

  /** List of Entities to recognize
    *
    * @group param
    */
  val entities = new StringArrayParam(this, "entities", "List of Entities to recognize")

  /** The CRF model
    *
    * @group param
    */
  val model: StructFeature[LinearChainCrfModel] =
    new StructFeature[LinearChainCrfModel](this, "crfModel")

  /** Additional dictionary to use as for features (Default: `Map.empty[String, String]`)
    *
    * @group param
    */
  val dictionaryFeatures: MapFeature[String, String] =
    new MapFeature[String, String](this, "dictionaryFeatures")

  /** Whether or not to calculate prediction confidence by token, included in metadata (Default:
    * `false`)
    *
    * @group param
    */
  val includeConfidence = new BooleanParam(
    this,
    "includeConfidence",
    "whether or not to calculate prediction confidence by token, includes in metadata")

  /** @group setParam */
  def setModel(crf: LinearChainCrfModel): NerCrfModel = set(model, crf)

  /** @group setParam */
  def setDictionaryFeatures(dictFeatures: DictionaryFeatures): this.type =
    set(dictionaryFeatures, dictFeatures.dict)

  /** @group setParam */
  def setEntities(toExtract: Array[String]): NerCrfModel = set(entities, toExtract)

  /** @group setParam */
  def setIncludeConfidence(c: Boolean): this.type = set(includeConfidence, c)

  /** @group getParam */
  def getIncludeConfidence: Boolean = $(includeConfidence)

  setDefault(dictionaryFeatures, () => Map.empty[String, String])
  setDefault(includeConfidence, false)

  /** Predicts Named Entities in input sentences
    *
    * @param sentences
    *   POS tagged and WordpieceEmbeddings sentences
    * @return
    *   sentences with recognized Named Entities
    */
  def tag(sentences: Seq[(PosTaggedSentence, WordpieceEmbeddingsSentence)])
      : Seq[NerTaggedSentence] = {
    require(model.isSet, "model must be set before tagging")

    val crf = $$(model)

    val fg = FeatureGenerator(new DictionaryFeatures($$(dictionaryFeatures)))
    sentences.map { case (sentence, withEmbeddings) =>
      val instance = fg.generate(sentence, withEmbeddings, crf.metadata)

      lazy val confidenceValues = {
        val fb = new FbCalculator(instance.items.length, crf.metadata)
        fb.calculate(instance, $$(model).weights, 1)
        fb.alpha
      }

      val labelIds = crf.predict(instance)

      val words = sentence.indexedTaggedWords
        .zip(labelIds.labels)
        .zipWithIndex
        .flatMap { case ((word, labelId), idx) =>
          val label = crf.metadata.labels(labelId)

          val alpha = if ($(includeConfidence)) {
            val scores = Some(confidenceValues.apply(idx))
            Some(
              crf.metadata.labels.zipWithIndex
                .filter(x => x._2 != 0)
                .map { case (t, i) =>
                  Map(
                    t -> scores
                      .getOrElse(Array.empty[String])
                      .lift(i)
                      .getOrElse(0.0f)
                      .toString)
                })
          } else None

          if (!isDefined(entities) || $(entities).isEmpty || $(entities).contains(label))
            Some(IndexedTaggedWord(word.word, label, word.begin, word.end, alpha))
          else
            None
        }

      TaggedSentence(words)
    }
  }

  override protected def beforeAnnotate(dataset: Dataset[_]): Dataset[_] = {
    validateStorageRef(dataset, $(inputCols), AnnotatorType.WORD_EMBEDDINGS)
    dataset
  }

  override def annotate(annotations: Seq[Annotation]): Seq[Annotation] = {
    val sourceSentences = PosTagged.unpack(annotations)
    val withEmbeddings = WordpieceEmbeddingsSentence.unpack(annotations)
    val taggedSentences = tag(sourceSentences.zip(withEmbeddings))
    NerTagged.pack(taggedSentences)
  }

  def shrink(minW: Float): NerCrfModel = set(model, $$(model).shrink(minW))

  /** Input Annotator Types: DOCUMENT, TOKEN, POS, WORD_EMBEDDINGS
    * @group anno
    */
  override val inputAnnotatorTypes: Array[String] = Array(DOCUMENT, TOKEN, POS, WORD_EMBEDDINGS)

  /** Output Annotator Types: NAMED_ENTITY
    * @group anno
    */
  override val outputAnnotatorType: AnnotatorType = NAMED_ENTITY

}

trait ReadablePretrainedNerCrf
    extends ParamsAndFeaturesReadable[NerCrfModel]
    with HasPretrained[NerCrfModel] {
  override val defaultModelName: Option[String] = Some("ner_crf")

  /** Java compliant-overrides */
  override def pretrained(): NerCrfModel = super.pretrained()

  override def pretrained(name: String): NerCrfModel = super.pretrained(name)

  override def pretrained(name: String, lang: String): NerCrfModel = super.pretrained(name, lang)

  override def pretrained(name: String, lang: String, remoteLoc: String): NerCrfModel =
    super.pretrained(name, lang, remoteLoc)
}

/** This is the companion object of [[NerCrfModel]]. Please refer to that class for the
  * documentation.
  */
object NerCrfModel extends ReadablePretrainedNerCrf

1	/*
2	* Copyright 2017-2022 John Snow Labs
3	*
4	* Licensed under the Apache License, Version 2.0 (the "License");
5	* you may not use this file except in compliance with the License.
6	* You may obtain a copy of the License at
7	*
8	* http://www.apache.org/licenses/LICENSE-2.0
9	*
10	* Unless required by applicable law or agreed to in writing, software
11	* distributed under the License is distributed on an "AS IS" BASIS,
12	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13	* See the License for the specific language governing permissions and
14	* limitations under the License.
15	*/
16
17	package com.johnsnowlabs.nlp.annotators.ner.crf
18
19	import com.johnsnowlabs.ml.crf.{FbCalculator, LinearChainCrfModel}
20	import com.johnsnowlabs.nlp.AnnotatorType._
21	import com.johnsnowlabs.nlp._
22	import com.johnsnowlabs.nlp.annotators.common.Annotated.{NerTaggedSentence, PosTaggedSentence}
23	import com.johnsnowlabs.nlp.annotators.common._
24	import com.johnsnowlabs.nlp.serialization.{MapFeature, StructFeature}
25	import com.johnsnowlabs.storage.HasStorageRef
26	import org.apache.spark.ml.param.{BooleanParam, StringArrayParam}
27	import org.apache.spark.ml.util._
28	import org.apache.spark.sql.Dataset
29
30	import scala.collection.Map
31
32	/** Extracts Named Entities based on a CRF Model.
33	*
34	* This Named Entity recognition annotator allows for a generic model to be trained by utilizing
35	* a CRF machine learning algorithm. The data should have columns of type `DOCUMENT, TOKEN, POS,
36	* WORD_EMBEDDINGS`. These can be extracted with for example
37	* - a [[com.johnsnowlabs.nlp.annotators.sbd.pragmatic.SentenceDetector SentenceDetector]],
38	* - a [[com.johnsnowlabs.nlp.annotators.Tokenizer Tokenizer]] and
39	* - a [[com.johnsnowlabs.nlp.annotators.pos.perceptron.PerceptronModel PerceptronModel]].
40	*
41	* This is the instantiated model of the [[NerCrfApproach]]. For training your own model, please
42	* see the documentation of that class.
43	*
44	* Pretrained models can be loaded with `pretrained` of the companion object:
45	* {{{
46	* val nerTagger = NerCrfModel.pretrained()
47	* .setInputCols("sentence", "token", "word_embeddings", "pos")
48	* .setOutputCol("ner"
49	* }}}
50	* The default model is `"ner_crf"`, if no name is provided. For available pretrained models
51	* please see the [[https://sparknlp.org/models?task=Named+Entity+Recognition Models Hub]].
52	*
53	* For extended examples of usage, see the
54	* [[https://github.com/JohnSnowLabs/spark-nlp/blob/master/examples/python/training/english/crf-ner/ner_dl_crf.ipynb Examples]].
55	*
56	* ==Example==
57	* {{{
58	* import spark.implicits._
59	* import com.johnsnowlabs.nlp.base.DocumentAssembler
60	* import com.johnsnowlabs.nlp.annotators.Tokenizer
61	* import com.johnsnowlabs.nlp.annotators.sbd.pragmatic.SentenceDetector
62	* import com.johnsnowlabs.nlp.embeddings.WordEmbeddingsModel
63	* import com.johnsnowlabs.nlp.annotators.pos.perceptron.PerceptronModel
64	* import com.johnsnowlabs.nlp.annotators.ner.crf.NerCrfModel
65	* import org.apache.spark.ml.Pipeline
66	*
67	* // First extract the prerequisites for the NerCrfModel
68	* val documentAssembler = new DocumentAssembler()
69	* .setInputCol("text")
70	* .setOutputCol("document")
71	*
72	* val sentence = new SentenceDetector()
73	* .setInputCols("document")
74	* .setOutputCol("sentence")
75	*
76	* val tokenizer = new Tokenizer()
77	* .setInputCols("sentence")
78	* .setOutputCol("token")
79	*
80	* val embeddings = WordEmbeddingsModel.pretrained()
81	* .setInputCols("sentence", "token")
82	* .setOutputCol("word_embeddings")
83	*
84	* val posTagger = PerceptronModel.pretrained()
85	* .setInputCols("sentence", "token")
86	* .setOutputCol("pos")
87	*
88	* // Then NER can be extracted
89	* val nerTagger = NerCrfModel.pretrained()
90	* .setInputCols("sentence", "token", "word_embeddings", "pos")
91	* .setOutputCol("ner")
92	*
93	* val pipeline = new Pipeline().setStages(Array(
94	* documentAssembler,
95	* sentence,
96	* tokenizer,
97	* embeddings,
98	* posTagger,
99	* nerTagger
100	* ))
101	*
102	* val data = Seq("U.N. official Ekeus heads for Baghdad.").toDF("text")
103	* val result = pipeline.fit(data).transform(data)
104	*
105	* result.select("ner.result").show(false)
106	* +------------------------------------+
107	* \|result \|
108	* +------------------------------------+
109	* \|[I-ORG, O, O, I-PER, O, O, I-LOC, O]\|
110	* +------------------------------------+
111	* }}}
112	*
113	* @see
114	* [[com.johnsnowlabs.nlp.annotators.ner.dl.NerDLModel NerDLModel]] for a deep learning based
115	* approach
116	* @see
117	* [[com.johnsnowlabs.nlp.annotators.ner.NerConverter NerConverter]] to further process the
118	* results
119	* @param uid
120	* required uid for storing annotator to disk
121	* @groupname anno Annotator types
122	* @groupdesc anno
123	* Required input and expected output annotator types
124	* @groupname Ungrouped Members
125	* @groupname param Parameters
126	* @groupname setParam Parameter setters
127	* @groupname getParam Parameter getters
128	* @groupname Ungrouped Members
129	* @groupprio param 1
130	* @groupprio anno 2
131	* @groupprio Ungrouped 3
132	* @groupprio setParam 4
133	* @groupprio getParam 5
134	* @groupdesc param
135	* A list of (hyper-)parameter keys this annotator can take. Users can set and get the
136	* parameter values through setters and getters, respectively.
137	*/
138	class NerCrfModel(override val uid: String)
139	extends AnnotatorModel[NerCrfModel]
140	with HasSimpleAnnotate[NerCrfModel]
141	with HasStorageRef {
142
143	def this() = this(Identifiable.randomUID("NER"))	1✔
144
145	/** List of Entities to recognize
146	*
147	* @group param
148	*/
149	val entities = new StringArrayParam(this, "entities", "List of Entities to recognize")	1✔
150
151	/** The CRF model
152	*
153	* @group param
154	*/
155	val model: StructFeature[LinearChainCrfModel] =
156	new StructFeature[LinearChainCrfModel](this, "crfModel")	1✔
157
158	/** Additional dictionary to use as for features (Default: `Map.empty[String, String]`)
159	*
160	* @group param
161	*/
162	val dictionaryFeatures: MapFeature[String, String] =
163	new MapFeature[String, String](this, "dictionaryFeatures")	1✔
164
165	/** Whether or not to calculate prediction confidence by token, included in metadata (Default:
166	* `false`)
167	*
168	* @group param
169	*/
170	val includeConfidence = new BooleanParam(	1✔
171	this,
172	"includeConfidence",	1✔
173	"whether or not to calculate prediction confidence by token, includes in metadata")	1✔
174
175	/** @group setParam */
176	def setModel(crf: LinearChainCrfModel): NerCrfModel = set(model, crf)	1✔
177
178	/** @group setParam */
179	def setDictionaryFeatures(dictFeatures: DictionaryFeatures): this.type =
180	set(dictionaryFeatures, dictFeatures.dict)	1✔
181
182	/** @group setParam */
183	def setEntities(toExtract: Array[String]): NerCrfModel = set(entities, toExtract)	1✔
184
185	/** @group setParam */
186	def setIncludeConfidence(c: Boolean): this.type = set(includeConfidence, c)	1✔
187
188	/** @group getParam */
189	def getIncludeConfidence: Boolean = $(includeConfidence)	×
190
191	setDefault(dictionaryFeatures, () => Map.empty[String, String])	1✔
192	setDefault(includeConfidence, false)	1✔
193
194	/** Predicts Named Entities in input sentences
195	*
196	* @param sentences
197	* POS tagged and WordpieceEmbeddings sentences
198	* @return
199	* sentences with recognized Named Entities
200	*/
201	def tag(sentences: Seq[(PosTaggedSentence, WordpieceEmbeddingsSentence)])
202	: Seq[NerTaggedSentence] = {
203	require(model.isSet, "model must be set before tagging")	1✔
204
205	val crf = $$(model)	1✔
206
207	val fg = FeatureGenerator(new DictionaryFeatures($$(dictionaryFeatures)))	1✔
208	sentences.map { case (sentence, withEmbeddings) =>	1✔
209	val instance = fg.generate(sentence, withEmbeddings, crf.metadata)	1✔
210
211	lazy val confidenceValues = {
212	val fb = new FbCalculator(instance.items.length, crf.metadata)
213	fb.calculate(instance, $$(model).weights, 1)
214	fb.alpha
215	}
216
217	val labelIds = crf.predict(instance)	1✔
218
219	val words = sentence.indexedTaggedWords	1✔
220	.zip(labelIds.labels)	1✔
221	.zipWithIndex	1✔
222	.flatMap { case ((word, labelId), idx) =>	1✔
223	val label = crf.metadata.labels(labelId)	1✔
224
225	val alpha = if ($(includeConfidence)) {	×
226	val scores = Some(confidenceValues.apply(idx))	×
227	Some(	×
228	crf.metadata.labels.zipWithIndex	×
229	.filter(x => x._2 != 0)	×
230	.map { case (t, i) =>	×
231	Map(	×
232	t -> scores	×
233	.getOrElse(Array.empty[String])
234	.lift(i)
235	.getOrElse(0.0f)
236	.toString)	×
237	})
238	} else None	1✔
239
240	if (!isDefined(entities) \|\| $(entities).isEmpty \|\| $(entities).contains(label))	1✔
241	Some(IndexedTaggedWord(word.word, label, word.begin, word.end, alpha))	1✔
242	else
243	None	1✔
244	}
245
246	TaggedSentence(words)	1✔
247	}
248	}
249
250	override protected def beforeAnnotate(dataset: Dataset[_]): Dataset[_] = {
251	validateStorageRef(dataset, $(inputCols), AnnotatorType.WORD_EMBEDDINGS)	1✔
252	dataset
253	}
254
255	override def annotate(annotations: Seq[Annotation]): Seq[Annotation] = {
256	val sourceSentences = PosTagged.unpack(annotations)	1✔
257	val withEmbeddings = WordpieceEmbeddingsSentence.unpack(annotations)	1✔
258	val taggedSentences = tag(sourceSentences.zip(withEmbeddings))	1✔
259	NerTagged.pack(taggedSentences)	1✔
260	}
261
262	def shrink(minW: Float): NerCrfModel = set(model, $$(model).shrink(minW))	×
263
264	/** Input Annotator Types: DOCUMENT, TOKEN, POS, WORD_EMBEDDINGS
265	* @group anno
266	*/
267	override val inputAnnotatorTypes: Array[String] = Array(DOCUMENT, TOKEN, POS, WORD_EMBEDDINGS)	1✔
268
269	/** Output Annotator Types: NAMED_ENTITY
270	* @group anno
271	*/
272	override val outputAnnotatorType: AnnotatorType = NAMED_ENTITY	1✔
273
274	}
275
276	trait ReadablePretrainedNerCrf
277	extends ParamsAndFeaturesReadable[NerCrfModel]
278	with HasPretrained[NerCrfModel] {
279	override val defaultModelName: Option[String] = Some("ner_crf")	×
280
281	/** Java compliant-overrides */
282	override def pretrained(): NerCrfModel = super.pretrained()	×
283
284	override def pretrained(name: String): NerCrfModel = super.pretrained(name)	×
285
286	override def pretrained(name: String, lang: String): NerCrfModel = super.pretrained(name, lang)	×
287
288	override def pretrained(name: String, lang: String, remoteLoc: String): NerCrfModel =
289	super.pretrained(name, lang, remoteLoc)	×
290	}
291
292	/** This is the companion object of [[NerCrfModel]]. Please refer to that class for the
293	* documentation.
294	*/
295	object NerCrfModel extends ReadablePretrainedNerCrf

JohnSnowLabs / spark-nlp / 4992350528

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous