2497970187

Build Type

github

Committed by GitHub

Commit Message

<a href="https://github.com/JohnSnowLabs/spark-nlp/commit/<a class=hub.com/JohnSnowLabs/spark-nlp/commit/<a class="double-link" href="https://git"><a class=hub.com/JohnSnowLabs/spark-nlp/commit/<a class="double-link" href="https://git"><a class=hub.com/JohnSnowLabs/spark-nlp/commit/<a class="double-link" href="https://git"><a class=hub.com/JohnSnowLabs/spark-nlp/commit/b2b53a16772702b76e598b1ded0359f43b7526b0">b2b53a167">&lt;a href=&quot;https://github.com/JohnSnowLabs/spark-nlp/commit/</a><a class="double-link" href="https://github.com/JohnSnowLabs/spark-nlp/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/JohnSnowLabs/spark-nlp/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/JohnSnowLabs/spark-nlp/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/JohnSnowLabs/spark-nlp/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/JohnSnowLabs/spark-nlp/commit/b2b53a16772702b76e598b1ded0359f43b7526b0">b2b53a167</a><a href="https://github.com/JohnSnowLabs/spark-nlp/commit/b2b53a16772702b76e598b1ded0359f43b7526b0">&quot;&gt;&amp;lt;a href=&amp;quot;https://github.com/JohnSnowLabs/spark-nlp/commit/&lt;/a&gt;&lt;a class=&quot;double-link&quot; href=&quot;https://github.com/JohnSnowLabs/spark-nlp/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/JohnSnowLabs/spark-nlp/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/JohnSnowLabs/spark-nlp/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/JohnSnowLabs/spark-nlp/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/JohnSnowLabs/spark-nlp/commit/b2b53a16772702b76e598b1ded0359f43b7526b0&quot;&gt;b2b53a167&lt;/a&gt;&lt;a href=&quot;https://github.com/JohnSnowLabs/spark-nlp/commit/b2b53a16772702b76e598b1ded0359f43b7526b0&quot;&gt;&amp;quot;&amp;gt;&amp;amp;lt;a href=&amp;amp;quot;https://github.com/JohnSnowLabs/spark-nlp/commit/&amp;lt;/a&amp;gt;&amp;lt;a class=&amp;quot;double-link&amp... (continued)

Pull Request Pull Request #8320: Release/400 release candidate

Run Details

337 of 337 new or added lines in 18 files covered. (100.0%)

8202 of 11235 relevant lines covered (73.0%)

0.73 hits per line

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

0.0

/src/main/scala/com/johnsnowlabs/nlp/annotators/coref/SpanBertCorefModel.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.coref

import com.johnsnowlabs.ml.tensorflow.{
  ReadTensorflowModel,
  TensorflowSpanBertCoref,
  TensorflowWrapper,
  WriteTensorflowModel
}
import com.johnsnowlabs.nlp.annotators.common._
import com.johnsnowlabs.nlp.annotators.tokenizer.wordpiece.{BasicTokenizer, WordpieceEncoder}
import com.johnsnowlabs.nlp.embeddings.HasEmbeddingsProperties
import com.johnsnowlabs.nlp.serialization.MapFeature
import com.johnsnowlabs.nlp.util.io.{ExternalResource, ReadAs, ResourceHelper}
import com.johnsnowlabs.nlp._
import com.johnsnowlabs.storage.HasStorageRef
import org.apache.spark.broadcast.Broadcast
import org.apache.spark.ml.param.{IntArrayParam, IntParam, Param}
import org.apache.spark.ml.util.Identifiable
import org.apache.spark.sql.SparkSession
import org.slf4j.{Logger, LoggerFactory}

import java.io.File

/** A coreference resolution model based on SpanBert
  *
  * A coreference resolution model identifies expressions which refer to the same entity in a
  * text. For example, given a sentence "John told Mary he would like to borrow a book from her."
  * the model will link "he" to "John" and "her" to "Mary".
  *
  * This model is based on SpanBert, which is fine-tuned on the OntoNotes 5.0 data set.
  *
  * Pretrained models can be loaded with `pretrained` of the companion object:
  * {{{
  * val dependencyParserApproach = SpanBertCorefModel.pretrained()
  *   .setInputCols("sentence", "token")
  *   .setOutputCol("corefs")
  * }}}
  * The default model is `"spanbert_base_coref"`, if no name is provided. For available pretrained
  * models please see the [[https://nlp.johnsnowlabs.com/models Models Hub]].
  *
  * '''Sources:'''
  *   - [[https://github.com/mandarjoshi90/coref]]
  *
  * ==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.annotators.coref.SpanBertCorefModel*
  * import org.apache.spark.ml.Pipeline
  *
  * 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 corefResolution = SpanBertCorefModel.pretrained()
  *   .setInputCols("sentence", "token")
  *   .setOutputCol("corefs")
  *
  * val pipeline = new Pipeline().setStages(Array(
  *   documentAssembler,
  *   sentence,
  *   tokenizer,
  *   corefResolution
  * ))
  *
  * val data = Seq(
  *   "John told Mary he would like to borrow a book from her."
  * ).toDF("text")
  *
  * val result = pipeline.fit(data).transform(data)
  *
  * result.selectExpr(""explode(corefs) AS coref"")
  *   .selectExpr("coref.result as token", "coref.metadata").show(truncate = false)
  * +-----+------------------------------------------------------------------------------------+
  * |token|metadata                                                                            |
  * +-----+------------------------------------------------------------------------------------+
  * |John |{head.sentence -> -1, head -> ROOT, head.begin -> -1, head.end -> -1, sentence -> 0}|
  * |he   |{head.sentence -> 0, head -> John, head.begin -> 0, head.end -> 3, sentence -> 0}   |
  * |Mary |{head.sentence -> -1, head -> ROOT, head.begin -> -1, head.end -> -1, sentence -> 0}|
  * |her  |{head.sentence -> 0, head -> Mary, head.begin -> 10, head.end -> 13, sentence -> 0} |
  * +-----+------------------------------------------------------------------------------------+
  * }}}
  *
  * @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 SpanBertCorefModel(override val uid: String)
    extends AnnotatorModel[SpanBertCorefModel]
    with HasSimpleAnnotate[SpanBertCorefModel]
    with WriteTensorflowModel
    with HasEmbeddingsProperties
    with HasStorageRef
    with HasCaseSensitiveProperties {

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

  override val inputAnnotatorTypes: Array[String] =
    Array(AnnotatorType.DOCUMENT, AnnotatorType.TOKEN)
  override val outputAnnotatorType: AnnotatorType = AnnotatorType.DEPENDENCY

  def sentenceStartTokenId: Int = {
    $$(vocabulary)("[CLS]")
  }

  /** @group setParam */
  def sentenceEndTokenId: Int = {
    $$(vocabulary)("[SEP]")
  }

  /** Vocabulary used to encode the words to ids with WordPieceEncoder
    *
    * @group param
    */
  val vocabulary: MapFeature[String, Int] = new MapFeature(this, "vocabulary")

  /** @group setParam */
  def setVocabulary(value: Map[String, Int]): this.type = set(vocabulary, value)

  /** ConfigProto from tensorflow, serialized into byte array. Get with
    * `config_proto.SerializeToString()`
    *
    * @group param
    */
  val configProtoBytes = new IntArrayParam(
    this,
    "configProtoBytes",
    "ConfigProto from tensorflow, serialized into byte array. Get with config_proto.SerializeToString()")

  /** @group setParam */
  def setConfigProtoBytes(bytes: Array[Int]): SpanBertCorefModel.this.type =
    set(this.configProtoBytes, bytes)

  /** @group getParam */
  def getConfigProtoBytes: Option[Array[Byte]] = get(this.configProtoBytes).map(_.map(_.toByte))

  /** Max sentence length to process (Default: `128`)
    *
    * @group param
    */
  val maxSentenceLength =
    new IntParam(this, "maxSentenceLength", "Max sentence length to process")

  /** @group setParam */
  def setMaxSentenceLength(value: Int): this.type = {
    require(
      value <= 512,
      "BERT models do not support sequences longer than 512 because of trainable positional embeddings.")
    require(value >= 1, "The maxSentenceLength must be at least 1")
    set(maxSentenceLength, value)
    this
  }

  /** @group getParam */
  def getMaxSentenceLength: Int = $(maxSentenceLength)

  /** It contains TF model signatures for the laded saved model
    *
    * @group param
    */
  val signatures = new MapFeature[String, String](model = this, name = "signatures")

  /** @group setParam */
  def setSignatures(value: Map[String, String]): this.type = {
    if (get(signatures).isEmpty)
      set(signatures, value)
    this
  }

  /** @group getParam */
  def getSignatures: Option[Map[String, String]] = get(this.signatures)

  val _textGenres: Array[String] = Array(
    "bc", // Broadcast conversation, default
    "bn", // Broadcast news
    "mz", //
    "nw", // News wire
    "pt", // Pivot text: Old Testament and New Testament text
    "tc", // Telephone conversation
    "wb" // Web data
  )

  /** Text genre, one of the following values: `bc`: Broadcast conversation, default `bn`:
    * Broadcast news `nw`: News wire `pt`: Pivot text: Old Testament and New Testament text `tc`:
    * Telephone conversation `wb`: Web data
    *
    * @group param
    */
  val textGenre =
    new Param[String](
      this,
      "textGenre",
      s"Text genre, one of %s. Default is 'bc'.".format(
        _textGenres.map("\"" + _ + "\"").mkString(", ")))

  /** @group setParam */
  def setTextGenre(value: String): this.type = {
    require(
      Array().contains(value.toLowerCase),
      s"Text text genre must be one of %s".format(
        _textGenres.map("\"" + _ + "\"").mkString(", ")))
    set(textGenre, value.toLowerCase)
    this
  }

  /** @group getParam */
  def getTextGenre: String = $(textGenre)

  /** Max segment length to process (Read-only, depends on model)
    *
    * @group param
    */
  val maxSegmentLength = new IntParam(this, "maxSegmentLength", "Maximum segment length")

  /** @group setParam */
  def setMaxSegmentLength(value: Int): this.type = {
    if (get(maxSegmentLength).isEmpty)
      set(maxSegmentLength, value)
    this
  }

  /** @group getParam */
  def getMaxSegmentLength: Int = $(maxSegmentLength)

  private var _model: Option[Broadcast[TensorflowSpanBertCoref]] = None

  setDefault(
    maxSentenceLength -> 512,
    caseSensitive -> true,
    textGenre -> _textGenres(0)
//    maxSegmentLength -> 384,
  )

  def setModelIfNotSet(
      spark: SparkSession,
      tensorflowWrapper: TensorflowWrapper): SpanBertCorefModel = {
    if (_model.isEmpty) {
      _model = Some(
        spark.sparkContext.broadcast(
          new TensorflowSpanBertCoref(
            tensorflowWrapper,
            sentenceStartTokenId,
            sentenceEndTokenId,
            configProtoBytes = getConfigProtoBytes,
            signatures = getSignatures)))
    }

    this
  }

  def getModelIfNotSet: TensorflowSpanBertCoref = _model.get.value

  def tokenizeSentence(tokens: Seq[TokenizedSentence]): Seq[WordpieceTokenizedSentence] = {
    val basicTokenizer = new BasicTokenizer($(caseSensitive))
    val encoder = new WordpieceEncoder($$(vocabulary))

    tokens.map { tokenIndex =>
      // filter empty and only whitespace tokens
      val bertTokens =
        tokenIndex.indexedTokens.filter(x => x.token.nonEmpty && !x.token.equals(" ")).map {
          token =>
            val content = if ($(caseSensitive)) token.token else token.token.toLowerCase()
            val sentenceBegin = token.begin
            val sentenceEnd = token.end
            val sentenceIndex = tokenIndex.sentenceIndex
            val result = basicTokenizer.tokenize(
              Sentence(content, sentenceBegin, sentenceEnd, sentenceIndex))
            if (result.nonEmpty) result.head else IndexedToken("")
        }
      val wordPieceTokens =
        bertTokens.flatMap(token => encoder.encode(token)).take($(maxSentenceLength) - 2)
      WordpieceTokenizedSentence(wordPieceTokens)
    }
  }

  override def annotate(annotations: Seq[Annotation]): Seq[Annotation] = {

    val sentencesWithRow = TokenizedWithSentence.unpack(annotations)
    val tokenizedSentences = tokenizeSentence(sentencesWithRow).toArray
    val inputIds = tokenizedSentences.map(x => x.tokens.map(_.pieceId))

    if (inputIds.map(x => x.length).sum < 2) {
      return Seq()
    }

    val predictedClusters = getModelIfNotSet.predict(
      inputIds = inputIds,
      genre = _textGenres.indexOf($(textGenre)),
      maxSegmentLength = $(maxSegmentLength))

    def getTokensFromSpan(span: ((Int, Int), (Int, Int))): Array[(TokenPiece, Int)] = {
      val sentence1 = span._1._1
      val sentence2 = span._2._1
      val tokenStart = span._1._2
      val tokenEnd = span._2._2
      if (sentence1 == sentence2) {
        tokenizedSentences(sentence1).tokens.slice(tokenStart, tokenEnd + 1).map((_, sentence1))
      } else {
        (tokenizedSentences(sentence1).tokens
          .slice(tokenStart, tokenizedSentences(sentence1).tokens.length - 1)
          .map((_, sentence1))
          ++
            tokenizedSentences(sentence2).tokens.slice(0, tokenEnd + 1).map((_, sentence2)))
      }
    }

//    predictedClusters.zipWithIndex.foreach{
//      case (cluster, i) =>
//        print(s"Cluster #$i\n")
//        print(s"\t%s\n".format(
//          cluster.map(
//            xy =>
//              getTokensFromSpan(xy).map(x => (if (x.isWordStart) " " else "") + x.wordpiece.replaceFirst("##", "") ).mkString("").trim,
//            ).mkString(", ")))
//    }
    predictedClusters.flatMap(cluster => {

      val clusterSpans = cluster.map(xy => getTokensFromSpan(xy))
      val clusterHeadSpan = clusterSpans.head
      val clusterHeadSpanText = clusterHeadSpan
        .map(x => (if (x._1.isWordStart) " " else "") + x._1.wordpiece.replaceFirst("##", ""))
        .mkString("")
        .trim
      Array(
        Annotation(
          annotatorType = AnnotatorType.DEPENDENCY,
          begin = clusterHeadSpan.head._1.begin,
          end = clusterHeadSpan.last._1.end,
          result = clusterHeadSpanText,
          metadata = Map(
            "head" -> "ROOT",
            "head.begin" -> "-1",
            "head.end" -> "-1",
            "head.sentence" -> "-1",
            "sentence" -> clusterHeadSpan.head._2.toString))) ++ clusterSpans.tail.map(span => {
        Annotation(
          annotatorType = AnnotatorType.DEPENDENCY,
          begin = span.head._1.begin,
          end = span.last._1.end,
          result = span
            .map(x => (if (x._1.isWordStart) " " else "") + x._1.wordpiece.replaceFirst("##", ""))
            .mkString("")
            .trim,
          metadata = Map(
            "head" -> clusterHeadSpanText,
            "head.begin" -> clusterHeadSpan.head._1.begin.toString,
            "head.end" -> clusterHeadSpan.last._1.end.toString,
            "head.sentence" -> clusterHeadSpan.head._2.toString,
            "sentence" -> span.head._2.toString))
      })
    })
  }

  override def onWrite(path: String, spark: SparkSession): Unit = {
    super.onWrite(path, spark)
    writeTensorflowModelV2(
      path,
      spark,
      getModelIfNotSet.tensorflowWrapper,
      "_bert",
      SpanBertCorefModel.tfFile,
      configProtoBytes = getConfigProtoBytes)
  }
}

trait ReadablePretrainedSpanBertCorefModel
    extends ParamsAndFeaturesReadable[SpanBertCorefModel]
    with HasPretrained[SpanBertCorefModel] {
  override val defaultModelName: Some[String] = Some("spanbert_base_coref")

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

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

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

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

trait ReadSpanBertCorefTensorflowModel extends ReadTensorflowModel {
  this: ParamsAndFeaturesReadable[SpanBertCorefModel] =>

  override val tfFile: String = "spanbert_tensorflow"

  def readTensorflow(instance: SpanBertCorefModel, path: String, spark: SparkSession): Unit = {

    val tf = readTensorflowModel(path, spark, "_bert_tf", initAllTables = false)
    instance.setModelIfNotSet(spark, tf)
  }

  addReader(readTensorflow)

  def loadSavedModel(tfModelPath: String, spark: SparkSession): SpanBertCorefModel = {

    val f = new File(tfModelPath)
    val savedModel = new File(tfModelPath, "saved_model.pb")

    require(f.exists, s"Folder $tfModelPath not found")
    require(f.isDirectory, s"File $tfModelPath is not folder")
    require(
      savedModel.exists(),
      s"savedModel file saved_model.pb not found in folder $tfModelPath")

    val vocab = new File(tfModelPath + "/assets", "vocab.txt")

    require(f.exists, s"Folder $tfModelPath not found")
    require(f.isDirectory, s"File $tfModelPath is not folder")
    require(vocab.exists(), s"Vocabulary file vocab.txt not found in folder $tfModelPath")

    val vocabResource =
      new ExternalResource(vocab.getAbsolutePath, ReadAs.TEXT, Map("format" -> "text"))
    val words = ResourceHelper.parseLines(vocabResource).zipWithIndex.toMap

    val (wrapper, signatures) =
      TensorflowWrapper.read(tfModelPath, zipped = false, useBundle = true)

    val _signatures = signatures match {
      case Some(s) => s
      case None => throw new Exception("Cannot load signature definitions from model!")
    }

    /** the order of setSignatures is important if we use getSignatures inside setModelIfNotSet */
    new SpanBertCorefModel()
      .setVocabulary(words)
      .setSignatures(_signatures)
      .setModelIfNotSet(spark, wrapper)
  }
}

/** This is the companion object of [[SpanBertCorefModel]]. Please refer to that class for the
  * documentation.
  */
object SpanBertCorefModel
    extends ReadablePretrainedSpanBertCorefModel
    with ReadSpanBertCorefTensorflowModel {
  private[SpanBertCorefModel] val logger: Logger = LoggerFactory.getLogger("SpanBertCorefModel")
}

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	package com.johnsnowlabs.nlp.annotators.coref
17
18	import com.johnsnowlabs.ml.tensorflow.{
19	ReadTensorflowModel,
20	TensorflowSpanBertCoref,
21	TensorflowWrapper,
22	WriteTensorflowModel
23	}
24	import com.johnsnowlabs.nlp.annotators.common._
25	import com.johnsnowlabs.nlp.annotators.tokenizer.wordpiece.{BasicTokenizer, WordpieceEncoder}
26	import com.johnsnowlabs.nlp.embeddings.HasEmbeddingsProperties
27	import com.johnsnowlabs.nlp.serialization.MapFeature
28	import com.johnsnowlabs.nlp.util.io.{ExternalResource, ReadAs, ResourceHelper}
29	import com.johnsnowlabs.nlp._
30	import com.johnsnowlabs.storage.HasStorageRef
31	import org.apache.spark.broadcast.Broadcast
32	import org.apache.spark.ml.param.{IntArrayParam, IntParam, Param}
33	import org.apache.spark.ml.util.Identifiable
34	import org.apache.spark.sql.SparkSession
35	import org.slf4j.{Logger, LoggerFactory}
36
37	import java.io.File
38
39	/** A coreference resolution model based on SpanBert
40	*
41	* A coreference resolution model identifies expressions which refer to the same entity in a
42	* text. For example, given a sentence "John told Mary he would like to borrow a book from her."
43	* the model will link "he" to "John" and "her" to "Mary".
44	*
45	* This model is based on SpanBert, which is fine-tuned on the OntoNotes 5.0 data set.
46	*
47	* Pretrained models can be loaded with `pretrained` of the companion object:
48	* {{{
49	* val dependencyParserApproach = SpanBertCorefModel.pretrained()
50	* .setInputCols("sentence", "token")
51	* .setOutputCol("corefs")
52	* }}}
53	* The default model is `"spanbert_base_coref"`, if no name is provided. For available pretrained
54	* models please see the [[https://nlp.johnsnowlabs.com/models Models Hub]].
55	*
56	* '''Sources:'''
57	* - [[https://github.com/mandarjoshi90/coref]]
58	*
59	* ==Example==
60	* {{{
61	* import spark.implicits._
62	* import com.johnsnowlabs.nlp.base.DocumentAssembler
63	* import com.johnsnowlabs.nlp.annotators.Tokenizer
64	* import com.johnsnowlabs.nlp.annotators.sbd.pragmatic.SentenceDetector
65	* import com.johnsnowlabs.nlp.annotators.coref.SpanBertCorefModel*
66	* import org.apache.spark.ml.Pipeline
67	*
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 corefResolution = SpanBertCorefModel.pretrained()
81	* .setInputCols("sentence", "token")
82	* .setOutputCol("corefs")
83	*
84	* val pipeline = new Pipeline().setStages(Array(
85	* documentAssembler,
86	* sentence,
87	* tokenizer,
88	* corefResolution
89	* ))
90	*
91	* val data = Seq(
92	* "John told Mary he would like to borrow a book from her."
93	* ).toDF("text")
94	*
95	* val result = pipeline.fit(data).transform(data)
96	*
97	* result.selectExpr(""explode(corefs) AS coref"")
98	* .selectExpr("coref.result as token", "coref.metadata").show(truncate = false)
99	* +-----+------------------------------------------------------------------------------------+
100	* \|token\|metadata \|
101	* +-----+------------------------------------------------------------------------------------+
102	* \|John \|{head.sentence -> -1, head -> ROOT, head.begin -> -1, head.end -> -1, sentence -> 0}\|
103	* \|he \|{head.sentence -> 0, head -> John, head.begin -> 0, head.end -> 3, sentence -> 0} \|
104	* \|Mary \|{head.sentence -> -1, head -> ROOT, head.begin -> -1, head.end -> -1, sentence -> 0}\|
105	* \|her \|{head.sentence -> 0, head -> Mary, head.begin -> 10, head.end -> 13, sentence -> 0} \|
106	* +-----+------------------------------------------------------------------------------------+
107	* }}}
108	*
109	* @groupname anno Annotator types
110	* @groupdesc anno
111	* Required input and expected output annotator types
112	* @groupname Ungrouped Members
113	* @groupname param Parameters
114	* @groupname setParam Parameter setters
115	* @groupname getParam Parameter getters
116	* @groupname Ungrouped Members
117	* @groupprio param 1
118	* @groupprio anno 2
119	* @groupprio Ungrouped 3
120	* @groupprio setParam 4
121	* @groupprio getParam 5
122	* @groupdesc param
123	* A list of (hyper-)parameter keys this annotator can take. Users can set and get the
124	* parameter values through setters and getters, respectively.
125	*/
126	class SpanBertCorefModel(override val uid: String)
127	extends AnnotatorModel[SpanBertCorefModel]
128	with HasSimpleAnnotate[SpanBertCorefModel]
129	with WriteTensorflowModel
130	with HasEmbeddingsProperties
131	with HasStorageRef
132	with HasCaseSensitiveProperties {
133
134	def this() = this(Identifiable.randomUID("SPANBERTCOREFMODEL"))	×
135
136	override val inputAnnotatorTypes: Array[String] =
137	Array(AnnotatorType.DOCUMENT, AnnotatorType.TOKEN)	×
138	override val outputAnnotatorType: AnnotatorType = AnnotatorType.DEPENDENCY	×
139
140	def sentenceStartTokenId: Int = {
141	$$(vocabulary)("[CLS]")	×
142	}
143
144	/** @group setParam */
145	def sentenceEndTokenId: Int = {
146	$$(vocabulary)("[SEP]")	×
147	}
148
149	/** Vocabulary used to encode the words to ids with WordPieceEncoder
150	*
151	* @group param
152	*/
153	val vocabulary: MapFeature[String, Int] = new MapFeature(this, "vocabulary")	×
154
155	/** @group setParam */
156	def setVocabulary(value: Map[String, Int]): this.type = set(vocabulary, value)	×
157
158	/** ConfigProto from tensorflow, serialized into byte array. Get with
159	* `config_proto.SerializeToString()`
160	*
161	* @group param
162	*/
163	val configProtoBytes = new IntArrayParam(	×
164	this,
165	"configProtoBytes",	×
166	"ConfigProto from tensorflow, serialized into byte array. Get with config_proto.SerializeToString()")	×
167
168	/** @group setParam */
169	def setConfigProtoBytes(bytes: Array[Int]): SpanBertCorefModel.this.type =
170	set(this.configProtoBytes, bytes)	×
171
172	/** @group getParam */
173	def getConfigProtoBytes: Option[Array[Byte]] = get(this.configProtoBytes).map(_.map(_.toByte))	×
174
175	/** Max sentence length to process (Default: `128`)
176	*
177	* @group param
178	*/
179	val maxSentenceLength =
180	new IntParam(this, "maxSentenceLength", "Max sentence length to process")	×
181
182	/** @group setParam */
183	def setMaxSentenceLength(value: Int): this.type = {
184	require(	×
185	value <= 512,	×
186	"BERT models do not support sequences longer than 512 because of trainable positional embeddings.")	×
187	require(value >= 1, "The maxSentenceLength must be at least 1")	×
188	set(maxSentenceLength, value)	×
189	this
190	}
191
192	/** @group getParam */
193	def getMaxSentenceLength: Int = $(maxSentenceLength)	×
194
195	/** It contains TF model signatures for the laded saved model
196	*
197	* @group param
198	*/
199	val signatures = new MapFeature[String, String](model = this, name = "signatures")	×
200
201	/** @group setParam */
202	def setSignatures(value: Map[String, String]): this.type = {
203	if (get(signatures).isEmpty)	×
204	set(signatures, value)	×
205	this
206	}
207
208	/** @group getParam */
209	def getSignatures: Option[Map[String, String]] = get(this.signatures)	×
210
211	val _textGenres: Array[String] = Array(	×
212	"bc", // Broadcast conversation, default	×
213	"bn", // Broadcast news	×
214	"mz", //	×
215	"nw", // News wire	×
216	"pt", // Pivot text: Old Testament and New Testament text	×
217	"tc", // Telephone conversation	×
218	"wb" // Web data	×
219	)
220
221	/** Text genre, one of the following values: `bc`: Broadcast conversation, default `bn`:
222	* Broadcast news `nw`: News wire `pt`: Pivot text: Old Testament and New Testament text `tc`:
223	* Telephone conversation `wb`: Web data
224	*
225	* @group param
226	*/
227	val textGenre =
228	new Param[String](	×
229	this,
230	"textGenre",	×
231	s"Text genre, one of %s. Default is 'bc'.".format(	×
232	_textGenres.map("\"" + _ + "\"").mkString(", ")))	×
233
234	/** @group setParam */
235	def setTextGenre(value: String): this.type = {
236	require(	×
237	Array().contains(value.toLowerCase),	×
238	s"Text text genre must be one of %s".format(	×
239	_textGenres.map("\"" + _ + "\"").mkString(", ")))	×
240	set(textGenre, value.toLowerCase)	×
241	this
242	}
243
244	/** @group getParam */
245	def getTextGenre: String = $(textGenre)	×
246
247	/** Max segment length to process (Read-only, depends on model)
248	*
249	* @group param
250	*/
251	val maxSegmentLength = new IntParam(this, "maxSegmentLength", "Maximum segment length")	×
252
253	/** @group setParam */
254	def setMaxSegmentLength(value: Int): this.type = {
255	if (get(maxSegmentLength).isEmpty)	×
256	set(maxSegmentLength, value)	×
257	this
258	}
259
260	/** @group getParam */
261	def getMaxSegmentLength: Int = $(maxSegmentLength)	×
262
263	private var _model: Option[Broadcast[TensorflowSpanBertCoref]] = None	×
264
265	setDefault(	×
266	maxSentenceLength -> 512,	×
267	caseSensitive -> true,	×
268	textGenre -> _textGenres(0)	×
269	// maxSegmentLength -> 384,
270	)
271
272	def setModelIfNotSet(
273	spark: SparkSession,
274	tensorflowWrapper: TensorflowWrapper): SpanBertCorefModel = {
275	if (_model.isEmpty) {	×
276	_model = Some(	×
277	spark.sparkContext.broadcast(	×
278	new TensorflowSpanBertCoref(	×
279	tensorflowWrapper,
280	sentenceStartTokenId,	×
281	sentenceEndTokenId,	×
282	configProtoBytes = getConfigProtoBytes,	×
283	signatures = getSignatures)))	×
284	}
285
286	this
287	}
288
289	def getModelIfNotSet: TensorflowSpanBertCoref = _model.get.value	×
290
291	def tokenizeSentence(tokens: Seq[TokenizedSentence]): Seq[WordpieceTokenizedSentence] = {
292	val basicTokenizer = new BasicTokenizer($(caseSensitive))	×
293	val encoder = new WordpieceEncoder($$(vocabulary))	×
294
295	tokens.map { tokenIndex =>	×
296	// filter empty and only whitespace tokens
297	val bertTokens =
298	tokenIndex.indexedTokens.filter(x => x.token.nonEmpty && !x.token.equals(" ")).map {	×
299	token =>
300	val content = if ($(caseSensitive)) token.token else token.token.toLowerCase()	×
301	val sentenceBegin = token.begin	×
302	val sentenceEnd = token.end	×
303	val sentenceIndex = tokenIndex.sentenceIndex	×
304	val result = basicTokenizer.tokenize(	×
305	Sentence(content, sentenceBegin, sentenceEnd, sentenceIndex))	×
306	if (result.nonEmpty) result.head else IndexedToken("")	×
307	}
308	val wordPieceTokens =
309	bertTokens.flatMap(token => encoder.encode(token)).take($(maxSentenceLength) - 2)	×
310	WordpieceTokenizedSentence(wordPieceTokens)	×
311	}
312	}
313
314	override def annotate(annotations: Seq[Annotation]): Seq[Annotation] = {
315
316	val sentencesWithRow = TokenizedWithSentence.unpack(annotations)	×
317	val tokenizedSentences = tokenizeSentence(sentencesWithRow).toArray	×
318	val inputIds = tokenizedSentences.map(x => x.tokens.map(_.pieceId))	×
319
320	if (inputIds.map(x => x.length).sum < 2) {	×
321	return Seq()	×
322	}
323
324	val predictedClusters = getModelIfNotSet.predict(	×
325	inputIds = inputIds,
326	genre = _textGenres.indexOf($(textGenre)),	×
327	maxSegmentLength = $(maxSegmentLength))	×
328
329	def getTokensFromSpan(span: ((Int, Int), (Int, Int))): Array[(TokenPiece, Int)] = {
330	val sentence1 = span._1._1	×
331	val sentence2 = span._2._1	×
332	val tokenStart = span._1._2	×
333	val tokenEnd = span._2._2	×
334	if (sentence1 == sentence2) {	×
335	tokenizedSentences(sentence1).tokens.slice(tokenStart, tokenEnd + 1).map((_, sentence1))	×
336	} else {
337	(tokenizedSentences(sentence1).tokens	×
338	.slice(tokenStart, tokenizedSentences(sentence1).tokens.length - 1)	×
339	.map((_, sentence1))	×
340	++	×
341	tokenizedSentences(sentence2).tokens.slice(0, tokenEnd + 1).map((_, sentence2)))	×
342	}
343	}
344
345	// predictedClusters.zipWithIndex.foreach{
346	// case (cluster, i) =>
347	// print(s"Cluster #$i\n")
348	// print(s"\t%s\n".format(
349	// cluster.map(
350	// xy =>
351	// getTokensFromSpan(xy).map(x => (if (x.isWordStart) " " else "") + x.wordpiece.replaceFirst("##", "") ).mkString("").trim,
352	// ).mkString(", ")))
353	// }
354	predictedClusters.flatMap(cluster => {	×
355
356	val clusterSpans = cluster.map(xy => getTokensFromSpan(xy))	×
357	val clusterHeadSpan = clusterSpans.head	×
358	val clusterHeadSpanText = clusterHeadSpan
359	.map(x => (if (x._1.isWordStart) " " else "") + x._1.wordpiece.replaceFirst("##", ""))
360	.mkString("")
361	.trim	×
362	Array(	×
363	Annotation(	×
364	annotatorType = AnnotatorType.DEPENDENCY,	×
365	begin = clusterHeadSpan.head._1.begin,	×
366	end = clusterHeadSpan.last._1.end,	×
367	result = clusterHeadSpanText,
368	metadata = Map(	×
369	"head" -> "ROOT",	×
370	"head.begin" -> "-1",	×
371	"head.end" -> "-1",	×
372	"head.sentence" -> "-1",	×
373	"sentence" -> clusterHeadSpan.head._2.toString))) ++ clusterSpans.tail.map(span => {	×
374	Annotation(	×
375	annotatorType = AnnotatorType.DEPENDENCY,	×
376	begin = span.head._1.begin,	×
377	end = span.last._1.end,	×
378	result = span
379	.map(x => (if (x._1.isWordStart) " " else "") + x._1.wordpiece.replaceFirst("##", ""))
380	.mkString("")
381	.trim,	×
382	metadata = Map(	×
383	"head" -> clusterHeadSpanText,	×
384	"head.begin" -> clusterHeadSpan.head._1.begin.toString,	×
385	"head.end" -> clusterHeadSpan.last._1.end.toString,	×
386	"head.sentence" -> clusterHeadSpan.head._2.toString,	×
387	"sentence" -> span.head._2.toString))	×
388	})
389	})
390	}
391
392	override def onWrite(path: String, spark: SparkSession): Unit = {
393	super.onWrite(path, spark)	×
394	writeTensorflowModelV2(	×
395	path,
396	spark,
397	getModelIfNotSet.tensorflowWrapper,	×
398	"_bert",	×
399	SpanBertCorefModel.tfFile,	×
400	configProtoBytes = getConfigProtoBytes)	×
401	}
402	}
403
404	trait ReadablePretrainedSpanBertCorefModel
405	extends ParamsAndFeaturesReadable[SpanBertCorefModel]
406	with HasPretrained[SpanBertCorefModel] {
407	override val defaultModelName: Some[String] = Some("spanbert_base_coref")	×
408
409	/** Java compliant-overrides */
410	override def pretrained(): SpanBertCorefModel = super.pretrained()	×
411
412	override def pretrained(name: String): SpanBertCorefModel = super.pretrained(name)	×
413
414	override def pretrained(name: String, lang: String): SpanBertCorefModel =
415	super.pretrained(name, lang)	×
416
417	override def pretrained(name: String, lang: String, remoteLoc: String): SpanBertCorefModel =
418	super.pretrained(name, lang, remoteLoc)	×
419	}
420
421	trait ReadSpanBertCorefTensorflowModel extends ReadTensorflowModel {
422	this: ParamsAndFeaturesReadable[SpanBertCorefModel] =>
423
424	override val tfFile: String = "spanbert_tensorflow"	×
425
426	def readTensorflow(instance: SpanBertCorefModel, path: String, spark: SparkSession): Unit = {
427
428	val tf = readTensorflowModel(path, spark, "_bert_tf", initAllTables = false)	×
429	instance.setModelIfNotSet(spark, tf)	×
430	}
431
432	addReader(readTensorflow)	×
433
434	def loadSavedModel(tfModelPath: String, spark: SparkSession): SpanBertCorefModel = {
435
436	val f = new File(tfModelPath)	×
437	val savedModel = new File(tfModelPath, "saved_model.pb")	×
438
439	require(f.exists, s"Folder $tfModelPath not found")	×
440	require(f.isDirectory, s"File $tfModelPath is not folder")	×
441	require(	×
442	savedModel.exists(),	×
443	s"savedModel file saved_model.pb not found in folder $tfModelPath")	×
444
445	val vocab = new File(tfModelPath + "/assets", "vocab.txt")	×
446
447	require(f.exists, s"Folder $tfModelPath not found")	×
448	require(f.isDirectory, s"File $tfModelPath is not folder")	×
449	require(vocab.exists(), s"Vocabulary file vocab.txt not found in folder $tfModelPath")	×
450
451	val vocabResource =
452	new ExternalResource(vocab.getAbsolutePath, ReadAs.TEXT, Map("format" -> "text"))	×
453	val words = ResourceHelper.parseLines(vocabResource).zipWithIndex.toMap	×
454
455	val (wrapper, signatures) =	×
456	TensorflowWrapper.read(tfModelPath, zipped = false, useBundle = true)
457
458	val _signatures = signatures match {
459	case Some(s) => s
460	case None => throw new Exception("Cannot load signature definitions from model!")	×
461	}
462
463	/** the order of setSignatures is important if we use getSignatures inside setModelIfNotSet */
464	new SpanBertCorefModel()
465	.setVocabulary(words)
466	.setSignatures(_signatures)
467	.setModelIfNotSet(spark, wrapper)	×
468	}
469	}
470
471	/** This is the companion object of [[SpanBertCorefModel]]. Please refer to that class for the
472	* documentation.
473	*/
474	object SpanBertCorefModel
475	extends ReadablePretrainedSpanBertCorefModel
476	with ReadSpanBertCorefTensorflowModel {
477	private[SpanBertCorefModel] val logger: Logger = LoggerFactory.getLogger("SpanBertCorefModel")	×
478	}

JohnSnowLabs / spark-nlp / 2497970187

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