Now supporting pathlike in pipelines too. (#20030)

This commit is contained in:
Nicolas Patry
2022-11-03 09:14:45 +01:00
committed by GitHub
parent aa39967b28
commit ec6878f6ca
2 changed files with 12 additions and 0 deletions

View File

@@ -356,6 +356,15 @@ class CommonPipelineTest(unittest.TestCase):
self.assertEqual(pipe._batch_size, 2)
self.assertEqual(pipe._num_workers, 1)
@require_torch
def test_pipeline_pathlike(self):
pipe = pipeline(model="hf-internal-testing/tiny-random-distilbert")
with tempfile.TemporaryDirectory() as d:
pipe.save_pretrained(d)
path = Path(d)
newpipe = pipeline(task="text-classification", model=path)
self.assertIsInstance(newpipe, TextClassificationPipeline)
@require_torch
def test_pipeline_override(self):
class MyPipeline(TextClassificationPipeline):