Parametrize similarity

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

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

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

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
+

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
+

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
-

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

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_
 

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_
 

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

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

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

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

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

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

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

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

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

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

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

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_
 

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_
 

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

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

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

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

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

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

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

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

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

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

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

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
+

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
+

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)