diff --git a/src/transformers/hf_api.py b/src/transformers/hf_api.py index f9fd5450340..1bacd68a8cf 100644 --- a/src/transformers/hf_api.py +++ b/src/transformers/hf_api.py @@ -80,6 +80,7 @@ class ModelInfo: author: Optional[str] = None, downloads: Optional[int] = None, tags: List[str] = [], + pipeline_tag: Optional[str] = None, siblings: Optional[List[Dict]] = None, # list of files that constitute the model **kwargs ): @@ -88,6 +89,7 @@ class ModelInfo: self.author = author self.downloads = downloads self.tags = tags + self.pipeline_tag = pipeline_tag self.siblings = [S3Object(**x) for x in siblings] if siblings is not None else None for k, v in kwargs.items(): setattr(self, k, v)