data_juicer.ops.mapper.query_topic_detection_mapper module

class data_juicer.ops.mapper.query_topic_detection_mapper.QueryTopicDetectionMapper(hf_model: str = 'dstefa/roberta-base_topic_classification_nyt_news', zh_to_en_hf_model: str | None = 'Helsinki-NLP/opus-mt-zh-en', model_params: Dict = {}, zh_to_en_model_params: Dict = {}, *, label_key: str = 'query_topic_label', score_key: str = 'query_topic_label_score', **kwargs)

Bases: Mapper

Mapper to predict user’s topic label in query. Input from query_key. Output topic label and corresponding score for the query, which is store in ‘query_topic_label’ and ‘query_topic_label_score’ in Data-Juicer meta field.

__init__(hf_model: str = 'dstefa/roberta-base_topic_classification_nyt_news', zh_to_en_hf_model: str | None = 'Helsinki-NLP/opus-mt-zh-en', model_params: Dict = {}, zh_to_en_model_params: Dict = {}, *, label_key: str = 'query_topic_label', score_key: str = 'query_topic_label_score', **kwargs)

Initialization method.

Parameters:

• hf_model – Huggingface model ID to predict topic label.

• zh_to_en_hf_model – Translation model from Chinese to English. If not None, translate the query from Chinese to English.

• model_params – model param for hf_model.

• zh_to_en_model_params – model param for zh_to_hf_model.

• label_key – The key name in the meta field to store the output label. It is ‘query_topic_label’ in default.

• score_key – The key name in the meta field to store the corresponding label score. It is ‘query_topic_label_score’ in default.

• kwargs – Extra keyword arguments.

process_batched(samples, rank=None)