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

Author: RyanMullins

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