SPARKNLP-867 Refactors getActivation to take multilabel into account (#…

…13888)
JohnSnowLabs · Jul 18, 2023 · 31e5bde · 31e5bde
1 parent c07a4a7
commit 31e5bde
Show file tree

Hide file tree

Showing 5 changed files with 13 additions and 7 deletions.
diff --git a/python/sparknlp/annotator/classifier_dl/bert_for_zero_shot_classification.py b/python/sparknlp/annotator/classifier_dl/bert_for_zero_shot_classification.py
@@ -164,7 +164,8 @@ def __init__(self, classname="com.johnsnowlabs.nlp.annotators.classifier.dl.Bert
             maxSentenceLength=128,
             caseSensitive=True,
             coalesceSentences=False,
-            activation="softmax"
+            activation="softmax",
+            multilabel=False
         )
 
     @staticmethod

diff --git a/src/main/scala/com/johnsnowlabs/nlp/HasClassifierActivationProperties.scala b/src/main/scala/com/johnsnowlabs/nlp/HasClassifierActivationProperties.scala
@@ -52,10 +52,13 @@ trait HasClassifierActivationProperties extends ParamsAndFeaturesWritable {
     "multilabel",
     "Whether or not the result should be multi-class (the sum of all probabilities is 1.0) or multi-label (each label has a probability between 0.0 to 1.0). Default is False i.e. multi-class")
 
-  setDefault(activation -> ActivationFunction.softmax, threshold -> 0.5f, multilabel -> false)
-
   /** @group getParam */
-  def getActivation: String = $(activation)
+  def getActivation: String = {
+    val activation =
+      if ($(multilabel)) ActivationFunction.sigmoid else ActivationFunction.softmax
+
+    if ($(multilabel)) activation else $(this.activation)
+  }
 
   /** @group setParam */
   def setActivation(value: String): this.type = {
@@ -94,6 +97,8 @@ trait HasClassifierActivationProperties extends ParamsAndFeaturesWritable {
     set(this.multilabel, value)
   }
 
+  setDefault(activation -> ActivationFunction.softmax, threshold -> 0.5f, multilabel -> false)
+
 }
 
 object ActivationFunction {

diff --git a/...n/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/BertForZeroShotClassification.scala b/...n/scala/com/johnsnowlabs/nlp/annotators/classifier/dl/BertForZeroShotClassification.scala
@@ -325,7 +325,7 @@ class BertForZeroShotClassification(override val uid: String)
           $(caseSensitive),
           $(coalesceSentences),
           $$(labels),
-          $(activation))
+          getActivation)
 
       } else {
         Seq.empty[Annotation]

diff --git a/...a/com/johnsnowlabs/nlp/annotators/classifier/dl/DistilBertForZeroShotClassification.scala b/...a/com/johnsnowlabs/nlp/annotators/classifier/dl/DistilBertForZeroShotClassification.scala
@@ -325,7 +325,7 @@ class DistilBertForZeroShotClassification(override val uid: String)
           $(caseSensitive),
           $(coalesceSentences),
           $$(labels),
-          $(activation))
+          getActivation)
 
       } else {
         Seq.empty[Annotation]

diff --git a/...cala/com/johnsnowlabs/nlp/annotators/classifier/dl/RoBertaForZeroShotClassification.scala b/...cala/com/johnsnowlabs/nlp/annotators/classifier/dl/RoBertaForZeroShotClassification.scala
@@ -339,7 +339,7 @@ class RoBertaForZeroShotClassification(override val uid: String)
           $(caseSensitive),
           $(coalesceSentences),
           $$(labels),
-          $(activation))
+          getActivation)
 
       } else {
         Seq.empty[Annotation]