Pipelines: fix crash when modelcard is None

cc @mfuntowicz does this seem correct?

Pipelines: fix crash when modelcard is None
cc @mfuntowicz does this seem correct?
9a50828b · Julien Chaumond · 6c1b2355 · 9a50828b
Commit 9a50828b authored Feb 03, 2020 by Julien Chaumond
Show whitespace changes
Inline Side-by-side

Showing with 7 additions and 6 deletions

src/transformers/pipelines.py src/transformers/pipelines.py +7 -6

No files found.
--- a/src/transformers/pipelines.py
+++ b/src/transformers/pipelines.py
@@ -326,7 +326,7 @@ class Pipeline(_ScikitCompat):
        self,
        model,
        tokenizer: PreTrainedTokenizer = None,
-        modelcard: ModelCard = None,
+        modelcard: Optional[ModelCard] = None,
        framework: Optional[str] = None,
        args_parser: ArgumentHandler = None,
        device: int = -1,
@@ -358,6 +358,7 @@ class Pipeline(_ScikitCompat):

        self.model.save_pretrained(save_directory)
        self.tokenizer.save_pretrained(save_directory)
+        if self.modelcard is not None:
            self.modelcard.save_pretrained(save_directory)

    def transform(self, X):
@@ -476,7 +477,7 @@ class FeatureExtractionPipeline(Pipeline):
        self,
        model,
        tokenizer: PreTrainedTokenizer = None,
-        modelcard: ModelCard = None,
+        modelcard: Optional[ModelCard] = None,
        framework: Optional[str] = None,
        args_parser: ArgumentHandler = None,
        device: int = -1,
@@ -515,7 +516,7 @@ class FillMaskPipeline(Pipeline):
        self,
        model,
        tokenizer: PreTrainedTokenizer = None,
-        modelcard: ModelCard = None,
+        modelcard: Optional[ModelCard] = None,
        framework: Optional[str] = None,
        args_parser: ArgumentHandler = None,
        device: int = -1,
@@ -582,7 +583,7 @@ class NerPipeline(Pipeline):
        self,
        model,
        tokenizer: PreTrainedTokenizer = None,
-        modelcard: ModelCard = None,
+        modelcard: Optional[ModelCard] = None,
        framework: Optional[str] = None,
        args_parser: ArgumentHandler = None,
        device: int = -1,
@@ -721,7 +722,7 @@ class QuestionAnsweringPipeline(Pipeline):
        self,
        model,
        tokenizer: Optional[PreTrainedTokenizer],
-        modelcard: Optional[ModelCard],
+        modelcard: Optional[ModelCard] = None,
        framework: Optional[str] = None,
        device: int = -1,
        **kwargs