Move model wrapper definitions to separate modules

README.md

@@ -40,7 +40,7 @@ pip install lm-polygraph
 1. Initialize the base model (encoder-decoder or decoder-only) and tokenizer from HuggingFace or a local file, and use them to initialize the WhiteboxModel for evaluation:
 ```python
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from lm_polygraph.utils.model import WhiteboxModel
+from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel
 
 model_path = "Qwen/Qwen2.5-0.5B-Instruct"
 base_model = AutoModelForCausalLM.from_pretrained(model_path, device_map="cuda:0")
@@ -77,10 +77,9 @@ LM-Polygraph can work with any OpenAI-compatible API services:
 from lm_polygraph import BlackboxModel
 from lm_polygraph.estimators import Perplexity, MaximumSequenceProbability
 
-model = BlackboxModel.from_openai(
-    openai_api_key='YOUR_API_KEY',
+model = BlackboxModel(
     model_path='gpt-4o',
-    supports_logprobs=True  # Enable for deployments 
+    api_provider_name='openai'
 )
 
 ue_method = Perplexity()  # or DetMat(), MeanTokenEntropy(), EigValLaplacian(), etc.

docs/usage.rst

@@ -62,7 +62,7 @@ Uncertainty for single input
     .. code-block:: python
 
         from transformers import AutoModelForCausalLM, AutoTokenizer
-        from lm_polygraph.utils.model import WhiteboxModel
+        from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel
 
         model_path = "bigscience/bloomz-560m"
         base_model = AutoModelForCausalLM.from_pretrained(model_path, device_map="cuda:0")
@@ -74,7 +74,7 @@ Uncertainty for single input
 
     .. code-block:: python
 
-        from lm_polygraph.utils.model import WhiteboxModel
+        from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel
 
         model = WhiteboxModel.from_pretrained(
             "bigscience/bloomz-3b",

examples/basic_example.ipynb

@@ -61,7 +61,8 @@
     "%autoreload 2\n",
     "\n",
     "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
-    "from lm_polygraph.utils.model import WhiteboxModel, BlackboxModel\n",
+    "from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel\n",
+    "from lm_polygraph.model_adapters.blackbox_model import BlackboxModel\n",
     "from lm_polygraph import estimate_uncertainty\n",
     "from lm_polygraph.estimators import MaximumTokenProbability, MaximumSequenceProbability, SemanticEntropy, EigValLaplacian\n",
     "from lm_polygraph.utils.generation_parameters import GenerationParameters"
@@ -381,7 +382,10 @@
     "\n",
     "HUGGINGFACE_API_TOKEN = '<Your HuggingFace API token>'\n",
     "MODEL_ID = 'meta-llama/Llama-3.3-70B-Instruct'\n",
-    "model = BlackboxModel.from_huggingface(hf_api_token=HUGGINGFACE_API_TOKEN, hf_model_id=MODEL_ID, openai_api_key=None, openai_model_path=None)"
+    "model = BlackboxModel(\n",
+    "    model_path=MODEL_ID,\n",
+    "    api_provider_name='huggingface'\n",
+    ")\n"
    ]
   },
   {

examples/configs/estimators/default_claim_estimators.yaml

@@ -25,4 +25,4 @@
     idf_dataset_size: -1
     spacy_path: "en_core_web_sm"
 - name: FrequencyScoringClaim
-- name: TokenSARClaim
+- name: TokenSARClaim

examples/configs/estimators/default_estimators_blackbox.yaml

@@ -0,0 +1,16 @@
+- name: LexicalSimilarity
+  cfg:
+    metric: "rougeL"
+- name: NumSemSets
+- name: EigValLaplacian
+  cfg:
+    similarity_score: "NLI_score"
+    affinity: "entail"
+- name: DegMat
+  cfg:
+    similarity_score: "NLI_score"
+    affinity: "entail"
+- name: Eccentricity
+  cfg:
+    similarity_score: "NLI_score"
+    affinity: "entail"

examples/configs/estimators/default_estimators_greybox.yaml

@@ -0,0 +1,27 @@
+- name: MaximumSequenceProbability
+- name: Perplexity
+- name: MeanTokenEntropy
+- name: MonteCarloSequenceEntropy
+- name: MonteCarloNormalizedSequenceEntropy
+- name: LexicalSimilarity
+  cfg:
+    metric: "rougeL"
+- name: NumSemSets
+- name: EigValLaplacian
+  cfg:
+    similarity_score: "NLI_score"
+    affinity: "entail"
+- name: DegMat
+  cfg:
+    similarity_score: "NLI_score"
+    affinity: "entail"
+- name: Eccentricity
+  cfg:
+    similarity_score: "NLI_score"
+    affinity: "entail"
+- name: SemanticEntropy
+- name: SentenceSAR
+- name: CocoaMSP
+- name: CocoaPPL
+- name: CocoaMTE
+- name: SemanticDensity

examples/configs/model/gpt_4.1.yaml

@@ -0,0 +1,3 @@
+path: gpt-4.1
+type: Blackbox
+provider: openai

examples/configs/model/llama_huggingface.yaml

@@ -0,0 +1,3 @@
+path: meta-llama/Meta-Llama-3.1-8B-Instruct
+type: Blackbox
+provider: huggingface

examples/configs/model/llama_together.yaml

@@ -0,0 +1,3 @@
+path: meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo
+type: Blackbox
+provider: together_ai

examples/other/ats_example.ipynb

@@ -9,7 +9,7 @@
    "source": [
     "from transformers import AutoModelForSeq2SeqLM, AutoTokenizer\n",
     "from lm_polygraph.estimators import *\n",
-    "from lm_polygraph.utils.model import WhiteboxModel\n",
+    "from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel\n",
     "from lm_polygraph.utils.dataset import Dataset\n",
     "from lm_polygraph.utils.processor import Logger\n",
     "from lm_polygraph.utils.manager import UEManager\n",

examples/other/claim_level_example_ar.ipynb

@@ -135,7 +135,7 @@
    "source": [
     "import os\n",
     "\n",
-    "from lm_polygraph.utils.model import WhiteboxModel\n",
+    "from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel\n",
     "from lm_polygraph.estimators import *\n",
     "from lm_polygraph.stat_calculators import *\n",
     "from lm_polygraph.utils.openai_chat import OpenAIChat\n",

examples/other/claim_level_example_en.ipynb

@@ -9,7 +9,7 @@
    "source": [
     "import os\n",
     "\n",
-    "from lm_polygraph.utils.model import WhiteboxModel\n",
+    "from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel\n",
     "from lm_polygraph.estimators import MaximumClaimProbability, ClaimConditionedProbabilityClaim\n",
     "from lm_polygraph.stat_calculators import *\n",
     "from lm_polygraph.utils.openai_chat import OpenAIChat\n",

examples/other/claim_level_example_ru.ipynb

@@ -136,7 +136,7 @@
    "source": [
     "import os\n",
     "\n",
-    "from lm_polygraph.utils.model import WhiteboxModel\n",
+    "from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel\n",
     "from lm_polygraph.estimators import *\n",
     "from lm_polygraph.stat_calculators import *\n",
     "from lm_polygraph.utils.openai_chat import OpenAIChat\n",

examples/other/claim_level_example_zh.ipynb

@@ -22,7 +22,7 @@
     "import sys\n",
     "sys.path.append('../src')\n",
     "\n",
-    "from lm_polygraph.utils.model import WhiteboxModel\n",
+    "from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel\n",
     "from lm_polygraph.estimators import *\n",
     "from lm_polygraph.stat_calculators import *\n",
     "from lm_polygraph.utils.openai_chat import OpenAIChat\n",

examples/other/mt_example.ipynb

@@ -14,7 +14,7 @@
     "from datasets import load_dataset\n",
     "from transformers import FSMTForConditionalGeneration, FSMTTokenizer\n",
     "from lm_polygraph.estimators import *\n",
-    "from lm_polygraph.utils.model import WhiteboxModel\n",
+    "from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel\n",
     "from lm_polygraph.utils.dataset import Dataset\n",
     "from lm_polygraph.utils.processor import Logger\n",
     "from lm_polygraph.utils.manager import UEManager\n",

examples/other/qa_example.ipynb

@@ -13,7 +13,7 @@
     "\n",
     "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
     "from lm_polygraph.estimators import *\n",
-    "from lm_polygraph.utils.model import WhiteboxModel\n",
+    "from lm_polygraph.model_adapters.whitebox_model import WhiteboxModel\n",
     "from lm_polygraph.utils.dataset import Dataset\n",
     "from lm_polygraph.utils.processor import Logger\n",
     "from lm_polygraph.utils.manager import UEManager\n",

requirements.txt

@@ -35,3 +35,4 @@ evaluate>=0.4.2
 spacy>=3.4.0,<3.8.0
 fastchat
 diskcache>=5.6.3
+together>=1.5

scripts/polygraph_eval

@@ -16,8 +16,7 @@ logging.getLogger("httpx").setLevel(logging.WARNING)
 
 from lm_polygraph.utils.manager import UEManager
 from lm_polygraph.utils.dataset import Dataset
-from lm_polygraph.utils.model import WhiteboxModel, BlackboxModel
-from lm_polygraph.model_adapters import WhiteboxModelvLLM
+from lm_polygraph.model_adapters import *
 from lm_polygraph.utils.processor import Logger
 from lm_polygraph.generation_metrics import *
 from lm_polygraph.estimators import *
@@ -120,7 +119,7 @@ def main(args):
         ue_metrics = get_ue_metrics(args)
 
         builder_env_stat_calc = BuilderEnvironmentStatCalculator(model=model)
-        available_stat_calculators = get_stat_calculator_names(args)
+        available_stat_calculators = get_stat_calculator_names(args, model)
 
         man = UEManager(
             data=dataset,
@@ -171,16 +170,13 @@ def get_ue_metrics(args):
     return ue_metrics
 
 
-def get_stat_calculator_names(config):
+def get_stat_calculator_names(config, model):
     model_type = "Whitebox" if getattr(config.model, "type", "Whitebox") != "Blackbox" else "Blackbox"
     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
-    )
 
     all_stat_calculators = []
     if "auto" in config.stat_calculators:
@@ -190,8 +186,8 @@ def get_stat_calculator_names(config):
             hf_cache,
             output_attentions=output_attentions, 
             output_hidden_states=output_hidden_states,
-            blackbox_supports_logprobs=blackbox_supports_logprobs,
             deberta_batch_size=deberta_batch_size,
+            model=model,
         )
 
     for stat_calculator in config.stat_calculators:
@@ -231,7 +227,7 @@ def get_generation_metrics(args):
             RougeMetric("rouge2"),
             RougeMetric("rougeL"),
             BLEUMetric(),
-            BertScoreMetric("rh"),
+            BertScoreMetric(),
             SbertMetric(),
             AccuracyMetric(
                 target_ignore_regex=getattr(args, "target_ignore_regex", None),
@@ -316,32 +312,17 @@ def get_blackbox_model(args):
     provider = getattr(args.model, "provider", "")
     if provider is None or provider.strip() == "":
         raise ValueError(
-            "Blackbox model provider cannot be empty or None. Please specify a valid provider."
-        )
-
-    supports_logprobs = getattr(args.model, "supports_logprobs", False)
-
-    if provider == "openai":
-        openai_api_key = os.environ.get("OPENAI_API_KEY")
-        if openai_api_key is None:
-            raise ValueError("OpenAI API key is not set in the environment variables.")
-        return BlackboxModel.from_openai(
-            openai_api_key=openai_api_key,
-            model_path=args.model.path,
-            supports_logprobs=supports_logprobs,
+            "Blackbox model provider cannot be empty or None. Please specify a valid provider (e.g. 'openai', 'huggingface', 'together_ai')."
         )
-    elif provider == "huggingface":
-        hf_api_key = os.environ.get("HUGGINGFACE_API_KEY")
-        if hf_api_key is None:
-            raise ValueError(
-                "HuggingFace API key is not set in the environment variables."
-            )
-        return BlackboxModel.from_huggingface(
-            hf_api_token=hf_api_key, hf_model_id=args.model.path
-        )
-    else:
+    elif provider not in ["openai", "huggingface", "together_ai"]:
         raise ValueError(f"Unsupported black-box model provider: {provider}")
 
+    return BlackboxModel(
+        model_path=args.model.path,
+        api_provider_name=provider,
+        supports_logprobs=getattr(args.model, "supports_logprobs", None),
+    )
+
 
 def get_whitebox_model(args, cache_kwargs={}):
     if not "path_to_load_script" in args.model or not args.model.path_to_load_script:

src/lm_polygraph/__init__.py

@@ -1,4 +1,8 @@
-from .utils.model import WhiteboxModel, BlackboxModel
+from .model_adapters.whitebox_model import WhiteboxModel
+from .model_adapters.blackbox_model import BlackboxModel
 from .utils.manager import UEManager
 from .utils.estimate_uncertainty import estimate_uncertainty
 from .utils.dataset import Dataset
+
+# Import model adapters to ensure API provider adapters are registered
+from . import model_adapters