From 7140cdb670532db401117456ec8d386663ed5092 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=80=9D=E5=AE=8F?= <fubang.zfb@alibaba-inc.com>
Date: Tue, 14 Jun 2022 13:58:43 +0800
Subject: [PATCH] [to #42322933] init

---
 .../nlp/sentiment_classification_model.py     |  2 +-
 .../test_sentiment_classification.py          | 52 +++++++++++++++++++
 2 files changed, 53 insertions(+), 1 deletion(-)
 create mode 100644 tests/pipelines/test_sentiment_classification.py

diff --git a/modelscope/models/nlp/sentiment_classification_model.py b/modelscope/models/nlp/sentiment_classification_model.py
index fb32aff2..d0ab6698 100644
--- a/modelscope/models/nlp/sentiment_classification_model.py
+++ b/modelscope/models/nlp/sentiment_classification_model.py
@@ -55,7 +55,7 @@ class SbertForSentimentClassification(Model):
         self.model_dir = model_dir
 
         self.model = SbertTextClassifier.from_pretrained(
-            model_dir, num_labels=3)
+            model_dir, num_labels=2)
         self.model.eval()
 
     def forward(self, input: Dict[str, Any]) -> Dict[str, np.ndarray]:
diff --git a/tests/pipelines/test_sentiment_classification.py b/tests/pipelines/test_sentiment_classification.py
new file mode 100644
index 00000000..1576b335
--- /dev/null
+++ b/tests/pipelines/test_sentiment_classification.py
@@ -0,0 +1,52 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import unittest
+
+from maas_hub.snapshot_download import snapshot_download
+
+from modelscope.models import Model
+from modelscope.models.nlp import SbertForSentimentClassification
+from modelscope.pipelines import SentimentClassificationPipeline, pipeline
+from modelscope.preprocessors import SentimentClassificationPreprocessor
+from modelscope.utils.constant import Tasks
+
+
+class SentimentClassificationTest(unittest.TestCase):
+    model_id = 'damo/nlp_structbert_sentence-similarity_chinese-base'
+    sentence1 = '四川商务职业学院和四川财经职业学院哪个好？'
+
+    def test_run_from_local(self):
+        cache_path = snapshot_download(self.model_id)
+        tokenizer = SentimentClassificationPreprocessor(cache_path)
+        model = SbertForSentimentClassification(
+            cache_path, tokenizer=tokenizer)
+        pipeline1 = SentimentClassificationPipeline(
+            model, preprocessor=tokenizer)
+        pipeline2 = pipeline(
+            Tasks.sentence_similarity, model=model, preprocessor=tokenizer)
+        print(f'sentence1: {self.sentence1}\n'
+              f'pipeline1:{pipeline1(input=self.sentence1)}')
+        print()
+        print(f'sentence1: {self.sentence1}\n'
+              f'pipeline1: {pipeline2(input=self.sentence1)}')
+
+    def test_run_with_model_from_modelhub(self):
+        model = Model.from_pretrained(self.model_id)
+        tokenizer = SentimentClassificationPreprocessor(model.model_dir)
+        pipeline_ins = pipeline(
+            task=Tasks.sentence_similarity,
+            model=model,
+            preprocessor=tokenizer)
+        print(pipeline_ins(input=self.sentence1))
+
+    def test_run_with_model_name(self):
+        pipeline_ins = pipeline(
+            task=Tasks.sentence_similarity, model=self.model_id)
+        print(pipeline_ins(input=self.sentence1))
+
+    def test_run_with_default_model(self):
+        pipeline_ins = pipeline(task=Tasks.sentence_similarity)
+        print(pipeline_ins(input=self.sentence1))
+
+
+if __name__ == '__main__':
+    unittest.main()