diff --git a/examples/configs/instruct/polygraph_eval_coqa_default_instruct.yaml b/examples/configs/instruct/polygraph_eval_coqa_default_instruct.yaml index 188022b86..644665200 100644 --- a/examples/configs/instruct/polygraph_eval_coqa_default_instruct.yaml +++ b/examples/configs/instruct/polygraph_eval_coqa_default_instruct.yaml @@ -4,6 +4,7 @@ hydra: defaults: - model: mistral-7b-instruct-v0.2 + - similarity_model: default_en - _self_ cache_path: ./workdir/output diff --git a/examples/configs/instruct/polygraph_eval_mmlu_default_instruct.yaml b/examples/configs/instruct/polygraph_eval_mmlu_default_instruct.yaml index 3b5ed57fe..3a79d6a39 100644 --- a/examples/configs/instruct/polygraph_eval_mmlu_default_instruct.yaml +++ b/examples/configs/instruct/polygraph_eval_mmlu_default_instruct.yaml @@ -4,6 +4,7 @@ hydra: defaults: - model: mistral-7b-instruct-v0.2 + - similarity_model: default_en - _self_ cache_path: ./workdir/output @@ -37,7 +38,6 @@ ens_type: ignore_exceptions: false batch_size: 1 -deberta_batch_size: 10 seed: - 1 diff --git a/examples/configs/instruct/polygraph_eval_triviaqa_default_instruct.yaml b/examples/configs/instruct/polygraph_eval_triviaqa_default_instruct.yaml index ee4273a0a..cdf155139 100644 --- a/examples/configs/instruct/polygraph_eval_triviaqa_default_instruct.yaml +++ b/examples/configs/instruct/polygraph_eval_triviaqa_default_instruct.yaml @@ -4,6 +4,7 @@ hydra: defaults: - model: mistral-7b-instruct-v0.2 + - similarity_model: default_en - _self_ cache_path: ./workdir/output @@ -38,7 +39,6 @@ ens_type: ignore_exceptions: false batch_size: 2 -deberta_batch_size: 10 seed: - 1 diff --git a/examples/configs/instruct/similarity_model/default_en.yaml b/examples/configs/instruct/similarity_model/default_en.yaml new file mode 100644 index 000000000..2d71307c9 --- /dev/null +++ b/examples/configs/instruct/similarity_model/default_en.yaml @@ -0,0 +1,7 @@ +nli: + deberta_path: microsoft/deberta-large-mnli + batch_size: 10 +cross_encoder: + cross_encoder_name: cross-encoder/stsb-roberta-large + batch_size: 10 + diff --git a/examples/configs/instruct/similarity_model/default_multilingual.yaml b/examples/configs/instruct/similarity_model/default_multilingual.yaml new file mode 100644 index 000000000..277289de0 --- /dev/null +++ b/examples/configs/instruct/similarity_model/default_multilingual.yaml @@ -0,0 +1,7 @@ +nli: + deberta_path: MoritzLaurer/mDeBERTa-v3-base-xnli-multilingual-nli-2mil7 + batch_size: 10 +cross_encoder: + cross_encoder_name: cross-encoder/stsb-roberta-large + batch_size: 10 + diff --git a/examples/configs/polygraph_eval_aeslc.yaml b/examples/configs/polygraph_eval_aeslc.yaml index c90a5e011..89455aadc 100644 --- a/examples/configs/polygraph_eval_aeslc.yaml +++ b/examples/configs/polygraph_eval_aeslc.yaml @@ -6,6 +6,7 @@ defaults: - model: bloomz-560m - estimators: default_estimators - stat_calculators: default_calculators + - similarity_model: default_en - _self_ cache_path: ./workdir/output @@ -38,4 +39,3 @@ generation_metrics: null ignore_exceptions: false seed: - 1 - diff --git a/examples/configs/polygraph_eval_babiqa.yaml b/examples/configs/polygraph_eval_babiqa.yaml index b9a368c14..1db376f2f 100644 --- a/examples/configs/polygraph_eval_babiqa.yaml +++ b/examples/configs/polygraph_eval_babiqa.yaml @@ -6,6 +6,7 @@ defaults: - model: bloomz-560m - estimators: default_estimators - stat_calculators: default_calculators + - similarity_model: default_en - _self_ cache_path: ./workdir/output diff --git a/examples/configs/polygraph_eval_coqa.yaml b/examples/configs/polygraph_eval_coqa.yaml index e33401ed1..9ac58b144 100644 --- a/examples/configs/polygraph_eval_coqa.yaml +++ b/examples/configs/polygraph_eval_coqa.yaml @@ -6,6 +6,7 @@ defaults: - model: bloomz-560m - estimators: default_estimators - stat_calculators: default_calculators + - similarity_model: default_en - base_processing_coqa - _self_ diff --git a/examples/configs/polygraph_eval_coqa_simple_instruct.yaml b/examples/configs/polygraph_eval_coqa_simple_instruct.yaml index 93010974d..038280f80 100644 --- a/examples/configs/polygraph_eval_coqa_simple_instruct.yaml +++ b/examples/configs/polygraph_eval_coqa_simple_instruct.yaml @@ -6,6 +6,7 @@ defaults: - model: bloomz-560m - estimators: default_estimators - stat_calculators: default_calculators + - similarity_model: default_en - base_processing_coqa - _self_ diff --git a/examples/configs/polygraph_eval_gsm8k.yaml b/examples/configs/polygraph_eval_gsm8k.yaml index 57954e0fa..5047e5bc9 100644 --- a/examples/configs/polygraph_eval_gsm8k.yaml +++ b/examples/configs/polygraph_eval_gsm8k.yaml @@ -6,6 +6,7 @@ defaults: - model: bloomz-560m - estimators: default_estimators - stat_calculators: default_calculators + - similarity_model: default_en - _self_ cache_path: ./workdir/output diff --git a/examples/configs/polygraph_eval_gsm8k_simple_instruct.yaml b/examples/configs/polygraph_eval_gsm8k_simple_instruct.yaml index 91ead5514..e35a57d6d 100644 --- a/examples/configs/polygraph_eval_gsm8k_simple_instruct.yaml +++ b/examples/configs/polygraph_eval_gsm8k_simple_instruct.yaml @@ -6,6 +6,7 @@ defaults: - model: bloomz-560m - estimators: default_estimators - stat_calculators: default_calculators + - similarity_model: default_en - _self_ cache_path: ./workdir/output diff --git a/examples/configs/polygraph_eval_mmlu.yaml b/examples/configs/polygraph_eval_mmlu.yaml index 92cbc9c36..818ee3d81 100644 --- a/examples/configs/polygraph_eval_mmlu.yaml +++ b/examples/configs/polygraph_eval_mmlu.yaml @@ -6,6 +6,7 @@ defaults: - model: bloomz-560m - estimators: default_estimators - stat_calculators: default_calculators + - similarity_model: default_en - _self_ cache_path: ./workdir/output diff --git a/examples/configs/polygraph_eval_mmlu_simple_instruct.yaml b/examples/configs/polygraph_eval_mmlu_simple_instruct.yaml index 256f00c06..9f9e6a203 100644 --- a/examples/configs/polygraph_eval_mmlu_simple_instruct.yaml +++ b/examples/configs/polygraph_eval_mmlu_simple_instruct.yaml @@ -6,6 +6,7 @@ defaults: - model: bloomz-560m - estimators: default_estimators - stat_calculators: default_calculators + - similarity_model: default_en - _self_ cache_path: ./workdir/output diff --git a/examples/configs/polygraph_eval_person_bio_ar.yaml b/examples/configs/polygraph_eval_person_bio_ar.yaml index 813c39b00..53e7b6afd 100644 --- a/examples/configs/polygraph_eval_person_bio_ar.yaml +++ b/examples/configs/polygraph_eval_person_bio_ar.yaml @@ -5,6 +5,7 @@ hydra: defaults: - model: bloomz-560m - estimators: default_claim_estimators + - similarity_model: default_multilingual - _self_ cache_path: ./workdir/output diff --git a/examples/configs/polygraph_eval_person_bio_en_jais.yaml b/examples/configs/polygraph_eval_person_bio_en_jais.yaml index e2782a7ee..50a1bc6e6 100644 --- a/examples/configs/polygraph_eval_person_bio_en_jais.yaml +++ b/examples/configs/polygraph_eval_person_bio_en_jais.yaml @@ -5,6 +5,7 @@ hydra: defaults: - model: bloomz-560m - estimators: default_claim_estimators + - similarity_model: default_en - _self_ cache_path: ./workdir/output diff --git a/examples/configs/polygraph_eval_person_bio_en_mistral.yaml b/examples/configs/polygraph_eval_person_bio_en_mistral.yaml index 5f999cb45..a1d642e04 100644 --- a/examples/configs/polygraph_eval_person_bio_en_mistral.yaml +++ b/examples/configs/polygraph_eval_person_bio_en_mistral.yaml @@ -5,6 +5,7 @@ hydra: defaults: - model: bloomz-560m - estimators: default_claim_estimators + - similarity_model: default_en - _self_ cache_path: ./workdir/output diff --git a/examples/configs/polygraph_eval_person_bio_ru.yaml b/examples/configs/polygraph_eval_person_bio_ru.yaml index ff06f2b55..9ea2f08cb 100644 --- a/examples/configs/polygraph_eval_person_bio_ru.yaml +++ b/examples/configs/polygraph_eval_person_bio_ru.yaml @@ -5,6 +5,7 @@ hydra: defaults: - model: bloomz-560m - estimators: default_claim_estimators + - similarity_model: default_multilingual - _self_ cache_path: ./workdir/output diff --git a/examples/configs/polygraph_eval_person_bio_zh.yaml b/examples/configs/polygraph_eval_person_bio_zh.yaml index a4e8d750d..bbc018eee 100644 --- a/examples/configs/polygraph_eval_person_bio_zh.yaml +++ b/examples/configs/polygraph_eval_person_bio_zh.yaml @@ -5,6 +5,7 @@ hydra: defaults: - model: bloomz-560m - estimators: default_claim_estimators + - similarity_model: default_multilingual - _self_ cache_path: ./workdir/output diff --git a/examples/configs/polygraph_eval_samsum.yaml b/examples/configs/polygraph_eval_samsum.yaml index 9cc7240fa..3f3d36c1d 100644 --- a/examples/configs/polygraph_eval_samsum.yaml +++ b/examples/configs/polygraph_eval_samsum.yaml @@ -6,6 +6,7 @@ defaults: - model: bloomz-560m - estimators: default_estimators - stat_calculators: default_calculators + - similarity_model: default_en - _self_ cache_path: ./workdir/output diff --git a/examples/configs/polygraph_eval_samsum_simple_instruct.yaml b/examples/configs/polygraph_eval_samsum_simple_instruct.yaml index 89648827a..8f9d2e827 100644 --- a/examples/configs/polygraph_eval_samsum_simple_instruct.yaml +++ b/examples/configs/polygraph_eval_samsum_simple_instruct.yaml @@ -6,6 +6,7 @@ defaults: - model: bloomz-560m - estimators: default_estimators - stat_calculators: default_calculators + - similarity_model: default_en - _self_ cache_path: ./workdir/output diff --git a/examples/configs/polygraph_eval_triviaqa.yaml b/examples/configs/polygraph_eval_triviaqa.yaml index 0379659b2..63779143b 100644 --- a/examples/configs/polygraph_eval_triviaqa.yaml +++ b/examples/configs/polygraph_eval_triviaqa.yaml @@ -6,6 +6,7 @@ defaults: - model: bloomz-560m - estimators: default_estimators - stat_calculators: default_calculators + - similarity_model: default_en - base_processing_triviaqa - _self_ diff --git a/examples/configs/polygraph_eval_triviaqa_simple_instruct.yaml b/examples/configs/polygraph_eval_triviaqa_simple_instruct.yaml index f712a6a63..230cb0acf 100644 --- a/examples/configs/polygraph_eval_triviaqa_simple_instruct.yaml +++ b/examples/configs/polygraph_eval_triviaqa_simple_instruct.yaml @@ -6,6 +6,7 @@ defaults: - model: bloomz-560m - estimators: default_estimators - stat_calculators: default_calculators + - similarity_model: default_en - base_processing_triviaqa - _self_ diff --git a/examples/configs/polygraph_eval_truthfulqa.yaml b/examples/configs/polygraph_eval_truthfulqa.yaml index 964c41b29..deeb9de28 100644 --- a/examples/configs/polygraph_eval_truthfulqa.yaml +++ b/examples/configs/polygraph_eval_truthfulqa.yaml @@ -6,6 +6,7 @@ defaults: - model: bloomz-560m - estimators: default_estimators - stat_calculators: default_calculators + - similarity_model: default_en - _self_ cache_path: ./workdir/output diff --git a/examples/configs/polygraph_eval_truthfulqa_simple_instruct.yaml b/examples/configs/polygraph_eval_truthfulqa_simple_instruct.yaml index 60f72eb05..94ba5f001 100644 --- a/examples/configs/polygraph_eval_truthfulqa_simple_instruct.yaml +++ b/examples/configs/polygraph_eval_truthfulqa_simple_instruct.yaml @@ -6,6 +6,7 @@ defaults: - model: bloomz-560m - estimators: default_estimators - stat_calculators: default_calculators + - similarity_model: default_en - _self_ cache_path: ./workdir/output diff --git a/examples/configs/polygraph_eval_wmt14_deen.yaml b/examples/configs/polygraph_eval_wmt14_deen.yaml index 738ac9ca9..38ebf8909 100644 --- a/examples/configs/polygraph_eval_wmt14_deen.yaml +++ b/examples/configs/polygraph_eval_wmt14_deen.yaml @@ -6,6 +6,7 @@ defaults: - model: bloomz-560m - estimators: default_estimators - stat_calculators: default_calculators + - similarity_model: default_en - _self_ cache_path: ./workdir/output diff --git a/examples/configs/polygraph_eval_wmt14_fren.yaml b/examples/configs/polygraph_eval_wmt14_fren.yaml index c5a171fee..485b93e6e 100644 --- a/examples/configs/polygraph_eval_wmt14_fren.yaml +++ b/examples/configs/polygraph_eval_wmt14_fren.yaml @@ -6,6 +6,7 @@ defaults: - model: bloomz-560m - estimators: default_estimators - stat_calculators: default_calculators + - similarity_model: default_en - _self_ cache_path: ./workdir/output diff --git a/examples/configs/polygraph_eval_wmt14_fren_simple_instruct.yaml b/examples/configs/polygraph_eval_wmt14_fren_simple_instruct.yaml index 76adb4746..251b78279 100644 --- a/examples/configs/polygraph_eval_wmt14_fren_simple_instruct.yaml +++ b/examples/configs/polygraph_eval_wmt14_fren_simple_instruct.yaml @@ -6,6 +6,7 @@ defaults: - model: bloomz-560m - estimators: default_estimators - stat_calculators: default_calculators + - similarity_model: default_en - _self_ cache_path: ./workdir/output diff --git a/examples/configs/polygraph_eval_wmt19_deen.yaml b/examples/configs/polygraph_eval_wmt19_deen.yaml index 71f06ea02..1f08c2859 100644 --- a/examples/configs/polygraph_eval_wmt19_deen.yaml +++ b/examples/configs/polygraph_eval_wmt19_deen.yaml @@ -6,6 +6,7 @@ defaults: - model: bloomz-560m - estimators: default_estimators - stat_calculators: default_calculators + - similarity_model: default_en - _self_ cache_path: ./workdir/output diff --git a/examples/configs/polygraph_eval_wmt19_deen_simple_instruct.yaml b/examples/configs/polygraph_eval_wmt19_deen_simple_instruct.yaml index 6f43d4b6f..c2d636e94 100644 --- a/examples/configs/polygraph_eval_wmt19_deen_simple_instruct.yaml +++ b/examples/configs/polygraph_eval_wmt19_deen_simple_instruct.yaml @@ -6,6 +6,7 @@ defaults: - model: bloomz-560m - estimators: default_estimators - stat_calculators: default_calculators + - similarity_model: default_en - _self_ cache_path: ./workdir/output diff --git a/examples/configs/polygraph_eval_wmt19_ruen.yaml b/examples/configs/polygraph_eval_wmt19_ruen.yaml index 952f952bb..76bf72080 100644 --- a/examples/configs/polygraph_eval_wmt19_ruen.yaml +++ b/examples/configs/polygraph_eval_wmt19_ruen.yaml @@ -6,6 +6,7 @@ defaults: - model: bloomz-560m - estimators: default_estimators - stat_calculators: default_calculators + - similarity_model: default_en - _self_ cache_path: ./workdir/output diff --git a/examples/configs/polygraph_eval_wmt19_ruen_simple_instruct.yaml b/examples/configs/polygraph_eval_wmt19_ruen_simple_instruct.yaml index 6b1d13912..2cb5fb08d 100644 --- a/examples/configs/polygraph_eval_wmt19_ruen_simple_instruct.yaml +++ b/examples/configs/polygraph_eval_wmt19_ruen_simple_instruct.yaml @@ -6,6 +6,7 @@ defaults: - model: bloomz-560m - estimators: default_estimators - stat_calculators: default_calculators + - similarity_model: default_en - _self_ cache_path: ./workdir/output diff --git a/examples/configs/polygraph_eval_xsum.yaml b/examples/configs/polygraph_eval_xsum.yaml index bf4c9676f..572f4883b 100644 --- a/examples/configs/polygraph_eval_xsum.yaml +++ b/examples/configs/polygraph_eval_xsum.yaml @@ -6,6 +6,7 @@ defaults: - model: bloomz-560m - estimators: default_estimators - stat_calculators: default_calculators + - similarity_model: default_en - _self_ cache_path: ./workdir/output diff --git a/examples/configs/polygraph_eval_xsum_simple_instruct.yaml b/examples/configs/polygraph_eval_xsum_simple_instruct.yaml index 9364f8d5c..332251efb 100644 --- a/examples/configs/polygraph_eval_xsum_simple_instruct.yaml +++ b/examples/configs/polygraph_eval_xsum_simple_instruct.yaml @@ -6,6 +6,7 @@ defaults: - model: bloomz-560m - estimators: default_estimators - stat_calculators: default_calculators + - similarity_model: default_en - _self_ cache_path: ./workdir/output diff --git a/examples/configs/similarity_model/default_en.yaml b/examples/configs/similarity_model/default_en.yaml new file mode 100644 index 000000000..2d71307c9 --- /dev/null +++ b/examples/configs/similarity_model/default_en.yaml @@ -0,0 +1,7 @@ +nli: + deberta_path: microsoft/deberta-large-mnli + batch_size: 10 +cross_encoder: + cross_encoder_name: cross-encoder/stsb-roberta-large + batch_size: 10 + diff --git a/examples/configs/similarity_model/default_multilingual.yaml b/examples/configs/similarity_model/default_multilingual.yaml new file mode 100644 index 000000000..277289de0 --- /dev/null +++ b/examples/configs/similarity_model/default_multilingual.yaml @@ -0,0 +1,7 @@ +nli: + deberta_path: MoritzLaurer/mDeBERTa-v3-base-xnli-multilingual-nli-2mil7 + batch_size: 10 +cross_encoder: + cross_encoder_name: cross-encoder/stsb-roberta-large + batch_size: 10 + diff --git a/scripts/polygraph_eval b/scripts/polygraph_eval index 8163e9ca0..3ccc3388c 100755 --- a/scripts/polygraph_eval +++ b/scripts/polygraph_eval @@ -176,22 +176,32 @@ def get_stat_calculator_names(config): language = getattr(config, "language", "en") output_attentions = getattr(config, "output_attentions", True) and (getattr(config.model, "type", "Whitebox") != "vLLMCausalLM") output_hidden_states = False if getattr(config.model, "type", "Whitebox") == "vLLMCausalLM" else True - hf_cache = getattr(config, "hf_cache", None) - deberta_batch_size = getattr(config, "deberta_batch_size", 10) blackbox_supports_logprobs = model_type == "Blackbox" and getattr( config.model, "supports_logprobs", False ) + sim_cfg = getattr(config, "similarity_model", {}) + + nli_cfg = getattr(sim_cfg, "nli", {}) + nli_deberta_path = nli_cfg.get("deberta_path", None) + nli_batch_size = nli_cfg.get("batch_size", None) + + cross_encoder_cfg = getattr(sim_cfg, "cross_encoder", {}) + cross_encoder_name = cross_encoder_cfg.get("cross_encoder_name", None) + cross_encoder_batch_size = cross_encoder_cfg.get("batch_size", None) + all_stat_calculators = [] if "auto" in config.stat_calculators: all_stat_calculators += register_default_stat_calculators( model_type, language, - hf_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, blackbox_supports_logprobs=blackbox_supports_logprobs, - deberta_batch_size=deberta_batch_size, + nli_deberta_path=nli_deberta_path, + nli_batch_size=nli_batch_size, + cross_encoder_name=cross_encoder_name, + cross_encoder_batch_size=cross_encoder_batch_size, ) for stat_calculator in config.stat_calculators: @@ -306,7 +316,6 @@ def get_model(args): return get_vllm_model(args) else: cache_kwargs = { - "cache_dir": getattr(args, "hf_cache", None), "token": getattr(args, "hf_token", None), } return get_whitebox_model(args, cache_kwargs) diff --git a/src/lm_polygraph/defaults/register_default_stat_calculators.py b/src/lm_polygraph/defaults/register_default_stat_calculators.py index b7fa94fc7..1a582790c 100644 --- a/src/lm_polygraph/defaults/register_default_stat_calculators.py +++ b/src/lm_polygraph/defaults/register_default_stat_calculators.py @@ -11,11 +11,13 @@ def register_default_stat_calculators( model_type: str, language: str = "en", - hf_cache: Optional[str] = None, blackbox_supports_logprobs: bool = False, output_attentions: bool = True, output_hidden_states: bool = True, - deberta_batch_size: int = 10, + nli_deberta_path: Optional[str] = None, + nli_batch_size: Optional[int] = None, + cross_encoder_name: Optional[str] = None, + cross_encoder_batch_size: Optional[int] = None, ) -> List[StatCalculatorContainer]: """ Specifies the list of the default stat_calculators that could be used in the evaluation scripts and @@ -43,16 +45,10 @@ def _register( ) all_stat_calculators.append(sc) - if language == "en": - deberta_model_path = "microsoft/deberta-large-mnli" - else: - deberta_model_path = "MoritzLaurer/mDeBERTa-v3-base-xnli-multilingual-nli-2mil7" - # Shared NLI model config nli_model_cfg = { - "deberta_path": deberta_model_path, - "hf_cache": hf_cache, - "batch_size": deberta_batch_size, + "deberta_path": nli_deberta_path, + "batch_size": nli_batch_size, "device": None, } @@ -111,16 +107,16 @@ def _register( CrossEncoderSimilarityMatrixCalculator, "lm_polygraph.defaults.stat_calculator_builders.default_CrossEncoderSimilarityMatrixCalculator", { - "batch_size": deberta_batch_size, - "cross_encoder_name": "cross-encoder/stsb-roberta-large", + "batch_size": cross_encoder_batch_size, + "cross_encoder_name": cross_encoder_name, }, ) _register( GreedyCrossEncoderSimilarityMatrixCalculator, "lm_polygraph.defaults.stat_calculator_builders.default_GreedyCrossEncoderSimilarityMatrixCalculator", { - "batch_size": 10, - "cross_encoder_name": "cross-encoder/stsb-roberta-large", + "batch_size": cross_encoder_batch_size, + "cross_encoder_name": cross_encoder_name, }, ) _register( @@ -170,8 +166,8 @@ def _register( CrossEncoderSimilarityMatrixVisualCalculator, "lm_polygraph.defaults.stat_calculator_builders.default_CrossEncoderSimilarityMatrixVisualCalculator", { - "batch_size": deberta_batch_size, - "cross_encoder_name": "cross-encoder/stsb-roberta-large", + "batch_size": cross_encoder_batch_size, + "cross_encoder_name": cross_encoder_name, }, ) _register( diff --git a/src/lm_polygraph/defaults/stat_calculator_builders/utils.py b/src/lm_polygraph/defaults/stat_calculator_builders/utils.py index bff96354e..f82d4b8ee 100644 --- a/src/lm_polygraph/defaults/stat_calculator_builders/utils.py +++ b/src/lm_polygraph/defaults/stat_calculator_builders/utils.py @@ -9,14 +9,11 @@ def load_nli_model( deberta_path="microsoft/deberta-large-mnli", batch_size=10, device=None, - hf_cache: str = None, ): if deberta_path.startswith("microsoft"): - nli_model = Deberta(deberta_path, batch_size, device=device, hf_cache=hf_cache) + nli_model = Deberta(deberta_path, batch_size, device=device) else: - nli_model = MultilingualDeberta( - deberta_path, batch_size, device=device, hf_cache=hf_cache - ) + nli_model = MultilingualDeberta(deberta_path, batch_size, device=device) log.info( f"Initialized {nli_model.deberta_path} on {nli_model.device} with batch_size={nli_model.batch_size}" ) diff --git a/src/lm_polygraph/utils/deberta.py b/src/lm_polygraph/utils/deberta.py index e6e4ac681..572be6e95 100644 --- a/src/lm_polygraph/utils/deberta.py +++ b/src/lm_polygraph/utils/deberta.py @@ -2,6 +2,8 @@ from transformers import ( DebertaForSequenceClassification, DebertaTokenizer, + DebertaV2ForSequenceClassification, + DebertaV2Tokenizer, AutoTokenizer, AutoModelForSequenceClassification, ) @@ -18,7 +20,6 @@ def __init__( deberta_path: str = "microsoft/deberta-large-mnli", batch_size: int = 10, device: str = None, - hf_cache: str = None, ): """ Parameters @@ -36,7 +37,6 @@ def __init__( self.device = "cuda:0" if torch.cuda.is_available() else "cpu" else: self.device = device - self.hf_cache = hf_cache self.setup() @property @@ -64,15 +64,26 @@ def setup(self): """ if self._deberta is not None: return - self._deberta = DebertaForSequenceClassification.from_pretrained( - self.deberta_path, - problem_type="multi_label_classification", - cache_dir=self.hf_cache, - ) - self._deberta_tokenizer = DebertaTokenizer.from_pretrained( - self.deberta_path, cache_dir=self.hf_cache - ) - self._deberta.to(self.device) + + if "v2" in self.deberta_path or "v3" in self.deberta_path: + self._deberta = DebertaV2ForSequenceClassification.from_pretrained( + self.deberta_path, + problem_type="multi_label_classification", + device_map="auto", + ) + self._deberta_tokenizer = DebertaV2Tokenizer.from_pretrained( + self.deberta_path + ) + else: + self._deberta = DebertaForSequenceClassification.from_pretrained( + self.deberta_path, + problem_type="multi_label_classification", + ) + self._deberta_tokenizer = DebertaTokenizer.from_pretrained( + self.deberta_path + ) + self._deberta.to(self.device) + self._deberta.eval() @@ -87,7 +98,6 @@ def __init__( deberta_path: str = "MoritzLaurer/mDeBERTa-v3-base-xnli-multilingual-nli-2mil7", batch_size: int = 10, device: str = None, - hf_cache: str = None, ): """ Parameters @@ -106,7 +116,6 @@ def __init__( self.device = "cuda:0" if torch.cuda.is_available() else "cpu" else: self.device = device - self.hf_cache = hf_cache self.setup() def setup(self): @@ -115,12 +124,9 @@ def setup(self): """ if self._deberta is not None: return - self._deberta_tokenizer = AutoTokenizer.from_pretrained( - self.deberta_path, cache_dir=self.hf_cache - ) + self._deberta_tokenizer = AutoTokenizer.from_pretrained(self.deberta_path) self._deberta = AutoModelForSequenceClassification.from_pretrained( self.deberta_path, - cache_dir=self.hf_cache, ) self._deberta.to(self.device) self._deberta.eval() diff --git a/src/lm_polygraph/utils/estimate_uncertainty.py b/src/lm_polygraph/utils/estimate_uncertainty.py index 8fac44bdf..44bcfc7d4 100644 --- a/src/lm_polygraph/utils/estimate_uncertainty.py +++ b/src/lm_polygraph/utils/estimate_uncertainty.py @@ -1,4 +1,4 @@ -from typing import List, Union +from typing import List, Union, Optional from dataclasses import dataclass from lm_polygraph.utils.model import Model, WhiteboxModel @@ -35,7 +35,13 @@ class UncertaintyOutput: def estimate_uncertainty( - model: Model, estimator: Estimator, input_text: str + model: Model, + estimator: Estimator, + input_text: str, + nli_deberta_path: Optional[str] = "microsoft/deberta-large-mnli", + nli_batch_size: Optional[int] = 10, + cross_encoder_name: Optional[str] = "cross-encoder/stsb-roberta-large", + cross_encoder_batch_size: Optional[int] = 10, ) -> UncertaintyOutput: """ Estimated uncertainty of the model generation using the provided esitmator. @@ -87,7 +93,11 @@ def estimate_uncertainty( model, [estimator], available_stat_calculators=register_default_stat_calculators( - model_type + model_type, + nli_deberta_path=nli_deberta_path, + nli_batch_size=nli_batch_size, + cross_encoder_name=cross_encoder_name, + cross_encoder_batch_size=cross_encoder_batch_size, ), # TODO: builder_env_stat_calc=BuilderEnvironmentStatCalculator(model), generation_metrics=[], diff --git a/test/configs/similarity_model/default_en.yaml b/test/configs/similarity_model/default_en.yaml new file mode 100644 index 000000000..d6e5d8737 --- /dev/null +++ b/test/configs/similarity_model/default_en.yaml @@ -0,0 +1,6 @@ +nli: + deberta_path: microsoft/deberta-large-mnli + batch_size: 1 +cross_encoder: + cross_encoder_name: cross-encoder/stsb-roberta-large + batch_size: 1 diff --git a/test/configs/test_polygraph_eval.yaml b/test/configs/test_polygraph_eval.yaml index 2ad115bfc..2b7e054b4 100644 --- a/test/configs/test_polygraph_eval.yaml +++ b/test/configs/test_polygraph_eval.yaml @@ -4,6 +4,7 @@ hydra: defaults: - model: bloomz-560m + - similarity_model: default_en - _self_ cache_path: ./workdir/output diff --git a/test/configs/test_polygraph_eval_blackbox.yaml b/test/configs/test_polygraph_eval_blackbox.yaml index 62f2f4777..07ea42e7b 100644 --- a/test/configs/test_polygraph_eval_blackbox.yaml +++ b/test/configs/test_polygraph_eval_blackbox.yaml @@ -6,6 +6,7 @@ defaults: - model: gpt-4o-mini - test_default_blackbox_estimators - test_default_calculators + - similarity_model: default_en - _self_ @@ -44,7 +45,6 @@ ens_type: ignore_exceptions: false batch_size: 1 -deberta_batch_size: 1 seed: - 1 diff --git a/test/configs/test_polygraph_eval_greybox.yaml b/test/configs/test_polygraph_eval_greybox.yaml index ed1f46dbf..2eb4dabb5 100644 --- a/test/configs/test_polygraph_eval_greybox.yaml +++ b/test/configs/test_polygraph_eval_greybox.yaml @@ -6,7 +6,7 @@ defaults: - model: gpt-4o-mini-greybox - test_default_greybox_estimators - test_default_calculators - + - similarity_model: default_en - _self_ cache_path: ./workdir/output @@ -44,7 +44,6 @@ ens_type: ignore_exceptions: false batch_size: 1 -deberta_batch_size: 1 seed: - 1 diff --git a/test/configs/test_polygraph_eval_seq_ue.yaml b/test/configs/test_polygraph_eval_seq_ue.yaml index b259f071f..6991819f9 100644 --- a/test/configs/test_polygraph_eval_seq_ue.yaml +++ b/test/configs/test_polygraph_eval_seq_ue.yaml @@ -5,6 +5,7 @@ hydra: defaults: - model: bloomz-560m - test_default_estimators + - similarity_model: default_en - _self_ cache_path: ./workdir/output diff --git a/test/configs/test_polygraph_eval_vllm.yaml b/test/configs/test_polygraph_eval_vllm.yaml index 44cc38893..848566681 100644 --- a/test/configs/test_polygraph_eval_vllm.yaml +++ b/test/configs/test_polygraph_eval_vllm.yaml @@ -5,6 +5,7 @@ hydra: defaults: - model: vllm - test_default_estimators_vllm + - similarity_model: default_en - _self_ cache_path: ./workdir/output