From 48846602c117c51e66d098ddd6ae360f088df906 Mon Sep 17 00:00:00 2001
From: Matt Watson <1389937+mattdangerw@users.noreply.github.com>
Date: Wed, 3 Jul 2024 09:44:46 -0700
Subject: [PATCH] Remove Keras pin, fix tests (#1681)

* Remove Keras pin, see what breaks

* Fix tests
---
 keras_nlp/src/models/deberta_v3/deberta_v3_backbone_test.py   | 2 --
 keras_nlp/src/models/distil_bert/distil_bert_backbone_test.py | 2 --
 keras_nlp/src/models/roberta/roberta_backbone_test.py         | 2 --
 keras_nlp/src/models/xlm_roberta/xlm_roberta_backbone_test.py | 2 --
 requirements-jax-cuda.txt                                     | 1 -
 requirements-tensorflow-cuda.txt                              | 1 -
 requirements-torch-cuda.txt                                   | 1 -
 7 files changed, 11 deletions(-)

diff --git a/keras_nlp/src/models/deberta_v3/deberta_v3_backbone_test.py b/keras_nlp/src/models/deberta_v3/deberta_v3_backbone_test.py
index 0f878c749e..3001098003 100644
--- a/keras_nlp/src/models/deberta_v3/deberta_v3_backbone_test.py
+++ b/keras_nlp/src/models/deberta_v3/deberta_v3_backbone_test.py
@@ -33,7 +33,6 @@ def setUp(self):
         }
         self.input_data = {
             "token_ids": ops.ones((2, 5), dtype="int32"),
-            "segment_ids": ops.zeros((2, 5), dtype="int32"),
             "padding_mask": ops.ones((2, 5), dtype="int32"),
         }
 
@@ -60,7 +59,6 @@ def test_smallest_preset(self):
             preset="deberta_v3_extra_small_en",
             input_data={
                 "token_ids": ops.array([[0, 581, 63773, 2]], dtype="int32"),
-                "segment_ids": ops.zeros((1, 4), dtype="int32"),
                 "padding_mask": ops.ones((1, 4), dtype="int32"),
             },
             expected_output_shape=(1, 4, 384),
diff --git a/keras_nlp/src/models/distil_bert/distil_bert_backbone_test.py b/keras_nlp/src/models/distil_bert/distil_bert_backbone_test.py
index 3b282bba29..96e48a35d7 100644
--- a/keras_nlp/src/models/distil_bert/distil_bert_backbone_test.py
+++ b/keras_nlp/src/models/distil_bert/distil_bert_backbone_test.py
@@ -33,7 +33,6 @@ def setUp(self):
         }
         self.input_data = {
             "token_ids": ops.ones((2, 5), dtype="int32"),
-            "segment_ids": ops.zeros((2, 5), dtype="int32"),
             "padding_mask": ops.ones((2, 5), dtype="int32"),
         }
 
@@ -60,7 +59,6 @@ def test_smallest_preset(self):
             preset="distil_bert_base_en_uncased",
             input_data={
                 "token_ids": ops.array([[101, 1996, 4248, 102]], dtype="int32"),
-                "segment_ids": ops.zeros((1, 4), dtype="int32"),
                 "padding_mask": ops.ones((1, 4), dtype="int32"),
             },
             expected_output_shape=(1, 4, 768),
diff --git a/keras_nlp/src/models/roberta/roberta_backbone_test.py b/keras_nlp/src/models/roberta/roberta_backbone_test.py
index 0631a47560..a2108aec46 100644
--- a/keras_nlp/src/models/roberta/roberta_backbone_test.py
+++ b/keras_nlp/src/models/roberta/roberta_backbone_test.py
@@ -31,7 +31,6 @@ def setUp(self):
         }
         self.input_data = {
             "token_ids": ops.ones((2, 5), dtype="int32"),
-            "segment_ids": ops.zeros((2, 5), dtype="int32"),
             "padding_mask": ops.ones((2, 5), dtype="int32"),
         }
 
@@ -58,7 +57,6 @@ def test_smallest_preset(self):
             preset="roberta_base_en",
             input_data={
                 "token_ids": ops.array([[0, 133, 2119, 2]], dtype="int32"),
-                "segment_ids": ops.zeros((1, 4), dtype="int32"),
                 "padding_mask": ops.ones((1, 4), dtype="int32"),
             },
             expected_output_shape=(1, 4, 768),
diff --git a/keras_nlp/src/models/xlm_roberta/xlm_roberta_backbone_test.py b/keras_nlp/src/models/xlm_roberta/xlm_roberta_backbone_test.py
index 46f2a9bdeb..853ef5f4e4 100644
--- a/keras_nlp/src/models/xlm_roberta/xlm_roberta_backbone_test.py
+++ b/keras_nlp/src/models/xlm_roberta/xlm_roberta_backbone_test.py
@@ -33,7 +33,6 @@ def setUp(self):
         }
         self.input_data = {
             "token_ids": ops.ones((2, 5), dtype="int32"),
-            "segment_ids": ops.zeros((2, 5), dtype="int32"),
             "padding_mask": ops.ones((2, 5), dtype="int32"),
         }
 
@@ -60,7 +59,6 @@ def test_smallest_preset(self):
             preset="xlm_roberta_base_multi",
             input_data={
                 "token_ids": ops.array([[0, 581, 63773, 2]], dtype="int32"),
-                "segment_ids": ops.zeros((1, 4), dtype="int32"),
                 "padding_mask": ops.ones((1, 4), dtype="int32"),
             },
             expected_output_shape=(1, 4, 768),
diff --git a/requirements-jax-cuda.txt b/requirements-jax-cuda.txt
index f5ab7b7acc..2ded131217 100644
--- a/requirements-jax-cuda.txt
+++ b/requirements-jax-cuda.txt
@@ -1,4 +1,3 @@
-keras==3.3.3
 # Tensorflow cpu-only version.
 tensorflow-cpu~=2.16.1  # Pin to TF 2.16
 tensorflow-text~=2.16.1
diff --git a/requirements-tensorflow-cuda.txt b/requirements-tensorflow-cuda.txt
index 1079cde640..5426beb5a3 100644
--- a/requirements-tensorflow-cuda.txt
+++ b/requirements-tensorflow-cuda.txt
@@ -1,4 +1,3 @@
-keras==3.3.3
 # Tensorflow with cuda support.
 tensorflow[and-cuda]~=2.16.1  # Pin to TF 2.16
 tensorflow-text~=2.16.1
diff --git a/requirements-torch-cuda.txt b/requirements-torch-cuda.txt
index dee39b8f7d..c1ea85b184 100644
--- a/requirements-torch-cuda.txt
+++ b/requirements-torch-cuda.txt
@@ -1,4 +1,3 @@
-keras==3.3.3
 # Tensorflow cpu-only version.
 tensorflow-cpu~=2.16.1  # Pin to TF 2.16
 tensorflow-text~=2.16.1