fix file swap

meedan · Aug 23, 2024 · 1fab4dd · 1fab4dd
1 parent 7a1f37e
commit 1fab4dd
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/test/lib/model/test_indian_sbert.py b/test/lib/model/test_indian_sbert.py
@@ -7,13 +7,13 @@
 from lib.model.generic_transformer import GenericTransformerModel
 from lib import schemas
 
-class TestParaphraseMultilingual(unittest.TestCase):
+class TestIndianSbert(unittest.TestCase):
     def setUp(self):
         self.model = GenericTransformerModel(None)
         self.mock_model = MagicMock()
 
     def test_vectorize(self):
-        texts = [schemas.parse_message({"body": {"id": "123", "callback_url": "http://example.com/callback", "text": "Hello, how are you?"}, "model_name": "paraphrase_multilingual__Model"}), schemas.parse_message({"body": {"id": "123", "callback_url": "http://example.com/callback", "text": "I'm doing great, thanks!"}, "model_name": "paraphrase_multilingual__Model"})]
+        texts = [schemas.parse_message({"body": {"id": "123", "callback_url": "http://example.com/callback", "text": "Hello, how are you?"}, "model_name": "indian_sbert__Model"}), schemas.parse_message({"body": {"id": "123", "callback_url": "http://example.com/callback", "text": "I'm doing great, thanks!"}, "model_name": "indian_sbert__Model"})]
         self.model.model = self.mock_model
         self.model.model.encode = MagicMock(return_value=np.array([[4, 5, 6], [7, 8, 9]]))
         vectors = self.model.vectorize(texts)
@@ -22,11 +22,11 @@ def test_vectorize(self):
         self.assertEqual(vectors[1], [7, 8, 9])
 
     def test_respond(self):
-        query = schemas.parse_message({"body": {"id": "123", "callback_url": "http://example.com/callback", "text": "What is the capital of India?"}, "model_name": "paraphrase_multilingual__Model"})
+        query = schemas.parse_message({"body": {"id": "123", "callback_url": "http://example.com/callback", "text": "What is the capital of India?"}, "model_name": "indian_sbert__Model"})
         self.model.vectorize = MagicMock(return_value=[[1, 2, 3]])
         response = self.model.respond(query)
         self.assertEqual(len(response), 1)
         self.assertEqual(response[0].body.result, [1, 2, 3])
 
 if __name__ == '__main__':
-    unittest.main()
+    unittest.main()
diff --git a/test/lib/model/test_paraphrase_multilingual.py b/test/lib/model/test_paraphrase_multilingual.py
@@ -7,13 +7,13 @@
 from lib.model.generic_transformer import GenericTransformerModel
 from lib import schemas
 
-class TestIndianSbert(unittest.TestCase):
+class TestParaphraseMultilingual(unittest.TestCase):
     def setUp(self):
         self.model = GenericTransformerModel(None)
         self.mock_model = MagicMock()
 
     def test_vectorize(self):
-        texts = [schemas.parse_message({"body": {"id": "123", "callback_url": "http://example.com/callback", "text": "Hello, how are you?"}, "model_name": "indian_sbert__Model"}), schemas.parse_message({"body": {"id": "123", "callback_url": "http://example.com/callback", "text": "I'm doing great, thanks!"}, "model_name": "indian_sbert__Model"})]
+        texts = [schemas.parse_message({"body": {"id": "123", "callback_url": "http://example.com/callback", "text": "Hello, how are you?"}, "model_name": "paraphrase_multilingual__Model"}), schemas.parse_message({"body": {"id": "123", "callback_url": "http://example.com/callback", "text": "I'm doing great, thanks!"}, "model_name": "paraphrase_multilingual__Model"})]
         self.model.model = self.mock_model
         self.model.model.encode = MagicMock(return_value=np.array([[4, 5, 6], [7, 8, 9]]))
         vectors = self.model.vectorize(texts)
@@ -22,7 +22,7 @@ def test_vectorize(self):
         self.assertEqual(vectors[1], [7, 8, 9])
 
     def test_respond(self):
-        query = schemas.parse_message({"body": {"id": "123", "callback_url": "http://example.com/callback", "text": "What is the capital of India?"}, "model_name": "indian_sbert__Model"})
+        query = schemas.parse_message({"body": {"id": "123", "callback_url": "http://example.com/callback", "text": "What is the capital of India?"}, "model_name": "paraphrase_multilingual__Model"})
         self.model.vectorize = MagicMock(return_value=[[1, 2, 3]])
         response = self.model.respond(query)
         self.assertEqual(len(response), 1)