1576910141

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/784e6b9419532ab8390c81fb0f76c23001d7f3db">784e6b941">&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/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/JohnSnowLabs/spark-nlp/commit/784e6b9419532ab8390c81fb0f76c23001d7f3db">784e6b941</a><a href="https://github.com/JohnSnowLabs/spark-nlp/commit/784e6b9419532ab8390c81fb0f76c23001d7f3db">&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/&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/784e6b9419532ab8390c81fb0f76c23001d7f3db&quot;&gt;784e6b941&lt;/a&gt;&lt;a href=&q... (continued)

Pull Request Pull Request #6607: Feature/sbt 156

Run Details

1356 of 1356 new or added lines in 68 files covered. (100.0%)

7939 of 16034 relevant lines covered (49.51%)

0.5 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/classifier/dl/RoBertaForSequenceClassification.scala

/*
 * Copyright 2017-2021 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.classifier.dl

import com.johnsnowlabs.ml.tensorflow._
import com.johnsnowlabs.nlp._
import com.johnsnowlabs.nlp.annotators.common._
import com.johnsnowlabs.nlp.serialization.MapFeature
import com.johnsnowlabs.nlp.util.io.{ExternalResource, ReadAs, ResourceHelper}
import org.apache.spark.broadcast.Broadcast
import org.apache.spark.ml.param.{BooleanParam, IntArrayParam, IntParam}
import org.apache.spark.ml.util.Identifiable
import org.apache.spark.sql.SparkSession

import java.io.File

/**
 * RoBertaForSequenceClassification can load RoBERTa Models with sequence classification/regression head on top
 * (a linear layer on top of the pooled output) e.g. for multi-class document classification tasks.
 *
 * Pretrained models can be loaded with `pretrained` of the companion object:
 * {{{
 * val sequenceClassifier = RoBertaForSequenceClassification.pretrained()
 *   .setInputCols("token", "document")
 *   .setOutputCol("label")
 * }}}
 * The default model is `"roberta_base_sequence_classifier_imdb"`, if no name is provided.
 *
 * For available pretrained models please see the [[https://nlp.johnsnowlabs.com/models?task=Text+Classification Models Hub]].
 *
 * Models from the HuggingFace 🤗 Transformers library are also compatible with Spark NLP 🚀. The Spark NLP Workshop
 * example shows how to import them [[https://github.com/JohnSnowLabs/spark-nlp/discussions/5669]].
 * and the [[https://github.com/JohnSnowLabs/spark-nlp/blob/master/src/test/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/RoBertaForSequenceClassificationTestSpec.scala RoBertaForSequenceClassification]].
 *
 * ==Example==
 * {{{
 * import spark.implicits._
 * import com.johnsnowlabs.nlp.base._
 * import com.johnsnowlabs.nlp.annotator._
 * import org.apache.spark.ml.Pipeline
 *
 * val documentAssembler = new DocumentAssembler()
 *   .setInputCol("text")
 *   .setOutputCol("document")
 *
 * val tokenizer = new Tokenizer()
 *   .setInputCols("document")
 *   .setOutputCol("token")
 *
 * val sequenceClassifier = RoBertaForSequenceClassification.pretrained()
 *   .setInputCols("token", "document")
 *   .setOutputCol("label")
 *   .setCaseSensitive(true)
 *
 * val pipeline = new Pipeline().setStages(Array(
 *   documentAssembler,
 *   tokenizer,
 *   sequenceClassifier
 * ))
 *
 * val data = Seq("John Lenon was born in London and lived in Paris. My name is Sarah and I live in London").toDF("text")
 * val result = pipeline.fit(data).transform(data)
 *
 * result.select("label.result").show(false)
 * +--------------------+
 * |result              |
 * +--------------------+
 * |[neg, neg]          |
 * |[pos, pos, pos, pos]|
 * +--------------------+
 * }}}
 *
 * @see [[RoBertaForSequenceClassification]] for sequence-level classification
 * @see [[https://nlp.johnsnowlabs.com/docs/en/annotators Annotators Main Page]] for a list of transformer based classifiers
 * @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 RoBertaForSequenceClassification(override val uid: String)
  extends AnnotatorModel[RoBertaForSequenceClassification]
    with HasBatchedAnnotate[RoBertaForSequenceClassification]
    with WriteTensorflowModel
    with HasCaseSensitiveProperties {

  /** Annotator reference id. Used to identify elements in metadata or to refer to this annotator type */
  def this() = this(Identifiable.randomUID("RoBertaForSequenceClassification"))

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

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

  def sentenceStartTokenId: Int = {
    $$(vocabulary)("<s>")
  }

  def sentenceEndTokenId: Int = {
    $$(vocabulary)("</s>")
  }

  def padTokenId: Int = {
    $$(vocabulary)("<pad>")
  }


  /**
   * 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)


  /**
   * Labels used to decode predicted IDs back to string tags
   *
   * @group param
   * */
  val labels: MapFeature[String, Int] = new MapFeature(this, "labels")

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

  /**
   * Returns labels used to train this model
   */
  def getClasses: Array[String] = {
    $$(labels).keys.toArray
  }

  /**
   * Holding merges.txt coming from RoBERTa model
   *
   * @group param
   */
  val merges: MapFeature[(String, String), Int] = new MapFeature(this, "merges")

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

  /** Instead of 1 class per sentence (if inputCols is '''sentence''') output 1 class per document by averaging probabilities in all sentences.
   * Due to max sequence length limit in almost all transformer models such as BERT (512 tokens), this parameter helps feeding all the sentences
   * into the model and averaging all the probabilities for the entire document instead of probabilities per sentence. (Default: true)
   *
   * @group param
   * */
  val coalesceSentences = new BooleanParam(this, "coalesceSentences", "If sets to true the output of all sentences will be averaged to one output instead of one output per sentence. Default to true.")

  /** @group setParam */
  def setCoalesceSentences(value: Boolean): this.type = set(coalesceSentences, value)

  /** @group getParam */
  def getCoalesceSentences: Boolean = $(coalesceSentences)

  /** 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]): RoBertaForSequenceClassification.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, "RoBERTa 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)

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

  /** @group setParam */
  def setModelIfNotSet(spark: SparkSession, tensorflowWrapper: TensorflowWrapper): RoBertaForSequenceClassification = {
    if (_model.isEmpty) {
      _model = Some(
        spark.sparkContext.broadcast(
          new TensorflowRoBertaClassification(
            tensorflowWrapper,
            sentenceStartTokenId,
            sentenceEndTokenId,
            padTokenId,
            configProtoBytes = getConfigProtoBytes,
            tags = $$(labels),
            signatures = getSignatures,
            $$(merges),
            $$(vocabulary)
          )
        )
      )
    }

    this
  }

  /** @group getParam */
  def getModelIfNotSet: TensorflowRoBertaClassification = _model.get.value


  /** Whether to lowercase tokens or not
   *
   * @group setParam
   * */
  override def setCaseSensitive(value: Boolean): this.type = {
    if (get(caseSensitive).isEmpty)
      set(this.caseSensitive, value)
    this
  }

  setDefault(
    batchSize -> 8,
    maxSentenceLength -> 128,
    caseSensitive -> true,
    coalesceSentences -> false
  )

  /**
   * takes a document and annotations and produces new annotations of this annotator's annotation type
   *
   * @param batchedAnnotations Annotations that correspond to inputAnnotationCols generated by previous annotators if any
   * @return any number of annotations processed for every input annotation. Not necessary one to one relationship
   */
  override def batchAnnotate(batchedAnnotations: Seq[Array[Annotation]]): Seq[Seq[Annotation]] = {
    batchedAnnotations.map(annotations => {
      val sentences = SentenceSplit.unpack(annotations).toArray
      val tokenizedSentences = TokenizedWithSentence.unpack(annotations).toArray

      if (tokenizedSentences.nonEmpty) {
        getModelIfNotSet.predictSequence(
          tokenizedSentences,
          sentences,
          $(batchSize),
          $(maxSentenceLength),
          $(caseSensitive),
          $(coalesceSentences),
          $$(labels)
        )
      }
      else {
        Seq.empty[Annotation]
      }
    }

    )
  }


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

}

trait ReadablePretrainedRoBertaForSequenceModel extends ParamsAndFeaturesReadable[RoBertaForSequenceClassification] with HasPretrained[RoBertaForSequenceClassification] {
  override val defaultModelName: Some[String] = Some("roberta_base_sequence_classifier_imdb")

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

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

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

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

trait ReadRoBertaForSequenceTensorflowModel extends ReadTensorflowModel {
  this: ParamsAndFeaturesReadable[RoBertaForSequenceClassification] =>

  override val tfFile: String = "roberta_classification_tensorflow"

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

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

  addReader(readTensorflow)

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

    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 vocabFile = new File(tfModelPath + "/assets", "vocab.txt")
    require(f.exists, s"Folder $tfModelPath not found")
    require(f.isDirectory, s"File $tfModelPath is not folder")
    require(vocabFile.exists(), s"Vocabulary file vocab.txt not found in folder $tfModelPath")

    val mergesFile = new File(tfModelPath + "/assets", "merges.txt")
    require(f.exists, s"Folder $tfModelPath not found")
    require(f.isDirectory, s"File $tfModelPath is not folder")
    require(mergesFile.exists(), s"merges file merges.txt not found in folder $tfModelPath")

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

    val mergesResource = new ExternalResource(mergesFile.getAbsolutePath, ReadAs.TEXT, Map("format" -> "text"))
    val merges = ResourceHelper.parseLines(mergesResource)

    val bytePairs: Map[(String, String), Int] = merges.map(_.split(" "))
      .filter(w => w.length == 2)
      .map { case Array(c1, c2) => (c1, c2) }
      .zipWithIndex.toMap

    val labelsPath = new File(tfModelPath + "/assets", "labels.txt")
    require(labelsPath.exists(), s"Labels file labels.txt not found in folder $tfModelPath/assets/")

    val labelsResource = new ExternalResource(labelsPath.getAbsolutePath, ReadAs.TEXT, Map("format" -> "text"))
    val labels = ResourceHelper.parseLines(labelsResource).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 RoBertaForSequenceClassification()
      .setVocabulary(words)
      .setMerges(bytePairs)
      .setLabels(labels)
      .setSignatures(_signatures)
      .setModelIfNotSet(spark, wrapper)
  }
}

/**
 * This is the companion object of [[RoBertaForSequenceClassification]]. Please refer to that class for the documentation.
 */
object RoBertaForSequenceClassification extends ReadablePretrainedRoBertaForSequenceModel with ReadRoBertaForSequenceTensorflowModel

1	/*
2	* Copyright 2017-2021 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.classifier.dl
18
19	import com.johnsnowlabs.ml.tensorflow._
20	import com.johnsnowlabs.nlp._
21	import com.johnsnowlabs.nlp.annotators.common._
22	import com.johnsnowlabs.nlp.serialization.MapFeature
23	import com.johnsnowlabs.nlp.util.io.{ExternalResource, ReadAs, ResourceHelper}
24	import org.apache.spark.broadcast.Broadcast
25	import org.apache.spark.ml.param.{BooleanParam, IntArrayParam, IntParam}
26	import org.apache.spark.ml.util.Identifiable
27	import org.apache.spark.sql.SparkSession
28
29	import java.io.File
30
31	/**
32	* RoBertaForSequenceClassification can load RoBERTa Models with sequence classification/regression head on top
33	* (a linear layer on top of the pooled output) e.g. for multi-class document classification tasks.
34	*
35	* Pretrained models can be loaded with `pretrained` of the companion object:
36	* {{{
37	* val sequenceClassifier = RoBertaForSequenceClassification.pretrained()
38	* .setInputCols("token", "document")
39	* .setOutputCol("label")
40	* }}}
41	* The default model is `"roberta_base_sequence_classifier_imdb"`, if no name is provided.
42	*
43	* For available pretrained models please see the [[https://nlp.johnsnowlabs.com/models?task=Text+Classification Models Hub]].
44	*
45	* Models from the HuggingFace 🤗 Transformers library are also compatible with Spark NLP 🚀. The Spark NLP Workshop
46	* example shows how to import them [[https://github.com/JohnSnowLabs/spark-nlp/discussions/5669]].
47	* and the [[https://github.com/JohnSnowLabs/spark-nlp/blob/master/src/test/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/RoBertaForSequenceClassificationTestSpec.scala RoBertaForSequenceClassification]].
48	*
49	* ==Example==
50	* {{{
51	* import spark.implicits._
52	* import com.johnsnowlabs.nlp.base._
53	* import com.johnsnowlabs.nlp.annotator._
54	* import org.apache.spark.ml.Pipeline
55	*
56	* val documentAssembler = new DocumentAssembler()
57	* .setInputCol("text")
58	* .setOutputCol("document")
59	*
60	* val tokenizer = new Tokenizer()
61	* .setInputCols("document")
62	* .setOutputCol("token")
63	*
64	* val sequenceClassifier = RoBertaForSequenceClassification.pretrained()
65	* .setInputCols("token", "document")
66	* .setOutputCol("label")
67	* .setCaseSensitive(true)
68	*
69	* val pipeline = new Pipeline().setStages(Array(
70	* documentAssembler,
71	* tokenizer,
72	* sequenceClassifier
73	* ))
74	*
75	* val data = Seq("John Lenon was born in London and lived in Paris. My name is Sarah and I live in London").toDF("text")
76	* val result = pipeline.fit(data).transform(data)
77	*
78	* result.select("label.result").show(false)
79	* +--------------------+
80	* \|result \|
81	* +--------------------+
82	* \|[neg, neg] \|
83	* \|[pos, pos, pos, pos]\|
84	* +--------------------+
85	* }}}
86	*
87	* @see [[RoBertaForSequenceClassification]] for sequence-level classification
88	* @see [[https://nlp.johnsnowlabs.com/docs/en/annotators Annotators Main Page]] for a list of transformer based classifiers
89	* @param uid required uid for storing annotator to disk
90	* @groupname anno Annotator types
91	* @groupdesc anno Required input and expected output annotator types
92	* @groupname Ungrouped Members
93	* @groupname param Parameters
94	* @groupname setParam Parameter setters
95	* @groupname getParam Parameter getters
96	* @groupname Ungrouped Members
97	* @groupprio param 1
98	* @groupprio anno 2
99	* @groupprio Ungrouped 3
100	* @groupprio setParam 4
101	* @groupprio getParam 5
102	* @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.
103	* */
104	class RoBertaForSequenceClassification(override val uid: String)
105	extends AnnotatorModel[RoBertaForSequenceClassification]
106	with HasBatchedAnnotate[RoBertaForSequenceClassification]
107	with WriteTensorflowModel
108	with HasCaseSensitiveProperties {
109
110	/** Annotator reference id. Used to identify elements in metadata or to refer to this annotator type */
111	def this() = this(Identifiable.randomUID("RoBertaForSequenceClassification"))	×
112
113	/**
114	* Input Annotator Types: DOCUMENT, TOKEN
115	*
116	* @group anno
117	*/
118	override val inputAnnotatorTypes: Array[String] = Array(AnnotatorType.DOCUMENT, AnnotatorType.TOKEN)	×
119
120	/**
121	* Output Annotator Types: CATEGORY
122	*
123	* @group anno
124	*/
125	override val outputAnnotatorType: AnnotatorType = AnnotatorType.CATEGORY	×
126
127	def sentenceStartTokenId: Int = {
128	$$(vocabulary)("<s>")	×
129	}
130
131	def sentenceEndTokenId: Int = {
132	$$(vocabulary)("</s>")	×
133	}
134
135	def padTokenId: Int = {
136	$$(vocabulary)("<pad>")	×
137	}
138
139
140	/**
141	* Vocabulary used to encode the words to ids with WordPieceEncoder
142	*
143	* @group param
144	* */
145	val vocabulary: MapFeature[String, Int] = new MapFeature(this, "vocabulary")	×
146
147	/** @group setParam */
148	def setVocabulary(value: Map[String, Int]): this.type = set(vocabulary, value)	×
149
150
151	/**
152	* Labels used to decode predicted IDs back to string tags
153	*
154	* @group param
155	* */
156	val labels: MapFeature[String, Int] = new MapFeature(this, "labels")	×
157
158	/** @group setParam */
159	def setLabels(value: Map[String, Int]): this.type = set(labels, value)	×
160
161	/**
162	* Returns labels used to train this model
163	*/
164	def getClasses: Array[String] = {
165	$$(labels).keys.toArray	×
166	}
167
168	/**
169	* Holding merges.txt coming from RoBERTa model
170	*
171	* @group param
172	*/
173	val merges: MapFeature[(String, String), Int] = new MapFeature(this, "merges")	×
174
175	/** @group setParam */
176	def setMerges(value: Map[(String, String), Int]): this.type = set(merges, value)	×
177
178	/** Instead of 1 class per sentence (if inputCols is '''sentence''') output 1 class per document by averaging probabilities in all sentences.
179	* Due to max sequence length limit in almost all transformer models such as BERT (512 tokens), this parameter helps feeding all the sentences
180	* into the model and averaging all the probabilities for the entire document instead of probabilities per sentence. (Default: true)
181	*
182	* @group param
183	* */
184	val coalesceSentences = new BooleanParam(this, "coalesceSentences", "If sets to true the output of all sentences will be averaged to one output instead of one output per sentence. Default to true.")	×
185
186	/** @group setParam */
187	def setCoalesceSentences(value: Boolean): this.type = set(coalesceSentences, value)	×
188
189	/** @group getParam */
190	def getCoalesceSentences: Boolean = $(coalesceSentences)	×
191
192	/** ConfigProto from tensorflow, serialized into byte array. Get with `config_proto.SerializeToString()`
193	*
194	* @group param
195	* */
196	val configProtoBytes = new IntArrayParam(this, "configProtoBytes", "ConfigProto from tensorflow, serialized into byte array. Get with config_proto.SerializeToString()")	×
197
198	/** @group setParam */
199	def setConfigProtoBytes(bytes: Array[Int]): RoBertaForSequenceClassification.this.type = set(this.configProtoBytes, bytes)	×
200
201	/** @group getParam */
202	def getConfigProtoBytes: Option[Array[Byte]] = get(this.configProtoBytes).map(_.map(_.toByte))	×
203
204	/** Max sentence length to process (Default: `128`)
205	*
206	* @group param
207	* */
208	val maxSentenceLength = new IntParam(this, "maxSentenceLength", "Max sentence length to process")	×
209
210	/** @group setParam */
211	def setMaxSentenceLength(value: Int): this.type = {
212	require(value <= 512, "RoBERTa models do not support sequences longer than 512 because of trainable positional embeddings.")	×
213	require(value >= 1, "The maxSentenceLength must be at least 1")	×
214	set(maxSentenceLength, value)	×
215	this
216	}
217
218	/** @group getParam */
219	def getMaxSentenceLength: Int = $(maxSentenceLength)	×
220
221	/**
222	* It contains TF model signatures for the laded saved model
223	*
224	* @group param
225	* */
226	val signatures = new MapFeature[String, String](model = this, name = "signatures")	×
227
228	/** @group setParam */
229	def setSignatures(value: Map[String, String]): this.type = {
230	if (get(signatures).isEmpty)	×
231	set(signatures, value)	×
232	this
233	}
234
235	/** @group getParam */
236	def getSignatures: Option[Map[String, String]] = get(this.signatures)	×
237
238	private var _model: Option[Broadcast[TensorflowRoBertaClassification]] = None	×
239
240	/** @group setParam */
241	def setModelIfNotSet(spark: SparkSession, tensorflowWrapper: TensorflowWrapper): RoBertaForSequenceClassification = {
242	if (_model.isEmpty) {	×
243	_model = Some(	×
244	spark.sparkContext.broadcast(	×
245	new TensorflowRoBertaClassification(	×
246	tensorflowWrapper,
247	sentenceStartTokenId,	×
248	sentenceEndTokenId,	×
249	padTokenId,	×
250	configProtoBytes = getConfigProtoBytes,	×
251	tags = $$(labels),	×
252	signatures = getSignatures,	×
253	$$(merges),	×
254	$$(vocabulary)	×
255	)
256	)
257	)
258	}
259
260	this
261	}
262
263	/** @group getParam */
264	def getModelIfNotSet: TensorflowRoBertaClassification = _model.get.value	×
265
266
267	/** Whether to lowercase tokens or not
268	*
269	* @group setParam
270	* */
271	override def setCaseSensitive(value: Boolean): this.type = {
272	if (get(caseSensitive).isEmpty)	×
273	set(this.caseSensitive, value)	×
274	this
275	}
276
277	setDefault(	×
278	batchSize -> 8,	×
279	maxSentenceLength -> 128,	×
280	caseSensitive -> true,	×
281	coalesceSentences -> false	×
282	)
283
284	/**
285	* takes a document and annotations and produces new annotations of this annotator's annotation type
286	*
287	* @param batchedAnnotations Annotations that correspond to inputAnnotationCols generated by previous annotators if any
288	* @return any number of annotations processed for every input annotation. Not necessary one to one relationship
289	*/
290	override def batchAnnotate(batchedAnnotations: Seq[Array[Annotation]]): Seq[Seq[Annotation]] = {
291	batchedAnnotations.map(annotations => {	×
292	val sentences = SentenceSplit.unpack(annotations).toArray	×
293	val tokenizedSentences = TokenizedWithSentence.unpack(annotations).toArray	×
294
295	if (tokenizedSentences.nonEmpty) {	×
296	getModelIfNotSet.predictSequence(	×
297	tokenizedSentences,	×
298	sentences,	×
299	$(batchSize),	×
300	$(maxSentenceLength),	×
301	$(caseSensitive),	×
302	$(coalesceSentences),	×
303	$$(labels)	×
304	)
305	}
306	else {
307	Seq.empty[Annotation]	×
308	}
309	}
310
311	)
312	}
313
314
315	override def onWrite(path: String, spark: SparkSession): Unit = {
316	super.onWrite(path, spark)	×
317	writeTensorflowModelV2(path, spark, getModelIfNotSet.tensorflowWrapper, "_roberta_classification", RoBertaForSequenceClassification.tfFile, configProtoBytes = getConfigProtoBytes)	×
318	}
319
320	}
321
322	trait ReadablePretrainedRoBertaForSequenceModel extends ParamsAndFeaturesReadable[RoBertaForSequenceClassification] with HasPretrained[RoBertaForSequenceClassification] {
323	override val defaultModelName: Some[String] = Some("roberta_base_sequence_classifier_imdb")	×
324
325	/** Java compliant-overrides */
326	override def pretrained(): RoBertaForSequenceClassification = super.pretrained()	×
327
328	override def pretrained(name: String): RoBertaForSequenceClassification = super.pretrained(name)	×
329
330	override def pretrained(name: String, lang: String): RoBertaForSequenceClassification = super.pretrained(name, lang)	×
331
332	override def pretrained(name: String, lang: String, remoteLoc: String): RoBertaForSequenceClassification = super.pretrained(name, lang, remoteLoc)	×
333	}
334
335	trait ReadRoBertaForSequenceTensorflowModel extends ReadTensorflowModel {
336	this: ParamsAndFeaturesReadable[RoBertaForSequenceClassification] =>
337
338	override val tfFile: String = "roberta_classification_tensorflow"	×
339
340	def readTensorflow(instance: RoBertaForSequenceClassification, path: String, spark: SparkSession): Unit = {
341
342	val tf = readTensorflowModel(path, spark, "_roberta_classification_tf", initAllTables = false)	×
343	instance.setModelIfNotSet(spark, tf)	×
344	}
345
346	addReader(readTensorflow)	×
347
348	def loadSavedModel(tfModelPath: String, spark: SparkSession): RoBertaForSequenceClassification = {
349
350	val f = new File(tfModelPath)	×
351	val savedModel = new File(tfModelPath, "saved_model.pb")	×
352
353	require(f.exists, s"Folder $tfModelPath not found")	×
354	require(f.isDirectory, s"File $tfModelPath is not folder")	×
355	require(	×
356	savedModel.exists(),	×
357	s"savedModel file saved_model.pb not found in folder $tfModelPath"	×
358	)
359
360	val vocabFile = new File(tfModelPath + "/assets", "vocab.txt")	×
361	require(f.exists, s"Folder $tfModelPath not found")	×
362	require(f.isDirectory, s"File $tfModelPath is not folder")	×
363	require(vocabFile.exists(), s"Vocabulary file vocab.txt not found in folder $tfModelPath")	×
364
365	val mergesFile = new File(tfModelPath + "/assets", "merges.txt")	×
366	require(f.exists, s"Folder $tfModelPath not found")	×
367	require(f.isDirectory, s"File $tfModelPath is not folder")	×
368	require(mergesFile.exists(), s"merges file merges.txt not found in folder $tfModelPath")	×
369
370	val vocabResource = new ExternalResource(vocabFile.getAbsolutePath, ReadAs.TEXT, Map("format" -> "text"))	×
371	val words = ResourceHelper.parseLines(vocabResource).zipWithIndex.toMap	×
372
373	val mergesResource = new ExternalResource(mergesFile.getAbsolutePath, ReadAs.TEXT, Map("format" -> "text"))	×
374	val merges = ResourceHelper.parseLines(mergesResource)	×
375
376	val bytePairs: Map[(String, String), Int] = merges.map(_.split(" "))	×
377	.filter(w => w.length == 2)	×
378	.map { case Array(c1, c2) => (c1, c2) }	×
379	.zipWithIndex.toMap	×
380
381	val labelsPath = new File(tfModelPath + "/assets", "labels.txt")	×
382	require(labelsPath.exists(), s"Labels file labels.txt not found in folder $tfModelPath/assets/")	×
383
384	val labelsResource = new ExternalResource(labelsPath.getAbsolutePath, ReadAs.TEXT, Map("format" -> "text"))	×
385	val labels = ResourceHelper.parseLines(labelsResource).zipWithIndex.toMap	×
386
387	val (wrapper, signatures) = TensorflowWrapper.read(tfModelPath, zipped = false, useBundle = true)	×
388
389	val _signatures = signatures match {
390	case Some(s) => s
391	case None => throw new Exception("Cannot load signature definitions from model!")	×
392	}
393
394	/** the order of setSignatures is important if we use getSignatures inside setModelIfNotSet */
395	new RoBertaForSequenceClassification()
396	.setVocabulary(words)
397	.setMerges(bytePairs)
398	.setLabels(labels)
399	.setSignatures(_signatures)
400	.setModelIfNotSet(spark, wrapper)	×
401	}
402	}
403
404	/**
405	* This is the companion object of [[RoBertaForSequenceClassification]]. Please refer to that class for the documentation.
406	*/
407	object RoBertaForSequenceClassification extends ReadablePretrainedRoBertaForSequenceModel with ReadRoBertaForSequenceTensorflowModel

JohnSnowLabs / spark-nlp / 1576910141

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