From 6e016634f1f58af66d6af27e272b5d9ad406e0f4 Mon Sep 17 00:00:00 2001 From: Joao Gante Date: Thu, 1 Sep 2022 10:53:39 +0100 Subject: [PATCH] Generate: smaller TF serving test (#18840) --- tests/test_modeling_tf_common.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_modeling_tf_common.py b/tests/test_modeling_tf_common.py index 3372e96454..a0413afdfa 100644 --- a/tests/test_modeling_tf_common.py +++ b/tests/test_modeling_tf_common.py @@ -75,7 +75,7 @@ if is_tf_available(): TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING, BertConfig, TFAutoModel, - TFAutoModelForSeq2SeqLM, + TFAutoModelForCausalLM, TFAutoModelForSequenceClassification, TFBertModel, TFSharedEmbeddings, @@ -2180,8 +2180,8 @@ class UtilsFunctionsTest(unittest.TestCase): self.assertTrue(np.allclose(p1.numpy(), p2.numpy())) def test_generate_tf_function_export(self): - test_model = TFAutoModelForSeq2SeqLM.from_pretrained("hf-internal-testing/tiny-random-t5") - max_length = 8 + test_model = TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2") + max_length = 2 class DummyModel(tf.Module): def __init__(self, model): @@ -2204,8 +2204,8 @@ class UtilsFunctionsTest(unittest.TestCase): ) return {"sequences": outputs["sequences"]} - dummy_input_ids = [[2, 3, 4, 1, 0, 0, 0, 0], [102, 103, 104, 105, 1, 0, 0, 0]] - dummy_attention_masks = [[1, 1, 1, 1, 0, 0, 0, 0], [1, 1, 1, 1, 1, 0, 0, 0]] + dummy_input_ids = [[2, 0], [102, 103]] + dummy_attention_masks = [[1, 0], [1, 1]] dummy_model = DummyModel(model=test_model) with tempfile.TemporaryDirectory() as tmp_dir: tf.saved_model.save(dummy_model, tmp_dir, signatures={"serving_default": dummy_model.serving})