From af45ec0a1611062929ddbf6f15935e01e6cbf1af Mon Sep 17 00:00:00 2001 From: Samin Yasar Date: Tue, 30 May 2023 16:05:58 +0600 Subject: [PATCH] add type hint in pipeline model argument (#23740) * add type hint in pipeline model argument * add pretrainedmodel and tfpretainedmodel type hint * make type hints string --- src/transformers/pipelines/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/pipelines/__init__.py b/src/transformers/pipelines/__init__.py index 84d461cd1a..818164f3c2 100755 --- a/src/transformers/pipelines/__init__.py +++ b/src/transformers/pipelines/__init__.py @@ -505,7 +505,7 @@ def clean_custom_task(task_info): def pipeline( task: str = None, - model: Optional = None, + model: Optional[Union[str, "PreTrainedModel", "TFPreTrainedModel"]] = None, config: Optional[Union[str, PretrainedConfig]] = None, tokenizer: Optional[Union[str, PreTrainedTokenizer, "PreTrainedTokenizerFast"]] = None, feature_extractor: Optional[Union[str, PreTrainedFeatureExtractor]] = None,