Browse Source

add default args

master
雨泓 3 years ago
parent
commit
96e25be7d2
2 changed files with 2 additions and 2 deletions
  1. +1
    -1
      modelscope/models/nlp/sbert_for_sentence_similarity.py
  2. +1
    -1
      modelscope/models/nlp/sbert_for_sentiment_classification.py

+ 1
- 1
modelscope/models/nlp/sbert_for_sentence_similarity.py View File

@@ -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

+ 1
- 1
modelscope/models/nlp/sbert_for_sentiment_classification.py View File

@@ -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

Loading…
Cancel
Save