saved with correct format

Author: philschmid

tests/data/huggingface/run_tf.py

@@ -57,10 +57,8 @@
     )
     train_dataset.set_format(type="tensorflow", columns=["input_ids", "attention_mask", "label"])
 
-    train_features = {
-        x: train_dataset[x].to_tensor(default_value=0, shape=[None, tokenizer.model_max_length])
-        for x in ["input_ids", "attention_mask"]
-    }
+    train_features = {x: train_dataset[x] for x in ["input_ids", "attention_mask"]}
+
     tf_train_dataset = tf.data.Dataset.from_tensor_slices(
         (train_features, train_dataset["label"])
     ).batch(args.per_device_train_batch_size)
@@ -71,10 +69,8 @@
     )
     test_dataset.set_format(type="tensorflow", columns=["input_ids", "attention_mask", "label"])
 
-    test_features = {
-        x: test_dataset[x].to_tensor(default_value=0, shape=[None, tokenizer.model_max_length])
-        for x in ["input_ids", "attention_mask"]
-    }
+    test_features = {x: test_dataset[x] for x in ["input_ids", "attention_mask"]}
+
     tf_test_dataset = tf.data.Dataset.from_tensor_slices(
         (test_features, test_dataset["label"])
     ).batch(args.per_device_eval_batch_size)

tests/integ/test_huggingface.py

@@ -158,7 +158,7 @@ def test_huggingface_inference(
     huggingface_pytorch_latest_inference_py_version,
 ):
     env = {
-        "HF_MODEL_ID": "sshleifer/tiny-distilbert-base-uncased-finetuned-sst-2-english",
+        "HF_MODEL_ID": "philschmid/tiny-distilbert-classification",
         "HF_TASK": "text-classification",
     }
     endpoint_name = unique_name_from_base("test-hf-inference")