From 96e25be7d250fcfdb3e748efb6b24c9666aa570b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=A8=E6=B3=93?= Date: Wed, 22 Jun 2022 21:54:41 +0800 Subject: [PATCH] add default args --- modelscope/models/nlp/sbert_for_sentence_similarity.py | 2 +- modelscope/models/nlp/sbert_for_sentiment_classification.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/modelscope/models/nlp/sbert_for_sentence_similarity.py b/modelscope/models/nlp/sbert_for_sentence_similarity.py index db469f4f..e893a301 100644 --- a/modelscope/models/nlp/sbert_for_sentence_similarity.py +++ b/modelscope/models/nlp/sbert_for_sentence_similarity.py @@ -18,6 +18,6 @@ class SbertForSentenceSimilarity(SbertForSequenceClassificationBase): model_cls (Optional[Any], optional): model loader, if None, use the default loader to load model weights, by default None. """ - super().__init__(model_dir, *args, **kwargs) + super().__init__(model_dir, *args, model_args={"num_labels": 2}, **kwargs) self.model_dir = model_dir assert self.model.config.num_labels == 2 diff --git a/modelscope/models/nlp/sbert_for_sentiment_classification.py b/modelscope/models/nlp/sbert_for_sentiment_classification.py index 4683b079..10bfaa0f 100644 --- a/modelscope/models/nlp/sbert_for_sentiment_classification.py +++ b/modelscope/models/nlp/sbert_for_sentiment_classification.py @@ -19,5 +19,5 @@ class SbertForSentimentClassification(SbertForSequenceClassificationBase): model_cls (Optional[Any], optional): model loader, if None, use the default loader to load model weights, by default None. """ - super().__init__(model_dir, *args, **kwargs) + super().__init__(model_dir, *args, model_args={"num_labels": 2}, **kwargs) assert self.model.config.num_labels == 2