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
2023-05-30 16:05:58 +06:00
parent 4b6a5a7caa
commit af45ec0a16
1 changed files with 1 additions and 1 deletions
--- 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,