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
sheonhan · Jun 1, 2023 · ba7afe8 · ba7afe8
1 parent 0e1d17b
commit ba7afe8
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,