diff --git a/data_juicer/ops/mapper/query_intent_detection_mapper.py b/data_juicer/ops/mapper/query_intent_detection_mapper.py index 66290532..badc7f49 100644 --- a/data_juicer/ops/mapper/query_intent_detection_mapper.py +++ b/data_juicer/ops/mapper/query_intent_detection_mapper.py @@ -25,7 +25,8 @@ class QueryIntentDetectionMapper(Mapper): def __init__( self, - hf_model: str = 'Falconsai/intent_classification', + hf_model: + str = 'bespin-global/klue-roberta-small-3i4k-intent-classification', # noqa: E501 E131 zh_to_en_hf_model: Optional[str] = 'Helsinki-NLP/opus-mt-zh-en', model_params: Dict = {}, zh_to_en_model_params: Dict = {}, diff --git a/tests/ops/mapper/test_query_intent_detection_mapper.py b/tests/ops/mapper/test_query_intent_detection_mapper.py index 5fac5ffc..592d3d0c 100644 --- a/tests/ops/mapper/test_query_intent_detection_mapper.py +++ b/tests/ops/mapper/test_query_intent_detection_mapper.py @@ -12,7 +12,7 @@ class TestQueryIntentDetectionMapper(DataJuicerTestCaseBase): - hf_model = '/mnt/workspace/shared/checkpoints/huggingface/Falconsai/intent_classification' + hf_model = '/mnt/workspace/shared/checkpoints/huggingface/bespin-global/klue-roberta-small-3i4k-intent-classification' zh_to_en_hf_model = '/mnt/workspace/shared/checkpoints/huggingface/Helsinki-NLP/opus-mt-zh-en' def _run_op(self, op, samples, intensity_key, targets): @@ -27,11 +27,11 @@ def _run_op(self, op, samples, intensity_key, targets): def test_default(self): samples = [{ - 'query': '我要一个汉堡。' + 'query': '这样好吗?' },{ - 'query': '你最近过得怎么样?' + 'query': '把那只笔递给我。' },{ - 'query': '它是正方形的。' + 'query': '难道不是这样的吗?' } ] targets = [1, 0, -1] @@ -45,9 +45,9 @@ def test_default(self): def test_no_zh_to_en(self): samples = [{ - 'query': '它是正方形的。' + 'query': '这样好吗?' },{ - 'query': 'It is square.' + 'query': 'Is this okay?' } ] targets = [0, 1] diff --git a/tests/ops/mapper/test_query_sentiment_detection_mapper.py b/tests/ops/mapper/test_query_sentiment_detection_mapper.py index dcc29e25..62ed0f38 100644 --- a/tests/ops/mapper/test_query_sentiment_detection_mapper.py +++ b/tests/ops/mapper/test_query_sentiment_detection_mapper.py @@ -12,8 +12,8 @@ class TestQuerySentimentDetectionMapper(DataJuicerTestCaseBase): - hf_model = '/mnt/workspace/shared/checkpoints/huggingface/mrm8488/distilroberta-finetuned-financial-news-sentiment-analysis' - zh_to_en_hf_model = '/mnt/workspace/shared/checkpoints/huggingface/Helsinki-NLP/opus-mt-zh-en' + hf_model = 'mrm8488/distilroberta-finetuned-financial-news-sentiment-analysis' + zh_to_en_hf_model = 'Helsinki-NLP/opus-mt-zh-en' def _run_op(self, op, samples, label_key, targets): dataset = Dataset.from_list(samples)