Gemma: Add logit soft-capping to score function. (#1712)

keras-team · Jul 26, 2024 · fa0fbb7 · fa0fbb7
1 parent 8824388
commit fa0fbb7
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/keras_nlp/src/models/gemma/gemma_causal_lm.py b/keras_nlp/src/models/gemma/gemma_causal_lm.py
@@ -445,6 +445,12 @@ def default_layer_intercept_fn(x, unused_i):
         x = self.backbone.layer_norm(x)
         logits = self.backbone.token_embedding(x, reverse=True)
 
+        if self.backbone.final_logit_soft_cap is not None:
+            logits = ops.divide(logits, self.backbone.final_logit_soft_cap)
+            logits = ops.multiply(
+                ops.tanh(logits), self.backbone.final_logit_soft_cap
+            )
+
         if scoring_mode == "logits":
             return logits