Fix RequestCounter to make it more future-proof (#27406)
* Fix RequestCounter to make it more future-proof * code quality
This commit is contained in:
@@ -301,14 +301,14 @@ class TFAutoModelTest(unittest.TestCase):
|
||||
_ = TFAutoModel.from_pretrained("hf-internal-testing/tiny-random-bert")
|
||||
with RequestCounter() as counter:
|
||||
_ = TFAutoModel.from_pretrained("hf-internal-testing/tiny-random-bert")
|
||||
self.assertEqual(counter.get_request_count, 0)
|
||||
self.assertEqual(counter.head_request_count, 1)
|
||||
self.assertEqual(counter.other_request_count, 0)
|
||||
self.assertEqual(counter["GET"], 0)
|
||||
self.assertEqual(counter["HEAD"], 1)
|
||||
self.assertEqual(counter.total_calls, 1)
|
||||
|
||||
# With a sharded checkpoint
|
||||
_ = TFAutoModel.from_pretrained("ArthurZ/tiny-random-bert-sharded")
|
||||
with RequestCounter() as counter:
|
||||
_ = TFAutoModel.from_pretrained("ArthurZ/tiny-random-bert-sharded")
|
||||
self.assertEqual(counter.get_request_count, 0)
|
||||
self.assertEqual(counter.head_request_count, 1)
|
||||
self.assertEqual(counter.other_request_count, 0)
|
||||
self.assertEqual(counter["GET"], 0)
|
||||
self.assertEqual(counter["HEAD"], 1)
|
||||
self.assertEqual(counter.total_calls, 1)
|
||||
|
||||
Reference in New Issue
Block a user