output double logits and softmax

modeling5.py

@@ -87,9 +87,10 @@ class MiniSunModel(tf.keras.Model):
         # LM Head for token generation
         logits = self.lm_head(hidden_states)
 
+        # Softmax layer for confidence
         softmax_output = tf.nn.softmax(logits, axis=-1)
 
-        return softmax_output
+        return logits, softmax_output
 
     def get_config(self):
         # Return the configuration of the model