add type hint in pipeline model argument (huggingface#23740)

* add type hint in pipeline model argument * add pretrainedmodel and tfpretainedmodel type hint * make type hints string
novice03 · Jun 23, 2023 · 579771e · 579771e
1 parent e38055f
commit 579771e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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,