for test

modelscope · Dec 17, 2024 · 9b6652d · 9b6652d
1 parent 6f44ec0
commit 9b6652d
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 5 deletions.
diff --git a/data_juicer/ops/mapper/query_intent_detection_mapper.py b/data_juicer/ops/mapper/query_intent_detection_mapper.py
@@ -76,7 +76,6 @@ def process_batched(self, samples, rank=None):
                                       self.use_cuda())
             results = translater(queries)
             queries = [item['translation_text'] for item in results]
-        print(queries)
 
         classifier, _ = get_model(self.model_key, rank, self.use_cuda())
         results = classifier(queries)

diff --git a/tests/ops/mapper/test_query_intent_detection_mapper.py b/tests/ops/mapper/test_query_intent_detection_mapper.py
@@ -21,8 +21,7 @@ def _run_op(self, op, samples, label_key, targets):
 
         for sample, target in zip(dataset, targets):
             label = nested_access(sample[Fields.meta], label_key)
-            print(label)
-            # self.assertEqual(label, target)
+            self.assertEqual(label, target)
 
     def test_default(self):
 
@@ -34,7 +33,7 @@ def test_default(self):
             'query': '今天阳光灿烂。'
         }
         ]
-        targets = [1, 0, -1]
+        targets = ['question', 'command', 'statement']
 
         op = QueryIntentDetectionMapper(
             hf_model = self.hf_model,
@@ -50,7 +49,7 @@ def test_no_zh_to_en(self):
             'query': 'Is this okay?'
         }
         ]
-        targets = [0, 1]
+        targets = ['question', 'rhetorical question']
 
         op = QueryIntentDetectionMapper(
             hf_model = self.hf_model,