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
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 12 additions and 0 deletions

View File

@ -21,6 +21,7 @@ import os
# See the License for the specific language governing permissions and
# limitations under the License.
import warnings
from pathlib import Path
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
from numpy import isin
@ -638,6 +639,8 @@ def pipeline(
" feature_extractor may not be compatible with the default model. Please provide a PreTrainedModel class"
" or a path/identifier to a pretrained model when providing feature_extractor."
)
if isinstance(model, Path):
model = str(model)
# Config is the primordial information item.
# Instantiate config if needed

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):