Fix TextVectorization tf-idf mode deserialization failure (#22330)

* fix(saving): fix TextVectorization tf-idf mode deserialization

* test(saving): skip tf-idf save/load test on non-tensorflow backends

Author: rstar327

keras/src/layers/preprocessing/index_lookup.py

@@ -861,7 +861,11 @@ def save_own_variables(self, store):
 
     def load_own_variables(self, store):
         if self.output_mode == "tf_idf":
-            self.idf_weights.assign(store["idf_weights"])
+            idf_weights = store["idf_weights"]
+            if hasattr(self, "idf_weights"):
+                self.idf_weights.assign(idf_weights)
+            else:
+                self.idf_weights = tf.Variable(idf_weights, trainable=False)
             self.idf_weights_const = self.idf_weights.value()
 
     def save_assets(self, dir_path):
@@ -889,7 +893,8 @@ def load_assets(self, dir_path):
             else:
                 values = [int(line) for line in lines]
             if self.output_mode == "tf_idf":
-                self.set_vocabulary(values, idf_weights=False)
+                idf_weights = self.idf_weights_const.numpy()
+                self.set_vocabulary(values, idf_weights=idf_weights)
             else:
                 self.set_vocabulary(values)
 

keras/src/layers/preprocessing/text_vectorization_test.py

@@ -85,6 +85,24 @@ def test_save_load_with_ngrams_flow(self):
         model = saving.load_model(temp_filepath)
         self.assertAllClose(output, model(input_data))
 
+    @pytest.mark.skipif(
+        backend.backend() != "tensorflow", reason="Requires string input dtype"
+    )
+    def test_save_load_tf_idf_mode(self):
+        input_data = np.array(["foo bar", "bar baz", "baz bada boom"])
+        model = Sequential(
+            [
+                layers.Input(dtype="string", shape=()),
+                layers.TextVectorization(max_tokens=100, output_mode="tf_idf"),
+            ]
+        )
+        model.layers[0].adapt(input_data)
+        output = model(input_data)
+        temp_filepath = os.path.join(self.get_temp_dir(), "model.keras")
+        model.save(temp_filepath)
+        loaded_model = saving.load_model(temp_filepath)
+        self.assertAllClose(output, loaded_model(input_data))
+
     def test_tf_data_compatibility(self):
         max_tokens = 5000
         max_len = 4