Add from_config to sampler

Every class with a get_config should have a from_config somewhere in
the super() chain. This is actually causing test failures on nightly.

Also fix some unit testing to not rely on the specific dictionary
structure of a serialized object, which is changing across keras
versions.

Author: mattdangerw

keras_nlp/samplers/sampler.py

@@ -377,6 +377,10 @@ def one_step(
         )
         return prompt
 
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
+
     def get_config(self):
         return {
             "jit_compile": self.jit_compile,

keras_nlp/samplers/sampler_test.py

@@ -17,21 +17,18 @@
 
 import keras_nlp
 from keras_nlp.samplers.greedy_sampler import GreedySampler
+from keras_nlp.samplers.top_k_sampler import TopKSampler
 
 
 class SamplerTest(tf.test.TestCase):
     def test_serialization(self):
-        sampler = keras_nlp.samplers.GreedySampler()
-        config = keras_nlp.samplers.serialize(sampler)
-        expected_config = {
-            "class_name": "keras_nlp>GreedySampler",
-            "config": {
-                "jit_compile": True,
-            },
-        }
-        self.assertDictEqual(expected_config, config)
+        sampler = TopKSampler(k=5)
+        restored = keras_nlp.samplers.deserialize(
+            keras_nlp.samplers.serialize(sampler)
+        )
+        self.assertDictEqual(sampler.get_config(), restored.get_config())
 
-    def test_deserialization(self):
+    def test_get(self):
         # Test get from string.
         identifier = "greedy"
         sampler = keras_nlp.samplers.get(identifier)