SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Return input_ids in ImageGPT feature extractor (#16872)

Browse Source
This commit is contained in:
Sylvain Gugger
2022-04-21 09:09:00 -04:00
committed by GitHub
parent e789418ebe
commit cb555af2c7
2 changed files with 9 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
tests/imagegpt/test_feature_extraction_imagegpt.py
View File

@@ -161,17 +161,17 @@ class ImageGPTFeatureExtractorIntegrationTest(unittest.TestCase):
# test non-batched
encoding = feature_extractor(images[0], return_tensors="pt")
self.assertIsInstance(encoding.pixel_values, torch.LongTensor)
self.assertEqual(encoding.pixel_values.shape, (1, 1024))
self.assertIsInstance(encoding.input_ids, torch.LongTensor)
self.assertEqual(encoding.input_ids.shape, (1, 1024))
expected_slice = [306, 191, 191]
self.assertEqual(encoding.pixel_values[0, :3].tolist(), expected_slice)
self.assertEqual(encoding.input_ids[0, :3].tolist(), expected_slice)
# test batched
encoding = feature_extractor(images, return_tensors="pt")
self.assertIsInstance(encoding.pixel_values, torch.LongTensor)
self.assertEqual(encoding.pixel_values.shape, (2, 1024))
self.assertIsInstance(encoding.input_ids, torch.LongTensor)
self.assertEqual(encoding.input_ids.shape, (2, 1024))
expected_slice = [303, 13, 13]
self.assertEqual(encoding.pixel_values[1, -3:].tolist(), expected_slice)
self.assertEqual(encoding.input_ids[1, -3:].tolist(), expected_slice)