modelscope/tests/utils/test_sentence_embedding_utils.py

import unittest

from modelscope.utils.test_utils import test_level

from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

class LLMPipelineTest(unittest.TestCase):
    def setUp(self) -> None:
        self.model_id = 'Qwen/Qwen3-Embedding-0.6B'
        self.querys = [
            "What is the capital of China?",
            "Explain gravity",
        ]
        self.documents = [
            "The capital of China is Beijing.",
            "Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun.",
        ]

    def test_ori_pipeline(self):
        ppl = pipeline(
            Tasks.sentence_embedding,
            model=self.model_id,
            model_revision='master',
        )
        inputs = {"source_sentence": self.documents}
        embeddings = ppl(input=inputs)["text_embedding"]
        self.assertEqual(embeddings.shape[0], len(self.documents))
        assert((embeddings[0][0]+0.0471825)<0.01) # check value

    def test_sentence_embedding_input(self):
        ppl = pipeline(
            Tasks.sentence_embedding,
            model=self.model_id,
            model_revision='master',
        )
        embeddings = ppl(self.documents, prompt_name="query")
        self.assertEqual(embeddings.shape[0], len(self.documents))
        assert ((embeddings[0][0] + 0.0471825) < 0.01)  # check value


if __name__ == '__main__':
    unittest.main()